Repository Public
1"""Minimum-risk training (Shen et al., ACL 2016) for HTR: after maximum-likelihood fine-tuning, directly minimise the2expected *risk* of the model's own samples, risk = 1 − chrF++(candidate, original)/100 (+ a fact-preservation penalty),3over a candidate subspace {gold, greedy, N nucleus samples}; Q(y) ∝ p(y|x)^α. A CE term keeps the MLE anchor.4Rationale: HTR's lexical gain over identity is small (+0.4 chrF++) and MBR over samples adds only +0.1–0.2; MRT lets the5training signal reward the *metric* instead of token likelihood (classic +1–2 BLEU/chrF in MT)."""6from __future__ import annotations78import torch910from htr.evaluation.metrics import _chrf, factual11from htr.models.losses import LossWeights, compute_loss121314def seq_logprob(model, src: torch.Tensor, meta: dict | None, cand: torch.Tensor, pad_id: int) -> torch.Tensor:15 """Sum of token log-probs of each candidate (B, L) under teacher forcing (BOS … EOS, pad after EOS)."""16 tgt_in, tgt_out = cand[:, :-1], cand[:, 1:]17 out = model(src, tgt_in, meta)18 lp = out["logp"].gather(-1, tgt_out.unsqueeze(-1)).squeeze(-1)19 mask = tgt_out.ne(pad_id)20 return (lp * mask).sum(-1)212223def pad_to(ids: list[list[int]], pad_id: int, L: int | None = None) -> torch.Tensor:24 L = L or max(len(x) for x in ids)25 return torch.tensor([x + [pad_id] * (L - len(x)) for x in ids], dtype=torch.long)262728def strip(ids: list[int], eos_id: int, pad_id: int) -> list[int]:29 out = []30 for t in ids:31 out.append(t)32 if t == eos_id:33 break34 return [t for t in out if t != pad_id] if out and out[-1] == eos_id else [t for t in out if t != pad_id] + [eos_id]353637class Discriminator:38 """Frozen rewrite-vs-original classifier (scripts/train_discriminator.py) used as a *reward* term: P(original | text).39 In-house proxy only; the held-out detector (Pangram) is never used for training or tuning."""4041 def __init__(self, path: str, device, max_len: int = 384, batch_size: int = 64):42 from transformers import AutoModelForSequenceClassification, AutoTokenizer4344 self.tok = AutoTokenizer.from_pretrained(path)45 self.model = AutoModelForSequenceClassification.from_pretrained(path).to(device).eval()46 for p in self.model.parameters():47 p.requires_grad_(False)48 self.device, self.max_len, self.batch_size = device, max_len, batch_size4950 @torch.no_grad()51 def p_human(self, texts: list[str]) -> list[float]:52 out: list[float] = []53 for i in range(0, len(texts), self.batch_size):54 enc = self.tok(texts[i : i + self.batch_size], return_tensors="pt", padding=True, truncation=True, max_length=self.max_len).to(self.device)55 with torch.autocast("cuda", dtype=torch.bfloat16, enabled=self.device.type == "cuda"):56 lg = self.model(**enc).logits.float()57 out += lg.softmax(-1)[:, 1].tolist()58 return out596061def mrt_loss(model, tok, batch: dict, n_samples: int = 8, temperature: float = 0.8, top_p: float = 0.95, alpha: float = 0.5, w_ce: float = 0.3, fact_weight: float = 0.5, include_gold: bool = True, loss_w: LossWeights | None = None, disc: "Discriminator | None" = None, disc_weight: float = 0.0, restore=None) -> tuple[torch.Tensor, dict]:62 dev = batch["src"].device63 pad, eos, bos = model.cfg.pad_id, model.cfg.eos_id, model.cfg.bos_id64 src, meta = batch["src"], batch.get("meta")65 B = src.size(0)66 src_len = int(src.ne(pad).sum(1).max())67 mt = min(model.cfg.max_len - 2, int(src_len * 1.6) + 16)68 model.eval()69 with torch.no_grad():70 g = model.generate(src, meta, max_new_tokens=mt, beam=1)71 srcN = src.repeat_interleave(n_samples, 0)72 metaN = {k: v.repeat_interleave(n_samples, 0) for k, v in meta.items()} if meta else None73 s = model.generate(srcN, metaN, max_new_tokens=mt, beam=1, sample=True, temperature=temperature, top_p=top_p)74 model.train()75 gold_full = [[bos] + [t for t in row if t != pad] for row in batch["tgt_full"].tolist()] # tgt_full = ids after BOS incl. EOS76 cands: list[list[list[int]]] = []77 for b in range(B):78 c = [strip(g[b].tolist(), eos, pad)] + [strip(s[b * n_samples + k].tolist(), eos, pad) for k in range(n_samples)]79 if include_gold:80 c.append(gold_full[b])81 cands.append(c)82 K = len(cands[0])83 tgt_text = [tok.decode(x[1:]) for x in gold_full]84 src_text = [tok.decode([t for t in row if t != pad]) for row in src.tolist()]85 risk = torch.zeros(B, K, device=dev)86 texts: list[str] = []87 for b in range(B):88 for k, c in enumerate(cands[b]):89 txt = tok.decode(c[1:])90 texts.append(txt)91 r = 1.0 - float(_chrf.sentence_score(txt, [tgt_text[b]]).score) / 100.092 if fact_weight > 0:93 f = factual(txt, src_text[b])94 r += fact_weight * (1.0 - f["number_recall"]) + fact_weight * (1.0 - f["placeholders_exact"])95 risk[b, k] = r96 d_mean = None97 if disc is not None and disc_weight > 0:98 ph = torch.tensor(disc.p_human(texts), device=dev).view(B, K)99 risk = risk + disc_weight * (1.0 - ph) # reward human-likeness as judged by the in-house discriminator100 d_mean = float(ph[:, 0].mean())101 flat = [c for cs in cands for c in cs]102 L = min(model.cfg.max_len, max(len(c) for c in flat))103 cand_t = pad_to([c[:L] for c in flat], pad, L).to(dev)104 srcK = src.repeat_interleave(K, 0)105 metaK = {k: v.repeat_interleave(K, 0) for k, v in meta.items()} if meta else None106 lp = seq_logprob(model, srcK, metaK, cand_t, pad).view(B, K)107 q = torch.softmax(alpha * lp, dim=-1)108 l_mrt = (q * risk).sum(-1).mean()109 logs = {"mrt": float(l_mrt.detach()), "risk_greedy": float(risk[:, 0].mean()), "risk_min": float(risk.min(-1).values.mean()), "risk_q": float((q * risk).sum(-1).mean().detach()), "q_gold": float(q[:, -1].mean().detach()) if include_gold else 0.0}110 if d_mean is not None:111 logs["disc_p_human_greedy"] = d_mean112 total = l_mrt113 if w_ce > 0:114 l_ce, lg = compute_loss(model, batch, loss_w or LossWeights(fact=0.5))115 total = total + w_ce * l_ce116 logs["ce"] = lg["ce"]117 logs["total"] = float(total.detach())118 return total, logs119