Repository Public Research preview
1"""Inference API for the non-generative HTR-Edit model: ``HTREditor.from_pretrained(dir).edit(text, metadata)``.23Pipeline: protect spans (citations, math, URLs → placeholders, exactly as ``htr.inference.pipeline``) → tokenise →4one encoder forward pass → argmax KEEP/DELETE per source token and argmax phrase per insertion slot → protected5tokens (placeholders, digit-bearing pieces) and their glue pieces forced to KEEP, insertions inside/next to them6blocked (``protection_plan``) → realise edits → decode → restore placeholders →7preservation checks (``htr.evaluation.metrics.factual``) and warnings. Deterministic; no sampling, no decoding loop.8Sources longer than ``max_src`` tokens are edited on their first ``max_src`` tokens; the remainder is passed through9unchanged (warning added).10"""11from __future__ import annotations1213import json14from pathlib import Path1516import torch1718from htr.evaluation.metrics import factual19from htr.inference.pipeline import NOTICE, ReconstructionResult20from htr.models.tagger import HTREditModel, load_edit_checkpoint21from htr.models.tokenizer import HTRTokenizer22from htr.parse.protect import protect, restore_placeholders23from htr.training.edits import PhraseVocab, apply_edits, protected_ids, protection_plan, space_ids2425EDIT_NOTICE = NOTICE.replace("produced by HTR", "produced by HTR-Edit (non-generative editor)")262728def _warnings(pres: dict, hyp: str, text: str) -> list[str]:29 w = []30 if pres["placeholders_exact"] < 1.0:31 w.append("citations/equations/links not all preserved — verify protected spans")32 if pres["number_recall"] < 1.0:33 w.append(f"numbers from the input missing in the output (recall {pres['number_recall']:.2f})")34 if pres["number_precision"] < 1.0:35 w.append("output contains numbers absent from the input")36 if abs(pres["negation_delta"]) >= 1:37 w.append(f"negation count changed by {int(pres['negation_delta'])}")38 if abs(pres["hedge_delta"]) >= 3:39 w.append(f"hedging count changed by {int(pres['hedge_delta'])}")40 lr = len(hyp.split()) / max(1, len(text.split()))41 if lr < 0.6 or lr > 1.7:42 w.append(f"length ratio {lr:.2f} outside the expected range")43 return w444546class HTREditor:47 def __init__(self, model: HTREditModel, tok: HTRTokenizer, phrases: PhraseVocab, release_meta: dict | None = None, none_bias: float | None = None):48 self.model = model.eval()49 self.tok = tok50 self.phrases = phrases51 self.meta = release_meta or {}52 self.none_bias = none_bias53 self.device = next(model.parameters()).device54 self.keep_ids = protected_ids(tok.sp)55 self.spaced = space_ids(tok.sp)56 if len(phrases) != model.cfg.n_phrases:57 raise ValueError(f"phrase vocabulary ({len(phrases)}) does not match the model head ({model.cfg.n_phrases})")5859 @classmethod60 def from_pretrained(cls, path: str | Path, device: str = "auto", none_bias: float | None = None) -> "HTREditor":61 path = Path(path)62 if device == "auto":63 device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")64 model = load_edit_checkpoint(path, device)65 tok = HTRTokenizer(path / "tokenizer.model")66 phrases = PhraseVocab.load(path / "phrases.json")67 rm = json.loads((path / "release_meta.json").read_text()) if (path / "release_meta.json").exists() else {}68 return cls(model, tok, phrases, rm, none_bias)6970 def edit(self, text: str, metadata: dict | None = None) -> ReconstructionResult:71 return self.edit_many([text], [metadata or {}], batch_size=1)[0]7273 @torch.no_grad()74 def edit_many(self, texts: list[str], metadatas: list[dict] | None = None, batch_size: int = 64) -> list[ReconstructionResult]:75 metas = metadatas or [{} for _ in texts]76 cfg = self.model.cfg77 prepared = []78 for i, (t, m) in enumerate(zip(texts, metas)):79 prot = protect(t)80 ids = self.tok.sp.encode(prot.text, out_type=int)81 prepared.append((i, t, m or {}, prot, ids))82 order = sorted(range(len(prepared)), key=lambda k: len(prepared[k][4]))83 results: list[ReconstructionResult | None] = [None] * len(texts)84 for b in range(0, len(order), batch_size):85 chunk = [prepared[k] for k in order[b : b + batch_size]]86 heads = [x[4][: cfg.max_src] for x in chunk]87 L = max(len(h) for h in heads) + 288 inp = torch.tensor([[cfg.bos_id] + h + [cfg.eos_id] + [cfg.pad_id] * (L - len(h) - 2) for h in heads], device=self.device)89 tags, ins = self.model.predict(inp, self.none_bias)90 tags, ins = tags.cpu().tolist(), ins.cpu().tolist()91 for (i, t, m, prot, ids), h, tg, it in zip(chunk, heads, tags, ins):92 n = len(h)93 keep, block = protection_plan(h, self.keep_ids, self.spaced)94 if len(ids) > n:95 block[n] = True # never insert at the truncation boundary96 out_ids = apply_edits(h, tg[1 : n + 1], it[: n + 1], self.phrases, force_keep=keep, block_slot=block) + ids[n:]97 hyp_prot = self.tok.sp.decode(out_ids)98 hyp = restore_placeholders(hyp_prot, prot)99 pres = factual(hyp_prot, prot.text)100 warnings = _warnings(pres, hyp, t)101 if len(ids) > n:102 warnings.append(f"input longer than {cfg.max_src} tokens: the last {len(ids) - n} tokens were passed through unchanged")103 n_del = sum(1 for k in range(n) if tg[k + 1] == 1 and not keep[k])104 n_ins = sum(1 for k in range(n + 1) if it[k] != 0 and not block[k])105 results[i] = ReconstructionResult(t, hyp, warnings, pres, {**m, "arch": "htr_edit", "edits": {"deleted_tokens": n_del, "insertions": n_ins, "source_tokens": len(ids)}}, self.meta.get("name", "htr-edit"), EDIT_NOTICE)106 return results # type: ignore[return-value]107