training/align_encoder.py

"""© KAND CA 2026 - Sentence embeddings from the Jisr MT encoder.

Surgery, not a new model: Emhotob-MT-50M-FINAL's encoder is 29.28M of the 49.06M
total (14.34M of that the shared 32000x448 table). Drop the decoder and lm_head,
mean-pool over the attention mask, L2 normalise. Nothing is trained here.

FLORES+ devtest, 1012-way ar->en retrieval, untrained, measured before any of
this was designed - "centred" subtracts the per-language mean before normalising:

    tag (en / ar)              raw    doc-ctr   global-ctr
    (none)                     0.134    0.217     0.163
    >>eng<< / >>ara<< (own)    0.252    0.976     0.943
    >>ara<< / >>ara<<          0.939    0.973     0.960
    >>ara<< / >>eng<<          0.916    0.970     0.908

Three decisions come out of that table, and the first two are easy to get
backwards - I did, before measuring the no-tag row properly:

1. A TAG IS NECESSARY; WHICH TAG IS A TIE-BREAK. Untagged, centring only reaches
   0.217 - the two languages are simply not in a comparable space. Tagged, every
   combination lands at 0.970-0.976 once centred, a spread of about 6 sentences
   in 1012. So the tag and centring are complementary, not substitutes: the tag
   makes the spaces comparable, centring removes what is left. Reading the raw
   column alone suggests ">>ara<< by 70 points" and that conclusion is an
   artifact of skipping the centring step.
2. CENTRE PER DOCUMENT, NOT GLOBALLY. doc beats global at every tag (0.976 vs
   0.943 at the default) and at every document size measured:

       N          4      8     16     32     64    128
       none   0.881  0.816  0.718  0.608  0.519  0.438
       doc    1.000  1.000  0.998  0.997  0.994  0.992
       global 0.998  0.999  0.994  0.990  0.984  0.976

   The worry that a 4-sentence document cannot estimate its own mean is wrong:
   the offset being removed is large and shared, so even a 4-sample estimate of
   it helps. This is not leakage - an aligner always holds the whole document,
   so per-document statistics are the real operating condition. The fitted
   global means stay in align_config.json as the fallback for one-off calls
   that have no document context.
3. LAYER 5 BEATS LAYER 6 everywhere, centred and raw. The last layer is
   specialised for decoder cross-attention rather than for representing the
   sentence, so pool from the penultimate one.

Rejected: max pooling (0.527 at L6 best) and first-token pooling (0.096) - token
0 IS the tag, so first-token pooling reads the prefix and nothing else.

The tag changes the first word's tokenisation (leading space: "The" -> "ĠThe"),
so it has to be byte-identical between fitting the means and using them. That is
why TAG lives here as one dict and is written into align_config.json, never
hard-coded at a second call site.
"""
import argparse
import json
import os
import warnings

warnings.filterwarnings("ignore")

import numpy as np
import torch
from transformers import AutoTokenizer, MarianMTModel

TAG = {"en": ">>eng<<", "ar": ">>ara<<"}     # tag = the text's own language
POOL_LAYER = 5                               # penultimate, measured
CFG_NAME = "align_config.json"


def load_encoder(path, device="cuda"):
    """Return (tok, encoder, cfg). Works for the raw MT model and a trained one."""
    tok = AutoTokenizer.from_pretrained(path)
    full = MarianMTModel.from_pretrained(path)
    enc = full.model.encoder
    del full.model.decoder, full.lm_head
    enc = enc.to(device).eval()
    cfg = {"tag": dict(TAG), "pool_layer": POOL_LAYER, "mu": None}
    p = os.path.join(path, CFG_NAME)
    if os.path.exists(p):
        with open(p) as f:
            cfg.update(json.load(f))
    return tok, enc, cfg


@torch.no_grad()
def encode_raw(tok, enc, texts, lang, bs=256, layer=POOL_LAYER, tag=None,
               max_length=256):
    """Masked mean-pooled hidden states. No centring, no normalisation."""
    tag = TAG[lang] if tag is None else tag
    out = []
    for i in range(0, len(texts), bs):
        b = texts[i:i + bs]
        b = [f"{tag} {t}" for t in b] if tag else list(b)
        e = tok(b, return_tensors="pt", padding=True, truncation=True,
                max_length=max_length).to(enc.device)
        h = enc(**e, output_hidden_states=True).hidden_states[layer]
        m = e["attention_mask"].unsqueeze(-1).to(h.dtype)
        out.append(((h * m).sum(1) / m.sum(1)).float().cpu().numpy())
    return np.concatenate(out) if out else np.zeros((0, enc.config.d_model),
                                                    dtype=np.float32)


def normalize(x):
    n = np.linalg.norm(x, axis=1, keepdims=True)
    return x / np.maximum(n, 1e-9)


def embed(tok, enc, texts, lang, cfg=None, bs=256, center="doc", mu=None,
          layer=None, max_length=256):
    """Embeddings ready for cosine: centred then L2-normalised.

    center="global" uses the fitted per-language mean (the default, and the only
    mode that is safe on short documents); "doc" re-estimates from these texts,
    which is better when a document is long and hopeless when it is four
    sentences; "none" disables it and is there so the ablation is runnable.
    """
    cfg = cfg or {}
    layer = layer if layer is not None else cfg.get("pool_layer", POOL_LAYER)
    tag = (cfg.get("tag") or TAG).get(lang)
    x = encode_raw(tok, enc, texts, lang, bs=bs, layer=layer, tag=tag,
                   max_length=max_length)
    if len(x) == 0:
        return x
    if center == "doc":
        x = x - x.mean(0, keepdims=True)
    elif center == "global":
        m = mu if mu is not None else (cfg.get("mu") or {}).get(lang)
        if m is None:
            raise ValueError(
                f"center='global' needs a fitted mean for {lang!r}; run "
                f"align_encoder.py --fit-means, or pass center='doc'")
        x = x - np.asarray(m, dtype=np.float32)
    return normalize(x)


def fit_means(tok, enc, pairs, cfg=None, bs=256, layer=None):
    """Per-language mean vectors, from unit-normalised embeddings.

    Normalise first, then average: the mean of raw vectors is dominated by
    whichever sentences happen to have the largest norm, and norm tracks length.
    """
    cfg = cfg or {}
    layer = layer if layer is not None else cfg.get("pool_layer", POOL_LAYER)
    mu = {}
    for lang in ("en", "ar"):
        x = normalize(encode_raw(tok, enc, pairs[lang], lang, bs=bs, layer=layer,
                                 tag=(cfg.get("tag") or TAG)[lang]))
        mu[lang] = x.mean(0).astype(np.float32).tolist()
    return mu


def _flores(split="devtest"):
    from datasets import load_dataset
    ds = load_dataset("openlanguagedata/flores_plus", split=split)
    en, ar = {}, {}
    for r in ds:
        # iso_15924 guard: arb_Latn is also in this split, 1012 rows of
        # ROMANISED Arabic. Matching on iso_639_3 alone silently swaps it in
        # and drops P@1 from 0.87 to 0.07.
        if r["iso_639_3"] == "eng" and r["iso_15924"] == "Latn":
            en[r["id"]] = r["text"]
        elif r["iso_639_3"] == "arb" and r["iso_15924"] == "Arab":
            ar[r["id"]] = r["text"]
    ids = sorted(set(en) & set(ar))
    return [en[i] for i in ids], [ar[i] for i in ids]


def p_at_k(xe, xa, k=1):
    s = xa @ xe.T
    idx = np.argsort(-s, axis=1)[:, :k]
    return float(np.mean([i in row for i, row in enumerate(idx)]))


def _msa_pairs(n):
    from datasets import load_dataset
    ds = load_dataset("oddadmix/quick-mt-en-ar-5m", split="train")
    ds = ds.select(range(min(n * 4, len(ds))))
    en, ar = [], []
    for msgs in ds["messages"]:
        m = {x["role"]: x["content"] for x in msgs}
        if m.get("system") == "Translate to Arabic":
            s, t = m.get("user"), m.get("assistant")
        elif m.get("system") == "Translate to English":
            t, s = m.get("user"), m.get("assistant")
        else:
            continue
        if s and t:
            en.append(s.strip())
            ar.append(t.strip())
        if len(en) >= n:
            break
    return {"en": en, "ar": ar}


def probe(tok, enc, cfg, args):
    EN, AR = _flores(args.split)
    print(f"[*] FLORES+ {args.split}: {len(EN)} pairs", flush=True)
    print(f"\n{'layer':>5} {'en tag':>9} {'ar tag':>9} {'centre':>7} "
          f"{'P@1':>7} {'P@10':>7}")
    print("-" * 50)
    # "" is the explicit no-tag sentinel. None must NOT be used here: encode_raw
    # reads None as "fall back to TAG[lang]", so a (None, None) row silently
    # re-measures the default and reports it as the untagged baseline.
    combos = [(">>eng<<", ">>ara<<"), (">>ara<<", ">>ara<<"),
              (">>ara<<", ">>eng<<"), ("", "")]
    mu = (cfg.get("mu") or {})
    for layer in args.layers:
        for te, ta in combos:
            xe = encode_raw(tok, enc, EN, "en", args.bs, layer, te)
            xa = encode_raw(tok, enc, AR, "ar", args.bs, layer, ta)
            for mode in ("none", "doc", "global"):
                if mode == "none":
                    a, b = normalize(xe), normalize(xa)
                elif mode == "doc":
                    a = normalize(xe - xe.mean(0, keepdims=True))
                    b = normalize(xa - xa.mean(0, keepdims=True))
                else:
                    if not mu:
                        continue
                    a = normalize(normalize(xe) - np.asarray(mu["en"]))
                    b = normalize(normalize(xa) - np.asarray(mu["ar"]))
                print(f"{layer:>5} {str(te):>9} {str(ta):>9} {mode:>7} "
                      f"{p_at_k(a,b,1):>7.3f} {p_at_k(a,b,10):>7.3f}", flush=True)
        print("-" * 50)

    # the question the retrieval table cannot answer: per-document centring is
    # estimated from N sentences, and real documents are short.
    if args.doc_sizes:
        print(f"\n[*] per-doc centring vs document size (mean over samples)")
        print(f"{'N':>5} {'none':>8} {'doc':>8} {'global':>8}")
        print("-" * 32)
        xe = encode_raw(tok, enc, EN, "en", args.bs, cfg.get("pool_layer", 5),
                        TAG["en"])
        xa = encode_raw(tok, enc, AR, "ar", args.bs, cfg.get("pool_layer", 5),
                        TAG["ar"])
        rng = np.random.default_rng(1234)
        for N in args.doc_sizes:
            acc = {"none": [], "doc": [], "global": []}
            for _ in range(200):
                i = rng.choice(len(EN), size=min(N, len(EN)), replace=False)
                e, a = xe[i], xa[i]
                acc["none"].append(p_at_k(normalize(e), normalize(a)))
                acc["doc"].append(p_at_k(normalize(e - e.mean(0, keepdims=True)),
                                         normalize(a - a.mean(0, keepdims=True))))
                if mu:
                    acc["global"].append(p_at_k(
                        normalize(normalize(e) - np.asarray(mu["en"])),
                        normalize(normalize(a) - np.asarray(mu["ar"]))))
            g = f"{np.mean(acc['global']):>8.3f}" if mu else f"{'-':>8}"
            print(f"{N:>5} {np.mean(acc['none']):>8.3f} "
                  f"{np.mean(acc['doc']):>8.3f} {g}", flush=True)
        print("-" * 32)
        print("[*] Retrieval within a document is the aligner's actual job, so")
        print("[*] these rows - not the 1012-way table - decide the default.")


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="./Emhotob-MT-50M-FINAL")
    p.add_argument("--probe", action="store_true")
    p.add_argument("--fit-means", action="store_true")
    p.add_argument("--fit-n", type=int, default=100000)
    p.add_argument("--split", default="devtest")
    p.add_argument("--bs", type=int, default=256)
    p.add_argument("--layers", type=int, nargs="+", default=[4, 5, 6])
    p.add_argument("--doc-sizes", type=int, nargs="*",
                   default=[4, 8, 16, 32, 64, 128])
    a = p.parse_args()

    tok, enc, cfg = load_encoder(a.model)
    print(f"[*] {a.model}: encoder "
          f"{sum(x.numel() for x in enc.parameters())/1e6:.2f}M params, "
          f"d_model {enc.config.d_model}, layer {cfg.get('pool_layer')}",
          flush=True)

    if a.fit_means:
        print(f"[*] fitting per-language means on {a.fit_n:,} MSA pairs",
              flush=True)
        cfg["mu"] = fit_means(tok, enc, _msa_pairs(a.fit_n), cfg, a.bs)
        out = os.path.join(a.model, CFG_NAME)
        with open(out, "w") as f:
            json.dump({"tag": cfg["tag"], "pool_layer": cfg["pool_layer"],
                       "mu": cfg["mu"], "fit_n": a.fit_n}, f)
        print(f"[+] wrote {out}", flush=True)

    if a.probe:
        probe(tok, enc, cfg, a)


if __name__ == "__main__":
    main()