Blame for app.py

"""
Arabic Isolated-Letter Pronunciation Verifier
=============================================
ZeroGPU Gradio Space serving
`masumtechnonext/wav2vec2-arabic-letter-verifier`.

The model is a 29-way classifier: the 28 Arabic letters plus a trained
"Unknown" class (silence, noise, time-reversed clips, cross-letter mixups).
A prediction is only *accepted* when the top softmax probability clears the
calibrated threshold from `calibration.json` AND the argmax is not Unknown —
so out-of-vocabulary audio gets rejected instead of being force-fitted onto
a real letter.

Free-tier notes
---------------
This Space runs on ZeroGPU (`zero-a10g`). GPU seconds are deducted from the
*caller's* daily quota (anonymous: 2 min/day, signed-in free account:
5 min/day, PRO: 25 min/day). One prediction costs ~0.09 GPU-seconds, so a
signed-in free account gets roughly 3,500 predictions per day.
"""
import json
import os

import gradio as gr
import numpy as np
import spaces
import torch
import torchaudio
from huggingface_hub import hf_hub_download
from transformers import Wav2Vec2FeatureExtractor, Wav2Vec2ForSequenceClassification

MODEL_ID = os.environ.get(
    "HF_MODEL_ID", "masumtechnonext/wav2vec2-arabic-letter-verifier"
)
HF_TOKEN = os.environ.get("HF_TOKEN") or None
SR = 16000
MAX_SECONDS = 8.0
DETECT_ONLY = "__detect_only__"

# Canonical teaching order (the dataset's own `alphabet_order`) with the
# Arabic glyph for each class. The names are the model's own label strings,
# so whatever the UI sends is exactly what the model was trained on.
LETTERS = [
    ("Alif", "ا"), ("Ba", "ب"), ("Ta", "ت"), ("Tha", "ث"), ("Jeem", "ج"),
    ("Haa", "ح"), ("Kha", "خ"), ("Dal", "د"), ("Dhal", "ذ"), ("Raa", "ر"),
    ("Zay", "ز"), ("Seen", "س"), ("Sheen", "ش"), ("Saad", "ص"), ("Daad", "ض"),
    ("Taa", "ط"), ("Dhaa", "ظ"), ("Ain", "ع"), ("Ghain", "غ"), ("Faa", "ف"),
    ("Qaf", "ق"), ("Kaf", "ك"), ("Laam", "ل"), ("Meem", "م"), ("Noon", "ن"),
    ("Waw", "و"), ("Ha", "ه"), ("Yaa", "ي"),
]
GLYPH = dict(LETTERS)

# ---------------------------------------------------------------- model load
feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(MODEL_ID, token=HF_TOKEN)
model = Wav2Vec2ForSequenceClassification.from_pretrained(MODEL_ID, token=HF_TOKEN)
model.eval()
id2label = {int(k): v for k, v in model.config.id2label.items()}
unknown_id = next(i for i, label in id2label.items() if label == "Unknown")

with open(hf_hub_download(MODEL_ID, "calibration.json", token=HF_TOKEN)) as fh:
    CALIBRATION = json.load(fh)
THRESHOLD = float(CALIBRATION["confidence_threshold"])
METRICS = CALIBRATION.get("validation_metrics_at_threshold", {})


# ------------------------------------------------------------------- helpers
def _to_mono_float32(waveform: np.ndarray) -> np.ndarray:
    waveform = np.asarray(waveform)
    if waveform.ndim > 1:
        waveform = waveform.mean(axis=-1)
    if np.issubdtype(waveform.dtype, np.integer):
        waveform = waveform.astype(np.float32) / float(np.iinfo(waveform.dtype).max)
    return waveform.astype(np.float32)


def _resample(waveform: np.ndarray, orig_sr: int) -> np.ndarray:
    """Resample to 16 kHz. torchaudio when available, numpy fallback."""
    if orig_sr == SR:
        return waveform
    try:
        tensor = torch.from_numpy(waveform)
        return (
            torchaudio.functional.resample(tensor, orig_sr, SR)
            .contiguous()
            .numpy()
            .astype(np.float32)
        )
    except Exception:
        duration = waveform.shape[0] / float(orig_sr)
        target_len = max(1, int(round(duration * SR)))
        source_idx = np.linspace(0.0, waveform.shape[0] - 1, num=target_len)
        return np.interp(
            source_idx, np.arange(waveform.shape[0], dtype=np.float64), waveform
        ).astype(np.float32)


def _prepare(audio):
    """Gradio (sample_rate, waveform) -> 16 kHz mono float32 waveform."""
    sample_rate, waveform = audio
    waveform = _to_mono_float32(waveform)
    waveform = _resample(waveform, int(sample_rate))
    limit = int(MAX_SECONDS * SR)
    if waveform.shape[0] > limit:
        waveform = waveform[:limit]
    return waveform


@spaces.GPU
def _infer(waveform: np.ndarray) -> np.ndarray:
    """Run the model on the GPU and return the 29-class probability vector."""
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    inputs = feature_extractor(waveform, sampling_rate=SR, return_tensors="pt")
    inputs = {k: v.to(device) for k, v in inputs.items()}
    with torch.inference_mode():
        logits = model(**inputs).logits[0].float().cpu().numpy()
    probs = np.exp(logits - logits.max())
    probs /= probs.sum()
    return probs


def _letter_md(label: str) -> str:
    glyph = GLYPH.get(label)
    return f"{glyph} **{label}**" if glyph else f"**{label}**"


def _prob_table(probs: np.ndarray, top_n: int = 5) -> str:
    order = np.argsort(probs)[::-1][:top_n]
    rows = [
        "| হরফ · Letter | সম্ভাবনা · Probability |",
        "|---|---|",
    ]
    for idx in order:
        label = id2label[int(idx)]
        glyph = GLYPH.get(label, "")
        name = f"{glyph} {label}" if glyph else label
        rows.append(f"| {name} | {probs[idx]:.2%} |")
    return "\n".join(rows)


def _analyse(audio, target_letter):
    """Shared core: returns a result dict, never raises."""
    if audio is None:
        return {
            "ok": False,
            "verdict_md": "### ⚠️ অডিও নেই · No audio\n\n"
            "আগে রেকর্ড করুন বা একটি অডিও ফাইল আপলোড করুন।",
            "details_md": "",
            "payload": {"error": "no_audio"},
        }

    try:
        waveform = _prepare(audio)
        if waveform.size == 0:
            raise ValueError("empty waveform")
        probs = _infer(waveform)
    except Exception as exc:  # noqa: BLE001 - surface a friendly message
        message = str(exc)
        if "ZeroGPU" in message or "quota" in message.lower():
            verdict_md = (
                "### ⏳ আজকের ফ্রি কোটা শেষ · Free GPU quota exhausted\n\n"
                "আজকের ZeroGPU কোটা ফুরিয়ে গেছে। Hugging Face-এ **লগইন** করলে "
                "আরও বেশি রান পাওয়া যাবে, অথবা কাল আবার চেষ্টা করুন।\n\n"
                f"<sub>{message[:160]}</sub>"
            )
        else:
            verdict_md = (
                "### ⚠️ প্রক্রিয়ায় সমস্যা · Processing error\n\n"
                f"<sub>{message[:220]}</sub>"
            )
        return {
            "ok": False,
            "verdict_md": verdict_md,
            "details_md": "",
            "payload": {"error": message[:300]},
        }

    pred_id = int(np.argmax(probs))
    confidence = float(probs[pred_id])
    accepted = confidence >= THRESHOLD and pred_id != unknown_id
    predicted_label = id2label[pred_id] if accepted else "Unrecognized"

    target = None if target_letter in (None, "", DETECT_ONLY) else target_letter
    is_correct = bool(accepted and target is not None and predicted_label == target)

    if not accepted:
        if pred_id == unknown_id:
            verdict_md = (
                "### 🚫 হরফ নয় · Not a letter\n\n"
                "মডেল এই অডিওটিকে **Unknown** শ্রেণিতে রেখেছে — নীরবতা, শব্দ, "
                "অথবা অন্য কোনো শব্দ হওয়ার সম্ভাবনা আছে।"
            )
        else:
            verdict_md = (
                "### ⚠️ চেনা যায়নি · Unrecognized\n\n"
                f"সর্বোচ্চ সম্ভাবনা **{confidence:.1%}**, কিন্তু গ্রহণযোগ্য থ্রেশহোল্ড "
                f"**{THRESHOLD:.0%}** — তাই ফলাফল গ্রহণ করা হয়নি।"
            )
    elif target is None:
        verdict_md = (
            f"### ✅ শনাক্ত হয়েছে · Detected: {_letter_md(predicted_label)}\n\n"
            f"বিশ্বাসযোগ্যতা **{confidence:.1%}**"
        )
    elif is_correct:
        verdict_md = (
            f"### ✅ সঠিক · Correct — {_letter_md(predicted_label)}\n\n"
            f"বিশ্বাসযোগ্যতা **{confidence:.1%}** (থ্রেশহোল্ড {THRESHOLD:.0%})"
        )
    else:
        verdict_md = (
            f"### ❌ ভুল · Incorrect\n\n"
            f"আপনি {_letter_md(predicted_label)} বলেছেন, লক্ষ্য ছিল "
            f"{_letter_md(target)}।\n\n"
            f"বিশ্বাসযোগ্যতা **{confidence:.1%}**"
        )

    details_md = "\n\n".join(
        [
            f"**🔎 ভবিষ্যৎ · Predicted:** {_letter_md(predicted_label)}  \n"
            f"**🎯 লক্ষ্য · Target:** "
            f"{_letter_md(target) if target else '_কোনোটিই নয় · none_'}  \n"
            f"**📊 বিশ্বাসযোগ্যতা · Confidence:** {confidence:.1%}  \n"
            f"**🚪 গ্রহণযোগ্য · Accepted:** {'হ্যাঁ · yes' if accepted else 'না · no'}",
            _prob_table(probs),
        ]
    )

    payload = {
        "predicted_label": predicted_label,
        "confidence": round(confidence, 4),
        "accepted": accepted,
        "threshold": THRESHOLD,
    }
    if target is not None:
        payload["target_letter"] = target
        payload["is_correct"] = is_correct
    payload["top5"] = {
        id2label[int(i)]: round(float(probs[i]), 4)
        for i in np.argsort(probs)[::-1][:5]
    }
    payload["probs"] = {
        id2label[int(i)]: round(float(probs[i]), 6) for i in range(len(probs))
    }

    return {
        "ok": True,
        "verdict_md": verdict_md,
        "details_md": details_md,
        "payload": payload,
        "prediction_text": f"{predicted_label} ({confidence:.1%} confidence)",
        "verdict_text": (
            "✅ Correct" if is_correct else ("❌ Incorrect" if target else "—")
        ),
    }


# ------------------------------------------------------------------- handlers
def verify(audio, target_letter):
    result = _analyse(audio, target_letter)
    return result["verdict_md"], result["details_md"], result["payload"]


def predict(audio, target_letter):
    """Legacy 2-string endpoint, kept for backwards compatibility."""
    result = _analyse(audio, target_letter)
    if not result["ok"]:
        return result["verdict_md"], ""
    return result["prediction_text"], result["verdict_text"]


# ------------------------------------------------------------------------ ui
AR_FONT = (
    "https://fonts.googleapis.com/css2?family=Amiri:wght@400;700&display=swap"
)


def _letter_grid_html() -> str:
    cards = []
    for name, glyph in LETTERS:
        cards.append(
            f'<div class="hcell"><span class="hg">{glyph}</span>'
            f'<span class="hn">{name}</span></div>'
        )
    return (
        f'<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>'
        f'<link href="{AR_FONT}" rel="stylesheet">'
        "<style>"
        ".hgrid{display:grid;grid-template-columns:repeat(auto-fill,minmax(96px,1fr));"
        "gap:8px;margin:8px 0 4px}"
        ".hcell{border:1px solid var(--border-color-primary,#e2e8f0);border-radius:10px;"
        "padding:8px 4px;text-align:center;background:rgba(127,127,127,.06)}"
        ".hg{font-family:'Amiri','Noto Naskh Arabic',serif;font-size:30px;line-height:1.35;"
        "display:block;direction:rtl}"
        ".hn{font-size:12px;opacity:.78;display:block;margin-top:2px}"
        "</style>"
        f'<div class="hgrid">{"".join(cards)}</div>'
    )


def _how_it_works_html() -> str:
    leak_unknown = METRICS.get("unknown_leak_rate")
    leak_wrong = METRICS.get("wrong_letter_leak_rate")
    accept = METRICS.get("correct_accept_rate")

    def pct(value):
        return f"{value:.2%}" if isinstance(value, (int, float)) else "—"

    return (
        "<div style='line-height:1.65'>"
        "<p><b>২৯টি ক্লাস · 29 classes</b> — ২৮টি আরবি হরফের পাশে একটি "
        "<b>Unknown</b> ক্লাস (নীরবতা, শব্দ, উল্টো করা অডিও, ভিন্ন হরফের মিশ্রণ)। "
        "ফলে অচেনা অডিওকে জোর করে কোনো হরফে ঠেলে দেওয়া হয় না।</p>"
        f"<p><b>থ্রেশহোল্ড · Threshold = {THRESHOLD:.0%}</b> — "
        f"ক্যালিব্রেশনে এই সীমার উপরে থাকলেই ফলাফল গ্রহণ করা হয়। "
        f"এই সীমায়: সঠিক গ্রহণ <b>{pct(accept)}</b>, "
        f"ভুল হরফ চুঁকে যাওয়া <b>{pct(leak_wrong)}</b>, "
        f"Unknown চুঁকে যাওয়া <b>{pct(leak_unknown)}</b>।</p>"
        "<p><b>⚠️ সীমাবদ্ধতা · Known limitation</b> — সোর্স রেকর্ডিংয়ে "
        "স্পিকার আইডি নেই, তাই ট্রেন/টেস্ট ভাগ হয়েছে ক্লিপ অনুযায়ী, স্পিকার "
        "অনুযায়ী নয়। একই স্পিকার ট্রেন ও টেস্টে থাকলে held-out নির্ভুলতা "
        "আসল চেয়ে বেশি দেখাতে পারে। নতুন স্পিকার/মাইক্রোফোনে নিজে যাচাই করুন।</p>"
        "<p><b>🔒 ফ্রি কোটা · Free quota</b> — এই Space-টি ZeroGPU-তে চলে, "
        "GPU-সময় ব্যবহারকারীর দৈনিক কোটা থেকে কাটা হয় (লগইন করা ফ্রি অ্যাকাউন্টে "
        "দৈনিক ~৫ মিনিট ≈ ৩,৫০০টি যাচাই)।</p>"
        f"<p><b>মডেল · Model</b> — "
        f"<a href='https://huggingface.co/{MODEL_ID}'>{MODEL_ID}</a> "
        "(wav2vec2-large-xlsr-53-arabic, Apache-2.0)</p>"
        "</div>"
    )


with gr.Blocks(title="Arabic Letter Pronunciation Verifier") as demo:
    gr.HTML(
        "<div style='text-align:center;padding:6px 0 2px'>"
        "<div style='font-size:30px;font-weight:700'>আরবি হরফ উচ্চারণ যাচাইকারী</div>"
        "<div style='font-size:17px;opacity:.75'>Arabic Letter Pronunciation Verifier</div>"
        "<div style='font-size:13px;opacity:.6;margin-top:4px'>"
        "রেকর্ড করুন বা অডিও আপলোড করুন → লক্ষ্য হরফ বেছে নিন → ফলাফল দেখুন"
        "</div></div>"
    )

    with gr.Row():
        with gr.Column(scale=5):
            audio = gr.Audio(
                sources=["microphone", "upload"],
                type="numpy",
                label="🎙️ আপনার অডিও · Your recording",
            )
            target = gr.Dropdown(
                choices=[
                    ("— শুধু শনাক্ত করুন · Detect only —", DETECT_ONLY),
                ]
                + [(f"{glyph}  {name}", name) for name, glyph in LETTERS],
                value=DETECT_ONLY,
                label="🎯 লক্ষ্য হরফ · Target letter",
            )
            run = gr.Button("যাচাই করুন · Verify", variant="primary", size="lg")
        with gr.Column(scale=6):
            verdict = gr.Markdown(
                "### ⏳ প্রস্তুত · Ready\n\n"
                "অডিও দিয়ে **যাচাই করুন** চাপুন।",
                label="ফলাফল · Verdict",
            )
            with gr.Accordion("বিস্তারিত · Details", open=True):
                details = gr.Markdown(label="বিস্তারিত · Details")
                payload = gr.JSON(label="🔧 API ফলাফল · Raw result")

    run.click(verify, inputs=[audio, target], outputs=[verdict, details, payload])
    audio.change(verify, inputs=[audio, target], outputs=[verdict, details, payload])

    with gr.Accordion("আলফাবেট · The 28 letters", open=False):
        gr.HTML(_letter_grid_html())

    with gr.Accordion("কীভাবে কাজ করে · How it works", open=False):
        gr.HTML(_how_it_works_html())

    # Backwards-compatible endpoint: (audio, target_letter) -> (text, text)
    _legacy_a = gr.Textbox(visible=False)
    _legacy_b = gr.Textbox(visible=False)
    _legacy_a.change(
        predict, inputs=[audio, target], outputs=[_legacy_a, _legacy_b], api_name="predict"
    )


if __name__ == "__main__":
    demo.queue().launch()