Function test_packer_masking_randomly

tests/test_sequence_packer.py:251–354

()

def test_packer_masking_randomly():
    input_batch_size = 4
    xds = [{"input_ids": seq} for seq in xs]
    input_batches = list(batched(xds, input_batch_size))
    d = {
        "src_iterable": input_batches,
        "src_batch_size": input_batch_size,
        "src_max_seq_len": max_seq_len,
        "out_batch_size": 5,
        "out_pseq_len": 15,
        "buffer_size": 5,
        "pad_token_id": -1,
        "mask_token_id": -2,
        "ignore_token_id": -3,
        "mask_prob": 0.5,  # <== extensive masking and word randomization
        "seed": 42,
        "suppress_masking": False,
    }

    out_batches = list(GreedyBestFitSequencePacker(**d))

    out_expected = [
        {
            "input_ids": tensor(
                [
                    [0, 4, 0, 0, -2, 1, 1, 1, -2, 4, -2, 4, 4, 4, 6],
                    [-2, 2, -2, 2, 2, 2, -2, 2, 2, 2, -2, 3, -2, -2, 8],
                    [5, 5, -2, -2, 5, -2, -2, 6, -2, 6, 10, 6, 6, -2, -1],
                    [7, -2, -2, 10, -2, -2, -2, 10, 10, 10, 11, 11, 11, -2, -2],
                    [12, 12, 12, 12, 13, 13, 13, -2, -2, 13, -2, 6, 14, -2, -2],
                ]
            ),
            "labels": tensor(
                [
                    [-3, 0, -3, -3, 1, -3, -3, -3, 1, 4, 4, -3, -3, -3, 4],
                    [2, -3, 2, -3, -3, -3, 2, -3, -3, -3, 3, 3, 3, 8, -3],
                    [-3, -3, 5, 5, 5, 5, 6, 6, 6, -3, 6, -3, -3, -3, -3],
                    [-3, 7, 7, -3, 10, 10, 10, -3, -3, -3, -3, 11, -3, 11, 18],
                    [-3, 12, -3, -3, -3, -3, -3, 13, 13, 13, 13, 13, -3, 14, 14],
                ]
            ),
            "cu_seqlens": [
                tensor([0, 4, 9, 15], dtype=torch.int32),
                tensor([0, 10, 13, 15], dtype=torch.int32),
                tensor([0, 6, 13, 15], dtype=torch.int32),
                tensor([0, 3, 10, 14, 15], dtype=torch.int32),
                tensor([0, 4, 12, 15], dtype=torch.int32),
            ],
            "max_seqlen": [6, 10, 7, 7, 8],
            "attention_mask": tensor(
                [
                    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
                    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                ]
            ),