Upload eval_and_compare.py with huggingface_hub

eval_and_compare.py

@@ -0,0 +1,277 @@
+#!/usr/bin/env python3 -u
+"""eval_and_compare.py — Evaluate all ours (28) + externals, generate plot, save results."""
+
+import json, os, sys, time, csv, gc, warnings
+from collections import Counter
+from dataclasses import dataclass, field, asdict
+from typing import List, Dict
+
+import numpy as np
+warnings.filterwarnings("ignore")
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+BASE = "/root/oiq_cc_tokenizer/results"
+CORPORA = os.path.join(BASE, "corpora")
+TOK_DIR = os.path.join(BASE, "tokenizers")
+PLOTS_DIR = os.path.join(BASE, "plots")
+
+import regex
+_WORD_PAT = regex.compile(r"[\p{L}\p{M}\p{N}]+", regex.UNICODE)
+_AR_PAT = regex.compile(r"[\u0600-\u06FF\u0750-\u077F]")
+_SPECIAL = {"<unk>", "<s>", "</s>", "[CLS]", "[SEP]", "[PAD]", "[UNK]", "<pad>"}
+
+def segment_words(t): return _WORD_PAT.findall(t)
+def count_graphemes(t): return len(regex.findall(r"\X", t))
+def detect_script(t): return "ar" if len(_AR_PAT.findall(t)) > len(t) * 0.3 else "az"
+def filter_sp(tokens): return [t for t in tokens if t not in _SPECIAL]
+
+@dataclass
+class M:
+    name: str = ""
+    source: str = ""
+    algorithm: str = ""
+    architecture: str = ""
+    vocab_size: int = 0
+    fertility_ar: float = 0.0
+    fertility_az: float = 0.0
+    fertility_overall: float = 0.0
+    disparity: float = 0.0
+    cpt_ar: float = 0.0
+    cpt_az: float = 0.0
+    exact_match_ar: float = 0.0
+    exact_match_az: float = 0.0
+
+
+class RawConcat:
+    def __init__(self, ar_j, az_j):
+        from tokenizers import Tokenizer
+        self.ar = Tokenizer.from_file(ar_j)
+        self.az = Tokenizer.from_file(az_j)
+
+    def encode(self, text):
+        s = detect_script(text)
+        t = self.ar if s == "ar" else self.az
+        enc = t.encode(text)
+        return enc.tokens, enc.ids, s
+
+    def decode(self, ids, script):
+        t = self.ar if script == "ar" else self.az
+        return t.decode(ids, skip_special_tokens=True)
+
+
+class RawShared:
+    def __init__(self, j):
+        from tokenizers import Tokenizer
+        self.tok = Tokenizer.from_file(j)
+
+    def encode(self, text):
+        enc = self.tok.encode(text)
+        return enc.tokens, enc.ids, detect_script(text)
+
+    def decode(self, ids, script):
+        return self.tok.decode(ids, skip_special_tokens=True)
+
+
+class HFTok:
+    def __init__(self, repo):
+        from transformers import AutoTokenizer
+        self.tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
+
+    def encode(self, text):
+        ids = self.tok.encode(text, add_special_tokens=False)
+        return self.tok.convert_ids_to_tokens(ids), ids, detect_script(text)
+
+    def decode(self, ids, script):
+        return self.tok.decode(ids, skip_special_tokens=True)
+
+
+def evaluate(tok, name, source, algo, arch, vsz, texts):
+    m = M(name=name, source=source, algorithm=algo, architecture=arch, vocab_size=vsz)
+    ar_f, az_f, all_f = [], [], []
+    ar_c, az_c = [], []
+    ar_ok, az_ok, ar_n, az_n = 0, 0, 0, 0
+
+    for i, text in enumerate(texts):
+        if (i + 1) % 5000 == 0:
+            print(f"    [{i+1}/{len(texts)}] {name}", flush=True)
+        try:
+            tokens, ids, script = tok.encode(text)
+            content = filter_sp(tokens)
+            words = segment_words(text)
+            if not words:
+                continue
+            fert = len(content) / len(words)
+            all_f.append(fert)
+            cpt = count_graphemes(text) / max(len(content), 1)
+            try:
+                dec = tok.decode(ids, script)
+                exact = dec.strip() == text.strip()
+            except:
+                exact = False
+            if script == "ar":
+                ar_f.append(fert); ar_c.append(cpt); ar_n += 1
+                if exact: ar_ok += 1
+            else:
+                az_f.append(fert); az_c.append(cpt); az_n += 1
+                if exact: az_ok += 1
+        except:
+            pass
+
+    m.fertility_ar = float(np.mean(ar_f)) if ar_f else 0
+    m.fertility_az = float(np.mean(az_f)) if az_f else 0
+    m.fertility_overall = float(np.mean(all_f)) if all_f else 0
+    mx = max(m.fertility_ar, m.fertility_az, 1e-9)
+    m.disparity = abs(m.fertility_ar - m.fertility_az) / mx
+    m.cpt_ar = float(np.mean(ar_c)) if ar_c else 0
+    m.cpt_az = float(np.mean(az_c)) if az_c else 0
+    m.exact_match_ar = ar_ok / max(ar_n, 1)
+    m.exact_match_az = az_ok / max(az_n, 1)
+    return m
+
+
+def main():
+    # Load test texts
+    texts = []
+    for s in ("test_ar", "test_az", "test_mi"):
+        p = os.path.join(CORPORA, f"{s}.txt")
+        if os.path.exists(p):
+            with open(p) as f:
+                texts.extend(l.strip() for l in f if l.strip())
+    print(f"{len(texts)} texts", flush=True)
+
+    results = []
+
+    # --- Our tokenizers ---
+    for vsz in (8000, 16000, 32000):
+        for algo in ("bpe", "unigram", "wordpiece", "bbpe"):
+            # Shared
+            jp = os.path.join(TOK_DIR, f"shared_{algo}_{vsz}.json")
+            if os.path.exists(jp):
+                name = f"shared_{algo}_{vsz}"
+                print(f"\n{name}", flush=True)
+                tok = RawShared(jp)
+                r = evaluate(tok, name, "ours", algo, "shared", vsz, texts)
+                print(f"  F={r.fertility_overall:.3f} D={r.disparity:.3f} EM_ar={r.exact_match_ar:.2%}", flush=True)
+                results.append(r)
+                del tok; gc.collect()
+
+            # Concat
+            ar_j = os.path.join(TOK_DIR, f"concat_ar_{algo}_{vsz//2}.json")
+            az_j = os.path.join(TOK_DIR, f"concat_az_{algo}_{vsz//2}.json")
+            if os.path.exists(ar_j) and os.path.exists(az_j):
+                name = f"concat_{algo}_{vsz}"
+                print(f"\n{name}", flush=True)
+                tok = RawConcat(ar_j, az_j)
+                r = evaluate(tok, name, "ours", algo, "concatenated", vsz, texts)
+                print(f"  F={r.fertility_overall:.3f} D={r.disparity:.3f} EM_ar={r.exact_match_ar:.2%}", flush=True)
+                results.append(r)
+                del tok; gc.collect()
+
+    # --- External ---
+    externals = [
+        ("CaMeLBERT-MSA", "external_msa", "WordPiece", "shared", 30000, "CAMeL-Lab/bert-base-arabic-camelbert-msa"),
+        ("Asafaya-BERT", "external_msa", "WordPiece", "shared", 32000, "asafaya/bert-base-arabic"),
+        ("Aranizer-SP-86k", "external_msa", "SentencePiece", "shared", 86000, "riotu-lab/Aranizer-SP-86k"),
+        ("DarijaBERT-ar", "external_darija", "WordPiece", "shared", 80000, "SI2M-Lab/DarijaBERT"),
+        ("DarijaBERT-az", "external_darija", "WordPiece", "shared", 110000, "SI2M-Lab/DarijaBERT-arabizi"),
+    ]
+    for name, src, algo, arch, vsz, repo in externals:
+        print(f"\n{name} ({repo})", flush=True)
+        try:
+            tok = HFTok(repo)
+            r = evaluate(tok, name, src, algo, arch, vsz, texts)
+            print(f"  F={r.fertility_overall:.3f} D={r.disparity:.3f} EM_ar={r.exact_match_ar:.2%}", flush=True)
+            results.append(r)
+            del tok; gc.collect()
+        except Exception as e:
+            print(f"  FAILED: {e}", flush=True)
+
+    # Save
+    out_csv = os.path.join(BASE, "external_comparison.csv")
+    out_json = os.path.join(BASE, "external_comparison.json")
+    with open(out_csv, "w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=list(asdict(results[0]).keys()))
+        w.writeheader()
+        for r in results: w.writerow(asdict(r))
+    with open(out_json, "w") as f:
+        json.dump([asdict(r) for r in results], f, indent=2)
+
+    # Print table
+    print("\n" + "=" * 130, flush=True)
+    hdr = f"{'Name':<30} {'Source':<16} {'V':>7} {'Fert':>7} {'F_ar':>7} {'F_az':>7} {'Disp':>7} {'CPT_ar':>7} {'CPT_az':>7} {'EM_ar':>7} {'EM_az':>7}"
+    print(hdr, flush=True)
+    print("-" * 130, flush=True)
+    for r in sorted(results, key=lambda x: (0 if x.source == "ours" else 1, x.vocab_size)):
+        print(f"{r.name:<30} {r.source:<16} {r.vocab_size:>7,} {r.fertility_overall:>7.3f} {r.fertility_ar:>7.3f} {r.fertility_az:>7.3f} {r.disparity:>7.3f} {r.cpt_ar:>7.3f} {r.cpt_az:>7.3f} {r.exact_match_ar:>7.2%} {r.exact_match_az:>7.2%}", flush=True)
+    print("=" * 130, flush=True)
+
+    # Generate comparison plot
+    ours_best = []
+    for vsz in (8000, 16000, 32000):
+        cands = [r for r in results if r.vocab_size == vsz and r.architecture == "concatenated" and r.source == "ours"]
+        if cands:
+            best = min(cands, key=lambda x: x.fertility_overall)
+            ours_best.append(best)
+    ext = [r for r in results if r.source != "ours"]
+    plot_data = ours_best + ext
+
+    fig, axes = plt.subplots(2, 2, figsize=(16, 11))
+    colors = {"8000": "#E69F00", "16000": "#009E73", "32000": "#0072B2",
+             "external_msa": "#CC79A7", "external_darija": "#D55E00"}
+    labels = [f"Ours\n{r.name}\n({r.vocab_size:,})" for r in ours_best] + \
+            [f"{r.name}\n({r.vocab_size:,})" for r in ext]
+    bar_c = [colors[str(r.vocab_size)] for r in ours_best] + \
+            [colors.get(r.source, "#999") for r in ext]
+    n = len(plot_data)
+
+    for idx, (key, vals_fn, title, ylabel) in enumerate([
+        ("fert", lambda r: r.fertility_overall, "Overall Fertility (Lower = Better)", "Fertility"),
+        ("disp", lambda r: r.disparity, "Cross-Script Disparity (Lower = Better)", "Disparity"),
+    ]):
+        ax = axes[0, idx]
+        vals = [vals_fn(r) for r in plot_data]
+        bars = ax.bar(range(n), vals, color=bar_c, edgecolor="gray", linewidth=0.5)
+        ax.set_xticks(range(n))
+        ax.set_xticklabels(labels, fontsize=6, ha="center")
+        ax.set_ylabel(ylabel, fontsize=9)
+        ax.set_title(title, fontsize=10, fontweight="bold")
+        for b, v in zip(bars, vals):
+            ax.text(b.get_x()+b.get_width()/2, b.get_height()+0.005, f"{v:.3f}", ha="center", va="bottom", fontsize=6)
+
+    # Exact match grouped
+    ax = axes[1, 0]
+    x = np.arange(n)
+    w = 0.35
+    ax.bar(x-w/2, [r.exact_match_ar*100 for r in plot_data], w, label="Arabic", color="#56B4E9")
+    ax.bar(x+w/2, [r.exact_match_az*100 for r in plot_data], w, label="Arabizi", color="#333")
+    ax.set_xticks(x); ax.set_xticklabels(labels, fontsize=6, ha="center")
+    ax.set_ylabel("Exact Match (%)"); ax.set_title("Exact Reconstruction", fontsize=10, fontweight="bold")
+    ax.legend(fontsize=7); ax.set_ylim(0, 108)
+
+    # CPT grouped
+    ax = axes[1, 1]
+    ax.bar(x-w/2, [r.cpt_ar for r in plot_data], w, label="Arabic", color="#56B4E9")
+    ax.bar(x+w/2, [r.cpt_az for r in plot_data], w, label="Arabizi", color="#333")
+    ax.set_xticks(x); ax.set_xticklabels(labels, fontsize=6, ha="center")
+    ax.set_ylabel("CPT"); ax.set_title("Characters Per Token (Higher = Better)", fontsize=10, fontweight="bold")
+    ax.legend(fontsize=7)
+
+    from matplotlib.patches import Patch
+    fig.legend(handles=[
+        Patch(fc=colors["8000"], label="Ours (8K)"),
+        Patch(fc=colors["16000"], label="Ours (16K)"),
+        Patch(fc=colors["32000"], label="Ours (32K)"),
+        Patch(fc=colors["external_msa"], label="External (MSA)"),
+        Patch(fc=colors["external_darija"], label="External (Darija)"),
+    ], loc="upper center", ncol=5, fontsize=8, bbox_to_anchor=(0.5, 0.98), frameon=True)
+    plt.tight_layout(rect=[0, 0, 1, 0.95])
+    fig.savefig(os.path.join(PLOTS_DIR, "external_comparison.png"), dpi=150, bbox_inches="tight")
+    plt.close(fig)
+    print(f"\nPlot: {os.path.join(PLOTS_DIR, 'external_comparison.png')}", flush=True)
+    print("DONE!", flush=True)
+
+if __name__ == "__main__":
+    main()