Update app.py

app.py

@@ -4,24 +4,25 @@ 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
 
-# Models: Wav2Vec2 for both Arabic and English
+# --- Device setup ---
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# --- Load faster/smaller models for English & Arabic ---
 MODELS = {
     "Arabic": {
         "processor": Wav2Vec2Processor.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-arabic"),
-        "model": Wav2Vec2ForCTC.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-arabic"),
+        "model": Wav2Vec2ForCTC.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-arabic").to(device),
         "epitran": epitran.Epitran("ara-Arab")
     },
     "English": {
-        "processor": Wav2Vec2Processor.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-english"),
-        "model": Wav2Vec2ForCTC.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-english"),
+        "processor": Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h"),
+        "model": Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h").to(device),
         "epitran": epitran.Epitran("eng-Latn")
     }
 }
@@ -29,27 +30,23 @@ 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:
-        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}")
+        ipa_text = epi.transliterate(word.strip())
+        return clean_phonemes(ipa_text)
+    except:
         return ""
 
 def transliterate_english(word):
     try:
-        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 LETTER_IPA.get(word.upper(), "")
+    except:
         return ""
 
 def analyze_phonemes(language, reference_text, audio_file):
@@ -60,160 +57,61 @@ def analyze_phonemes(language, reference_text, audio_file):
 
     transliterate_fn = safe_transliterate_arabic if language == "Arabic" else transliterate_english
 
-    ref_phonemes = [list(transliterate_fn(word)) for word in reference_text.split()]
-
-    # Load audio
+    # --- Load & normalize audio ---
     audio, sr = librosa.load(audio_file, sr=16000)
+    if len(audio) < sr * 0.1:
+        return {"language": language, "transcription": "No speech detected", "correct": False}
 
-    # Normalize volume
-    max_amp = np.max(np.abs(audio))
-    if max_amp > 0:
-        audio = audio / max_amp  # Normalize to [-1, 1]
-
-    # Stricter silence trimming
+    audio = audio / max(np.max(np.abs(audio)), 1e-9)
     trimmed_audio, _ = librosa.effects.trim(audio, top_db=30)
-    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
+    trimmed_audio = trimmed_audio[:int(sr*0.75)]  # max 0.75s
 
+    # --- 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()
 
-    # Stricter confidence check
+    # --- Quick confidence check ---
     probs = torch.softmax(logits, dim=-1)
-    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": {}
-    }
+    if probs.max(dim=-1).values.mean().item() < 0.6:
+        return {"language": language, "transcription": "Low confidence", "correct": False}
 
-    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
-        })
+    # --- 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}
 
-        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
-    }
+    # --- 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)
+        })
 
-    return json.dumps(results, indent=2, ensure_ascii=False)
+    return {"language": language, "reference_text": reference_text, "transcription": transcription, "word_alignment": results}
 
+# --- 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("# 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")
-
+    gr.Markdown("# Fast Multilingual Letter & Word Phoneme Analysis")
+    language = gr.Dropdown(["Arabic","English"], value="English", label="Language")
     reference_text = gr.Textbox(label="Reference Text", value=get_default_text("English"))
-    audio_input = gr.Audio(label="Upload Audio File", type="filepath")
+    audio_input = gr.Audio(label="Record Audio", type="filepath")
     submit_btn = gr.Button("Analyze")
-    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()
+    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)
+
+demo.launch()