def test_packing_with_leftovers():
    """Tests that packing handles leftovers and computes waste correctly."""
    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,
        [4] * 4,
    ], max_seq_len, pad_token_id)

    packed_batch = packer.pack(batch)

    assert torch.equal(
        packed_batch['input_ids'],
        torch.Tensor([[4, 4, 4, 4, 1], [4, 4, 4, 4, 0]]),
    )
    assert torch.equal(
        packed_batch['attention_mask'],
        torch.Tensor([[1, 1, 1, 1, 1], [1, 1, 1, 1, 0]]),
    )

    # Check leftovers and waste.
    assert len(packer._leftover_bins) == 1
    leftover_size, leftover = packer._leftover_bins[0]
    assert leftover_size == 2
    assert torch.equal(leftover['input_ids'], torch.Tensor([2, 2]))
    assert torch.equal(leftover['attention_mask'], torch.Tensor([1, 1]))
    assert packer.waste == approx(2 / 11)  # 2 tokens wasted of 11 tokens total

    # Ensure that leftovers are used in the next batch if possible.
    batch = _data_to_batch([[1]], max_seq_len, pad_token_id)
    packed_batch = packer.pack(batch)
    assert torch.equal(
        packed_batch['input_ids'],
        torch.Tensor([[2, 2, 0, 0, 0], [1, 0, 0, 0, 0]]),
    )
    assert torch.equal(
        packed_batch['attention_mask'],
        torch.Tensor([[1, 1, 0, 0, 0], [1, 0, 0, 0, 0]]),
    )
Tests that packing handles leftovers and computes waste correctly.
tests/data/test_packing.py