File size: 2,196 Bytes
dc72977 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import pytest
from datasets import Dataset
from dlb.data import build_owt_split, pack_tokens, preprocess_split
class TinyBatchTokenizer:
def __call__(self, texts, **kwargs):
assert kwargs == {
"add_special_tokens": False,
"return_attention_mask": False,
"return_token_type_ids": False,
}
table = {"first": [10, 11], "second": [12, 13, 14]}
return {"input_ids": [table[text] for text in texts]}
def test_owt_split_is_last_100k_documents():
"""Catch shuffling or selecting any validation slice except the final 100k."""
split = build_owt_split(total_documents=8_013_769)
assert split.train == "train[:-100000]"
assert split.validation == "train[-100000:]"
def test_owt_split_rejects_a_source_too_small_for_both_splits():
"""Catch silently producing an empty training split from a truncated snapshot."""
with pytest.raises(ValueError, match="more than 100,000"):
build_owt_split(total_documents=100_000)
def test_pack_tokens_reserves_boundaries():
"""Catch a packed block ending in document content rather than EOS."""
blocks = list(
pack_tokens([[10, 11], [12, 13, 14]], length=6, bos_id=101, eos_id=102)
)
assert blocks[0] == [101, 10, 11, 102, 12, 102]
assert all(len(block) == 6 for block in blocks)
def test_pack_tokens_carries_overflow_and_drops_only_the_incomplete_tail():
"""Catch token loss when one document crosses a fixed-length boundary."""
blocks = list(
pack_tokens([[10, 11, 12, 13, 14, 15]], length=6, bos_id=101, eos_id=102)
)
assert blocks == [[101, 10, 11, 12, 13, 102]]
def test_preprocess_split_packs_real_dataset_rows_in_source_order(tmp_path):
"""Catch preprocessing that reorders documents or packs each batch separately."""
source = Dataset.from_dict({"text": ["first", "second"]})
processed = preprocess_split(
source,
tokenizer=TinyBatchTokenizer(),
length=6,
bos_id=101,
eos_id=102,
cache_dir=tmp_path,
batch_documents=1,
)
assert processed.to_list() == [
{"input_ids": [101, 10, 11, 102, 12, 102]},
]
|