# ── 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 … | 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 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".*?","",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()}  |  " f"**Index:** `{INDEX_VERSION}`  |  " "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, )