Blame for code/repr_probe.py

#!/usr/bin/env python3
"""
Representation-layer probe: what does MiniCPM5-2B's tokenizer actually see?

This is the mechanistic half of the study. The behavioural half asks whether
the model *reveals* that a question has several legitimate answers. This half
asks whether the tokenizer could even carry the distinctions those answers
depend on.

Everything here is computed from the tokenizer and the vocabulary, so the
numbers are exact and independent of quantisation or of prompting.

Method note. MiniCPM5-2B uses a GPT-2-style byte-level BPE. Tokens that are
genuine text units ("prayer") live in the vocabulary directly; anything else is
stored as *bytes* of UTF-8, each mapped into a printable stand-in character.
Arabic therefore has to be reassembled from raw byte tokens. To measure that
honestly we decode the byte alphabet back to bytes and then to text, rather
than inspecting the stand-in characters.

Measurements
------------
fertility_by_script        tokens per character and per word, by script
byte_level_profile         what share of the vocabulary and of each script is
                            byte-level fallback rather than a text unit
vocab_script_profile       how many consecutive Arabic *characters* any single
                            token covers
term_morphology            how a set of Arabic religious terms actually
                            tokenises: characters per token, fused or detached
                            short vowels
diacritic_cost             what a vocalised spelling costs over the unvocalised
                            one, and how the vowel is encoded
diacritic_discrimination   whether two spellings differing only in a short
                            vowel survive as different token sequences
"""
from __future__ import annotations

import json
import pathlib
from collections import Counter

from transformers import AutoTokenizer

ROOT = pathlib.Path(__file__).resolve().parent.parent
OUT = ROOT / "results" / "repr_probe.json"

REPO = "openbpm/MiniCPM5-2B".replace("openbpm", "openbmb")

TEXTS = {
    "en": "A traveller with no water performs tayammum by striking the hands on bare stone, and this is valid in the view of the school.",
    "zh": "一名没有水的旅行者在裸石上拍打双手进行土净礼，这在教派的观点中是被允许的。",
    "ar_vocalized": "السَّابِحُ الَّذِي لَا مَاءَ عِنْدَهُ يَتَضَمَّضُ بِالْحِجَارَةِ وَيَجُوزُ فِي مَذْهَبِ الْفَقِيهِ.",
    "ar_plain": "السابح الذي لا ماء عنده يتضض بالحجارة ويجوز في مذهب الفقيه.",
    "ur": "سفر کرنے والا مسافر جب پر پتھر پر ہاتھ مارنا اور مذہب کے مطابق یہ جائز ہے۔",
    "fa": "مسافری که آب ندارد بر سنگ خام با دستهای خود تیمم می‌کند و در مذهب فقیه جایز است.",
    "hi": "एक यात्री जिसके पास पानी नहीं है वह नंगे पत्थर पर अपने हाथ मारता है और मजहब के अनुसार यह जायज है।",
    "arabizi": "Al musafir ila laa maa lahu maa yatamassabu bi al-hijarati wa yajuzu fi madhhab al-faqih.",
}

TERMS = [
    ("صَلَاة", "صلاه"), ("زَكَاة", "زكاه"), ("تَيَمُّم", "تيمم"),
    ("وُضُوء", "وضوء"), ("رَكْعَة", "ركعه"), ("سُجُود", "سجود"),
    ("خُطْبَة", "خطبه"), ("جَنَازَة", "جنازه"), ("مَذْهَب", "مذهب"),
    ("حَجّ", "حج"), ("قَصْر", "قصر"), ("نِيَّة", "نية"),
]

HARAKAT_PAIRS = [
    ("نِيَّة", "نَيَّة"), ("زَكَاة", "زِكَاة"), ("حَلَال", "حِلَال"),
    ("رَكْعَة", "رَكَعَة"), ("طَهَارَة", "طُهَارَة"),
]

HARAKAT = set("ًٌٍَُِّْٰٕٓٔ")
ENGLISH_EQUIV = [("prayer", "صلاة"), ("purification", "طهارة"),
                 ("zakat", "زكاة"), ("fasting", "صوم"), ("pilgrimage", "حج")]


# ------------------------------------------------------- byte-level alphabet --
def bytes_to_unicode() -> dict[int, str]:
    bs = (list(range(ord("!"), ord("~") + 1)) +
          list(range(ord("\xa1"), ord("\xac") + 1)) +
          list(range(ord("\xae"), ord("\xff") + 1)))
    cs = bs[:]
    n = 0
    for b in range(256):
        if b not in bs:
            bs.append(b)
            cs.append(256 + n)
            n += 1
    return dict(zip(bs, (chr(c) for c in cs)))


B2U = bytes_to_unicode()
U2B = {v: k for k, v in B2U.items()}


def token_to_bytes(tok_str: str) -> bytes | None:
    try:
        return bytes(U2B[c] for c in tok_str)
    except KeyError:
        return None


def is_byte_token(tok_str: str) -> bool:
    """A token is byte-level iff its characters decode to bytes that are NOT
    valid UTF-8 on their own -- i.e. it is a fragment of a character rather
    than a character (or a word) in the vocabulary."""
    b = token_to_bytes(tok_str)
    if b is None:
        return False
    try:
        b.decode("utf-8")
        return False
    except UnicodeDecodeError:
        return True


def byte_token_to_text(tok_str: str) -> str:
    b = token_to_bytes(tok_str)
    if b is None:
        return ""
    return b.decode("utf-8", errors="replace")


def is_arabic_letter(ch: str) -> bool:
    return "؀" <= ch <= "ۿ"


def longest_arabic_run(s: str) -> int:
    best = run = 0
    for ch in s:
        run = run + 1 if is_arabic_letter(ch) else 0
        best = max(best, run)
    return best


def token_texts(tok, ids):
    """Decode each token id back to the text it contributes."""
    out = []
    for i in ids:
        t = tok.convert_ids_to_tokens(i)
        if is_byte_token(t):
            out.append(byte_token_to_text(t))
        else:
            out.append(t.replace("Ġ", " "))
    return out


# ----------------------------------------------------------------- measures --
def fertility(tok, text: str) -> dict:
    ids = tok(text, add_special_tokens=False)["input_ids"]
    toks = [tok.convert_ids_to_tokens(i) for i in ids]
    n_byte = sum(1 for t in toks if is_byte_token(t))
    words = len(text.split())
    chars = len([c for c in text if not c.isspace()])
    return {
        "n_tokens": len(ids),
        "n_words": words,
        "n_chars": chars,
        "tokens_per_word": round(len(ids) / max(words, 1), 3),
        "tokens_per_char": round(len(ids) / max(chars, 1), 4),
        "byte_level_tokens": n_byte,
        "byte_level_rate": round(n_byte / max(len(ids), 1), 4),
    }


def vocab_script_profile(tok) -> dict:
    vocab = tok.get_vocab()
    n_total = len(vocab)
    n_byte = 0
    n_text = 0
    n_arabic_units = 0
    n_latin_units = 0
    n_cjk_units = 0
    n_devanagari_units = 0
    n_other_units = 0
    run_hist: Counter[int] = Counter()
    for tok_str in vocab:
        if is_byte_token(tok_str):
            n_byte += 1
            continue
        n_text += 1
        s = tok_str
        if any(c in U2B for c in tok_str):
            s = byte_token_to_text(tok_str)
        s = s.replace("Ġ", " ")
        if any(is_arabic_letter(c) for c in s):
            n_arabic_units += 1
            run_hist[longest_arabic_run(s)] += 1
        elif any("a" <= c <= "z" or "A" <= c <= "Z" for c in s):
            n_latin_units += 1
        elif any(0x4E00 <= ord(c) <= 0x9FFF for c in s):
            n_cjk_units += 1
        elif any(0x0900 <= ord(c) <= 0x097F for c in s):
            n_devanagari_units += 1
        else:
            n_other_units += 1
    tot = max(n_arabic_units, 1)
    return {
        "vocab_size": n_total,
        "n_byte_level_tokens": n_byte,
        "share_byte_level": round(n_byte / n_total, 4),
        "n_text_unit_tokens": n_text,
        "n_arabic_text_units": n_arabic_units,
        "share_arabic_text_units": round(n_arabic_units / n_total, 6),
        "n_latin_text_units": n_latin_units,
        "n_cjk_text_units": n_cjk_units,
        "n_devanagari_text_units": n_devanagari_units,
        "n_other_text_units": n_other_units,
        "arabic_unit_run_histogram": {str(k): v for k, v in sorted(run_hist.items())},
        "max_arabic_chars_in_one_unit": max(run_hist) if run_hist else 0,
        "mean_arabic_chars_per_unit": round(
            sum(k * v for k, v in run_hist.items()) / tot, 3),
    }


def text_of_ids(tok, ids) -> str:
    """Reconstruct the original string from token ids by joining their bytes."""
    buf = bytearray()
    for i in ids:
        t = tok.convert_ids_to_tokens(i)
        b = token_to_bytes(t)
        if b is None:
            buf += t.encode("utf-8")
        else:
            buf += b
    return buf.decode("utf-8", errors="replace")


def token_char_attribution(tok, ids) -> list[str]:
    """Attribute decoded characters to the token that completed them.

    A byte-level token may hold half a character; the character is attributed
    to whichever token finally completes it. This lets us ask, for example,
    whether a short vowel ever rides along with its letter inside one token.
    """
    pending = b""
    out = []
    for i in ids:
        t = tok.convert_ids_to_tokens(i)
        b = token_to_bytes(t)
        if b is None:
            b = t.encode("utf-8")
        pending += b
        got = ""
        while pending:
            try:
                got += pending.decode("utf-8")
                pending = b""
            except UnicodeDecodeError as e:
                if e.reason == "unexpected end of data":
                    break
                got += "�"
                pending = pending[e.end:]
        out.append(got)
    return out


def term_morphology(tok) -> dict:
    rows = []
    for voc, plain in TERMS:
        vi = tok(voc, add_special_tokens=False)["input_ids"]
        pi = tok(plain, add_special_tokens=False)["input_ids"]
        vt = token_char_attribution(tok, vi)
        pt = token_char_attribution(tok, pi)
        runs = [longest_arabic_run(t) for t in vt]
        fused = sum(1 for t in vt
                    if any(is_arabic_letter(c) for c in t) and
                    any(c in HARAKAT for c in t))
        detached_v = sum(1 for t in vt if t and all(c in HARAKAT for c in t))
        rows.append({
            "vocalized": voc, "plain": plain,
            "n_tokens_vocalized": len(vi), "n_tokens_plain": len(pi),
            "diacritic_cost_ratio": round(len(vi) / max(len(pi), 1), 3),
            "extra_tokens_for_vocalization": len(vi) - len(pi),
            "max_arabic_chars_in_one_token": max(runs) if runs else 0,
            "n_tokens_fusing_letter_and_vowel": fused,
            "n_tokens_that_are_only_a_vowel": detached_v,
            "tokens_vocalized": vt,
            "tokens_plain": pt,
        })
    return {
        "per_term": rows,
        "n_terms": len(rows),
        "mean_tokens_plain": round(sum(r["n_tokens_plain"] for r in rows) / len(rows), 2),
        "mean_tokens_vocalized": round(sum(r["n_tokens_vocalized"] for r in rows) / len(rows), 2),
        "mean_diacritic_cost_ratio": round(
            sum(r["diacritic_cost_ratio"] for r in rows) / len(rows), 3),
        "max_arabic_chars_in_one_token": max(
            r["max_arabic_chars_in_one_token"] for r in rows),
        "n_terms_with_fused_letter_and_vowel": sum(
            r["n_tokens_fusing_letter_and_vowel"] > 0 for r in rows),
        "n_terms_with_detached_vowel_tokens": sum(
            r["n_tokens_that_are_only_a_vowel"] > 0 for r in rows),
    }


def diacritic_cost(tok) -> dict:
    fused = detached = total = 0
    examples = []
    for voc, _ in TERMS:
        for t in token_char_attribution(tok, tok(voc, add_special_tokens=False)["input_ids"]):
            if any(c in HARAKAT for c in t):
                total += 1
                if any(is_arabic_letter(c) for c in t):
                    fused += 1
                else:
                    detached += 1
                    if len(examples) < 8:
                        examples.append(t)
    return {
        "tokens_containing_a_harakat": total,
        "fused_with_letter": fused,
        "detached_harakat_tokens": detached,
        "share_detached": round(detached / max(total, 1), 4),
        "examples_detached": examples,
    }


def diacritic_discrimination(tok) -> dict:
    rows = []
    for a, b in HARAKAT_PAIRS:
        ai = tok(a, add_special_tokens=False)["input_ids"]
        bi = tok(b, add_special_tokens=False)["input_ids"]
        rows.append({"a": a, "b": b, "a_tokens": ai, "b_tokens": bi,
                     "identical": ai == bi})
    return {
        "pairs": rows,
        "n_pairs": len(rows),
        "n_merged": sum(1 for r in rows if r["identical"]),
        "merge_rate": round(sum(1 for r in rows if r["identical"]) / len(rows), 3),
    }


def english_equivalent_cost(tok) -> list[dict]:
    rows = []
    for en, ar in ENGLISH_EQUIV:
        e = tok(en, add_special_tokens=False)["input_ids"]
        a = tok(ar, add_special_tokens=False)["input_ids"]
        rows.append({"english": en, "arabic": ar,
                     "english_tokens": len(e), "arabic_tokens": len(a),
                     "arabic_over_english": round(len(a) / max(len(e), 1), 2)})
    return rows


def main() -> None:
    tok = AutoTokenizer.from_pretrained(REPO, trust_remote_code=True)
    report = {
        "repo": REPO,
        "vocab_size": len(tok),
        "fertility_by_script": {k: fertility(tok, v) for k, v in TEXTS.items()},
        "vocab_script_profile": vocab_script_profile(tok),
        "term_morphology": term_morphology(tok),
        "diacritic_cost": diacritic_cost(tok),
        "diacritic_discrimination": diacritic_discrimination(tok),
        "english_equivalent_cost": english_equivalent_cost(tok),
    }
    OUT.parent.mkdir(parents=True, exist_ok=True)
    OUT.write_text(json.dumps(report, ensure_ascii=False, indent=2))

    print("=== fertility (byte-level share = share of tokens that are raw bytes) ===")
    for k, v in report["fertility_by_script"].items():
        print(f"  {k:14s} tok/char={v['tokens_per_char']:.4f}  tok/word={v['tokens_per_word']:6.2f}"
              f"  byte-level={v['byte_level_rate']*100:5.1f}%")
    vp = report["vocab_script_profile"]
    print("\n=== vocabulary profile ===")
    print(f"  vocab={vp['vocab_size']}  byte-level tokens={vp['n_byte_level_tokens']} "
          f"({vp['share_byte_level']*100:.1f}%)")
    print(f"  tokens containing Arabic: {vp['n_arabic_text_units']} "
          f"({vp['share_arabic_text_units']*100:.3f}%)   Latin {vp['n_latin_text_units']}   "
          f"CJK {vp['n_cjk_text_units']}")
    print(f"  max Arabic characters in one unit: {vp['max_arabic_chars_in_one_unit']}  "
          f"mean {vp['mean_arabic_chars_per_unit']}")
    tm = report["term_morphology"]
    print("\n=== religious terms ===")
    print(f"  mean tokens: plain {tm['mean_tokens_plain']}  vocalised {tm['mean_tokens_vocalized']} "
          f"(ratio {tm['mean_diacritic_cost_ratio']})")
    print(f"  max Arabic characters in one token across terms: {tm['max_arabic_chars_in_one_token']}")
    print(f"  terms where any token fuses a letter with its vowel: "
          f"{tm['n_terms_with_fused_letter_and_vowel']}/{tm['n_terms']}")
    dc = report["diacritic_cost"]
    print(f"  harakat tokens: {dc['tokens_containing_a_harakat']}  detached: "
          f"{dc['detached_harakat_tokens']} ({dc['share_detached']*100:.0f}%)")
    dd = report["diacritic_discrimination"]
    print(f"  harakat-only pairs merged by the tokenizer: {dd['n_merged']}/{dd['n_pairs']}")
    print("\n=== same content, two scripts ===")
    for r in report["english_equivalent_cost"]:
        print(f"  {r['english']:14s} {r['english_tokens']} tok   |  {r['arabic']} "
              f"{r['arabic_tokens']} tok  ({r['arabic_over_english']}x)")


if __name__ == "__main__":
    main()