Repository Public Research preview
1"""Usable inference API: ``HTRReconstructor.from_pretrained(dir).reconstruct(text, metadata)``.23Pipeline: protect spans (citations, math, URLs → placeholders) → tokenise with optional metadata prefix → decode4(greedy or beam, deterministic) → restore placeholders → preservation warnings (numbers, placeholders, negation,5hedging, length) → result object. Outputs are synthetic reconstructions; the result carries an explicit notice.6"""7from __future__ import annotations89import json10from dataclasses import asdict, dataclass, field11from pathlib import Path1213import torch1415from htr.data.taxonomy import era_of16from htr.evaluation.metrics import factual17from htr.models.tokenizer import HTRTokenizer18from htr.parse.protect import protect, restore_placeholders19from htr.training.data import DISC_IDX, ERA_IDX, ROLE_IDX, SRC_IDX, meta_prefix20from htr.training.trainer import load_checkpoint2122NOTICE = "Synthetic reconstruction produced by HTR. It is a model's estimate of pre-2018-style scientific prose given an LLM rewrite; it is not the original text and is not evidence of human authorship."232425@dataclass26class ReconstructionResult:27 input_text: str28 reconstruction: str29 warnings: list[str] = field(default_factory=list)30 preservation: dict = field(default_factory=dict)31 metadata_used: dict = field(default_factory=dict)32 model: str = ""33 notice: str = NOTICE3435 def to_dict(self) -> dict:36 return asdict(self)373839class HTRReconstructor:40 def __init__(self, model, tok: HTRTokenizer, release_meta: dict, device: str = "auto"):41 self.model = model42 self.tok = tok43 self.meta = release_meta44 self.text_meta_prefix = bool(release_meta.get("text_meta_prefix", False))45 self.device = next(model.parameters()).device4647 @classmethod48 def from_pretrained(cls, path: str | Path, device: str = "auto") -> "HTRReconstructor":49 path = Path(path)50 if device == "auto":51 device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")52 model = load_checkpoint(path, device)53 tok = HTRTokenizer(path / "tokenizer.model")54 rm = json.loads((path / "release_meta.json").read_text()) if (path / "release_meta.json").exists() else {}55 return cls(model, tok, rm, device)5657 def _meta_tensors(self, metadata: dict | None) -> dict | None:58 if self.model.cfg.arch not in ("m2", "m3"):59 return None60 m = metadata or {}61 era = m.get("era") or era_of(m.get("pub_year"))62 return {63 "disc": torch.tensor([DISC_IDX.get(m.get("discipline") or "unknown", DISC_IDX["unknown"])], device=self.device),64 "era": torch.tensor([ERA_IDX.get(era or "unknown", ERA_IDX["unknown"])], device=self.device),65 "role": torch.tensor([ROLE_IDX.get(m.get("section_role") or "other", ROLE_IDX["other"])], device=self.device),66 "src": torch.tensor([SRC_IDX.get(m.get("source") or "local", SRC_IDX["local"])], device=self.device),67 }6869 @torch.no_grad()70 def reconstruct(self, text: str, metadata: dict | None = None, beam: int = 1, max_new_tokens: int | None = None, seed: int = 0) -> ReconstructionResult:71 torch.manual_seed(seed)72 metadata = metadata or {}73 prot = protect(text)74 src_text = prot.text75 if self.text_meta_prefix:76 src_text = meta_prefix({"metadata": {**metadata, "era": metadata.get("era") or era_of(metadata.get("pub_year"))}}) + src_text77 ids = self.tok.encode(src_text, add_eos=True, max_len=self.model.cfg.max_len)78 src = torch.tensor([ids], device=self.device)79 mt = max_new_tokens or min(self.model.cfg.max_len - 2, int(len(ids) * 1.6) + 16)80 out = self.model.generate(src, self._meta_tensors(metadata), max_new_tokens=mt, beam=beam)81 hyp_prot = self.tok.decode(out[0].tolist())82 hyp = restore_placeholders(hyp_prot, prot)83 pres = factual(hyp_prot, prot.text)84 warnings = []85 if pres["placeholders_exact"] < 1.0:86 warnings.append("citations/equations/links not all preserved — verify protected spans")87 if pres["number_recall"] < 1.0:88 warnings.append(f"numbers from the input missing in the output (recall {pres['number_recall']:.2f})")89 if pres["number_precision"] < 1.0:90 warnings.append("output contains numbers absent from the input")91 if abs(pres["negation_delta"]) >= 1:92 warnings.append(f"negation count changed by {int(pres['negation_delta'])}")93 if abs(pres["hedge_delta"]) >= 3:94 warnings.append(f"hedging count changed by {int(pres['hedge_delta'])}")95 lr = len(hyp.split()) / max(1, len(text.split()))96 if lr < 0.6 or lr > 1.7:97 warnings.append(f"length ratio {lr:.2f} outside the expected range")98 if out.shape[1] >= mt:99 warnings.append("generation hit the length limit (possible truncation)")100 return ReconstructionResult(text, hyp, warnings, pres, {**metadata, "arch": self.model.cfg.arch}, self.meta.get("name", f"htr-{self.model.cfg.arch}"))101102 def reconstruct_batch(self, texts: list[str], metadatas: list[dict] | None = None, beam: int = 1) -> list[ReconstructionResult]:103 return [self.reconstruct(t, (metadatas or [{}] * len(texts))[i], beam) for i, t in enumerate(texts)]104105 @torch.no_grad()106 def reconstruct_many(self, texts: list[str], metadatas: list[dict] | None = None, batch_size: int = 32, seed: int = 0) -> list[ReconstructionResult]:107 """Batched greedy reconstruction (same outputs as ``reconstruct`` item by item, far faster): items are sorted by108 length, padded, decoded together with the KV cache, then restored and checked individually."""109 torch.manual_seed(seed)110 metas = metadatas or [{}] * len(texts)111 prepared = []112 for i, (t, m) in enumerate(zip(texts, metas)):113 prot = protect(t)114 src_text = prot.text115 if self.text_meta_prefix:116 src_text = meta_prefix({"metadata": {**m, "era": m.get("era") or era_of(m.get("pub_year"))}}) + src_text117 ids = self.tok.encode(src_text, add_eos=True, max_len=self.model.cfg.max_len)118 prepared.append((i, t, m, prot, ids))119 prepared.sort(key=lambda x: len(x[4]))120 results: list[ReconstructionResult | None] = [None] * len(texts)121 pad = self.model.cfg.pad_id122 for b in range(0, len(prepared), batch_size):123 chunk = prepared[b : b + batch_size]124 L = max(len(x[4]) for x in chunk)125 src = torch.tensor([x[4] + [pad] * (L - len(x[4])) for x in chunk], device=self.device)126 meta_t = None127 if self.model.cfg.arch in ("m2", "m3"):128 parts = [self._meta_tensors(x[2]) for x in chunk]129 meta_t = {k: torch.cat([p[k] for p in parts]) for k in parts[0]}130 mt = min(self.model.cfg.max_len - 2, int(L * 1.6) + 16)131 out = self.model.generate(src, meta_t, max_new_tokens=mt, beam=1)132 for row, (i, t, m, prot, ids) in zip(out, chunk):133 hyp_prot = self.tok.decode(row.tolist())134 hyp = restore_placeholders(hyp_prot, prot)135 pres = factual(hyp_prot, prot.text)136 warnings = []137 if pres["placeholders_exact"] < 1.0:138 warnings.append("citations/equations/links not all preserved — verify protected spans")139 if pres["number_recall"] < 1.0:140 warnings.append(f"numbers from the input missing in the output (recall {pres['number_recall']:.2f})")141 if pres["number_precision"] < 1.0:142 warnings.append("output contains numbers absent from the input")143 if abs(pres["negation_delta"]) >= 1:144 warnings.append(f"negation count changed by {int(pres['negation_delta'])}")145 lr = len(hyp.split()) / max(1, len(t.split()))146 if lr < 0.6 or lr > 1.7:147 warnings.append(f"length ratio {lr:.2f} outside the expected range")148 if int((row != pad).sum()) >= mt:149 warnings.append("generation hit the length limit (possible truncation)")150 results[i] = ReconstructionResult(t, hyp, warnings, pres, {**m, "arch": self.model.cfg.arch}, self.meta.get("name", f"htr-{self.model.cfg.arch}"))151 return results # type: ignore[return-value]152