spb/wpai0002_uqo

1"""Datasets and collation for pairs (source = LLM rewrite, target = original human passage), with metadata2label maps, protected-token masks for L_factual and length-bucketed batching."""3from __future__ import annotations45import json6import functools7import random8import re9from dataclasses import dataclass, field10from pathlib import Path1112import torch1314from htr.data.taxonomy import DISCIPLINES, ERAS15from htr.models.tokenizer import ROLES, SOURCES, HTRTokenizer16from htr.utils.io import read_jsonl1718DISC_IDX = {d: i for i, d in enumerate(DISCIPLINES + ["unknown"])}19ERA_IDX = {e[0]: i for i, e in enumerate(ERAS)}20ERA_IDX.update({"unknown": len(ERA_IDX), "post2017": len(ERA_IDX) + 1})21ROLE_IDX = {r: i for i, r in enumerate(ROLES)}22SRC_IDX = {s: i for i, s in enumerate(SOURCES)}23_DIGIT = re.compile(r"\d")242526@dataclass27class StyleLabelMap:28    """Style classes for M3 = observed (discipline, era) buckets in the training split (training-only supervision)."""29    classes: list[str] = field(default_factory=list)3031    @classmethod32    def from_pairs(cls, pairs: list[dict], min_count: int = 20) -> "StyleLabelMap":33        from collections import Counter3435        c = Counter(f"{p['metadata'].get('discipline', 'unknown')}|{p['metadata'].get('era', 'unknown')}" for p in pairs)36        keep = sorted(k for k, n in c.items() if n >= min_count)37        return cls(classes=keep + ["other"])3839    def index(self, p: dict) -> int:40        k = f"{p['metadata'].get('discipline', 'unknown')}|{p['metadata'].get('era', 'unknown')}"41        return self.classes.index(k) if k in self.classes else len(self.classes) - 14243    def save(self, path: Path) -> None:44        path.write_text(json.dumps({"classes": self.classes}, indent=2))4546    @classmethod47    def load(cls, path: Path) -> "StyleLabelMap":48        return cls(**json.loads(path.read_text()))495051def meta_prefix(p: dict) -> str:52    m = p["metadata"]53    return f"<disc:{m.get('discipline') or 'unknown'}> <era:{m.get('era') or 'unknown'}> <role:{m.get('section_role') or 'other'}> <src:{m.get('source') or 'local'}> "545556class PairDataset(torch.utils.data.Dataset):57    def __init__(self, pairs_path: Path, tok: HTRTokenizer, max_src: int = 512, max_tgt: int = 512, style_map: StyleLabelMap | None = None, text_meta_prefix: bool = False, limit: int | None = None, split_doc_ids: set[str] | None = None, variants: set[str] | None = None):58        self.pairs = [p for p in read_jsonl(pairs_path) if (split_doc_ids is None or p["doc_id"] in split_doc_ids) and (variants is None or p["variant"] in variants)]59        if limit:60            self.pairs = self.pairs[:limit]61        self.tok, self.max_src, self.max_tgt = tok, max_src, max_tgt62        self.style_map = style_map63        self.text_meta_prefix = text_meta_prefix6465    def __len__(self) -> int:66        return len(self.pairs)6768    def __getitem__(self, i: int) -> dict:69        p = self.pairs[i]70        m = p["metadata"]71        src_text = (meta_prefix(p) if self.text_meta_prefix else "") + p["source_text"]72        src = self.tok.encode(src_text, add_bos=False, add_eos=True, max_len=self.max_src)73        tgt = self.tok.encode(p["target_text"], add_bos=True, add_eos=True, max_len=self.max_tgt)74        pieces = self.tok.sp.id_to_piece(tgt[1:])75        fact = [1.0 if (("⟦" in pc) or _DIGIT.search(pc)) else 0.0 for pc in pieces]76        return {77            "src": src, "tgt_in": tgt[:-1], "tgt_out": tgt[1:], "fact_mask": fact, "tgt_full": tgt[1:],78            "disc": DISC_IDX.get(m.get("discipline") or "unknown", DISC_IDX["unknown"]), "era": ERA_IDX.get(m.get("era") or "unknown", ERA_IDX["unknown"]),79            "role": ROLE_IDX.get(m.get("section_role") or "other", ROLE_IDX["other"]), "src_id": SRC_IDX.get(m.get("source") or "local", SRC_IDX["local"]),80            "style_label": self.style_map.index(p) if self.style_map else 0, "pair_id": p["pair_id"], "passage_id": p.get("passage_id", p["pair_id"]),81        }828384PAD_MULTIPLE = 32  # pad lengths to a multiple → far fewer distinct tensor shapes (MPS/CUDA kernel caches stay small)858687def _pad(seqs: list[list], pad_val, dtype=torch.long) -> torch.Tensor:88    L = max(len(s) for s in seqs)89    L = ((L + PAD_MULTIPLE - 1) // PAD_MULTIPLE) * PAD_MULTIPLE90    return torch.tensor([list(s) + [pad_val] * (L - len(s)) for s in seqs], dtype=dtype)919293def collate(items: list[dict], pad_id: int = 0) -> dict:94    return {95        "src": _pad([x["src"] for x in items], pad_id),96        "tgt_in": _pad([x["tgt_in"] for x in items], pad_id),97        "tgt_out": _pad([x["tgt_out"] for x in items], pad_id),98        "tgt_full": _pad([x["tgt_full"] for x in items], pad_id),99        "fact_mask": _pad([x["fact_mask"] for x in items], 0.0, torch.float),100        "meta": {"disc": torch.tensor([x["disc"] for x in items]), "era": torch.tensor([x["era"] for x in items]), "role": torch.tensor([x["role"] for x in items]), "src": torch.tensor([x["src_id"] for x in items])},101        "style_label": torch.tensor([x["style_label"] for x in items]),102        "pair_ids": [x["pair_id"] for x in items],103        "passage_ids": [x.get("passage_id", x["pair_id"]) for x in items],104    }105106107class BucketSampler(torch.utils.data.Sampler):108    """Groups examples of similar source length into batches with a token budget (dynamic batch size)."""109110    def __init__(self, lengths: list[int], max_tokens: int, max_batch: int = 64, seed: int = 0, shuffle: bool = True):111        self.lengths, self.max_tokens, self.max_batch, self.seed, self.shuffle = lengths, max_tokens, max_batch, seed, shuffle112        self.epoch = 0113114    def set_epoch(self, e: int) -> None:115        self.epoch = e116117    def _batches(self) -> list[list[int]]:118        idx = sorted(range(len(self.lengths)), key=lambda i: self.lengths[i])119        batches, cur, cur_max = [], [], 0120        for i in idx:121            L = self.lengths[i]122            if cur and (max(cur_max, L) * (len(cur) + 1) > self.max_tokens or len(cur) >= self.max_batch):123                batches.append(cur)124                cur, cur_max = [], 0125            cur.append(i)126            cur_max = max(cur_max, L)127        if cur:128            batches.append(cur)129        if self.shuffle:130            random.Random(self.seed + self.epoch).shuffle(batches)131        return batches132133    def __iter__(self):134        return iter(self._batches())135136    def __len__(self) -> int:137        return len(self._batches())138139140def make_loader(ds: PairDataset, max_tokens: int = 8192, max_batch: int = 64, seed: int = 0, shuffle: bool = True, num_workers: int = 0) -> torch.utils.data.DataLoader:141    lengths = [min(ds.max_src, len(p["source_text"]) // 4 + 8) for p in ds.pairs]142    sampler = BucketSampler(lengths, max_tokens, max_batch, seed, shuffle)143    return torch.utils.data.DataLoader(ds, batch_sampler=sampler, collate_fn=functools.partial(collate, pad_id=ds.tok.pad_id), num_workers=num_workers, prefetch_factor=(4 if num_workers else None))144