Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,660 @@
+"""
+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)