training/train-align-encoder.py

"""© 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())