ks-graphrag / app.py
seg007's picture
Upload app.py with huggingface_hub
1b7622e verified
Raw
History Blame Contribute Delete
24.2 kB
"""
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)