Update app.py

app.py

@@ -1,4 +1,4 @@
-import os, re, math, json, tempfile
+import os, re, json, math, tempfile, traceback
 import numpy as np
 import pandas as pd
 import torch
@@ -9,12 +9,22 @@ from faster_whisper import WhisperModel
 
 from sentence_transformers import SentenceTransformer, util
 from transformers import AutoTokenizer, AutoModel
+import soundfile as sf
 
 # =========================
-# Device & Lazy-loaded models
+# Device & global config
 # =========================
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+CPU_MODE = (DEVICE != "cuda")
+
+# أمان الذاكرة على CPU
+DEFAULT_WHISPER_CPU = "small"
+DEFAULT_COMPUTE_CPU = "int8"
+DEFAULT_USE_MARBERT_CPU = False
 
+# =========================
+# Lazy models
+# =========================
 _SBERT = None
 _MARBERT_TOK = None
 _MARBERT = None
@@ -26,13 +36,23 @@ def load_models(
     whisper_name="small",
     whisper_compute="int8"
 ):
+    """Load models only once."""
     global _SBERT, _MARBERT_TOK, _MARBERT, _WHISPER
+
+    # حماية على CPU: اجبار نماذج أخف
+    if CPU_MODE:
+        whisper_name = DEFAULT_WHISPER_CPU
+        whisper_compute = DEFAULT_COMPUTE_CPU
+
     if _SBERT is None:
         _SBERT = SentenceTransformer(sbert_name, device=DEVICE)
-    if _MARBERT is None:
+
+    # حمّل MARBERT فقط عند الحاجة (قد يستهلك RAM)
+    if _MARBERT is None and (not CPU_MODE):
         _MARBERT_TOK = AutoTokenizer.from_pretrained(marbert_name)
         _MARBERT = AutoModel.from_pretrained(marbert_name).to(DEVICE)
         _MARBERT.eval()
+
     if _WHISPER is None:
         _WHISPER = WhisperModel(whisper_name, device=DEVICE, compute_type=whisper_compute)
 
@@ -48,12 +68,17 @@ def normalize_ar_orth(text: str) -> str:
     return text
 
 def simple_tokenize(text: str):
-    import nltk
+    """يحاول punkt؛ وإن فشل يستخدم تجزئة بسيطة بالمسافات."""
+    t = normalize_ar_orth(text)
     try:
-        nltk.data.find('tokenizers/punkt')
-    except LookupError:
-        nltk.download('punkt')
-    return nltk.word_tokenize(normalize_ar_orth(text))
+        import nltk
+        try:
+            nltk.data.find('tokenizers/punkt')
+        except LookupError:
+            nltk.download('punkt', quiet=True)
+        return nltk.word_tokenize(t)
+    except Exception:
+        return t.split()
 
 def align_texts(ref_tokens, hyp_tokens):
     import difflib
@@ -95,7 +120,7 @@ def is_levenshtein_1(w1, w2):
     return textdistance.levenshtein(w1, w2) == 1
 
 # =========================
-# Numbers (digits & word-numbers)
+# Numbers
 # =========================
 AR_DIGITS = str.maketrans("٠١٢٣٤٥٦٧٨٩", "0123456789")
 UNITS = {"صفر":0,"واحد":1,"واحدة":1,"اثنان":2,"اثنين":2,"اثنتان":2,"اثنتين":2,
@@ -136,10 +161,12 @@ def to_numeric_value(token: str):
     return words_to_number(toks)
 
 # =========================
-# Semantic similarities (SBERT + MARBERT CLS)
+# Semantic similarities
 # =========================
 def marbert_cls_similarity(a: str, b: str) -> float:
     if not a or not b: return 0.0
+    if _MARBERT is None:
+        return 0.0
     with torch.no_grad():
         ta = _MARBERT_TOK(a, return_tensors='pt', truncation=True, padding=True).to(DEVICE)
         tb = _MARBERT_TOK(b, return_tensors='pt', truncation=True, padding=True).to(DEVICE)
@@ -231,20 +258,17 @@ def gate_by_word_conf(base_decision: str, prob: float, sbert_sim: float,
     return base_decision
 
 # =========================
-# Pair classifier + main alignment classifier
+# Pair + main classifiers
 # =========================
 def classify_pair(ref_w, hyp_w, bert_scores, phon_sim, lev1, short_word,
                   bert_thresh=0.75, max_bert=0.85):
-    # 1) numbers
     ref_num = to_numeric_value(ref_w)
     hyp_num = to_numeric_value(hyp_w)
     if (ref_num is not None) or (hyp_num is not None):
         if (ref_num is not None) and (hyp_num is not None) and (ref_num == hyp_num):
             return 'ASR error (numbers equal)'
-    # 2) short + lev1
     if short_word and lev1:
         return 'ASR error (short+lev1)'
-    # 3) semantic
     avg_ok = bert_scores["avg"] >= bert_thresh
     max_ok = bert_scores["max"] >  max_bert
     if ((phon_sim or lev1) and avg_ok) or max_ok:
@@ -254,7 +278,6 @@ def classify_pair(ref_w, hyp_w, bert_scores, phon_sim, lev1, short_word,
 def classify_alignment_optimized(aligned, ref_tokens, hyp_tokens,
                                  bert_thresh=0.75, max_bert=0.85,
                                  asr_token_conf=None, low_high=None):
-    # thresholds
     if low_high is None:
         if asr_token_conf:
             probs = [v["prob"] for v in asr_token_conf.values() if v["prob"] is not None]
@@ -268,8 +291,7 @@ def classify_alignment_optimized(aligned, ref_tokens, hyp_tokens,
     else:
         low_t, high_t = low_high
 
-    results = []
-    corrected_words = []
+    results, corrected_words = [], []
 
     for entry in aligned:
         tag = entry['type']
@@ -285,7 +307,8 @@ def classify_alignment_optimized(aligned, ref_tokens, hyp_tokens,
             for k in range(max_len):
                 ref_w = entry['ref'][k] if k < len(entry['ref']) else ''
                 hyp_w = entry['hyp'][k] if k < len(entry['hyp']) else ''
-                if not ref_w and not hyp_w: continue
+                if not ref_w and not hyp_w: 
+                    continue
 
                 phon_sim = phonetic_similarity(ref_w, hyp_w) if ref_w and hyp_w else False
                 lev1 = is_levenshtein_1(ref_w, hyp_w) if ref_w and hyp_w else False
@@ -302,7 +325,6 @@ def classify_alignment_optimized(aligned, ref_tokens, hyp_tokens,
                 else:
                     base_status = 'Undefined Case'
 
-                # word-level confidence
                 word_prob = None; word_dur = None
                 if (j1 is not None) and (j2 is not None):
                     hyp_abs_idx = j1 + k
@@ -320,14 +342,13 @@ def classify_alignment_optimized(aligned, ref_tokens, hyp_tokens,
                         low_t=low_t, high_t=high_t, sbert_lo=0.60
                     )
 
+                used = hyp_w
                 if ref_w and hyp_w:
                     used = ref_w if final_status.startswith("ASR error") else hyp_w
                 elif hyp_w == '':
                     used = ''
                 elif ref_w == '':
                     used = hyp_w
-                else:
-                    used = hyp_w
 
                 reason = (f'Phonetic={phon_sim}, Lev1={lev1}, '
                           f'SBERT={bert_scores["sbert"]:.2f}, '
@@ -347,15 +368,9 @@ def classify_alignment_optimized(aligned, ref_tokens, hyp_tokens,
     return results, corrected_text
 
 # =========================
-# Literal / Semantic final scores
+# Scores
 # =========================
 def literal_similarity(original, recited):
-    import nltk
-    try:
-        nltk.data.find('tokenizers/punkt')
-    except LookupError:
-        nltk.download('punkt')
-
     def norm(t):
         t = re.sub(r'[ًٌٍَُِّْـ]', '', t)
         t = re.sub(r'[“”",:؛؟.!()\[\]{}،\-–—_]', ' ', t)
@@ -363,41 +378,33 @@ def literal_similarity(original, recited):
         return t
     o = norm(original); r = norm(recited)
     lev = textdistance.levenshtein.normalized_similarity(o, r)
-    ot = nltk.word_tokenize(o); rt = nltk.word_tokenize(r)
+    ot = simple_tokenize(o); rt = simple_tokenize(r)
     common = sum(1 for w1, w2 in zip(ot, rt) if w1 == w2)
     word_overlap = common / max(len(ot), 1)
-    import nltk.translate.bleu_score as bleu
-    bleu1 = bleu.sentence_bleu([ot], rt, weights=(1,0,0,0)) if (ot and rt) else 0.0
+    try:
+        import nltk.translate.bleu_score as bleu
+        bleu1 = bleu.sentence_bleu([ot], rt, weights=(1,0,0,0)) if (ot and rt) else 0.0
+    except Exception:
+        bleu1 = 0.0
     final_score = 0.5*lev + 0.3*word_overlap + 0.2*bleu1
     return {"levenshtein": round(lev,3), "word_overlap": round(word_overlap,3),
             "bleu1": round(bleu1,3), "literal_score": round(final_score,3)}
 
-def semantic_similarity(original, recited):
+def semantic_similarity(original, recited, use_marbert=True):
     sbert_sim = float(util.pytorch_cos_sim(_SBERT.encode(original, convert_to_tensor=True),
                                            _SBERT.encode(recited,  convert_to_tensor=True)))
-    with torch.no_grad():
-        ta = _MARBERT_TOK(original, return_tensors='pt', truncation=True, padding=True).to(DEVICE)
-        tb = _MARBERT_TOK(recited,  return_tensors='pt', truncation=True, padding=True).to(DEVICE)
-        ea = _MARBERT(**ta).last_hidden_state[:,0,:]
-        eb = _MARBERT(**tb).last_hidden_state[:,0,:]
-        sim = util.cos_sim(ea, eb).item()
-        marbert_sim = (sim + 1)/2
+    marbert_sim = marbert_cls_similarity(original, recited) if use_marbert else 0.0
     return {"sbert_sim": round(sbert_sim,3), "marbert_sim": round(marbert_sim,3),
             "semantic_score": round(max(sbert_sim, marbert_sim),3)}
 
 # =========================
-# Audio input helper (filepath or numpy)
+# Audio input helper
 # =========================
-import soundfile as sf
-
 def ensure_audio_path(audio):
-    """
-    Accepts:
-      - str (filepath)
-      - tuple (numpy_array, sample_rate) if Gradio Audio type='numpy'
-    Returns a filepath suitable for faster-whisper.
-    """
+    """Accepts filepath (str) OR (numpy_array, sr). Returns a valid filepath."""
     if isinstance(audio, str):
+        if not os.path.exists(audio):
+            raise FileNotFoundError(f"Audio path not found: {audio}")
         return audio
     if isinstance(audio, tuple) and len(audio) == 2:
         data, sr = audio
@@ -408,79 +415,87 @@ def ensure_audio_path(audio):
     raise ValueError("Unsupported audio input format")
 
 # =========================
-# Transcribe + Evaluate
+# Pipeline (with robust error reporting)
 # =========================
-def transcribe_and_evaluate(audio, original_text, whisper_size="small",
-                            compute_type=("float16" if DEVICE=="cuda" else "int8"),
-                            vad=True, use_marbert=True):
-    # Load models lazily
-    load_models(whisper_name=whisper_size, whisper_compute=compute_type)
-
-    # Transcribe (word_timestamps=True for word-level probs)
-    audio_path = ensure_audio_path(audio)
-    segments, info = _WHISPER.transcribe(
-        audio_path, word_timestamps=True,
-        vad_filter=vad, vad_parameters={"min_silence_duration_ms": 200}
-    )
-    segments = list(segments)
+def transcribe_and_evaluate(audio, original_text, whisper_size=None,
+                            compute_type=None, vad=True, use_marbert=True):
+    try:
+        if not original_text or not original_text.strip():
+            raise ValueError("Original text is empty.")
+
+        # Defaults per device
+        if CPU_MODE:
+            whisper_size = DEFAULT_WHISPER_CPU
+            compute_type = DEFAULT_COMPUTE_CPU
+            use_marbert = DEFAULT_USE_MARBERT_CPU
+        else:
+            whisper_size = whisper_size or "large-v3"
+            compute_type = compute_type or "float16"
 
-    # ASR text from words (cleaned)
-    words = []
-    for seg in segments:
-        for w in (seg.words or []):
-            tok = clean_ar_token(w.word)
-            if tok: words.append(tok)
-    asr_text = " ".join(words)
-
-    # Tokens + align
-    ref_tokens = simple_tokenize(original_text)
-    hyp_tokens = simple_tokenize(asr_text)
-    aligned = align_texts(ref_tokens, hyp_tokens)
-
-    # Word confidence map
-    df_words = extract_word_conf_table(segments)
-    asr_token_conf, low_t, high_t = build_asr_token_conf(df_words, hyp_tokens)
-
-    # Classify + corrected text
-    results, corrected_text = classify_alignment_optimized(
-        aligned, ref_tokens, hyp_tokens,
-        bert_thresh=0.75, max_bert=0.85,
-        asr_token_conf=asr_token_conf, low_high=(low_t, high_t)
-    )
+        load_models(whisper_name=whisper_size, whisper_compute=compute_type)
 
-    # Scores
-    lit = literal_similarity(original_text, corrected_text)
-    if use_marbert:
-        sem = semantic_similarity(original_text, corrected_text)
-    else:
-        sbert_sim = float(util.pytorch_cos_sim(_SBERT.encode(original_text, convert_to_tensor=True),
-                                               _SBERT.encode(corrected_text,  convert_to_tensor=True)))
-        sem = {"sbert_sim": round(sbert_sim,3), "semantic_score": round(sbert_sim,3)}
-
-    df = pd.DataFrame(results)
-
-    report = {
-        "whisper_model": whisper_size,
-        "compute_type": compute_type,
-        "original_text": original_text,
-        "asr_text": asr_text,
-        "corrected_text": corrected_text,
-        "literal": lit,
-        "semantic": sem,
-        "low_t": low_t, "high_t": high_t,
-    }
-    return corrected_text, asr_text, json.dumps(report, ensure_ascii=False, indent=2), df
+        audio_path = ensure_audio_path(audio)
+        segments, info = _WHISPER.transcribe(
+            audio_path, word_timestamps=True,
+            vad_filter=vad, vad_parameters={"min_silence_duration_ms": 200}
+        )
+        segments = list(segments)
+
+        words = []
+        for seg in segments:
+            for w in (seg.words or []):
+                tok = clean_ar_token(w.word)
+                if tok: words.append(tok)
+        asr_text = " ".join(words)
+
+        ref_tokens = simple_tokenize(original_text)
+        hyp_tokens = simple_tokenize(asr_text)
+        aligned = align_texts(ref_tokens, hyp_tokens)
+
+        df_words = extract_word_conf_table(segments)
+        asr_token_conf, low_t, high_t = build_asr_token_conf(df_words, hyp_tokens)
+
+        results, corrected_text = classify_alignment_optimized(
+            aligned, ref_tokens, hyp_tokens,
+            bert_thresh=0.75, max_bert=0.85,
+            asr_token_conf=asr_token_conf, low_high=(low_t, high_t)
+        )
 
-# =========================
-# JSON-only API wrapper (optional)
-# =========================
-def api_predict(audio, original_text, whisper_size="small",
-                compute_type=("float16" if DEVICE=="cuda" else "int8"),
-                vad=True, use_marbert=True):
+        lit = literal_similarity(original_text, corrected_text)
+        sem = semantic_similarity(original_text, corrected_text, use_marbert=(use_marbert and not CPU_MODE))
+
+        df = pd.DataFrame(results)
+
+        report = {
+            "whisper_model": whisper_size,
+            "compute_type": compute_type,
+            "original_text": original_text,
+            "asr_text": asr_text,
+            "corrected_text": corrected_text,
+            "literal": lit,
+            "semantic": sem,
+            "low_t": low_t, "high_t": high_t,
+        }
+        return corrected_text, asr_text, json.dumps(report, ensure_ascii=False, indent=2), df
+
+    except Exception as e:
+        tb = traceback.format_exc()
+        print("ERROR in transcribe_and_evaluate:\n", tb, flush=True)
+        # نرجع JSON بالخطأ بدل ما نفجّر الواجهة
+        empty_df = pd.DataFrame([{"ASR_word":"","GT_word":"","status":"ERROR","reason":str(e),"used":""}])
+        err_json = json.dumps({"error": str(e), "traceback": tb}, ensure_ascii=False, indent=2)
+        gr.Warning(str(e))
+        return "", "", err_json, empty_df
+
+def api_predict(audio, original_text, whisper_size=None, compute_type=None, vad=True, use_marbert=True):
+    # نفس الدالة لكن ترجع JSON فقط
     corrected_text, asr_text, report_json, df = transcribe_and_evaluate(
         audio, original_text, whisper_size, compute_type, vad, use_marbert
     )
-    return json.loads(report_json)
+    try:
+        return json.loads(report_json)
+    except Exception:
+        return {"error": "Failed to parse report_json."}
 
 # =========================
 # Gradio UI
@@ -488,23 +503,25 @@ def api_predict(audio, original_text, whisper_size="small",
 def build_ui():
     with gr.Blocks(title="Samaali ASR Post-Processing", theme=gr.themes.Soft()) as demo:
         gr.Markdown("## Samaali — ASR Post-Processing (Whisper + Alignment + Confidence + Semantics)")
+
         with gr.Row():
+            # filepath أسلم للـ Spaces
             audio = gr.Audio(sources=["microphone","upload"], type="filepath", label="Audio")
             original = gr.Textbox(lines=8, label="Original Text (Ground Truth)")
 
         with gr.Row():
             whisper_size = gr.Dropdown(
                 choices=["tiny","base","small","medium","large-v3"],
-                value=("large-v3" if DEVICE=="cuda" else "small"),
+                value=("large-v3" if not CPU_MODE else DEFAULT_WHISPER_CPU),
                 label="Whisper model size"
             )
             compute_type = gr.Dropdown(
                 choices=["int8", "int8_float16", "float16", "float32"],
-                value=("float16" if DEVICE=="cuda" else "int8"),
+                value=("float16" if not CPU_MODE else DEFAULT_COMPUTE_CPU),
                 label="compute_type"
             )
             vad = gr.Checkbox(value=True, label="VAD filter")
-            use_marbert = gr.Checkbox(value=(DEVICE=="cuda"), label="Use MARBERT (semantic)")
+            use_marbert = gr.Checkbox(value=(not CPU_MODE), label="Use MARBERT (semantic)")
 
         btn = gr.Button("Transcribe & Evaluate", variant="primary")
 
@@ -515,20 +532,18 @@ def build_ui():
         table = gr.Dataframe(headers=["ASR_word","GT_word","status","reason","used"],
                              label="Token-level Decisions", wrap=True)
 
-        # UI action + API endpoint
         btn.click(
             fn=transcribe_and_evaluate,
             inputs=[audio, original, whisper_size, compute_type, vad, use_marbert],
             outputs=[corrected, asr_out, report, table],
-            api_name="evaluate"  # ← Inference API endpoint
+            api_name="evaluate"
         )
 
-        # JSON-only endpoint (hidden button)
         gr.Button(visible=False).click(
             fn=api_predict,
             inputs=[audio, original, whisper_size, compute_type, vad, use_marbert],
             outputs=gr.JSON(),
-            api_name="predict"   # ← Inference API endpoint (JSON only)
+            api_name="predict"
         )
 
     return demo