#!/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 /../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())