Function test_packer_masking

tests/test_sequence_packer.py:148–248

()

def test_packer_masking():
    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.0,  # 0 mask probability, like no masking
        "seed": 42,
        "suppress_masking": False,
    }

    out_batches = list(GreedyBestFitSequencePacker(**d))

    # expected given a seed of 42!
    out_expected = [
        {
            "input_ids": tensor(
                [
                    [0, 0, 0, 0, 1, 1, 1, 1, 1, 4, 4, 4, 4, 4, 4],
                    [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 8, 8],
                    [5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, -1, -1],
                    [7, 7, 7, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 18],
                    [12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13, 14, 14, 14],
                ]
            ),
            "labels": tensor(
                [
                    [-3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3],
                    [-3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3],
                    [-3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3],
                    [-3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3],
                    [-3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3, -3],
                ]
            ),
            "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],
                ]