Update app.py

app.py

@@ -1,17 +1,16 @@
 import gradio as gr
 import torch
 import torchaudio
-from transformers import WhisperProcessor, WhisperForConditionalGeneration
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
 import epitran
 import re
 import editdistance
-from jiwer import wer
 import orjson
-import eng_to_ipa as ipa
+from jiwer import wer
 
-# --- Device setup ---
+# --- Device ---
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print("Using device:", device)
+print("Using:", device)
 
 # --- WordMap ---
 WORD_MAP = {
@@ -43,142 +42,83 @@ WORD_MAP = {
     'Z': {'word': 'Zebra', 'phonetic': 'ˈziːbrə'}
 }
 
-# --- Load Whisper tiny ---
-processor = WhisperProcessor.from_pretrained("openai/whisper-tiny.en")
-model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en").to(device).eval()
-epi = epitran.Epitran("eng-Latn")
+# --- Load wav2vec2 (smaller + faster than Whisper) ---
+processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
+model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h").to(device).eval()
 
-# --- Precompute IPA cache ---
+epi = epitran.Epitran("eng-Latn")
 IPA_CACHE = {v['word'].lower(): re.sub(r'[^\w\s]', '', v['phonetic']) for v in WORD_MAP.values()}
 
 # --- Helpers ---
-def clean_phonemes(ipa_text):
-    return re.sub(r'[^\w\s]', '', ipa_text)
-
-def transliterate_english(word):
+def transliterate(word):
     word_lower = word.lower()
     if word_lower in IPA_CACHE:
         return IPA_CACHE[word_lower]
     try:
-        return clean_phonemes(ipa.convert(word)) or ""
+        return epi.transliterate(word_lower)
     except Exception:
         return ""
 
-def find_closest_word(transcription, reference_word):
-    if not transcription:
-        return reference_word, 0.0
-    transcription = transcription.lower().strip()
-    distances = {entry['word'].lower(): editdistance.eval(transcription, entry['word'].lower()) for entry in WORD_MAP.values()}
-    closest_word = min(distances, key=distances.get)
-    max_len = max(len(transcription), len(closest_word))
-    similarity = round((1 - distances[closest_word] / max(1, max_len)) * 100, 2)
-    return closest_word, similarity
-
-def fast_transcribe(audio_path):
+def transcribe(audio_path):
     waveform, sr = torchaudio.load(audio_path)
     if sr != 16000:
         waveform = torchaudio.functional.resample(waveform, sr, 16000)
-
-    input_features = processor(waveform.squeeze(), sampling_rate=16000, return_tensors="pt").input_features.to(device)
+    inputs = processor(waveform.squeeze(), sampling_rate=16000, return_tensors="pt", padding=True).to(device)
 
     with torch.no_grad():
-        predicted_ids = model.generate(input_features)
+        logits = model(**inputs).logits
+    pred_ids = torch.argmax(logits, dim=-1)
+    return processor.decode(pred_ids[0]).lower()
 
-    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
-    return transcription.strip().lower()
-
-# --- Main function ---
-def analyze_phonemes(language, reference_text, audio_input, detailed=True):
+def analyze(language, reference_text, audio_input, detailed=True):
     try:
-        transcription = fast_transcribe(audio_input)
+        transcription = transcribe(audio_input)
+
+        # match closest word from WORD_MAP
+        distances = {entry['word'].lower(): editdistance.eval(transcription, entry['word'].lower()) for entry in WORD_MAP.values()}
+        closest_word = min(distances, key=distances.get)
+        similarity = round((1 - distances[closest_word] / max(1, len(closest_word))) * 100, 2)
 
         if not detailed:
-            return orjson.dumps({
-                "language": language,
-                "reference_text": reference_text,
-                "transcription": transcription
-            }).decode()
+            return {"language": language, "reference": reference_text, "transcription": closest_word}
 
-        # Detailed phoneme alignment
-        closest_word, similarity = find_closest_word(transcription, reference_text.lower())
-        transcription_clean = closest_word
+        # phoneme-level alignment
+        ref_ph = list(transliterate(reference_text))
+        obs_ph = list(transliterate(closest_word))
 
-        obs_phonemes = [list(transliterate_english(word)) for word in transcription_clean.split()]
-        ref_words = reference_text.lower().split()
-        ref_phonemes = [list(transliterate_english(word)) for word in ref_words]
+        edits = editdistance.eval(ref_ph, obs_ph)
+        phon_acc = round((1 - edits / max(1, len(ref_ph))) * 100, 2)
 
-        results = {
+        return {
             "language": language,
-            "reference_text": reference_text,
-            "transcription": transcription_clean,
-            "word_alignment": [],
-            "metrics": {"similarity": similarity}
+            "reference": reference_text,
+            "transcription": closest_word,
+            "metrics": {
+                "similarity": similarity,
+                "phoneme_accuracy": phon_acc,
+                "asr_word_error_rate": round(wer(reference_text, closest_word) * 100, 2)
+            },
+            "alignment": {
+                "reference_phonemes": "".join(ref_ph),
+                "observed_phonemes": "".join(obs_ph),
+                "edit_distance": edits
+            }
         }
-
-        total_phoneme_errors = 0
-        total_phoneme_length = 0
-        correct_words = 0
-        total_word_length = len(ref_phonemes)
-
-        for i, (ref, obs) in enumerate(zip(ref_phonemes, obs_phonemes)):
-            ref_str = ''.join(ref)
-            obs_str = ''.join(obs)
-            edits = editdistance.eval(ref, obs)
-            acc = round((1 - edits / max(1, len(ref))) * 100, 2)
-            results["word_alignment"].append({
-                "word_index": i,
-                "reference_phonemes": ref_str,
-                "observed_phonemes": obs_str,
-                "edit_distance": edits,
-                "accuracy": acc,
-                "is_correct": edits == 0
-            })
-            total_phoneme_errors += edits
-            total_phoneme_length += len(ref)
-            correct_words += int(edits == 0)
-
-        phoneme_acc = round((1 - total_phoneme_errors / max(1, total_phoneme_length)) * 100, 2)
-        phoneme_er = round((total_phoneme_errors / max(1, total_phoneme_length)) * 100, 2)
-        word_acc = round((correct_words / max(1, total_word_length)) * 100, 2)
-        word_er = round(((total_word_length - correct_words) / max(1, total_word_length)) * 100, 2)
-        text_wer = round(wer(reference_text, transcription_clean) * 100, 2)
-
-        results["metrics"].update({
-            "word_accuracy": word_acc,
-            "word_error_rate": word_er,
-            "phoneme_accuracy": phoneme_acc,
-            "phoneme_error_rate": phoneme_er,
-            "asr_word_error_rate": text_wer
-        })
-
-        return orjson.dumps(results).decode()
-
     except Exception as e:
-        return orjson.dumps({
-            "language": language,
-            "reference_text": reference_text,
-            "transcription": "Error processing audio",
-            "word_alignment": [],
-            "metrics": {"message": f"Error: {str(e)}"}
-        }).decode()
+        return {"error": str(e)}
 
 # --- Gradio UI ---
-def get_default_text(language):
-    return "A" if language == "English" else ""
-
 with gr.Blocks() as demo:
-    gr.Markdown("# Multilingual Phoneme Alignment (Fast Whisper Backend)")
-    gr.Markdown("Compare audio pronunciation with reference text at phoneme level. Toggle fast vs detailed mode.")
+    gr.Markdown("## Fast wav2vec2-based Phoneme Checker")
 
     with gr.Row():
-        language = gr.Dropdown(["English"], label="Language", value="English")
-    reference_text = gr.Textbox(label="Reference Text", value="A")
-    audio_input = gr.Audio(label="Record Audio", type="filepath")   # ⚡ filepath, not numpy
-    detailed = gr.Checkbox(label="Detailed Mode (phoneme analysis)", value=True)
-    submit_btn = gr.Button("Analyze")
-    output = gr.JSON(label="Results")
-
-    language.change(fn=get_default_text, inputs=language, outputs=reference_text)
-    submit_btn.click(fn=analyze_phonemes, inputs=[language, reference_text, audio_input, detailed], outputs=output)
+        lang = gr.Dropdown(["English"], value="English", label="Language")
+        ref = gr.Textbox(value="A", label="Reference Word")
+    audio = gr.Audio(label="Record Audio", type="filepath")
+    detailed = gr.Checkbox(value=True, label="Detailed Mode")
+    out = gr.JSON(label="Results")
+
+    demo_btn = gr.Button("Analyze")
+    demo_btn.click(analyze, inputs=[lang, ref, audio, detailed], outputs=out)
 
 demo.launch()