Function test_split_packed_batch

tests/test_sequence_packer.py:462–613

()

def test_split_packed_batch():
    input_batches = [
        {
            "input_ids": tensor(
                [
                    [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],
                ]
            ),
            "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],
                ]
            ),
            "cu_seqlens": [
                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),
            ],
            "max_seqlen": [10, 7, 7],
            "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, 0, 0],
                    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                ]
            ),
        },
        {
            "input_ids": tensor(
                [
                    [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": 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],
                ]
            ),
            "cu_seqlens": [
                tensor([0, 10, 14, 15], dtype=torch.int32),
                tensor([0, 8, 15], dtype=torch.int32),
                tensor([0, 10, 15], dtype=torch.int32),
            ],
            "max_seqlen": [10, 8, 10],
            "attention_mask": tensor(
                [
                    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 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, 0, 0, 0, 0, 0],
                ]