Update app.py

app.py

@@ -3,7 +3,6 @@ import os
 import re
 import gradio as gr
 import torch
-torch.backends.cuda.matmul.allow_tf32 = True
 import faiss
 from dataclasses import dataclass
 from typing import Optional
@@ -11,18 +10,18 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 from sentence_transformers import SentenceTransformer
 from huggingface_hub import login
 
-# ── Auth ───────────────────────────────────────────────────────────────────────
+# ── Auth ─────────────────────────────────────────────
 _token = os.environ.get("HF_TOKEN")
 if _token:
     login(token=_token)
 
-# ── Config ─────────────────────────────────────────────────────────────────────
+# ── Config ───────────────────────────────────────────
 CHECKPOINT   = "selfrag/selfrag_llama2_7b"
 K_PASSAGES   = 5
-MAX_NEW_TOKS = 150
+MAX_NEW_TOKS = 120   # reduced for stability
 W_REL = 1.0; W_SUP = 1.0; W_USE = 0.5
 
-# ── Reflection tokens ──────────────────────────────────────────────────────────
+# ── Reflection Tokens ─────────────────────────────────
 class RT:
     RETRIEVAL    = "[Retrieval]"
     NO_RETRIEVAL = "[No Retrieval]"
@@ -42,307 +41,98 @@ class RT:
 ABSTAIN_PHRASES = [
     "not specified in the input","not provided in the input",
     "not mentioned in the input","no information","cannot be determined",
-    "is not specified","are not specified","not available in the input",
-    "cannot answer","no relevant information","not found in",
-    "does not contain","not in the context","not mentioned in the context",
 ]
 
-# ── Corpus ─────────────────────────────────────────────────────────────────────
+# ── Documents ─────────────────────────────────────────
 DOCUMENTS = [
-    {"id":"doc_00","title":"RAG Overview",
-     "text":"Retrieval-Augmented Generation (RAG) improves factual correctness in large language "
-            "models by combining a neural generator with a non-parametric retrieval system. The model "
-            "is given relevant documents at inference time instead of relying solely on memorised "
-            "parametric knowledge. RAG was introduced by Lewis et al. in 2020."},
-    {"id":"doc_01","title":"Vanilla RAG Limitations",
-     "text":"Vanilla RAG always retrieves a fixed number of top-k documents regardless of whether "
-            "retrieval is actually necessary for the query. This indiscriminate retrieval can introduce "
-            "noise when the documents are not relevant to the question, and the model may produce "
-            "outputs that are inconsistent with retrieved passages."},
-    {"id":"doc_02","title":"SELF-RAG Framework",
-     "text":"SELF-RAG is a framework that trains a single language model to retrieve, generate, and "
-            "critique text through self-reflection. It uses four special reflection tokens: Retrieve "
-            "(yes/no/continue), IsRel (relevant/irrelevant), IsSupport (fully/partially/no support), "
-            "and IsUse (1 to 5). These tokens allow the model to decide adaptively when to retrieve "
-            "and evaluate the quality of its own generation at inference time. SELF-RAG was published "
-            "at ICLR 2024 by Asai et al."},
-    {"id":"doc_03","title":"CRAG — Corrective RAG",
-     "text":"Corrective Retrieval-Augmented Generation (CRAG) addresses the problem of inaccurate or "
-            "irrelevant retrieved documents by adding a lightweight evaluator that grades each retrieved "
-            "document's relevance. When all retrieved documents are judged irrelevant, CRAG triggers a "
-            "corrective web search to obtain better evidence before generation."},
-    {"id":"doc_04","title":"RAGAs Evaluation Framework",
-     "text":"RAGAs is an automated evaluation framework for retrieval-augmented generation systems. "
-            "It provides metrics including faithfulness, which measures whether the generated answer "
-            "is supported by the retrieved context, and answer relevance, which measures whether the "
-            "answer addresses the original question."},
-    {"id":"doc_05","title":"Hallucination in LLMs",
-     "text":"Hallucination in large language models refers to the generation of fluent and confident "
-            "statements that are not supported by evidence or are factually incorrect. RAG reduces "
-            "hallucination by grounding generation in retrieved documents, but does not eliminate it "
-            "entirely when the model ignores the retrieved context."},
-    {"id":"doc_06","title":"Adaptive Retrieval in SELF-RAG",
-     "text":"SELF-RAG uses an adaptive retrieval threshold delta to control how often retrieval is "
-            "triggered. The model predicts a Retrieve token and if the probability of Retrieve=Yes "
-            "exceeds the threshold delta, retrieval is triggered. The default threshold is 0.2."},
-    {"id":"doc_07","title":"Segment-Level Beam Search",
-     "text":"SELF-RAG uses segment-level beam search with beam width B=2 during inference. For each "
-            "retrieved passage, the model generates a candidate scored using IsRel, IsSupport, and "
-            "IsUse token probabilities combined with log-probability of the generated segment. "
-            "Weights: 1.0 for IsRel, 1.0 for IsSupport, 0.5 for IsUse."},
+    {"id":"doc_00","title":"RAG Overview","text":"Retrieval-Augmented Generation improves factual correctness."},
+    {"id":"doc_01","title":"Vanilla RAG Limitations","text":"Vanilla RAG retrieves fixed documents even when unnecessary."},
+    {"id":"doc_02","title":"SELF-RAG Framework","text":"SELF-RAG uses reflection tokens to improve responses."},
+    {"id":"doc_03","title":"CRAG","text":"CRAG adds evaluator for retrieval correction."},
+    {"id":"doc_04","title":"RAGAs","text":"Evaluation framework for RAG systems."},
+    {"id":"doc_05","title":"Hallucination","text":"LLMs may generate unsupported content."},
+    {"id":"doc_06","title":"Adaptive Retrieval","text":"SELF-RAG triggers retrieval dynamically."},
+    {"id":"doc_07","title":"Beam Search","text":"Uses scoring with reflection tokens."},
 ]
 
 @dataclass
 class Passage:
     id: str; title: str; text: str
 
-@dataclass
-class Candidate:
-    text: str; passage: Passage
-    isrel: Optional[str] = None
-    issup: Optional[str] = None
-    isuse: Optional[str] = None
-    score: float = 0.0
-
-# ── Retriever — CPU, loaded at startup ────────────────────────────────────────
+# ── Retriever ─────────────────────────────────────────
 print("Loading retriever...")
-_retriever = SentenceTransformer("all-MiniLM-L6-v2", device="cpu")
-_passages  = [Passage(**d) for d in DOCUMENTS]
-_embs      = _retriever.encode(
-    [f"{p.title} {p.text}" for p in _passages],
-    convert_to_numpy=True, normalize_embeddings=True
-).astype("float32")
-_index = faiss.IndexFlatIP(_embs.shape[1])
-_index.add(_embs)
-print(f"Retriever ready: {_index.ntotal} passages")
-
-def retrieve(query: str):
-    q = _retriever.encode([query], convert_to_numpy=True,
-                          normalize_embeddings=True).astype("float32")
-    _, idxs = _index.search(q, K_PASSAGES)
-    return [_passages[i] for i in idxs[0] if i < len(_passages)]
+retriever = SentenceTransformer("all-MiniLM-L6-v2", device="cpu")
 
-# ── Generator — loaded at module level on CUDA (ZeroGPU emulates outside GPU fn)
-# Per official ZeroGPU docs: place model on cuda at module level.
-# ZeroGPU emulation mode allows this even when no real GPU is present at startup.
-print(f"Loading {CHECKPOINT}...")
-_tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
-_model = AutoModelForCausalLM.from_pretrained(
-    CHECKPOINT,
-    torch_dtype=torch.float16,
-    low_cpu_mem_usage=True,
-)
-_model = AutoModelForCausalLM.from_pretrained(
-    CHECKPOINT,
-    torch_dtype=torch.float16,
-    device_map="auto",
-)   
-_model.eval()
-print("Generator ready.")
-
-# ── GPU-decorated inference ────────────────────────────────────────────────────
-@spaces.GPU
-def generate_candidates(query: str, passage_texts: list) -> list:
-    """Runs inside real ZeroGPU GPU context."""
-    results = []
-    for pt in passage_texts:
-        prompt = (f"### Instruction:\n{query}\n\n"
-                  f"### Input:\n<p>{pt}</p>\n\n"
-                  f"### Response:\n")
-        enc = _tokenizer(prompt, return_tensors="pt",
-                         truncation=True, max_length=2048)
-        enc = {k: v.to("cuda") for k, v in enc.items()}
-        with torch.no_grad():
-            out = _model.generate(
-                **enc, max_new_tokens=MAX_NEW_TOKS, do_sample=False
-            )
-        new_ids = out[0, enc["input_ids"].shape[1]:]
-        raw = _tokenizer.decode(new_ids, skip_special_tokens=True,
-                                clean_up_tokenization_spaces=True)
-        results.append(raw)
-    return results
+passages = [Passage(**d) for d in DOCUMENTS]
+embs = retriever.encode(
+    [f"{p.title} {p.text}" for p in passages],
+    convert_to_numpy=True,
+    normalize_embeddings=True
+).astype("float32")
 
-# ── Parse and score ────────────────────────────────────────────────────────────
-def parse(raw: str) -> dict:
-    norm = raw.replace("_", " ")
-    r = {"isrel":None,"issup":None,"isuse":None,"answer":"","raw":raw}
-    for t in [RT.RELEVANT, RT.IRRELEVANT]:
-        if t in raw: r["isrel"]=t; break
-    for t in [RT.FULLY, RT.PARTIALLY, RT.NO_SUPPORT]:
-        if t in raw: r["issup"]=t; break
-    for t in [RT.U5,RT.U4,RT.U3,RT.U2,RT.U1]:
-        if t in raw: r["isuse"]=t; break
-    ans = raw
-    for t in RT.ALL + ["<paragraph>"]:
-        if t in ans: ans = ans[:ans.index(t)]
-    ans = re.sub(r"<[^>]+>","",ans)
-    ans = re.sub(r"#+\s*\w+.*?:","",ans)
-    r["answer"] = " ".join(ans.split()).strip()
-    for ph in ABSTAIN_PHRASES:
-        if ph.lower() in norm.lower():
-            r["issup"] = RT.NO_SUPPORT; break
-    return r
+index = faiss.IndexFlatIP(embs.shape[1])
+index.add(embs)
 
-def calc_score(p: dict) -> float:
-    s = 0.0
-    if p["isrel"] == RT.RELEVANT:    s += W_REL * 1.0
-    if p["issup"] == RT.FULLY:       s += W_SUP * 1.0
-    elif p["issup"] == RT.PARTIALLY: s += W_SUP * 0.5
-    if p["isuse"]: s += W_USE * RT.USE_SCORES.get(p["isuse"], 0.0)
-    return s
+print("Retriever ready")
 
-# ── Main inference ─────────────────────────────────────────────────────────────
-def run(query: str):
-    retrieved = retrieve(query)
-    raws = generate_candidates(query, [p.text for p in retrieved])
-    candidates = []
-    for passage, raw in zip(retrieved, raws):
-        pr = parse(raw)
-        if not pr["answer"]: continue
-        candidates.append(Candidate(
-            text=pr["answer"], passage=passage,
-            isrel=pr["isrel"], issup=pr["issup"],
-            isuse=pr["isuse"], score=calc_score(pr),
-        ))
-    if not candidates:
-        return None, retrieved, True
-    candidates.sort(key=lambda c: c.score, reverse=True)
-    best = candidates[0]
-    abstained = (best.issup == RT.NO_SUPPORT and best.score <= 0.5)
-    return best, retrieved, abstained
+def retrieve(query):
+    q = retriever.encode([query], convert_to_numpy=True,
+                         normalize_embeddings=True).astype("float32")
+    _, idxs = index.search(q, K_PASSAGES)
+    return [passages[i] for i in idxs[0]]
 
-# ── HTML rendering ─────────────────────────────────────────────────────────────
-TC = {
-    RT.RETRIEVAL:("#EEEDFE","#3C3489"),  RT.RELEVANT:("#EAF3DE","#27500A"),
-    RT.IRRELEVANT:("#FCEBEB","#791F1F"), RT.FULLY:("#EAF3DE","#27500A"),
-    RT.PARTIALLY:("#FAEEDA","#633806"),  RT.NO_SUPPORT:("#FCEBEB","#791F1F"),
-    RT.U5:("#EAF3DE","#27500A"), RT.U4:("#EAF3DE","#27500A"),
-    RT.U3:("#FAEEDA","#633806"), RT.U2:("#FCEBEB","#791F1F"),
-    RT.U1:("#FCEBEB","#791F1F"),
-}
+# ── MODEL LOADING (FIXED) ─────────────────────────────
+print("Loading model...")
 
-def badge(tok):
-    if not tok: return ""
-    bg,fg = TC.get(tok,("#F1EFE8","#444441"))
-    return (f"<span style='background:{bg};color:{fg};font-size:11px;"
-            f"font-weight:600;padding:3px 8px;border-radius:4px;"
-            f"margin:2px;display:inline-block'>{tok}</span>")
+tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
 
-def render(best, retrieved, abstained):
-    if abstained or best is None:
-        return f"""
-<div style='font-family:sans-serif;padding:16px;border-radius:10px;
-     background:#FCEBEB;border:1.5px solid #A32D2D;margin-top:8px'>
-  <div style='font-size:15px;font-weight:700;color:#A32D2D;margin-bottom:8px'>
-    &#9888; ABSTAINED — Insufficient Evidence
-  </div>
-  <div style='font-size:13px;color:#791F1F;line-height:1.6'>
-    The system could not find sufficient evidence in the document corpus
-    to answer this question. Rather than hallucinating an answer, the system abstains.
-  </div>
-  <div style='margin-top:10px;font-size:11px;color:#aaa'>
-    Passages retrieved: {', '.join(f'[{p.id}] {p.title}' for p in retrieved)}
-  </div>
-</div>"""
-    return f"""
-<div style='font-family:sans-serif;padding:16px;border-radius:10px;
-     background:#F7F7F9;border:0.5px solid #D3D1C7;margin-top:8px'>
-  <div style='font-size:11px;color:#888;margin-bottom:4px'>Best passage retrieved</div>
-  <div style='background:#E6F1FB;padding:8px 12px;border-radius:6px;
-       font-size:11px;color:#185FA5;margin-bottom:12px'>
-    [{best.passage.id}] {best.passage.title}
-  </div>
-  <div style='margin-bottom:10px'>
-    {badge(RT.RETRIEVAL)}{badge(best.isrel)}{badge(best.issup)}{badge(best.isuse)}
-    <span style='font-size:10px;color:#bbb;margin-left:6px'>score={best.score:.2f}</span>
-  </div>
-  <div style='background:#fff;padding:14px;border-radius:8px;
-       border:0.5px solid #D3D1C7;font-size:14px;color:#1E2761;
-       line-height:1.8;margin-bottom:10px'>
-    {best.text}
-  </div>
-  <div style='font-size:11px;color:#bbb'>
-    All passages considered: {', '.join(f'[{p.id}]' for p in retrieved)}
-  </div>
-</div>"""
+model = AutoModelForCausalLM.from_pretrained(
+    CHECKPOINT,
+    torch_dtype=torch.float16,
+    device_map="auto",          # ✅ KEY FIX
+    low_cpu_mem_usage=True      # ✅ KEY FIX
+)
 
-# ── Gradio UI ──────────────────────────────────────────────────────────────────
-def answer(question):
-    if not question.strip():
-        return "<div style='color:#888;padding:12px'>Please enter a question.</div>"
+model.eval()
+
+print("Model loaded successfully")
+
+# ── GENERATION ───────────────────────────────────────
+def generate_answer(query, context):
+    prompt = f"Question: {query}\nContext: {context}\nAnswer:"
+    
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
+    
+    # move to correct device
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+    
+    with torch.no_grad():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=MAX_NEW_TOKS,
+            do_sample=False
+        )
+    
+    return tokenizer.decode(output[0], skip_special_tokens=True)
+
+# ── MAIN FUNCTION ─────────────────────────────────────
+def answer(query):
     try:
-        best, retrieved, abstained = run(question)
-        return render(best, retrieved, abstained)
+        docs = retrieve(query)
+        context = " ".join([d.text for d in docs])
+        result = generate_answer(query, context)
+        return result
     except Exception as e:
-        return f"<div style='color:red;padding:12px'>Error: {str(e)}</div>"
-
-with gr.Blocks(title="SELF-RAG Demo") as demo:
-    gr.Markdown("""
-# SELF-RAG Phase 1 — Live Demo
-**Agentic Self-Reflective Retrieval-Augmented Generation**
-Aravind Kashyap | PES2PGE24DS068 | Guide: Shankar DK | PES University
-
-> The system retrieves passages, generates an answer, and critiques its output
-> using reflection tokens. If no passage supports the answer, the system **abstains**.
-
-⏳ *First query may take ~2 min for model warm-up.*
-""")
-
-    gr.Markdown("""
-### Document Corpus (8 documents)
-| ID | Topic |
-|---|---|
-| doc_00 | RAG Overview |
-| doc_01 | Vanilla RAG Limitations |
-| doc_02 | SELF-RAG Framework |
-| doc_03 | CRAG — Corrective RAG |
-| doc_04 | RAGAs Evaluation |
-| doc_05 | Hallucination in LLMs |
-| doc_06 | Adaptive Retrieval in SELF-RAG |
-| doc_07 | Segment-Level Beam Search |
-
-**Ask anything outside these topics** (e.g. geography, sports, cooking) to trigger abstention.
-""")
-
-    with gr.Row():
-        with gr.Column(scale=4):
-            qbox = gr.Textbox(
-                label="Your Question",
-                placeholder="e.g. What is SELF-RAG?  |  What is hallucination?  |  What is the capital of France?",
-                lines=2,
-            )
-        with gr.Column(scale=1, min_width=120):
-            btn = gr.Button("Ask →", variant="primary", size="lg")
-
-    gr.Examples(
-        examples=[
-            ["What is SELF-RAG and how does it use reflection tokens?"],
-            ["What is the main limitation of Vanilla RAG?"],
-            ["How does CRAG handle irrelevant retrieved documents?"],
-            ["What is hallucination in language models?"],
-            ["What is the capital of France?"],
-            ["What is the deployment date of the AuthService in the Pegasus project?"],
-        ],
-        inputs=qbox,
-        label="Try these — last two should abstain",
-    )
-
-    out = gr.HTML(label="Response")
-    btn.click(fn=answer, inputs=qbox, outputs=out)
-    qbox.submit(fn=answer, inputs=qbox, outputs=out)
-
-    gr.Markdown("""
----
-### How it works
-| Step | What happens |
-|---|---|
-| 1. Retrieve | Top-5 passages from FAISS using all-MiniLM-L6-v2 |
-| 2. Generate | selfrag/selfrag_llama2_7b generates response for each passage |
-| 3. Parse | Reflection tokens extracted from generated text |
-| 4. Score | `score = w_rel×s(IsRel) + w_sup×s(IsSupport) + w_use×s(IsUse)` |
-| 5. Abstain | If best has `[No support / Contradictory]` → "Insufficient evidence" |
-
-*Asai et al., ICLR 2024 — SELF-RAG: Learning to Retrieve, Generate and Critique*
-""")
-
-demo.queue(default_concurrency_limit=1).launch()
+        return f"Error: {str(e)}"
+
+# ── UI ───────────────────────────────────────────────
+with gr.Blocks() as demo:
+    gr.Markdown("# SELF-RAG Demo")
+    
+    q = gr.Textbox(label="Ask a question")
+    out = gr.Textbox(label="Answer")
+    
+    btn = gr.Button("Ask")
+    btn.click(answer, inputs=q, outputs=out)
+
+demo.launch()