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