Spaces:
Running
Running
| # ── 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()} | " | |
| 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, | |
| ) | |