Add training test suite — 19 tests covering data pipeline, packing, collation, integration, Trainer smoke test

tests/test_training.py

@@ -0,0 +1,207 @@
+"""
+Tests for domainTokenizer Phase 2C: Pre-training Pipeline.
+19 tests covering tokenization, packing, collation, integration, and Trainer smoke test.
+
+Run: pytest tests/test_training.py -v
+"""
+
+import logging
+import random
+from datetime import datetime, timedelta
+from typing import Any, Dict, List
+
+import numpy as np
+import torch
+import pytest
+
+from datasets import Dataset as HFDataset
+from transformers import DataCollatorForLanguageModeling
+
+from domain_tokenizer.schemas.predefined import FINANCE_SCHEMA
+from domain_tokenizer.tokenizers.domain_tokenizer import DomainTokenizerBuilder
+from domain_tokenizer.models.configuration import DomainTransformerConfig
+from domain_tokenizer.models.modeling import DomainTransformerForCausalLM
+from domain_tokenizer.training.data_pipeline import (
+    tokenize_user_sequences, pack_sequences, prepare_clm_dataset,
+)
+from domain_tokenizer.training.pretrain import pretrain_domain_model
+
+logging.basicConfig(level=logging.INFO)
+
+
+def make_finance_events(n_events=10, seed=42):
+    rng = random.Random(seed)
+    merchants = ["AMAZON", "UBER", "SALARY", "GROCERY", "NETFLIX", "GAS", "RESTAURANT", "PHARMACY"]
+    base_date = datetime(2025, 1, 1)
+    return [
+        {"amount_sign": (amt := rng.uniform(5, 5000) * rng.choice([1, -1])),
+         "amount": amt,
+         "timestamp": base_date + timedelta(days=rng.randint(0, 365), hours=rng.randint(0, 23)),
+         "description": rng.choice(merchants)}
+        for _ in range(n_events)
+    ]
+
+
+def make_user_sequences(n_users=20, min_events=5, max_events=30, seed=42):
+    rng = random.Random(seed)
+    return [make_finance_events(rng.randint(min_events, max_events), seed + i) for i in range(n_users)]
+
+
+def build_finance_tokenizer(events_flat):
+    builder = DomainTokenizerBuilder(FINANCE_SCHEMA)
+    builder.fit(events_flat)
+    text_corpus = list(set(e["description"] for e in events_flat)) * 20
+    return builder, builder.build(text_corpus=text_corpus, bpe_vocab_size=500)
+
+
+class TestTokenizeUserSequences:
+    @pytest.fixture
+    def setup(self):
+        seqs = make_user_sequences(5, 3, 10)
+        flat = [e for s in seqs for e in s]
+        b, t = build_finance_tokenizer(flat)
+        return seqs, b, t
+
+    def test_returns_list_of_lists(self, setup):
+        seqs, b, t = setup
+        r = tokenize_user_sequences(seqs, b, t)
+        assert len(r) == 5 and all(isinstance(s, list) for s in r)
+
+    def test_variable_lengths(self, setup):
+        seqs, b, t = setup
+        r = tokenize_user_sequences(seqs, b, t)
+        assert len(set(len(s) for s in r)) > 1
+
+    def test_bos_eos_present(self, setup):
+        seqs, b, t = setup
+        r = tokenize_user_sequences(seqs, b, t, add_bos=True, add_eos=True)
+        bos = t.convert_tokens_to_ids("[BOS]")
+        assert all(bos in s[:5] for s in r)
+
+    def test_no_bos_eos(self, setup):
+        seqs, b, t = setup
+        with_ = tokenize_user_sequences(seqs[:1], b, t, add_bos=True, add_eos=True)
+        without = tokenize_user_sequences(seqs[:1], b, t, add_bos=False, add_eos=False)
+        assert len(without[0]) < len(with_[0])
+
+
+class TestPackSequences:
+    def test_fixed_length(self):
+        ds = pack_sequences([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15]], block_size=5)
+        assert len(ds) == 3 and all(len(r["input_ids"]) == 5 for r in ds)
+
+    def test_concat(self):
+        ds = pack_sequences([[1, 2, 3], [4, 5, 6]], block_size=3)
+        assert ds[0]["input_ids"] == [1, 2, 3] and ds[1]["input_ids"] == [4, 5, 6]
+
+    def test_drops_remainder(self):
+        ds = pack_sequences([[1, 2, 3, 4, 5, 6, 7]], block_size=3)
+        assert len(ds) == 2
+
+    def test_too_few(self):
+        with pytest.raises(ValueError):
+            pack_sequences([[1, 2]], block_size=10)
+
+    def test_hf_dataset(self):
+        assert isinstance(pack_sequences([[i for i in range(100)]], block_size=10), HFDataset)
+
+    def test_no_padding(self):
+        ds = pack_sequences([[i for i in range(50)] for _ in range(10)], block_size=25)
+        assert all(len(r["input_ids"]) == 25 for r in ds)
+
+
+class TestPrepareCLMDataset:
+    def test_full(self):
+        seqs = make_user_sequences(10, 5, 15)
+        flat = [e for s in seqs for e in s]
+        b, t = build_finance_tokenizer(flat)
+        ds = prepare_clm_dataset(seqs, b, t, block_size=64)
+        assert len(ds) > 0 and all(len(r["input_ids"]) == 64 for r in ds)
+
+    def test_block_sizes(self):
+        seqs = make_user_sequences(10, 10, 20)
+        flat = [e for s in seqs for e in s]
+        b, t = build_finance_tokenizer(flat)
+        ds32 = prepare_clm_dataset(seqs, b, t, block_size=32)
+        ds64 = prepare_clm_dataset(seqs, b, t, block_size=64)
+        assert len(ds32) > len(ds64)
+
+
+class TestDataCollator:
+    @pytest.fixture
+    def setup(self):
+        seqs = make_user_sequences(5, 5, 15)
+        flat = [e for s in seqs for e in s]
+        b, t = build_finance_tokenizer(flat)
+        ds = prepare_clm_dataset(seqs, b, t, block_size=32)
+        return ds, DataCollatorForLanguageModeling(tokenizer=t, mlm=False), t
+
+    def test_adds_labels(self, setup):
+        ds, c, _ = setup
+        batch = c([ds[i] for i in range(min(4, len(ds)))])
+        assert all(k in batch for k in ["input_ids", "labels", "attention_mask"])
+
+    def test_labels_eq_ids(self, setup):
+        ds, c, _ = setup
+        batch = c([ds[0]])
+        assert torch.equal(batch["input_ids"], batch["labels"])
+
+    def test_shapes(self, setup):
+        ds, c, _ = setup
+        n = min(4, len(ds))
+        batch = c([ds[i] for i in range(n)])
+        assert batch["input_ids"].shape == (n, 32)
+
+    def test_all_ones_mask(self, setup):
+        ds, c, _ = setup
+        batch = c([ds[0]])
+        assert batch["attention_mask"].sum() == 32
+
+
+class TestTrainingIntegration:
+    def test_forward(self):
+        seqs = make_user_sequences(10, 5, 15)
+        flat = [e for s in seqs for e in s]
+        b, t = build_finance_tokenizer(flat)
+        ds = prepare_clm_dataset(seqs, b, t, block_size=32)
+        c = DataCollatorForLanguageModeling(tokenizer=t, mlm=False)
+        batch = c([ds[i] for i in range(min(4, len(ds)))])
+        config = DomainTransformerConfig(vocab_size=t.vocab_size, hidden_size=64,
+                                          num_hidden_layers=2, num_attention_heads=4, intermediate_size=128)
+        model = DomainTransformerForCausalLM(config)
+        out = model(**batch)
+        assert out.loss.item() > 0
+        out.loss.backward()
+        assert sum(p.grad.norm().item() for p in model.parameters() if p.grad is not None) > 0
+
+
+class TestPretrainDomainModel:
+    def test_smoke(self, tmp_path):
+        seqs = make_user_sequences(20, 5, 15)
+        flat = [e for s in seqs for e in s]
+        b, t = build_finance_tokenizer(flat)
+        ds = prepare_clm_dataset(seqs, b, t, block_size=32)
+        config = DomainTransformerConfig(vocab_size=t.vocab_size, hidden_size=64,
+                                          num_hidden_layers=2, num_attention_heads=4, intermediate_size=128)
+        model = DomainTransformerForCausalLM(config)
+        trainer = pretrain_domain_model(
+            model=model, tokenizer=t, train_dataset=ds,
+            output_dir=str(tmp_path / "ck"), hub_model_id=None,
+            num_epochs=1, per_device_batch_size=4, gradient_accumulation_steps=1,
+            learning_rate=1e-3, warmup_steps=0, logging_steps=1,
+            save_steps=999999, report_to="none", seed=42,
+        )
+        assert trainer.state.global_step > 0
+
+    def test_no_pad_raises(self, tmp_path):
+        from transformers import PreTrainedTokenizerFast
+        from tokenizers import Tokenizer
+        from tokenizers.models import BPE
+        hf = PreTrainedTokenizerFast(tokenizer_object=Tokenizer(BPE(unk_token="[UNK]")), unk_token="[UNK]")
+        config = DomainTransformerConfig(vocab_size=100, hidden_size=32, num_hidden_layers=1, num_attention_heads=2)
+        with pytest.raises(ValueError, match="pad_token"):
+            pretrain_domain_model(
+                model=DomainTransformerForCausalLM(config), tokenizer=hf,
+                train_dataset=HFDataset.from_dict({"input_ids": [[1, 2, 3]]}),
+                output_dir=str(tmp_path),
+            )