Update app.py

app.py

@@ -14,19 +14,28 @@ import soundfile as sf
 # =========================
 # Global config (forced per your request)
 # =========================
-# نثبّت الإعدادات المطلوبة على CPU
 FORCE_WHISPER_NAME = "large-v3"
 FORCE_COMPUTE_TYPE = "int8"
 FORCE_USE_MARBERT = True
 
-# خيارات تفريغ ثابتة لتقليل الفروقات مع النوتبوك
+# ======= Budget Config =======
+# "auto": يعتمد على الحارس العالمي (SBERT/ROUGE/WER)
+# "fixed": عدد ثابت من الاستبدالات (0 يعني عدم استبدال مطلقًا)
+# "ratio": نسبة من طول النص المنطوق
+# "off": بدون سقف (سلوك قديم)
+FORCE_BUDGET_MODE   = "auto"   # "auto" | "fixed" | "ratio" | "off"
+FIXED_BUDGET_TOKENS = 0
+BUDGET_RATIO        = 0.15
+# =============================
+
+# خيارات تفريغ ثابتة لتقليل الفروقات
 ASR_OPTS = dict(
     word_timestamps=True,
     vad_filter=True,
     vad_parameters={"min_silence_duration_ms": 200},
     beam_size=5,
     best_of=5,
-    temperature=0.0,          # جعل فك التشفير حتمي قدر الإمكان
+    temperature=0.0,
 )
 
 # =========================
@@ -57,7 +66,6 @@ def load_models(
         _SBERT = SentenceTransformer(sbert_name, device=("cuda" if DEVICE=="cuda" else "cpu"))
         print(f"[LOAD] SBERT: {sbert_name}", flush=True)
 
-    # مفعّل على CPU حسب رغبتك
     if _MARBERT is None and use_marbert:
         _MARBERT_TOK = AutoTokenizer.from_pretrained(marbert_name)
         _MARBERT = AutoModel.from_pretrained(marbert_name).to(("cuda" if DEVICE=="cuda" else "cpu"))
@@ -73,11 +81,23 @@ def load_models(
 # Normalization / Tokenization / Alignment
 # =========================
 def normalize_ar_orth(text: str) -> str:
+    # تطبيع عام للمحاذاة
     text = re.sub(r"[ًٌٍَُِّْـ]", "", text)
     text = re.sub(r"[“”\"',:؛؟.!()\[\]{}،\-–—_]", " ", text)
+    text = re.sub(r"[إأٱآا]", "ا", text)
+    text = text.replace("ة", "ه").replace("ى", "ي")
     text = re.sub(r"\s+", " ", text).strip()
     return text
 
+def _normalize_for_models(s: str) -> str:
+    # تطبيع خاص لمدخلات SBERT/MARBERT
+    s = re.sub(r"[ًٌٍَُِّْـ]", "", s)
+    s = re.sub(r"[“”\"',:؛؟.!()\[\]{}،\-–—_]", " ", s)
+    s = re.sub(r"[إأٱآا]", "ا", s)
+    s = s.replace("ة", "ه").replace("ى", "ي")
+    s = re.sub(r"\s+", " ", s).strip()
+    return s
+
 def simple_tokenize(text: str):
     t = normalize_ar_orth(text)
     try:
@@ -118,8 +138,7 @@ def arabic_soundex(word):
     for ch in w:
         for rep, chars in groups.items():
             if ch in chars:
-                code.append(rep)
-                break
+                code.append(rep); break
     return "".join(code)
 
 def phonetic_similarity(w1, w2):
@@ -172,28 +191,60 @@ def to_numeric_value(token: str):
     return words_to_number(toks)
 
 # =========================
-# Semantic similarities
+# Semantic similarities (MARBERT fixed)
 # =========================
+def _mean_pool(last_hidden_state, attention_mask):
+    mask = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
+    summed = (last_hidden_state * mask).sum(dim=1)
+    counts = mask.sum(dim=1).clamp(min=1e-9)
+    return summed / counts
+
 def marbert_cls_similarity(a: str, b: str) -> float:
-    if not a or not b: return 0.0
-    if _MARBERT is None:
+    """Return 0 when [UNK] dominates; use mean pooling instead of CLS only."""
+    if not a or not b or _MARBERT is None:
+        return 0.0
+
+    a_n = _normalize_for_models(a)
+    b_n = _normalize_for_models(b)
+
+    # UNK ratio check
+    ids_a = _MARBERT_TOK(a_n, add_special_tokens=False).input_ids
+    ids_b = _MARBERT_TOK(b_n, add_special_tokens=False).input_ids
+    unk_id = _MARBERT_TOK.unk_token_id
+    if len(ids_a) == 0 or len(ids_b) == 0:
+        return 0.0
+    unk_ratio_a = (ids_a.count(unk_id) / len(ids_a)) if unk_id is not None else 0.0
+    unk_ratio_b = (ids_b.count(unk_id) / len(ids_b)) if unk_id is not None else 0.0
+    if max(unk_ratio_a, unk_ratio_b) > 0.5:
+        # too many unknowns → ignore MARBERT
         return 0.0
+
     with torch.no_grad():
-        ta = _MARBERT_TOK(a, return_tensors='pt', truncation=True, padding=True).to(("cuda" if DEVICE=="cuda" else "cpu"))
-        tb = _MARBERT_TOK(b, return_tensors='pt', truncation=True, padding=True).to(("cuda" if DEVICE=="cuda" else "cpu"))
-        ea = _MARBERT(**ta).last_hidden_state[:,0,:]
-        eb = _MARBERT(**tb).last_hidden_state[:,0,:]
-        sim = util.cos_sim(ea, eb).item()
-        return (sim + 1) / 2
+        ta = _MARBERT_TOK(a_n, return_tensors='pt', truncation=True, padding=True).to(("cuda" if DEVICE=="cuda" else "cpu"))
+        tb = _MARBERT_TOK(b_n, return_tensors='pt', truncation=True, padding=True).to(("cuda" if DEVICE=="cuda" else "cpu"))
+        ea = _mean_pool(_MARBERT(**ta).last_hidden_state, ta["attention_mask"])
+        eb = _mean_pool(_MARBERT(**tb).last_hidden_state, tb["attention_mask"])
+        sim = util.cos_sim(ea, eb).item()    # -1..1
+        return (sim + 1) / 2                  # 0..1
 
 def multi_bert_similarity(a: str, b: str):
     if not a or not b:
-        return {"sbert":0.0, "marbert":0.0, "max":0.0, "avg":0.0}
-    sbert_sim = float(util.pytorch_cos_sim(_SBERT.encode(a, convert_to_tensor=True),
-                                           _SBERT.encode(b, convert_to_tensor=True)))
-    marbert_sim = marbert_cls_similarity(a, b)
+        return {"sbert":0.0, "marbert":0.0, "max":0.0, "avg":0.0, "note":"empty"}
+
+    a_n = _normalize_for_models(a); b_n = _normalize_for_models(b)
+    sbert_sim = float(util.pytorch_cos_sim(
+        _SBERT.encode(a_n, convert_to_tensor=True),
+        _SBERT.encode(b_n, convert_to_tensor=True)
+    ))
+    marbert_sim = marbert_cls_similarity(a_n, b_n)
+
+    note = None
+    if abs(sbert_sim - marbert_sim) > 0.35:
+        note = "models_disagree"
+
     vals = [sbert_sim, marbert_sim]
-    return {"sbert": sbert_sim, "marbert": marbert_sim, "max": max(vals), "avg": sum(vals)/len(vals)}
+    return {"sbert": sbert_sim, "marbert": marbert_sim,
+            "max": max(vals), "avg": sum(vals)/len(vals), "note": note}
 
 # =========================
 # Faster-Whisper helpers
@@ -269,38 +320,48 @@ def gate_by_word_conf(base_decision: str, prob: float, sbert_sim: float,
     return base_decision
 
 # =========================
-# Pair + main classifiers
+# Pair + main classifiers (tightened)
 # =========================
 def classify_pair(ref_w, hyp_w, bert_scores, phon_sim, lev1, short_word,
                   bert_thresh=0.75, max_bert=0.85):
+    # numbers equal
     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)'
+
+    # short+lev1
     if short_word and lev1:
         return 'ASR error (short+lev1)'
-    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:
-        return 'ASR error (semantic/phonetic)'
+
+    # semantic/phonetic
+    sbert_ok = bert_scores["sbert"] >= 0.70
+    avg_ok   = bert_scores["avg"]   >= bert_thresh
+    max_ok   = (bert_scores["max"]  >  max_bert) and sbert_ok
+    disagree = (bert_scores.get("note") == "models_disagree")
+
+    if not disagree:
+        if ((phon_sim or lev1) and avg_ok) or max_ok:
+            return 'ASR error (semantic/phonetic)'
+    else:
+        if phon_sim or lev1:
+            if sbert_ok and avg_ok:
+                return 'ASR error (semantic/phonetic)'
+        else:
+            if bert_scores["sbert"] >= 0.78:
+                return 'ASR error (semantic)'
+
     return 'Memorization error'
 
 def classify_alignment_optimized(
     aligned, ref_tokens, hyp_tokens,
     bert_thresh=0.75, max_bert=0.85,
     asr_token_conf=None, low_high=None,
-    replace_budget_tokens=None,   # NEW: سقف الاستبدال (int أو None)
-    guard_note=None               # NEW: وسم حر (مثلاً: "off-topic" أو "ok")
+    replace_budget_tokens=None,   # سقف الاستبدال
+    guard_note=None               # وسم مثل "off-topic"/"ok"/"budget_off"
 ):
-    """
-    مصنّف المحاذاة مع دعم 'سقف الاستبدال'.
-    - إذا replace_budget_tokens=None → لا يوجد سقف.
-    - إذا replace_budget_tokens=0 → لا يتم أي استبدال حتى لو كانت الحالة ASR error.
-    - عند بلوغ السقف نحتفظ بكلمة الطالب ونضيف "[guard: budget reached]" على الحالة.
-    - guard_note (اختياري) يُضاف للـ reason لتوثيق قرار الحارس العالمي.
-    """
-    # --- thresholds من احتمالات الكلمات ---
+    # 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]
@@ -315,7 +376,7 @@ def classify_alignment_optimized(
         low_t, high_t = low_high
 
     results, corrected_words = [], []
-    replaced_count = 0  # NEW: عدّاد الاستبدالات الفعلية
+    replaced_count = 0
 
     for entry in aligned:
         tag = entry['type']
@@ -324,7 +385,7 @@ def classify_alignment_optimized(
 
         if tag == 'equal':
             for ref_w, hyp_w in zip(entry['ref'], entry['hyp']):
-                results.append({'ASR_word': hyp_w, 'GT_word': ref_w, 'status': 'Correct', 'reason': ''})
+                results.append({'ASR_word': hyp_w, 'GT_word': ref_w, 'status': 'Correct', 'reason': '', 'used': hyp_w})
                 corrected_words.append(hyp_w)
 
         elif tag in ['replace', 'delete', 'insert']:
@@ -335,13 +396,13 @@ def classify_alignment_optimized(
                 if not ref_w and not hyp_w:
                     continue
 
-                # --- similarities ---
+                # similarities
                 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
                 bert_scores = multi_bert_similarity(ref_w, hyp_w) if ref_w and hyp_w else {"sbert":0,"marbert":0,"max":0,"avg":0}
                 short_word = bool(ref_w and hyp_w and max(len(ref_w), len(hyp_w)) <= 6)
 
-                # --- base status ---
+                # base status
                 if ref_w and hyp_w:
                     base_status = classify_pair(ref_w, hyp_w, bert_scores, phon_sim, lev1, short_word,
                                                 bert_thresh, max_bert)
@@ -352,7 +413,7 @@ def classify_alignment_optimized(
                 else:
                     base_status = 'Undefined Case'
 
-                # --- word-level confidence gate ---
+                # word-level confidence gate
                 word_prob = None; word_dur = None
                 if (j1 is not None) and (j2 is not None):
                     hyp_abs_idx = j1 + k
@@ -370,30 +431,27 @@ def classify_alignment_optimized(
                         low_t=low_t, high_t=high_t, sbert_lo=0.60
                     )
 
-                # --- choose token to use (with budget) ---
+                # choose token with budget
                 used = hyp_w
                 budget_info = ""
                 if ref_w and hyp_w:
                     if final_status.startswith("ASR error"):
-                        # نتحقق من السقف
                         if (replace_budget_tokens is None) or (replaced_count < replace_budget_tokens):
                             used = ref_w
                             replaced_count += 1
                             if replace_budget_tokens is not None:
                                 budget_info = f", budget={replaced_count}/{replace_budget_tokens}"
                         else:
-                            # تجاوز السقف → لا نستبدل
                             used = hyp_w
                             final_status += " [guard: budget reached]"
                             budget_info = f", budget={replaced_count}/{replace_budget_tokens}"
                     else:
                         used = hyp_w
                 elif hyp_w == '':
-                    used = ''          # حذف
+                    used = ''
                 elif ref_w == '':
-                    used = hyp_w       # إدراج
+                    used = hyp_w
 
-                # --- reason string ---
                 reason = (f'Phonetic={phon_sim}, Lev1={lev1}, '
                           f'SBERT={bert_scores["sbert"]:.2f}, '
                           f'MARBERT={bert_scores["marbert"]:.2f}, '
@@ -403,6 +461,8 @@ def classify_alignment_optimized(
                           f'dur_ms={None if word_dur is None else int(word_dur)}, '
                           f'low_t={round(low_t,2)}, high_t={round(high_t,2)}')
 
+                if bert_scores.get("note"):
+                    reason += f", note={bert_scores['note']}"
                 if guard_note:
                     reason += f", guard='{guard_note}'"
                 if budget_info:
@@ -416,62 +476,26 @@ def classify_alignment_optimized(
                     corrected_words.append(used)
 
     corrected_text = " ".join([w for w in corrected_words if w])
-    return results, corrected_text
-
-# =========================
-# Scores
-# =========================
-def literal_similarity(original, recited):
-    def norm(t):
-        t = re.sub(r'[ًٌٍَُِّْـ]', '', t)
-        t = re.sub(r'[“”",:؛؟.!()\[\]{}،\-–—_]', ' ', t)
-        t = re.sub(r'\s+', ' ', t).strip()
-        return t
-    o = norm(original); r = norm(recited)
-    lev = textdistance.levenshtein.normalized_similarity(o, 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)
-    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, use_marbert=FORCE_USE_MARBERT):
-    sbert_sim = float(util.pytorch_cos_sim(_SBERT.encode(original, convert_to_tensor=True),
-                                           _SBERT.encode(recited,  convert_to_tensor=True)))
-    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 helper
-# =========================
-def ensure_audio_path(audio):
-    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
-        if isinstance(data, np.ndarray):
-            tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
-            sf.write(tmp.name, data, sr)
-            return tmp.name
-    raise ValueError("Unsupported audio input format")
+    # إحصاءات محلية مفيدة للتقرير
+    stats = {
+        "replacements_made": sum(1 for r in results
+                                 if r.get("used") and r.get("GT_word") and r["used"] == r["GT_word"]
+                                 and r.get("ASR_word") and r["ASR_word"] != r["GT_word"]),
+        "budget_reached_count": sum(1 for r in results if isinstance(r.get("status"), str) and "budget reached" in r["status"]),
+        "asr_error_count": sum(1 for r in results if isinstance(r.get("status"), str) and r["status"].startswith("ASR error")),
+        "memorization_error_count": sum(1 for r in results if r.get("status") == "Memorization error"),
+        "missing_count": sum(1 for r in results if r.get("status","").startswith("Missing")),
+        "extra_count": sum(1 for r in results if r.get("status","").startswith("Extra")),
+        "total_tokens": len(results)
+    }
 
+    return results, corrected_text, stats
 
 # =========================
-# 
+# ROUGE-L / WER-like / Guard
 # =========================
-
-
 def lcs_len(a, b):
-    """Longest Common Subsequence length على مستوى التوكنات."""
     m, n = len(a), len(b)
     dp = [[0]*(n+1) for _ in range(m+1)]
     for i in range(1, m+1):
@@ -484,7 +508,6 @@ def lcs_len(a, b):
     return dp[m][n]
 
 def rouge_l_f1_tokens(ref_tokens, hyp_tokens, beta=1.2):
-    """تقريب ROUGE-L F1 على مستوى التوكنات."""
     if not ref_tokens or not hyp_tokens:
         return 0.0, 0.0, 0.0
     lcs = lcs_len(ref_tokens, hyp_tokens)
@@ -496,7 +519,6 @@ def rouge_l_f1_tokens(ref_tokens, hyp_tokens, beta=1.2):
     return float(f1), float(prec), float(rec)
 
 def compute_wer_like(aligned, ref_tokens_len):
-    """WER مبسط من opcodes: (S+D+I)/N."""
     S = D = I = 0
     for op in aligned:
         if op['type'] == 'replace':
@@ -509,36 +531,18 @@ def compute_wer_like(aligned, ref_tokens_len):
     return (S + D + I) / N
 
 def global_offtopic_guard(original_text, asr_text, ref_tokens, hyp_tokens, aligned, sbert_model):
-    """
-    يعيد dict يحوي:
-      off_topic: bool
-      budget_tokens: int (سقف الاستبدالات المسموح)
-      metrics: كل المقاييس للتقرير
-    """
-    # SBERT للنص الكامل
     sbert_sim_text = float(util.pytorch_cos_sim(
-        sbert_model.encode(original_text, convert_to_tensor=True),
-        sbert_model.encode(asr_text,      convert_to_tensor=True)
+        sbert_model.encode(_normalize_for_models(original_text), convert_to_tensor=True),
+        sbert_model.encode(_normalize_for_models(asr_text),      convert_to_tensor=True)
     ))
 
-    # ROUGE-L(F1) و LCS بنسخة توكنات
     rouge_f1, rouge_p, rouge_r = rouge_l_f1_tokens(ref_tokens, hyp_tokens)
-
-    # نسبة التطابق المباشر (equal) من المحاذاة
     equal_tokens = sum(len(op['ref']) for op in aligned if op['type'] == 'equal')
     equal_ratio  = equal_tokens / max(len(ref_tokens), 1)
-
-    # WER مبسّط
     wer = compute_wer_like(aligned, len(ref_tokens))
 
-    # قاعدة قرار Off-topic (حذرين)
-    # نعتبر خارج النص إذا: SBERT<0.70 و ROUGE_F1<0.45 و equal_ratio<0.25 أو WER>0.65
     off_topic = ((sbert_sim_text < 0.70 and rouge_f1 < 0.45 and equal_ratio < 0.25) or (wer > 0.65))
 
-    # ميزانية الاستبدال (عدد الكلمات كحد أقصى يُسمح باستبدالها بـ GT)
-    #  - خارج النص: 0
-    #  - تشابه متوسط: 15% من طول Hyp
-    #  - تشابه مرتفع: 40% من طول Hyp
     L = len(hyp_tokens)
     if off_topic:
         budget = 0
@@ -558,18 +562,64 @@ def global_offtopic_guard(original_text, asr_text, ref_tokens, hyp_tokens, align
     print(f"[GUARD] off_topic={off_topic}, budget={budget}, metrics={metrics}", flush=True)
     return {"off_topic": off_topic, "budget_tokens": budget, "metrics": metrics}
 
+# =========================
+# Scores
+# =========================
+def literal_similarity(original, recited):
+    def norm(t):
+        t = re.sub(r'[ًٌٍَُِّْـ]', '', t)
+        t = re.sub(r'[“”",:؛؟.!()\[\]{}،\-–—_]', ' ', t)
+        t = re.sub(r'\s+', ' ', t).strip()
+        return t
+    o = norm(original); r = norm(recited)
+    lev = textdistance.levenshtein.normalized_similarity(o, 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)
+    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, use_marbert=FORCE_USE_MARBERT):
+    sbert_sim = float(util.pytorch_cos_sim(
+        _SBERT.encode(_normalize_for_models(original), convert_to_tensor=True),
+        _SBERT.encode(_normalize_for_models(recited),  convert_to_tensor=True)
+    ))
+    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)}
 
 # =========================
-# Pipeline (robust errors + logs)
+# Audio helper
 # =========================
+def ensure_audio_path(audio):
+    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
+        if isinstance(data, np.ndarray):
+            tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
+            sf.write(tmp.name, data, sr)
+            return tmp.name
+    raise ValueError("Unsupported audio input format")
 
+# =========================
+# Pipeline (robust errors + logs)
+# =========================
 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.")
 
-        # نُهمل اختيارات الواجهة ونفرض إعداداتك
+        # Forced settings
         whisper_size = FORCE_WHISPER_NAME
         compute_type = FORCE_COMPUTE_TYPE
         use_marbert = FORCE_USE_MARBERT
@@ -585,7 +635,7 @@ def transcribe_and_evaluate(audio, original_text, whisper_size=None,
         segments = list(segments)
         print(f"[ASR] segments={len(segments)}", flush=True)
 
-        # Build ASR text from words (more control)
+        # Build ASR text from words
         words = []
         for seg in segments:
             for w in (seg.words or []):
@@ -594,37 +644,66 @@ def transcribe_and_evaluate(audio, original_text, whisper_size=None,
                     words.append(tok)
         asr_text = " ".join(words)
 
+        # Tokens & alignment
         ref_tokens = simple_tokenize(original_text)
         hyp_tokens = simple_tokenize(asr_text)
         aligned = align_texts(ref_tokens, hyp_tokens)
 
-        # --- Global guard ---
+        # Guard & budget
         guard = global_offtopic_guard(original_text, asr_text, ref_tokens, hyp_tokens, aligned, _SBERT)
         off_topic = guard["off_topic"]
-        budget_tokens = guard["budget_tokens"]
         guard_metrics = guard["metrics"]
 
+        if FORCE_BUDGET_MODE == "off":
+            budget_tokens = None
+            guard_note = "budget_off"
+        elif FORCE_BUDGET_MODE == "fixed":
+            budget_tokens = int(FIXED_BUDGET_TOKENS)
+            guard_note = f"budget_fixed_{budget_tokens}"
+        elif FORCE_BUDGET_MODE == "ratio":
+            budget_tokens = int(BUDGET_RATIO * len(hyp_tokens))
+            guard_note = f"budget_ratio_{BUDGET_RATIO}"
+        else:
+            budget_tokens = guard["budget_tokens"]
+            guard_note = "off-topic" if off_topic else "ok"
+
+        print(f"[BUDGET] mode={FORCE_BUDGET_MODE}, budget={budget_tokens}, note={guard_note}", flush=True)
+
+        # Word-level confidences
         df_words = extract_word_conf_table(segments)
         asr_token_conf, low_t, high_t = build_asr_token_conf(df_words, hyp_tokens)
         print(f"[CONF] low_t={low_t:.3f}, high_t={high_t:.3f}", flush=True)
 
-        results, corrected_text = classify_alignment_optimized(
+        # Classification
+        results, corrected_text, local_stats = 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),
-            replace_budget_tokens=budget_tokens,   # ← عدد استبدالات أقصى
-            guard_note=("off-topic" if off_topic else "ok")
+            replace_budget_tokens=budget_tokens,
+            guard_note=guard_note
         )
 
+        # Scores
         lit = literal_similarity(original_text, corrected_text)
         sem = semantic_similarity(original_text, corrected_text, use_marbert=use_marbert)
 
+        # Extra global metrics for report
+        all_probs = df_words["prob"].dropna().tolist()
+        conf_summary = {
+            "num_words_with_prob": int(len(all_probs)),
+            "avg_prob": None if not all_probs else float(np.mean(all_probs)),
+            "p15": None if not all_probs else float(np.quantile(all_probs, 0.15)),
+            "p70": None if not all_probs else float(np.quantile(all_probs, 0.70)),
+        }
+
         df = pd.DataFrame(results)
 
         report = {
             "requested": {"whisper_model": whisper_size, "compute_type": compute_type, "use_marbert": use_marbert},
             "effective": {"whisper_model": whisper_size, "compute_type": compute_type, "use_marbert": use_marbert},
-            "guard": {"off_topic": off_topic,"budget_tokens": int(budget_tokens),**guard_metrics},
+            "guard": {"mode": FORCE_BUDGET_MODE, "off_topic": off_topic, "budget_tokens": None if budget_tokens is None else int(budget_tokens), **guard_metrics},
+            "local_stats": local_stats,
+            "confidence_summary": conf_summary,
             "original_text": original_text,
             "asr_text": asr_text,
             "corrected_text": corrected_text,
@@ -663,7 +742,6 @@ def build_ui():
             original = gr.Textbox(lines=8, label="Original Text (Ground Truth)")
 
         with gr.Row():
-            # واجهة ثابتة حسب طلبك (تُهمل في الدالة لكن نعرضها)
             whisper_size = gr.Dropdown(choices=["large-v3"], value="large-v3", label="Whisper model size (forced)")
             compute_type = gr.Dropdown(choices=["int8"], value="int8", label="compute_type (forced)")
             vad = gr.Checkbox(value=True, label="VAD filter")