rzen
c202bada6a
- src/dataset.py: memmap uint16 token loader + random-crop batch sampler (x, y shifted-by-one, reproducible via generator) - scripts/encode_corpus.py: stream stories -> encode -> uint16 .bin, one EOT token appended per story, flat memory. encode_file() factored out for testing. - tests/test_dataset.py: batch shapes/shift/bounds/reproducibility + encoder round-trip/separators/uint16 on synthetic data (5 tests)
100 lines
3.2 KiB
Python
100 lines
3.2 KiB
Python
"""Dataset + corpus-encoding tests on synthetic data (no real corpus needed).
|
|
|
|
Run directly: `.venv/bin/python tests/test_dataset.py`
|
|
"""
|
|
|
|
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(ROOT / "src"))
|
|
sys.path.insert(0, str(ROOT / "scripts"))
|
|
|
|
from dataset import TOKEN_DTYPE, get_batch, load_tokens # noqa: E402
|
|
from tokenizer import BPETokenizer # noqa: E402
|
|
from encode_corpus import encode_file # noqa: E402
|
|
|
|
EOT = "<|endoftext|>"
|
|
|
|
|
|
def test_get_batch_shapes_and_shift():
|
|
# Tokens 0..999 in order, so the shift property is easy to verify.
|
|
data = np.arange(1000, dtype=TOKEN_DTYPE)
|
|
gen = torch.Generator().manual_seed(0)
|
|
x, y = get_batch(data, batch_size=8, context_length=16, device="cpu", generator=gen)
|
|
assert x.shape == (8, 16) and y.shape == (8, 16)
|
|
assert x.dtype == torch.long
|
|
# y is x shifted by one; with sequential data each y == x + 1.
|
|
assert torch.equal(y, x + 1)
|
|
|
|
|
|
def test_get_batch_is_reproducible():
|
|
data = np.arange(1000, dtype=TOKEN_DTYPE)
|
|
a = get_batch(data, 4, 8, "cpu", torch.Generator().manual_seed(42))
|
|
b = get_batch(data, 4, 8, "cpu", torch.Generator().manual_seed(42))
|
|
assert torch.equal(a[0], b[0]) and torch.equal(a[1], b[1])
|
|
|
|
|
|
def test_get_batch_in_bounds():
|
|
data = np.arange(50, dtype=TOKEN_DTYPE)
|
|
x, y = get_batch(data, 32, 8, "cpu", torch.Generator().manual_seed(1))
|
|
assert int(x.max()) < 50 and int(y.max()) < 50
|
|
|
|
|
|
def _tiny_tokenizer() -> BPETokenizer:
|
|
corpus = ("The cat sat. Tom ran fast. Lily played.\n" + EOT + "\n") * 40
|
|
return BPETokenizer.train(corpus, vocab_size=350, special_tokens=[EOT])
|
|
|
|
|
|
def test_encode_file_roundtrip_and_separators():
|
|
tok = _tiny_tokenizer()
|
|
eot_id = tok.special_tokens[EOT]
|
|
with tempfile.TemporaryDirectory() as d:
|
|
src = Path(d) / "mini.txt"
|
|
dst = Path(d) / "mini.bin"
|
|
# Two stories, separator on its own line, with stray blank lines.
|
|
src.write_text(
|
|
"The cat sat.\nTom ran fast.\n" + EOT + "\n\n"
|
|
"Lily played.\n" + EOT + "\n"
|
|
)
|
|
n_stories, n_tokens = encode_file(tok, src, dst, progress=False)
|
|
|
|
assert n_stories == 2
|
|
ids = load_tokens(dst).tolist()
|
|
assert len(ids) == n_tokens
|
|
# Exactly one EOT per story, and each story ends with one.
|
|
assert ids.count(eot_id) == 2
|
|
assert ids[-1] == eot_id
|
|
|
|
# Content between separators decodes back to the original stories.
|
|
first = ids[: ids.index(eot_id)]
|
|
assert tok.decode(first) == "The cat sat.\nTom ran fast."
|
|
|
|
|
|
def test_encoded_ids_fit_uint16():
|
|
tok = _tiny_tokenizer()
|
|
with tempfile.TemporaryDirectory() as d:
|
|
src = Path(d) / "m.txt"
|
|
dst = Path(d) / "m.bin"
|
|
src.write_text("The cat sat.\n" + EOT + "\n")
|
|
encode_file(tok, src, dst, progress=False)
|
|
arr = load_tokens(dst)
|
|
assert arr.dtype == TOKEN_DTYPE # stored as uint16
|
|
assert int(arr.max()) < 65536
|
|
|
|
|
|
def main() -> None:
|
|
tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
|
|
for t in tests:
|
|
t()
|
|
print(f"ok {t.__name__}")
|
|
print(f"\n{len(tests)} passed")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|