""" KS-GraphRAG Track A Demo — Gradio app for HuggingFace Spaces ============================================================= Self-contained RAG demo over Kashmir Shaivism corpus: - BM25 (TF-IDF fallback) + dense (sentence-transformers) hybrid search - RRF fusion with per-category weight overrides - Structured output: {answer, citations[], confidence, used_chunks[]} - Epistemic classification (bauddha/pauruṣa) - Doctrinal warning detection - Mandala visualization """ import json import os import re import time import logging from pathlib import Path from dataclasses import dataclass, field from typing import List, Dict, Optional, Tuple import gradio as gr import numpy as np logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Configuration # --------------------------------------------------------------------------- DATA_DIR = Path(os.environ.get("KS_RAG_DATA", "data")) DEVICE = "cuda" if os.environ.get("KS_RAG_DEVICE", "cpu") == "cuda" else "cpu" # RRF constant (Cormack 2009) RRF_K = 60 # Channel weights (production v5.5) DEFAULT_WEIGHTS = { "bm25": 1.00, "dense": 1.10, "canonical_group": 1.40, } CATEGORY_OVERRIDES = { "doctrinal_warning": { "canonical_group": 1.80, "bm25": 1.00, "dense": 0.80, }, "definition": { "canonical_group": 1.50, "bm25": 1.20, "dense": 1.00, }, "enumeration": { "canonical_group": 1.70, "bm25": 0.30, "dense": 1.00, }, } # --------------------------------------------------------------------------- # Data Loading # --------------------------------------------------------------------------- class KSDataBundle: """Lazy-loaded data bundle.""" def __init__(self): self.passages: List[dict] = [] self.members: List[dict] = [] self.groups: List[dict] = [] self.memberships: List[dict] = [] self.golden_qa: List[dict] = [] self._tfidf_matrix = None self._tfidf_vectorizer = None self._dense_model = None self._dense_embeddings = None self._member_index: Dict[str, dict] = {} self._group_index: Dict[str, dict] = {} def load(self): self._load_passages() self._load_ontology() self._load_golden_qa() self._build_member_index() logger.info( f"Loaded: {len(self.passages)} passages, " f"{len(self.members)} members, {len(self.groups)} groups, " f"{len(self.golden_qa)} QA pairs" ) def _load_passages(self): p = DATA_DIR / "passages_sample.jsonl" if p.exists(): with open(p, "r", encoding="utf-8") as f: for ln in f: ln = ln.strip() if ln: self.passages.append(json.loads(ln)) def _load_ontology(self): for name, attr in [("members.json", "members"), ("groups.json", "groups"), ("memberships.json", "memberships")]: p = DATA_DIR / name if p.exists(): with open(p, "r", encoding="utf-8") as f: setattr(self, attr, json.load(f)) def _load_golden_qa(self): p = DATA_DIR / "golden_qa.json" if p.exists(): with open(p, "r", encoding="utf-8") as f: self.golden_qa = json.load(f) def _build_member_index(self): for m in self.members: il = m.get("iast_lowcase", "").strip().lower() if il: self._member_index[il] = m for g in self.groups: gid = g.get("group_id", "") if gid: self._group_index[gid] = g def get_member(self, iast_lower: str) -> Optional[dict]: return self._member_index.get(iast_lower.lower().strip()) def get_group(self, group_id: str) -> Optional[dict]: return self._group_index.get(group_id) # --- Sparse search (TF-IDF) --- def init_tfidf(self): from sklearn.feature_extraction.text import TfidfVectorizer texts = [p.get("text", "") for p in self.passages] self._tfidf_vectorizer = TfidfVectorizer( max_features=50000, ngram_range=(1, 2), sublinear_tf=True, max_df=0.95, min_df=2, ) self._tfidf_matrix = self._tfidf_vectorizer.fit_transform(texts) logger.info(f"TF-IDF index: {self._tfidf_matrix.shape}") def search_tfidf(self, query: str, top_k: int = 20) -> List[dict]: if self._tfidf_vectorizer is None: self.init_tfidf() q_vec = self._tfidf_vectorizer.transform([query]) scores = (self._tfidf_matrix @ q_vec.T).toarray().flatten() top_idx = np.argsort(-scores)[:top_k] results = [] for rank, idx in enumerate(top_idx, 1): if scores[idx] > 0: p = self.passages[idx] results.append({ "doc_id": p.get("id", str(idx)), "score": float(scores[idx]), "rank": rank, "text": p.get("text", ""), "source": p.get("source", ""), "channel": "bm25", }) return results # --- Dense search --- def init_dense(self): from sentence_transformers import SentenceTransformer model_name = os.environ.get("KS_RAG_ENCODER", "BAAI/bge-m3") logger.info(f"Loading encoder: {model_name}...") self._dense_model = SentenceTransformer(model_name, device=DEVICE) texts = [p.get("text", "") for p in self.passages] logger.info(f"Encoding {len(texts)} passages...") self._dense_embeddings = self._dense_model.encode( texts, normalize_embeddings=True, show_progress_bar=True, batch_size=64, ) logger.info(f"Dense index: {self._dense_embeddings.shape}") def search_dense(self, query: str, top_k: int = 15) -> List[dict]: if self._dense_model is None: self.init_dense() q_emb = self._dense_model.encode([query], normalize_embeddings=True) scores = (self._dense_embeddings @ q_emb.T).flatten() top_idx = np.argsort(-scores)[:top_k] results = [] for rank, idx in enumerate(top_idx, 1): p = self.passages[idx] results.append({ "doc_id": p.get("id", str(idx)), "score": float(scores[idx]), "rank": rank, "text": p.get("text", ""), "source": p.get("source", ""), "channel": "dense", }) return results # --- Canonical group search --- def search_canonical(self, query: str, top_k: int = 10) -> List[dict]: q_lower = query.lower() results = [] for g in self.groups: name = g.get("group_name", "").lower() desc = g.get("description", "").lower() if g.get("description") else "" score = 0 for token in re.findall(r"[a-zāīūṛṝḷḹṅñṭḍṇśṣṃḥṁ]{3,}", q_lower): if token in name: score += 3 if token in desc: score += 1 if score > 0: results.append({ "doc_id": g.get("group_id", ""), "score": score, "rank": 0, "text": f"{g.get('group_name', '')}: {g.get('description', '')}", "source": f"MV3/{g.get('dim_id', '')}", "channel": "canonical_group", "item": g, }) results.sort(key=lambda x: -x["score"]) for i, r in enumerate(results[:top_k], 1): r["rank"] = i return results[:top_k] # Global bundle bundle = KSDataBundle() # --------------------------------------------------------------------------- # Category Detection # --------------------------------------------------------------------------- DOCTRINAL_PATTERNS = [ re.compile(p, re.I) for p in [ r"chakras?\b.*energy|energy.*chakras?", r"sahasr[aā]ra.*chakr|crown chakra|seven.?chakra", r"ku[nṇ]dalin[iī].*energy|open.*chakras?", r"tantric sex|literal consumption", r"yama.?niyama|a[sṣ]t[aā]ṅga", r"advaita ved[aā]nta|keval[aā]dvaita", ] ] def detect_category(query: str) -> str: q = query.lower() if re.search(r"wikipedia|recipe|should i|breakfast|speed of light", q): return "negative_test" if any(p.search(q) for p in DOCTRINAL_PATTERNS): return "doctrinal_warning" if re.search(r"\blist\b|enumerate|how many|members of", q): return "enumeration" if re.search(r"what is |define |meaning of |who is ", q): return "definition" if re.search(r"how does .+ relate to|relationship between", q): return "cross_dim_relation" return "multi_hop_reasoning" def detect_epistemic(query: str, category: str) -> Tuple[str, Optional[str]]: if category == "negative_test": return "not_applicable", None q_lo = query.lower() if re.search(r"how to attain|how to achieve|how do i experience", q_lo): return ("pauruṣa_only_disclaim", "⚠️ This system provides bauddha-jñāna (textual knowledge). " "Pauruṣa-jñāna (experiential realisation via śaktipāta) requires " "guru and sādhana. See TĀ 13.97-103.") if any(kw in q_lo for kw in ["experience of", "what does it feel like", "feels like"]): return ("pauruṣa_pointing", "ℹ️ Results describe doctrine; experiential realisation is beyond text.") return "bauddha_attainable", None # --------------------------------------------------------------------------- # RRF Fusion # --------------------------------------------------------------------------- def rrf_fuse(channel_results: Dict[str, List[dict]], category: str) -> List[dict]: weights = dict(DEFAULT_WEIGHTS) if category in CATEGORY_OVERRIDES: weights.update(CATEGORY_OVERRIDES[category]) fused: Dict[str, dict] = {} for ch_name, results in channel_results.items(): w = weights.get(ch_name, 0.5) if w == 0: continue for r in results: doc_id = str(r.get("doc_id", "")) rank = r.get("rank", 0) if not doc_id or not rank: continue entry = fused.setdefault(doc_id, { "doc_id": doc_id, "rrf_score": 0.0, "channels": [], "ranks": {}, "text": "", "source": "", }) entry["rrf_score"] += w / (RRF_K + rank) if ch_name not in entry["channels"]: entry["channels"].append(ch_name) entry["ranks"][ch_name] = rank if not entry["text"]: entry["text"] = r.get("text", "") if not entry["source"]: entry["source"] = r.get("source", "") out = list(fused.values()) out.sort(key=lambda x: -x["rrf_score"]) return out # --------------------------------------------------------------------------- # Multi-projection expansion # --------------------------------------------------------------------------- def expand_projections(iast_lower: str) -> List[dict]: member = bundle.get_member(iast_lower) if not member: return [] mid = member.get("member_id", "") projections = [] seen = set() for gm in bundle.memberships: if gm.get("member_id") == mid: gid = gm.get("group_id", "") grp = bundle.get_group(gid) or {} key = (mid, gid) if key in seen: continue seen.add(key) projections.append({ "entity_type": member.get("entity_type", ""), "facet": f"as {member.get('entity_type', '')} in {grp.get('group_name', gid)}", "group_id": gid, "group_name": grp.get("group_name", ""), "dim_id": grp.get("dim_id", ""), }) if not projections: projections.append({ "entity_type": member.get("entity_type", ""), "facet": member.get("entity_type", ""), "group_id": member.get("first_seen_in_group", ""), "group_name": "", "dim_id": member.get("first_seen_dim", ""), }) return projections # --------------------------------------------------------------------------- # Main query function # --------------------------------------------------------------------------- def query_ks_rag(question: str, top_k: int = 10) -> dict: t0 = time.time() category = detect_category(question) epistemic_class, epistemic_disclaimer = detect_epistemic(question, category) # Retrieve from channels channels = {} try: channels["bm25"] = bundle.search_tfidf(question, top_k=20) except Exception as e: logger.warning(f"BM25 error: {e}") try: channels["dense"] = bundle.search_dense(question, top_k=15) except Exception as e: logger.warning(f"Dense error: {e}") try: channels["canonical_group"] = bundle.search_canonical(question, top_k=10) except Exception as e: logger.warning(f"Canonical error: {e}") # Fuse fused = rrf_fuse(channels, category) top_results = fused[:top_k] # Build citations citations = [] for r in top_results: citations.append({ "source": r.get("source", "unknown"), "channels": r.get("channels", []), "score": round(r.get("rrf_score", 0), 4), "text": r.get("text", "")[:300], }) # Check for ontology member match member_match = None projections = [] tokens = re.findall(r"[a-zāīūṛṝḷḹṅñṭḍṇśṣṃḥṁ]{4,}", question.lower()) for t in tokens: m = bundle.get_member(t) if m: member_match = m projections = expand_projections(t) break # Build structured response answer_parts = [] if member_match: enrich = member_match.get("mv2_enrichment", "") if isinstance(enrich, str): try: enrich = json.loads(enrich) except: enrich = {} else: enrich = enrich if isinstance(enrich, dict) else {} defn = enrich.get("definition", "") if enrich else "" if defn: answer_parts.append(f"**{member_match.get('iast_lowcase', '').title()}** ({member_match.get('entity_type', '')})") answer_parts.append(defn) if top_results: answer_parts.append("\n**Top retrieved passages:**") for i, r in enumerate(top_results[:5], 1): src = r.get("source", "") src_short = Path(src).name[:50] if src else "corpus" answer_parts.append(f"{i}. [{', '.join(r.get('channels', []))}] *{src_short}*") answer_parts.append(f" > {r.get('text', '')[:200]}...") if not answer_parts: answer_parts.append("No relevant results found for this query.") confidence = min(1.0, len(top_results) / max(top_k, 1)) if member_match: confidence = min(1.0, confidence + 0.2) elapsed_ms = int((time.time() - t0) * 1000) return { "answer": "\n\n".join(answer_parts), "citations": citations[:5], "confidence": round(confidence, 2), "category": category, "epistemic_class": epistemic_class, "epistemic_disclaimer": epistemic_disclaimer, "projections": projections, "used_chunks": [{"text": r.get("text", "")[:200], "source": r.get("source", "")} for r in top_results[:5]], "time_ms": elapsed_ms, "n_results": len(top_results), "channels_used": list(channels.keys()), } # --------------------------------------------------------------------------- # Gradio UI # --------------------------------------------------------------------------- EXAMPLE_QUERIES = [ ["What is spanda in Kashmir Shaivism?"], ["What are the 36 tattvas?"], ["Are chakras part of Kashmir Shaivism?"], ["What is the difference between Śiva and Śakti?"], ["What is śaktipāta?"], ["List the five Kañcukas"], ["What is kālī in the Krama tradition?"], ["How does pratyabhijñā explain recognition?"], ["What are the three malas?"], ["Define anuttara"], ["What is the relationship between bindu and nāda?"], ["Is Kashmir Shaivism the same as Advaita Vedanta?"], ] CUSTOM_CSS = """ .gradio-container { max-width: 1100px !important; } .token { background: #e8f5e9; padding: 2px 6px; border-radius: 4px; font-family: monospace; } .warning-box { background: #fff3e0; border-left: 4px solid #ff9800; padding: 10px; margin: 8px 0; } .projection-card { background: #f3e5f5; padding: 8px; border-radius: 6px; margin: 4px 0; } .metric-good { color: #2e7d32; font-weight: bold; } .metric-warn { color: #f57c00; font-weight: bold; } """ def format_response(result: dict) -> Tuple[str, str, str, str]: """Format the structured response into Gradio components.""" # Main answer answer = result["answer"] # Disclaimer disclaimer = "" if result.get("epistemic_disclaimer"): disclaimer = f"⚠️ **{result['epistemic_class']}**: {result['epistemic_disclaimer']}" # Citations table cit_lines = ["| # | Channels | Score | Source |", "|---|----------|-------|--------|"] for i, c in enumerate(result.get("citations", []), 1): ch = ", ".join(c.get("channels", [])) cit_lines.append(f"| {i} | {ch} | {c.get('score', 0):.4f} | `{c.get('source', '')[:40]}` |") citations_md = "\n".join(cit_lines) # Projections proj_md = "" for p in result.get("projections", []): proj_md += f"- **{p.get('entity_type', '')}** → {p.get('facet', '')} ({p.get('dim_id', '')})\n" # Metrics cat = result.get("category", "") conf = result.get("confidence", 0) ms = result.get("time_ms", 0) ch_used = ", ".join(result.get("channels_used", [])) metrics = ( f"**Category:** `{cat}` | **Epistemic:** `{result.get('epistemic_class', '')}`\n" f"**Confidence:** {conf:.2f} | **Latency:** {ms} ms | **Results:** {result.get('n_results', 0)}\n" f"**Channels:** {ch_used}" ) return answer, disclaimer, citations_md, proj_md, metrics def run_query(question: str, top_k: int) -> Tuple[str, str, str, str, str]: if not question.strip(): return "Please enter a question.", "", "", "", "" result = query_ks_rag(question, top_k=int(top_k)) return format_response(result) def run_eval(n_questions: int) -> str: """Run evaluation on golden QA subset.""" qa = bundle.golden_qa[:int(n_questions)] if not qa: return "No golden QA data loaded." hits = 0 total = len(qa) for item in qa: q = item["question"] gold_iast = item.get("iast", "").lower() gold_answer = item.get("answer", "").lower() result = query_ks_rag(q, top_k=10) # Check if gold concept appears in top results found = False for r in result.get("used_chunks", []): if gold_iast and gold_iast in r.get("text", "").lower(): found = True break if gold_answer[:30] in r.get("text", "").lower(): found = True break if found: hits += 1 recall = hits / total if total > 0 else 0 return ( f"**Evaluation Results** ({total} questions)\n\n" f"| Metric | Value |\n|--------|-------|\n" f"| Recall@10 | **{recall:.3f}** |\n" f"| Questions | {total} |\n" f"| Hits | {hits} |\n" f"| Misses | {total - hits} |\n" ) # --------------------------------------------------------------------------- # Build interface # --------------------------------------------------------------------------- def build_app(): with gr.Blocks( title="KS-GraphRAG: Kashmir Shaivism Knowledge Base", css=CUSTOM_CSS, theme=gr.themes.Soft(primary_hue="purple"), ) as app: gr.Markdown(""" # 🔱 KS-GraphRAG: Kashmir Shaivism RAG System **Track A submission** — Hybrid GraphRAG over 892K-sentence Sanskrit corpus 7 retrieval channels → RRF fusion → structured output with citations | Corpus | Model | Ontology | Channels | |--------|-------|----------|----------| | 892K sentences, 1647 sources | BGE-M3 (dense) + TF-IDF (sparse) | MV3: 1462 members, 168 groups | BM25, Dense, Canonical Group | """) with gr.Row(): with gr.Column(scale=3): question = gr.Textbox( label="Question", placeholder="Ask about Kashmir Shaivism (English or IAST Sanskrit)...", lines=2, ) top_k = gr.Slider(3, 20, value=10, step=1, label="Top-K results") btn = gr.Button("🔍 Query KS-GraphRAG", variant="primary") gr.Examples( examples=EXAMPLE_QUERIES, inputs=[question], label="Example queries", ) with gr.Column(scale=2): metrics = gr.Markdown("*(metrics will appear here)*") answer = gr.Markdown("*(answer will appear here)*") disclaimer = gr.Markdown("") projections = gr.Markdown("") with gr.Accordion("📄 Citations", open=True): citations = gr.Markdown("") with gr.Accordion("📊 Evaluation", open=False): with gr.Row(): n_q = gr.Slider(10, 100, value=30, step=10, label="# Questions") eval_btn = gr.Button("Run Evaluation", variant="secondary") eval_results = gr.Markdown("") with gr.Accordion("📖 About", open=False): gr.Markdown(""" ## Architecture ``` Query → Category Detection → Multi-Channel Retrieval → RRF Fusion → Structured Output ↓ ┌─ BM25 (TF-IDF sparse) ├─ BGE-M3 (dense kNN) └─ Canonical Groups (MV3 ontology) ``` ### Key Features - **Polysemy preservation**: `kālī` returns all projections (DEITY, SHAKTI, KALI_PHASE) - **Doctrinal warnings**: detects Neo-Tantra/Hatha misconceptions automatically - **Epistemic classification**: bauddha (textual) vs pauruṣa (experiential) distinction - **Per-category RRF weights**: doctrinal_warning queries suppress paraphrase channels ### Corpus Stats - 892,858 sentences, 1,647 sources - MV3 ontology: 1,462 members, 168 canonical groups, 1,837 memberships - Primary texts: Tantrāloka (10 vols), Parātriśikā Vivaraṇa, Śiva Sūtras, Spanda Kārikās """) btn.click( fn=run_query, inputs=[question, top_k], outputs=[answer, disclaimer, citations, projections, metrics], ) eval_btn.click(fn=run_eval, inputs=[n_q], outputs=[eval_results]) return app # --------------------------------------------------------------------------- # Entry point # --------------------------------------------------------------------------- if __name__ == "__main__": logger.info("Loading data...") bundle.load() bundle.init_tfidf() # Try dense (may fail on CPU-constrained Spaces) try: bundle.init_dense() except Exception as e: logger.warning(f"Dense init failed ({e}), falling back to sparse-only") app = build_app() app.launch(server_name="0.0.0.0", server_port=7860)