eval_doda_independent.py

#!/usr/bin/env python3 -u
"""
eval_doda_independent.py — Evaluate on atlasia/DODa (87K Arabizi entries, independent dataset).
Avoids any contamination since DODa was not used in training.
"""

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")
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]")
_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 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]

def normalize_decode(s):
    s = s.replace("##", "")
    s = " ".join(s.split())
    return s


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(tok, name, source, algo, arch, vsz, texts):
    all_f, all_c = [], []
    em_ok, em_n = 0, 0

    for i, text in enumerate(texts):
        if (i + 1) % 10000 == 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)
            all_c.append(cpt)
            try:
                dec = tok.decode(ids, script)
                if normalize_decode(dec) == normalize_decode(text):
                    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 DODa
    from datasets import load_dataset
    # Use subset for speed
    N_EVAL = 10000
    print(f"Loading atlasia/DODa (evaluating on {N_EVAL} random subset)...", flush=True)
    ds = load_dataset("atlasia/DODa", split="train", token=HF_TOKEN, trust_remote_code=True)
    import random; random.seed(42)
    all_texts = [row["darija"].strip() for row in ds if row["darija"].strip()]
    texts = random.sample(all_texts, min(N_EVAL, len(all_texts)))
    del ds, all_texts; gc.collect()
    print(f"Evaluating on {len(texts)} DODa texts (100% Arabizi/Latin)", flush=True)

    results = []

    # Our best 3
    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(tok, name, "ours", algo, arch, vsz, texts)
            results.append(r)
            print(f"  F={r['fertility']:.3f} CPT={r['cpt']:.3f} EM={r['exact_match']:.2%}", flush=True)
            del tok; gc.collect()

    # All external tokenizers
    externals = [
        ("CaMeLBERT-MSA", "external_msa", "WordPiece", "shared", 30000,
         "CAMeL-Lab/bert-base-arabic-camelbert-msa", False),
        ("Asafaya-BERT", "external_msa", "WordPiece", "shared", 32000,
         "asafaya/bert-base-arabic", False),
        ("Aranizer-SP-86k", "external_msa", "SentencePiece", "shared", 86000,
         "riotu-lab/Aranizer-SP-86k", False),
        ("B2BERT", "external_msa", "WordPiece", "shared", 30000,
         "AHAAM/B2BERT", False),
        ("DarijaBERT-ar", "external_darija", "WordPiece", "shared", 80000,
         "SI2M-Lab/DarijaBERT", False),
        ("DarijaBERT-az", "external_darija", "WordPiece", "shared", 110000,
         "SI2M-Lab/DarijaBERT-arabizi", False),
        ("DarijaBERT-mix", "external_darija", "WordPiece", "shared", 160000,
         "SI2M-Lab/DarijaBERT-mix", False),
        ("Moroccan-Darija-Tokenizer", "external_darija", "BPE", "shared", 30000,
         "BounharAbdelaziz/Moroccan-Darija-Tokenizer", True),
        ("Translit-Darija", "external_darija", "BPE", "shared", 30000,
         "atlasia/Transliteration-Moroccan-Darija", True),
        ("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:
        print(f"\n{name} ({repo})", flush=True)
        try:
            tok = HFTok(repo, use_token=gated)
            r = evaluate(tok, name, src, algo, arch, vsz, texts)
            results.append(r)
            print(f"  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"  FAILED: {e}", flush=True)

    # Save
    out_csv = os.path.join(BASE, "doda_independent_results.csv")
    out_json = os.path.join(BASE, "doda_independent_results.json")
    with open(out_csv, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(results[0].keys()))
        w.writeheader()
        for r in results:
            w.writerow(r)
    with open(out_json, "w") as f:
        json.dump(results, f, indent=2)

    # Print
    print("\n" + "=" * 100, flush=True)
    hdr = f"{'Name':<35} {'V':>7} {'Fert':>7} {'CPT':>7} {'EM':>7} {'n':>8}"
    print(hdr, flush=True)
    print("-" * 100, flush=True)
    for r in sorted(results, key=lambda x: (0 if x["source"]=="ours" else 1, x["fertility"])):
        print(f"{r['name']:<35} {r['vocab_size']:>7,} {r['fertility']:>7.3f} {r['cpt']:>7.3f} {r['exact_match']:>7.2%} {r['n_texts']:>8,}", flush=True)
    print("=" * 100, flush=True)
    print(f"\nSaved: {out_csv}", flush=True)
    print("DONE!", flush=True)


if __name__ == "__main__":
    main()