add: phoneme eval code

Makefile

@@ -11,3 +11,8 @@ quality:
 	python -m black --check --line-length 119 .
 	python -m isort --check-only .
 	ruff check .
+
+
+.PHONY: eval
+eval:
+	python -m src.phoneme_eval

requirements.txt

@@ -13,4 +13,8 @@ python-dateutil
 tqdm
 transformers
 tokenizers>=0.15.0
-sentencepiece
+sentencepiece
+torchaudio
+torch
+nltk
+g2p-en

src/phoneme_eval.py

@@ -0,0 +1,186 @@
+import os
+import json
+import time
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+import torchaudio
+from datasets import load_dataset
+from transformers import (
+    Wav2Vec2Processor,
+    HubertForCTC,
+    Wav2Vec2ForCTC,
+)
+
+
+@dataclass
+class EvalConfig:
+    dataset_name: str = "mirfan899/phoneme_asr"
+    split: str = "train"
+    max_examples: int = 100
+    results_dir: str = "eval-results"  # relative to CWD
+    model_sha: str = ""
+    model_dtype: str = "float16"
+
+
+def load_audio_array(example):
+    return example["audio"]["array"]
+
+
+def load_models(device: torch.device):
+    base_proc = Wav2Vec2Processor.from_pretrained("facebook/hubert-large-ls960-ft")
+    base_model = HubertForCTC.from_pretrained("facebook/hubert-large-ls960-ft").to(device).eval()
+
+    timit_proc = Wav2Vec2Processor.from_pretrained("vitouphy/wav2vec2-xls-r-300m-timit-phoneme")
+    timit_model = Wav2Vec2ForCTC.from_pretrained("vitouphy/wav2vec2-xls-r-300m-timit-phoneme").to(device).eval()
+
+    return (base_proc, base_model), (timit_proc, timit_model)
+
+
+def clean_cmu(text: str) -> str:
+    res = text.replace("0", "").replace("1", "").replace("2", "").replace("-", "").strip()
+    return res.lower()
+
+
+def cmu_to_ipa(cmu_sentence: str) -> str:
+    cmu_map = {
+        "AA": "ɑ", "AE": "æ", "AH": "ʌ", "AH0": "ə", "AO": "ɔ", "AW": "aʊ", "AY": "aɪ",
+        "EH": "ɛ", "ER": "ɝ", "ER0": "ɚ", "EY": "eɪ", "IH": "ɪ", "IY": "i", "OW": "oʊ",
+        "OY": "ɔɪ", "UH": "ʊ", "UW": "u", "B": "b", "CH": "tʃ", "D": "d", "DH": "ð",
+        "F": "f", "G": "ɡ", "HH": "h", "JH": "dʒ", "K": "k", "L": "l", "M": "m",
+        "N": "n", "NG": "ŋ", "P": "p", "R": "r", "S": "s", "SH": "ʃ", "T": "t",
+        "TH": "θ", "V": "v", "W": "w", "Y": "j", "Z": "z", "ZH": "ʒ",
+    }
+    ipa_tokens = []
+    for word in cmu_sentence.strip().split():
+        i = 0
+        while i < len(word):
+            if i + 2 <= len(word) and word[i:i+2].upper() in cmu_map:
+                ipa_tokens.append(cmu_map[word[i:i+2].upper()]); i += 2
+            elif word[i].upper() in cmu_map:
+                ipa_tokens.append(cmu_map[word[i].upper()]); i += 1
+            else:
+                ipa_tokens.append(word[i].lower()); i += 1
+        ipa_tokens.append(" ")
+    return "".join(ipa_tokens)
+
+
+def align_sequences(seq1: str, seq2: str):
+    n, m = len(seq1), len(seq2)
+    dp = np.zeros((n + 1, m + 1), dtype=np.float32)
+    back = np.empty((n + 1, m + 1), dtype="U1")
+    dp[:, 0] = np.arange(n + 1)
+    dp[0, :] = np.arange(m + 1)
+    back[:, 0] = "D"; back[0, :] = "I"; back[0, 0] = ""
+    for i in range(1, n + 1):
+        for j in range(1, m + 1):
+            cost = 0.0 if seq1[i - 1] == seq2[j - 1] else 1.0
+            opts = [(dp[i - 1][j] + 1, "D"), (dp[i][j - 1] + 1, "I"), (dp[i - 1][j - 1] + cost, "M")]
+            dp[i][j], back[i][j] = min(opts, key=lambda x: x[0])
+    i, j = n, m; a1, a2 = [], []
+    while i > 0 or j > 0:
+        mv = back[i][j]
+        if mv == "M": a1.append(seq1[i - 1]); a2.append(seq2[j - 1]); i -= 1; j -= 1
+        elif mv == "D": a1.append(seq1[i - 1]); a2.append("-"); i -= 1
+        elif mv == "I": a1.append("-"); a2.append(seq2[j - 1]); j -= 1
+        else: break
+    a1.reverse(); a2.reverse(); return a1, a2
+
+
+def calculate_per(ref_seq: str, hyp_seq: str) -> float:
+    ref_seq = ref_seq.replace(" ", ""); hyp_seq = hyp_seq.replace(" ", "")
+    aligned_ref, aligned_hyp = align_sequences(ref_seq, hyp_seq)
+    s = d = i = 0
+    for r, h in zip(aligned_ref, aligned_hyp):
+        if r == h: continue
+        if r == "-": i += 1
+        elif h == "-": d += 1
+        else: s += 1
+    n = len(ref_seq)
+    return ((s + d + i) / n) * 100.0 if n > 0 else 0.0
+
+
+def run_hubert_base(proc, model, wav, device):
+    inputs = proc(wav, sampling_rate=16000, return_tensors="pt", padding=True).input_values.to(device)
+    with torch.no_grad():
+        logits = model(inputs).logits
+    ids = torch.argmax(logits, dim=-1)
+    text = proc.batch_decode(ids)[0]
+    return text
+
+
+def run_timit(proc, model, wav, device):
+    inputs = proc(wav, sampling_rate=16000, return_tensors="pt", padding=True).to(device)
+    with torch.no_grad():
+        logits = model(inputs.input_values, attention_mask=inputs.attention_mask).logits
+    ids = torch.argmax(logits, dim=-1)
+    ph = proc.batch_decode(ids)
+    return "".join(ph)
+
+
+def evaluate(config: EvalConfig):
+    os.makedirs(config.results_dir, exist_ok=True)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    (base_proc, base_model), (timit_proc, timit_model) = load_models(device)
+
+    ds = load_dataset(config.dataset_name, split=config.split)
+    uniq = set(ds.unique("phonetic"))
+    ds = ds.filter(lambda x: x["phonetic"] in uniq)
+    ds = ds.filter(lambda x: len(x["phonetic"].split()) >= 10)
+    ds = ds.shuffle(seed=42).select(range(min(config.max_examples, len(ds))))
+
+    results = {
+        "results": {
+            "phoneme_dev": {},
+            "phoneme_test": {},
+        },
+        "config": {
+            "model_name": "phoneme/baselines",
+            "model_sha": config.model_sha,
+            "model_dtype": config.model_dtype,
+        },
+    }
+
+    # Simple split into dev/test halves
+    mid = len(ds) // 2
+    halves = [("phoneme_dev", ds.select(range(0, mid))), ("phoneme_test", ds.select(range(mid, len(ds))))]
+
+    for split_key, subset in halves:
+        per_scores_hubert = []
+        per_scores_timit = []
+        for ex in subset:
+            wav = ex["audio"]["array"]
+            ref = cmu_to_ipa(clean_cmu(ex["phonetic"]))
+
+            # HuBERT base → CMU→IPA
+            base_pred_cmu = run_hubert_base(base_proc, base_model, wav, device)
+            base_pred_ipa = cmu_to_ipa(base_pred_cmu)
+            per_scores_hubert.append(calculate_per(ref, base_pred_ipa))
+
+            # TIMIT phoneme model (already phoneme-like)
+            timit_pred = run_timit(timit_proc, timit_model, wav, device)
+            timit_pred_ipa = timit_pred  # leave as-is
+            per_scores_timit.append(calculate_per(ref, timit_pred_ipa))
+
+        # record mean PER per model under this split
+        results["results"][split_key] = {
+            "hubert_base": {"per": float(np.mean(per_scores_hubert)) if per_scores_hubert else None},
+            "timit_model": {"per": float(np.mean(per_scores_timit)) if per_scores_timit else None},
+        }
+
+    # Save a single combined result file
+    ts = int(time.time())
+    out_path = os.path.join(config.results_dir, f"results_{ts}.json")
+    with open(out_path, "w", encoding="utf-8") as f:
+        json.dump(results, f, ensure_ascii=False, indent=2)
+    return out_path
+
+
+if __name__ == "__main__":
+    cfg = EvalConfig()
+    path = evaluate(cfg)
+    print(f"Saved results to {path}")
+
+