Upload src/xscript/eval/bts.py with huggingface_hub

src/xscript/eval/bts.py

@@ -0,0 +1,136 @@
+"""Bilingual Transfer Score and the thesis's interaction estimate.
+
+BTS(L) = (BPB_mono(L) - BPB_bi(L)) / BPB_mono(L)   (positive = bilingual helps)
+
+Two matchings (the plan asks for both):
+  - matched-total   : mono-final vs bilingual-final (both at the same TOTAL
+                      tokens). This is ATLAS's framing; under a starved tokenizer
+                      the cross-script partner saw fewer of its own tokens.
+  - matched-lang    : bilingual-final vs the mono checkpoint at the same number
+                      of *that-language* tokens the bilingual run actually saw
+                      (= total x mixing-prob). Removes the token-count confound.
+
+Headline interaction (the untested contribution):
+  penalty(C) = mean BTS(same-script) - mean BTS(cross-script)   under tokenizer C
+  interaction = penalty(starved) - penalty(destarved)
+A large positive interaction => the cross-script "penalty" is largely a
+tokenizer starvation artifact, not intrinsic script transfer.
+
+Run naming convention (see configs/matrix.py): "<mix>__<tok_name>", where mix is
+one language ("en") or "en-<partner>", and tok_name is e.g. "bl_destarved".
+"""
+import json
+from pathlib import Path
+
+from ..langs import ANCHOR, PARTNERS, LANGS
+from ..paths import RUNS, RESULTS, ensure
+
+
+def _read_evals(name: str) -> list[dict]:
+    """[(tokens, {source: bpb})...] sorted by tokens, from a run's train.jsonl."""
+    log = RUNS / name / "train.jsonl"
+    if not log.exists():
+        return []
+    series = {}
+    for line in log.read_text().splitlines():
+        try:
+            rec = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+        ev = rec.get("eval") or rec.get("eval_final")
+        if ev:
+            series[rec["tokens"]] = {s: v["bpb"] for s, v in ev.items()}
+    return [{"tokens": t, "bpb": series[t]} for t in sorted(series)]
+
+
+def _final_bpb(name: str, source: str) -> float | None:
+    ev = _read_evals(name)
+    for rec in reversed(ev):
+        if source in rec["bpb"]:
+            return rec["bpb"][source]
+    return None
+
+
+def _bpb_at_lang_tokens(name: str, source: str, lang_tokens: float) -> float | None:
+    """Mono BPB at the checkpoint closest to `lang_tokens` (matched-lang)."""
+    best, bd = None, None
+    for rec in _read_evals(name):
+        if source in rec["bpb"]:
+            d = abs(rec["tokens"] - lang_tokens)
+            if bd is None or d < bd:
+                bd, best = d, rec["bpb"][source]
+    return best
+
+
+def compute(tok_name: str, source_kind: str = "flores",
+            mix_prob: float = 0.5, total_tokens: float = 30e9) -> dict:
+    """BTS for every partner under one tokenizer condition."""
+    rows = {}
+    for p in PARTNERS:
+        src = f"{source_kind}_{p}"
+        mono, bi = f"{p}__{tok_name}", f"{ANCHOR}-{p}__{tok_name}"
+        bpb_mono = _final_bpb(mono, src)
+        bpb_bi = _final_bpb(bi, src)
+        entry = {"same_script": LANGS[p].same_script_as_en,
+                 "bpb_mono_final": bpb_mono, "bpb_bi_final": bpb_bi}
+        if bpb_mono and bpb_bi:
+            entry["bts_matched_total"] = (bpb_mono - bpb_bi) / bpb_mono
+            bpb_mono_lang = _bpb_at_lang_tokens(mono, src, total_tokens * mix_prob)
+            if bpb_mono_lang:
+                entry["bpb_mono_at_lang_tokens"] = bpb_mono_lang
+                entry["bts_matched_lang"] = (bpb_mono_lang - bpb_bi) / bpb_mono_lang
+        rows[p] = entry
+    return rows
+
+
+def _penalty(rows: dict, key: str) -> float | None:
+    same = [r[key] for r in rows.values() if r.get("same_script") and key in r]
+    cross = [r[key] for r in rows.values() if not r.get("same_script") and key in r]
+    if not same or not cross:
+        return None
+    return sum(same) / len(same) - sum(cross) / len(cross)
+
+
+def run(flavor: str = "unigram", source_kind: str = "flores",
+        total_tokens: float = 30e9, mix_prob: float = 0.5,
+        out_dir: Path | None = None) -> dict:
+    out_dir = ensure(Path(out_dir) if out_dir else RESULTS / "bts")
+    conds = {c: compute(f"{flavor}_{c}", source_kind, mix_prob, total_tokens)
+             for c in ("starved", "destarved")}
+    inter = {}
+    for key in ("bts_matched_total", "bts_matched_lang"):
+        ps = {c: _penalty(conds[c], key) for c in conds}
+        if ps["starved"] is not None and ps["destarved"] is not None:
+            inter[key] = {"penalty_starved": ps["starved"],
+                          "penalty_destarved": ps["destarved"],
+                          "interaction": ps["starved"] - ps["destarved"]}
+    result = {"flavor": flavor, "source": source_kind, "by_condition": conds,
+              "interaction": inter}
+    (out_dir / f"bts_{flavor}_{source_kind}.json").write_text(json.dumps(result, indent=2))
+
+    md = [f"# BTS ({flavor}, eval on {source_kind})", ""]
+    for c, rows in conds.items():
+        md += [f"## {c}", "",
+               "| partner | script | BPB mono | BPB bi | BTS (total) | BTS (lang) |",
+               "|---|---|---|---|---|---|"]
+        for p, r in rows.items():
+            md.append("| {} | {} | {} | {} | {} | {} |".format(
+                p, "same" if r["same_script"] else "cross",
+                _f(r.get("bpb_mono_final")), _f(r.get("bpb_bi_final")),
+                _f(r.get("bts_matched_total")), _f(r.get("bts_matched_lang"))))
+        md.append("")
+    md += ["## Interaction (same-script penalty - cross-script penalty)", ""]
+    for key, v in inter.items():
+        md.append(f"- **{key}**: penalty(starved)={v['penalty_starved']:.4f}, "
+                  f"penalty(destarved)={v['penalty_destarved']:.4f}, "
+                  f"**interaction={v['interaction']:.4f}**")
+    md += ["", "> interaction >> 0  =>  cross-script penalty is a tokenizer-"
+           "starvation artifact.", "> interaction ~ 0  =>  penalty persists "
+           "under a fair tokenizer (genuine script effect)."]
+    (out_dir / f"bts_{flavor}_{source_kind}.md").write_text("\n".join(md) + "\n")
+    print(f"[bts] wrote {out_dir}/bts_{flavor}_{source_kind}.md")
+    return result
+
+
+def _f(x):
+    return f"{x:.4f}" if isinstance(x, float) else "-"