Repository Public Research preview
1"""Encoder–decoder Transformer family for HTR (PyTorch reference implementation; MPS/CUDA/CPU).23M1: plain pre-LN encoder–decoder with RoPE and tied embeddings.4M2: M1 + metadata conditioning (discipline, era, section role, source embeddings added to the encoder input and5 prepended as one <meta> token to the decoder memory).6M3: dual-latent HTR — content encoder E_c(x) (the standard encoder), a style branch E_s(pool(E_c(x)), metadata)7 predicting a distribution over K historical style classes (training-only supervision from metadata), style8 memory tokens built from learned style prototypes weighted by the predicted distribution, a gradient-reversal9 adversary that discourages content representations from encoding style, an orthogonality penalty between10 pooled content and style vectors, and an optional pointer/copy gate over encoder positions. At inference only11 x and (optionally) metadata are used; the original target is never fed.12"""13from __future__ import annotations1415import math1617import torch18import torch.nn as nn19import torch.nn.functional as F2021from htr.models.config import HTRConfig222324# ---------------------------------------------------------------- positional encodings25def rope_cache(seq_len: int, dim: int, device, base: float = 10000.0) -> tuple[torch.Tensor, torch.Tensor]:26 inv = 1.0 / (base ** (torch.arange(0, dim, 2, device=device).float() / dim))27 t = torch.arange(seq_len, device=device).float()28 freqs = torch.outer(t, inv) # (L, dim/2)29 return freqs.cos(), freqs.sin()303132def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, offset: int = 0) -> torch.Tensor:33 # x: (B, H, L, D); positions offset..offset+L (incremental decoding)34 x1, x2 = x[..., ::2], x[..., 1::2]35 L = x.shape[-2]36 c, s = cos[offset : offset + L].unsqueeze(0).unsqueeze(0), sin[offset : offset + L].unsqueeze(0).unsqueeze(0)37 out = torch.stack([x1 * c - x2 * s, x1 * s + x2 * c], dim=-1)38 return out.flatten(-2)394041# ---------------------------------------------------------------- attention42class MultiHeadAttention(nn.Module):43 def __init__(self, d_model: int, n_heads: int, dropout: float, rope: bool):44 super().__init__()45 assert d_model % n_heads == 046 self.h, self.d = n_heads, d_model // n_heads47 self.q, self.k, self.v, self.o = (nn.Linear(d_model, d_model, bias=False) for _ in range(4))48 self.drop = dropout49 self.rope = rope5051 def forward(self, x: torch.Tensor, mem: torch.Tensor | None = None, key_pad: torch.Tensor | None = None, causal: bool = False, rope_cs: tuple | None = None, need_weights: bool = False, cache: dict | None = None, offset: int = 0):52 """``cache`` (incremental decoding): for self-attention holds past k/v and is extended with the new53 positions; for cross-attention holds the projected memory k/v computed once. ``offset`` = number of54 already-decoded positions (RoPE phase and causal mask)."""55 B, Lq, _ = x.shape56 q = self.q(x).view(B, Lq, self.h, self.d).transpose(1, 2)57 if mem is None:58 k = self.k(x).view(B, Lq, self.h, self.d).transpose(1, 2)59 v = self.v(x).view(B, Lq, self.h, self.d).transpose(1, 2)60 if self.rope and rope_cs is not None:61 q = apply_rope(q, *rope_cs, offset=offset)62 k = apply_rope(k, *rope_cs, offset=offset)63 if cache is not None:64 if "k" in cache:65 k = torch.cat([cache["k"], k], 2)66 v = torch.cat([cache["v"], v], 2)67 cache["k"], cache["v"] = k, v68 else:69 if cache is not None and "k" in cache:70 k, v = cache["k"], cache["v"]71 else:72 Lk0 = mem.shape[1]73 k = self.k(mem).view(B, Lk0, self.h, self.d).transpose(1, 2)74 v = self.v(mem).view(B, Lk0, self.h, self.d).transpose(1, 2)75 if cache is not None:76 cache["k"], cache["v"] = k, v77 Lk = k.shape[2]78 mask = None79 if key_pad is not None: # (B, Lk) True = pad80 mask = ~key_pad[:, None, None, :] # True = attend81 if causal:82 cm = torch.tril(torch.ones(Lq, Lk, dtype=torch.bool, device=x.device), diagonal=Lk - Lq)83 mask = cm[None, None] if mask is None else (mask & cm[None, None])84 if need_weights:85 att = (q @ k.transpose(-1, -2)) / math.sqrt(self.d)86 if mask is not None:87 att = att.masked_fill(~mask, float("-inf"))88 w = att.softmax(-1)89 out = F.dropout(w, self.drop, self.training) @ v90 weights = w.mean(1) # (B, Lq, Lk)91 else:92 out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=self.drop if self.training else 0.0)93 weights = None94 out = out.transpose(1, 2).reshape(B, Lq, self.h * self.d)95 return self.o(out), weights969798class FeedForward(nn.Module):99 def __init__(self, d_model: int, d_ff: int, dropout: float):100 super().__init__()101 self.w1, self.w2, self.w3 = nn.Linear(d_model, d_ff, bias=False), nn.Linear(d_ff, d_model, bias=False), nn.Linear(d_model, d_ff, bias=False)102 self.drop = nn.Dropout(dropout)103104 def forward(self, x): # SwiGLU105 return self.w2(self.drop(F.silu(self.w1(x)) * self.w3(x)))106107108class EncoderLayer(nn.Module):109 def __init__(self, cfg: HTRConfig):110 super().__init__()111 self.ln1, self.ln2 = nn.LayerNorm(cfg.d_model), nn.LayerNorm(cfg.d_model)112 self.attn = MultiHeadAttention(cfg.d_model, cfg.n_heads, cfg.dropout, cfg.pos == "rope")113 self.ff = FeedForward(cfg.d_model, cfg.d_ff, cfg.dropout)114 self.drop = nn.Dropout(cfg.dropout)115116 def forward(self, x, pad, rope_cs):117 a, _ = self.attn(self.ln1(x), key_pad=pad, rope_cs=rope_cs)118 x = x + self.drop(a)119 return x + self.drop(self.ff(self.ln2(x)))120121122class DecoderLayer(nn.Module):123 def __init__(self, cfg: HTRConfig):124 super().__init__()125 self.ln1, self.ln2, self.ln3 = nn.LayerNorm(cfg.d_model), nn.LayerNorm(cfg.d_model), nn.LayerNorm(cfg.d_model)126 self.self_attn = MultiHeadAttention(cfg.d_model, cfg.n_heads, cfg.dropout, cfg.pos == "rope")127 self.cross_attn = MultiHeadAttention(cfg.d_model, cfg.n_heads, cfg.dropout, False)128 self.ff = FeedForward(cfg.d_model, cfg.d_ff, cfg.dropout)129 self.drop = nn.Dropout(cfg.dropout)130131 def forward(self, y, mem, y_pad, mem_pad, rope_cs, need_weights=False, cache: dict | None = None, offset: int = 0):132 sc = cache.setdefault("self", {}) if cache is not None else None133 cc = cache.setdefault("cross", {}) if cache is not None else None134 a, _ = self.self_attn(self.ln1(y), key_pad=y_pad, causal=True, rope_cs=rope_cs, cache=sc, offset=offset)135 y = y + self.drop(a)136 c, w = self.cross_attn(self.ln2(y), mem=mem, key_pad=mem_pad, need_weights=need_weights, cache=cc)137 y = y + self.drop(c)138 return y + self.drop(self.ff(self.ln3(y))), w139140141class GradReverse(torch.autograd.Function):142 @staticmethod143 def forward(ctx, x, lam):144 ctx.lam = lam145 return x.view_as(x)146147 @staticmethod148 def backward(ctx, g):149 return -ctx.lam * g, None150151152def grad_reverse(x, lam: float):153 return GradReverse.apply(x, lam)154155156# ---------------------------------------------------------------- model157class HTRModel(nn.Module):158 def __init__(self, cfg: HTRConfig):159 super().__init__()160 self.cfg = cfg161 self.emb = nn.Embedding(cfg.vocab_size, cfg.d_model, padding_idx=cfg.pad_id)162 self.emb_scale = math.sqrt(cfg.d_model)163 if cfg.pos == "learned":164 self.pos_enc = nn.Embedding(cfg.max_len, cfg.d_model)165 self.pos_dec = nn.Embedding(cfg.max_len, cfg.d_model)166 self.enc = nn.ModuleList(EncoderLayer(cfg) for _ in range(cfg.n_enc_layers))167 self.dec = nn.ModuleList(DecoderLayer(cfg) for _ in range(cfg.n_dec_layers))168 self.enc_ln, self.dec_ln = nn.LayerNorm(cfg.d_model), nn.LayerNorm(cfg.d_model)169 self.drop = nn.Dropout(cfg.dropout)170 self.out = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)171 if cfg.tie_embeddings:172 self.out.weight = self.emb.weight173 # metadata conditioning (m2, m3)174 if cfg.arch in ("m2", "m3"):175 self.e_disc = nn.Embedding(cfg.n_disc, cfg.d_model)176 self.e_era = nn.Embedding(cfg.n_era, cfg.d_model)177 self.e_role = nn.Embedding(cfg.n_role, cfg.d_model)178 self.e_src = nn.Embedding(cfg.n_source, cfg.d_model)179 self.meta_proj = nn.Linear(cfg.d_model, cfg.d_model)180 # dual-latent (m3)181 if cfg.arch == "m3":182 K = max(2, cfg.n_styles)183 self.style_mlp = nn.Sequential(nn.Linear(2 * cfg.d_model, cfg.d_model), nn.GELU(), nn.Dropout(cfg.dropout), nn.Linear(cfg.d_model, K))184 self.style_prototypes = nn.Parameter(torch.randn(K, cfg.style_tokens, cfg.d_model) * 0.02)185 self.style_adv = nn.Sequential(nn.Linear(cfg.d_model, cfg.d_model // 2), nn.GELU(), nn.Linear(cfg.d_model // 2, K))186 self.content_proj = nn.Linear(cfg.d_model, cfg.d_model) # for contrastive / orthogonality187 self.style_vec_proj = nn.Linear(cfg.d_model, cfg.d_model)188 if cfg.use_copy:189 self.copy_gate = nn.Linear(cfg.d_model, 1)190 self.apply(self._init)191 self.n_params = sum(p.numel() for p in self.parameters())192193 @staticmethod194 def _init(m):195 if isinstance(m, nn.Linear):196 nn.init.normal_(m.weight, std=0.02)197 if m.bias is not None:198 nn.init.zeros_(m.bias)199 elif isinstance(m, nn.Embedding):200 nn.init.normal_(m.weight, std=0.02)201202 # ---- helpers203 def _rope(self, L: int, device):204 if self.cfg.pos != "rope":205 return None206 return rope_cache(L, self.cfg.d_model // self.cfg.n_heads, device)207208 def _meta_vec(self, meta: dict[str, torch.Tensor] | None, B: int, device) -> torch.Tensor | None:209 if self.cfg.arch not in ("m2", "m3") or meta is None:210 return None211 v = self.e_disc(meta["disc"]) + self.e_era(meta["era"]) + self.e_role(meta["role"]) + self.e_src(meta["src"])212 return self.meta_proj(v) # (B, D)213214 def encode(self, src: torch.Tensor, meta: dict | None = None) -> dict:215 B, L = src.shape216 pad = src.eq(self.cfg.pad_id)217 x = self.emb(src) * self.emb_scale218 if self.cfg.pos == "learned":219 x = x + self.pos_enc(torch.arange(L, device=src.device))[None]220 mv = self._meta_vec(meta, B, src.device)221 if mv is not None:222 x = x + mv[:, None, :]223 x = self.drop(x)224 rc = self._rope(L, src.device)225 for layer in self.enc:226 x = layer(x, pad, rc)227 h = self.enc_ln(x)228 # masked mean + max pooling → content summary229 m = (~pad).float().unsqueeze(-1)230 mean = (h * m).sum(1) / m.sum(1).clamp(min=1.0)231 mx = h.masked_fill(pad.unsqueeze(-1), -1e4).max(1).values232 out = {"mem": h, "mem_pad": pad, "pooled": torch.cat([mean, mx], -1), "meta_vec": mv}233 if self.cfg.arch == "m3":234 content = self.content_proj(mean) # (B, D)235 style_logits = self.style_mlp(out["pooled"] + (torch.cat([mv, mv], -1) if mv is not None else 0))236 style_probs = style_logits.softmax(-1)237 style_mem = torch.einsum("bk,ktd->btd", style_probs, self.style_prototypes) # (B, T, D)238 style_vec = self.style_vec_proj(style_mem.mean(1))239 adv_logits = self.style_adv(grad_reverse(content, self.cfg.grl_lambda)) # adversary: predict style from content240 out.update({"content": content, "style_logits": style_logits, "style_probs": style_probs, "style_mem": style_mem, "style_vec": style_vec, "adv_logits": adv_logits})241 mem = torch.cat([style_mem, h], 1)242 mem_pad = torch.cat([torch.zeros(B, style_mem.shape[1], dtype=torch.bool, device=src.device), pad], 1)243 out["mem"], out["mem_pad"], out["n_prefix"] = mem, mem_pad, style_mem.shape[1]244 elif self.cfg.arch == "m2" and mv is not None:245 out["mem"] = torch.cat([mv[:, None, :], h], 1)246 out["mem_pad"] = torch.cat([torch.zeros(B, 1, dtype=torch.bool, device=src.device), pad], 1)247 out["n_prefix"] = 1248 else:249 out["n_prefix"] = 0250 return out251252 def decode(self, tgt_in: torch.Tensor, enc: dict, src: torch.Tensor | None = None, meta: dict | None = None, need_copy: bool = True, cache: list[dict] | None = None, offset: int = 0) -> torch.Tensor:253 """Returns log-probabilities (B, Lt, V). With copy gate: p = (1-g)·softmax(logits) + g·copy_dist.254 Incremental decoding: pass ``cache`` (one dict per layer, created empty by the caller) and feed only the new255 token(s) in ``tgt_in`` with ``offset`` = number of tokens already decoded."""256 B, Lt = tgt_in.shape257 y_pad = None if cache is not None else tgt_in.eq(self.cfg.pad_id)258 y = self.emb(tgt_in) * self.emb_scale259 if self.cfg.pos == "learned":260 y = y + self.pos_dec(torch.arange(offset, offset + Lt, device=tgt_in.device))[None]261 mv = enc.get("meta_vec")262 if mv is not None:263 y = y + mv[:, None, :]264 y = self.drop(y)265 rc = self._rope(offset + Lt, tgt_in.device)266 w_last = None267 for i, layer in enumerate(self.dec):268 need = need_copy and self.cfg.use_copy and i == len(self.dec) - 1269 y, w = layer(y, enc["mem"], y_pad, enc["mem_pad"], rc, need_weights=need, cache=(cache[i] if cache is not None else None), offset=offset)270 if w is not None:271 w_last = w272 h = self.dec_ln(y)273 logits = self.out(h)274 if self.cfg.use_copy and src is not None and w_last is not None:275 gate = torch.sigmoid(self.copy_gate(h)) # (B, Lt, 1)276 n_prefix = enc.get("n_prefix", 0)277 att = w_last[:, :, n_prefix:] # attention over source positions (B, Lt, Ls)278 att = att.masked_fill(src.eq(self.cfg.pad_id)[:, None, :], 0.0)279 att = att / att.sum(-1, keepdim=True).clamp(min=1e-9)280 copy = torch.zeros_like(logits).scatter_add(-1, src[:, None, :].expand(-1, Lt, -1), att)281 probs = (1 - gate) * logits.softmax(-1) + gate * copy282 return (probs + 1e-9).log()283 return logits.log_softmax(-1)284285 def forward(self, src: torch.Tensor, tgt_in: torch.Tensor, meta: dict | None = None) -> dict:286 enc = self.encode(src, meta)287 logp = self.decode(tgt_in, enc, src=src, meta=meta)288 return {"logp": logp, **{k: v for k, v in enc.items() if k in ("content", "style_logits", "style_probs", "style_vec", "adv_logits", "pooled")}}289290 @torch.no_grad()291 def generate(self, src: torch.Tensor, meta: dict | None = None, max_new_tokens: int = 512, beam: int = 1, length_penalty: float = 1.0, repetition_penalty: float = 1.0, no_repeat_ngram: int = 4) -> torch.Tensor:292 """Greedy (beam=1) or beam-search decoding with KV cache and the copy mechanism. Returns ids (B, L) incl. BOS/EOS.293 ``no_repeat_ngram`` blocks tokens that would complete an n-gram already present in the output (greedy path)."""294 self.eval()295 enc = self.encode(src, meta)296 B = src.shape[0]297 dev = src.device298 if beam <= 1:299 cache = [dict() for _ in self.dec]300 ys = torch.full((B, 1), self.cfg.bos_id, dtype=torch.long, device=dev)301 seqs = [[self.cfg.bos_id] for _ in range(B)] # CPU mirror of ys (one device→host transfer per step)302 ngram_index: list[dict[tuple, set[int]]] = [dict() for _ in range(B)] # (n-1)-gram prefix → next tokens seen303 done = torch.zeros(B, dtype=torch.bool, device=dev)304 cur = ys305 n = no_repeat_ngram306 for t in range(max_new_tokens):307 logp = self.decode(cur, enc, src=src, meta=meta, cache=cache, offset=t)[:, -1, :]308 if repetition_penalty != 1.0:309 logp = logp.scatter(1, ys, logp.gather(1, ys) * repetition_penalty)310 if n and t + 1 >= n:311 rows, cols = [], []312 for b in range(B):313 banned = ngram_index[b].get(tuple(seqs[b][-(n - 1):]))314 if banned:315 for tok in banned:316 if tok != self.cfg.eos_id:317 rows.append(b)318 cols.append(tok)319 if rows:320 logp[torch.tensor(rows, device=dev), torch.tensor(cols, device=dev)] = -1e9321 nxt = logp.argmax(-1)322 nxt = torch.where(done, torch.full_like(nxt, self.cfg.pad_id), nxt)323 ys = torch.cat([ys, nxt[:, None]], 1)324 cur = nxt[:, None]325 done = done | nxt.eq(self.cfg.eos_id)326 nl = nxt.tolist()327 for b in range(B):328 seqs[b].append(nl[b])329 if n and len(seqs[b]) >= n:330 ngram_index[b].setdefault(tuple(seqs[b][-n:-1]), set()).add(seqs[b][-1])331 if done.all():332 break333 return ys334 # batched beam search: expand encoder outputs to (B*beam)335 K = beam336 def rep(v):337 return v.repeat_interleave(K, 0) if torch.is_tensor(v) else v338 enc_k = {k: rep(v) for k, v in enc.items()}339 src_k = rep(src)340 cache = [dict() for _ in self.dec]341 ys = torch.full((B * K, 1), self.cfg.bos_id, dtype=torch.long, device=dev)342 scores = torch.zeros(B, K, device=dev)343 scores[:, 1:] = -1e9 # only the first beam is alive at t=0344 finished = torch.zeros(B * K, dtype=torch.bool, device=dev)345 cur = ys346 for t in range(max_new_tokens):347 logp = self.decode(cur, enc_k, src=src_k, meta=None, cache=cache, offset=t)[:, -1, :] # (B*K, V)348 V = logp.size(-1)349 # finished beams can only emit pad with zero cost350 logp = torch.where(finished[:, None], torch.full_like(logp, -1e9), logp)351 logp[finished, self.cfg.pad_id] = 0.0352 cand = (scores.view(-1, 1) + logp).view(B, K * V)353 top_s, top_i = cand.topk(K, dim=-1)354 beam_idx = top_i // V355 tok_idx = top_i % V356 flat = (torch.arange(B, device=dev)[:, None] * K + beam_idx).view(-1)357 ys = torch.cat([ys[flat], tok_idx.view(-1, 1)], 1)358 finished = finished[flat] | tok_idx.view(-1).eq(self.cfg.eos_id) | tok_idx.view(-1).eq(self.cfg.pad_id)359 scores = top_s360 for layer_cache in cache:361 for part in ("self", "cross"):362 for kk in ("k", "v"):363 if kk in layer_cache.get(part, {}):364 layer_cache[part][kk] = layer_cache[part][kk][flat]365 cur = tok_idx.view(-1, 1)366 if finished.all():367 break368 lengths = (ys.ne(self.cfg.pad_id)).sum(1).view(B, K).float()369 norm = scores / lengths.pow(length_penalty)370 best = norm.argmax(-1)371 out = ys.view(B, K, -1)[torch.arange(B, device=dev), best]372 return out373