Add code-switching eval script and results

eval_codeswitch_and_new_baselines.py

@@ -0,0 +1,369 @@
+#!/usr/bin/env python3 -u
+"""
+eval_codeswitch_and_new_baselines.py
+1) Evaluate mixed-script texts SEPARATELY (not forced into ar/az binary)
+2) Add atlasia/darija_bpe_tokenizer baseline
+3) Evaluate on independent DODa dataset (Arabic-only)
+"""
+
+import json, os, sys, time, csv, gc, warnings
+from collections import Counter
+from dataclasses import dataclass, asdict
+from typing import List
+
+import numpy as np
+import regex
+warnings.filterwarnings("ignore")
+
+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")
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
+
+_WORD_PAT = regex.compile(r"[\p{L}\p{M}\p{N}]+", regex.UNICODE)
+_AR_PAT = regex.compile(r"[\u0600-\u06FF\u0750-\u077F]")
+_LAT_PAT = regex.compile(r"[a-zA-Z]")
+_SPECIAL = {"<unk>", "<s>", "</s>", "[CLS]", "[SEP]", "[PAD]", "[UNK]", "<pad>", "",
+            "<|im_start|>", "<|im_end|>"}
+
+def segment_words(t): return _WORD_PAT.findall(t)
+def count_graphemes(t): return len(regex.findall(r"\X", t))
+def filter_sp(tokens): return [t for t in tokens if t not in _SPECIAL]
+
+def classify_script_detailed(t):
+    """Classify as 'ar', 'az', or 'mi' (mixed)."""
+    ar_chars = len(_AR_PAT.findall(t))
+    lat_chars = len(_LAT_PAT.findall(t))
+    total_alpha = ar_chars + lat_chars
+    if total_alpha == 0:
+        return "ar"
+    ar_ratio = ar_chars / total_alpha
+    lat_ratio = lat_chars / total_alpha
+    # Pure = >90% one script, mixed = both scripts present with >10% each
+    if ar_ratio > 0.9 and lat_ratio < 0.1:
+        return "ar"
+    elif lat_ratio > 0.9 and ar_ratio < 0.1:
+        return "az"
+    else:
+        return "mi"
+
+def detect_script(t): return "ar" if len(_AR_PAT.findall(t)) > len(t) * 0.3 else "az"
+
+
+@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_mi: float = 0.0
+    fertility_overall: float = 0.0
+    disparity: float = 0.0
+    cpt_ar: float = 0.0
+    cpt_az: float = 0.0
+    cpt_mi: float = 0.0
+    exact_match_ar: float = 0.0
+    exact_match_az: float = 0.0
+    exact_match_mi: 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 HFTok:
+    def __init__(self, repo, use_token=False):
+        from transformers import AutoTokenizer
+        kwargs = {"trust_remote_code": True}
+        if use_token:
+            kwargs["token"] = HF_TOKEN
+        self.tok = AutoTokenizer.from_pretrained(repo, **kwargs)
+
+    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_with_mixed(tok, name, source, algo, arch, vsz, texts):
+    """Evaluate with separate mi bucket for code-switched texts."""
+    m = M(name=name, source=source, algorithm=algo, architecture=arch, vocab_size=vsz)
+    buckets = {"ar": [], "az": [], "mi": []}
+    cpt_buckets = {"ar": [], "az": [], "mi": []}
+    em_buckets = {"ar": {"ok": 0, "n": 0}, "az": {"ok": 0, "n": 0}, "mi": {"ok": 0, "n": 0}}
+    all_f = []
+
+    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)
+
+            # Use DETAILED classification for separate mi bucket
+            sc = classify_script_detailed(text)
+
+            buckets[sc].append(fert)
+            cpt_buckets[sc].append(cpt)
+            em_buckets[sc]["n"] += 1
+
+            try:
+                dec = tok.decode(ids, script)
+                if dec.strip() == text.strip():
+                    em_buckets[sc]["ok"] += 1
+            except:
+                pass
+        except:
+            pass
+
+    for sc in ("ar", "az", "mi"):
+        setattr(m, f"fertility_{sc}", float(np.mean(buckets[sc])) if buckets[sc] else 0)
+        setattr(m, f"cpt_{sc}", float(np.mean(cpt_buckets[sc])) if cpt_buckets[sc] else 0)
+        b = em_buckets[sc]
+        setattr(m, f"exact_match_{sc}", b["ok"] / max(b["n"], 1))
+
+    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
+    return m
+
+
+def evaluate_on_doda(tok, name, source, algo, arch, vsz, texts):
+    """Evaluate on independent DODa data (Arabic only, no ar/az split)."""
+    all_f, all_c = [], []
+    em_ok, em_n = 0, 0
+
+    for i, text in enumerate(texts):
+        if (i + 1) % 5000 == 0:
+            print(f"    [{i+1}/{len(texts)}] {name} (doda)", 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)
+            all_c.append(cpt)
+            try:
+                dec = tok.decode(ids, script)
+                if dec.strip() == text.strip():
+                    em_ok += 1
+            except:
+                pass
+            em_n += 1
+        except:
+            pass
+
+    return {
+        "name": name, "source": source, "algorithm": algo,
+        "architecture": arch, "vocab_size": vsz,
+        "n_texts": em_n,
+        "fertility": float(np.mean(all_f)) if all_f else 0,
+        "cpt": float(np.mean(all_c)) if all_c else 0,
+        "exact_match": em_ok / max(em_n, 1),
+    }
+
+
+def main():
+    # Load test texts
+    test_ar, test_az, test_mi = [], [], []
+    for s, lst in [("test_ar", test_ar), ("test_az", test_az), ("test_mi", test_mi)]:
+        p = os.path.join(CORPORA, f"{s}.txt")
+        if os.path.exists(p):
+            with open(p) as f:
+                lst.extend(l.strip() for l in f if l.strip())
+
+    # Check script distribution with DETAILED classification
+    print("=== Script distribution (detailed) ===", flush=True)
+    from collections import Counter
+    dist = Counter()
+    for f in [test_ar, test_az, test_mi]:
+        for t in f:
+            dist[classify_script_detailed(t)] += 1
+    total = sum(dist.values())
+    for sc in ("ar", "az", "mi"):
+        print(f"  {sc}: {dist[sc]} ({dist[sc]/total*100:.1f}%)", flush=True)
+    print(f"  Total: {total}", flush=True)
+
+    all_texts = test_ar + test_az + test_mi
+    mi_texts = test_mi  # Only mixed-script texts for dedicated eval
+
+    print(f"\n=== 1. Code-switching evaluation (3 best ours + key externals) ===", flush=True)
+    cs_results = []
+
+    ours_cfg = [
+        ("concat_bpe_8000", "concat_ar_bpe_4000", "concat_az_bpe_4000", "bpe", "concatenated", 8000),
+        ("concat_wordpiece_16000", "concat_ar_wordpiece_8000", "concat_az_wordpiece_8000", "wordpiece", "concatenated", 16000),
+        ("concat_bpe_32000", "concat_ar_bpe_16000", "concat_az_bpe_16000", "bpe", "concatenated", 32000),
+    ]
+    for name, ar_sub, az_sub, algo, arch, vsz in ours_cfg:
+        ar_j = os.path.join(TOK_DIR, f"{ar_sub}.json")
+        az_j = os.path.join(TOK_DIR, f"{az_sub}.json")
+        if os.path.exists(ar_j) and os.path.exists(az_j):
+            print(f"\n{name}", flush=True)
+            tok = RawConcat(ar_j, az_j)
+            r = evaluate_with_mixed(tok, name, "ours", algo, arch, vsz, all_texts)
+            cs_results.append(r)
+            print(f"  F_ar={r.fertility_ar:.3f} F_az={r.fertility_az:.3f} F_mi={r.fertility_mi:.3f} EM_mi={r.exact_match_mi:.2%}", flush=True)
+            del tok; gc.collect()
+
+    # Key externals for code-switching
+    externals_cs = [
+        ("DarijaBERT-ar", "external_darija", "WordPiece", "shared", 80000,
+         "SI2M-Lab/DarijaBERT", False),
+        ("Qwen2.5-Darija", "external_darija", "SentencePiece", "shared", 151643,
+         "GemMaroc/Qwen2.5-7B-Instruct-darija", False),
+    ]
+    for name, src, algo, arch, vsz, repo, gated in externals_cs:
+        print(f"\n{name} ({repo})", flush=True)
+        try:
+            tok = HFTok(repo, use_token=gated)
+            r = evaluate_with_mixed(tok, name, src, algo, arch, vsz, all_texts)
+            cs_results.append(r)
+            print(f"  F_ar={r.fertility_ar:.3f} F_az={r.fertility_az:.3f} F_mi={r.fertility_mi:.3f} EM_mi={r.exact_match_mi:.2%}", flush=True)
+            del tok; gc.collect()
+        except Exception as e:
+            print(f"  FAILED: {e}", flush=True)
+
+    # Save code-switching results
+    cs_csv = os.path.join(BASE, "codeswitch_results.csv")
+    cs_json = os.path.join(BASE, "codeswitch_results.json")
+    with open(cs_csv, "w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=list(asdict(cs_results[0]).keys()))
+        w.writeheader()
+        for r in cs_results:
+            w.writerow(asdict(r))
+    with open(cs_json, "w") as f:
+        json.dump([asdict(r) for r in cs_results], f, indent=2)
+    print(f"\nCode-switching results saved to {cs_csv}", flush=True)
+
+    # Print code-switching table
+    print("\n" + "=" * 130, flush=True)
+    hdr = f"{'Name':<30} {'F_ar':>7} {'F_az':>7} {'F_mi':>7} {'CPT_ar':>7} {'CPT_az':>7} {'CPT_mi':>7} {'EM_ar':>7} {'EM_az':>7} {'EM_mi':>7}"
+    print(hdr, flush=True)
+    print("-" * 130, flush=True)
+    for r in cs_results:
+        print(f"{r.name:<30} {r.fertility_ar:>7.3f} {r.fertility_az:>7.3f} {r.fertility_mi:>7.3f} {r.cpt_ar:>7.3f} {r.cpt_az:>7.3f} {r.cpt_mi:>7.3f} {r.exact_match_ar:>7.2%} {r.exact_match_az:>7.2%} {r.exact_match_mi:>7.2%}", flush=True)
+    print("=" * 130, flush=True)
+
+    # ========== 2. atlasia/darija_bpe_tokenizer ==========
+    print("\n=== 2. Evaluating atlasia/darija_bpe_tokenizer ===", flush=True)
+    try:
+        tok = HFTok("atlasia/darija_bpe_tokenizer", use_token=True)
+        r = evaluate_with_mixed(tok, "atlasia_darija_bpe", "external_darija", "BPE", "shared", 0, all_texts)
+        # Also need vocab size
+        r.vocab_size = tok.tok.vocab_size
+        print(f"  Vocab size: {r.vocab_size}", flush=True)
+        print(f"  F={r.fertility_overall:.3f} F_ar={r.fertility_ar:.3f} F_az={r.fertility_az:.3f} ΔF={r.disparity:.3f}", flush=True)
+        cs_results.append(r)
+        del tok; gc.collect()
+        print("  atlasia/darija_bpe_tokenizer evaluated successfully", flush=True)
+    except Exception as e:
+        print(f"  atlasia/darija_bpe_tokenizer FAILED: {e}", flush=True)
+        import traceback; traceback.print_exc()
+
+    # ========== 3. Independent DODa evaluation ==========
+    print("\n=== 3. Independent dataset evaluation (DODa) ===", flush=True)
+
+    # Try to load DODa from HF
+    doda_texts = []
+    try:
+        from datasets import load_dataset
+        print("  Loading DODa from HuggingFace...", flush=True)
+        ds = load_dataset("OussamaElbaz/DODa", split="train", token=HF_TOKEN, trust_remote_code=True)
+        if ds is not None:
+            # Extract Arabic text
+            for row in ds:
+                t = row.get("text", "") or row.get("arabic", "") or row.get("word", "") or row.get("sentence", "")
+                if t and len(t.strip()) > 5:
+                    doda_texts.append(t.strip())
+            print(f"  Loaded {len(doda_texts)} DODa entries", flush=True)
+    except Exception as e:
+        print(f"  DODa load failed: {e}", flush=True)
+        import traceback; traceback.print_exc()
+
+    # Try alternative DODa repos
+    if not doda_texts:
+        for repo in ["OussamaElbaz/DODa", "DODa"]:
+            try:
+                from datasets import load_dataset
+                ds = load_dataset(repo, split="train", token=HF_TOKEN, trust_remote_code=True)
+                for row in ds:
+                    for k, v in row.items():
+                        if isinstance(v, str) and len(v.strip()) > 5 and any(c in v for c in "ابتثج"):
+                            doda_texts.append(v.strip())
+                            break
+                if doda_texts:
+                    print(f"  Loaded {len(doda_texts)} from {repo}", flush=True)
+                    break
+            except:
+                continue
+
+    if not doda_texts:
+        print("  No DODa data available locally. Skipping independent evaluation.", flush=True)
+        print("  (Would need to download DODa separately)", flush=True)
+    else:
+        print(f"  Evaluating {len(doda_texts)} DODa texts...", flush=True)
+        doda_results = []
+        for name, ar_sub, az_sub, algo, arch, vsz in ours_cfg:
+            ar_j = os.path.join(TOK_DIR, f"{ar_sub}.json")
+            az_j = os.path.join(TOK_DIR, f"{az_sub}.json")
+            if os.path.exists(ar_j) and os.path.exists(az_j):
+                tok = RawConcat(ar_j, az_j)
+                r = evaluate_on_doda(tok, name, "ours", algo, arch, vsz, doda_texts)
+                doda_results.append(r)
+                print(f"  {name}: F={r['fertility']:.3f} CPT={r['cpt']:.3f} EM={r['exact_match']:.2%}", flush=True)
+                del tok; gc.collect()
+
+        # Key externals
+        for name, repo in [("CaMeLBERT-MSA", "CAMeL-Lab/bert-base-arabic-camelbert-msa"),
+                           ("Qwen2.5-Darija", "GemMaroc/Qwen2.5-7B-Instruct-darija")]:
+            try:
+                tok = HFTok(repo, use_token=False)
+                r = evaluate_on_doda(tok, name, "external", "WordPiece", "shared", 0, doda_texts)
+                doda_results.append(r)
+                print(f"  {name}: F={r['fertility']:.3f} CPT={r['cpt']:.3f} EM={r['exact_match']:.2%}", flush=True)
+                del tok; gc.collect()
+            except Exception as e:
+                print(f"  {name} FAILED: {e}", flush=True)
+
+        doda_csv = os.path.join(BASE, "doda_independent_results.csv")
+        with open(doda_csv, "w", newline="") as f:
+            w = csv.DictWriter(f, fieldnames=list(doda_results[0].keys()))
+            w.writeheader()
+            for r in doda_results:
+                w.writerow(r)
+        print(f"\nDODa results saved to {doda_csv}", flush=True)
+
+    print("\n=== ALL DONE ===", flush=True)
+
+
+if __name__ == "__main__":
+    main()