Update app.py

app.py

@@ -7,120 +7,114 @@ import re
 import difflib
 import editdistance
 from jiwer import wer
-import json
+import orjson
 import string
 import eng_to_ipa as ipa
 import numpy as np
+from functools import lru_cache
 
-# Models: Wav2Vec2 for both Arabic and English
+# Check for GPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Lazy-load models
 MODELS = {
     "Arabic": {
-        "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")
+        "processor_path": "jonatasgrosman/wav2vec2-large-xlsr-53-arabic",
+        "model_path": "jonatasgrosman/wav2vec2-large-xlsr-53-arabic",
+        "epitran": lambda: epitran.Epitran("ara-Arab"),
+        "processor": None,
+        "model": None,
+        "epitran_instance": None
     },
     "English": {
-        "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")
+        "processor_path": "jonatasgrosman/wav2vec2-large-xlsr-53-english",
+        "model_path": "jonatasgrosman/wav2vec2-large-xlsr-53-english",
+        "epitran": lambda: epitran.Epitran("eng-Latn"),
+        "processor": None,
+        "model": None,
+        "epitran_instance": None
     }
 }
 
-for lang in MODELS.values():
-    lang["model"].config.ctc_loss_reduction = "mean"
+def load_model(language):
+    if MODELS[language]["processor"] is None:
+        MODELS[language]["processor"] = Wav2Vec2Processor.from_pretrained(MODELS[language]["processor_path"])
+        MODELS[language]["model"] = Wav2Vec2ForCTC.from_pretrained(MODELS[language]["model_path"]).to(device)
+        MODELS[language]["model"].config.ctc_loss_reduction = "mean"
+        MODELS[language]["epitran_instance"] = MODELS[language]["epitran"]()
 
+@lru_cache(maxsize=1000)
 def clean_phonemes(ipa_text):
-    return re.sub(r'[\u064B-\u0652\u02D0]', '', ipa_text)
+    return re.sub(r'[^\w\s]', '', ipa_text)
 
-def safe_transliterate_arabic(epi, word):
+@lru_cache(maxsize=1000)
+def safe_transliterate_arabic(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 = MODELS["Arabic"]["epitran_instance"].transliterate(word)
+        return clean_phonemes(ipa) if ipa.strip() else ""
+    except Exception:
         return ""
 
+@lru_cache(maxsize=1000)
 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 clean_phonemes(ipa.convert(word))
+    except Exception:
         return ""
 
 def analyze_phonemes(language, reference_text, audio_file):
+    load_model(language)
     lang_models = MODELS[language]
     processor = lang_models["processor"]
     model = lang_models["model"]
-    epi = lang_models["epitran"]
 
     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
-    audio, sr = librosa.load(audio_file, sr=16000)
+    ref_phonemes = [list(transliterate_fn(word)) for word in reference_text.split() if transliterate_fn(word)]
 
-    # Normalize volume
+    # Load and preprocess audio
+    audio, _ = librosa.load(audio_file, sr=16000)
     max_amp = np.max(np.abs(audio))
     if max_amp > 0:
-        audio = audio / max_amp  # Normalize to [-1, 1]
+        audio = audio / max_amp
 
-    # Stricter silence trimming
-    trimmed_audio, _ = librosa.effects.trim(audio, top_db=30)
-    if len(trimmed_audio) < (sr * 0.15):
-        return json.dumps({
+    trimmed_audio, _ = librosa.effects.trim(audio, top_db=25)
+    if len(trimmed_audio) < 2400:  # 0.15s at 16kHz
+        return orjson.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)
+            "metrics": {"message": "Audio too short or silent."}
+        }).decode()
 
-    # 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)]
+    # Cap audio length to 0.75s
+    if len(trimmed_audio) > 12000:
+        trimmed_audio = trimmed_audio[:12000]
 
-    # 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
+    input_values = processor(trimmed_audio, sampling_rate=16000, 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()
+        transcription = processor.batch_decode(pred_ids)[0].strip().lower()
 
-    # Stricter confidence check
+    # Combined validation
     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({
+    transcription_clean = transcription.replace("the", "").strip()
+    if max_probs < 0.6 or len(transcription_clean) > 3 or re.match(r'^[aeiou]+$', transcription_clean):
+        return orjson.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)
+            "metrics": {"message": "Unclear or noisy speech."}
+        }).decode()
 
-    obs_phonemes = [list(transliterate_fn(word)) for word in transcription_clean.split()]
+    obs_phonemes = [list(transliterate_fn(word)) for word in transcription_clean.split() if transliterate_fn(word)]
 
     results = {
         "language": language,
@@ -142,17 +136,10 @@ def analyze_phonemes(language, reference_text, audio_file):
         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
-                })
+        error_details = [
+            {"type": tag.upper(), "reference": ''.join(ref[i1:i2]) or '-', "observed": ''.join(obs[j1:j2]) or '-'}
+            for tag, i1, i2, j1, j2 in matcher.get_opcodes() if tag != 'equal'
+        ]
 
         results["word_alignment"].append({
             "word_index": i,
@@ -182,7 +169,7 @@ def analyze_phonemes(language, reference_text, audio_file):
         "asr_word_error_rate": text_wer
     }
 
-    return json.dumps(results, indent=2, ensure_ascii=False)
+    return orjson.dumps(results).decode()
 
 def get_default_text(language):
     return {
@@ -192,28 +179,17 @@ def get_default_text(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.")
+    gr.Markdown("Compare audio pronunciation with reference text at phoneme level.")
 
     with gr.Row():
         language = gr.Dropdown(["Arabic", "English"], label="Language", value="English")
 
-    reference_text = gr.Textbox(label="Reference Text", value=get_default_text("English"))
+    reference_text = gr.Textbox(label="Reference Text", value="A")
     audio_input = gr.Audio(label="Upload Audio File", 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"
-    )
+    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()