Upload app.py with huggingface_hub

app.py

@@ -1,59 +1,68 @@
+import json
 import os
 
 import gradio as gr
 import torch
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
+import torchaudio
+from huggingface_hub import hf_hub_download
+from transformers import Wav2Vec2FeatureExtractor, Wav2Vec2ForSequenceClassification
 
 MODEL_ID = os.environ.get("HF_ASR_MODEL", "masumtechnonext/wav2vec2-arabic-letter-verifier")
 HF_TOKEN = os.environ.get("HF_TOKEN")
+SAMPLE_RATE = 16000
 
-processor = Wav2Vec2Processor.from_pretrained(MODEL_ID, token=HF_TOKEN)
-model = Wav2Vec2ForCTC.from_pretrained(MODEL_ID, token=HF_TOKEN)
+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()}
+calibration_path = hf_hub_download(MODEL_ID, "calibration.json", token=HF_TOKEN)
+with open(calibration_path) as f:
+    THRESHOLD = json.load(f)["confidence_threshold"]
+UNKNOWN_ID = next(i for i, label in id2label.items() if label == "Unknown")
+LETTERS = sorted(label for label in id2label.values() if label != "Unknown")
 
-def transcribe(audio):
+
+def predict(audio, target_letter):
     if audio is None:
-        return "", "Record or upload audio first."
+        return "Record or upload audio first.", ""
 
     sample_rate, waveform = audio
     waveform = torch.tensor(waveform, dtype=torch.float32)
     if waveform.ndim > 1:
         waveform = waveform.mean(dim=-1)
+    if sample_rate != SAMPLE_RATE:
+        waveform = torchaudio.functional.resample(waveform, sample_rate, SAMPLE_RATE)
 
-    inputs = processor(
-        waveform.numpy(),
-        sampling_rate=sample_rate,
-        return_tensors="pt",
-        padding=True,
-    )
+    inputs = feature_extractor(waveform.numpy(), sampling_rate=SAMPLE_RATE, return_tensors="pt")
 
     with torch.no_grad():
-        logits = model(inputs.input_values).logits
+        logits = model(**inputs).logits[0]
 
-    predicted_ids = torch.argmax(logits, dim=-1)
-    transcription = processor.batch_decode(predicted_ids)[0].strip()
-    return transcription
+    probs = torch.softmax(logits, dim=-1)
+    pred_id = int(torch.argmax(probs))
+    confidence = float(probs[pred_id])
 
+    accepted = confidence >= THRESHOLD and pred_id != UNKNOWN_ID
+    predicted_label = id2label[pred_id] if accepted else "Unrecognized"
+    prediction = f"{predicted_label} ({confidence:.1%} confidence)"
 
-def verify(audio, expected_letter):
-    transcription = transcribe(audio)
-    if not expected_letter:
-        return transcription, "Enter an expected letter to verify."
+    if not target_letter:
+        return prediction, "Pick an expected letter to verify."
 
-    is_match = transcription.strip() == expected_letter.strip()
-    verdict = "✅ Match" if is_match else "❌ No match"
-    return transcription, verdict
+    is_correct = accepted and predicted_label == target_letter
+    verdict = "✅ Correct" if is_correct else "❌ Incorrect"
+    return prediction, verdict
 
 
 demo = gr.Interface(
-    fn=verify,
+    fn=predict,
     inputs=[
         gr.Audio(sources=["microphone", "upload"], type="numpy", label="Speak the letter"),
-        gr.Textbox(label="Expected letter (optional)"),
+        gr.Dropdown(choices=LETTERS, label="Expected letter", value=None),
     ],
     outputs=[
-        gr.Textbox(label="Transcription"),
+        gr.Textbox(label="Prediction"),
         gr.Textbox(label="Verification"),
     ],
     title="Arabic Letter Verifier",