Repository Public
1"""HTR-Edit: encoder-only, non-autoregressive text-editing model (LaserTagger/Felix-style).23Input = [BOS] src_1 … src_n [EOS] (positions 0..n+1)4Encoder = the project's pre-LN Transformer encoder (``EncoderLayer``: RoPE self-attention + SwiGLU), same parameter5 names as ``HTRModel`` (``emb``, ``enc``, ``enc_ln``) so a denoiser/HTR checkpoint's encoder loads by name.6Heads:7* ``tag`` : Linear(d, 2) on the state of each source token (positions 1..n) → KEEP / DELETE;8* ``insert`` : MLP(concat(h_k, h_{k+1})) → phrase id over V_phr + 1 for each of the n+1 insertion slots9 (slot k sits between position k and k+1; slot 0 uses the BOS state, slot n the learned EOS state).1011No decoder, no free generation: the output can only consist of kept source tokens and phrases of a closed,12training-built vocabulary (``htr.training.edits.PhraseVocab``). One forward pass per input.13"""14from __future__ import annotations1516import json17import math18from dataclasses import asdict, dataclass, field19from pathlib import Path2021import torch22import torch.nn as nn23import torch.nn.functional as F24from safetensors.torch import load_file, save_file2526from htr.models.config import SIZES, HTRConfig27from htr.models.transformer import EncoderLayer, rope_cache2829INSERT_W_EDIT, INSERT_W_NONE = 1.0, 0.1 # insertion-loss weights on slots with / without an insertion303132@dataclass33class EditConfig:34 vocab_size: int = 3200035 d_model: int = 38436 n_heads: int = 637 d_ff: int = 153638 n_enc_layers: int = 439 dropout: float = 0.140 max_src: int = 510 # source tokens per forward (BOS/EOS added)41 pos: str = "rope"42 n_phrases: int = 4001 # V_phr + <none>43 insert_w_none: float = INSERT_W_NONE44 insert_w_edit: float = INSERT_W_EDIT45 # inference-time logit adjustment that undoes the class weighting of the insertion loss (see HTR_EDIT.md):46 # with weights w_c the optimum of weighted CE is q_c ∝ w_c p_c, hence p ∝ q / w → add log(w_edit / w_none) to <none>47 none_logit_bias: float = math.log(INSERT_W_EDIT / INSERT_W_NONE)48 pad_id: int = 049 bos_id: int = 150 eos_id: int = 25152 def to_dict(self) -> dict:53 return asdict(self)5455 @classmethod56 def from_dict(cls, d: dict) -> "EditConfig":57 return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})5859 @classmethod60 def from_size(cls, size: str, vocab_size: int, **over) -> "EditConfig":61 s = SIZES[size]62 return cls(vocab_size=vocab_size, d_model=s["d_model"], n_heads=s["n_heads"], d_ff=s["d_ff"], n_enc_layers=s["n_enc_layers"], **over)6364 def layer_cfg(self) -> HTRConfig:65 return HTRConfig(vocab_size=self.vocab_size, d_model=self.d_model, n_heads=self.n_heads, d_ff=self.d_ff, n_enc_layers=self.n_enc_layers, n_dec_layers=0, dropout=self.dropout, pos=self.pos, pad_id=self.pad_id)666768class HTREditModel(nn.Module):69 def __init__(self, cfg: EditConfig):70 super().__init__()71 self.cfg = cfg72 lc = cfg.layer_cfg()73 self.emb = nn.Embedding(cfg.vocab_size, cfg.d_model, padding_idx=cfg.pad_id)74 self.emb_scale = math.sqrt(cfg.d_model)75 self.enc = nn.ModuleList(EncoderLayer(lc) for _ in range(cfg.n_enc_layers))76 self.enc_ln = nn.LayerNorm(cfg.d_model)77 self.drop = nn.Dropout(cfg.dropout)78 self.tag_head = nn.Linear(cfg.d_model, 2)79 self.ins_head = nn.Sequential(nn.Linear(2 * cfg.d_model, cfg.d_model), nn.GELU(), nn.Dropout(cfg.dropout), nn.Linear(cfg.d_model, cfg.n_phrases))80 self.apply(self._init)81 self.n_params = sum(p.numel() for p in self.parameters())8283 @staticmethod84 def _init(m):85 if isinstance(m, nn.Linear):86 nn.init.normal_(m.weight, std=0.02)87 if m.bias is not None:88 nn.init.zeros_(m.bias)89 elif isinstance(m, nn.Embedding):90 nn.init.normal_(m.weight, std=0.02)9192 def forward(self, ids: torch.Tensor) -> dict:93 """ids (B, L) = [BOS] src [EOS] + pad. Returns tag logits (B, L, 2) (meaningful at 1..n) and insertion logits94 (B, L-1, n_phrases) (slot k = between positions k and k+1, meaningful at 0..n)."""95 B, L = ids.shape96 pad = ids.eq(self.cfg.pad_id)97 x = self.drop(self.emb(ids) * self.emb_scale)98 rc = rope_cache(L, self.cfg.d_model // self.cfg.n_heads, ids.device) if self.cfg.pos == "rope" else None99 for layer in self.enc:100 x = layer(x, pad, rc)101 h = self.enc_ln(x)102 return {"tag": self.tag_head(h), "ins": self.ins_head(torch.cat([h[:, :-1], h[:, 1:]], -1))}103104 def loss(self, batch: dict) -> tuple[torch.Tensor, dict]:105 out = self(batch["ids"])106 tag_l = F.cross_entropy(out["tag"].flatten(0, 1), batch["tags"].flatten(), ignore_index=-100)107 ins_t = batch["ins"].flatten()108 ins_ce = F.cross_entropy(out["ins"].flatten(0, 1), ins_t.clamp(min=0), reduction="none")109 valid = ins_t.ne(-100)110 w = torch.where(ins_t.eq(0), self.cfg.insert_w_none, self.cfg.insert_w_edit) * valid.float()111 ins_l = (ins_ce * w).sum() / w.sum().clamp(min=1e-6)112 total = tag_l + ins_l113 return total, {"total": float(total.detach()), "tag_ce": float(tag_l.detach()), "ins_ce": float(ins_l.detach()), "n_tok": int(batch["tags"].ne(-100).sum())}114115 @torch.no_grad()116 def predict(self, ids: torch.Tensor, none_bias: float | None = None) -> tuple[torch.Tensor, torch.Tensor]:117 """Deterministic argmax edits: (tags (B, L), inserts (B, L-1)) aligned like ``forward``."""118 out = self(ids)119 ins = out["ins"]120 b = self.cfg.none_logit_bias if none_bias is None else none_bias121 if b:122 ins = ins.clone()123 ins[..., 0] += b124 return out["tag"].argmax(-1), ins.argmax(-1)125126127def init_from_checkpoint(model: HTREditModel, ckpt_dir: str | Path) -> dict:128 """Load ``emb.*``, ``enc.*``, ``enc_ln.*`` tensors by name (shapes must match) from an HTR checkpoint."""129 state = load_file(str(Path(ckpt_dir) / "model.safetensors"))130 own = model.state_dict()131 loadable = {k: v for k, v in state.items() if k in own and own[k].shape == v.shape and k.split(".")[0] in ("emb", "enc", "enc_ln")}132 model.load_state_dict({**own, **loadable})133 enc_keys = [k for k in own if k.split(".")[0] in ("emb", "enc", "enc_ln")]134 return {"init_from": str(ckpt_dir), "loaded": len(loadable), "encoder_tensors": len(enc_keys), "not_loaded_encoder": [k for k in enc_keys if k not in loadable], "new_tensors": [k for k in own if k not in loadable]}135136137def save_edit_checkpoint(model: HTREditModel, path: Path, extra: dict | None = None) -> None:138 from htr.utils.io import ensure_dir, write_json139 from htr.utils.provenance import run_header140141 ensure_dir(path)142 state = {k: v.detach().cpu().contiguous() for k, v in model.state_dict().items()}143 save_file(state, str(path / "model.safetensors"), metadata={"format": "pt", "model_type": "htr_edit", "htr_version": "0.1.0"})144 write_json(path / "config.json", {"model_type": "htr_edit", **model.cfg.to_dict(), "n_params": model.n_params})145 write_json(path / "checkpoint_meta.json", {**run_header(), **(extra or {})})146147148def load_edit_checkpoint(path: str | Path, device: str | torch.device = "cpu") -> HTREditModel:149 d = json.loads((Path(path) / "config.json").read_text())150 if d.get("model_type") != "htr_edit":151 raise ValueError(f"{path} is not an htr_edit checkpoint (model_type={d.get('model_type')})")152 model = HTREditModel(EditConfig.from_dict(d))153 model.load_state_dict(load_file(str(Path(path) / "model.safetensors")), strict=True)154 return model.to(device).eval()155