src/fertility.py

#!/usr/bin/env python3
"""
Tokenizer fertility across the 14 Indic languages.

WHY THIS IS THE FIRST EXPERIMENT
--------------------------------
Multilingual tokenizers fragment Indic scripts far more aggressively than Latin.
If fertility is ~2x for Hindi, then a fixed 256-token chunk budget cuts an
English passage into 1 chunk and the SAME passage in Hindi into 2 -- so chunk
boundaries stop being comparable across languages, canonical cross-lingual chunk
IDs break, and retrieval is silently biased by language.

MSMARCO-XI is a PARALLEL corpus: every passage exists in English and in each
Indic language. That makes fertility directly measurable rather than estimated --
we compare token counts for the *same content*.

OUTPUT
  results/fertility.json  - per-language stats
  results/fertility.png   - the plot (chunk-count skew is the money panel)

DECISION THIS DRIVES
  If max/min fertility ratio > 1.3, do NOT budget chunks in token space.
  Chunk in word/sentence space so boundaries are language-invariant, and derive
  per-language token bounds from the measured fertility.
"""
from __future__ import annotations

import argparse
import json
import os
from collections import defaultdict
from pathlib import Path

# ISO codes present in ai4bharat/MSMARCO-XI
LANGS = {
    "as": "Assamese", "bn": "Bengali", "gu": "Gujarati", "hi": "Hindi",
    "kn": "Kannada", "ml": "Malayalam", "mr": "Marathi", "ne": "Nepali",
    "or": "Odia", "pa": "Punjabi", "sa": "Sanskrit", "ta": "Tamil",
    "te": "Telugu", "ur": "Urdu",
}


import sys
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from src.schema_utils import iter_passages, norm_lang, default_root, load_report  # noqa: E402


def load_pairs(root: Path, n_per_lang: int) -> tuple[dict[str, list[tuple[str, str]]], str]:
    """Return ({lang: [(english, translated), ...]}, level).

    Prefers PASSAGE-level pairs (long text, most reliable). Falls back to
    QUERY-level pairs (Eng_Query vs query) when the passages struct carries no
    English counterpart -- still a valid parallel measurement, just shorter text,
    so it needs more samples.
    """
    import polars as pl

    report = load_report(root)
    fmap, pmap = report["field_mapping"], report["passage_mapping"]
    pcol, langcol = fmap.get("passages"), fmap.get("lang")
    t_key, en_key = pmap.get("text"), pmap.get("text_en")

    use_passages = bool(pcol and t_key and en_key)
    if use_passages:
        level = "passage"
        cols = [c for c in (pcol, langcol) if c]
        min_words = 15
    else:
        level = "query"
        q_col, qen_col = fmap.get("query"), fmap.get("query_en")
        if not (q_col and qen_col):
            raise SystemExit(
                "\nCannot find a parallel text pair.\n"
                f"  passages: {pcol}  text: {t_key}  text_en: {en_key}\n"
                f"  query: {q_col}  query_en: {qen_col}\n"
                f"  inner keys seen: {report.get('passage_inner_keys')}\n"
                "Send me schema_report.json -> passage_inner_keys / passage_inner_stats.\n"
            )
        print(f"  note: no English passage field ({en_key=}); "
              "falling back to QUERY-level pairs")
        cols = [c for c in (q_col, qen_col, langcol) if c]
        min_words = 3
        n_per_lang = max(n_per_lang, 5000)   # shorter text -> need more samples

    pairs: dict[str, list[tuple[str, str]]] = defaultdict(list)
    n_failed = 0

    for fp in report["files"]:
        if len(pairs) >= len(LANGS) and all(len(v) >= n_per_lang for v in pairs.values()):
            break
        try:
            df = pl.read_parquet(fp, columns=cols, n_rows=60_000)
        except Exception as exc:
            n_failed += 1
            if n_failed <= 3:      # show the first few instead of hiding all of them
                print(f"    !! cannot read {Path(fp).name}: {exc}")
            continue

        langs_col = df[langcol].to_list() if langcol else [None] * len(df)

        if use_passages:
            for row, lang in zip(df[pcol].to_list(), langs_col):
                lang = norm_lang(lang)
                if lang not in LANGS or len(pairs[lang]) >= n_per_lang:
                    continue
                for _i, tr, en, _sel, _u in iter_passages(row, t_key, en_key, None, None):
                    if isinstance(en, str) and isinstance(tr, str) and len(en.split()) >= min_words:
                        pairs[lang].append((en, tr))
                        if len(pairs[lang]) >= n_per_lang:
                            break
        else:
            q_col, qen_col = fmap["query"], fmap["query_en"]
            for tr, en, lang in zip(df[q_col].to_list(), df[qen_col].to_list(), langs_col):
                lang = norm_lang(lang)
                if lang not in LANGS or len(pairs[lang]) >= n_per_lang:
                    continue
                if isinstance(en, str) and isinstance(tr, str) and len(en.split()) >= min_words:
                    pairs[lang].append((en, tr))

    if n_failed:
        print(f"    !! {n_failed} of {len(report['files'])} files unreadable")
    if not pairs:
        raise SystemExit(
            "\nNo parallel text extracted.\n"
            f"  files listed : {len(report['files'])}\n"
            f"  unreadable   : {n_failed}\n"
            f"  level        : {level}\n"
            f"  columns used : {cols}\n\n"
            "If the files are unreadable, the schema report was written on the other\n"
            f"host. Re-run:  python scripts/02_inspect_schema.py --root {root}\n"
        )
    return dict(pairs), level


def measure(pairs: dict[str, list[tuple[str, str]]], model: str) -> dict:
    from transformers import AutoTokenizer

    tok = AutoTokenizer.from_pretrained(model)
    out: dict[str, dict] = {}

    for lang, items in sorted(pairs.items()):
        if not items:
            continue
        en_txt = [a for a, _ in items]
        tr_txt = [b for _, b in items]
        en_tok = [len(x) for x in tok(en_txt, add_special_tokens=False)["input_ids"]]
        tr_tok = [len(x) for x in tok(tr_txt, add_special_tokens=False)["input_ids"]]
        en_wrd = [len(t.split()) for t in en_txt]
        tr_wrd = [len(t.split()) for t in tr_txt]

        # Fertility = tokens for the translation / tokens for the same English content.
        # Ratio of sums, not mean of ratios: weights by length, avoids short-text noise.
        fert = sum(tr_tok) / max(1, sum(en_tok))
        out[lang] = {
            "language": LANGS[lang],
            "n_samples": len(items),
            "fertility_vs_english": round(fert, 4),
            "tokens_per_word": round(sum(tr_tok) / max(1, sum(tr_wrd)), 4),
            "en_tokens_per_word": round(sum(en_tok) / max(1, sum(en_wrd)), 4),
            "mean_tokens": round(sum(tr_tok) / len(tr_tok), 1),
            "mean_words": round(sum(tr_wrd) / len(tr_wrd), 1),
            "en_mean_tokens": round(sum(en_tok) / len(en_tok), 1),
            # If we fix a 256-token budget, how many chunks does this language get
            # for content that is 1 chunk in English?
            "chunks_at_256_rel_english": round(fert, 3),
        }
    return out


def plot(stats: dict, out_png: Path) -> None:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    langs = sorted(stats, key=lambda k: stats[k]["fertility_vs_english"])
    names = [stats[k]["language"] for k in langs]
    fert = [stats[k]["fertility_vs_english"] for k in langs]

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6.5), constrained_layout=True)

    bars = ax1.barh(names, fert, color="#4C78A8")
    ax1.axvline(1.0, color="#333", ls="--", lw=1.2, label="English baseline")
    ax1.set_xlabel("Tokenizer fertility  (tokens vs. same content in English)")
    ax1.set_title("Indic scripts cost more tokens per unit of meaning", fontsize=12, weight="bold")
    ax1.legend(frameon=False, loc="lower right")
    for b, v in zip(bars, fert):
        ax1.text(v + 0.02, b.get_y() + b.get_height() / 2, f"{v:.2f}x", va="center", fontsize=9)
    ax1.margins(x=0.14)

    # The money panel: a fixed token budget produces different chunk counts per language.
    chunks = [max(1, round(f)) for f in fert]
    colors = ["#54A24B" if c == 1 else "#E45756" for c in chunks]
    ax2.barh(names, chunks, color=colors)
    ax2.set_xlabel("Chunks produced at a fixed 256-token budget\n(for content that is 1 chunk in English)")
    ax2.set_title("Why a fixed token budget breaks cross-lingual chunk IDs",
                  fontsize=12, weight="bold")
    ax2.set_xticks(range(0, max(chunks) + 2))

    ratio = max(fert) / min(fert)
    fig.suptitle(
        f"Tokenizer fertility, MSMARCO-XI  |  spread {min(fert):.2f}x - {max(fert):.2f}x "
        f"(ratio {ratio:.2f}x)  ->  "
        + ("chunk in WORD space" if ratio > 1.3 else "token budgets are comparable"),
        fontsize=13, weight="bold",
    )
    fig.savefig(out_png, dpi=160)
    print(f"==> wrote {out_png}")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", type=Path, default=None,
                    help="data root; defaults to $VOICERAG_ROOT or <repo>/../voicerag_data")
    ap.add_argument("--model", default="BAAI/bge-m3")
    ap.add_argument("--n-per-lang", type=int, default=2000)
    args = ap.parse_args()

    root = (args.root.expanduser().resolve() if args.root else default_root())
    print(f"==> data root: {root}")
    res_dir = root / "results"
    res_dir.mkdir(parents=True, exist_ok=True)

    print(f"==> loading parallel text ({args.n_per_lang}/language)")
    pairs, level = load_pairs(root, args.n_per_lang)
    print(f"    measurement level: {level}")
    print(f"    got {len(pairs)} languages: {sorted(pairs)}")
    missing = sorted(set(LANGS) - set(pairs))
    if missing:
        print(f"    !! no samples for: {missing}")

    print(f"==> tokenizing with {args.model}")
    stats = measure(pairs, args.model)

    print(f"\n{'lang':<12}{'fertility':>11}{'tok/word':>11}{'mean tok':>11}{'mean words':>12}")
    print("-" * 57)
    for k in sorted(stats, key=lambda x: -stats[x]["fertility_vs_english"]):
        s = stats[k]
        print(f"{s['language']:<12}{s['fertility_vs_english']:>11.2f}"
              f"{s['tokens_per_word']:>11.2f}{s['mean_tokens']:>11.1f}{s['mean_words']:>12.1f}")

    if not stats:
        raise SystemExit("no languages measured — see the diagnostics above")
    fert = [s["fertility_vs_english"] for s in stats.values()]
    ratio = max(fert) / min(fert)
    verdict = (
        "CHUNK IN WORD/SENTENCE SPACE. A fixed token budget is not comparable across languages."
        if ratio > 1.3 else
        "Token budgets are roughly comparable; a shared token budget is defensible."
    )
    print(f"\n==> spread {min(fert):.2f}x - {max(fert):.2f}x (ratio {ratio:.2f}x)\n==> {verdict}")

    payload = {
        "model": args.model,
        "level": level,
        "n_per_lang": args.n_per_lang,
        "per_language": stats,
        "fertility_ratio_max_over_min": round(ratio, 4),
        "verdict": verdict,
    }
    (res_dir / "fertility.json").write_text(json.dumps(payload, indent=2))
    print(f"==> wrote {res_dir/'fertility.json'}")
    plot(stats, res_dir / "fertility.png")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())