Repository Public Research preview
1import json2import tempfile3from pathlib import Path45import pytest6import torch78from htr.models.config import make_config9from htr.models.losses import LossWeights, compute_loss10from htr.models.tokenizer import HTRTokenizer, train_tokenizer11from htr.models.transformer import HTRModel12from htr.training.data import PairDataset, StyleLabelMap, collate, make_loader13from htr.training.trainer import load_checkpoint, overfit_one_batch, save_checkpoint14from htr.utils.io import write_jsonl1516SENTS = ["the model reconstructs scientific prose from rewritten passages", "we measured 12 samples at 37 degrees ⟦C1⟧ and found no effect", "economic growth was estimated with panel regressions in 2010", "the decoder attends to encoder states and copies protected spans", "results suggest that hedging decreased after the intervention"]171819@pytest.fixture(scope="module")20def tmp_tok(tmp_path_factory):21 d = tmp_path_factory.mktemp("tok")22 txt = d / "t.txt"23 txt.write_text("\n".join(SENTS * 40))24 path = train_tokenizer(txt, d, vocab_size=1000, model_type="unigram")25 return HTRTokenizer(path)262728def test_tokenizer_roundtrip_and_control_tokens(tmp_tok):29 ids = tmp_tok.encode("we measured 12 samples ⟦C1⟧ <disc:economics>", add_bos=True)30 assert ids[0] == 1 and ids[-1] == 231 assert tmp_tok.piece_id("⟦C1⟧") != tmp_tok.unk_id and tmp_tok.piece_id("<disc:economics>") != tmp_tok.unk_id32 assert "⟦C1⟧" in tmp_tok.decode(ids)333435def _pairs(n=16):36 rows = []37 for i in range(n):38 s = SENTS[i % len(SENTS)]39 rows.append({"pair_id": f"p{i}", "doc_id": f"d{i % 4}", "rewrite_id": f"r{i}", "passage_id": f"g{i}", "variant": "A", "sample_index": 0, "source_text": s.replace("the", "a"), "target_text": s, "metadata": {"discipline": "economics" if i % 2 else "biomedicine", "era": "2010-2017", "section_role": "results", "source": "mock", "pub_year": 2012}, "quarantined": False})40 return rows414243@pytest.mark.parametrize("arch", ["m1", "m2", "m3"])44def test_forward_backward_and_overfit(tmp_tok, arch, tmp_path):45 pairs_path = tmp_path / "pairs.jsonl"46 write_jsonl(pairs_path, _pairs())47 sm = StyleLabelMap.from_pairs(_pairs(), min_count=1) if arch == "m3" else None48 ds = PairDataset(pairs_path, tmp_tok, style_map=sm)49 batch = collate([ds[i] for i in range(8)], 0)50 cfg = make_config("nano", arch, tmp_tok.vocab_size, n_styles=(len(sm.classes) if sm else 0), d_model=128, d_ff=256, n_heads=4, n_enc_layers=2, n_dec_layers=2)51 model = HTRModel(cfg)52 w = LossWeights(style=0.3, adv=0.3, con=0.2, sem=0.1, orth=0.1, fact=0.5) if arch == "m3" else LossWeights()53 loss, logs = compute_loss(model, batch, w)54 loss.backward()55 assert torch.isfinite(loss) and logs["n_tok"] > 056 res = overfit_one_batch(model, batch, steps=120, lr=2e-3, device="cpu", loss_w=LossWeights(label_smoothing=0.0))57 assert res["last_ce"] < 0.5 * res["first_ce"], res585960def test_checkpoint_roundtrip_and_generate(tmp_tok, tmp_path):61 cfg = make_config("nano", "m3", tmp_tok.vocab_size, n_styles=3, d_model=64, d_ff=128, n_heads=4, n_enc_layers=1, n_dec_layers=1)62 model = HTRModel(cfg)63 save_checkpoint(model, tmp_path / "ck", {"test": True})64 m2 = load_checkpoint(tmp_path / "ck", "cpu")65 src = torch.tensor([tmp_tok.encode("we measured 12 samples")])66 meta = {"disc": torch.tensor([1]), "era": torch.tensor([5]), "role": torch.tensor([4]), "src": torch.tensor([7])}67 a = model.eval().generate(src, meta, max_new_tokens=6)68 b = m2.generate(src, meta, max_new_tokens=6)69 assert torch.equal(a, b)70 assert json.loads((tmp_path / "ck" / "config.json").read_text())["arch"] == "m3"71 bm = m2.generate(src, meta, max_new_tokens=6, beam=2)72 assert bm.shape[0] == 1737475def test_loader_buckets(tmp_tok, tmp_path):76 pairs_path = tmp_path / "pairs.jsonl"77 write_jsonl(pairs_path, _pairs(32))78 ds = PairDataset(pairs_path, tmp_tok)79 dl = make_loader(ds, max_tokens=256, max_batch=8)80 n = sum(b["src"].shape[0] for b in dl)81 assert n == 3282