Update app.py

app.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import os
 import re
 import time
+import math
 from functools import lru_cache
 from typing import List, Dict, Any, Tuple, Optional
 
@@ -24,18 +25,26 @@ MODEL_NAME = os.getenv("HADITH_MODEL_NAME", "intfloat/multilingual-e5-base")
 DEFAULT_TOP_K = 10
 MAX_TOP_K = 50
 
-# pull more from FAISS then rerank by evidence
 DEFAULT_RERANK_K = 35
 MAX_RERANK_K = 120
-MIN_RERANK_K = 20
+MIN_RERANK_K = 10
 
-DEFAULT_HL_TOPN = 6        # 0 = disable highlighting (FAST)
+DEFAULT_HL_TOPN = 6        # for /highlight and html responses
 MAX_HL_TOPN = 25
 
-DEFAULT_SEG_MAXLEN = 220   # segment size
+DEFAULT_SEG_MAXLEN = 220
 MAX_SEG_MAXLEN = 420
 MIN_SEG_MAXLEN = 120
 
+# Rerank speed/quality knobs (safe defaults)
+RERANK_MAX_SEGS_PER_DOC = int(os.getenv("RERANK_MAX_SEGS_PER_DOC", "10"))   # keep it small for speed
+RERANK_SEG_MAXLEN = int(os.getenv("RERANK_SEG_MAXLEN", "240"))             # segment length during rerank
+RERANK_WEIGHT = float(os.getenv("RERANK_WEIGHT", "0.65"))                  # 0..1 combine rerank with faiss
+RERANK_ENABLE = os.getenv("RERANK_ENABLE", "1").strip() != "0"
+
+# CORS
+CORS_ALLOW_ORIGIN = os.getenv("CORS_ALLOW_ORIGIN", "*")  # set to your domain if you want strict
+
 
 # =========================
 # Arabic normalization
@@ -47,6 +56,8 @@ _AR_DIACRITICS = re.compile(r"""
   | [\u06D6-\u06ED]
 """, re.VERBOSE)
 
+_AR_PUNCT = re.compile(r"[^\w\u0600-\u06FF]+", re.UNICODE)
+
 def normalize_ar(text: str) -> str:
     if text is None:
         return ""
@@ -60,6 +71,14 @@ def normalize_ar(text: str) -> str:
     text = re.sub(r"\s+", " ", text).strip()
     return text
 
+def ar_tokens(text: str) -> List[str]:
+    t = normalize_ar(text)
+    t = _AR_PUNCT.sub(" ", t)
+    toks = [x.strip() for x in t.split() if x.strip()]
+    # remove super short tokens
+    toks = [x for x in toks if len(x) >= 2]
+    return toks
+
 def escape_html(s: str) -> str:
     if s is None:
         return ""
@@ -74,42 +93,7 @@ def escape_html(s: str) -> str:
 
 
 # =========================
-# Lexical match helpers
-# =========================
-AR_STOPWORDS = {
-    "من","الى","إلى","عن","على","في","و","ثم","أو","او","كما","كان","كانت","يكون","تكون",
-    "هذا","هذه","ذلك","تلك","هناك","هنا","هو","هي","هم","هن","أنا","انت","أنت","نحن",
-    "ما","ماذا","هل","لماذا","لم","لن","لا","إن","أن","إنه","أنه","إلا","الا","حتى","قد",
-    "كل","أي","أيّ","اي","ايًّا","أيضا","أيضًا","مع","عند","بين","بعد","قبل","إذا","اذ","إذ",
-    "قال","وقالت","يقول","يقولون","رسول","الله","صلى","عليه","وسلم"
-}
-
-def ar_tokens(text_norm: str) -> List[str]:
-    if not text_norm:
-        return []
-    t = re.sub(r"[^\u0600-\u06FF0-9\s]", " ", text_norm)
-    t = re.sub(r"\s+", " ", t).strip()
-    toks = [x for x in t.split(" ") if x and x not in AR_STOPWORDS and len(x) >= 2]
-    seen = set()
-    out = []
-    for w in toks:
-        if w not in seen:
-            seen.add(w)
-            out.append(w)
-    return out
-
-def lexical_match_ratio(query_norm: str, doc_norm: str) -> Tuple[float, List[str]]:
-    q_toks = ar_tokens(query_norm)
-    if not q_toks:
-        return 0.0, []
-    doc = " " + (doc_norm or "") + " "
-    matched = [w for w in q_toks if f" {w} " in doc]
-    ratio = len(matched) / max(1, len(q_toks))
-    return float(ratio), matched
-
-
-# =========================
-# Segmenting + isnad detection
+# Segmenting
 # =========================
 def split_ar_segments(text: str, max_len: int) -> List[str]:
     if not text:
@@ -133,44 +117,26 @@ def split_ar_segments(text: str, max_len: int) -> List[str]:
     if buf:
         segs.append(buf)
 
+    # fallback chunking
     if len(segs) <= 1 and len(t) > max_len:
         segs = [t[i:i+max_len].strip() for i in range(0, len(t), max_len) if t[i:i+max_len].strip()]
     return segs
 
-ISNAD_START = re.compile(r"^\s*(حدثنا|أخبرنا|أنبأنا|سمعت|حدثني|أخبرني|قال|عن)\b")
-NAME_HEURISTIC = re.compile(r"(بن|ابن|أبو|أبي|بنت)\s+\S+")
-
-def is_isnad_segment(seg: str) -> bool:
-    if not seg:
-        return False
-    s = seg.strip()
-    if ISNAD_START.search(s):
-        hits = len(NAME_HEURISTIC.findall(s))
-        chain_markers = sum(s.count(x) for x in [" عن ", " قال ", " حدثنا ", " أخبرنا ", " سمعت "])
-        if hits >= 1 or chain_markers >= 2 or len(s) < 120:
-            return True
-    return False
-
-
-# =========================
-# Load model + index + meta (once)
-# =========================
-if not os.path.exists(INDEX_PATH):
-    raise FileNotFoundError(f"FAISS index not found: {INDEX_PATH}")
-if not os.path.exists(META_PATH):
-    raise FileNotFoundError(f"Meta parquet not found: {META_PATH}")
-
-model = SentenceTransformer(MODEL_NAME)
-index = faiss.read_index(INDEX_PATH)
-meta  = pd.read_parquet(META_PATH)
-
-required_cols = {"hadithID", "collection", "hadith_number", "arabic", "english"}
-missing = required_cols - set(meta.columns)
-if missing:
-    raise ValueError(f"Meta is missing required columns: {missing}")
-
-if "arabic_clean" not in meta.columns:
-    meta["arabic_clean"] = ""
+def pick_segs_for_rerank(segs: List[str], max_keep: int) -> List[str]:
+    """Pick up to max_keep segments spread out (for speed)."""
+    if len(segs) <= max_keep:
+        return segs
+    # spread indices evenly
+    idxs = np.linspace(0, len(segs) - 1, num=max_keep)
+    idxs = [int(round(x)) for x in idxs]
+    # unique preserve order
+    seen = set()
+    out = []
+    for i in idxs:
+        if i not in seen:
+            seen.add(i)
+            out.append(segs[i])
+    return out[:max_keep]
 
 
 # =========================
@@ -184,45 +150,103 @@ def cached_query_emb(query_norm: str) -> bytes:
 def get_query_emb(query_norm: str) -> np.ndarray:
     return np.frombuffer(cached_query_emb(query_norm), dtype=np.float32)
 
-def compute_segment_sims(query_emb: np.ndarray, segments: List[str]) -> np.ndarray:
-    if not segments:
-        return np.array([], dtype=np.float32)
-    seg_emb = model.encode(
-        ["passage: " + s for s in segments],
-        normalize_embeddings=True
-    ).astype("float32")
-    return (seg_emb @ query_emb).astype(np.float32)
-
 
 # =========================
-# Core search: candidates -> rerank by best MATN segment
+# Rerank + evidence HTML (no extra encode)
 # =========================
-def search_candidates_df(query_norm: str, rerank_k: int) -> pd.DataFrame:
-    q_emb = get_query_emb(query_norm).reshape(1, -1)
-    scores, idxs = index.search(q_emb, rerank_k)
+def build_heatmap_html(segs: List[str], sims: np.ndarray, top_n: int = 6) -> str:
+    """Small bar-like heatmap using segment similarity (already computed)."""
+    if not segs or sims.size == 0:
+        return ""
 
-    res = meta.iloc[idxs[0]].copy()
-    res["faiss_score"] = scores[0]
-    res["faiss_rank"] = np.arange(len(res), dtype=np.int32)
+    n = len(segs)
+    top_n = max(1, min(top_n, n))
 
-    res["arabic"] = res["arabic"].fillna("").astype(str)
-    res = res[res["arabic"].str.strip() != ""]
-    return res
+    s_min = float(np.min(sims))
+    s_max = float(np.max(sims))
+    denom = (s_max - s_min) if (s_max - s_min) > 1e-6 else 1.0
 
-def rerank_rows(query_norm: str, cand: pd.DataFrame, seg_maxlen: int) -> pd.DataFrame:
-    if cand.empty:
-        out = cand.copy()
-        out["score"] = np.nan
-        out["best_seg"] = ""
-        out["lex_ratio"] = 0.0
-        out["lex_terms"] = ""
-        return out
+    # choose top indices
+    order = np.argsort(-sims)
+    keep = set(order[:top_n])
+
+    blocks = []
+    for i in range(n):
+        w = (float(sims[i]) - s_min) / denom  # 0..1
+        # stronger for top segments
+        alpha = (0.20 + 0.60 * w) if i in keep else (0.08 + 0.18 * w)
+        alpha = max(0.06, min(alpha, 0.85))
+        blocks.append(
+            f'<span title="{escape_html(segs[i])}" '
+            f'style="display:inline-block;width:10px;height:10px;margin:0 3px 0 0;'
+            f'border-radius:4px;background:rgba(37,99,235,{alpha:.3f});border:1px solid rgba(37,99,235,0.20);"></span>'
+        )
 
-    q_emb = get_query_emb(query_norm)  # (d,)
+    return (
+        '<div style="margin:10px 0 0;direction:ltr;text-align:left;">'
+        '<div style="font-size:12px;color:#475569;margin-bottom:6px;">Evidence heatmap</div>'
+        + "".join(blocks) +
+        '</div>'
+    )
+
+def best_seg_html(segs: List[str], sims: np.ndarray) -> str:
+    if not segs or sims.size == 0:
+        return ""
+    i = int(np.argmax(sims))
+    return f'<span style="background:rgba(255,230,120,0.55);border:1px solid rgba(234,179,8,0.35);border-radius:12px;padding:3px 8px;display:inline;">{escape_html(segs[i])}</span>'
+
+def lexical_ratio(query_norm: str, doc_norm: str, max_terms: int = 10) -> Tuple[float, str]:
+    q_toks = ar_tokens(query_norm)
+    d_toks = set(ar_tokens(doc_norm))
+    if not q_toks:
+        return 0.0, ""
+    hit = [t for t in q_toks if t in d_toks]
+    ratio = len(hit) / max(1, len(set(q_toks)))
+    terms = " ".join(hit[:max_terms])
+    return float(ratio), terms
+
+def confidence_label(score: float) -> Tuple[str, str]:
+    """
+    Simple score->label mapping.
+    Assumes cosine-like range ~[0..1] after normalization & blending.
+    """
+    if score >= 0.78:
+        return "HIGH", "bHigh"
+    if score >= 0.62:
+        return "MED", "bMed"
+    return "LOW", "bLow"
+
+def rerank_rows(
+    query_norm: str,
+    df: pd.DataFrame,
+    k_final: int,
+) -> Tuple[pd.DataFrame, Dict[int, Dict[str, Any]]]:
+    """
+    Rerank using segment max similarity:
+    - Split each doc to segments (short)
+    - Pick a limited set of segments (speed)
+    - One encode call for all segments
+    Returns reranked df and per-hadith evidence dict (sims/segs + prebuilt html).
+    """
+    evidence: Dict[int, Dict[str, Any]] = {}
+
+    if (not RERANK_ENABLE) or df.empty:
+        # still fill basic fields
+        for _, row in df.iterrows():
+            hid = int(row["hadithID"]) if pd.notna(row.get("hadithID")) else -1
+            evidence[hid] = {"mode": "disabled"}
+        return df.head(k_final), evidence
+
+    # Collect segments for each candidate
+    cand_rows = df.copy()
+
+    per_doc_segs: List[List[str]] = []
+    doc_hids: List[int] = []
+
+    for _, row in cand_rows.iterrows():
+        hid = int(row["hadithID"]) if pd.notna(row.get("hadithID")) else -1
+        doc_hids.append(hid)
 
-    # Build per-row clean text
-    arabic_clean_list: List[str] = []
-    for _, row in cand.iterrows():
         ar = str(row.get("arabic", "") or "")
         ar_clean = row.get("arabic_clean", "")
         if ar_clean is None or (isinstance(ar_clean, float) and np.isnan(ar_clean)):
@@ -230,136 +254,114 @@ def rerank_rows(query_norm: str, cand: pd.DataFrame, seg_maxlen: int) -> pd.Data
         ar_clean = str(ar_clean).strip()
         if not ar_clean:
             ar_clean = normalize_ar(ar)
-        arabic_clean_list.append(ar_clean)
-
-    # Prepare segments for rerank (batch over all segments)
-    per_segments: List[List[str]] = []
-    per_deemph: List[np.ndarray] = []
-    all_segments: List[str] = []
-    seg_map: List[Tuple[int, int]] = []  # (row_i, seg_i_local)
 
-    for i, txt in enumerate(arabic_clean_list):
-        segs = split_ar_segments(txt, seg_maxlen)
+        segs = split_ar_segments(ar_clean, max_len=RERANK_SEG_MAXLEN)
+        segs = pick_segs_for_rerank(segs, max_keep=RERANK_MAX_SEGS_PER_DOC)
         if not segs:
-            segs = [txt] if txt else []
-        per_segments.append(segs)
-
-        deemph_mask = np.array([1.0 if not is_isnad_segment(s) else 0.0 for s in segs], dtype=np.float32)
-        per_deemph.append(deemph_mask)
-
-        for j, s in enumerate(segs):
-            all_segments.append(s)
-            seg_map.append((i, j))
-
-    if not all_segments:
-        out = cand.copy()
-        out["score"] = out["faiss_score"].astype(float)
-        out["best_seg"] = ""
-        out["lex_ratio"] = 0.0
-        out["lex_terms"] = ""
-        return out
-
-    # sims for all segments once
-    sims_all = compute_segment_sims(q_emb, all_segments)
-
-    # best segment per row (downweight isnad)
-    n_rows = len(per_segments)
-    best_sim = np.full((n_rows,), -1.0, dtype=np.float32)
-    best_local = np.full((n_rows,), -1, dtype=np.int32)
-
-    for k, (ri, sj) in enumerate(seg_map):
-        sim = float(sims_all[k])
-        deemph = float(per_deemph[ri][sj])  # 1 matn, 0 isnad-ish
-        sim_adj = sim * (0.70 + 0.30 * deemph)  # isnad gets downweighted
-        if sim_adj > best_sim[ri]:
-            best_sim[ri] = sim_adj
-            best_local[ri] = sj
-
-    # lexical match
-    lex_ratios: List[float] = []
-    lex_terms: List[str] = []
-    for txt in arabic_clean_list:
-        r, matched = lexical_match_ratio(query_norm, txt)
-        lex_ratios.append(r)
-        lex_terms.append("، ".join(matched[:10]) if matched else "")
-
-    out = cand.copy()
-    out["score"] = best_sim.astype(float)
-    out["best_seg_idx"] = best_local.astype(int)
-
-    # compute best_seg text
-    best_segs = []
-    for i, segs in enumerate(per_segments):
-        j = int(best_local[i])
-        best_segs.append(segs[j] if (0 <= j < len(segs)) else (segs[0] if segs else ""))
-    out["best_seg"] = best_segs
-
-    out["lex_ratio"] = np.array(lex_ratios, dtype=np.float32)
-    out["lex_terms"] = lex_terms
-
-    # Sort by evidence score, then faiss score, then original rank
-    out = out.sort_values(["score", "faiss_score", "faiss_rank"], ascending=[False, False, True])
-    return out
+            segs = [ar_clean[:RERANK_SEG_MAXLEN]] if ar_clean else []
+        per_doc_segs.append(segs)
+
+    # Flatten
+    all_segs: List[str] = []
+    offsets: List[Tuple[int, int]] = []
+    cur = 0
+    for segs in per_doc_segs:
+        start = cur
+        all_segs.extend(segs)
+        cur += len(segs)
+        offsets.append((start, cur))
+
+    if not all_segs:
+        # fallback: no rerank
+        for hid in doc_hids:
+            evidence[hid] = {"mode": "empty"}
+        return cand_rows.head(k_final), evidence
+
+    # Encode query once + all segments once
+    q_emb = get_query_emb(query_norm)  # (d,)
+    seg_emb = model.encode(
+        ["passage: " + s for s in all_segs],
+        normalize_embeddings=True
+    ).astype("float32")  # (N, d)
+
+    sims_all = (seg_emb @ q_emb).astype(np.float32)  # (N,)
+
+    # Compute per-doc rerank score = max(sim)
+    rr_scores: List[float] = []
+    for hid, (start, end), segs in zip(doc_hids, offsets, per_doc_segs):
+        if start == end:
+            rr = -1.0
+            sims = np.array([], dtype=np.float32)
+        else:
+            sims = sims_all[start:end]
+            rr = float(np.max(sims))
+        rr_scores.append(rr)
+
+        # Build evidence HTML now (no extra encode)
+        hm = build_heatmap_html(segs, sims, top_n=min(6, len(segs))) if sims.size else ""
+        best = best_seg_html(segs, sims) if sims.size else ""
+        evidence[hid] = {
+            "mode": "rerank",
+            "rerank_score": rr,
+            "heatmap_html": hm,
+            "best_seg_html": best,
+            "rerank_segs": segs,   # keep for debugging (can omit if you want)
+            "rerank_sims": None,   # don't ship full sims to client
+        }
+
+    cand_rows["rerank_score"] = rr_scores
+
+    # Blend: score_final = (1-w)*faiss + w*rerank
+    # Both are cosine-ish in [0,1] in your setup (normalize embeddings + IP index)
+    faiss_scores = cand_rows["score"].astype(float).to_numpy()
+    rr = cand_rows["rerank_score"].astype(float).to_numpy()
+
+    w = float(max(0.0, min(1.0, RERANK_WEIGHT)))
+    blended = (1.0 - w) * faiss_scores + w * rr
+    cand_rows["final_score"] = blended
+
+    cand_rows = cand_rows.sort_values("final_score", ascending=False).head(k_final)
+    return cand_rows, evidence
 
 
 # =========================
-# UI helpers for html mode
+# Full highlight for ONE hadith (on click)
 # =========================
-def confidence_badge(score1: float, score2: Optional[float]) -> Tuple[str, str]:
-    s1 = float(score1) if score1 is not None else 0.0
-    s2 = float(score2) if score2 is not None else None
-    margin = (s1 - s2) if s2 is not None else s1
-
-    if s1 >= 0.78 or margin >= 0.08:
-        return "High", "bHigh"
-    if s1 >= 0.68 or margin >= 0.04:
-        return "Medium", "bMed"
-    return "Low", "bLow"
-
-def heatmap_html(sims: np.ndarray, bins: int = 16) -> str:
-    if sims.size == 0:
-        return ""
-    s_min = float(np.min(sims))
-    s_max = float(np.max(sims))
-    denom = (s_max - s_min) if (s_max - s_min) > 1e-6 else 1.0
-
-    n = sims.size
-    if n <= bins:
-        take_idx = list(range(n))
-    else:
-        take_idx = [int(round(i)) for i in np.linspace(0, n - 1, bins)]
+def full_highlight_html(
+    query_norm: str,
+    arabic_clean_text: str,
+    hl_topn: int,
+    seg_maxlen: int,
+) -> Dict[str, str]:
+    segs = split_ar_segments(arabic_clean_text, max_len=seg_maxlen)
+    if not segs:
+        return {
+            "arabic_clean_html": escape_html(arabic_clean_text),
+            "heatmap_html": "",
+            "best_seg_html": "",
+        }
 
-    parts = []
-    for i in take_idx:
-        w = (float(sims[i]) - s_min) / denom
-        alpha = 0.10 + 0.75 * w
-        alpha = max(0.08, min(alpha, 0.90))
-        parts.append(
-            f'<span title="seg {i+1}" style="display:inline-block;width:10px;height:10px;'
-            f'margin:0 2px;border-radius:3px;background:rgba(37,99,235,{alpha:.3f});"></span>'
-        )
-    return '<div style="margin:10px 0 8px;direction:ltr;text-align:left;">' + "".join(parts) + "</div>"
+    q_emb = get_query_emb(query_norm)
+    seg_emb = model.encode(
+        ["passage: " + s for s in segs],
+        normalize_embeddings=True
+    ).astype("float32")
 
-def highlight_segments_html(segs: List[str], sims: np.ndarray, strong_topn: int, deemph_mask: np.ndarray) -> str:
-    if not segs or sims.size == 0:
-        return ""
+    sims = (seg_emb @ q_emb).astype(np.float32)
 
     s_min = float(np.min(sims))
     s_max = float(np.max(sims))
     denom = (s_max - s_min) if (s_max - s_min) > 1e-6 else 1.0
 
     order = np.argsort(-sims)
-    keep = set(order[:min(strong_topn, len(segs))])
+    keep = set(order[:max(0, min(hl_topn, len(segs)))])
 
     parts: List[str] = []
     for i, seg in enumerate(segs):
         w = (float(sims[i]) - s_min) / denom
-
-        deemph = float(deemph_mask[i])  # 1 matn, 0 isnad-like
         alpha = (0.18 + 0.62 * w) if i in keep else (0.06 + 0.20 * w)
-        alpha = alpha * (0.45 + 0.55 * deemph)
-        alpha = max(0.04, min(alpha, 0.82))
-        border_alpha = max(0.08, min(alpha * 0.75, 0.60))
+        alpha = max(0.05, min(alpha, 0.82))
+        border_alpha = max(0.10, min(alpha * 0.8, 0.65))
 
         style = (
             f"background: rgba(255, 230, 120, {alpha:.3f});"
@@ -370,285 +372,80 @@ def highlight_segments_html(segs: List[str], sims: np.ndarray, strong_topn: int,
             "display: inline;"
         )
         parts.append(f'<span style="{style}">{escape_html(seg)}</span> ')
-    return "".join(parts).strip()
 
-def build_html_extras_for_row(query_norm: str, arabic_clean_text: str, hl_topn: int, seg_maxlen: int) -> Dict[str, str]:
-    segs = split_ar_segments(arabic_clean_text, seg_maxlen)
-    if not segs:
-        segs = [arabic_clean_text] if arabic_clean_text else []
+    return {
+        "arabic_clean_html": "".join(parts).strip() or escape_html(arabic_clean_text),
+        "heatmap_html": build_heatmap_html(segs, sims, top_n=min(6, len(segs))),
+        "best_seg_html": best_seg_html(segs, sims),
+    }
 
-    deemph_mask = np.array([1.0 if not is_isnad_segment(s) else 0.0 for s in segs], dtype=np.float32)
-    q_emb = get_query_emb(query_norm)
-    sims = compute_segment_sims(q_emb, segs)
 
-    hm = heatmap_html(sims, bins=16) if hl_topn > 0 else ""
-    highlighted = highlight_segments_html(segs, sims, strong_topn=max(1, hl_topn), deemph_mask=deemph_mask) if hl_topn > 0 else escape_html(arabic_clean_text)
+# =========================
+# Load model + index + meta (once)
+# =========================
+if not os.path.exists(INDEX_PATH):
+    raise FileNotFoundError(f"FAISS index not found: {INDEX_PATH}")
+if not os.path.exists(META_PATH):
+    raise FileNotFoundError(f"Meta parquet not found: {META_PATH}")
 
-    # Best seg
-    best_seg = ""
-    if sims.size > 0:
-        best_i = int(np.argmax(sims))
-        best_seg = segs[best_i]
+model = SentenceTransformer(MODEL_NAME)
+index = faiss.read_index(INDEX_PATH)
+meta  = pd.read_parquet(META_PATH)
 
-    best_seg_html = (
-        f'<span style="background:rgba(255,230,120,.55);border:1px solid rgba(234,179,8,.40);'
-        f'border-radius:12px;padding:4px 10px;display:inline;">{escape_html(best_seg)}</span>'
-        if best_seg else ""
-    )
+required_cols = {"hadithID", "collection", "hadith_number", "arabic", "english"}
+missing = required_cols - set(meta.columns)
+if missing:
+    raise ValueError(f"Meta is missing required columns: {missing}")
 
-    return {
-        "heatmap_html": hm,
-        "arabic_clean_html": highlighted if highlighted else escape_html(arabic_clean_text),
-        "best_seg_html": best_seg_html or "",
-    }
+if "arabic_clean" not in meta.columns:
+    meta["arabic_clean"] = ""
 
 
 # =========================
-# Flask app
+# FAISS Search
 # =========================
-app = Flask(__name__)
+def semantic_search_df(query: str, top_k: int) -> pd.DataFrame:
+    q = str(query or "").strip()
+    if not q:
+        return meta.iloc[0:0].copy()
 
-# (UI is optional for HF; keep it for quick testing)
-UI_HTML = r"""
-<!doctype html>
-<html lang="ar" dir="rtl">
-<head>
-  <meta charset="utf-8">
-  <meta name="viewport" content="width=device-width,initial-scale=1">
-  <title>البحث الدلالي في الأحاديث</title>
-  <link rel="preconnect" href="https://fonts.googleapis.com">
-  <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
-  <link href="https://fonts.googleapis.com/css2?family=Amiri:wght@400;700&family=Tajawal:wght@400;700&display=swap" rel="stylesheet">
-  <style>
-    :root{
-      --bg:#f6f7fb; --card:#ffffff; --text:#0f172a; --muted:#475569;
-      --line:#e5e7eb; --accent:#2563eb; --shadow: 0 10px 30px rgba(15, 23, 42, .08);
-      --good:#16a34a; --warn:#f59e0b; --bad:#ef4444;
-    }
-    body{
-      margin:0; background: linear-gradient(180deg, #ffffff, var(--bg)); color: var(--text);
-      font-family: Tajawal, ui-sans-serif, system-ui, -apple-system, Segoe UI, Roboto, Arial;
-    }
-    .wrap{ max-width: 1100px; margin:0 auto; padding: 26px 16px 44px; }
-    .header{
-      background: var(--card); border: 1px solid var(--line); border-radius: 18px;
-      padding: 18px; box-shadow: var(--shadow);
-    }
-    .title{ font-family: Amiri, serif; font-size: 36px; font-weight: 700; margin:0; line-height:1.2; }
-    .sub{ margin: 6px 0 0; font-size: 18px; color: var(--muted); direction:ltr; text-align:left; }
-    .credit{ margin: 6px 0 0; font-size: 14px; color: var(--muted); direction:ltr; text-align:left; }
-
-    form{ display:flex; flex-wrap:wrap; gap:10px; align-items:center; margin-top: 14px; }
-    input[type="text"]{
-      flex: 1 1 620px; background:#fff; border:1px solid var(--line);
-      border-radius:14px; padding:14px; font-size:18px; outline:none;
-    }
-    input[type="number"]{
-      width: 92px; background:#fff; border:1px solid var(--line);
-      border-radius:14px; padding:14px 10px; font-size:16px; direction:ltr; text-align:left; outline:none;
-    }
-    button{
-      background: linear-gradient(180deg, #3b82f6, #2563eb);
-      border: 1px solid #1d4ed8; color:#fff; border-radius:14px;
-      padding:14px 16px; cursor:pointer; font-weight:700; font-size:16px;
-    }
-    button:hover{ filter: brightness(1.05); }
+    top_k = max(1, min(int(top_k), MAX_TOP_K))
+    q_norm = normalize_ar(q)
 
-    .controls{
-      margin-top: 12px; display:flex; gap:10px; flex-wrap:wrap; align-items:center;
-      direction:ltr; text-align:left; color: var(--muted); font-size: 13px;
-    }
-    .controls label{ display:flex; gap:8px; align-items:center; }
-    .controls input[type="range"]{ width: 200px; }
+    q_emb = get_query_emb(q_norm).reshape(1, -1)
+    scores, idx = index.search(q_emb, top_k)
 
-    .meta{ display:flex; gap:10px; flex-wrap:wrap; margin-top: 10px; color: var(--muted); font-size: 13px; direction:ltr; }
-    .pill{ border:1px solid var(--line); background:#fff; padding:6px 10px; border-radius:999px; }
+    res = meta.iloc[idx[0]].copy()
+    res["score"] = scores[0]
+    res = res.sort_values("score", ascending=False)
+
+    # ensure arabic
+    res["arabic"] = res["arabic"].fillna("").astype(str)
+    res = res[res["arabic"].str.strip() != ""]
+    return res
 
-    .grid{ display:grid; grid-template-columns: 1fr; gap:14px; margin-top:14px; }
-    .card{
-      background: var(--card); border: 1px solid var(--line); border-radius:18px;
-      padding: 16px; box-shadow: var(--shadow);
-    }
-    .row{ display:grid; grid-template-columns: 240px 1fr; gap:14px; }
-    @media (max-width: 900px){ .row{ grid-template-columns: 1fr; } }
-
-    .left{ color: var(--muted); font-size:14px; direction:ltr; text-align:left; }
-    .score{ font-weight:900; color: var(--accent); font-size:18px; }
-    .badge{
-      display:inline-flex; align-items:center; gap:6px; border-radius:999px;
-      padding:5px 10px; font-weight:900; font-size:12px; margin-top:10px;
-      border:1px solid var(--line); background:#fff;
-    }
-    .bHigh{ color: var(--good); border-color: rgba(22,163,74,.35); background: rgba(22,163,74,.08); }
-    .bMed{ color: var(--warn); border-color: rgba(245,158,11,.35); background: rgba(245,158,11,.10); }
-    .bLow{ color: var(--bad); border-color: rgba(239,68,68,.35); background: rgba(239,68,68,.08); }
-
-    .arabic{
-      direction: rtl; text-align:right; font-family: Amiri, serif; font-size:22px;
-      line-height: 2.05; background:#fbfcff; border:1px solid var(--line);
-      border-radius:16px; padding:14px; white-space: pre-wrap;
-    }
-    .evidence{
-      margin-top: 10px; border: 1px dashed rgba(37,99,235,.25);
-      background: rgba(37,99,235,.05); border-radius: 14px;
-      padding: 10px 12px; direction: rtl; text-align: right;
-      font-family: Amiri, serif; font-size: 18px; line-height: 1.95;
-    }
-    .evidence small{
-      display:block; margin-bottom:6px; font-family: Tajawal, sans-serif;
-      color: var(--muted); direction:ltr; text-align:left; font-size:12px;
-    }
-    .english{
-      direction:ltr; text-align:left; font-size:16px; line-height:1.8; color:#111827;
-      background:#fbfcff; border:1px solid var(--line); border-radius:16px; padding:14px; white-space: pre-wrap;
-    }
-    details summary{
-      cursor:pointer; color: var(--accent); margin-top:12px; user-select:none;
-      direction:ltr; text-align:left; font-weight:800;
-    }
-    .empty{ margin-top: 14px; color: var(--muted); font-size: 15px; direction:ltr; text-align:left; }
-    .tiny{ margin-top:8px; font-size:12px; color: var(--muted); direction:ltr; text-align:left; }
-  </style>
-</head>
-<body>
-  <div class="wrap">
-    <div class="header">
-      <h1 class="title">البحث الدلالي في الأحاديث</h1>
-      <div class="sub">search in hadeeth (API + UI)</div>
-      <div class="credit">Done by Dr Faisal Alshargi</div>
-
-      <form id="f">
-        <input id="q" type="text" placeholder="اكتب سؤالك هنا… مثال: الاستغفار بعد الذنب وركعتين">
-        <input id="k" type="number" min="1" max="50" value="10">
-        <button type="submit">Search</button>
-      </form>
-
-      <div class="controls">
-        <label>Highlight Top Segments:
-          <input id="hl" type="range" min="0" max="25" value="6"><b id="hlv">6</b>
-        </label>
-        <label>Segment Size:
-          <input id="seg" type="range" min="120" max="420" step="20" value="220"><b id="segv">220</b>
-        </label>
-        <label>Re-rank pool:
-          <input id="rk" type="range" min="20" max="120" step="5" value="35"><b id="rkv">35</b>
-        </label>
-      </div>
-
-      <div id="meta" class="meta" style="display:none;"></div>
-      <div id="msg" class="empty" style="display:none;"></div>
-    </div>
-
-    <div id="grid" class="grid"></div>
-  </div>
-
-<script>
-  const $ = (id)=>document.getElementById(id);
-  function esc(s){
-    return String(s??"")
-      .replaceAll("&","&amp;").replaceAll("<","&lt;").replaceAll(">","&gt;")
-      .replaceAll('"',"&quot;").replaceAll("'","&#39;");
-  }
-  function pill(k,v){ return `<div class="pill">${esc(k)}: <b>${esc(v)}</b></div>`; }
-
-  function sync(rangeId, labelId){
-    const r=$(rangeId), l=$(labelId);
-    l.textContent = r.value;
-    r.addEventListener("input", ()=> l.textContent = r.value);
-  }
-  sync("hl","hlv"); sync("seg","segv"); sync("rk","rkv");
-
-  $("f").addEventListener("submit", async (e)=>{
-    e.preventDefault();
-    const q = $("q").value.trim();
-    const k = parseInt($("k").value||"10",10);
-    const hl = parseInt($("hl").value||"6",10);
-    const seg = parseInt($("seg").value||"220",10);
-    const rk  = parseInt($("rk").value||"35",10);
-
-    $("msg").style.display="none";
-    $("grid").innerHTML = "";
-    $("meta").style.display="none";
-
-    if(!q){
-      $("msg").textContent="اكتب نص البحث أولًا.";
-      $("msg").style.display="block";
-      return;
-    }
 
-    $("msg").textContent="... جاري البحث";
-    $("msg").style.display="block";
+# =========================
+# Flask app
+# =========================
+app = Flask(__name__)
 
-    const url = `/search?q=${encodeURIComponent(q)}&k=${encodeURIComponent(k)}&rerank_k=${encodeURIComponent(rk)}&hl_topn=${encodeURIComponent(hl)}&seg_maxlen=${encodeURIComponent(seg)}&format=html`;
-    const res = await fetch(url);
-    const js = await res.json();
+def add_cors(resp):
+    resp.headers["Access-Control-Allow-Origin"] = CORS_ALLOW_ORIGIN
+    resp.headers["Access-Control-Allow-Methods"] = "GET, OPTIONS"
+    resp.headers["Access-Control-Allow-Headers"] = "Content-Type, Authorization"
+    resp.headers["Access-Control-Max-Age"] = "86400"
+    return resp
 
-    $("meta").style.display="flex";
-    $("meta").innerHTML =
-      pill("Rows", js.rows) + pill("Results", js.n) + pill("Time(ms)", js.took_ms) +
-      pill("TopK", js.k) + pill("ReRank", js.rerank_k) + pill("Query", js.query);
+@app.after_request
+def _after(resp):
+    return add_cors(resp)
 
-    if(!js.ok || !js.results || js.results.length===0){
-      $("msg").textContent="لا توجد نتائج. جرّب كلمات مختلفة.";
-      $("msg").style.display="block";
-      return;
-    }
-    $("msg").style.display="none";
-
-    const cards = js.results.map(r=>{
-      const hm = r.heatmap_html || "";
-      const best = r.best_seg_html || "";
-      const ar = r.arabic_clean_html || esc(r.arabic_clean||"");
-      const ar_tashkeel = esc(r.arabic||"");
-      const en = esc(r.english||"");
-
-      return `
-        <div class="card">
-          <div class="row">
-            <div class="left">
-              <div><span class="score">${Number(r.score||0).toFixed(4)}</span> evidence</div>
-              <div class="tiny">FAISS: <b>${Number(r.faiss_score||0).toFixed(4)}</b></div>
-
-              <div class="badge ${esc(r.conf_class||"")}">Confidence: <b>${esc(r.conf_label||"")}</b></div>
-
-              <div class="tiny" style="margin-top:10px;">
-                Lexical match: <b>${Math.round((r.lex_ratio||0)*100)}%</b>
-                ${r.lex_terms ? `<div style="margin-top:6px;">Matched: <b>${esc(r.lex_terms)}</b></div>` : ``}
-              </div>
-
-              <div style="margin-top:12px;">HadithID: <b>${esc(r.hadithID)}</b></div>
-              <div>Collection: <b>${esc(r.collection)}</b></div>
-              <div>No: <b>${esc(r.hadith_number)}</b></div>
-            </div>
-
-            <div>
-              ${hm}
-              <div class="evidence"><small>Top evidence snippet</small>${best}</div>
-              <div class="arabic" style="margin-top:10px;">${ar}</div>
-
-              <details>
-                <summary>Show Arabic with tashkeel</summary>
-                <div style="height:10px;"></div>
-                <div class="arabic">${ar_tashkeel}</div>
-              </details>
-
-              <details>
-                <summary>Show English</summary>
-                <div style="height:10px;"></div>
-                <div class="english">${en}</div>
-              </details>
-            </div>
-          </div>
-        </div>
-      `;
-    }).join("");
-
-    $("grid").innerHTML = cards;
-  });
-</script>
-</body>
-</html>
-"""
+@app.route("/search", methods=["OPTIONS"])
+@app.route("/highlight", methods=["OPTIONS"])
+@app.route("/", methods=["OPTIONS"])
+def options():
+    return add_cors(Response("", status=204))
 
 
 @app.get("/")
@@ -658,24 +455,25 @@ def health():
         "model": MODEL_NAME,
         "index_ntotal": int(getattr(index, "ntotal", -1)),
         "rows": int(len(meta)),
+        "rerank": {
+            "enabled": bool(RERANK_ENABLE),
+            "weight": RERANK_WEIGHT,
+            "max_segs_per_doc": RERANK_MAX_SEGS_PER_DOC,
+            "seg_maxlen": RERANK_SEG_MAXLEN,
+        },
         "endpoints": {
-            "ui": "/ui",
-            "search_json": "/search?q=...&k=10",
-            "search_html": "/search?q=...&k=10&format=html",
+            "search": "/search?q=...&k=10&rerank_k=35&format=json",
+            "search_html": "/search?q=...&k=10&rerank_k=35&format=html",
+            "highlight": "/highlight?q=...&hadithID=123&format=html&hl_topn=6&seg_maxlen=220",
         }
     })
 
 
-@app.get("/ui")
-def ui():
-    return Response(UI_HTML, mimetype="text/html; charset=utf-8")
-
-
 @app.get("/search")
 def search():
     q = request.args.get("q", "").strip()
 
-    # TopK
+    # TopK final
     k_raw = request.args.get("k", str(DEFAULT_TOP_K)).strip()
     try:
         k = int(k_raw) if k_raw else DEFAULT_TOP_K
@@ -690,9 +488,9 @@ def search():
     except Exception:
         rerank_k = DEFAULT_RERANK_K
     rerank_k = max(MIN_RERANK_K, min(rerank_k, MAX_RERANK_K))
-    rerank_k = max(rerank_k, k)  # must be >= k
+    rerank_k = max(rerank_k, k)
 
-    # Highlight controls
+    # Highlight controls (only used for format=html; for fast mode you can still send hl_topn=0)
     hl_raw = request.args.get("hl_topn", str(DEFAULT_HL_TOPN)).strip()
     seg_raw = request.args.get("seg_maxlen", str(DEFAULT_SEG_MAXLEN)).strip()
     try:
@@ -721,28 +519,26 @@ def search():
             "rows": int(len(meta)),
             "took_ms": 0,
             "format": "html" if want_html else "json",
+            "hl_topn": hl_topn,
+            "seg_maxlen": seg_maxlen,
             "results": [],
         })
 
     t0 = time.time()
-    q_norm = normalize_ar(q)
 
-    # 1) candidates from FAISS
-    cand = search_candidates_df(q_norm, rerank_k=rerank_k)
+    # 1) FAISS retrieve pool (rerank_k)
+    df_pool = semantic_search_df(q, top_k=rerank_k)
+    q_norm = normalize_ar(q)
 
-    # 2) rerank by MATN evidence
-    reranked = rerank_rows(q_norm, cand, seg_maxlen=seg_maxlen)
+    # 2) rerank to final k + evidence (no extra encode)
+    df_final, ev = rerank_rows(query_norm=q_norm, df=df_pool, k_final=k)
 
-    # 3) take top k
-    reranked = reranked.head(k).copy()
     took_ms = int((time.time() - t0) * 1000)
 
-    # confidence uses margin between first and second
-    scores_final = reranked["score"].astype(float).tolist()
-    top2 = scores_final[1] if len(scores_final) > 1 else None
-
+    # Build results
     results: List[Dict[str, Any]] = []
-    for pos, (_, row) in enumerate(reranked.iterrows()):
+    for _, row in df_final.iterrows():
+        hid = int(row.get("hadithID")) if pd.notna(row.get("hadithID")) else None
         arabic = str(row.get("arabic", "") or "")
         english = str(row.get("english", "") or "")
 
@@ -753,45 +549,61 @@ def search():
         if not ar_clean:
             ar_clean = normalize_ar(arabic)
 
-        score = float(row.get("score")) if pd.notna(row.get("score")) else 0.0
+        # lexical
+        lex_r, lex_terms = lexical_ratio(q_norm, ar_clean)
 
-        # confidence
-        if pos == 0:
-            conf_label, conf_class = confidence_badge(score, top2)
-        else:
-            # compare against top1 as rough indicator
-            conf_label, conf_class = confidence_badge(score, scores_final[0] if scores_final else None)
+        # scores
+        faiss_score = float(row.get("score")) if pd.notna(row.get("score")) else 0.0
+        rerank_score = float(row.get("rerank_score")) if pd.notna(row.get("rerank_score")) else faiss_score
+        final_score = float(row.get("final_score")) if pd.notna(row.get("final_score")) else faiss_score
+
+        conf_label, conf_class = confidence_label(final_score)
 
-        r: Dict[str, Any] = {
-            "hadithID": int(row.get("hadithID")) if pd.notna(row.get("hadithID")) else None,
+        e = ev.get(hid or -1, {})
+        heatmap_html = e.get("heatmap_html", "") if isinstance(e, dict) else ""
+        best_html = e.get("best_seg_html", "") if isinstance(e, dict) else ""
+
+        r = {
+            "hadithID": hid,
             "collection": str(row.get("collection", "") or ""),
             "hadith_number": int(row.get("hadith_number")) if pd.notna(row.get("hadith_number")) else None,
 
-            "score": score,                # evidence score (reranked)
-            "faiss_score": float(row.get("faiss_score") or 0.0),
-            "faiss_rank": int(row.get("faiss_rank") or 0),
+            # unified score the UI should use
+            "score": final_score,
 
-            "lex_ratio": float(row.get("lex_ratio") or 0.0),
-            "lex_terms": str(row.get("lex_terms", "") or ""),
+            # diagnostics
+            "faiss_score": faiss_score,
+            "rerank_score": rerank_score,
 
             "conf_label": conf_label,
             "conf_class": conf_class,
 
+            "lex_ratio": float(lex_r),
+            "lex_terms": lex_terms,
+
             "arabic": arabic,
             "arabic_clean": ar_clean,
             "english": english,
 
-            "best_seg": str(row.get("best_seg", "") or ""),
+            # Provide evidence html even in json (cheap: already computed in rerank)
+            "heatmap_html": heatmap_html,
+            "best_seg_html": best_html,
         }
 
-        if want_html:
-            extras = build_html_extras_for_row(
+        # If the caller asked for html AND did not disable highlight, also compute full highlight for each result.
+        # This is heavier. Recommended: keep hl_topn=0 for fast mode and use /highlight on click.
+        if want_html and hl_topn > 0:
+            extras = full_highlight_html(
                 query_norm=q_norm,
                 arabic_clean_text=ar_clean,
                 hl_topn=hl_topn,
                 seg_maxlen=seg_maxlen,
             )
-            r.update(extras)
+            r["arabic_clean_html"] = extras["arabic_clean_html"]
+            # You can overwrite with full-doc ones (optional):
+            r["heatmap_html"] = extras["heatmap_html"] or r["heatmap_html"]
+            r["best_seg_html"] = extras["best_seg_html"] or r["best_seg_html"]
+
         results.append(r)
 
     return jsonify({
@@ -810,5 +622,90 @@ def search():
     })
 
 
+@app.get("/highlight")
+def highlight():
+    """
+    Highlight a single hadith on-demand (for fast UI).
+    GET /highlight?q=...&hadithID=123&format=html&hl_topn=6&seg_maxlen=220
+    """
+    q = request.args.get("q", "").strip()
+    hid_raw = request.args.get("hadithID", "").strip()
+
+    hl_raw = request.args.get("hl_topn", str(DEFAULT_HL_TOPN)).strip()
+    seg_raw = request.args.get("seg_maxlen", str(DEFAULT_SEG_MAXLEN)).strip()
+    try:
+        hl_topn = int(hl_raw) if hl_raw else DEFAULT_HL_TOPN
+    except Exception:
+        hl_topn = DEFAULT_HL_TOPN
+    try:
+        seg_maxlen = int(seg_raw) if seg_raw else DEFAULT_SEG_MAXLEN
+    except Exception:
+        seg_maxlen = DEFAULT_SEG_MAXLEN
+
+    hl_topn = max(0, min(hl_topn, MAX_HL_TOPN))
+    seg_maxlen = max(MIN_SEG_MAXLEN, min(seg_maxlen, MAX_SEG_MAXLEN))
+
+    fmt = (request.args.get("format", "html") or "html").lower()
+    want_html = (fmt == "html")
+
+    if not q or not hid_raw:
+        return jsonify({"ok": False, "error": "q and hadithID are required"}), 400
+
+    try:
+        hid = int(hid_raw)
+    except Exception:
+        return jsonify({"ok": False, "error": "hadithID must be int"}), 400
+
+    row_df = meta[meta["hadithID"] == hid]
+    if row_df.empty:
+        return jsonify({"ok": False, "error": "hadithID not found"}), 404
+    row = row_df.iloc[0]
+
+    q_norm = normalize_ar(q)
+
+    arabic = str(row.get("arabic", "") or "")
+    english = str(row.get("english", "") or "")
+
+    ar_clean = row.get("arabic_clean", "")
+    if ar_clean is None or (isinstance(ar_clean, float) and np.isnan(ar_clean)):
+        ar_clean = ""
+    ar_clean = str(ar_clean).strip()
+    if not ar_clean:
+        ar_clean = normalize_ar(arabic)
+
+    # Always produce evidence + highlight here (one doc only)
+    extras = full_highlight_html(
+        query_norm=q_norm,
+        arabic_clean_text=ar_clean,
+        hl_topn=hl_topn if want_html else 0,
+        seg_maxlen=seg_maxlen,
+    )
+
+    # lexical
+    lex_r, lex_terms = lexical_ratio(q_norm, ar_clean)
+
+    return jsonify({
+        "ok": True,
+        "query": q,
+        "query_norm": q_norm,
+        "hadithID": hid,
+        "format": "html" if want_html else "json",
+        "hl_topn": hl_topn,
+        "seg_maxlen": seg_maxlen,
+
+        "lex_ratio": float(lex_r),
+        "lex_terms": lex_terms,
+
+        "arabic": arabic,
+        "arabic_clean": ar_clean,
+        "english": english,
+
+        "arabic_clean_html": extras.get("arabic_clean_html", "") if want_html else "",
+        "heatmap_html": extras.get("heatmap_html", ""),
+        "best_seg_html": extras.get("best_seg_html", ""),
+    })
+
+
 if __name__ == "__main__":
+    # local run only
     app.run(host="127.0.0.1", port=5000, debug=True)