""" ASR Benchmark — Gradio HuggingFace Space Evaluates Wav2Vec2-BERT and Gemma-3n models on Ghanaian language datasets. """ import os import gc import warnings import tempfile import torch import numpy as np import pandas as pd import torchaudio import jiwer import gradio as gr import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import matplotlib.patches as mpatches from datasets import load_dataset from transformers import ( Wav2Vec2ForCTC, Wav2Vec2Processor, Wav2Vec2BertForCTC, AutoProcessor, ) warnings.filterwarnings("ignore") # ── Device ────────────────────────────────────────────────────────────────── device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # ── Model & dataset registries ─────────────────────────────────────────────── WAV2VEC_MODELS = { "Dagbani (w2v-bert-2.0)": "FarmerlineML/w2v-bert-2.0_2026_dagbani_ASR", "Ewe (w2v-bert-2.0)": "FarmerlineML/w2v-bert-2.0_ewe_2", "Twi (w2v-bert-2.0)": "FarmerlineML/w2v-bert-2.0_twi_alpha_v1", } GEMMA_MODEL_ID = "unsloth/gemma-3n-E2B-it" DATASETS = { "UNICEF Dagbani": "ghananlpcommunity/UNICEF-Ghana-Dagbani-ASR", "UNICEF Twi": "ghananlpcommunity/UNICEF-Ghana-Twi-ASR", "UNICEF Ewe": "ghananlpcommunity/UNICEF-Ghana-Ewe-ASR", "Bible Ewe": "ghananlpcommunity/ewe-bible-audio-text-tts", "Bible Dagbani": "ghananlpcommunity/dagbani-bible-audio-text-tts", "Bible Twi": "ghananlpcommunity/asante-twi-bible-speech-text", } # Which wav2vec model is compatible with which dataset DATASET_MODEL_COMPAT = { "UNICEF Dagbani": ["Dagbani (w2v-bert-2.0)", "Gemma-3n (E2B-it)"], "UNICEF Twi": ["Twi (w2v-bert-2.0)", "Gemma-3n (E2B-it)"], "UNICEF Ewe": ["Ewe (w2v-bert-2.0)", "Gemma-3n (E2B-it)"], "Bible Dagbani": ["Dagbani (w2v-bert-2.0)", "Gemma-3n (E2B-it)"], "Bible Twi": ["Twi (w2v-bert-2.0)", "Gemma-3n (E2B-it)"], "Bible Ewe": ["Ewe (w2v-bert-2.0)", "Gemma-3n (E2B-it)"], } BLUE = "#185FA5" ORANGE = "#D85A30" # ── Helpers ────────────────────────────────────────────────────────────────── def load_dataset_samples(dataset_name: str, n_samples: int) -> list[dict]: ds_path = DATASETS[dataset_name] ds = load_dataset(ds_path, split="train", streaming=True, trust_remote_code=True) samples = [] for i, sample in enumerate(ds): if i >= n_samples: break audio = sample.get("audio") if audio is None: continue transcript = sample.get("text") or sample.get("transcript") if not transcript: continue samples.append({ "audio": audio["array"], "sampling_rate": audio["sampling_rate"], "text": transcript.lower().strip(), }) return samples def calculate_metrics(reference: str, hypothesis: str): if not hypothesis or not reference: return None, None return jiwer.cer(reference, hypothesis), jiwer.wer(reference, hypothesis) def resample_to_16k(audio_array: np.ndarray, sr: int) -> np.ndarray: if sr == 16000: return audio_array t = torch.from_numpy(audio_array).float() return torchaudio.transforms.Resample(sr, 16000)(t).numpy() # ── Wav2Vec2 inference ─────────────────────────────────────────────────────── class Wav2VecModel: def __init__(self, model_id: str): self.processor = AutoProcessor.from_pretrained(model_id) self.model = Wav2Vec2BertForCTC.from_pretrained( model_id, ignore_mismatched_sizes=True, torch_dtype=torch.float16, ).to(device) self.model.eval() def transcribe(self, audio_array: np.ndarray, sampling_rate: int) -> str: audio_16k = resample_to_16k(audio_array, sampling_rate) inputs = self.processor( audio_16k, sampling_rate=16000, return_tensors="pt", padding=True, ) inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(): logits = self.model(**inputs).logits predicted_ids = torch.argmax(logits, dim=-1) return self.processor.decode(predicted_ids[0]).lower().strip() def unload(self): del self.model, self.processor torch.cuda.empty_cache() gc.collect() # ── Gemma-3n inference (Unsloth) ───────────────────────────────────────────── class Gemma3nModel: def __init__(self): from unsloth import FastModel self.model, self.processor = FastModel.from_pretrained( model_name = GEMMA_MODEL_ID, dtype = None, max_seq_length = 1024, load_in_4bit = True, full_finetuning = False, ) self.model.eval() def transcribe(self, audio_array: np.ndarray, sampling_rate: int) -> str: audio_16k = resample_to_16k(audio_array, sampling_rate) messages = [ {"role": "system", "content": [{"type": "text", "text": "You are an assistant that transcribes speech accurately."}]}, {"role": "user", "content": [ {"type": "audio", "audio": audio_16k}, {"type": "text", "text": "Please transcribe this audio."}, ]}, ] inputs = self.processor.apply_chat_template( messages, add_generation_prompt=True, tokenize=True, return_tensors="pt", return_dict=True, ).to("cuda") inputs = { k: v.to(torch.bfloat16) if v.dtype in (torch.float32, torch.float64) else v for k, v in inputs.items() } with torch.inference_mode(): gen = self.model.generate(**inputs, max_new_tokens=200, do_sample=False) tokens = gen[0][inputs["input_ids"].shape[-1]:] return self.processor.decode(tokens, skip_special_tokens=True).lower().strip() def unload(self): del self.model, self.processor torch.cuda.empty_cache() gc.collect() # ── Chart helpers ───────────────────────────────────────────────────────────── def make_cer_wer_chart(df: pd.DataFrame) -> str: df = df.dropna(subset=["Avg WER", "Avg CER"]).copy() df["CER %"] = df["Avg CER"] * 100 df["WER %"] = df["Avg WER"] * 100 df["Model Short"] = df["Model"].apply(lambda x: x.split("/")[-1] if "/" in x else x) models = df["Model"].unique() ncols = min(len(models), 2) nrows = (len(models) + 1) // 2 fig, axes = plt.subplots(nrows, ncols, figsize=(8 * ncols, 5 * nrows)) axes = np.array(axes).flatten() if len(models) > 1 else [axes] fig.suptitle("CER & WER per Model / Dataset", fontsize=16, fontweight="bold") for i, model in enumerate(models): ax = axes[i] sub = df[df["Model"] == model].sort_values("Dataset") x = np.arange(len(sub)) w = 0.35 ax.bar(x - w / 2, sub["CER %"], w, label="CER %", color=BLUE, alpha=0.88) ax.bar(x + w / 2, sub["WER %"], w, label="WER %", color=ORANGE, alpha=0.88) ax.set_title(sub["Model Short"].iloc[0], fontsize=12, fontweight="bold") ax.set_xticks(x) ax.set_xticklabels(sub["Dataset"], rotation=30, ha="right", fontsize=9) ax.set_ylabel("Error Rate (%)") mx = max(sub["CER %"].max(), sub["WER %"].max()) * 1.15 or 100 ax.set_ylim(0, mx) for j in range(len(x)): ax.text(j - w/2, sub["CER %"].iloc[j] + mx*0.01, f"{sub['CER %'].iloc[j]:.1f}", ha="center", fontsize=8) ax.text(j + w/2, sub["WER %"].iloc[j] + mx*0.01, f"{sub['WER %'].iloc[j]:.1f}", ha="center", fontsize=8) ax.legend(fontsize=9) ax.spines[["top","right"]].set_visible(False) ax.grid(axis="y", linestyle="--", alpha=0.3) for j in range(len(models), len(axes)): axes[j].axis("off") plt.tight_layout() path = tempfile.mktemp(suffix="_cer_wer.png") plt.savefig(path, dpi=130, bbox_inches="tight") plt.close(fig) return path def make_domain_chart(df: pd.DataFrame) -> str: df = df.dropna(subset=["Avg WER"]).copy() df["Domain"] = df["Dataset"].apply(lambda x: "UNICEF" if "unicef" in x.lower() or "UNICEF" in x else "Bible") df["Language"] = df["Dataset"].apply( lambda x: x.replace("UNICEF ","").replace("Bible ","").strip() ) df["WER %"] = df["Avg WER"] * 100 df["Model Short"] = df["Model"].apply(lambda x: x.split("/")[-1] if "/" in x else x) models = df["Model Short"].unique() fig, axes = plt.subplots(1, 2, figsize=(14, 5)) fig.suptitle("Domain Comparison — UNICEF vs Bible", fontsize=15, fontweight="bold") # Avg WER per model × domain ax = axes[0] u_avg = [df[(df["Model Short"]==m) & (df["Domain"]=="UNICEF")]["WER %"].mean() for m in models] b_avg = [df[(df["Model Short"]==m) & (df["Domain"]=="Bible")]["WER %"].mean() for m in models] x = np.arange(len(models)) w = 0.35 ax.bar(x - w/2, u_avg, w, color=BLUE, alpha=0.88, label="UNICEF") ax.bar(x + w/2, b_avg, w, color=ORANGE, alpha=0.88, label="Bible") ax.set_title("Avg WER per Model × Domain", fontsize=11) ax.set_xticks(x) ax.set_xticklabels(models, rotation=20, ha="right", fontsize=9) ax.set_ylabel("WER %"); ax.set_ylim(0, 100) ax.legend(); ax.spines[["top","right"]].set_visible(False) # Domain gap bar ax2 = axes[1] gaps = [b - u for u, b in zip(u_avg, b_avg)] colors = [ORANGE if g > 0 else BLUE for g in gaps] ax2.bar(range(len(models)), gaps, color=colors, alpha=0.88) ax2.axhline(0, color="#888", linewidth=0.8, linestyle="--") ax2.set_title("Domain gap: Bible − UNICEF (pp)", fontsize=11) ax2.set_xticks(range(len(models))) ax2.set_xticklabels(models, rotation=20, ha="right", fontsize=9) ax2.set_ylabel("Percentage points"); ax2.spines[["top","right"]].set_visible(False) plt.tight_layout() path = tempfile.mktemp(suffix="_domain.png") plt.savefig(path, dpi=130, bbox_inches="tight") plt.close(fig) return path # ── Core evaluation function ────────────────────────────────────────────────── def run_evaluation( selected_models: list[str], selected_datasets: list[str], n_samples: int, progress=gr.Progress(track_tqdm=True), ): if not selected_models: return "⚠️ Please select at least one model.", None, None, None if not selected_datasets: return "⚠️ Please select at least one dataset.", None, None, None results_data = [] log_lines = [] all_model_labels = list(WAV2VEC_MODELS.keys()) + ["Gemma-3n (E2B-it)"] models_to_run = [m for m in all_model_labels if m in selected_models] for model_label in progress.tqdm(models_to_run, desc="Models"): log_lines.append(f"\n🔄 Loading model: {model_label}") yield "\n".join(log_lines), None, None, None # Load model try: if model_label == "Gemma-3n (E2B-it)": model = Gemma3nModel() model_id = GEMMA_MODEL_ID else: model_id = WAV2VEC_MODELS[model_label] model = Wav2VecModel(model_id) log_lines.append(f" ✅ Loaded: {model_id}") except Exception as e: log_lines.append(f" ❌ Failed to load {model_label}: {e}") yield "\n".join(log_lines), None, None, None continue # Evaluate on each selected compatible dataset for ds_label in selected_datasets: compat = DATASET_MODEL_COMPAT.get(ds_label, []) if model_label not in compat: log_lines.append(f" ⏭️ Skipping {ds_label} (incompatible language)") continue log_lines.append(f" 📂 Dataset: {ds_label} ({n_samples} samples)") yield "\n".join(log_lines), None, None, None try: samples = load_dataset_samples(ds_label, n_samples) except Exception as e: log_lines.append(f" ❌ Failed to load dataset: {e}") continue cers, wers, details = [], [], [] for j, sample in enumerate(samples): try: hyp = model.transcribe(sample["audio"], sample["sampling_rate"]) cer, wer = calculate_metrics(sample["text"], hyp) if cer is not None: cers.append(cer); wers.append(wer) details.append({ "ref": sample["text"][:80], "hyp": hyp[:80], "CER": round(cer, 4), "WER": round(wer, 4), }) except Exception as e: log_lines.append(f" ⚠️ Sample {j} error: {e}") if cers: avg_cer = float(np.mean(cers)) avg_wer = float(np.mean(wers)) results_data.append({ "Model": model_id, "Dataset": ds_label, "Avg CER": round(avg_cer, 4), "Avg WER": round(avg_wer, 4), "Num Samples": len(cers), }) log_lines.append(f" ✅ CER={avg_cer*100:.1f}% WER={avg_wer*100:.1f}%") else: log_lines.append(" ⚠️ No valid results for this combo.") # Unload model to free memory model.unload() log_lines.append(f" 🧹 {model_label} unloaded.\n") yield "\n".join(log_lines), None, None, None # Build outputs if not results_data: yield "\n".join(log_lines) + "\n\n❌ No results collected.", None, None, None return df = pd.DataFrame(results_data) log_lines.append("✅ Evaluation complete!") try: chart1 = make_cer_wer_chart(df) chart2 = make_domain_chart(df) except Exception as e: log_lines.append(f"⚠️ Chart generation failed: {e}") chart1, chart2 = None, None yield "\n".join(log_lines), df, chart1, chart2 # ── Gradio UI ───────────────────────────────────────────────────────────────── ALL_MODELS = list(WAV2VEC_MODELS.keys()) + ["Gemma-3n (E2B-it)"] ALL_DATASETS = list(DATASETS.keys()) with gr.Blocks(title="ASR Benchmark — Ghanaian Languages", theme=gr.themes.Soft()) as demo: gr.Markdown( """ # 🎙️ ASR Benchmark — Ghanaian Languages Evaluate **Wav2Vec2-BERT** fine-tuned models and **Gemma-3n** on Dagbani, Twi, and Ewe speech datasets. Models are loaded one at a time and unloaded after evaluation to minimise VRAM usage. """ ) with gr.Row(): with gr.Column(scale=1): gr.Markdown("### ⚙️ Configuration") model_selector = gr.CheckboxGroup( choices=ALL_MODELS, value=list(WAV2VEC_MODELS.keys()), label="Models to evaluate", ) dataset_selector = gr.CheckboxGroup( choices=ALL_DATASETS, value=ALL_DATASETS, label="Datasets to include", ) n_samples_slider = gr.Slider( minimum=5, maximum=200, value=50, step=5, label="Samples per dataset", info="More samples = better estimates but longer runtime.", ) run_btn = gr.Button("▶ Run Evaluation", variant="primary") with gr.Column(scale=2): gr.Markdown("### 📋 Evaluation Log") log_box = gr.Textbox( label="", lines=20, max_lines=40, interactive=False, placeholder="Logs will appear here once evaluation starts...", ) gr.Markdown("### 📊 Results") results_table = gr.DataFrame( label="Summary (CER & WER per model × dataset)", interactive=False, wrap=True, ) with gr.Row(): chart_cer_wer = gr.Image(label="CER & WER per Model", type="filepath") chart_domain = gr.Image(label="Domain Comparison", type="filepath") run_btn.click( fn=run_evaluation, inputs=[model_selector, dataset_selector, n_samples_slider], outputs=[log_box, results_table, chart_cer_wer, chart_domain], ) gr.Markdown( """ --- **Models:** [FarmerlineML](https://huggingface.co/FarmerlineML) | **Datasets:** [ghananlpcommunity](https://huggingface.co/ghananlpcommunity) | **Framework:** [Unsloth](https://github.com/unslothai/unsloth) + [🤗 Transformers](https://huggingface.co/transformers) """ ) if __name__ == "__main__": demo.launch()