Repository Public
1"""Edit supervision for HTR-Edit (non-generative, LaserTagger/Felix-style text editing).23A pair (source = LLM rewrite, target = original passage), both in placeholder form, is tokenised with the project4SentencePiece model and the two token-id sequences are aligned with ``difflib.SequenceMatcher(autojunk=False)``.5The alignment is converted into per-source-token labels:67* ``tags[i]`` ∈ {KEEP=0, DELETE=1} for every source token ``i`` (n labels);8* ``inserts[k]`` = phrase id inserted *before* source position ``k`` (k = 0..n-1) or after the last token (k = n),9 0 = ``<none>``; phrase ids index a closed vocabulary of frequent inserted token sequences (length 1..max_len)10 built on the training split.1112Edit realisation: ``out = ins[0] (src[0] if KEEP) ins[1] (src[1] if KEEP) … ins[n]``.1314Opcode → labels:15* ``equal`` → KEEP;16* ``delete`` → DELETE;17* ``insert`` (i1 == i2) → one slot (k = i1) receives the inserted sequence, which must be a single in-vocabulary18 phrase;19* ``replace`` (src[i1:i2] → tgt[j1:j2]) → the k = i2 − i1 source tokens are deleted and the inserted sequence is20 segmented (minimum number of pieces, dynamic programming) into ≤ k in-vocabulary phrases placed in the slots21 i1, i1+1, … (each slot sits before a deleted token, so the order of the output is preserved).2223When an inserted sequence cannot be represented (out of vocabulary, too long for the available slots, or it24contains a protected token), the edit is *not representable*: a ``replace`` falls back to KEEP on its source span25(no insertion — locally the identity), an ``insert`` falls back to ``<none>``. Such pairs are not exactly26reconstructible; ``coverage`` = fraction of pairs whose labels reproduce the target ids exactly.2728Protected tokens: pieces containing a placeholder bracket ``⟦`` or a digit. Phrases containing a protected token are29never admitted to the vocabulary (the editor cannot write a number or a citation that is not in its input), and at30inference protected source tokens are forced to KEEP.3132Phrase vocabulary counts (training split only): every inserted sequence of length 1..max_len counts once as a whole;33every token of a longer inserted sequence counts once as a unigram (so long rewrites can still be segmented into34frequent units). The ``V_phr`` most frequent admissible sequences are kept; id 0 is ``<none>``.35"""36from __future__ import annotations3738import json39import re40from collections import Counter41from dataclasses import dataclass, field42from difflib import SequenceMatcher43from pathlib import Path44from typing import Iterable, Sequence4546KEEP, DELETE = 0, 147NONE_PHRASE = 048_DIGIT = re.compile(r"\d")495051# ---------------------------------------------------------------- protected tokens52def is_protected_piece(piece: str) -> bool:53 return ("⟦" in piece) or ("⟧" in piece) or bool(_DIGIT.search(piece))545556def protected_ids(sp) -> set[int]:57 """Ids of SentencePiece pieces that must never be deleted nor inserted (placeholders, digit-bearing pieces)."""58 return {i for i in range(sp.get_piece_size()) if is_protected_piece(sp.id_to_piece(i))}596061# ---------------------------------------------------------------- alignment62def edit_spans(src: Sequence[int], tgt: Sequence[int]) -> list[tuple[str, int, int, tuple[int, ...]]]:63 """Non-equal opcodes as (op, i1, i2, inserted target ids). ``op`` ∈ {delete, insert, replace}."""64 sm = SequenceMatcher(a=list(src), b=list(tgt), autojunk=False)65 return [(op, i1, i2, tuple(tgt[j1:j2])) for op, i1, i2, j1, j2 in sm.get_opcodes() if op != "equal"]666768# ---------------------------------------------------------------- phrase vocabulary69@dataclass70class PhraseVocab:71 phrases: list[tuple[int, ...]] = field(default_factory=lambda: [()]) # index 0 = <none>72 max_len: int = 473 meta: dict = field(default_factory=dict)7475 def __post_init__(self) -> None:76 self.index = {p: i for i, p in enumerate(self.phrases)}7778 def __len__(self) -> int:79 return len(self.phrases)8081 def get(self, seq: tuple[int, ...]) -> int | None:82 return self.index.get(tuple(seq))8384 @classmethod85 def build(cls, span_lists: Iterable[list[tuple[str, int, int, tuple[int, ...]]]], size: int = 4000, max_len: int = 4, protected: set[int] | None = None) -> "PhraseVocab":86 protected = protected or set()87 c: Counter = Counter()88 for spans in span_lists:89 for op, _i1, _i2, ins in spans:90 if not ins:91 continue92 if len(ins) <= max_len:93 c[ins] += 194 else:95 for t in ins:96 c[(t,)] += 197 admissible = [(p, n) for p, n in c.items() if not any(t in protected for t in p)]98 admissible.sort(key=lambda x: (-x[1], len(x[0]), x[0])) # deterministic ties99 kept = [p for p, _ in admissible[:size]]100 total = sum(n for _, n in admissible)101 meta = {"size": size, "max_len": max_len, "n_candidates": len(c), "n_admissible": len(admissible), "kept": len(kept), "count_mass_kept": (sum(n for _, n in admissible[:size]) / total) if total else 0.0}102 return cls([()] + kept, max_len, meta)103104 def save(self, path: Path, sp=None) -> None:105 rows = [{"id": i, "ids": list(p), **({"text": "".join(sp.id_to_piece(list(p))) if p else "<none>"} if sp is not None else {})} for i, p in enumerate(self.phrases)]106 Path(path).write_text(json.dumps({"max_len": self.max_len, "meta": self.meta, "phrases": rows}, ensure_ascii=False, indent=1))107108 @classmethod109 def load(cls, path: Path) -> "PhraseVocab":110 d = json.loads(Path(path).read_text())111 rows = sorted(d["phrases"], key=lambda r: r["id"])112 return cls([tuple(r["ids"]) for r in rows], d.get("max_len", 4), d.get("meta", {}))113114115def segment(seq: tuple[int, ...], vocab: PhraseVocab, max_pieces: int) -> list[int] | None:116 """Minimum-piece segmentation of ``seq`` into in-vocabulary phrases (≤ ``max_pieces``); None if impossible."""117 m = len(seq)118 INF = 10**9119 best = [INF] * (m + 1)120 back: list[tuple[int, int]] = [(-1, -1)] * (m + 1)121 best[0] = 0122 for j in range(1, m + 1):123 for L in range(1, min(vocab.max_len, j) + 1):124 pid = vocab.get(seq[j - L : j])125 if pid is not None and best[j - L] + 1 < best[j]:126 best[j] = best[j - L] + 1127 back[j] = (j - L, pid)128 if best[m] > max_pieces:129 return None130 out, j = [], m131 while j > 0:132 j0, pid = back[j]133 out.append(pid)134 j = j0135 return out[::-1]136137138# ---------------------------------------------------------------- labels139@dataclass140class EditLabels:141 tags: list[int] # n142 inserts: list[int] # n + 1143 covered: bool # labels reproduce the target exactly144 n_spans: int # non-equal opcodes145 n_unrepresentable: int # spans that fell back (OOV / too long / protected insertion)146147148def labels_from_spans(src: Sequence[int], tgt: Sequence[int], spans, vocab: PhraseVocab) -> EditLabels:149 n = len(src)150 tags = [KEEP] * n151 ins = [NONE_PHRASE] * (n + 1)152 bad = 0153 for op, i1, i2, seq in spans:154 if op == "delete":155 for i in range(i1, i2):156 tags[i] = DELETE157 continue158 slots = 1 if op == "insert" else (i2 - i1)159 pieces = segment(seq, vocab, slots)160 if pieces is None or any(ins[i1 + s] != NONE_PHRASE for s in range(len(pieces))):161 bad += 1 # not representable → replace keeps its source span, insert is dropped162 continue163 for s, pid in enumerate(pieces):164 ins[i1 + s] = pid165 if op == "replace":166 for i in range(i1, i2):167 tags[i] = DELETE168 lab = EditLabels(tags, ins, False, len(spans), bad)169 lab.covered = bad == 0 and apply_edits(src, tags, ins, vocab) == list(tgt)170 return lab171172173def make_labels(src: Sequence[int], tgt: Sequence[int], vocab: PhraseVocab) -> EditLabels:174 return labels_from_spans(src, tgt, edit_spans(src, tgt), vocab)175176177def space_ids(sp) -> set[int]:178 """Ids of pieces that start with the SentencePiece word-boundary marker (a space precedes them when decoded)."""179 return {i for i in range(sp.get_piece_size()) if sp.id_to_piece(i).startswith("▁")}180181182def protection_plan(src: Sequence[int], prot: set[int], spaced: set[int]) -> tuple[list[bool], list[bool]]:183 """Inference-time constraints around protected pieces (placeholders, digit-bearing pieces).184185 Returns (force_keep[n], block_slot[n+1]):186 * every protected piece is kept;187 * a protected piece glued to its left neighbour (no leading ``▁``, e.g. the 2nd piece of a multi-piece number or188 a number after a bare ``▁``/``(``) keeps that neighbour and no phrase may be inserted before it;189 * a non-protected piece glued to the right of a protected piece (``%``, ``.5``, ``mm`` …) is kept and no phrase190 may be inserted before it.191 Without these rules an argmax insertion between ``▁3`` and ``7`` or the deletion of the space before ``37`` would192 corrupt a number although every digit piece is kept."""193 n = len(src)194 keep = [False] * n195 block = [False] * (n + 1)196 for i, t in enumerate(src):197 if t not in prot:198 continue199 keep[i] = True200 if t not in spaced:201 block[i] = True202 if i > 0:203 keep[i - 1] = True204 if i + 1 < n and src[i + 1] not in prot and src[i + 1] not in spaced:205 keep[i + 1] = True206 block[i + 1] = True207 return keep, block208209210def apply_edits(src: Sequence[int], tags: Sequence[int], inserts: Sequence[int], vocab: PhraseVocab, keep_ids: set[int] | None = None, force_keep: Sequence[bool] | None = None, block_slot: Sequence[bool] | None = None) -> list[int]:211 """Realise edits. ``keep_ids``: ids forced to KEEP whatever the tag; ``force_keep``/``block_slot``: per-position212 constraints from ``protection_plan`` (blocked slots receive no insertion)."""213 out: list[int] = []214 for i, t in enumerate(src):215 if not (block_slot is not None and block_slot[i]):216 out.extend(vocab.phrases[inserts[i]])217 if tags[i] == KEEP or (keep_ids is not None and t in keep_ids) or (force_keep is not None and force_keep[i]):218 out.append(t)219 if len(inserts) > len(src) and not (block_slot is not None and block_slot[len(src)]):220 out.extend(vocab.phrases[inserts[len(src)]])221 return out222223224def coverage_stats(labels: Iterable[EditLabels]) -> dict:225 n = cov = oov_pairs = spans = bad = deletes = ins = toks = 0226 for lab in labels:227 n += 1228 cov += lab.covered229 oov_pairs += lab.n_unrepresentable > 0230 spans += lab.n_spans231 bad += lab.n_unrepresentable232 deletes += sum(1 for t in lab.tags if t == DELETE)233 ins += sum(1 for p in lab.inserts if p != NONE_PHRASE)234 toks += len(lab.tags)235 return {236 "n_pairs": n,237 "coverage_exact": cov / n if n else 0.0,238 "pairs_with_unrepresentable_insertion": oov_pairs / n if n else 0.0,239 "edit_spans": spans,240 "unrepresentable_spans": bad,241 "unrepresentable_span_rate": bad / spans if spans else 0.0,242 "delete_rate": deletes / toks if toks else 0.0,243 "insert_slot_rate": ins / (toks + n) if toks else 0.0,244 "source_tokens": toks,245 }246247248# ---------------------------------------------------------------- parallel tokenisation + alignment249_SP = None250251252def _init_worker(model_path: str) -> None:253 global _SP254 import sentencepiece as spm255256 _SP = spm.SentencePieceProcessor(model_file=model_path)257258259def tokenize_align(item: tuple[str, str]) -> tuple[list[int], list[int], list]:260 src_text, tgt_text = item261 src = _SP.encode(src_text, out_type=int)262 tgt = _SP.encode(tgt_text, out_type=int)263 return src, tgt, edit_spans(src, tgt)264265266def tokenize_align_many(texts: list[tuple[str, str]], model_path: str, workers: int = 8, chunksize: int = 64) -> list[tuple[list[int], list[int], list]]:267 if workers <= 1:268 _init_worker(model_path)269 return [tokenize_align(t) for t in texts]270 import multiprocessing as mp271272 with mp.get_context("spawn").Pool(workers, initializer=_init_worker, initargs=(model_path,)) as pool:273 return pool.map(tokenize_align, texts, chunksize=chunksize)274275276