ludwig-ai--ludwig
wehub-resource-sync
593b94c120
pytest / Unit Tests (push) Has been cancelled
pytest / Integration (integration_tests_a) (push) Has been cancelled
pytest / Integration (integration_tests_b) (push) Has been cancelled
pytest / Integration (integration_tests_c) (push) Has been cancelled
pytest / Integration (integration_tests_d) (push) Has been cancelled
pytest / Integration (integration_tests_e) (push) Has been cancelled
pytest / Integration (integration_tests_f) (push) Has been cancelled
pytest / Integration (integration_tests_g) (push) Has been cancelled
pytest / Integration (integration_tests_h) (push) Has been cancelled
pytest / Integration (integration_tests_i) (push) Has been cancelled
pytest / Integration (integration_tests_j) (push) Has been cancelled
pytest / Distributed (distributed_a) (push) Has been cancelled
pytest / Distributed (distributed_b) (push) Has been cancelled
pytest / Distributed (distributed_c) (push) Has been cancelled
pytest / Distributed (distributed_d) (push) Has been cancelled
pytest / Distributed (distributed_e) (push) Has been cancelled
pytest / Distributed (distributed_f) (push) Has been cancelled
pytest / Minimal Install (push) Has been cancelled
pytest / Event File (push) Has been cancelled
pytest (slow) / py-slow (push) Has been cancelled
Publish JSON Schema / publish-schema (push) Has been cancelled
112 行
4.3 KiB
Python
112 行
4.3 KiB
Python
"""Tests for sequence packing utilities."""
|
|
|
|
import torch
|
|
|
|
from ludwig.utils.sequence_packing import create_block_diagonal_mask, pack_sequences
|
|
|
|
|
|
class TestPackSequences:
|
|
def test_basic_packing(self):
|
|
seqs = [torch.tensor([1, 2, 3]), torch.tensor([4, 5]), torch.tensor([6])]
|
|
masks = [torch.ones(3), torch.ones(2), torch.ones(1)]
|
|
labels = [torch.tensor([10, 20, 30]), torch.tensor([40, 50]), torch.tensor([60])]
|
|
|
|
packed_ids, packed_mask, packed_labels, seq_ids = pack_sequences(
|
|
seqs, masks, labels, max_length=8, pad_token_id=0
|
|
)
|
|
|
|
# All 3 short sequences should fit in one pack (3+2+1=6 <= 8)
|
|
assert packed_ids.shape[0] == 1
|
|
assert packed_ids.shape[1] == 8
|
|
|
|
def test_attention_is_block_diagonal(self):
|
|
seqs = [torch.tensor([1, 2, 3]), torch.tensor([4, 5])]
|
|
masks = [torch.ones(3), torch.ones(2)]
|
|
|
|
packed_ids, packed_mask, _, seq_ids = pack_sequences(seqs, masks, max_length=8, pad_token_id=0)
|
|
|
|
# Attention mask should be 2D (block diagonal)
|
|
assert packed_mask.dim() == 3 # [num_packs, max_length, max_length]
|
|
|
|
# Tokens from sequence 1 should NOT attend to tokens from sequence 2
|
|
# seq1 is at positions 0-2, seq2 is at positions 3-4
|
|
# Position 3 (first token of seq2) should NOT attend to position 2 (last token of seq1)
|
|
assert not packed_mask[0, 3, 2].item()
|
|
|
|
# But position 1 SHOULD attend to position 0 (same sequence, causal)
|
|
assert packed_mask[0, 1, 0].item()
|
|
|
|
def test_separate_packs_when_too_long(self):
|
|
seqs = [torch.tensor([1, 2, 3, 4, 5]), torch.tensor([6, 7, 8, 9, 10])]
|
|
masks = [torch.ones(5), torch.ones(5)]
|
|
|
|
packed_ids, _, _, _ = pack_sequences(seqs, masks, max_length=6, pad_token_id=0)
|
|
|
|
# Each sequence is 5 tokens, max_length=6, so they can't fit together
|
|
assert packed_ids.shape[0] == 2
|
|
|
|
def test_labels_masked_correctly(self):
|
|
seqs = [torch.tensor([1, 2, 3])]
|
|
masks = [torch.ones(3)]
|
|
labels = [torch.tensor([10, 20, 30])]
|
|
|
|
_, _, packed_labels, _ = pack_sequences(seqs, masks, labels, max_length=6, pad_token_id=0)
|
|
|
|
# Labels for actual tokens should be set, padding should be -100
|
|
assert packed_labels[0, 0] == 10
|
|
assert packed_labels[0, 1] == 20
|
|
assert packed_labels[0, 2] == 30
|
|
assert packed_labels[0, 3] == -100 # padding
|
|
|
|
def test_sequence_ids(self):
|
|
seqs = [torch.tensor([1, 2]), torch.tensor([3, 4, 5])]
|
|
masks = [torch.ones(2), torch.ones(3)]
|
|
|
|
_, _, _, seq_ids = pack_sequences(seqs, masks, max_length=8, pad_token_id=0)
|
|
|
|
# Should have different sequence IDs for each packed sequence
|
|
# and -1 for padding
|
|
assert (seq_ids[0, :5] >= 0).all() # 2+3 = 5 tokens have valid seq IDs
|
|
assert (seq_ids[0, 5:] == -1).all() # rest is padding
|
|
|
|
def test_max_sequences_per_pack(self):
|
|
seqs = [torch.tensor([i]) for i in range(10)]
|
|
masks = [torch.ones(1) for _ in range(10)]
|
|
|
|
_, _, _, _ = pack_sequences(seqs, masks, max_length=100, pad_token_id=0, max_sequences_per_pack=3)
|
|
# 10 sequences, max 3 per pack = at least 4 packs (ceil(10/3))
|
|
|
|
|
|
class TestBlockDiagonalMask:
|
|
def test_basic(self):
|
|
mask = create_block_diagonal_mask([3, 2], max_length=5, causal=True)
|
|
assert mask.shape == (5, 5)
|
|
|
|
# First block: 3x3 lower triangular
|
|
assert mask[0, 0].item()
|
|
assert mask[1, 0].item()
|
|
assert mask[1, 1].item()
|
|
assert not mask[0, 1].item() # causal: can't look ahead
|
|
|
|
# Second block: positions 3-4
|
|
assert mask[3, 3].item()
|
|
assert mask[4, 3].item()
|
|
assert mask[4, 4].item()
|
|
|
|
# Cross-block: should be zero
|
|
assert not mask[3, 2].item()
|
|
assert not mask[4, 0].item()
|
|
|
|
def test_non_causal(self):
|
|
mask = create_block_diagonal_mask([2, 2], max_length=4, causal=False)
|
|
# Non-causal: full attention within blocks
|
|
assert mask[0, 1].item() # position 0 can see position 1
|
|
assert mask[1, 0].item()
|
|
# Cross-block still zero
|
|
assert not mask[2, 1].item()
|
|
|
|
def test_exceeds_max_length(self):
|
|
mask = create_block_diagonal_mask([3, 3], max_length=4)
|
|
# Second sequence gets truncated at max_length
|
|
assert mask.shape == (4, 4)
|