Function test_packer

tests/test_sequence_packer.py:80–145

()

def test_packer():
    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,
        "seed": 42,
        "suppress_masking": True,
    }

    out_batches = list(GreedyBestFitSequencePacker(**d))

    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": None,
            "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],
        },
        {
            "input_ids": tensor(
                [
                    [19, 19, 19, 19, 17, 17, 17, 17, 17, 17, 17, 17, 23, 23, -1],
                    [16, 16, 16, 16, 16, 16, 16, 15, 15, 15, 15, 15, 15, 15, -1],
                    [20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 21, 21, 21, 21, -1],
                    [22, 22, 22, 22, 22, 22, 22, 22, 25, 25, 25, 25, 25, 25, 25],
                    [24, 24, 24, 24, 24, 24, 24, 24, 24, 24, -1, -1, -1, -1, -1],
                ]
            ),
            "labels": None,
            "cu_seqlens": [
                tensor([0, 4, 12, 14, 15], dtype=torch.int32),
                tensor([0, 7, 14, 15], dtype=torch.int32),
                tensor([0, 10, 14, 15], dtype=torch.int32),
                tensor([0, 8, 15], dtype=torch.int32),