"""© KAND CA 2026 - Contrastive finetune of the Jisr MT encoder for alignment. The warm start already works: 0.976 FLORES P@1 and 0.805 strict alignment F1 with nothing trained at all. So this run is an improvement step on a strong system, not a way to make one, and it ships only if it beats BOTH of those. Anything less and the right answer is to keep the untrained encoder. Symmetric InfoNCE over in-batch negatives, warm-started from Emhotob-MT-50M-FINAL's encoder. Only the encoder and the shared embedding table train; the decoder rides along so the checkpoint stays loadable by align_encoder.load_encoder(), which reconstructs a MarianMTModel and takes .model.encoder from it. The decoder's weights are stale in the saved model and that is fine - nothing ever runs it - but it does mean this directory is NOT an MT model any more. It is written to a separate path for that reason. Three decisions that are not obvious, in the order they bite: 1. THE POOLING MUST MATCH INFERENCE, INCLUDING THE LAYER. Training on the last layer and pooling from layer 5 at inference optimises a representation nobody reads. hidden_states[5] here, POOL_LAYER there, one constant in align_encoder. 2. CENTRE THE BATCH. Every measured number in this project comes from centred embeddings - it is worth +73 points at the worst tag and +3.4 at the best - so training on uncentred ones optimises a space the aligner never sees. The honest caveat: inference centres over ONE DOCUMENT, and a batch of 512 random pairs is not a document, so what this actually reproduces is closer to global centring. Closer to the operating condition than not centring at all, which is the argument for it. --center none runs the ablation. 3. THE TEMPERATURE IS LEARNABLE AND CLAMPED. Free, it runs away to a degenerate scale; fixed, it is one more thing to tune. CLIP's answer - a learnable log-scale clamped at 100 - is the one borrowed here, initialised at 1/0.05. Hard-negative mining is deliberately absent from this first run. The plan calls for it after epoch 1, and epoch 1 is the whole run: the Dialects2En ablation scored 3 epochs at -2.63 wavg chrF++ against v1 while 1 epoch scored +1.23, so with a warm start this good, over-training is the likelier failure mode. Mining becomes a measured second run, not an assumption baked into the first. Data is oddadmix/quick-mt-en-ar-5m alone. UN is available and excluded on purpose - it cost this project 3.65 BLEU on MT, and every scale lever here has lost while every signal lever has won. It is a separate ablation, not a default. """ import argparse import json import math import os import warnings warnings.filterwarnings("ignore") import numpy as np import torch import torch.nn.functional as F from torch.utils.data import DataLoader from transformers import (AutoTokenizer, MarianMTModel, get_cosine_schedule_with_warmup) import align_dp import align_encoder from align_encoder import TAG, POOL_LAYER, _flores, normalize, p_at_k def pairs_dataset(name, n_proc=8): """{"en","ar"} columns, cached by datasets so this pass happens once. The rows are chat messages, not columns, and the direction is carried in the system prompt - "Translate to Arabic" means user=en, the other way round means user=ar. Reading user/assistant positionally puts half the corpus in backwards, which in a symmetric loss is not an error you can see in the curve: it just quietly trains the model that en->en is a valid pair. """ from datasets import load_dataset ds = load_dataset(name, split="train") def conv(b): en, ar = [], [] for msgs in b["messages"]: m = {x["role"]: x["content"] for x in msgs} s = m.get("system") if s == "Translate to Arabic": e, a = m.get("user"), m.get("assistant") elif s == "Translate to English": a, e = m.get("user"), m.get("assistant") else: e = a = None en.append((e or "").strip()) ar.append((a or "").strip()) return {"en": en, "ar": ar} ds = ds.map(conv, batched=True, num_proc=n_proc, remove_columns=ds.column_names, desc="extract pairs") return ds.filter(lambda b: [bool(x) and bool(y) for x, y in zip(b["en"], b["ar"])], batched=True, num_proc=n_proc, desc="drop empties") class Collate: def __init__(self, tok, max_len): self.tok, self.max_len = tok, max_len def __call__(self, rows): out = [] for lang in ("en", "ar"): txt = [f"{TAG[lang]} {r[lang]}" for r in rows] out.append(self.tok(txt, return_tensors="pt", padding=True, truncation=True, max_length=self.max_len)) return out def pool(enc, batch, layer, center): h = enc(**batch, output_hidden_states=True).hidden_states[layer] m = batch["attention_mask"].unsqueeze(-1).to(h.dtype) x = (h * m).sum(1) / m.sum(1) x = x.float() if center == "batch": x = x - x.mean(0, keepdim=True) return F.normalize(x, dim=-1) @torch.no_grad() def flores_p1(tok, enc, EN, AR, layer, bs=256): """The 1012-way retrieval number, computed exactly as the 0.976 was. Centred over the whole 1012, which is what "doc" centring reduces to when the set IS the document. Keeping it identical is the point - a P@1 measured a different way cannot be compared to the baseline it has to beat. """ was = enc.training enc.eval() xe = align_encoder.encode_raw(tok, enc, EN, "en", bs=bs, layer=layer) xa = align_encoder.encode_raw(tok, enc, AR, "ar", bs=bs, layer=layer) enc.train(was) a = normalize(xe - xe.mean(0, keepdims=True)) b = normalize(xa - xa.mean(0, keepdims=True)) return p_at_k(a, b, 1) @torch.no_grad() def dev_f1(tok, enc, docs, del_cost, layer, k=4, bs=256): """Strict alignment F1 on held-out dev gold - the model-selection metric. P@1 cannot do this job. It is 0.9763 at the warm start and hits 1.0000 by step 1000 of 19511, so after the first twentieth of an epoch every later checkpoint ties and `p1 > best` never fires again: selection silently freezes on step 1000 and the rest of the run is discarded. A saturated metric does not mean the model stopped improving, only that this ruler ran out. Strict F1 is what the ship/no-ship gate is written in anyway, and it sits at 0.805/0.739 with room in both directions. Uses DEV gold only (flores-dev, tier-B dev half) so the reported devtest and tier-B test numbers stay clean. """ was = enc.training enc.eval() cfg = {"pool_layer": layer, "tag": dict(TAG)} tp = fp = fn = 0 for d in docs: ar, en = d["ar"], d["en"] kk = min(k, len(ar), len(en)) if kk < 1: continue Xs = {n: align_encoder.embed( tok, enc, [" ".join(ar[i:i + n]) for i in range(len(ar) - n + 1)], "ar", cfg=cfg, bs=bs, center="doc") for n in range(1, kk + 1)} Xt = {n: align_encoder.embed( tok, enc, [" ".join(en[j:j + n]) for j in range(len(en) - n + 1)], "en", cfg=cfg, bs=bs, center="doc") for n in range(1, kk + 1)} links = align_dp.align(Xs, Xt, len(ar), len(en), k=kk, del_cost=del_cost, nm_pow=align_dp.NM_POW) from collections import Counter g = Counter((tuple(x["ar"]), tuple(x["en"])) for x in d["links"] if not x.get("swapped", False)) p = Counter((tuple(s), tuple(t)) for s, t, _ in links) for key_, c in g.items(): h = min(c, p.get(key_, 0)) tp += h fn += c - h for key_, c in p.items(): fp += max(0, c - g.get(key_, 0)) enc.train(was) pr = tp / (tp + fp) if tp + fp else 0.0 rc = tp / (tp + fn) if tp + fn else 0.0 return 2 * pr * rc / (pr + rc) if pr + rc else 0.0 def save(model, tok, out, cfg_layer): os.makedirs(out, exist_ok=True) model.save_pretrained(out, safe_serialization=True) tok.save_pretrained(out) with open(os.path.join(out, align_encoder.CFG_NAME), "w") as f: # mu stays null: this checkpoint's embedding space is not the one the # old means were fitted in, and a stale mean is worse than none - it # silently shifts every vector. center="doc" needs no mu; refit with # align_encoder.py --fit-means if center="global" is ever wanted. json.dump({"tag": dict(TAG), "pool_layer": cfg_layer, "mu": None}, f) def main(): p = argparse.ArgumentParser() p.add_argument("--base", default="./Emhotob-MT-50M-FINAL") p.add_argument("--out", default="./Jisr-Align-29M") p.add_argument("--final", default="./Jisr-Align-29M-FINAL") p.add_argument("--dataset", default="oddadmix/quick-mt-en-ar-5m") p.add_argument("--bs", type=int, default=512) p.add_argument("--epochs", type=int, default=1) p.add_argument("--lr", type=float, default=3e-4) p.add_argument("--warmup-ratio", type=float, default=0.03) p.add_argument("--weight-decay", type=float, default=0.01) p.add_argument("--max-grad-norm", type=float, default=1.0) p.add_argument("--max-len", type=int, default=128) p.add_argument("--temp", type=float, default=0.05) p.add_argument("--center", default="batch", choices=["batch", "none"]) p.add_argument("--layer", type=int, default=POOL_LAYER) p.add_argument("--dev-gold", nargs="+", default=["align-gold/flores-dev-d5-heavy.jsonl", "align-gold/tierb-dialect-dev.jsonl"]) # del_cost is corpus-dependent (see align_dp.DEL_COST): 0.45 for gold with # real omissions, 0.70 for the near-parallel machine-translated tier B. One # value for both would measure each set at the other's operating point. p.add_argument("--dev-del-cost", type=float, nargs="+", default=[0.45, 0.70]) p.add_argument("--eval-steps", type=int, default=1000) p.add_argument("--log-steps", type=int, default=50) p.add_argument("--limit", type=int, default=0, help="0 = all pairs") p.add_argument("--workers", type=int, default=8) p.add_argument("--seed", type=int, default=1234) a = p.parse_args() torch.manual_seed(a.seed) dev = "cuda" tok = AutoTokenizer.from_pretrained(a.base) model = MarianMTModel.from_pretrained(a.base).to(dev) enc = model.model.encoder enc.train() # The decoder is dead weight here and must not collect gradients - it would # cost memory and, through the tied embedding table, let decoder-side # gradients into a table this objective is supposed to own. for prm in model.model.decoder.parameters(): prm.requires_grad_(False) trainable = [prm for prm in enc.parameters() if prm.requires_grad] logit_scale = torch.nn.Parameter( torch.tensor(math.log(1.0 / a.temp), device=dev)) ds = pairs_dataset(a.dataset) if a.limit: ds = ds.select(range(min(a.limit, len(ds)))) dl = DataLoader(ds, batch_size=a.bs, shuffle=True, drop_last=True, num_workers=a.workers, collate_fn=Collate(tok, a.max_len), pin_memory=True, persistent_workers=a.workers > 0) steps = len(dl) * a.epochs opt = torch.optim.AdamW([{"params": trainable, "weight_decay": a.weight_decay}, {"params": [logit_scale], "weight_decay": 0.0}], lr=a.lr, betas=(0.9, 0.98), fused=True) sch = get_cosine_schedule_with_warmup( opt, int(a.warmup_ratio * steps), steps) EN, AR = _flores("devtest") dev_sets = [ (os.path.basename(f).replace("-heavy.jsonl", "").replace(".jsonl", ""), [json.loads(l) for l in open(f)], a.dev_del_cost[min(i, len(a.dev_del_cost) - 1)]) for i, f in enumerate(a.dev_gold)] def dev_scores(): return [(nm, dev_f1(tok, enc, docs, dc, a.layer)) for nm, docs, dc in dev_sets] base_p1 = flores_p1(tok, enc, EN, AR, a.layer) base_dev = dev_scores() base_sel = sum(f for _, f in base_dev) / len(base_dev) print(f"[*] {len(ds):,} pairs, bs {a.bs}, {steps:,} steps, " f"{sum(x.numel() for x in trainable)/1e6:.2f}M trainable", flush=True) print(f"[*] warm start: P@1 {base_p1:.4f}, " + ", ".join(f"{nm} F1 {f:.4f}" for nm, f in base_dev) + f" -> select on mean dev F1 {base_sel:.4f}", flush=True) best, step = base_sel, 0 for ep in range(a.epochs): for be, ba in dl: be = {k: v.to(dev, non_blocking=True) for k, v in be.items()} ba = {k: v.to(dev, non_blocking=True) for k, v in ba.items()} with torch.autocast("cuda", dtype=torch.bfloat16): xe = pool(enc, be, a.layer, a.center) xa = pool(enc, ba, a.layer, a.center) scale = logit_scale.exp().clamp(max=100.0) logits = scale * xe @ xa.T tgt = torch.arange(len(xe), device=dev) loss = 0.5 * (F.cross_entropy(logits, tgt) + F.cross_entropy(logits.T, tgt)) loss.backward() torch.nn.utils.clip_grad_norm_(trainable, a.max_grad_norm) opt.step() sch.step() opt.zero_grad(set_to_none=True) step += 1 if step % a.log_steps == 0: print(f"[*] ep {ep} step {step}/{steps} loss {loss.item():.4f} " f"temp {1/scale.item():.4f} lr {sch.get_last_lr()[0]:.2e}", flush=True) if step % a.eval_steps == 0: sc = dev_scores() sel = sum(f for _, f in sc) / len(sc) mark = "" if sel > best: best = sel save(model, tok, a.out, a.layer) mark = " <- best, saved" print(f"[+] step {step} P@1 {flores_p1(tok, enc, EN, AR, a.layer):.4f} " + " ".join(f"{nm} {f:.4f}" for nm, f in sc) + f" | mean {sel:.4f} (best {best:.4f}, " f"warm start {base_sel:.4f}){mark}", flush=True) sc = dev_scores() sel = sum(f for _, f in sc) / len(sc) print(f"[+] end of training: P@1 {flores_p1(tok, enc, EN, AR, a.layer):.4f} " + " ".join(f"{nm} {f:.4f}" for nm, f in sc) + f" | mean {sel:.4f}, best {best:.4f}, warm start {base_sel:.4f}", flush=True) if sel > best: best = sel save(model, tok, a.out, a.layer) if best <= base_sel: # Not a crash, and not something to paper over: the run finished and # lost. Write nothing to --final so the eval cannot silently pick up a # worse encoder than the one already shipping. print(f"[!] mean dev F1 {best:.4f} does not beat the warm start " f"{base_sel:.4f} - NOT writing {a.final}", flush=True) return 1 import shutil if os.path.exists(a.final): shutil.rmtree(a.final) shutil.copytree(a.out, a.final) print(f"[+] {a.final}: mean dev F1 {best:.4f} vs warm start {base_sel:.4f}", flush=True) print(f"[+] next: .venv/bin/python eval-align.py --encoder {a.final} " f"--gold align-gold/flores-devtest-d5-heavy.jsonl " f"align-gold/tierb-dialect-test.jsonl", flush=True) return 0 if __name__ == "__main__": raise SystemExit(main())