def test_packing():
    """Tests that packing works for a single batch."""
    pad_token_id = 0
    max_seq_len = 5
    packer = BinPackCollator(
        collator=lambda x: x,
        target_batch_size=2,
        max_seq_len=max_seq_len,
        pad_token_id=pad_token_id,
        padding_side='right',
    )

    batch = _data_to_batch([
        [1],
        [2] * 2,
        [4] * 4,
        [3] * 3,
    ], max_seq_len, pad_token_id)

    packed_samples = packer.pack(batch)

    assert torch.equal(
        packed_samples['input_ids'],
        torch.Tensor([[3, 3, 3, 2, 2], [4, 4, 4, 4, 1]]),
    )
    assert torch.all(packed_samples['attention_mask'] == 1)
Tests that packing works for a single batch.
tests/data/test_packing.py