NaijaVox unified demo: V1+V2 model switcher, green Space Grotesk UI

README.md

@@ -1,21 +1,23 @@
 ---
-title: NaijaVox V1 Demo
+title: NaijaVox Demo
 emoji: 🎙
 colorFrom: green
-colorTo: yellow
+colorTo: green
 sdk: gradio
-python_version: '3.11'
+python_version: "3.11"
 app_file: app.py
 pinned: false
 license: apache-2.0
 ---
 
-# NaijaVox-V1 Demo
+# NaijaVox Demo — Nigerian Speech Recognition
 
-Try Nigeria's open-weight speech recognition model -- Yoruba, Hausa, Igbo, Nigerian Pidgin, and Nigerian English.
+Unified demo for **NaijaVox-V1** and **NaijaVox-2.0** — open-weight ASR for Yoruba, Hausa, Igbo, Nigerian Pidgin, and Nigerian English.
 
-Record or upload audio, pick the language, and see it transcribed live.
+Switch between models, pick your language, record or upload audio, and transcribe.
 
-[Model card ->](https://huggingface.co/Axiveri/NaijaVox-V1)
+**Models:**
+- [NaijaVox-V1](https://huggingface.co/Axiveri/NaijaVox-V1) — Avg WER 27.9%
+- [NaijaVox-2.0](https://huggingface.co/Axiveri/NaijaVox-2.0) — Avg WER 22.58% (+19.1% better)
 
-Built by [Emmanuel Ariyo (Ememzyvisuals)](https://huggingface.co/ememzyvisuals) - Axiveri
+Built by [Emmanuel Ariyo (Ememzyvisuals)](https://huggingface.co/ememzyvisuals) · Axiveri

app.py

@@ -2,99 +2,421 @@ import gradio as gr
 import torch
 import numpy as np
 import librosa
-from transformers import WhisperForConditionalGeneration, WhisperFeatureExtractor, PreTrainedTokenizerFast, WhisperProcessor
+from transformers import (
+    WhisperForConditionalGeneration, WhisperFeatureExtractor,
+    WhisperProcessor, PreTrainedTokenizerFast,
+)
 from huggingface_hub import hf_hub_download
 
-MODEL_ID = "Axiveri/NaijaVox-V1"
-TARGET_SR = 16000  # Whisper always expects 16 kHz
+TARGET_SR = 16000
 
-print("Loading NaijaVox-V1... (this takes a minute on first load)")
-model = WhisperForConditionalGeneration.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
-fe = WhisperFeatureExtractor.from_pretrained(MODEL_ID)
-tok_path = hf_hub_download(repo_id=MODEL_ID, filename="tokenizer.json")
-tokenizer = PreTrainedTokenizerFast(tokenizer_file=tok_path)
-tokenizer.add_special_tokens({"additional_special_tokens": [t for t in ["<|startoftranscript|>","<|endoftext|>","<|transcribe|>","<|notimestamps|>","<|en|>","<|yo|>","<|ha|>","<|ig|>","<|pcm|>"] if t not in tokenizer.get_vocab()]})
-processor = WhisperProcessor(feature_extractor=fe, tokenizer=tokenizer)
-model.eval()
-
-VOCAB = processor.tokenizer.get_vocab()
-START_OF_TRANSCRIPT = VOCAB["<|startoftranscript|>"]
-TRANSCRIBE = VOCAB["<|transcribe|>"]
-NOTIMESTAMPS = VOCAB["<|notimestamps|>"]
+MODEL_IDS = {
+    "NaijaVox-V1":  "Axiveri/NaijaVox-V1",
+    "NaijaVox-2.0": "Axiveri/NaijaVox-2.0",
+}
 
 LANGUAGES = {
-    "Nigerian English": "<|en|>",
-    "Nigerian Pidgin": "<|pcm|>",
-    "Yoruba": "<|yo|>",
-    "Hausa": "<|ha|>",
-    "Igbo": "<|ig|>",
+    "🇳🇬 Nigerian English": "<|en|>",
+    "🇳🇬 Nigerian Pidgin":  "<|pcm|>",
+    "🇳🇬 Yoruba":           "<|yo|>",
+    "🇳🇬 Hausa":            "<|ha|>",
+    "🇳🇬 Igbo":             "<|ig|>",
 }
 
-print("Model loaded. Ready.")
+MODEL_CACHE = {}
+
+
+def load_model(model_key):
+    if model_key in MODEL_CACHE:
+        return MODEL_CACHE[model_key]
+
+    model_id = MODEL_IDS[model_key]
+    print(f"Loading {model_key} from {model_id}...")
+    model = WhisperForConditionalGeneration.from_pretrained(
+        model_id, torch_dtype=torch.float32
+    )
+
+    # Try standard load first; fall back to manual tokenizer (same approach as V1 deploy)
+    try:
+        processor = WhisperProcessor.from_pretrained(model_id)
+        vocab = processor.tokenizer.get_vocab()
+        assert "<|pcm|>" in vocab and "<|ig|>" in vocab
+    except Exception as e:
+        print(f"  Standard load failed ({e}), using manual tokenizer...")
+        fe  = WhisperFeatureExtractor.from_pretrained(model_id)
+        tok = hf_hub_download(repo_id=model_id, filename="tokenizer.json")
+        tokenizer = PreTrainedTokenizerFast(tokenizer_file=tok)
+        tokenizer.add_special_tokens({
+            "additional_special_tokens": [
+                t for t in [
+                    "<|startoftranscript|>", "<|endoftext|>", "<|transcribe|>",
+                    "<|notimestamps|>", "<|en|>", "<|yo|>", "<|ha|>", "<|ig|>", "<|pcm|>",
+                ]
+                if t not in tokenizer.get_vocab()
+            ]
+        })
+        processor = WhisperProcessor(feature_extractor=fe, tokenizer=tokenizer)
+        vocab = processor.tokenizer.get_vocab()
+
+    model.eval()
+    MODEL_CACHE[model_key] = (model, processor, vocab)
+    print(f"  {model_key} ready.")
+    return model, processor, vocab
+
 
+# Pre-load V1 on startup so the demo is immediately responsive
+print("Pre-loading NaijaVox-V1...")
+load_model("NaijaVox-V1")
+print("Startup complete.")
 
-def transcribe(audio, language):
+
+def transcribe(audio, language, model_key):
     if audio is None:
-        return "Please record or upload audio first."
+        return "⚠️ Record or upload audio first, then tap Transcribe."
+
+    try:
+        model, processor, vocab = load_model(model_key)
+    except Exception as e:
+        return f"❌ Error loading {model_key}: {e}"
 
     sr, arr = audio
     arr = np.array(arr, dtype=np.float32)
-
-    # Convert stereo to mono
     if arr.ndim > 1:
         arr = arr.mean(axis=1)
-
-    # Normalise int16 PCM to float32 [-1, 1]
     if np.abs(arr).max() > 1.0:
-        arr = arr / 32768.0
-
-    # Resample to 16 kHz — browser mic records at 44100 Hz; Whisper needs 16000 Hz
+        arr /= 32768.0
     if sr != TARGET_SR:
         arr = librosa.resample(arr, orig_sr=sr, target_sr=TARGET_SR)
-        sr = TARGET_SR
 
-    lang_token = LANGUAGES[language]
-    lang_id = VOCAB[lang_token]
-    decoder_input_ids = torch.tensor([[START_OF_TRANSCRIPT, lang_id, TRANSCRIBE, NOTIMESTAMPS]])
+    lang_id = vocab[LANGUAGES[language]]
+    start   = vocab["<|startoftranscript|>"]
+    trans   = vocab["<|transcribe|>"]
+    nots    = vocab["<|notimestamps|>"]
 
-    inputs = processor.feature_extractor(
-        arr, sampling_rate=sr, return_tensors="pt"
+    dec_ids = torch.tensor([[start, lang_id, trans, nots]])
+    inputs  = processor.feature_extractor(
+        arr, sampling_rate=TARGET_SR, return_tensors="pt"
     ).input_features
 
     with torch.no_grad():
-        generated = model.generate(
+        gen = model.generate(
             input_features=inputs,
-            decoder_input_ids=decoder_input_ids,
+            decoder_input_ids=dec_ids,
             max_new_tokens=200,
         )
 
-    return processor.tokenizer.decode(generated[0], skip_special_tokens=True)
+    return processor.tokenizer.decode(gen[0], skip_special_tokens=True).strip()
+
+
+# ── Theme ──────────────────────────────────────────────────────────────────────
+GREEN = gr.themes.Color(
+    c50="#f0fdf4", c100="#dcfce7", c200="#bbf7d0",
+    c300="#86efac", c400="#4ade80", c500="#22c55e",
+    c600="#16a34a", c700="#15803d", c800="#166534",
+    c900="#14532d", c950="#052e16",
+)
+
+theme = gr.themes.Base(
+    primary_hue=GREEN,
+    secondary_hue=GREEN,
+    neutral_hue=GREEN,
+    font=gr.themes.GoogleFont("Space Grotesk"),
+    font_mono=gr.themes.GoogleFont("JetBrains Mono"),
+).set(
+    body_background_fill="#060d07",
+    body_background_fill_dark="#060d07",
+    block_background_fill="#0d1a0f",
+    block_background_fill_dark="#0d1a0f",
+    block_border_color="#1c3824",
+    block_border_color_dark="#1c3824",
+    block_border_width="1px",
+    block_radius="14px",
+    button_primary_background_fill="linear-gradient(135deg,#16a34a 0%,#22c55e 100%)",
+    button_primary_background_fill_hover="linear-gradient(135deg,#22c55e 0%,#4ade80 100%)",
+    button_primary_text_color="#ffffff",
+    button_primary_border_color="transparent",
+    button_primary_border_color_hover="transparent",
+    input_background_fill="#112015",
+    input_background_fill_dark="#112015",
+    input_border_color="#1c3824",
+    input_border_color_focus="#22c55e",
+    body_text_color="#dcfce7",
+    body_text_color_dark="#dcfce7",
+    body_text_color_subdued="#6ee7a0",
+    block_label_text_color="#4ade80",
+    block_label_text_color_dark="#4ade80",
+    block_label_text_weight="600",
+    block_title_text_color="#4ade80",
+    block_title_text_color_dark="#4ade80",
+    checkbox_label_background_fill="#112015",
+    checkbox_label_background_fill_hover="#1a3020",
+    checkbox_label_background_fill_selected="#0d2e1a",
+    checkbox_label_border_color="#1c3824",
+    checkbox_label_border_color_hover="#22c55e",
+    checkbox_label_text_color="#dcfce7",
+    table_even_background_fill="#0d1a0f",
+    table_odd_background_fill="#112015",
+    slider_color="#22c55e",
+)
+
+CSS = """
+@import url('https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@300;400;500;600;700&family=JetBrains+Mono:wght@400;600&display=swap');
 
+*, *::before, *::after {
+    font-family: 'Space Grotesk', system-ui, sans-serif !important;
+    box-sizing: border-box;
+}
+
+body, .gradio-container { background: #060d07 !important; }
+
+.gradio-container {
+    max-width: 820px !important;
+    margin: 0 auto !important;
+    padding: 0 1rem !important;
+}
+
+/* ── Header ── */
+.nv-header { text-align: center; padding: 2rem 0 0.5rem; }
+
+.nv-logo {
+    display: block;
+    margin: 0 auto 1.2rem;
+    max-width: 100%;
+    border-radius: 12px;
+}
+
+.nv-title {
+    font-size: clamp(2.4rem, 6vw, 3.6rem);
+    font-weight: 700;
+    letter-spacing: -0.03em;
+    background: linear-gradient(135deg, #22c55e 0%, #4ade80 55%, #86efac 100%);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    background-clip: text;
+    line-height: 1.1;
+    margin: 0 0 0.5rem;
+}
+
+.nv-sub {
+    color: #6ee7a0;
+    font-size: 0.95rem;
+    margin: 0 0 1.8rem;
+    letter-spacing: 0.01em;
+}
+
+/* ── Model info cards ── */
+.model-info-row {
+    display: flex;
+    gap: 0.75rem;
+    margin-bottom: 0.75rem;
+}
+
+.model-info-card {
+    flex: 1;
+    background: #0d1a0f;
+    border: 1px solid #1c3824;
+    border-radius: 12px;
+    padding: 0.9rem 1.1rem;
+    position: relative;
+}
+
+.model-info-card.featured { border-color: #22c55e; background: #0d2315; }
+
+.mic-badge {
+    position: absolute;
+    top: -10px; right: 12px;
+    background: #22c55e;
+    color: #052e16;
+    font-size: 0.65rem;
+    font-weight: 700;
+    padding: 2px 10px;
+    border-radius: 20px;
+    letter-spacing: 0.06em;
+}
 
-with gr.Blocks(title="NaijaVox-V1 Demo") as demo:
-    gr.Markdown(
-        """
-        # 🇳🇬 NaijaVox-V1 — Nigerian Speech Recognition
-        Open-weight speech-to-text for **Yoruba, Hausa, Igbo, Nigerian Pidgin, and Nigerian English**.
-        Record or upload audio, pick the language, and transcribe.
+.card-name { color: #4ade80; font-weight: 700; font-size: 0.95rem; margin: 0 0 0.25rem; }
+.card-stats { color: #6ee7a0; font-size: 0.78rem; }
+.card-stats strong { color: #22c55e; }
 
-        [Model card](https://huggingface.co/Axiveri/NaijaVox-V1) · Built by Emmanuel Ariyo (Ememzyvisuals) · Axiveri
-        """
+/* ── Section labels ── */
+.sl {
+    color: #4ade80;
+    font-size: 0.75rem;
+    font-weight: 700;
+    letter-spacing: 0.09em;
+    text-transform: uppercase;
+    margin: 1rem 0 0.35rem;
+}
+
+/* ── Model radio styled as pill tabs ── */
+#model-radio .wrap { gap: 0.6rem !important; }
+
+#model-radio label {
+    background: #0d1a0f !important;
+    border: 1.5px solid #1c3824 !important;
+    border-radius: 10px !important;
+    padding: 0.65rem 1rem !important;
+    color: #6ee7a0 !important;
+    font-weight: 500 !important;
+    cursor: pointer !important;
+    transition: all 0.16s !important;
+    flex: 1 !important;
+    text-align: center !important;
+}
+
+#model-radio label:hover {
+    border-color: #22c55e !important;
+    color: #dcfce7 !important;
+    background: #112015 !important;
+}
+
+#model-radio label:has(input:checked) {
+    border-color: #22c55e !important;
+    background: linear-gradient(135deg,#0d2e1a,#112015) !important;
+    color: #4ade80 !important;
+    box-shadow: 0 0 0 1px #22c55e, 0 4px 16px rgba(34,197,94,0.18) !important;
+}
+
+/* ── Transcribe button ── */
+#transcribe-btn {
+    margin-top: 0.5rem !important;
+}
+
+#transcribe-btn button {
+    font-size: 1.05rem !important;
+    font-weight: 700 !important;
+    letter-spacing: 0.04em !important;
+    padding: 0.9rem !important;
+    border-radius: 14px !important;
+    box-shadow: 0 4px 20px rgba(34,197,94,0.28) !important;
+    transition: all 0.18s ease !important;
+}
+
+#transcribe-btn button:hover {
+    transform: translateY(-2px) !important;
+    box-shadow: 0 8px 28px rgba(34,197,94,0.4) !important;
+}
+
+#transcribe-btn button:active { transform: translateY(0) !important; }
+
+/* ── Output ── */
+#output-box textarea {
+    font-size: 1.1rem !important;
+    line-height: 1.75 !important;
+    min-height: 96px !important;
+    color: #dcfce7 !important;
+}
+
+/* ── Footer ── */
+.nv-footer {
+    text-align: center;
+    padding: 1.5rem 0 2rem;
+    color: #6ee7a0;
+    font-size: 0.82rem;
+    border-top: 1px solid #1c3824;
+    margin-top: 1.5rem;
+}
+
+.nv-footer a { color: #4ade80; text-decoration: none; }
+.nv-footer a:hover { text-decoration: underline; }
+
+/* ── Scrollbar ── */
+::-webkit-scrollbar { width: 5px; height: 5px; }
+::-webkit-scrollbar-track { background: #060d07; }
+::-webkit-scrollbar-thumb { background: #1c3824; border-radius: 3px; }
+::-webkit-scrollbar-thumb:hover { background: #22c55e; }
+"""
+
+# ── UI ─────────────────────────────────────────────────────────────────────────
+with gr.Blocks(theme=theme, css=CSS, title="NaijaVox Demo — Nigerian ASR") as demo:
+
+    # Header
+    gr.HTML("""
+    <div class="nv-header">
+      <h1 class="nv-title">NaijaVox</h1>
+      <p class="nv-sub">
+        Open-weight Nigerian speech recognition &nbsp;·&nbsp;
+        Yoruba &nbsp;·&nbsp; Hausa &nbsp;·&nbsp; Igbo &nbsp;·&nbsp; Pidgin &nbsp;·&nbsp; Nigerian English
+      </p>
+    </div>
+    """)
+
+    # Model info cards (visual only)
+    gr.HTML("""
+    <div class="model-info-row">
+      <div class="model-info-card">
+        <p class="card-name">NaijaVox-V1</p>
+        <p class="card-stats">Avg WER <strong>27.9%</strong> &nbsp;·&nbsp; LoRA r=32 &nbsp;·&nbsp; 13,866 samples</p>
+      </div>
+      <div class="model-info-card featured">
+        <span class="mic-badge">IMPROVED</span>
+        <p class="card-name">NaijaVox-2.0</p>
+        <p class="card-stats">
+          Avg WER <strong>22.58%</strong> &nbsp;·&nbsp; LoRA r=64 + fc1/fc2 &nbsp;·&nbsp; 25,866 samples<br/>
+          <span style="color:#86efac;font-size:0.74rem;">SpecAugment · Noise augmentation · +19.1% better avg</span>
+        </p>
+      </div>
+    </div>
+    """)
+
+    # Model selector
+    gr.HTML('<p class="sl">Model</p>')
+    model_radio = gr.Radio(
+        choices=["NaijaVox-V1", "NaijaVox-2.0"],
+        value="NaijaVox-2.0",
+        label="",
+        elem_id="model-radio",
     )
-    with gr.Row():
-        with gr.Column():
-            audio_input = gr.Audio(sources=["microphone", "upload"], type="numpy", label="Audio")
-            lang_input = gr.Dropdown(
-                choices=list(LANGUAGES.keys()), value="Nigerian English", label="Language"
-            )
-            btn = gr.Button("Transcribe", variant="primary")
-        with gr.Column():
-            output = gr.Textbox(label="Transcription", lines=8)
-
-    btn.click(fn=transcribe, inputs=[audio_input, lang_input], outputs=output)
-
-    gr.Markdown(
-        "_Running on free CPU hardware — transcription may take 20-60 seconds per clip._"
+
+    # Language + Audio
+    gr.HTML('<p class="sl">Language</p>')
+    lang_dropdown = gr.Dropdown(
+        choices=list(LANGUAGES.keys()),
+        value="🇳🇬 Nigerian English",
+        label="",
     )
 
+    gr.HTML('<p class="sl">Audio</p>')
+    audio_input = gr.Audio(
+        sources=["microphone", "upload"],
+        type="numpy",
+        label="",
+    )
+
+    # Transcribe button
+    transcribe_btn = gr.Button(
+        "Transcribe",
+        variant="primary",
+        elem_id="transcribe-btn",
+    )
+
+    # Output
+    gr.HTML('<p class="sl">Transcription</p>')
+    output_box = gr.Textbox(
+        label="",
+        placeholder="Transcription will appear here...",
+        lines=4,
+        elem_id="output-box",
+    )
+
+    # Wire up
+    transcribe_btn.click(
+        fn=transcribe,
+        inputs=[audio_input, lang_dropdown, model_radio],
+        outputs=output_box,
+    )
+
+    # Footer
+    gr.HTML("""
+    <div class="nv-footer">
+      <a href="https://huggingface.co/Axiveri/NaijaVox-V1">NaijaVox-V1</a>
+      &nbsp;·&nbsp;
+      <a href="https://huggingface.co/Axiveri/NaijaVox-2.0">NaijaVox-2.0</a>
+      &nbsp;·&nbsp;
+      <a href="https://huggingface.co/collections/Axiveri/naijavox-nigerian-speech-recognition">NaijaVox Collection</a>
+      &nbsp;·&nbsp;
+      Built by <a href="https://huggingface.co/ememzyvisuals">Emmanuel Ariyo</a> · Axiveri
+    </div>
+    """)
+
 demo.launch()