Update app.py

app.py

@@ -1,27 +1,27 @@
 import gradio as gr
-from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
 import librosa
 import torch
 import epitran
 import re
+import difflib
 import editdistance
+from jiwer import wer
+import json
 import string
 import eng_to_ipa as ipa
 import numpy as np
 
-# --- Device setup ---
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# --- Load faster Wav2Vec2 models for English & Arabic ---
+# Models: Wav2Vec2 for both Arabic and English
 MODELS = {
     "Arabic": {
-        "processor": Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-xls-r-300m"),
-        "model": Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-xls-r-300m").to(device),
+        "processor": Wav2Vec2Processor.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-arabic"),
+        "model": Wav2Vec2ForCTC.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-arabic"),
         "epitran": epitran.Epitran("ara-Arab")
     },
     "English": {
-        "processor": Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-100k"),
-        "model": Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-100k").to(device),
+        "processor": Wav2Vec2Processor.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-english"),
+        "model": Wav2Vec2ForCTC.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-english"),
         "epitran": epitran.Epitran("eng-Latn")
     }
 }
@@ -29,23 +29,27 @@ MODELS = {
 for lang in MODELS.values():
     lang["model"].config.ctc_loss_reduction = "mean"
 
-# --- Precompute IPA mapping for single letters ---
-LETTER_IPA = {l: ipa.convert(l.lower()).replace(".", "") for l in string.ascii_uppercase}
-
 def clean_phonemes(ipa_text):
     return re.sub(r'[\u064B-\u0652\u02D0]', '', ipa_text)
 
 def safe_transliterate_arabic(epi, word):
     try:
-        ipa_text = epi.transliterate(word.strip())
-        return clean_phonemes(ipa_text)
-    except:
+        word = word.strip()
+        ipa = epi.transliterate(word)
+        if not ipa.strip():
+            raise ValueError("Empty IPA string")
+        return clean_phonemes(ipa)
+    except Exception as e:
+        print(f"[Warning] Arabic transliteration failed for '{word}': {e}")
         return ""
 
 def transliterate_english(word):
     try:
-        return LETTER_IPA.get(word.upper(), "")
-    except:
+        word = word.lower().translate(str.maketrans('', '', string.punctuation))
+        ipa_text = ipa.convert(word)
+        return clean_phonemes(ipa_text)
+    except Exception as e:
+        print(f"[Warning] English IPA conversion failed for '{word}': {e}")
         return ""
 
 def analyze_phonemes(language, reference_text, audio_file):
@@ -56,61 +60,160 @@ def analyze_phonemes(language, reference_text, audio_file):
 
     transliterate_fn = safe_transliterate_arabic if language == "Arabic" else transliterate_english
 
-    # --- Load & normalize audio ---
+    ref_phonemes = [list(transliterate_fn(word)) for word in reference_text.split()]
+
+    # Load audio
     audio, sr = librosa.load(audio_file, sr=16000)
-    if len(audio) < sr * 0.1:
-        return {"language": language, "transcription": "No speech detected", "correct": False}
 
-    audio = audio / max(np.abs(audio), 1e-9)
+    # Normalize volume
+    max_amp = np.max(np.abs(audio))
+    if max_amp > 0:
+        audio = audio / max_amp  # Normalize to [-1, 1]
+
+    # Stricter silence trimming
     trimmed_audio, _ = librosa.effects.trim(audio, top_db=30)
-    trimmed_audio = trimmed_audio[:int(sr*0.75)]  # max 0.75s
+    if len(trimmed_audio) < (sr * 0.15):
+        return json.dumps({
+            "language": language,
+            "reference_text": reference_text,
+            "transcription": "No speech detected",
+            "word_alignment": [],
+            "metrics": {"message": "Audio appears silent or too noisy. Try speaking louder or in a quieter environment."}
+        }, indent=2, ensure_ascii=False)
+
+    # Cap to 0.75s for single letters
+    max_duration = 0.75
+    if len(trimmed_audio) > int(sr * max_duration):
+        trimmed_audio = trimmed_audio[:int(sr * max_duration)]
+
+    # Noise gate
+    noise_gate_threshold = 0.02
+    trimmed_audio[np.abs(trimmed_audio) < noise_gate_threshold] = 0
+
+    input_values = processor(trimmed_audio, sampling_rate=sr, return_tensors="pt").input_values
 
-    # --- Wav2Vec2 inference ---
-    input_values = processor(trimmed_audio, sampling_rate=sr, return_tensors="pt").input_values.to(device)
     with torch.no_grad():
         logits = model(input_values).logits
-    pred_ids = torch.argmax(logits, dim=-1)
-    transcription = processor.batch_decode(pred_ids)[0].strip()
+        pred_ids = torch.argmax(logits, dim=-1)
+        transcription = processor.batch_decode(pred_ids)[0].strip()
 
-    # --- Quick confidence check ---
+    # Stricter confidence check
     probs = torch.softmax(logits, dim=-1)
-    if probs.max(dim=-1).values.mean().item() < 0.6:
-        return {"language": language, "transcription": "Low confidence", "correct": False}
-
-    # --- Single-letter optimization ---
-    if len(reference_text.strip()) == 1:
-        ref_ipa = transliterate_fn(reference_text.strip())
-        trans_ipa = transliterate_fn(transcription)
-        correct = ref_ipa == trans_ipa or reference_text.upper() == transcription.upper()
-        return {"language": language, "reference": reference_text, "transcription": transcription, "correct": correct}
+    max_probs = probs.max(dim=-1).values.mean().item()
+    if max_probs < 0.6:
+        return json.dumps({
+            "language": language,
+            "reference_text": reference_text,
+            "transcription": "No speech detected",
+            "word_alignment": [],
+            "metrics": {"message": "Low confidence transcription (possible noise). Try again with clearer speech."}
+        }, indent=2, ensure_ascii=False)
+
+    # Filter vowel-heavy or overly long transcriptions
+    transcription_clean = transcription.lower().replace("the", "").strip()
+    if len(transcription_clean) > 3 or re.match(r'^[aeiou]+$', transcription_clean):
+        return json.dumps({
+            "language": language,
+            "reference_text": reference_text,
+            "transcription": "No speech detected",
+            "word_alignment": [],
+            "metrics": {"message": "Detected noise or unclear speech. Try again with clear pronunciation."}
+        }, indent=2, ensure_ascii=False)
+
+    obs_phonemes = [list(transliterate_fn(word)) for word in transcription_clean.split()]
+
+    results = {
+        "language": language,
+        "reference_text": reference_text,
+        "transcription": transcription_clean or "No speech detected",
+        "word_alignment": [],
+        "metrics": {}
+    }
 
-    # --- Full phoneme alignment (for multi-letter words) ---
-    ref_phonemes = [list(transliterate_fn(word)) for word in reference_text.split()]
-    obs_phonemes = [list(transliterate_fn(word)) for word in transcription.split()]
-
-    results = []
-    for r, o in zip(ref_phonemes, obs_phonemes):
-        results.append({
-            "reference": ''.join(r),
-            "observed": ''.join(o),
-            "edit_distance": editdistance.eval(r, o)
+    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)
+
+        matcher = difflib.SequenceMatcher(None, ref, obs)
+        ops = matcher.get_opcodes()
+        error_details = []
+        for tag, i1, i2, j1, j2 in ops:
+            ref_seg = ''.join(ref[i1:i2]) or '-'
+            obs_seg = ''.join(obs[j1:j2]) or '-'
+            if tag != 'equal':
+                error_details.append({
+                    "type": tag.upper(),
+                    "reference": ref_seg,
+                    "observed": obs_seg
+                })
+
+        results["word_alignment"].append({
+            "word_index": i,
+            "reference_phonemes": ref_str,
+            "observed_phonemes": obs_str,
+            "edit_distance": edits,
+            "accuracy": acc,
+            "is_correct": edits == 0,
+            "errors": error_details
         })
 
-    return {"language": language, "reference_text": reference_text, "transcription": transcription, "word_alignment": results}
+        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 or "") * 100, 2)
+
+    results["metrics"] = {
+        "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 json.dumps(results, indent=2, ensure_ascii=False)
 
-# --- Gradio UI ---
 def get_default_text(language):
-    return {"Arabic": "فَبِأَيِّ آلَاءِ رَبِّكُمَا تُكَذِّبَانِ", "English": "A"}.get(language, "")
+    return {
+        "Arabic": "فَبِأَيِّ آلَاءِ رَبِّكُمَا تُكَذِّبَانِ",
+        "English": "A"
+    }.get(language, "")
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Fast Multilingual Letter & Word Phoneme Analysis (Wav2Vec2)")
-    language = gr.Dropdown(["Arabic", "English"], value="English", label="Language")
+    gr.Markdown("# Multilingual Phoneme Alignment Analysis")
+    gr.Markdown("Compare audio pronunciation with reference text at phoneme level. Tip: Speak clearly; silence or noise may cause errors.")
+
+    with gr.Row():
+        language = gr.Dropdown(["Arabic", "English"], label="Language", value="English")
+
     reference_text = gr.Textbox(label="Reference Text", value=get_default_text("English"))
-    audio_input = gr.Audio(label="Record Audio", type="filepath")
+    audio_input = gr.Audio(label="Upload Audio File", type="filepath")
     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], outputs=output)
+    output = gr.JSON(label="Phoneme Alignment Results")
+
+    language.change(
+        fn=get_default_text,
+        inputs=language,
+        outputs=reference_text,
+        api_name="/get_default_text"
+    )
+
+    submit_btn.click(
+        fn=analyze_phonemes,
+        inputs=[language, reference_text, audio_input],
+        outputs=output,
+        api_name="/analyze_phonemes"
+    )
 
 demo.launch()