| """ |
| 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), |
| ) |
|
|
