| """ |
| Tests for domainTokenizer Phase 2D: Fine-tuning Pipeline. |
| 15 tests covering dataset, batching, forward/backward, Trainer smoke, multiclass. |
| |
| Run: pytest tests/test_finetune.py -v |
| """ |
|
|
| import logging |
| import random |
| from datetime import datetime, timedelta |
|
|
| import numpy as np |
| import torch |
| import pytest |
|
|
| 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.models.joint_fusion import JointFusionModel |
| from domain_tokenizer.training.finetune_data import DomainFinetuneDataset, prepare_finetune_dataset |
| from domain_tokenizer.training.finetune import finetune_domain_model |
|
|
| logging.basicConfig(level=logging.INFO) |
|
|
|
|
| def make_events(n=10, seed=42): |
| rng = random.Random(seed) |
| merchants = ["AMAZON", "UBER", "SALARY", "GROCERY", "NETFLIX", "GAS"] |
| base = datetime(2025, 1, 1) |
| return [ |
| {"amount_sign": (a := rng.uniform(5, 5000) * rng.choice([1, -1])), |
| "amount": a, "timestamp": base + timedelta(days=rng.randint(0, 365), hours=rng.randint(0, 23)), |
| "description": rng.choice(merchants)} |
| for _ in range(n) |
| ] |
|
|
|
|
| def make_labeled_data(n_users=20, n_tabular=10, seed=42): |
| rng = random.Random(seed) |
| seqs = [make_events(rng.randint(3, 15), seed + i) for i in range(n_users)] |
| tab = np.random.RandomState(seed).randn(n_users, n_tabular).astype(np.float32) |
| labels = np.array([rng.choice([0.0, 1.0]) for _ in range(n_users)]) |
| return seqs, tab, labels |
|
|
|
|
| def build_tok(seqs): |
| flat = [e for s in seqs for e in s] |
| b = DomainTokenizerBuilder(FINANCE_SCHEMA) |
| b.fit(flat) |
| return b, b.build(text_corpus=list(set(e["description"] for e in flat)) * 20, bpe_vocab_size=300) |
|
|
|
|
| def tiny_cfg(v=128): |
| return DomainTransformerConfig(vocab_size=v, hidden_size=64, num_hidden_layers=2, |
| num_attention_heads=4, intermediate_size=128) |
|
|
|
|
| def make_fusion(v, nt=10, nc=1): |
| return JointFusionModel( |
| DomainTransformerForCausalLM(tiny_cfg(v)), nt, nc, |
| plr_frequencies=4, plr_embedding_dim=8, dcn_cross_layers=2, |
| dcn_deep_layers=1, dcn_deep_dim=32, head_hidden_dim=32, |
| ) |
|
|
|
|
| class TestDataset: |
| @pytest.fixture |
| def setup(self): |
| seqs, tab, labels = make_labeled_data(10, 5) |
| b, t = build_tok(seqs) |
| return DomainFinetuneDataset(seqs, tab, labels, b, t, 64), t |
|
|
| def test_len(self, setup): |
| assert len(setup[0]) == 10 |
|
|
| def test_keys(self, setup): |
| assert set(setup[0][0].keys()) == {"input_ids", "attention_mask", "tabular_features", "labels"} |
|
|
| def test_shapes(self, setup): |
| it = setup[0][0] |
| assert it["input_ids"].shape == (64,) and it["tabular_features"].shape == (5,) |
|
|
| def test_padding(self, setup): |
| it = setup[0][0] |
| assert (it["input_ids"] != setup[1].pad_token_id).any() |
|
|
| def test_mask_matches_pad(self, setup): |
| it = setup[0][0] |
| assert torch.equal(it["input_ids"] == setup[1].pad_token_id, it["attention_mask"] == 0) |
|
|
| def test_dtypes(self, setup): |
| it = setup[0][0] |
| assert it["labels"].dtype == torch.float32 and it["tabular_features"].dtype == torch.float32 |
|
|
| def test_mismatch(self): |
| seqs, tab, labels = make_labeled_data(10) |
| b, t = build_tok(seqs) |
| with pytest.raises(AssertionError): |
| DomainFinetuneDataset(seqs[:5], tab, labels, b, t) |
|
|
| def test_stats(self, setup): |
| assert setup[0].get_stats()["n_samples"] == 10 |
|
|
|
|
| class TestBatching: |
| def test_loader(self): |
| seqs, tab, labels = make_labeled_data(8, 5) |
| b, t = build_tok(seqs) |
| ds = DomainFinetuneDataset(seqs, tab, labels, b, t, 32) |
| batch = next(iter(torch.utils.data.DataLoader(ds, batch_size=4))) |
| assert batch["input_ids"].shape == (4, 32) and batch["tabular_features"].shape == (4, 5) |
|
|
|
|
| class TestForwardBackward: |
| def test_forward(self): |
| seqs, tab, labels = make_labeled_data(8, 5) |
| b, t = build_tok(seqs) |
| ds = DomainFinetuneDataset(seqs, tab, labels, b, t, 32) |
| batch = next(iter(torch.utils.data.DataLoader(ds, batch_size=4))) |
| out = make_fusion(t.vocab_size, 5)(** batch) |
| assert out["loss"].item() > 0 |
|
|
| def test_backward(self): |
| seqs, tab, labels = make_labeled_data(4, 5) |
| b, t = build_tok(seqs) |
| ds = DomainFinetuneDataset(seqs, tab, labels, b, t, 32) |
| batch = next(iter(torch.utils.data.DataLoader(ds, batch_size=4))) |
| model = make_fusion(t.vocab_size, 5) |
| model(**batch)["loss"].backward() |
| assert model.transformer.model.embed_tokens.weight.grad is not None |
| assert model.plr.frequencies.grad is not None |
|
|
| def test_multiclass(self): |
| seqs, tab, _ = make_labeled_data(8, 5) |
| labels = np.array([random.randint(0, 2) for _ in range(8)]) |
| b, t = build_tok(seqs) |
| ds = DomainFinetuneDataset(seqs, tab, labels, b, t, 32) |
| batch = next(iter(torch.utils.data.DataLoader(ds, batch_size=4))) |
| batch["labels"] = batch["labels"].long() |
| out = make_fusion(t.vocab_size, 5, 3)(**batch) |
| assert out["logits"].shape == (4, 3) and out["loss"] is not None |
|
|
|
|
| class TestTrainer: |
| def test_smoke(self, tmp_path): |
| seqs, tab, labels = make_labeled_data(20, 5) |
| b, t = build_tok(seqs) |
| ds = DomainFinetuneDataset(seqs, tab, labels, b, t, 32) |
| trainer = finetune_domain_model( |
| make_fusion(t.vocab_size, 5), ds, |
| output_dir=str(tmp_path), num_epochs=1, per_device_batch_size=4, |
| learning_rate=1e-3, warmup_steps=0, logging_steps=1, |
| save_strategy="no", report_to="none", seed=42, |
| ) |
| assert trainer.state.global_step > 0 |
| losses = [h["loss"] for h in trainer.state.log_history if "loss" in h] |
| assert len(losses) > 0 |
|
|
|
|
| class TestPrepare: |
| def test_prepare(self): |
| seqs, tab, labels = make_labeled_data(10, 5) |
| b, t = build_tok(seqs) |
| ds = prepare_finetune_dataset(seqs, tab, labels, b, t, 32) |
| assert len(ds) == 10 |
|
|
