model.py

import os, re, math, unicodedata, time, json, hashlib, importlib.util
from collections import defaultdict, Counter
from typing import List, Tuple, Dict, Optional
import numpy as np
import torch
from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
from sentence_transformers import SentenceTransformer
import sys, pathlib

HERE = pathlib.Path(__file__).resolve().parent
if str(HERE) not in sys.path:
    sys.path.insert(0, str(HERE))

# ======================= Tunables =======================
BM25_K1   = 1.3
BM25_B    = 0.7
RRF_K     = 35    # RRF constant
CE_MAXLEN = 640
CE_BATCH  = 128

TOP_BM25 = TOP_E5 = TOP_GEMMA = CE_POOL = 190

# Weighted RRF stage-1 fusion (BM25 + E5 + Gemma)
WRRF_BM25_W = 1.0
WRRF_E5_W   = 1.2
WRRF_GEMMA_W= 1.4

# Weighting for the final (reranker) fusion
FINAL_SCORE_BGE_WEIGHT = .07

# Model & cache dirs
USE_CACHE = True

BGE_DIR              = r"models/bge-reranker-hsrc-pairwise-rrf-V1.4".strip()

E5_DIR               = r"models/e5-large-ft_v6".strip()
E5_EVAL_CACHE_DIR    = r"".strip() if USE_CACHE else r""

GEMMA_DIR            = r"models/multilingual-e5-large".strip()
GEMMA_EVAL_CACHE_DIR = r"".strip() if USE_CACHE else r""

# Gemma dtype & max length (SentenceTransformers truncation)
PREFER_BF16_GEMMA = True
GEMMA_MAX_TOK     = 512
# GEMMA_QUERY_TASK  = "search result"   # used only when we must fall back to prompt=...

# ======================= Silence under eval_std ============================
_EVAL_SILENT = os.environ.get("EVAL_STD_MODE","").strip() == "1"
def _log(msg: str):
    if not _EVAL_SILENT:
        print(msg, flush=True)

# ======================= Normalization / Tokenization =======================
# Priority: 1) Relative import, 2) sys.path, 3) Dynamic import, 4) Fallback
try:
    from .text_utils import (  # type: ignore
        tok_he, norm_bm25,
        norm_e5_query, norm_e5_passage,
        norm_gemma_query, norm_gemma_passage,
        norm_bge_query, norm_bge_passage
    )
    _log("[Init] Loaded text_utils (relative import).")
except (ImportError, ModuleNotFoundError):
    try:
        from text_utils import (
            tok_he, norm_bm25,
            norm_e5_query, norm_e5_passage,
            norm_gemma_query, norm_gemma_passage,
            norm_bge_query, norm_bge_passage
        )
        _log("[Init] Loaded text_utils (sys.path import).")
    except (ImportError, ModuleNotFoundError):
        try:
            spec_path = HERE / "text_utils.py"
            if not spec_path.is_file():
                raise FileNotFoundError(f"{spec_path} not found.")
            spec = importlib.util.spec_from_file_location("text_utils", spec_path)
            text_utils_module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(text_utils_module)
            tok_he = text_utils_module.tok_he
            norm_bm25 = text_utils_module.norm_bm25
            norm_e5_query = text_utils_module.norm_e5_query
            norm_e5_passage = text_utils_module.norm_e5_passage
            # Gemma-specific normalizers if present; fallback to e5 norms
            norm_gemma_query = getattr(text_utils_module, "norm_gemma_query", text_utils_module.norm_e5_query)
            norm_gemma_passage = getattr(text_utils_module, "norm_gemma_passage", text_utils_module.norm_e5_passage)
            norm_bge_query = text_utils_module.norm_bge_query
            norm_bge_passage = text_utils_module.norm_bge_passage
            _log("[Init] Loaded text_utils (dynamic import).")
        except Exception:
            _log("[Init] `text_utils` not found. Using generic Hebrew-friendly normalizers for all components.")
            HEB_PREFIXES = ("ו","ה","ב","ל","כ","מ","ש")
            STOPWORDS = {}
            def _generic_norm_he(s:str)->str:
                if not s: return ""
                s=unicodedata.normalize("NFKC",s)
                s=re.sub(r"[\u0591-\u05BD\u05BF-\u05C7]","",s)
                s=(s.replace("״",'"').replace("׳","'")
                     .replace("”",'"').replace("“",'"')
                     .replace("–","-").replace("—","-"))
                return re.sub(r"\s+"," ",s).strip()
            norm_bm25 = norm_e5_query = norm_e5_passage = _generic_norm_he
            norm_gemma_query = norm_gemma_passage = _generic_norm_he
            norm_bge_query = norm_bge_passage = _generic_norm_he
            def tok_he(text: str) -> List[str]:
                s = norm_bm25(text)
                toks = re.findall(r"[A-Za-z0-9\u0590-\u05FF]+", s)
                out=[]
                for t in toks:
                    if len(t)>3 and t[0] in HEB_PREFIXES: out.append(t[1:])
                    out.append(t)
                return [t for t in out if t not in STOPWORDS]

# =========================== BM25 Backends ================================
get_bm25_backend = None
_HAS_BM25_BACKENDS = False
try:
    from .bm25_backends import get_bm25_backend
    _HAS_BM25_BACKENDS = True
    _log("[Init] Loaded bm25_backends (relative import).")
except (ImportError, ModuleNotFoundError):
    try:
        from bm25_backends import get_bm25_backend
        _HAS_BM25_BACKENDS = True
        _log("[Init] Loaded bm25_backends (sys.path import).")
    except (ImportError, ModuleNotFoundError):
        try:
            spec_path = HERE / "bm25_backends.py"
            if not spec_path.is_file():
                 raise FileNotFoundError(f"{spec_path} not found.")
            spec = importlib.util.spec_from_file_location("bm25_backends", spec_path)
            bm25_module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(bm25_module)
            get_bm25_backend = bm25_module.get_bm25_backend
            _HAS_BM25_BACKENDS = True
            _log("[Init] Loaded bm25_backends (dynamic import).")
        except Exception as e:
            _log(f"[Init] Could not load bm25_backends.py ({e}). Will use built-in fallbacks.")
            pass

class _LocalBM25SBackend:
    """Minimal local wrapper for bm25s when bm25_backends.py is missing."""
    def __init__(self, tokenizer, k1: float = 1.3, b: float = 0.7, logger=_log):
        import bm25s
        self._BM25 = bm25s.BM25
        self.tokenizer = tokenizer
        self.k1, self.b = k1, b
        self._bm25 = None
        self.doc_ids: List[str] = []
        self._logger = logger
    @property
    def name(self)->str:
        return f"LocalBM25S(k1={self.k1}, b={self.b})"
    def build(self, ids: List[str], texts: List[str]):
        t0=time.time()
        self.doc_ids = list(ids)
        tokenized = [self.tokenizer(t) for t in texts]
        self._bm25 = self._BM25(k1=self.k1, b=self.b)
        self._bm25.index(tokenized)
        if self._logger: self._logger(f"[{self.name}] Indexed {len(self.doc_ids):,} docs in {time.time()-t0:.2f}s")
    def search(self, query: str, topk: int = 300) -> List[str]:
        terms = self.tokenizer(query)
        if not terms or self._bm25 is None: return []
        k = min(topk, len(self.doc_ids))
        idxs, scores = self._bm25.retrieve([terms], k=k)
        idxs, scores = idxs[0], scores[0]
        mask = np.isfinite(scores) & (scores > 0)
        idxs, scores = idxs[mask], scores[mask]
        if idxs.size == 0: return []
        order = np.lexsort((idxs, -scores))
        idxs = idxs[order]
        return [self.doc_ids[int(i)] for i in idxs]

class _DeterministicBM25Backend:
    """Embedded pure-Python deterministic BM25. Guaranteed fallback."""
    def __init__(self, tokenizer, k1: float = 1.3, b: float = 0.7, logger=_log):
        self.tokenizer=tokenizer; self.k1=k1; self.b=b
        self.doc_ids: List[str]=[]; self.N=0; self.avgdl=0.0
        self.doc_lens=None; self.vocab: Dict[str,int]={}
        self.postings: Dict[int,Tuple[np.ndarray,np.ndarray]]={}
        self.idf=None; self._logger=logger
    @property
    def name(self)->str:
        return f"DeterministicBM25(k1={self.k1}, b={self.b})"
    def build(self, ids: List[str], texts: List[str]):
        self.doc_ids=list(ids); self.N=len(ids)
        lens=np.zeros(self.N,dtype=np.int32)
        tmp=defaultdict(list)
        t0=time.time()
        for i, text in enumerate(texts):
            terms=self.tokenizer(text); lens[i]=len(terms)
            if not terms: continue
            ctr=Counter(terms)
            for t,tf in ctr.items():
                tid=self.vocab.setdefault(t, len(self.vocab))
                tmp[tid].append((i, tf))
        self.doc_lens=lens; self.avgdl=float(np.maximum(1,lens).mean())
        V=len(self.vocab); self.idf=np.zeros(V,dtype=np.float32)
        self.postings={}
        for tid, pairs in tmp.items():
            docs=np.array([d for d,_ in pairs],dtype=np.int32)
            tfs =np.array([tf for _,tf in pairs],dtype=np.float32)
            df=float(len(docs))
            idf=math.log((self.N-df+0.5)/(df+0.5)+1.0)
            self.idf[tid]=idf
            self.postings[tid]=(docs,tfs)
        if self._logger: self._logger(f"[{self.name}] Indexed {self.N:,} docs in {time.time()-t0:.2f}s")
    def search(self, query: str, topk: int = 300) -> List[str]:
        terms=self.tokenizer(query)
        if not terms: return []
        seen: Dict[int,float] = {}
        for t in terms:
            tid=self.vocab.get(t)
            if tid is None: continue
            idf=float(self.idf[tid])
            docs,tfs=self.postings[tid]
            denom=tfs + self.k1*(1-self.b + self.b*(self.doc_lens[docs]/self.avgdl))
            contrib = idf * (tfs*(self.k1+1)) / denom
            for d, c in zip(docs, contrib):
                seen[d]=seen.get(d,0.0)+float(c)
        if not seen: return []
        idx=np.fromiter(seen.keys(),dtype=np.int32)
        scs=np.fromiter(seen.values(),dtype=np.float32)
        k=min(topk,len(scs))
        order = np.lexsort((idx, -scs))
        order = order[:k]
        idx = idx[order]
        return [self.doc_ids[i] for i in idx]

class BM25Index:
    """Unified BM25 wrapper. Returns List[str] of doc IDs."""
    def __init__(self, k1=1.3, b=0.70, logger=_log):
        self.k1, self.b = k1, b
        self.doc_ids: List[str] = []
        self._be = None; self._backend_name = "unset"; self._logger = logger
    def build(self, ids: List[str], texts_norm: List[str]):
        if _HAS_BM25_BACKENDS and callable(get_bm25_backend):
            try:
                self._be = get_bm25_backend(use_bm25s=True, tokenizer=tok_he, k1=self.k1, b=self.b, logger=self._logger)
                self._be.build(ids, texts_norm)
                self.doc_ids = list(self._be.doc_ids)
                self._backend_name = f"{self._be.name} (bm25_backends.py)"
                if self._logger: self._logger(f"[BM25] Using backend: {self._backend_name}")
                return
            except Exception as e:
                if self._logger: self._logger(f"[BM25] bm25_backends failed ({e}). Trying direct bm25s...)")
        try:
            self._be = _LocalBM25SBackend(tok_he, k1=self.k1, b=self.b, logger=self._logger)
            self._be.build(ids, texts_norm)
            self.doc_ids = list(self._be.doc_ids)
            self._backend_name = f"{self._be.name} (direct)"
            if self._logger: self._logger(f"[BM25] Using backend: {self._backend_name}")
            return
        except Exception as e:
            if self._logger: self._logger(f"[BM25] bm25s unavailable ({e}). Falling back to pure-Python).")
        self._be = _DeterministicBM25Backend(tok_he, k1=self.k1, b=self.b, logger=self._logger)
        self._be.build(ids, texts_norm)
        self.doc_ids = list(self._be.doc_ids)
        self._backend_name = f"{self._be.name} (embedded)"
        if self._logger: self._logger(f"[BM25] Using backend: {self._backend_name}")
    def search(self, query: str, topk: int = 200) -> List[str]:
        if self._be is None: return []
        return self._be.search(query, topk=topk)

# ======================= Model Path Resolution =======================
def _resolve_model_path(primary_path: str, fallback_names: List[str]) -> str:
    """
    Resolves a model path: checks primary_path, then HERE/models, HERE, CWD, CWD/models.
    Falls back to first fallback name (HF id/path).
    """
    if primary_path and pathlib.Path(primary_path).is_dir():
        return primary_path
    base_dirs = [HERE / "models", HERE, pathlib.Path.cwd(), pathlib.Path.cwd() / "models"]
    for base in base_dirs:
        for name in fallback_names:
            candidate = base / name
            if candidate.is_dir():
                return str(candidate)
    return fallback_names[0]

def model_name_key(s: str) -> str:
    if not s:
        return ""
    s = s.strip().rstrip("/\\")
    last = re.split(r"[\\/]+", s)[-1] or s
    return last.lower()

# ======================= E5 embedder =============================
class E5Embedder:
    def __init__(self, device=None):
        fallback_names = ["e5-large-ft_v4","multilingual-e5-large"]
        all_fallbacks = [pathlib.Path(E5_DIR).name] + fallback_names if E5_DIR else fallback_names
        self.model_path = _resolve_model_path(E5_DIR, all_fallbacks)
        self.model_name = model_name_key(self.model_path)
        self.device=device or ("cuda" if torch.cuda.is_available() else "cpu")
        _log(f"[E5] Loading encoder from: {self.model_path}  (device={self.device})")
        self.tok=AutoTokenizer.from_pretrained(self.model_path)
        self.mdl=AutoModel.from_pretrained(self.model_path, torch_dtype=torch.bfloat16 if self.device=="cuda" else None).to(self.device) # changed dtype to bf16
        self.mdl.eval()
    @torch.inference_mode()
    def encode(self, texts: List[str], is_query=False, batch=64, progress_desc="E5 encode"):
        # Expects already-normalized texts
        pref="query: " if is_query else "passage: "
        # pref="" if is_query else ""
        out=[]
        n=len(texts)
        if n==0: return np.zeros((0,768), dtype=np.float32)
        total_batches = (n + batch - 1)//batch
        t0=time.time()
        for bi in range(total_batches):
            i = bi*batch
            chunk = texts[i:i+batch]
            enc=self.tok([pref+t.strip() for t in chunk], padding=True, truncation=True, max_length=512, return_tensors="pt").to(self.device)
            hs=self.mdl(**enc).last_hidden_state
            mask=enc["attention_mask"].unsqueeze(-1).expand(hs.size()).float()
            embs=(hs*mask).sum(1)/mask.sum(1).clamp(min=1e-9)
            embs=torch.nn.functional.normalize(embs, p=2, dim=1)
            out.append(embs.detach().cpu().to(dtype=torch.float32))
            if not _EVAL_SILENT:
                if (bi+1)%50==0 or bi==0 or (bi+1)==total_batches:
                    pct = 100.0*(bi+1)/total_batches
                    elapsed = time.time()-t0
                    ips = (i+len(chunk))/max(elapsed,1e-6)
                    print(f"[{progress_desc}] batch {bi+1}/{total_batches}  ({pct:.1f}%)  ~{ips:.0f} items/s")
            del enc, hs, embs
            if torch.cuda.is_available(): torch.cuda.empty_cache()
        return torch.cat(out, dim=0).numpy()

# ======================= EmbeddingGemma embedder =====================
class GemmaEmbedder:
    """
    Uses SentenceTransformer('google/embeddinggemma-300m'), BF16 if available.
    Returns L2-normalized 768-dim numpy arrays.
    No manual prompt prefixing; let SentenceTransformers handle prompting.
    """
    def __init__(self, device=None):
        fallback_names = ["google/embeddinggemma-300m","embeddinggemma-300m"]
        all_fallbacks = [pathlib.Path(GEMMA_DIR).name] + fallback_names if GEMMA_DIR else fallback_names
        self.model_path = _resolve_model_path(GEMMA_DIR, all_fallbacks)
        self.model_name = model_name_key(self.model_path)
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        # dtype selection
        use_bf16 = bool(PREFER_BF16_GEMMA)
        if self.device == "cuda":
            try:
                use_bf16 = use_bf16 and torch.cuda.is_bf16_supported()
            except Exception:
                major, _ = torch.cuda.get_device_capability()
                use_bf16 = use_bf16 and (major >= 8)
        dtype = torch.bfloat16 if use_bf16 else torch.float16
        _log(f"[Gemma] Loading encoder from: {self.model_path}  (device={self.device}, dtype={'bf16' if use_bf16 else 'fp16'})")
        self.mdl = SentenceTransformer(
            self.model_path,
            device=self.device,
            model_kwargs={"torch_dtype": dtype},
        )
        # Tunable max tokens
        try:
            self.mdl.max_seq_length = int(GEMMA_MAX_TOK)
        except Exception:
            pass
        self.dim = 768
        self.mdl.eval()

    @torch.inference_mode()
    def encode(self, texts: List[str], is_query=False, batch=64, progress_desc="Gemma encode", max_length: Optional[int]=None):
        if not texts:
            return np.zeros((0, self.dim), dtype=np.float32)

        # Per-call max length override
        old_len = getattr(self.mdl, "max_seq_length", None)
        if isinstance(max_length, int) and max_length > 0:
            try:
                self.mdl.max_seq_length = max_length
            except Exception:
                pass

        show = not _EVAL_SILENT

        # DO NOT manually add prompts. Prefer encode_query / encode_document when available.
        try:
            if is_query and hasattr(self.mdl, "encode_query"):
                embs = self.mdl.encode_query(
                    texts, batch_size=batch, convert_to_numpy=True,
                    normalize_embeddings=True, show_progress_bar=show
                )
            elif (not is_query) and hasattr(self.mdl, "encode_document"):
                embs = self.mdl.encode_document(
                    texts, batch_size=batch, convert_to_numpy=True,
                    normalize_embeddings=True, show_progress_bar=show
                )
            else:
                # Fallback: use encode with prompt=... if supported (avoids manual concatenation)
                prompt = (f"{'query: ' if is_query else 'passage: '}")
                try:
                    embs = self.mdl.encode(
                        texts, batch_size=batch, convert_to_numpy=True,
                        normalize_embeddings=True, show_progress_bar=show,
                        prompt=prompt
                    )
                except TypeError:
                    # Last resort: plain encode (no prompt)
                    embs = self.mdl.encode(
                        texts, batch_size=batch, convert_to_numpy=True,
                        normalize_embeddings=True, show_progress_bar=show
                    )
        finally:
            if old_len is not None:
                try: self.mdl.max_seq_length = old_len
                except Exception: pass

        embs = np.asarray(embs)
        if embs.ndim == 1:
            embs = embs[None, :]
        return embs.astype(np.float32)

# ======================= BGE reranker ============================
class BGEReranker:
    def __init__(self, device=None):
        fallback_names = ["bge-reranker-hsrc-pairwise-rrf-V1.4","bge-v2-m3","bge-m3"]
        all_fallbacks = [pathlib.Path(BGE_DIR).name] + fallback_names if BGE_DIR else fallback_names
        self.model_path = _resolve_model_path(BGE_DIR, all_fallbacks)
        self.device=device or ("cuda" if torch.cuda.is_available() else "cpu")
        _log(f"[BGE] Loading reranker from: {self.model_path}  (device={self.device})")
        self.tok=AutoTokenizer.from_pretrained(self.model_path)
        self.mdl=AutoModelForSequenceClassification.from_pretrained(
            self.model_path, torch_dtype=torch.float16 if self.device=="cuda" else None, trust_remote_code=True
        ).to(self.device)
        self.mdl.eval()
    @torch.inference_mode()
    def score_pairs(self, q: str, passages: List[str], batch=32, max_len=512) -> List[float]:
        out=[]
        for i in range(0,len(passages), batch):
            enc=self.tok([q]*len(passages[i:i+batch]), passages[i:i+batch],
                         truncation="only_second", max_length=max_len, padding=True, return_tensors="pt").to(self.device)
            logits=self.mdl(**enc).logits
            if logits.ndim==1: s=logits
            elif logits.shape[1]==1: s=logits.squeeze(-1)
            else: s=logits[:,1]
            out += s.detach().float().cpu().tolist()
            del enc, logits
        return [float(x) for x in out]

# ======================== Hybrid Searcher ========================
class HybridSearcher:
    """
    Stage-1 retrieval: WRRF(BM25, E5, Gemma) → candidate ids + WRRF scores.
    Then stage-2 reranking is done outside in predict().
    """
    def __init__(self, bm25: BM25Index,
                 e5: E5Embedder, e5_corpus: np.ndarray,
                 gemma: GemmaEmbedder, gemma_corpus: np.ndarray,
                 id2text: Dict[str,str], id2norm: Dict[str,str]):
        self.bm25=bm25
        self.e5=e5; self.e5_corpus=e5_corpus
        self.gemma=gemma; self.gemma_corpus=gemma_corpus
        self.id2text=id2text; self.id2norm=id2norm
        self._last_q: Optional[str] = None
        self._last_fused: List[Tuple[str, float]] = []

    def _wrrf_fuse3(self, bm_ids: List[str], e5_ids: List[str], gm_ids: List[str], k=60,
                    w_bm25=1.0, w_e5=1.0, w_gm=1.0) -> List[Tuple[str, float]]:
        rankA={pid:i for i,pid in enumerate(bm_ids)}
        rankB={pid:i for i,pid in enumerate(e5_ids)}
        rankC={pid:i for i,pid in enumerate(gm_ids)}
        scores=defaultdict(float)
        for pid, r in rankA.items(): scores[pid]+=w_bm25*(1.0/(k+r+1))
        for pid, r in rankB.items(): scores[pid]+=w_e5  *(1.0/(k+r+1))
        for pid, r in rankC.items(): scores[pid]+=w_gm  *(1.0/(k+r+1))
        return sorted(scores.items(), key=lambda x:-x[1])

    def search(self, query: str, topk: int=200) -> List[Tuple[str, float]]:
        if self._last_q == query and self._last_fused:
            return self._last_fused[:topk]

        # BM25 list
        bm_ids = self.bm25.search(query, topk=TOP_BM25)

        # E5 list
        q_norm_e5 = norm_e5_query(query)          # per-query normalization
        qe = self.e5.encode([q_norm_e5], is_query=True, batch=1, progress_desc="E5 query")[0]
        sims_e5 = (self.e5_corpus @ qe)           # cosine (embeddings are L2-normalized)
        k2 = min(TOP_E5, len(sims_e5))
        top_idx_e5 = np.argpartition(-sims_e5, k2-1)[:k2]
        top_idx_e5 = top_idx_e5[np.argsort(-sims_e5[top_idx_e5])]
        e5_ids = [self.bm25.doc_ids[i] for i in top_idx_e5]

        # Gemma list
        q_norm_gm = norm_gemma_query(query)       # per-query normalization
        qg = self.gemma.encode([q_norm_gm], is_query=True, batch=1, progress_desc="Gemma query", max_length=GEMMA_MAX_TOK)[0]
        sims_gm = (self.gemma_corpus @ qg)        # cosine (normalized)
        k3 = min(TOP_GEMMA, len(sims_gm))
        top_idx_gm = np.argpartition(-sims_gm, k3-1)[:k3]
        top_idx_gm = top_idx_gm[np.argsort(-sims_gm[top_idx_gm])]
        gm_ids = [self.bm25.doc_ids[i] for i in top_idx_gm]

        fused_with_scores = self._wrrf_fuse3(
            bm_ids, e5_ids, gm_ids, k=RRF_K,
            w_bm25=WRRF_BM25_W, w_e5=WRRF_E5_W, w_gm=WRRF_GEMMA_W
        )
        # seen=set(); out=[]
        # for pid, score in fused_with_scores:
        #     key=self.id2norm.get(pid,"")
        #     if key in seen: continue
        #     seen.add(key)
        #     out.append((pid, score))
        #     if len(out)>=topk: break
        out = fused_with_scores[:topk]
        
        self._last_q = query
        self._last_fused = out[:]
        return out

# =========================== Globals ===========================
_STATE = {}

# =========================== Helpers ===========================
def _sha1_ids(ids: List[str]) -> str:
    h = hashlib.sha1()
    for pid in ids:
        h.update(pid.encode("utf-8")); h.update(b"\n")
    return h.hexdigest()

def _normalize_min_max(scores: List[float]) -> List[float]:
    """Scales a list of scores to the [0, 1] range."""
    if not scores or len(scores) < 2:
        return [0.5] * len(scores)
    min_s, max_s = min(scores), max(scores)
    delta = max_s - min_s
    if delta < 1e-9:
        return [0.5] * len(scores)
    return [(s - min_s) / delta for s in scores]

# =========================== API funcs =========================
def preprocess(corpus_dict: Dict[str, Dict]) -> Dict:
    ids, texts = [], []
    bm25_norms = []

    # -------- Per-paragraph normalization before indexing --------
    e5_passage_norms = []
    gm_passage_norms = []

    for pid,obj in corpus_dict.items():
        t = obj.get("passage") or obj.get("text") or ""
        pid = str(pid)
        ids.append(pid)
        texts.append(t)
        bm25_norms.append(norm_bm25(t))         # BM25 per paragraph
        e5_passage_norms.append(norm_e5_passage(t))
        gm_passage_norms.append(norm_gemma_passage(t))

    _log("="*60)
    _log(f"PREPROCESS: Building BM25 + E5 + Gemma embeddings + loading BGE")
    _log("="*60)

    # BM25
    bm25 = BM25Index(k1=BM25_K1, b=BM25_B, logger=_log)
    bm25.build(ids, bm25_norms)

    # E5 encoder + caching
    e5 = E5Embedder()
    e5_mat = None
    cache_note_e5 = None
    if E5_EVAL_CACHE_DIR:
        os.makedirs(E5_EVAL_CACHE_DIR, exist_ok=True)
        meta_p = os.path.join(E5_EVAL_CACHE_DIR, "e5_meta.json")
        npy_p  = os.path.join(E5_EVAL_CACHE_DIR, "e5_corpus.npy")
        sha = _sha1_ids(ids)
        if os.path.isfile(meta_p) and os.path.isfile(npy_p):
            try:
                with open(meta_p,"r",encoding="utf-8") as f: m=json.load(f)
                if m.get("sha1_ids")==sha and model_name_key(m.get("model_path",""))==e5.model_name and m.get("num_docs")==len(ids):
                    _log(f"[E5] Loading cached corpus embeddings from {npy_p}")
                    e5_mat = np.load(npy_p, mmap_mode=None)
                    cache_note_e5 = "loaded"
            except Exception as e: _log(f"[E5] Cache read failed: {e} — recomputing.")
    if e5_mat is None:
        _log("[E5] Computing corpus embeddings...")
        t0=time.time()
        e5_mat = e5.encode(e5_passage_norms, is_query=False, batch=64, progress_desc="E5 corpus")
        _log(f"[E5] Done in {time.time()-t0:.1f}s — shape={e5_mat.shape}")
        if E5_EVAL_CACHE_DIR:
            try:
                np.save(os.path.join(E5_EVAL_CACHE_DIR,"e5_corpus.npy"), e5_mat)
                meta = {"sha1_ids": _sha1_ids(ids), "num_docs": len(ids), "model_path": e5.model_path, "dim": int(e5_mat.shape[1]), "created": time.time()}
                with open(os.path.join(E5_EVAL_CACHE_DIR,"e5_meta.json"),"w",encoding="utf-8") as f: json.dump(meta,f,ensure_ascii=False, indent=2)
                cache_note_e5 = "saved"
                _log(f"[E5] Saved cache to {E5_EVAL_CACHE_DIR}")
            except Exception as e: _log(f"[E5] Cache save failed: {e}")

    # Gemma encoder + caching
    gemma = GemmaEmbedder()
    gemma_mat = None
    cache_note_gm = None
    if GEMMA_EVAL_CACHE_DIR:
        os.makedirs(GEMMA_EVAL_CACHE_DIR, exist_ok=True)
        meta_p_gm = os.path.join(GEMMA_EVAL_CACHE_DIR, "gemma_meta.json")
        npy_p_gm  = os.path.join(GEMMA_EVAL_CACHE_DIR, "gemma_corpus.npy")
        sha = _sha1_ids(ids)
        if os.path.isfile(meta_p_gm) and os.path.isfile(npy_p_gm):
            try:
                with open(meta_p_gm,"r",encoding="utf-8") as f: m=json.load(f)
                if m.get("sha1_ids")==sha and model_name_key(m.get("model_path",""))==gemma.model_name and m.get("num_docs")==len(ids):
                    _log(f"[Gemma] Loading cached corpus embeddings from {npy_p_gm}")
                    gemma_mat = np.load(npy_p_gm, mmap_mode=None)
                    cache_note_gm = "loaded"
            except Exception as e: _log(f"[Gemma] Cache read failed: {e} — recomputing.")
    if gemma_mat is None:
        _log("[Gemma] Computing corpus embeddings...")
        t0=time.time()
        gemma_mat = gemma.encode(gm_passage_norms, is_query=False, batch=64, progress_desc="Gemma corpus", max_length=GEMMA_MAX_TOK)
        _log(f"[Gemma] Done in {time.time()-t0:.1f}s — shape={gemma_mat.shape}")
        if GEMMA_EVAL_CACHE_DIR:
            try:
                np.save(os.path.join(GEMMA_EVAL_CACHE_DIR,"gemma_corpus.npy"), gemma_mat)
                meta_gm = {"sha1_ids": _sha1_ids(ids), "num_docs": len(ids), "model_path": gemma.model_path, "dim": int(gemma_mat.shape[1]), "created": time.time()}
                with open(os.path.join(GEMMA_EVAL_CACHE_DIR,"gemma_meta.json"),"w",encoding="utf-8") as f: json.dump(meta_gm,f,ensure_ascii=False, indent=2)
                cache_note_gm = "saved"
                _log(f"[Gemma] Saved cache to {GEMMA_EVAL_CACHE_DIR}")
            except Exception as e: _log(f"[Gemma] Cache save failed: {e}")

    # Reranker
    rr = BGEReranker()

    id2text = dict(zip(ids,texts))
    id2norm = dict(zip(ids,bm25_norms))

    hybrid = HybridSearcher(bm25, e5, e5_mat, gemma, gemma_mat, id2text, id2norm)
    _STATE.update({
        "bm25": bm25, "id2text": id2text, "id2norm": id2norm,
        "e5": e5, "e5_corpus": e5_mat,
        "gemma": gemma, "gemma_corpus": gemma_mat,
        "reranker": rr, "hybrid": hybrid
    })

    reranker_params = {
        "CE_POOL": CE_POOL, "CE_MAXLEN": CE_MAXLEN, "CE_BATCH": CE_BATCH,
        "FINAL_SCORE_BGE_WEIGHT": FINAL_SCORE_BGE_WEIGHT
    }

    meta = {
        "stage1_name": "WRRF(BM25, E5, Gemma)",
        "stage1_params": {
            "TOP_BM25": TOP_BM25, "TOP_E5": TOP_E5, "TOP_GEMMA": TOP_GEMMA, "RRF_K": RRF_K,
            "WRRF_WEIGHTS": {"bm25": WRRF_BM25_W, "e5": WRRF_E5_W, "gemma": WRRF_GEMMA_W}
        },
        "reranker_name": "BGE + Hybrid Fusion (Conditional Boost)",
        "reranker_params": reranker_params,
        "candidate_pool_cap": CE_POOL,
        "stage1_search_key": "bm25",
        "bm25_backend": getattr(bm25, "_backend_name", "unknown"),
        "e5_model_path": e5.model_path,
        "gemma_model_path": gemma.model_path,
        "bge_model_path": rr.model_path,
        "cache_dir_e5": E5_EVAL_CACHE_DIR or None,
        "cache_dir_gemma": GEMMA_EVAL_CACHE_DIR or None,
        "e5_cache": cache_note_e5 or ("unused" if not E5_EVAL_CACHE_DIR else "miss"),
        "gemma_cache": cache_note_gm or ("unused" if not GEMMA_EVAL_CACHE_DIR else "miss"),
    }

    _log("✓ PREPROCESS complete.")
    return {
        "bm25": hybrid, "id2text": id2text, "id2norm": id2norm,
        "reranker": rr, "num_documents": len(ids), "_eval": meta
    }

def predict(query: Dict, pre: Dict):
    q = query.get("query","")
    if not q: return []
    hyb = _STATE.get("hybrid") or pre["bm25"]
    rr  = _STATE.get("reranker") or pre["reranker"]
    id2text = _STATE.get("id2text") or pre["id2text"]

    # Stage-1: WRRF retrieval
    cand_id_scores = hyb.search(q, topk=CE_POOL)
    if not cand_id_scores: return []
    cand_ids, rrf_scores = zip(*cand_id_scores)
    passages = [id2text[pid] for pid in cand_ids]

    # Stage-2: BGE reranker (with its own normalizers)
    q_norm_bge = norm_bge_query(q)
    passages_norm_bge = [norm_bge_passage(p) for p in passages]
    bge_scores = rr.score_pairs(q_norm_bge, passages_norm_bge, batch=CE_BATCH, max_len=CE_MAXLEN)

    # Stage-3: Normalize and combine (conditional boost)
    norm_bge = _normalize_min_max(bge_scores)
    norm_rrf = _normalize_min_max(list(rrf_scores))
    final_scores = []
    w_rrf = 1.0 - FINAL_SCORE_BGE_WEIGHT
    for bge_score, rrf_score in zip(norm_bge, norm_rrf):
        boost = w_rrf * rrf_score * (1.0 - bge_score)
        final_scores.append(bge_score + boost)

    # Final output
    out = [{"paragraph_uuid": pid, "score": float(s)}
           for pid, s in sorted(zip(cand_ids, final_scores), key=lambda x: -x[1])]
    return out