Upload 2 files

app.py

@@ -0,0 +1,1778 @@
+# ── Cell 2: Imports ────────────────────────────────────────────────────────────
+import os, re, json, time, random, shutil, unicodedata, numpy as np, pandas as pd
+from getpass import getpass
+from pymilvus import MilvusClient
+from groq import Groq
+from openai import OpenAI
+from sentence_transformers import SentenceTransformer, CrossEncoder
+from rank_bm25 import BM25Okapi
+from sklearn.metrics import roc_auc_score
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from huggingface_hub import hf_hub_download, list_repo_files, HfFileSystem, login
+from datasets import load_dataset
+import gradio as gr
+
+
+# ── Cell 3: API Keys ───────────────────────────────────────────────────────────
+# Choose your provider: "groq" or "openrouter"
+LLM_PROVIDER = "openrouter"   # ← change to "groq" if preferred
+
+import os
+from getpass import getpass
+from huggingface_hub import login
+
+def get_secret_or_prompt(secret_name, prompt_text=None):
+    """
+    Try to read secret from Google Colab Secrets.
+    If not available, ask user securely using getpass().
+    """
+
+    value = None
+
+    # Try Colab Secrets first
+    try:
+        #from google.colab import userdata
+        value = os.environ.get(secret_name)
+    except Exception:
+        value = None
+
+    # Fallback to environment variable
+    if not value:
+        value = os.environ.get(secret_name)
+
+    # Fallback to manual secure input
+    if not value:
+        prompt_text = prompt_text or f"Enter {secret_name}: "
+        value = getpass(prompt_text)
+
+    return value
+
+
+# ── HuggingFace Token ─────────────────────────────────────────────────────────
+
+HF_TOKEN = get_secret_or_prompt(
+    "HF_TOKEN",
+    "Enter HuggingFace Token: "
+)
+
+login(token=HF_TOKEN)
+os.environ["HF_TOKEN"] = HF_TOKEN
+
+print("✅ HuggingFace token loaded and login completed")
+
+
+# ── LLM Provider API Key ──────────────────────────────────────────────────────
+
+if LLM_PROVIDER == "groq":
+
+    GROQ_API_KEY = get_secret_or_prompt(
+        "GROQ_API_KEY",
+        "Enter GROQ API Key: "
+    )
+
+    OPENROUTER_API_KEY = None
+
+    os.environ["GROQ_API_KEY"] = GROQ_API_KEY
+
+    print("✅ GROQ API key loaded")
+
+elif LLM_PROVIDER == "openrouter":
+
+    OPENROUTER_API_KEY = get_secret_or_prompt(
+        "OPENROUTER_API_KEY",
+        "Enter OpenRouter API Key: "
+    )
+
+    GROQ_API_KEY = None
+
+    os.environ["OPENROUTER_API_KEY"] = OPENROUTER_API_KEY
+
+    print("✅ OpenRouter API key loaded")
+
+else:
+    raise ValueError(f"Unknown LLM_PROVIDER: {LLM_PROVIDER}")
+
+
+# ── Cell 4: Global configuration ──────────────────────────────────────────────
+BUCKET_ID      = "Phani555/IIITH-Cohort26-RAG-Batch37-storage"
+BUCKET_PREFIX  = f"hf://buckets/{BUCKET_ID}/milvus_dbs"
+MILVUS_DIR     = "/content/milvus_store/milvus_dbs"
+HF_REPO_ID     = "Phani555/IIITH-Cohort26-RAG-Batch37-storage"
+HF_REPO_TYPE   = "dataset"
+HF_FOLDER      = "ablations"
+
+# ── Download mode ─────────────────────────────────────────────────────────────
+# "chunk_v5_domain_aware" : new advanced domain-aware chunk_v5 indexes (recommended)
+# "llm_embedder_default"  : legacy llm_embedder default indexes
+DOWNLOAD_MODE = "chunk_v5_domain_aware"
+
+if DOWNLOAD_MODE == "chunk_v5_domain_aware":
+    INDEX_VERSION = "chunk_v5_domain_aware"
+    DOMAIN_EMBEDDING_RECOMMENDATION = {
+        "Customer_Support": "qwen3_embedding_0_6b",
+        "Bio_Medical":      "bge_m3",
+        "General_Knowledge":"qwen3_embedding_0_6b",
+        "Legal_Contracts":  "bge_m3",
+        "Finance":          "bge_m3",
+    }
+    EMBEDDING_TYPE = "bge_m3"
+else:  # llm_embedder_default
+    INDEX_VERSION  = "default"
+    DOMAIN_EMBEDDING_RECOMMENDATION = None
+    EMBEDDING_TYPE = "llm_embedder"
+
+# ── Embedding models ───────────────────────────────────────────────────────────
+EMBED_MODELS = {
+    "bge_small":             "BAAI/bge-small-en-v1.5",
+    "llm_embedder":          "BAAI/llm-embedder",
+    "bge_m3":                "BAAI/bge-m3",
+    "qwen3_embedding_0_6b":  "Qwen/Qwen3-Embedding-0.6B",
+}
+EMBEDDING_CHOICES = list(EMBED_MODELS.keys())
+
+# ── Model lists per LLM provider ──────────────────────────────────────────────
+GROQ_LLM_CHOICES = [
+    "llama-3.1-8b-instant",
+    "gemma2-9b-it",
+    "llama-3.3-70b-versatile",
+    "mixtral-8x7b-32768",
+    "qwen/qwen3-32b",
+    "qwen-qwq-32b",
+    "deepseek-r1-distill-llama-70b",
+]
+OPENROUTER_LLM_CHOICES = [
+    "meta-llama/llama-3.1-8b-instruct",
+    "meta-llama/llama-3.3-70b-instruct",
+    "openai/gpt-oss-20b",
+    "openai/gpt-oss-120b",
+    "qwen/qwen3-32b",
+    "deepseek/deepseek-r1",
+    "moonshotai/kimi-k2-instruct",
+    "openai/gpt-oss-safeguard-20b",
+]
+LLM_CHOICES = OPENROUTER_LLM_CHOICES if LLM_PROVIDER == "openrouter" else GROQ_LLM_CHOICES
+
+# ── Runtime globals ────────────────────────────────────────────────────────────
+MODEL_NAME     = LLM_CHOICES[0]
+MODEL_NAME_BIG = LLM_CHOICES[4] if len(LLM_CHOICES) > 4 else LLM_CHOICES[-1]
+
+# ── Feature flags ──────────────────────────────────────────────────────────────
+ENABLE_HYBRID        = True
+ENABLE_HYDE          = False
+ENABLE_RERANKING     = False
+RERANKER_TYPE        = "monot5"    # monot5 | tilde
+ENABLE_RRF           = True        # Reciprocal Rank Fusion inside hybrid search
+RRF_K                = 60          # standard RRF constant
+PROMPT_STRATEGY      = "short"     # short | long | long_cot
+ENABLE_REPACKING     = False
+REPACK_STRATEGY      = "sides"     # forward | reverse | sides
+ENABLE_SUMMARIZATION = False
+SUMMARIZATION_TYPE   = "recomp"    # recomp | longllmlingua
+ENABLE_QUERY_REWRITING     = False
+ENABLE_QUERY_DECOMPOSITION = False
+ENABLE_QUERY_CLASSIFICATION= False
+MAX_SUBQUERIES             = 3
+QUERY_REWRITE_MODEL        = None
+QUERY_DECOMPOSE_MODEL      = None
+RETRIEVE_DEBUG             = False
+
+# ── Tunable knobs ─────────────────────────────────────────────────────────────
+RETRIEVE_TOP_K         = 10
+RERANK_TOP_K           = 3
+HYBRID_ALPHA           = 0.5
+MONOT5_MODEL           = "castorini/monot5-base-msmarco-10k"
+TILDE_MODEL            = "BAAI/bge-reranker-base"
+RECOMP_TOP_K_SENTS     = 6
+RECOMP_MIN_SCORE       = 0.00
+RECOMP_GROUNDING_BOOST = 0.15
+RECOMP_MIN_KEEP_RATIO  = 0.30
+RECOMP_KEEP_CRITICAL   = True
+LLMLINGUA_RATE         = 0.5
+
+# ── Runtime state ─────────────────────────────────────────────────────────────
+milvus_clients              = {}
+bm25_indexes                = {}
+loaded_embedding_models     = {}   # keyed by embedding_type string
+embed_model                 = None # single fallback embed model
+llm_client                  = None
+monot5_reranker             = None
+tilde_reranker              = None
+llmlingua_compressor        = None
+ragbench_by_domain          = {}
+LEGAL_SAMPLE_TO_CONTRACT_ID = {}
+
+DOMAIN_NAMES = [
+    "Bio_Medical",
+    "General_Knowledge",
+    "Customer_Support",
+    "Finance",
+    "Legal_Contracts",
+]
+
+GROUNDING_PATTERNS = [
+    r"\b(?:must|should|shall|cannot|can't|never|always|only|except|unless|required|recommended)\b",
+    r"\b(?:warning|caution|note|important|attention)\b",
+    r"\b(?:do not|don't|does not|did not|not allowed|not recommended|never)\b",
+    r"\b\d+(?:\.\d+)?\s*(?:%|percent|seconds?|minutes?|hours?|days?|weeks?|months?|years?)\b",
+    r"\b\d+(?:\.\d+)?\s*(?:GB|MB|KB|TB|kg|g|mg|mm|cm|m|km|degrees?|°C|°F)\b",
+    r"[$€£¥]\s*\d+(?:,\d{3})*(?:\.\d+)?",
+    r"\b\d+(?:,\d{3})*(?:\.\d+)?\s*(?:dollars?|rupees?|crores?|lakhs?|million|billion)\b",
+    r"\b(?:19|20)\d{2}\b",
+    r"\b\d{2,}\b",
+    r"\b[A-Z]{2,}[-_]?\d+[A-Z0-9-]*\b",
+    r"\b[A-Z0-9]{3,}[-_][A-Z0-9]{2,}\b",
+    r"\b[A-Z]{3,}\b",
+]
+GROUNDING_REGEX = re.compile("|".join(GROUNDING_PATTERNS), re.IGNORECASE)
+
+print(f"Config loaded.  Mode: {DOWNLOAD_MODE}  |  Provider: {LLM_PROVIDER}  |  Models: {len(LLM_CHOICES)}")
+
+
+# ── Cell 5: Pipeline functions (Advanced – chunk_v5 + RRF + Legal contract filtering) ──
+# NOTE: _hf_fs is initialized in Cell 7. hf_path_exists() uses globals() so it
+# safely resolves _hf_fs at call-time, not at definition-time.
+
+# ── Utilities ──────────────────────────────────────────────────────────────────
+
+def _safe_message_content(response):
+    try:
+        msg = response.choices[0].message
+        content = getattr(msg, "content", None)
+        return str(content).strip() if content else ""
+    except Exception:
+        return ""
+
+def _sanitize(text):
+    if not text: return text
+    text = unicodedata.normalize("NFC", str(text))
+    return text.encode("ascii", errors="replace").decode("ascii")
+
+def get_domain(dataset):
+    if dataset in ("covidqa","pubmedqa"):             return "Bio_Medical"
+    elif dataset in ("expertqa","hagrid","hotpotqa","msmarco"): return "General_Knowledge"
+    elif dataset in ("delucionqa","emanual","techqa"): return "Customer_Support"
+    elif dataset in ("finqa","tatqa"):                return "Finance"
+    else:                                             return "Legal_Contracts"
+
+def split_into_sentences(text):
+    return [s.strip() for s in re.split(r'(?<=[.!?])\s+', str(text).strip()) if s.strip()]
+
+def _tokenize(text):
+    return re.findall(r'\w+', str(text).lower())
+
+def _normalize(scores):
+    arr = np.array(scores, dtype=float)
+    if len(arr) == 0 or arr.max() == arr.min(): return np.zeros_like(arr)
+    return (arr - arr.min()) / (arr.max() - arr.min())
+
+def _count_grounding_signals(sentence):
+    return len(GROUNDING_REGEX.findall(str(sentence)))
+
+def _is_critical_sentence(sentence):
+    pat = re.compile(
+        r"\b(?:warning|caution|important|must|must not|cannot|can't|do not|don't|never|only|except|unless|required)\b",
+        re.IGNORECASE)
+    return bool(pat.search(str(sentence)))
+
+
+# ── LLM client factory ─────────────────────────────────────────────────────────
+
+def get_llm_client():
+    if LLM_PROVIDER == "groq":
+        return Groq(api_key=GROQ_API_KEY)
+    elif LLM_PROVIDER == "openrouter":
+        return OpenAI(api_key=OPENROUTER_API_KEY, base_url="https://openrouter.ai/api/v1")
+    raise ValueError(f"Unknown LLM_PROVIDER: {LLM_PROVIDER}")
+
+
+# ── DB path helpers ────────────────────────────────────────────────────────────
+
+def get_index_folder(embedding_type=None, index_version=None):
+    embedding_type = embedding_type or EMBEDDING_TYPE
+    index_version  = index_version  or INDEX_VERSION
+    return embedding_type if index_version == "default" else f"{embedding_type}_{index_version}"
+
+def get_db_path(domain_name, embedding_type=None, index_version=None):
+    folder = get_index_folder(embedding_type, index_version)
+    db_dir = os.path.join(MILVUS_DIR, folder)
+    os.makedirs(db_dir, exist_ok=True)
+    return os.path.join(db_dir, f"{domain_name}.db")
+
+def get_embedding_type_for_domain(domain_name):
+    rec = globals().get("DOMAIN_EMBEDDING_RECOMMENDATION")
+    if rec: return rec.get(domain_name, EMBEDDING_TYPE)
+    return EMBEDDING_TYPE
+
+def get_embed_model_for_domain(domain_name):
+    emb_type = get_embedding_type_for_domain(domain_name)
+    models   = globals().get("loaded_embedding_models", {})
+    if emb_type not in models:
+        raise ValueError(f"Embedding type '{emb_type}' not in loaded_embedding_models. Run Cell 7 first.")
+    return models[emb_type]
+
+
+# ── HF filesystem helper ───────────────────────────────────────────────────────
+# Uses globals() so _hf_fs is resolved at call-time (Cell 7), not import-time (Cell 5).
+
+def hf_path_exists(path):
+    fs = globals().get("_hf_fs")
+    if fs is None:
+        raise RuntimeError("_hf_fs not initialised — run Cell 7 before Cell 8.")
+    try:
+        fs.ls(path); return True
+    except Exception:
+        return False
+
+
+# ── Legal contract helpers ─────────────────────────────────────────────────────
+
+def build_contract_filter_expr(contract_id):
+    contract_id = str(contract_id).replace('"', '\\"')
+    return f'contract_id == "{contract_id}"'
+
+def load_legal_sample_to_contract_mapping(local_path=None):
+    if local_path is None:
+        legal_folder = get_index_folder(get_embedding_type_for_domain("Legal_Contracts"), INDEX_VERSION)
+        local_path   = os.path.join(MILVUS_DIR, legal_folder, "legal_sample_to_contract_id.json")
+    if not os.path.exists(local_path):
+        print(f"  WARNING: Legal mapping not found: {local_path}")
+        return {}
+    with open(local_path, "r") as f:
+        mapping = json.load(f)
+    print(f"  Legal sample->contract mapping loaded: {len(mapping):,} entries")
+    return mapping
+
+def get_contract_id_for_legal_sample(sample_id):
+    mapping = globals().get("LEGAL_SAMPLE_TO_CONTRACT_ID", {})
+    sid = str(sample_id)
+    if sid not in mapping:
+        raise ValueError(f"sample_id '{sid}' not found in Legal mapping.")
+    return mapping[sid]
+
+
+# ── Query Classification ───────────────────────────────────────────────────────
+
+def classify_query(query, domain_name=None):
+    if not ENABLE_QUERY_CLASSIFICATION: return "RAG"
+    rag_domains = {"Bio_Medical","General_Knowledge","Customer_Support","Finance","Legal_Contracts"}
+    if domain_name in rag_domains: return "RAG"
+    llm_keywords = ["who is","what is","when was","where is","define","explain",
+                    "tell me about","what are","why is","how does","what does"]
+    if any(kw in str(query).lower() for kw in llm_keywords): return "LLM"
+    return "RAG"
+
+
+# ── Query Rewriting ────────────────────────────────────────────────────────────
+
+def rewrite_query(query, domain_name, llm_client, model_name=None):
+    if not ENABLE_QUERY_REWRITING: return query
+    model  = model_name or QUERY_REWRITE_MODEL or MODEL_NAME
+    prompt = f"""Rewrite the question to improve document retrieval. Apply only when needed.
+Domain: {domain_name}
+Rules: Preserve meaning. Fix grammar. Expand abbreviations. Preserve all names/numbers/terms.
+Do not answer. Return ONLY the rewritten query.
+Original question: {query}""".strip()
+    try:
+        resp = llm_client.chat.completions.create(
+            model=model,
+            messages=[{"role":"system","content":"You rewrite questions to improve semantic document retrieval. Return only the rewritten question."},
+                      {"role":"user","content":_sanitize(prompt)}],
+            temperature=0.0, max_tokens=150,
+        )
+        rewritten = _safe_message_content(resp).strip()
+        return rewritten if rewritten and len(rewritten) < 600 else query
+    except Exception as e:
+        print(f"Query rewriting failed: {e}"); return query
+
+
+# ── Query Decomposition helpers ────────────────────────────────────────────────
+
+def _clean_subquery_text(text):
+    if text is None: return ""
+    text = str(text).strip().replace("```json","").replace("```","").strip()
+    text = text.rstrip(",").strip('"').strip("'").strip()
+    text = re.sub(r"^\s*[-*]\s*","",text); text = re.sub(r"^\s*\d+[\).\:\-]\s*","",text)
+    return text.strip()
+
+def _looks_like_explanation_line(text):
+    if not text: return True
+    tl = text.lower().strip()
+    bad = ["here are","here is","decomposed","search queries","the decomposed",
+           "queries:","subqueries:","output:","json:","answer:"]
+    if any(tl.startswith(p) for p in bad): return True
+    if tl in {"queries","subqueries","search queries","decomposed search queries"}: return True
+    return False
+
+def _parse_json_object_line(line):
+    line = _clean_subquery_text(line)
+    if not line: return None
+    try:
+        obj = json.loads(line)
+        if isinstance(obj, dict):
+            for k in ["query","question","subquery","search_query"]:
+                if k in obj and str(obj[k]).strip(): return str(obj[k]).strip()
+        if isinstance(obj, str): return obj.strip()
+    except Exception: pass
+    m = re.search(r'"(?:query|question|subquery|search_query)"\s*:\s*"([^"]+)"', line)
+    if m: return m.group(1).strip()
+    return None
+
+def _split_multi_question_locally(query, max_subqueries=None):
+    max_subqueries = max_subqueries or MAX_SUBQUERIES
+    parts = [p.strip() for p in re.split(r"\?\s*", str(query).strip()) if p.strip()]
+    if len(parts) <= 1: return None
+    return [(p+"?" if not p.endswith("?") else p) for p in parts[:max_subqueries]]
+
+def _parse_subqueries(raw_text, original_query, max_subqueries=None):
+    max_subqueries = max_subqueries or MAX_SUBQUERIES
+    if not raw_text: return [original_query]
+    text = str(raw_text).strip().replace("```json","").replace("```","").strip()
+    try:
+        parsed = json.loads(text)
+        if isinstance(parsed, list):
+            subs = []
+            for item in parsed:
+                if isinstance(item, dict):
+                    for k in ["query","question","subquery","search_query"]:
+                        if k in item and str(item[k]).strip(): subs.append(str(item[k]).strip()); break
+                elif isinstance(item, str): subs.append(item.strip())
+            subs = [_clean_subquery_text(q) for q in subs if _clean_subquery_text(q)]
+            return subs[:max_subqueries] or [original_query]
+        elif isinstance(parsed, dict):
+            raw_list = parsed.get("subqueries") or parsed.get("queries") or parsed.get("questions") or []
+            if isinstance(raw_list, list):
+                subs = [_clean_subquery_text(q) for q in raw_list if _clean_subquery_text(q)]
+                return subs[:max_subqueries] or [original_query]
+    except Exception: pass
+    subqueries = []
+    for raw_line in text.splitlines():
+        line = _clean_subquery_text(raw_line)
+        if not line or _looks_like_explanation_line(line): continue
+        obj_q = _parse_json_object_line(line)
+        if obj_q:
+            obj_q = _clean_subquery_text(obj_q)
+            if obj_q and not _looks_like_explanation_line(obj_q): subqueries.append(obj_q)
+            continue
+        if line.startswith("{") or line.endswith("}") or line in {"[","]","{","}"}: continue
+        subqueries.append(line)
+    deduped = []
+    for q in subqueries:
+        q = _clean_subquery_text(q)
+        if q and q not in deduped: deduped.append(q)
+    return deduped[:max_subqueries] or [original_query]
+
+def decompose_query(query, llm_client, domain=None, model=None, max_subqueries=None):
+    if not ENABLE_QUERY_DECOMPOSITION: return [query]
+    max_subqueries = max_subqueries or MAX_SUBQUERIES
+    local_split = _split_multi_question_locally(query, max_subqueries)
+    if local_split: return local_split
+    model = model or QUERY_DECOMPOSE_MODEL or MODEL_NAME
+    if not model: return [query]
+    prompt = f"""Decompose the question into at most {max_subqueries} retrieval-focused search queries.
+Return ONLY a valid JSON list of strings. No explanations. No markdown.
+Example: ["What caused the 2008 crisis?", "Which banks failed in 2008?"]
+Rules: If already simple return list with original. Preserve all technical terms. Do not answer.
+Domain: {domain}
+Question: {query}""".strip()
+    try:
+        resp = llm_client.chat.completions.create(
+            model=model,
+            messages=[{"role":"system","content":"You decompose complex questions into retrieval subqueries and return only a JSON list of strings."},
+                      {"role":"user","content":_sanitize(prompt)}],
+            temperature=0.0, max_tokens=300,
+        )
+        return _parse_subqueries(_safe_message_content(resp), original_query=query, max_subqueries=max_subqueries)
+    except Exception as e:
+        print(f"Query decomposition failed: {e}"); return [query]
+
+
+# ── Reranking ──────────────────────────────────────────────────────────────────
+
+class MonoT5Reranker:
+    def __init__(self, model_name=None):
+        model_name = model_name or MONOT5_MODEL
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model     = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        self.device    = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device); self.model.eval()
+        self.true_id  = self.tokenizer.convert_tokens_to_ids("▁true")
+        self.false_id = self.tokenizer.convert_tokens_to_ids("▁false")
+        if not self.true_id  or self.true_id  < 0: self.true_id  = self.tokenizer.encode("true",  add_special_tokens=False)[0]
+        if not self.false_id or self.false_id < 0: self.false_id = self.tokenizer.encode("false", add_special_tokens=False)[0]
+        print(f"MonoT5 loaded: {model_name} on {self.device}")
+
+    def score(self, query, document):
+        text = f"Query: {query} Document: {document} Relevant:"
+        enc  = self.tokenizer(text, return_tensors="pt", max_length=512, truncation=True).to(self.device)
+        with torch.no_grad():
+            out = self.model.generate(**enc, max_new_tokens=1, return_dict_in_generate=True, output_scores=True)
+        logits = out.scores[0][0]
+        probs  = torch.softmax(torch.stack([logits[self.false_id], logits[self.true_id]]), dim=0)
+        return float(probs[1].item())
+
+    def compute_scores(self, query, texts):
+        return np.array([self.score(query, t) for t in texts], dtype=float)
+
+def get_monot5_reranker():
+    global monot5_reranker
+    if monot5_reranker is None: monot5_reranker = MonoT5Reranker(MONOT5_MODEL)
+    return monot5_reranker
+
+def get_tilde_reranker():
+    global tilde_reranker
+    if tilde_reranker is None:
+        dev = "cuda" if torch.cuda.is_available() else "cpu"
+        tilde_reranker = CrossEncoder(TILDE_MODEL, device=dev)
+        print(f"TILDE reranker loaded on {dev}")
+    return tilde_reranker
+
+def rerank_documents(query, documents, top_k=3):
+    """Rerank documents; preserves all metadata fields including Legal contract_id."""
+    if not documents: return []
+    texts  = [d["text"] if isinstance(d, dict) else d for d in documents]
+    rtype  = RERANKER_TYPE.lower().strip()
+    scores = get_monot5_reranker().compute_scores(query, texts) if rtype == "monot5" \
+             else np.asarray(get_tilde_reranker().predict([(query, t) for t in texts], show_progress_bar=False), dtype=float).reshape(-1)
+    ranked_idx = np.argsort(scores)[::-1][:top_k]
+    reranked = []
+    for i in ranked_idx:
+        # dict() shallow-copies ALL fields (dense_score, bm25_score, contract_id, etc.)
+        item = dict(documents[i]) if isinstance(documents[i], dict) else {"text": documents[i]}
+        item["base_score"]    = item.get("score")   # preserve original retrieval score
+        item["score"]         = float(scores[i])
+        item["rerank_score"]  = float(scores[i])
+        item["reranker_type"] = rtype
+        reranked.append(item)
+    return reranked
+
+
+# ── BM25 (stores contract_ids for Legal to enable per-contract filtering) ─────
+
+def build_bm25_index(domain_name, clients):
+    client = clients[domain_name]; col = domain_name.lower()
+    try:
+        if "Loaded" not in str(client.get_load_state(col)): client.load_collection(col)
+    except Exception: pass
+    try:    n = int(client.get_collection_stats(col).get("row_count", 0))
+    except Exception: return
+    if n == 0: return
+    output_fields = ["text"]
+    if domain_name == "Legal_Contracts": output_fields.append("contract_id")
+    rows = client.query(collection_name=col, filter="", limit=n, output_fields=output_fields)
+    if not rows: return
+    texts = []; contract_ids = []
+    for r in rows:
+        t = r.get("text","")
+        if not t: continue
+        texts.append(t)
+        if domain_name == "Legal_Contracts": contract_ids.append(r.get("contract_id"))
+    if not texts: return
+    bm25_indexes[domain_name] = {
+        "bm25":         BM25Okapi([_tokenize(t) for t in texts]),
+        "texts":        texts,
+        "contract_ids": contract_ids if domain_name == "Legal_Contracts" else None,
+    }
+    print(f"  BM25 built: {len(texts)} docs  [{domain_name}]")
+    if domain_name == "Legal_Contracts":
+        valid = sum(1 for c in contract_ids if c is not None)
+        print(f"    Legal contract_ids: {valid:,}/{len(texts):,}")
+
+def build_all_bm25_indexes(clients):
+    bm25_indexes.clear()
+    for d in clients: build_bm25_index(d, clients)
+
+
+# ── HyDE ───────────────────────────────────────────────────────────────────────
+
+def generate_hyde(query, llm_client, model_name=None):
+    model  = model_name or MODEL_NAME
+    prompt = f"Write a brief factual passage answering this question (under 4 sentences).\nQuestion: {query}\nPassage:"
+    try:
+        resp = llm_client.chat.completions.create(
+            model=model,
+            messages=[{"role":"system","content":"You write hypothetical answer passages for retrieval."},
+                      {"role":"user","content":_sanitize(prompt)}],
+            temperature=0.2, max_tokens=300,
+        )
+        return _safe_message_content(resp)
+    except Exception as e:
+        print(f"HyDE failed: {e}"); return ""
+
+
+# ── Reciprocal Rank Fusion ─────────────────────────────────────────────────────
+
+def reciprocal_rank_fusion(dense_results, bm25_results, top_k=20, rrf_k=60,
+                           dense_meta=None, bm25_meta=None):
+    """
+    Fuse dense and BM25 ranked lists using RRF.
+    score(doc) = 1/(k + rank_dense) + 1/(k + rank_bm25)
+    dense_meta / bm25_meta: optional dicts of extra fields per text (e.g. Legal metadata).
+    """
+    dense_meta = dense_meta or {}; bm25_meta = bm25_meta or {}
+    rrf_scores = {}
+    for rank, (text, score) in enumerate(
+            sorted(dense_results.items(), key=lambda x: x[1], reverse=True), start=1):
+        rrf_scores.setdefault(text, {"text":text,"dense_score":float(score),"bm25_score":0.0,"score":0.0})
+        rrf_scores[text]["dense_score"] = float(score)
+        rrf_scores[text]["score"] += 1.0 / (rrf_k + rank)
+        if text in dense_meta: rrf_scores[text].update(dense_meta[text])
+    for rank, (text, score) in enumerate(
+            sorted(bm25_results.items(), key=lambda x: x[1], reverse=True), start=1):
+        rrf_scores.setdefault(text, {"text":text,"dense_score":0.0,"bm25_score":float(score),"score":0.0})
+        rrf_scores[text]["bm25_score"] = float(score)
+        rrf_scores[text]["score"] += 1.0 / (rrf_k + rank)
+        # dense_meta takes priority over bm25_meta for Legal metadata consistency
+        if text not in dense_meta and text in bm25_meta:
+            rrf_scores[text].update(bm25_meta[text])
+    fused = sorted(rrf_scores.values(), key=lambda x: x["score"], reverse=True)
+    return fused[:top_k]
+
+
+# ── Hybrid search (dense + BM25, Legal contract filtering, RRF or alpha fusion) ─
+
+def hybrid_search(query, domain_name, embed_model, top_k=20, alpha=0.5, contract_id=None):
+    client = milvus_clients[domain_name]; col = domain_name.lower()
+    try:
+        if "Loaded" not in str(client.get_load_state(col)): client.load_collection(col)
+    except Exception: pass
+
+    # ── Dense search ───────────────────────────────────────────���─────────────
+    q_emb = embed_model.encode([query], normalize_embeddings=True, convert_to_numpy=True).astype("float32")
+    output_fields = ["text"]
+    if domain_name == "Legal_Contracts": output_fields += ["contract_id","source_doc_id","source_hash"]
+    search_kwargs = dict(collection_name=col, data=q_emb.tolist(), limit=top_k,
+                         output_fields=output_fields, search_params={"metric_type":"IP","params":{}})
+    if domain_name == "Legal_Contracts" and contract_id is not None:
+        search_kwargs["filter"] = build_contract_filter_expr(contract_id)
+    dense_hits = client.search(**search_kwargs)
+    dense_results = {}; dense_meta = {}
+    hit_list = dense_hits[0] if (dense_hits and isinstance(dense_hits[0], (list,tuple))) else dense_hits
+    for hit in hit_list:
+        entity   = (hit.get("entity",{}) or hit) if isinstance(hit,dict) else (getattr(hit,"entity",{}) or {})
+        distance = hit.get("distance",0.0)        if isinstance(hit,dict) else getattr(hit,"distance",0.0)
+        text = entity.get("text","")
+        if not text: continue
+        dense_results[text] = float(distance)
+        if domain_name == "Legal_Contracts":
+            dense_meta[text] = {"contract_id":entity.get("contract_id"),
+                                 "source_doc_id":entity.get("source_doc_id"),
+                                 "source_hash":entity.get("source_hash")}
+
+    # ── BM25 search ───────────────────────────────────────────────────────────
+    bm25_obj = bm25_indexes.get(domain_name)
+    if not bm25_obj:  # fallback to dense-only
+        rows = [{"text":t,"score":s,"dense_score":s,"bm25_score":0.0,"hybrid_fallback":True}
+                for t,s in sorted(dense_results.items(),key=lambda x:-x[1])[:top_k]]
+        if domain_name == "Legal_Contracts":
+            for r in rows: r.update(dense_meta.get(r["text"],{}))
+        return rows
+
+    bm25_scores = bm25_obj["bm25"].get_scores(_tokenize(query))
+    bm25_texts  = bm25_obj["texts"]
+    bm25_cids   = bm25_obj.get("contract_ids")
+
+    # For Legal: filter BM25 candidates to the same contract before ranking
+    if domain_name == "Legal_Contracts" and contract_id is not None and bm25_cids:
+        candidate_idx = [i for i,cid in enumerate(bm25_cids) if str(cid)==str(contract_id)]
+    else:
+        candidate_idx = list(range(len(bm25_texts)))
+    top_bm25_idx = sorted(candidate_idx, key=lambda i: bm25_scores[i], reverse=True)[:top_k]
+
+    bm25_results = {}; bm25_meta = {}
+    for i in top_bm25_idx:
+        text = bm25_texts[i]
+        bm25_results[text] = float(bm25_scores[i])
+        if domain_name == "Legal_Contracts":
+            bm25_meta[text] = {"contract_id": str(contract_id) if contract_id else (bm25_cids[i] if bm25_cids else None)}
+
+    # ── Fusion ────────────────────────────────────────────────────────────────
+    if ENABLE_RRF:
+        return reciprocal_rank_fusion(dense_results, bm25_results, top_k=top_k, rrf_k=RRF_K,
+                                      dense_meta=dense_meta, bm25_meta=bm25_meta)
+
+    # Alpha-weighted min-max fusion (fallback when RRF disabled)
+    all_texts = sorted(set(dense_results)|set(bm25_results))
+    d_vals    = [dense_results.get(t,0.0) for t in all_texts]
+    b_vals    = [bm25_results.get(t,0.0) for t in all_texts]
+    d_norm, b_norm = _normalize(d_vals), _normalize(b_vals)
+    combined = []
+    for i, text in enumerate(all_texts):
+        row = {"text":text,"score":float(alpha*d_norm[i]+(1-alpha)*b_norm[i]),
+               "dense_score":float(d_vals[i]),"bm25_score":float(b_vals[i]),"alpha":alpha}
+        if domain_name == "Legal_Contracts":
+            row.update(dense_meta.get(text, bm25_meta.get(text,{})))
+            if "contract_id" not in row and contract_id is not None:
+                row["contract_id"] = str(contract_id)
+        combined.append(row)
+    combined.sort(key=lambda x: -x["score"])
+    return combined[:top_k]
+
+
+# ── Repacking ──────────────────────────────────────────────────────────────────
+
+def repack_documents(docs, strategy="sides"):
+    """
+    Reorder retrieved documents for LLM attention bias.
+    forward: most-relevant first (no change)
+    reverse: most-relevant last (benefits models that attend to end of context)
+    sides:   U-shape — highest-relevance at both ends, lowest in the middle
+    """
+    if not docs: return []
+    if strategy == "forward": return docs
+    if strategy == "reverse": return docs[::-1]
+    if strategy == "sides":
+        n, result, left, right = len(docs), [None]*len(docs), 0, len(docs)-1
+        for i, doc in enumerate(docs):
+            if i % 2 == 0: result[left]  = doc; left  += 1
+            else:           result[right] = doc; right -= 1
+        return result
+    raise ValueError(f"Unknown REPACK_STRATEGY: {strategy}")
+
+
+# ── Summarization ──────────────────────────────────────────────────────────────
+
+def recomp_summarize(query, docs, em, top_k=6, min_score=0.0,
+                     grounding_boost=0.15, min_keep_ratio=0.30, keep_critical=True):
+    texts     = [d.get("text","") if isinstance(d,dict) else d for d in docs]
+    sentences = [s.strip() for doc in texts for s in split_into_sentences(doc) if s.strip()]
+    if not sentences: return ""
+    q_emb  = em.encode([query], normalize_embeddings=True)
+    s_emb  = em.encode(sentences, normalize_embeddings=True)
+    scores = (q_emb @ s_emb.T).flatten() + np.array([_count_grounding_signals(s)*grounding_boost for s in sentences])
+    crits  = {i for i,s in enumerate(sentences) if keep_critical and _is_critical_sentence(s)}
+    valid  = np.where(scores >= min_score)[0]
+    if len(valid) == 0: valid = np.array([int(np.argmax(scores))])
+    keep   = min(max(top_k, int(np.ceil(len(sentences)*min_keep_ratio))), len(sentences))
+    chosen = sorted(set(list(valid[np.argsort(scores[valid])[::-1][:keep]])) | crits)
+    return " ".join(sentences[i] for i in chosen)
+
+def _get_llmlingua():
+    global llmlingua_compressor
+    if llmlingua_compressor is None:
+        from llmlingua import PromptCompressor
+        dev = "cuda" if torch.cuda.is_available() else "cpu"
+        llmlingua_compressor = PromptCompressor(
+            model_name="microsoft/llmlingua-2-bert-base-multilingual-cased-meetingbank",
+            use_llmlingua2=True, device_map=dev)
+        print(f"LLMLingua loaded on {dev}")
+    return llmlingua_compressor
+
+def llmlingua_compress(query, docs, rate=0.5):
+    texts = [d.get("text","") if isinstance(d,dict) else d for d in docs]
+    comp  = _get_llmlingua()
+    parts = []
+    for t in texts:
+        if not t or not t.strip(): continue
+        try:    parts.append(comp.compress_prompt(t, question=query, rate=rate)["compressed_prompt"])
+        except Exception as e: print(f"LLMLingua chunk failed: {e}"); parts.append(t)
+    return "\n\n".join(parts)
+
+def summarize_docs(query, docs, em=None, llm_client=None):
+    """
+    FIX: explicit None guard on em before calling encode().
+    Falls back to global embed_model, then raises a clear error.
+    """
+    if not ENABLE_SUMMARIZATION:
+        return [d.get("text","") if isinstance(d,dict) else d for d in docs]
+
+    # Resolve embed model — must be non-None before encode()
+    if em is None:
+        em = globals().get("embed_model")
+    if em is None:
+        raise RuntimeError("summarize_docs: no embed_model available. Run Cell 7 first.")
+
+    if SUMMARIZATION_TYPE == "recomp":
+        s = recomp_summarize(query, docs, em, RECOMP_TOP_K_SENTS, RECOMP_MIN_SCORE,
+                             RECOMP_GROUNDING_BOOST, RECOMP_MIN_KEEP_RATIO, RECOMP_KEEP_CRITICAL)
+        return [s] if s else []
+    elif SUMMARIZATION_TYPE == "longllmlingua":
+        c = llmlingua_compress(query, docs, LLMLINGUA_RATE)
+        return [c] if c else []
+    raise ValueError(f"Unknown SUMMARIZATION_TYPE: {SUMMARIZATION_TYPE}")
+
+
+# ── Main retrieve ──────────────────────────────────────────────────────────────
+# Pipeline order: HyDE → Retrieve → Rerank → Repack → Summarize
+
+def retrieve(query, domain_name, embed_model=None, llm_client=None, top_k=None,
+             rewritten_query=None, sample_id=None, contract_id=None):
+    """
+    FIX: top_k now defaults to RETRIEVE_TOP_K (10), not RERANK_TOP_K (3).
+    The fetch_k logic already enlarges the initial pool; top_k is the final
+    count after reranking.
+    """
+    if domain_name not in milvus_clients:
+        raise ValueError(f"Domain '{domain_name}' not loaded.")
+
+    # FIX: default to RETRIEVE_TOP_K for initial fetch, not RERANK_TOP_K
+    if top_k is None:
+        top_k = globals().get("RETRIEVE_TOP_K", 10)
+    top_k = int(top_k)
+
+    llm = llm_client or globals().get("llm_client")
+
+    # Resolve domain-specific embed model
+    em = embed_model
+    if em is None:
+        try:   em = get_embed_model_for_domain(domain_name)
+        except Exception: em = globals().get("embed_model")
+    if em is None:
+        raise ValueError(f"No embed_model available for domain '{domain_name}'. Run Cell 7 first.")
+
+    # Legal contract filtering (section 1.5)
+    legal_contract_id = None
+    if domain_name == "Legal_Contracts":
+        if contract_id is not None:
+            legal_contract_id = str(contract_id)
+        elif sample_id is not None:
+            try:   legal_contract_id = get_contract_id_for_legal_sample(sample_id)
+            except Exception as e:
+                print(f"  Could not resolve Legal contract_id for sample_id={sample_id}: {e}")
+        if legal_contract_id is None:
+            print("  WARNING: Legal_Contracts retrieval without contract filter — cross-contract contamination possible")
+
+    # Fetch more candidates if downstream processing will reduce count
+    fetch_k = max(int(RETRIEVE_TOP_K if (ENABLE_HYBRID or ENABLE_RERANKING or ENABLE_SUMMARIZATION) else top_k), top_k)
+
+    eff_q = rewritten_query or query; search_q = eff_q
+
+    # HyDE query expansion
+    if ENABLE_HYDE and llm:
+        try:
+            hyde = generate_hyde(eff_q, llm)
+            if hyde: search_q = f"{eff_q} {hyde}"
+        except Exception as e: print(f"HyDE failed: {e}")
+
+    # Retrieve
+    if ENABLE_HYBRID:
+        retrieved = hybrid_search(search_q, domain_name, em, top_k=fetch_k,
+                                  alpha=HYBRID_ALPHA, contract_id=legal_contract_id) or []
+        for d in retrieved:
+            if isinstance(d, dict): d.setdefault("retrieval_type","hybrid")
+    else:
+        client = milvus_clients[domain_name]; col = domain_name.lower()
+        try:
+            if "Loaded" not in str(client.get_load_state(col)): client.load_collection(col)
+        except Exception: pass
+        q_emb = em.encode([search_q], normalize_embeddings=True, convert_to_numpy=True).astype("float32")
+        output_fields = ["text"]
+        if domain_name == "Legal_Contracts": output_fields += ["contract_id","source_doc_id","source_hash"]
+        skw = dict(collection_name=col, data=q_emb.tolist(), limit=fetch_k,
+                   output_fields=output_fields, search_params={"metric_type":"IP","params":{}})
+        if domain_name == "Legal_Contracts" and legal_contract_id:
+            skw["filter"] = build_contract_filter_expr(legal_contract_id)
+        hits     = client.search(**skw)
+        hit_list = hits[0] if (hits and isinstance(hits[0],(list,tuple))) else hits
+        seen, retrieved = set(), []
+        for hit in hit_list:
+            entity   = (hit.get("entity",{}) or hit) if isinstance(hit,dict) else (getattr(hit,"entity",{}) or {})
+            distance = hit.get("distance",0.0)        if isinstance(hit,dict) else getattr(hit,"distance",0.0)
+            text = entity.get("text","")
+            if text and text not in seen:
+                item = {"text":text,"score":float(distance),"retrieval_type":"dense"}
+                if domain_name == "Legal_Contracts":
+                    item.update({"contract_id":entity.get("contract_id"),
+                                 "source_doc_id":entity.get("source_doc_id"),
+                                 "source_hash":entity.get("source_hash")})
+                retrieved.append(item); seen.add(text)
+
+    if not retrieved: return []
+
+    # Rerank → trim to top_k
+    if ENABLE_RERANKING: retrieved = rerank_documents(eff_q, retrieved, top_k)
+    else:                retrieved = retrieved[:top_k]
+
+    # Repack (reorder for LLM attention)
+    if ENABLE_REPACKING: retrieved = repack_documents(retrieved, REPACK_STRATEGY)
+
+    # Summarize / compress context
+    if ENABLE_SUMMARIZATION:
+        orig       = retrieved
+        summarized = summarize_docs(eff_q, retrieved, em=em, llm_client=llm)
+        if not summarized: return orig
+        scores_list = [d.get("rerank_score",d.get("score",0.0)) for d in retrieved if isinstance(d,dict)]
+        avg = float(np.mean(scores_list)) if scores_list else 1.0
+        mx  = float(max(scores_list))     if scores_list else avg
+        rtype   = retrieved[0].get("reranker_type")  if retrieved and isinstance(retrieved[0],dict) else None
+        rettype = retrieved[0].get("retrieval_type") if retrieved and isinstance(retrieved[0],dict) else None
+        # Preserve Legal metadata from the first (highest-relevance) source chunk
+        smeta = {}
+        if domain_name == "Legal_Contracts" and retrieved and isinstance(retrieved[0],dict):
+            smeta = {k: retrieved[0].get(k) for k in ("contract_id","source_doc_id","source_hash")}
+        retrieved = [{"text":s,"score":avg,"rerank_score":mx,"reranker_type":rtype,
+                      "retrieval_type":rettype,"summarized":True,"summary_type":SUMMARIZATION_TYPE,**smeta}
+                     for s in summarized if s and str(s).strip()]
+        if not retrieved: return orig
+
+    return retrieved
+
+
+# ── Prompt / generation ────────────────────────────────────────────────────────
+
+def _build_prompt(context, question, strategy="short"):
+    context = _sanitize(context); question = _sanitize(question)
+    if strategy == "short":
+        return f"Answer the question using the provided context.\n\nContext:\n{context}\n\nQuestion:\n{question}".strip()
+    elif strategy == "long":
+        return ("You are a chatbot providing answers to user queries. Use the context documents to answer the question.\n"
+                'If the documents do not provide enough information, say "The documents are missing some of the information required to answer the question."\n'
+                f"Do not use external knowledge. Do not make up an answer.\n\nContext Documents:\n{context}\n\nQuestion: {question}").strip()
+    elif strategy == "long_cot":
+        return ("You are a chatbot providing answers to user queries. Use the context documents to answer the question.\n"
+                'If the documents do not provide enough information, say "The documents are missing some of the information required to answer the question."\n'
+                f"Do not use external knowledge. Think step by step and quote documents when necessary.\n\nContext Documents:\n{context}\n\nQuestion: {question}").strip()
+    raise ValueError(f"Unknown PROMPT_STRATEGY: {strategy}")
+
+def ask_rag(context, question, llm_client, strategy=None):
+    strategy = strategy or PROMPT_STRATEGY
+    resp = llm_client.chat.completions.create(
+        model=MODEL_NAME,
+        messages=[{"role":"system","content":"You are a helpful RAG assistant"},
+                  {"role":"user","content":_build_prompt(context, question, strategy)}],
+        temperature=0.3,
+    )
+    return _safe_message_content(resp)
+
+
+print("Pipeline functions defined.")
+
+# ── Cell 6: Initialize LLM client ─────────────────────────────────────────────
+llm_client = get_llm_client()
+print(f"LLM client ready.  Provider: {LLM_PROVIDER}")
+
+
+# ── Cell 7: HF filesystem + embedding model loading ───────────────────────────
+# NOTE: _hf_fs must be initialized here before Cell 8 calls hf_path_exists()
+from sentence_transformers import SentenceTransformer
+import torch
+from huggingface_hub import HfFileSystem
+
+device   = "cuda" if torch.cuda.is_available() else "cpu"
+_hf_fs   = HfFileSystem()
+print(f"Device: {device}  |  HfFileSystem ready")
+
+# Determine which embedding types to load
+if DOMAIN_EMBEDDING_RECOMMENDATION:
+    models_to_load = sorted(set(DOMAIN_EMBEDDING_RECOMMENDATION.values()))
+    print(f"Domain-aware mode — loading: {models_to_load}")
+else:
+    # Legacy single-model mode: always load EMBEDDING_TYPE
+    models_to_load = [EMBEDDING_TYPE]
+    print(f"Single-model mode — loading: {models_to_load}")
+
+loaded_embedding_models = {}
+for emb_type in models_to_load:
+    if emb_type not in EMBED_MODELS:
+        raise ValueError(f"Unknown embedding type: {emb_type}.  Available: {list(EMBED_MODELS.keys())}")
+    model_name = EMBED_MODELS[emb_type]
+    print(f"  Loading {emb_type}: {model_name}")
+    m   = SentenceTransformer(model_name, device=device)
+    dim = m.get_sentence_embedding_dimension() if hasattr(m, "get_sentence_embedding_dimension") else getattr(m, "get_embedding_dimension", lambda: "?")()
+    loaded_embedding_models[emb_type] = m
+    print(f"    OK — dim={dim}")
+
+if not loaded_embedding_models:
+    raise RuntimeError("No embedding models were loaded. Check DOWNLOAD_MODE and EMBED_MODELS.")
+
+# Fallback single embed_model used by RECOMP summarization
+embed_model = loaded_embedding_models.get(EMBEDDING_TYPE) or next(iter(loaded_embedding_models.values()))
+print(f"\nAll embedding models ready.  Fallback embed_model: {EMBEDDING_TYPE}")
+
+# ── Cell 8: Download Milvus DBs + Legal mapping from HuggingFace ──────────────
+# Uses download_indexes() from Cell 5 which mirrors the Advanced notebook logic:
+# preferred path: BUCKET_PREFIX/{embedding_type}_{index_version}/{domain}.db
+# fallback path:  BUCKET_PREFIX/{embedding_type}/{domain}.db
+os.makedirs(MILVUS_DIR, exist_ok=True)
+
+def download_indexes():
+    """Download all domain DBs using domain-aware embedding types."""
+    report = []
+    for domain_name in DOMAIN_NAMES:
+        embedding_type = get_embedding_type_for_domain(domain_name)
+        local_target   = get_db_path(domain_name, embedding_type=embedding_type, index_version=INDEX_VERSION)
+        os.makedirs(os.path.dirname(local_target), exist_ok=True)
+
+        preferred_remote = f"{BUCKET_PREFIX}/{get_index_folder(embedding_type, INDEX_VERSION)}/{domain_name}.db"
+        fallback_remote  = f"{BUCKET_PREFIX}/{embedding_type}/{domain_name}.db"
+
+        # Remove stale file before re-download
+        if os.path.exists(local_target):
+            if os.path.isdir(local_target): shutil.rmtree(local_target)
+            else:                            os.remove(local_target)
+
+        selected_remote, source_type = None, None
+        if hf_path_exists(preferred_remote):
+            selected_remote = preferred_remote
+            source_type     = get_index_folder(embedding_type, INDEX_VERSION)
+        elif hf_path_exists(fallback_remote):
+            selected_remote = fallback_remote
+            source_type     = embedding_type
+
+        if selected_remote is None:
+            print(f"  MISSING: {domain_name} ({embedding_type})")
+            report.append({"domain":domain_name,"status":"missing","source_type":None,"local_target":local_target})
+            continue
+
+        print(f"  Downloading: {domain_name}  [{source_type}]")
+        try:
+            _hf_fs.get(selected_remote, local_target, recursive=True)
+            ok = os.path.exists(local_target) and os.path.getsize(local_target) > 0
+            status = "downloaded" if ok else "empty"
+            print(f"    {'OK' if ok else 'EMPTY'}: {local_target}")
+            report.append({"domain":domain_name,"status":status,"source_type":source_type,"local_target":local_target})
+        except Exception as e:
+            print(f"    FAILED: {e}")
+            report.append({"domain":domain_name,"status":"failed","source_type":source_type,"local_target":local_target,"error":str(e)})
+    return report
+
+print("Downloading domain DBs...")
+dl_report = download_indexes()
+
+# ── Download Legal sample→contract mapping ────────────────────────────────────
+legal_emb    = get_embedding_type_for_domain("Legal_Contracts")
+legal_folder = get_index_folder(legal_emb, INDEX_VERSION)
+remote_mapping = f"{BUCKET_PREFIX}/{legal_folder}/legal_sample_to_contract_id.json"
+local_mapping  = os.path.join(MILVUS_DIR, legal_folder, "legal_sample_to_contract_id.json")
+os.makedirs(os.path.dirname(local_mapping), exist_ok=True)
+print(f"\nDownloading Legal mapping: {remote_mapping}")
+try:
+    _hf_fs.get(remote_mapping, local_mapping)
+    if os.path.exists(local_mapping) and os.path.getsize(local_mapping) > 0:
+        print(f"  OK: {local_mapping}")
+    else:
+        print("  WARNING: Legal mapping download failed or empty")
+except Exception as e:
+    print(f"  WARNING: Could not download Legal mapping: {e}")
+
+# ── Sanity check ──────────────────────────────────────────────────────────────
+print("\nSanity check:")
+for domain in DOMAIN_NAMES:
+    p = get_db_path(domain, get_embedding_type_for_domain(domain), INDEX_VERSION)
+    print(f"  {'OK' if os.path.exists(p) else 'MISSING'}: {p}")
+
+
+    # ── Cell 9: Open Milvus clients + BM25 indexes + Legal mapping ────────────────
+def load_milvus_clients():
+    global milvus_clients, LEGAL_SAMPLE_TO_CONTRACT_ID
+    milvus_clients = {}
+    for domain_name in DOMAIN_NAMES:
+        embedding_type = get_embedding_type_for_domain(domain_name)
+        db_path        = get_db_path(domain_name, embedding_type=embedding_type, index_version=INDEX_VERSION)
+        col            = domain_name.lower()
+        if not os.path.exists(db_path):
+            print(f"  DB not found, skipping: {db_path}")
+            continue
+        try:
+            client = MilvusClient(db_path)
+            if not client.has_collection(col):
+                print(f"  Collection missing in {db_path}, skipping")
+                continue
+            client.load_collection(col)
+            stats = client.get_collection_stats(col)
+            rows  = int(stats.get("row_count", 0))
+            milvus_clients[domain_name] = client
+            print(f"  {domain_name}: {rows:,} rows  [{embedding_type}]")
+        except Exception as e:
+            print(f"  Failed to open {domain_name}: {e}")
+    print(f"\nLoaded {len(milvus_clients)} domain clients: {list(milvus_clients.keys())}")
+
+    # Legal sample→contract mapping
+    LEGAL_SAMPLE_TO_CONTRACT_ID = load_legal_sample_to_contract_mapping()
+
+load_milvus_clients()
+
+# Build BM25 indexes (stores contract_ids for Legal)
+build_all_bm25_indexes(milvus_clients)
+print("BM25 indexes ready.")
+
+# ── Cell 10: Load RAGBench (test split only) + sample catalogue ───────────────
+DATASET_BY_DOMAIN = {
+    "Bio_Medical":       ["covidqa", "pubmedqa"],
+    "General_Knowledge": ["expertqa", "hagrid", "hotpotqa", "msmarco"],
+    "Customer_Support":  ["delucionqa", "emanual", "techqa"],
+    "Finance":           ["finqa", "tatqa"],
+    "Legal_Contracts":   ["cuad"],
+}
+
+# sample_store[domain][dataset] = list of row dicts from the test split
+sample_store = {}
+
+def load_ragbench(domains=None):
+    global ragbench_by_domain, sample_store
+    domains = domains or list(DATASET_BY_DOMAIN.keys())
+    for domain in domains:
+        ragbench_by_domain[domain] = {}
+        sample_store[domain]       = {}
+        for ds_name in DATASET_BY_DOMAIN.get(domain, []):
+            try:
+                ds = load_dataset("rungalileo/ragbench", ds_name)
+                ragbench_by_domain[domain][ds_name] = ds
+                if "test" not in ds:
+                    print(f"  WARNING: no 'test' split for {domain}/{ds_name}, skipping")
+                    continue
+                rows = []
+                for idx, row in enumerate(ds["test"]):
+                    # For Legal_Contracts resolve contract_id from mapping
+                    contract_id = None
+                    if domain == "Legal_Contracts":
+                        try:   contract_id = get_contract_id_for_legal_sample(idx)
+                        except Exception: pass
+                    rows.append({
+                        "idx":         idx,
+                        "question":    row.get("question", ""),
+                        "response":    row.get("response", ""),
+                        "documents":   row.get("documents", []),
+                        "contract_id": contract_id,
+                        "gold_relevance":    row.get("relevance_score"),
+                        "gold_utilization":  row.get("utilization_score"),
+                        "gold_completeness": row.get("completeness_score"),
+                        "gold_adherence":    row.get("adherence_score"),
+                    })
+                sample_store[domain][ds_name] = rows
+                print(f"  Loaded: {domain}/{ds_name}  test rows={len(rows)}")
+            except Exception as e:
+                print(f"  Failed: {domain}/{ds_name}: {e}")
+    print(f"\nRAGBench loaded (test only). Domains: {list(sample_store.keys())}")
+
+load_ragbench()
+
+
+# ── Helpers for cascading dropdowns ───────────────────────────────────────────
+
+def get_datasets_for_domain(domain):
+    return list(sample_store.get(domain, {}).keys())
+
+def get_sample_ids_for_dataset(domain, dataset):
+    """
+    Returns label strings for the Sample ID dropdown.
+    For Legal_Contracts uses 'Contract ID' wording and shows contract hash.
+    """
+    rows   = sample_store.get(domain, {}).get(dataset, [])
+    is_legal = (domain == "Legal_Contracts")
+    labels = []
+    for r in rows:
+        q   = r["question"]
+        cid = r.get("contract_id")
+        if is_legal and cid:
+            prefix = f"Contract {str(cid)[:8]}…  |  idx={r['idx']}  –  "
+        else:
+            prefix = f"{r['idx']}  –  "
+        labels.append(f"{prefix}{q[:70]}{'…' if len(q)>70 else ''}")
+    return labels
+
+def get_row_by_label(domain, dataset, label):
+    """Retrieve a stored row dict from a label string."""
+    if not label: return None
+    rows     = sample_store.get(domain, {}).get(dataset, [])
+    is_legal = (domain == "Legal_Contracts")
+    # Legal labels: "Contract <hash8>…  |  idx=N  –  ..."
+    # Regular labels: "N  –  ..."
+    if is_legal:
+        m = re.search(r"idx=(\d+)", label)
+        try:   idx = int(m.group(1)) if m else int(label.split("–")[0].strip())
+        except ValueError: return None
+    else:
+        try:   idx = int(label.split("–")[0].strip())
+        except ValueError: return None
+    return next((r for r in rows if r["idx"] == idx), None)
+
+print("Sample catalogue ready.")
+
+# ── Cell 11: Evaluation helpers + all Gradio handlers ─────────────────────────
+
+# ── Judge / evaluation ─────────────────────────────────────────────────────────
+
+def build_keyed_response(answer):
+    return {f"r_{i}": s for i, s in enumerate(split_into_sentences(answer))}
+
+def build_sentence_keyed_docs(retrieved_docs):
+    keyed = {}
+    for di, doc in enumerate(retrieved_docs):
+        text = doc.get("text","") if isinstance(doc, dict) else doc
+        for si, s in enumerate(split_into_sentences(text)):
+            keyed[f"{di}_{si}"] = s
+    return keyed
+
+def build_evaluation_prompt(documents_text, question, answer_text):
+    return f"""Evaluate the RAG response using the provided documents.
+
+Documents (sentence-keyed):
+{documents_text}
+
+Question:
+{question}
+
+Response (sentence-keyed):
+{answer_text}
+
+Return ONLY valid JSON:
+{{
+  "overall_supported": true,
+  "all_relevant_sentence_keys": ["0_0"],
+  "all_utilized_sentence_keys": ["0_0"],
+  "sentence_support_information": [
+    {{"response_sentence_key": "r_0", "supporting_sentence_keys": ["0_0"], "fully_supported": true}}
+  ]
+}}
+Rules: document keys look like 0_0; response keys like r_0. Return only JSON.""".strip()
+
+def ask_judge(prompt, llm_client, judge_model, max_retries=5):
+    last_error = None
+    for attempt in range(max_retries):
+        try:
+            resp = llm_client.chat.completions.create(
+                model=judge_model,
+                messages=[
+                    {"role":"system","content":"You are a strict RAG evaluation judge. Return ONLY valid JSON. No markdown. No <think> tags."},
+                    {"role":"user","content":_sanitize(prompt)},
+                ],
+                temperature=0.0, max_tokens=3000,
+            )
+            return _safe_message_content(resp)
+        except Exception as e:
+            last_error = e; msg = str(e)
+            wait = 2**attempt
+            if "429" in msg or "rate_limit" in msg:
+                m = re.search(r"try again in ([\\d.]+)s", msg)
+                if m: wait = float(m.group(1))
+            elif not any(x in msg for x in ["503","502","504","over capacity","gateway"]):
+                raise
+            time.sleep(wait + random.uniform(0.1, 0.5))
+    raise RuntimeError(f"Judge failed after {max_retries} retries: {last_error}")
+
+def parse_judge_json(raw):
+    if not raw: raise ValueError("Judge output empty")
+    cleaned = re.sub(r"<think>.*?</think>","",str(raw),flags=re.DOTALL).strip()
+    cleaned = cleaned.replace("```json","").replace("```","").strip()
+    s, e = cleaned.find("{"), cleaned.rfind("}")
+    if s == -1 or e == -1: raise ValueError(f"No JSON: {cleaned[:300]}")
+    cleaned = cleaned[s:e+1]
+    cleaned = re.sub(r"}\s*{","}, {",cleaned)
+    cleaned = re.sub(r",\s*([}\]])",r"\1",cleaned)
+    return json.loads(cleaned)
+
+def evaluate_ragbench_json(judge_json, keyed_docs):
+    vk  = set(keyed_docs.keys())
+    rel = set(judge_json.get("all_relevant_sentence_keys", [])) & vk
+    utl = set(judge_json.get("all_utilized_sentence_keys", [])) & vk
+    ovl = rel & utl; n = len(vk)
+    return {
+        "adherence_score":    int(bool(judge_json.get("overall_supported", False))),
+        "hallucination_flag": 1 - int(bool(judge_json.get("overall_supported", False))),
+        "relevance_score":    float(np.clip(len(rel)/n        if n   else 0, 0, 1)),
+        "utilization_score":  float(np.clip(len(utl)/n        if n   else 0, 0, 1)),
+        "completeness_score": float(np.clip(len(ovl)/len(rel) if rel else 0, 0, 1)),
+    }
+
+
+# ── Source badge helper ────────────────────────────────────────────────────────
+
+def _source_badge(source, model, extra=None):
+    parts = [f"[Source: {source}  |  model: {model}"]
+    if extra: parts += [f"  |  {k}: {v}" for k, v in extra.items()]
+    parts.append("]")
+    return "".join(parts)
+
+
+# ── DB status helper ───────────────────────────────────────────────────────────
+
+def db_status_md():
+    if not milvus_clients:
+        return ("> **No vector DBs loaded.**  Re-run Cell 8 (download) then Cell 9 (open), then re-run Cell 12.")
+    rows = []
+    for d in sorted(milvus_clients.keys()):
+        emb = get_embedding_type_for_domain(d)
+        rows.append(f"`{d}` ({emb})")
+    return f"> **Loaded domains ({len(milvus_clients)}):** {', '.join(rows)}"
+
+
+# ── Config applier ─────────────────────────────────────────────────────────────
+
+def apply_config(llm_choice, embed_choice,
+                 enable_hybrid, enable_hyde, enable_reranking, reranker_type,
+                 enable_rrf, rrf_k,
+                 enable_repacking, repack_strategy,
+                 enable_summarization, summarization_type,
+                 prompt_strategy, hybrid_alpha, top_k,
+                 enable_query_classification, enable_query_rewriting, enable_query_decomp):
+    global MODEL_NAME, EMBEDDING_TYPE, embed_model
+    global ENABLE_HYBRID, ENABLE_HYDE, ENABLE_RERANKING, RERANKER_TYPE
+    global ENABLE_RRF, RRF_K
+    global ENABLE_REPACKING, REPACK_STRATEGY, ENABLE_SUMMARIZATION, SUMMARIZATION_TYPE
+    global PROMPT_STRATEGY, HYBRID_ALPHA
+    global ENABLE_QUERY_CLASSIFICATION, ENABLE_QUERY_REWRITING, ENABLE_QUERY_DECOMPOSITION
+
+    MODEL_NAME                  = llm_choice
+    ENABLE_HYBRID               = enable_hybrid
+    ENABLE_HYDE                 = enable_hyde
+    ENABLE_RERANKING            = enable_reranking
+    RERANKER_TYPE               = reranker_type
+    ENABLE_RRF                  = enable_rrf
+    RRF_K                       = int(rrf_k)
+    ENABLE_REPACKING            = enable_repacking
+    REPACK_STRATEGY             = repack_strategy
+    ENABLE_SUMMARIZATION        = enable_summarization
+    SUMMARIZATION_TYPE          = summarization_type
+    PROMPT_STRATEGY             = prompt_strategy
+    HYBRID_ALPHA                = float(hybrid_alpha)
+    ENABLE_QUERY_CLASSIFICATION = enable_query_classification
+    ENABLE_QUERY_REWRITING      = enable_query_rewriting
+    ENABLE_QUERY_DECOMPOSITION  = enable_query_decomp
+
+    # Update single fallback embed_model if user changes embedding choice
+    if embed_choice != EMBEDDING_TYPE:
+        EMBEDDING_TYPE = embed_choice
+        if embed_choice in loaded_embedding_models:
+            embed_model = loaded_embedding_models[embed_choice]
+        else:
+            print(f"Embedding type '{embed_choice}' not preloaded; loading now...")
+            embed_model = SentenceTransformer(EMBED_MODELS[embed_choice], device=device)
+            loaded_embedding_models[embed_choice] = embed_model
+
+
+# ── Cascading dropdown callbacks ───────────────────────────────────────────────
+
+_NONE_DOMAIN = "None (direct LLM, no retrieval)"
+
+def on_domain_change(domain):
+    if domain == _NONE_DOMAIN:
+        return gr.update(choices=[], value=None), gr.update(choices=[], value=None), gr.update()
+    datasets   = get_datasets_for_domain(domain)
+    ds         = datasets[0] if datasets else None
+    sample_ids = get_sample_ids_for_dataset(domain, ds) if ds else []
+    label_text = "Contract ID  (contract hash | idx – question preview)" if domain == "Legal_Contracts" else "Sample ID  (idx – question preview)"
+    return (
+        gr.update(choices=datasets,   value=ds),
+        gr.update(choices=sample_ids, value=None, label=label_text),
+        gr.update(value=""),
+    )
+
+def on_dataset_change(domain, dataset):
+    if domain == _NONE_DOMAIN or not dataset:
+        return gr.update(choices=[], value=None), gr.update(value="")
+    sample_ids = get_sample_ids_for_dataset(domain, dataset)
+    label_text = "Contract ID  (contract hash | idx – question preview)" if domain == "Legal_Contracts" else "Sample ID  (idx – question preview)"
+    return gr.update(choices=sample_ids, value=None, label=label_text), gr.update(value="")
+
+def on_sample_select(domain, dataset, label):
+    if domain == _NONE_DOMAIN or not label: return gr.update()
+    row = get_row_by_label(domain, dataset, label)
+    if row is None: return gr.update()
+    return gr.update(value=row["question"])
+
+
+# ── Chunk display helpers ──────────────────────────────────────────────────────
+
+def _format_chunks(docs, title="Retrieved"):
+    if not docs: return f"_No documents for {title}._"
+    parts = []
+    for i, doc in enumerate(docs):
+        if isinstance(doc, dict):
+            text  = doc.get("text", str(doc))
+            score = doc.get("rerank_score", doc.get("score", 0.0))
+            tags  = []
+            if doc.get("summarized"):    tags.append(f"summarized/{doc.get('summary_type','')}")
+            if doc.get("reranker_type"): tags.append(f"reranked/{doc.get('reranker_type','')}")
+            if doc.get("retrieval_type"): tags.append(doc.get("retrieval_type",""))
+            if ENABLE_HYBRID and ENABLE_RRF:
+                tags.append(f"RRF score={score:.4f}")
+            elif ENABLE_HYBRID:
+                tags.append(f"d={doc.get('dense_score',0):.3f} b={doc.get('bm25_score',0):.3f}")
+            if doc.get("contract_id"): tags.append(f"contract={str(doc.get('contract_id',''))[:8]}…")
+            tag_str = f"  `{'  |  '.join(tags)}`" if tags else ""
+        else:
+            text, score, tag_str = str(doc), 0.0, ""
+        parts.append(f"**{title} Chunk {i+1}** — score: `{score:.4f}`{tag_str}\n\n{text}")
+    return "\n\n---\n\n".join(parts)
+
+def _format_gt_docs(doc_list):
+    if not doc_list: return "_No ground-truth documents stored for this sample._"
+    parts = []
+    for i, text in enumerate(doc_list):
+        parts.append(f"**GT Doc {i+1}**\n\n{str(text)}")
+    return "\n\n---\n\n".join(parts)
+
+
+# ── Main run handler ───────────────────────────────────────────────────────────
+
+def run_query(
+    query, domain,
+    dataset_sel, sample_label,
+    llm_choice, judge_llm_choice, embed_choice,
+    enable_hybrid, enable_hyde, enable_reranking, reranker_type,
+    enable_rrf, rrf_k,
+    enable_repacking, repack_strategy,
+    enable_summarization, summarization_type,
+    prompt_strategy, hybrid_alpha, top_k,
+    enable_query_classification, enable_query_rewriting, enable_query_decomp,
+    run_judge,
+):
+    query = _sanitize(query)
+    if not query.strip():
+        return ("Please enter a query.",) + ("",)*4
+    if llm_client is None:
+        return ("LLM client not initialised. Re-run Cell 6 then Cell 12.",) + ("",)*4
+
+    apply_config(
+        llm_choice, embed_choice,
+        enable_hybrid, enable_hyde, enable_reranking, reranker_type,
+        enable_rrf, rrf_k,
+        enable_repacking, repack_strategy,
+        enable_summarization, summarization_type,
+        prompt_strategy, float(hybrid_alpha), int(top_k),
+        enable_query_classification, enable_query_rewriting, enable_query_decomp,
+    )
+
+    # ── Domain = None → direct LLM ──────────────────────────────���────────────
+    if domain == _NONE_DOMAIN or not domain:
+        try:
+            direct_ans = _safe_message_content(llm_client.chat.completions.create(
+                model=MODEL_NAME,
+                messages=[{"role":"system","content":"You are a helpful assistant."},
+                          {"role":"user","content":query}],
+                temperature=0.3, max_tokens=800,
+            ))
+        except Exception as e: direct_ans = f"Direct LLM error: {e}"
+        badge = _source_badge("Direct LLM (no retrieval)", MODEL_NAME)
+        note  = "_[Domain = None — answered directly by LLM without vector DB retrieval]_"
+        return f"{badge}\n\n{direct_ans}", note, note, note, note
+
+    # ── Guard ──────────────────────────────────────────────────────────────────
+    if not milvus_clients:
+        return ("No vector DBs loaded. Re-run Cell 8 then Cell 9, then re-run Cell 12.",) + ("",)*4
+    if domain not in milvus_clients:
+        return (f"Domain '{domain}' not loaded. Loaded: {list(milvus_clients.keys())}",) + ("",)*4
+
+    # ── Query Classification ──────────────────────────────────────────────────
+    route = classify_query(query, domain_name=domain)
+    if route == "LLM":
+        try:
+            direct_ans = _safe_message_content(llm_client.chat.completions.create(
+                model=MODEL_NAME,
+                messages=[{"role":"system","content":"You are a concise factual assistant."},
+                          {"role":"user","content":query}],
+                temperature=0.2, max_tokens=500,
+            ))
+        except Exception as e: direct_ans = f"Direct LLM error: {e}"
+        badge = _source_badge("Direct LLM", MODEL_NAME)
+        note  = "_[Query Classifier routed to direct LLM — no retrieval]_"
+        return f"{badge}\n\n{direct_ans}", note, note, note, note
+
+    # ── Query Rewriting + Decomposition ──────────────────────────────────────
+    rewritten  = rewrite_query(query, domain, llm_client) if ENABLE_QUERY_REWRITING else query
+    subqueries = decompose_query(rewritten, llm_client, domain=domain) if ENABLE_QUERY_DECOMPOSITION else [rewritten]
+
+    # ── Resolve row + contract_id for Legal ───────────────────────────────────
+    row = get_row_by_label(domain, dataset_sel, sample_label) if sample_label else None
+    if row is None:
+        for ds_name, rows in sample_store.get(domain, {}).items():
+            match = next((r for r in rows if r["question"].strip().lower() == query.strip().lower()), None)
+            if match: row = match; break
+
+    legal_contract_id = None
+    legal_sample_id   = None
+    if domain == "Legal_Contracts" and row is not None:
+        legal_contract_id = row.get("contract_id")
+        legal_sample_id   = row.get("idx")
+
+    # ── Retrieve + Generate ───────────────────────────────────────────────────
+    all_retrieved, all_answers = [], []
+    for sq in subqueries:
+        try:
+            docs = retrieve(sq, domain,
+                            llm_client=llm_client, top_k=int(top_k),
+                            sample_id=legal_sample_id, contract_id=legal_contract_id)
+        except Exception as e:
+            return (f"Retrieval error: {e}",) + ("",)*4
+        if not docs: continue
+        all_retrieved.extend(docs)
+        ctx = _sanitize("\n\n".join(d.get("text","") if isinstance(d,dict) else d for d in docs))
+        sq  = _sanitize(sq)
+        try:
+            all_answers.append(ask_rag(ctx, sq, llm_client, strategy=PROMPT_STRATEGY))
+        except Exception as e:
+            return (f"Generation error: {e}",) + ("",)*4
+
+    if not all_retrieved:
+        return ("No documents retrieved.",) + ("",)*4
+
+    raw_answer = "\n\n".join(all_answers)
+
+    # ── Source badge ──────────────────────────────────────────────────────────
+    active = {"prompt": PROMPT_STRATEGY, "chunks": len(all_retrieved)}
+    if ENABLE_HYBRID:
+        active["hybrid"] = f"RRF(k={RRF_K})" if ENABLE_RRF else f"alpha={HYBRID_ALPHA}"
+    if ENABLE_HYDE:          active["hyde"]   = "on"
+    if ENABLE_RERANKING:     active["rerank"] = RERANKER_TYPE
+    if ENABLE_SUMMARIZATION: active["summ"]   = SUMMARIZATION_TYPE
+    if ENABLE_REPACKING:     active["repack"] = REPACK_STRATEGY
+    if len(subqueries) > 1:  active["subq"]   = len(subqueries)
+    if legal_contract_id:    active["contract"] = str(legal_contract_id)[:8] + "…"
+    rag_response = f"{_source_badge('RAG', MODEL_NAME, extra=active)}\n\n{raw_answer}"
+
+    ground_truth = row["response"]  if row else "_(no matching sample found)_"
+    gt_docs_md   = _format_gt_docs(row["documents"] if row else [])
+    rag_docs_md  = _format_chunks(all_retrieved, title="RAG")
+
+    # ── Judge evaluation ──────────────────────────────────────────────────────
+    metrics_md = "_Judge evaluation not requested._"
+    if run_judge:
+        try:
+            keyed_docs   = build_sentence_keyed_docs(all_retrieved)
+            keyed_answer = build_keyed_response(raw_answer)
+            docs_text = _sanitize("\n".join(f"{k}: {v}" for k,v in keyed_docs.items()))
+            ans_text  = _sanitize("\n".join(f"{k}: {v}" for k,v in keyed_answer.items()))
+            raw       = ask_judge(build_evaluation_prompt(docs_text, query, ans_text), llm_client, judge_llm_choice)
+            pred      = evaluate_ragbench_json(parse_judge_json(raw), keyed_docs)
+            gold = {k: row.get(f"gold_{k}") for k in ("relevance","utilization","completeness","adherence")} if row else {}
+            def _f(v): return f"{v:.3f}" if isinstance(v, float) else (str(v) if v is not None else "—")
+            metrics_md = "\n".join([
+                "| Metric | Predicted | Gold |",
+                "|--------|-----------|------|",
+                f"| Relevance    | {_f(pred['relevance_score'])}    | {_f(gold.get('relevance'))} |",
+                f"| Utilization  | {_f(pred['utilization_score'])}  | {_f(gold.get('utilization'))} |",
+                f"| Completeness | {_f(pred['completeness_score'])} | {_f(gold.get('completeness'))} |",
+                f"| Adherence    | {_f(pred['adherence_score'])}    | {_f(gold.get('adherence'))} |",
+                f"| Hallucination| {_f(pred['hallucination_flag'])} | — |",
+            ])
+        except Exception as e:
+            metrics_md = f"Judge error: {e}"
+
+    return ground_truth, rag_response, gt_docs_md, rag_docs_md, metrics_md
+
+
+print("Handlers ready.")
+
+# ── Cell 12: Gradio UI ────────────────────────────────────────────────────────
+
+AVAILABLE_DOMAINS = list(milvus_clients.keys())
+DEFAULT_DOMAIN    = AVAILABLE_DOMAINS[0] if AVAILABLE_DOMAINS else None
+
+_init_datasets = get_datasets_for_domain(DEFAULT_DOMAIN) if DEFAULT_DOMAIN else []
+_init_ds       = _init_datasets[0] if _init_datasets else None
+_init_samples  = get_sample_ids_for_dataset(DEFAULT_DOMAIN, _init_ds) if _init_ds else []
+_legal_first   = DEFAULT_DOMAIN == "Legal_Contracts"
+
+CSS = """
+footer { display: none !important; }
+"""
+
+with gr.Blocks(title="RAG Capstone — Advanced Demo") as demo:
+
+    # ── Header ────────────────────────────────────────────────────────────────
+    gr.Markdown("# 🔍 RAG Capstone — Advanced Interactive Demo")
+    gr.Markdown(
+        f"**Provider:** {LLM_PROVIDER.upper()}  &nbsp;|&nbsp;  "
+        f"**Index:** `{INDEX_VERSION}`  &nbsp;|&nbsp;  "
+        "Type any question, or expand **Sample Selector** to load a test-split example."
+    )
+    gr.Markdown(db_status_md())
+
+    # ══════════════════════════════════════════════════════════════════════════
+    # SECTION 1 — Query + Domain  (always visible)
+    # ══════════════════════════════════════════════════════════════════════════
+    with gr.Row():
+        query_input = gr.Textbox(
+            lines=3,
+            placeholder="Type any question here… or expand Sample Selector below to auto-fill.",
+            label="Query",
+            scale=4,
+        )
+        domain_dd = gr.Dropdown(
+            choices=["None (direct LLM, no retrieval)"] + AVAILABLE_DOMAINS,
+            value="None (direct LLM, no retrieval)" if not AVAILABLE_DOMAINS else DEFAULT_DOMAIN,
+            label="Domain",
+            info="None = direct LLM; pick a domain to run full RAG retrieval",
+            scale=1,
+        )
+
+    # ══════════════════════════════════════════════════════════════════════════
+    # SECTION 2 — Sample Selector  (collapsed, optional)
+    # ═══════════════════════════════════════════════��══════════════════════════
+    with gr.Accordion("📋  Sample Selector  (optional — expand to load a test-split example)", open=False):
+        gr.Markdown(
+            "_Select a sample to auto-fill Query above.  "
+            "For **Legal_Contracts** the dropdown shows Contract ID (hash prefix) instead of plain Sample ID — "
+            "retrieval is automatically scoped to that contract._"
+        )
+        with gr.Row():
+            dataset_dd = gr.Dropdown(
+                choices=_init_datasets, value=_init_ds, label="Dataset", scale=1)
+            sample_dd  = gr.Dropdown(
+                choices=_init_samples, value=None,
+                label="Contract ID  (contract hash | idx – question preview)" if _legal_first else "Sample ID  (idx – question preview)",
+                scale=4)
+
+    # ══════════════════════════════════════════════════════════════════════════
+    # SECTION 3 — Control Panel
+    # ══════════════════════════════════════════════════════════════════════════
+    with gr.Accordion("⚙️  Control Panel", open=False):
+        with gr.Tabs():
+
+            # ── Models ───────────────────────────────────────────────────────
+            with gr.Tab("🤖 Models"):
+                gr.Markdown(
+                    f"**Domain embedding assignment** (chunk_v5_domain_aware):  \n"
+                    + "  \n".join(
+                        [f"- `{d}` → `{get_embedding_type_for_domain(d)}` ({EMBED_MODELS[get_embedding_type_for_domain(d)]})"
+                         for d in DOMAIN_NAMES]
+                    )
+                )
+                with gr.Row():
+                    llm_choice = gr.Dropdown(
+                        choices=LLM_CHOICES, value=LLM_CHOICES[0],
+                        label="Generator LLM", info="Produces the RAG answer")
+                    judge_llm_choice = gr.Dropdown(
+                        choices=LLM_CHOICES,
+                        value=LLM_CHOICES[4] if len(LLM_CHOICES) > 4 else LLM_CHOICES[-1],
+                        label="Judge LLM", info="Used for evaluation scoring")
+                    embed_choice = gr.Dropdown(
+                        choices=EMBEDDING_CHOICES, value=EMBEDDING_TYPE,
+                        label="Fallback Embedding Model",
+                        info="Used only when domain-specific model is unavailable")
+
+            # ── Query Processing ──────────────────────────────────────────────
+            with gr.Tab("🔄 Query Processing"):
+                gr.Markdown("Applied **before** retrieval: Classify → Rewrite → Decompose")
+                with gr.Row():
+                    enable_query_classification = gr.Checkbox(
+                        label="Query Classification", value=False,
+                        info="Route simple factual queries to LLM directly; benchmark domains always use RAG")
+                with gr.Row():
+                    enable_query_rewriting = gr.Checkbox(
+                        label="Query Rewriting", value=False,
+                        info="LLM rewrites the query for better retrieval")
+                    enable_query_decomp = gr.Checkbox(
+                        label="Query Decomposition", value=False,
+                        info="Break multi-part queries into subqueries")
+
+            # ── Retrieval ─────────────────────────────────────────────────────
+            with gr.Tab("🔎 Retrieval"):
+                with gr.Row():
+                    top_k = gr.Slider(minimum=1, maximum=10, step=1, value=3,
+                                      label="Top-K chunks returned")
+                    hybrid_alpha = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.5,
+                                             label="Hybrid Alpha  (1=dense, 0=BM25)  — used only when RRF is OFF")
+                with gr.Row():
+                    enable_hybrid = gr.Checkbox(label="Hybrid Search  (Dense + BM25)", value=True)
+                    enable_hyde   = gr.Checkbox(label="HyDE  (query expansion)", value=False)
+
+            # ── RRF ───────────────────────────────────────────────────────────
+            with gr.Tab("🔀 RRF"):
+                gr.Markdown(
+                    "**Reciprocal Rank Fusion** replaces the weighted alpha fusion inside Hybrid Search.  \n"
+                    "Score formula: `1/(k + rank_dense) + 1/(k + rank_bm25)`  \n"
+                    "Standard literature value for k is **60** — lower k boosts top-ranked docs more aggressively."
+                )
+                with gr.Row():
+                    enable_rrf = gr.Checkbox(
+                        label="Enable RRF  (replaces alpha fusion inside Hybrid Search)",
+                        value=True,
+                        info="RRF is only active when Hybrid Search is also enabled")
+                    rrf_k = gr.Slider(
+                        minimum=1, maximum=200, step=1, value=60,
+                        label="RRF k  (rank smoothing constant)")
+
+            # ── Reranking ─────────────────────────────────────────────────────
+            with gr.Tab("↕️ Reranking"):
+                with gr.Row():
+                    enable_reranking = gr.Checkbox(label="Enable Reranking", value=False)
+                    reranker_type    = gr.Radio(choices=["monot5","tilde"], value="monot5",
+                                               label="Reranker", info="MonoT5: seq2seq | TILDE: cross-encoder")
+
+            # ── Repacking ─────────────────────────────────────────────────────
+            with gr.Tab("📦 Repacking"):
+                with gr.Row():
+                    enable_repacking = gr.Checkbox(label="Enable Repacking", value=False)
+                    repack_strategy  = gr.Radio(choices=["forward","reverse","sides"], value="sides",
+                                               label="Strategy", info="forward | reverse | U-shape sides")
+
+            # ── Summarization ─────────────────────────────────────────────────
+            with gr.Tab("📝 Summarization"):
+                with gr.Row():
+                    enable_summarization = gr.Checkbox(label="Enable Summarization", value=False)
+                    summarization_type   = gr.Radio(choices=["recomp","longllmlingua"], value="recomp",
+                                                   label="Method", info="RECOMP: extractive | LLMLingua: token compression")
+
+            # ── Prompt ────────────────────────────────────────────────────────
+            with gr.Tab("💬 Prompt"):
+                prompt_strategy = gr.Radio(
+                    choices=["short","long","long_cot"], value="short",
+                    label="Prompt Strategy",
+                    info="short: minimal | long: strict no-hallucination | long_cot: step-by-step")
+
+            # ── Judge ─────────────────────────────────────────────────────────
+            with gr.Tab("⚖️ Judge"):
+                run_judge = gr.Checkbox(
+                    label="Run Judge evaluation after generation", value=False,
+                    info="~1 extra LLM call. Gold scores shown only for preloaded samples.")
+                gr.Markdown("_Judge LLM is configured in the **Models** tab._")
+
+    # ── Run button ────────────────────────────────────────────────────────────
+    run_btn = gr.Button("▶  Run Query", variant="primary", size="lg")
+
+    # ══════════════════════════════════════════════════════════════════════════
+    # SECTION 4 — Responses  (Ground Truth LEFT, RAG RIGHT)
+    # ══════════════════════════════════════════════════════════════════════════
+    gr.Markdown("## 💬 Responses")
+    with gr.Row(equal_height=True):
+        gt_out  = gr.Textbox(label="Ground Truth Response", lines=10, interactive=False, scale=1)
+        rag_out = gr.Textbox(label="RAG Response",          lines=10, interactive=False, scale=1)
+
+    # ══════════════════════════════════════════════════════════════════════════
+    # SECTION 5 — Retrieved Documents  (GT LEFT, RAG RIGHT)
+    # ══════════════════════════════════════════════════════════════════════════
+    gr.Markdown("## 📄 Retrieved Documents")
+    with gr.Row(equal_height=True):
+        with gr.Column(scale=1):
+            gr.Markdown("### Ground Truth Documents")
+            gt_docs_out  = gr.Markdown(value="_Select a preloaded sample to see GT documents._")
+        with gr.Column(scale=1):
+            gr.Markdown("### RAG Retrieved Documents")
+            rag_docs_out = gr.Markdown(value="_Run a query to see RAG retrieved chunks._")
+
+    # ══════════════════════════════════════════════════════════════════════════
+    # SECTION 6 — Metrics
+    # ══════════════════════════════════════════════════════════════════════════
+    with gr.Accordion("📊  Metrics  (Gold vs Predicted)", open=False):
+        metrics_out = gr.Markdown(value="_Enable the Judge in the Control Panel and run a query._")
+
+    # ── Cascading sample selector wiring ─────────────────────────────────────
+    domain_dd.change(
+        fn=on_domain_change, inputs=[domain_dd],
+        outputs=[dataset_dd, sample_dd, query_input],
+    )
+    dataset_dd.change(
+        fn=on_dataset_change, inputs=[domain_dd, dataset_dd],
+        outputs=[sample_dd, query_input],
+    )
+    sample_dd.change(
+        fn=on_sample_select, inputs=[domain_dd, dataset_dd, sample_dd],
+        outputs=[query_input],
+    )
+
+    # ── Run wiring ────────────────────────────────────────────────────────────
+    _config_inputs = [
+        llm_choice, judge_llm_choice, embed_choice,
+        enable_hybrid, enable_hyde, enable_reranking, reranker_type,
+        enable_rrf, rrf_k,
+        enable_repacking, repack_strategy,
+        enable_summarization, summarization_type,
+        prompt_strategy, hybrid_alpha, top_k,
+        enable_query_classification, enable_query_rewriting, enable_query_decomp,
+        run_judge,
+    ]
+    _all_inputs  = [query_input, domain_dd, dataset_dd, sample_dd] + _config_inputs
+    _all_outputs = [gt_out, rag_out, gt_docs_out, rag_docs_out, metrics_out]
+
+    run_btn.click(fn=run_query, inputs=_all_inputs, outputs=_all_outputs)
+    query_input.submit(fn=run_query, inputs=_all_inputs, outputs=_all_outputs)
+
+demo.launch(
+    share=True,
+    debug=True,
+    theme=gr.themes.Soft(),
+    css=CSS,
+)
+

requirements.txt

@@ -0,0 +1,15 @@
+numpy
+pandas
+pymilvus
+groq
+openai
+sentence_transformers
+rank_bm25
+scikit-learn
+torch
+transformers
+huggingface_hub
+datasets
+gradio
+llmlingua
+milvus-lite