Repository Public
1"""Reconstructors built on pretrained Hugging Face backbones fine-tuned on the HTR pairs (HTR 2.0 track T3):2``HFSeq2SeqReconstructor`` (flan-t5-large / -xl, system ``B2_flan-t5-*``) and ``HFCausalReconstructor`` (Qwen2.5-1.5B / -3B3SFT, system ``LLM_qwen2.5-*``). Same protocol as ``HTRReconstructor``: protect placeholders → model → restore → preservation4warnings → ``ReconstructionResult``. Placeholders are spelled ``[[C1]]`` for these tokenizers (``hf_finetune.ph_to_ascii``).5``load_reconstructor(dir)`` dispatches on ``model_card.json``/``config.json``. Weights are loaded lazily by ``ModelRegistry``6(bounded number of resident models, LRU eviction) so a 48 GB node can serve several backbones."""7from __future__ import annotations89import json10import threading11import time12from collections import OrderedDict13from pathlib import Path1415import torch1617from htr.data.taxonomy import era_of18from htr.evaluation.metrics import factual19from htr.inference.pipeline import NOTICE, HTRReconstructor, ReconstructionResult20from htr.parse.protect import protect, restore_placeholders21from htr.training.data import meta_prefix22from htr.training.hf_finetune import ascii_to_ph, ph_to_ascii, protect_unk_chars, restore_unk_chars2324if hasattr(torch.backends, "cuda") and hasattr(torch.backends.cuda, "enable_cudnn_sdp"):25 torch.backends.cuda.enable_cudnn_sdp(False)2627CAUSAL_PROMPT = "Below is a passage of scientific prose that was rewritten by a language model. Reconstruct the original passage as the human author wrote it before 2018, keeping every number, citation placeholder and fact unchanged.\n\n{meta}\n### Rewritten passage:\n{src}\n\n### Original passage:\n"282930def _pick_device(device: str) -> torch.device:31 if device != "auto":32 return torch.device(device)33 return torch.device("cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu"))343536def preservation_warnings(hyp_prot: str, prot, text: str, hyp: str, hit_limit: bool) -> tuple[list[str], dict]:37 pres = factual(hyp_prot, prot.text)38 w = []39 if pres["placeholders_exact"] < 1.0:40 w.append("citations/equations/links not all preserved — verify protected spans")41 if pres["number_recall"] < 1.0:42 w.append(f"numbers from the input missing in the output (recall {pres['number_recall']:.2f})")43 if pres["number_precision"] < 1.0:44 w.append("output contains numbers absent from the input")45 if abs(pres["negation_delta"]) >= 1:46 w.append(f"negation count changed by {int(pres['negation_delta'])}")47 if abs(pres.get("hedge_delta", 0)) >= 3:48 w.append(f"hedging count changed by {int(pres['hedge_delta'])}")49 lr = len(hyp.split()) / max(1, len(text.split()))50 if lr < 0.6 or lr > 1.7:51 w.append(f"length ratio {lr:.2f} outside the expected range")52 if hit_limit:53 w.append("generation hit the length limit (possible truncation)")54 return w, pres555657def _meta_dict(metadata: dict | None) -> dict:58 m = dict(metadata or {})59 m["era"] = m.get("era") or era_of(m.get("pub_year"))60 return {"metadata": m}616263class HFSeq2SeqReconstructor:64 """flan-T5 family fine-tuned with ``scripts/train_b2_hf.py`` (input: ``reconstruct: <meta prefix> <protected text>``)."""6566 family = "hf-seq2seq"6768 def __init__(self, path: str | Path, device: str = "auto", dtype: torch.dtype | None = None, name: str | None = None):69 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer7071 path = Path(path)72 self.device = _pick_device(device)73 self.dtype = dtype or (torch.bfloat16 if self.device.type in ("cuda", "mps") else torch.float32)74 self.tok = AutoTokenizer.from_pretrained(path)75 self.model = AutoModelForSeq2SeqLM.from_pretrained(path, dtype=self.dtype).to(self.device).eval()76 self.meta = json.loads((path / "model_card.json").read_text()) if (path / "model_card.json").exists() else {}77 self.name = name or self.meta.get("name") or path.name78 self.max_len = int(self.meta.get("max_len", 512))7980 @torch.no_grad()81 def reconstruct(self, text: str, metadata: dict | None = None, beam: int = 1, max_new_tokens: int | None = None, seed: int = 0) -> ReconstructionResult:82 torch.manual_seed(seed)83 prot = protect(text)84 safe, unk_map = protect_unk_chars(ph_to_ascii(prot.text), self.tok)85 src = "reconstruct: " + meta_prefix(_meta_dict(metadata)) + safe86 enc = self.tok([src], return_tensors="pt", truncation=True, max_length=self.max_len).to(self.device)87 mt = max_new_tokens or min(self.max_len, int(enc.input_ids.shape[1] * 1.3) + 16)88 gen = self.model.generate(**enc, max_new_tokens=mt, num_beams=max(1, beam), do_sample=False)89 hyp_prot = ascii_to_ph(restore_unk_chars(self.tok.decode(gen[0], skip_special_tokens=True), unk_map))90 hyp = restore_placeholders(hyp_prot, prot)91 warnings, pres = preservation_warnings(hyp_prot, prot, text, hyp, int(gen.shape[1]) >= mt)92 return ReconstructionResult(text, hyp, warnings, pres, {**(metadata or {}), "arch": self.family}, self.name)939495class HFCausalReconstructor:96 """Decoder-only LLM fine-tuned with ``scripts/train_llm_sft.py`` (prompt → original; loss on the original only)."""9798 family = "hf-causal"99100 def __init__(self, path: str | Path, device: str = "auto", dtype: torch.dtype | None = None, name: str | None = None):101 from transformers import AutoModelForCausalLM, AutoTokenizer102103 path = Path(path)104 self.device = _pick_device(device)105 self.dtype = dtype or (torch.bfloat16 if self.device.type in ("cuda", "mps") else torch.float32)106 self.tok = AutoTokenizer.from_pretrained(path)107 self.tok.padding_side = "left"108 if self.tok.pad_token is None:109 self.tok.pad_token = self.tok.eos_token110 self.model = AutoModelForCausalLM.from_pretrained(path, dtype=self.dtype).to(self.device).eval()111 self.meta = json.loads((path / "model_card.json").read_text()) if (path / "model_card.json").exists() else {}112 self.name = name or self.meta.get("name") or path.name113 self.max_new = int(self.meta.get("max_new_tokens", 768))114115 @torch.no_grad()116 def reconstruct(self, text: str, metadata: dict | None = None, beam: int = 1, max_new_tokens: int | None = None, seed: int = 0) -> ReconstructionResult:117 torch.manual_seed(seed)118 prot = protect(text)119 prompt = CAUSAL_PROMPT.format(meta=meta_prefix(_meta_dict(metadata)).strip(), src=ph_to_ascii(prot.text))120 enc = self.tok([prompt], return_tensors="pt", add_special_tokens=False).to(self.device)121 mt = max_new_tokens or min(self.max_new, int(enc.input_ids.shape[1] * 0.8) + 64)122 gen = self.model.generate(**enc, max_new_tokens=mt, do_sample=False, num_beams=max(1, beam), pad_token_id=self.tok.pad_token_id)123 new = gen[0, enc.input_ids.shape[1] :]124 txt = self.tok.decode(new, skip_special_tokens=True).split("\n### ")[0].strip()125 hyp_prot = ascii_to_ph(txt)126 hyp = restore_placeholders(hyp_prot, prot)127 warnings, pres = preservation_warnings(hyp_prot, prot, text, hyp, int(new.shape[0]) >= mt)128 return ReconstructionResult(text, hyp, warnings, pres, {**(metadata or {}), "arch": self.family}, self.name)129130131def load_reconstructor(path: str | Path, device: str = "auto"):132 """Dispatch on the directory contents: HTR checkpoint (``config.json`` with ``arch``/``d_model``) or an HF directory133 (``model_card.json`` ``family`` or the HF config ``architectures``)."""134 path = Path(path)135 card = json.loads((path / "model_card.json").read_text()) if (path / "model_card.json").exists() else {}136 fam = card.get("family")137 if fam is None:138 cfg = json.loads((path / "config.json").read_text()) if (path / "config.json").exists() else {}139 if "d_model" in cfg and "arch" in cfg and "architectures" not in cfg:140 fam = "htr"141 else:142 archs = " ".join(cfg.get("architectures", []))143 fam = "hf-causal" if "CausalLM" in archs else "hf-seq2seq"144 if fam == "htr":145 rec = HTRReconstructor.from_pretrained(path, device=device)146 rec.family = "htr" # type: ignore[attr-defined]147 rec.name = card.get("name") or rec.meta.get("name", path.name) # type: ignore[attr-defined]148 return rec149 if fam == "hf-seq2seq":150 return HFSeq2SeqReconstructor(path, device=device, name=card.get("name"))151 if fam == "hf-causal":152 return HFCausalReconstructor(path, device=device, name=card.get("name"))153 raise ValueError(f"unknown model family {fam!r} in {path}")154155156class ModelRegistry:157 """Lazy, bounded set of resident reconstructors keyed by model id (sub-directory name of ``root``). ``model_card.json``158 in each sub-directory describes the model (name, family, n_params, metrics, licence, notes); weights are loaded on first159 use and the least-recently-used model is evicted when more than ``max_resident`` are loaded (``pinned`` never evicted)."""160161 def __init__(self, root: str | Path, device: str = "auto", max_resident: int = 2, pinned: tuple[str, ...] = ()):162 self.root = Path(root)163 self.device = device164 self.max_resident = max_resident165 self.pinned = set(pinned)166 self.cards: dict[str, dict] = {}167 for d in sorted(p for p in self.root.iterdir() if p.is_dir() and not p.name.startswith(".")):168 if (d / "model_card.json").exists():169 self.cards[d.name] = {**json.loads((d / "model_card.json").read_text()), "id": d.name}170 elif (d / "config.json").exists():171 self.cards[d.name] = {"id": d.name, "name": d.name, "family": "htr" if (d / "tokenizer.model").exists() else "hf"}172 self.loaded: OrderedDict[str, object] = OrderedDict()173 self.load_s: dict[str, float] = {}174 self.lock = threading.Lock()175176 def ids(self) -> list[str]:177 return list(self.cards)178179 def default_id(self) -> str:180 for mid, c in self.cards.items():181 if c.get("default"):182 return mid183 return next(iter(self.cards))184185 def get(self, mid: str):186 if mid not in self.cards:187 raise KeyError(mid)188 with self.lock:189 if mid in self.loaded:190 self.loaded.move_to_end(mid)191 return self.loaded[mid]192 while len(self.loaded) >= self.max_resident + len(self.pinned & set(self.loaded)):193 victim = next((k for k in self.loaded if k not in self.pinned), None)194 if victim is None:195 break196 del self.loaded[victim]197 if torch.backends.mps.is_available():198 torch.mps.empty_cache()199 if torch.cuda.is_available():200 torch.cuda.empty_cache()201 t0 = time.time()202 rec = load_reconstructor(self.root / mid, device=self.device)203 self.load_s[mid] = round(time.time() - t0, 2)204 self.loaded[mid] = rec205 return rec206207 def is_loaded(self, mid: str) -> bool:208 return mid in self.loaded209210211__all__ = ["HFSeq2SeqReconstructor", "HFCausalReconstructor", "ModelRegistry", "load_reconstructor", "preservation_warnings", "NOTICE"]212