spb/wpai0002_uqo

1"""Denoising pretraining on *unlabelled* original passages (no API cost): the model learns to reconstruct a passage2from a corrupted version (BART-style span masking, word dropout, local shuffling, sentence permutation, light3"LLM-ish" synonym substitution via the mock rewriter). This teaches copying, fluency and the target register before4fine-tuning on (rewrite → original) pairs. Only TRAIN-split documents are used; eligibility and placeholder protocol5are the same as for pairs."""6from __future__ import annotations78import random9import re1011import torch1213from htr.models.tokenizer import HTRTokenizer14from htr.parse.protect import PROTECT_VERSION, protect15from htr.rewrite.mock_provider import MockProvider16from htr.rewrite.prompts import build_messages17from htr.training.data import DISC_IDX, ERA_IDX, ROLE_IDX, SRC_IDX, _DIGIT18from htr.utils.io import read_jsonl19from htr.utils.text import split_sentences2021_WORD = re.compile(r"\S+")222324class Noiser:25    def __init__(self, seed: int = 0, mask_ratio: float = 0.15, drop_ratio: float = 0.1, shuffle_window: int = 3, synonym_prob: float = 0.5):26        self.rng = random.Random(seed)27        self.mask_ratio, self.drop_ratio, self.shuffle_window, self.synonym_prob = mask_ratio, drop_ratio, shuffle_window, synonym_prob28        self.mock = MockProvider(seed=seed)2930    def __call__(self, text: str) -> str:31        rng = self.rng32        t = text33        if rng.random() < self.synonym_prob:  # rule-based paraphrase as a cheap proxy of LLM rewriting34            sys_, msgs = build_messages(t, rng.choice("ABCDE"))35            r = self.mock.rewrite(sys_, msgs, 2000, {"variant": "A", "sample_index": rng.randrange(1000)})36            t = r.text or t37        sents = split_sentences(t)38        if len(sents) > 2 and rng.random() < 0.3:  # sentence permutation (local)39            i = rng.randrange(len(sents) - 1)40            sents[i], sents[i + 1] = sents[i + 1], sents[i]41            t = " ".join(sents)42        toks = _WORD.findall(t)43        out: list[str] = []44        i = 045        while i < len(toks):46            w = toks[i]47            if "⟦" in w or _DIGIT.search(w):  # never corrupt protected spans or numbers48                out.append(w)49                i += 150                continue51            r = rng.random()52            if r < self.mask_ratio:  # span mask (Poisson-ish length 1–3) → single <mask>53                span = 1 + (rng.random() < 0.4) + (rng.random() < 0.2)54                out.append("<unk>")55                i += span56                continue57            if r < self.mask_ratio + self.drop_ratio:  # word dropout58                i += 159                continue60            out.append(w)61            i += 162        # local shuffling63        for j in range(0, len(out) - 1):64            if rng.random() < 0.05:65                k = min(len(out) - 1, j + rng.randrange(1, self.shuffle_window + 1))66                out[j], out[k] = out[k], out[j]67        return " ".join(out)686970class DenoiseDataset(torch.utils.data.Dataset):71    def __init__(self, passages_path: str, docs_path: str, tok: HTRTokenizer, train_doc_ids: set[str], max_len: int = 512, seed: int = 0, limit: int | None = None, variants_per_passage: int = 1):72        docs = {d["doc_id"]: d for d in read_jsonl(docs_path)}73        self.items = []74        for p in read_jsonl(passages_path):75            if not p["eligible"] or p["doc_id"] not in train_doc_ids:76                continue77            d = docs[p["doc_id"]]78            prot = protect(p["text"], version=PROTECT_VERSION).text79            for v in range(variants_per_passage):80                self.items.append({"text": prot, "disc": d.get("subject_area") or "unknown", "year": d.get("pub_year"), "role": p["section_role"], "src": d["source"], "seed": v})81            if limit and len(self.items) >= limit:82                break83        self.tok, self.max_len, self.seed = tok, max_len, seed84        self.epoch = 08586    def __len__(self) -> int:87        return len(self.items)8889    def __getitem__(self, i: int) -> dict:90        from htr.data.taxonomy import era_of9192        it = self.items[i]93        noiser = Noiser(seed=self.seed * 1_000_003 + self.epoch * 7919 + i * 31 + it["seed"])94        noised = noiser(it["text"])95        src = self.tok.encode(noised, add_eos=True, max_len=self.max_len)96        tgt = self.tok.encode(it["text"], add_bos=True, add_eos=True, max_len=self.max_len)97        pieces = self.tok.sp.id_to_piece(tgt[1:])98        return {99            "src": src, "tgt_in": tgt[:-1], "tgt_out": tgt[1:], "tgt_full": tgt[1:],100            "fact_mask": [1.0 if ("⟦" in pc or _DIGIT.search(pc)) else 0.0 for pc in pieces],101            "disc": DISC_IDX.get(it["disc"], DISC_IDX["unknown"]), "era": ERA_IDX.get(era_of(it["year"]), ERA_IDX["unknown"]),102            "role": ROLE_IDX.get(it["role"], ROLE_IDX["other"]), "src_id": SRC_IDX.get(it["src"], SRC_IDX["local"]), "style_label": 0, "pair_id": f"dn{i}",103        }104