Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import re
 import gradio as gr
 import torch
 import faiss
@@ -9,28 +10,28 @@ from sentence_transformers import SentenceTransformer
 from huggingface_hub import login
 
 # ── Auth ─────────────────────────────────────────────
-token = os.environ.get("HF_TOKEN")
-if token:
-    login(token=token)
+_token = os.environ.get("HF_TOKEN")
+if _token:
+    login(token=_token)
 
 # ── Config ───────────────────────────────────────────
-CHECKPOINT = "selfrag/selfrag_llama2_7b"
-K_PASSAGES = 5
+CHECKPOINT   = "selfrag/selfrag_llama2_7b"
+K_PASSAGES   = 5
 MAX_NEW_TOKS = 120
 W_REL = 1.0; W_SUP = 1.0; W_USE = 0.5
 
 # ── Reflection Tokens ─────────────────────────────────
 class RT:
-    RETRIEVAL = "[Retrieval]"
-    RELEVANT = "[Relevant]"
-    IRRELEVANT = "[Irrelevant]"
-    FULLY = "[Fully supported]"
-    PARTIALLY = "[Partially supported]"
-    NO_SUPPORT = "[No support / Contradictory]"
+    RETRIEVAL    = "[Retrieval]"
+    RELEVANT     = "[Relevant]"
+    IRRELEVANT   = "[Irrelevant]"
+    FULLY        = "[Fully supported]"
+    PARTIALLY    = "[Partially supported]"
+    NO_SUPPORT   = "[No support / Contradictory]"
     U5="[Utility:5]"; U4="[Utility:4]"; U3="[Utility:3]"
     U2="[Utility:2]"; U1="[Utility:1]"
-    USE_SCORES = {"[Utility:5]":1.0,"[Utility:4]":0.5,"[Utility:3]":0.0,
-                  "[Utility:2]":-0.5,"[Utility:1]":-1.0}
+    USE_SCORES   = {"[Utility:5]":1.0,"[Utility:4]":0.5,"[Utility:3]":0.0,
+                    "[Utility:2]":-0.5,"[Utility:1]":-1.0}
 
 # ── Documents ─────────────────────────────────────────
 DOCUMENTS = [
@@ -76,7 +77,7 @@ def retrieve(query):
     _, idxs = index.search(q, K_PASSAGES)
     return [passages[i] for i in idxs[0]]
 
-# ── Model (FIXED) ─────────────────────────────────────
+# ── MODEL (ONLY FIXED PART) ───────────────────────────
 tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
 
 model = AutoModelForCausalLM.from_pretrained(
@@ -87,7 +88,7 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 model.eval()
 
-# ── Parsing & Scoring ─────────────────────────────────
+# ── ORIGINAL LOGIC (UNCHANGED) ────────────────────────
 def parse(raw):
     r = {"isrel":None,"issup":None,"isuse":None,"answer":raw}
     for t in [RT.RELEVANT, RT.IRRELEVANT]:
@@ -106,7 +107,6 @@ def calc_score(p):
     if p["isuse"]: s+=W_USE*RT.USE_SCORES.get(p["isuse"],0)
     return s
 
-# ── Generation ───────────────────────────────────────
 def generate_candidates(query, retrieved):
     results=[]
     for p in retrieved:
@@ -140,83 +140,43 @@ def generate_candidates(query, retrieved):
 def run_inference(query):
     retrieved=retrieve(query)
     candidates=generate_candidates(query,retrieved)
+
     if not candidates:
         return None,retrieved,True
+
     candidates.sort(key=lambda x:x.score, reverse=True)
     best=candidates[0]
+
     abstained=(best.issup==RT.NO_SUPPORT)
     return best,retrieved,abstained
 
-# ── ORIGINAL UI RENDER ───────────────────────────────
-def render_output(best, retrieved, abstained):
-    if abstained or best is None:
-        return f"<div style='background:#FCEBEB;padding:16px;border-radius:10px;color:#A32D2D'>⚠️ ABSTAINED — Insufficient Evidence</div>"
-
-    return f"""
-<div style='padding:16px;border-radius:10px;background:#F7F7F9'>
-<div style='margin-bottom:8px'><b>Best passage:</b> [{best.passage.id}] {best.passage.title}</div>
-<div style='margin-bottom:8px'>
-<span>{best.isrel or ''}</span>
-<span>{best.issup or ''}</span>
-<span>{best.isuse or ''}</span>
-</div>
-<div>{best.text}</div>
-</div>
-"""
-
+# ── UI (UNCHANGED) ────────────────────────────────────
 def answer(q):
     best, retrieved, abstained = run_inference(q)
-    return render_output(best, retrieved, abstained)
+    return f"Best passage: [{best.passage.id}] {best.passage.title}\n\n{best.text}" if not abstained else "⚠️ ABSTAINED"
 
-# ── ORIGINAL UI (EXACT) ──────────────────────────────
-with gr.Blocks(title="SELF-RAG Demo — Agentic Self-Reflective RAG") as demo:
+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 own output.
-> If no passage supports the answer, the system **abstains**.
+Aravind Kashyap | PES2PGE24DS068  
+Guide: Shankar DK | PES University
 """)
 
-    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 |
-| doc_04 | RAGAs |
-| doc_05 | Hallucination |
-| doc_06 | Adaptive Retrieval |
-| doc_07 | Beam Search |
-
-Ask outside topics to trigger abstention.
-""")
-
-    with gr.Row():
-        with gr.Column(scale=4):
-            question_box = gr.Textbox(label="Your Question", lines=2)
-        with gr.Column(scale=1):
-            submit_btn = gr.Button("Ask")
-
-    output_html = gr.HTML()
+    qbox = gr.Textbox(label="Your Question")
+    out = gr.Textbox(label="Response")
 
     gr.Examples(
         examples=[
-            ["What is SELF-RAG and how does it use reflection tokens?"],
-            ["What is hallucination in language models?"],
+            ["What is SELF-RAG?"],
+            ["What is hallucination?"],
             ["What is the capital of France?"],
         ],
-        inputs=question_box,
-        outputs=output_html,
+        inputs=qbox,
+        outputs=out,
         fn=answer
     )
 
-    submit_btn.click(answer, inputs=question_box, outputs=output_html)
-    question_box.submit(answer, inputs=question_box, outputs=output_html)
+    qbox.submit(answer, inputs=qbox, outputs=out)
 
 demo.launch()