Added Generation

app.py

@@ -4,9 +4,10 @@ import pickle
 import sentence_transformers
 import faiss
 import gradio as gr
+from transformers import pipeline
 
 # ------------------------------
-# Configuration: URLs to your files
+# Configuration
 # ------------------------------
 INDEX_URL = "https://huggingface.co/LoneWolfgang/abalone-index/resolve/main/index.faiss"
 DOCSTORE_URL = "https://huggingface.co/LoneWolfgang/abalone-index/resolve/main/docstore.pkl"
@@ -31,21 +32,18 @@ def download_file(url, dest_path):
     else:
         print(f"{dest_path} already exists, skipping download.")
 
-# Download the FAISS index and docstore
+# Download index + docstore
 download_file(INDEX_URL, os.path.join(INDEX_DIR, "index.faiss"))
 download_file(DOCSTORE_URL, os.path.join(INDEX_DIR, "docstore.pkl"))
 
 # ------------------------------
-# Retriever class
+# Retriever
 # ------------------------------
 class Retriever:
     def __init__(self, index_dir, sbert_model="all-MiniLM-L12-v2"):
-        # Load FAISS index
         self.index = faiss.read_index(os.path.join(index_dir, "index.faiss"))
-        # Load docstore
         with open(os.path.join(index_dir, "docstore.pkl"), "rb") as f:
             self.segments = pickle.load(f)
-        # Load SentenceTransformer
         self.sbert = sentence_transformers.SentenceTransformer(sbert_model)
 
     def preprocess_query(self, query):
@@ -53,38 +51,66 @@ class Retriever:
         faiss.normalize_L2(embedding)
         return embedding
 
-    def retrieve(self, query, k=5):
+    def retrieve(self, query, k=1):
         embedding = self.preprocess_query(query)
         D, I = self.index.search(embedding, k)
-        results = []
-        for rank, (idx, score) in enumerate(zip(I[0], D[0]), start=1):
-            text = self.segments[idx]
-            results.append(f"**{rank}. (Score={score:.4f})**\n{text}")
-        return "\n\n".join(results)
+        top_docs = [self.segments[idx] for idx in I[0]]
+        return top_docs, D[0]
 
 # ------------------------------
-# Instantiate retriever
+# Lightweight Generator
 # ------------------------------
-retriever = Retriever(INDEX_DIR)
+# FLAN-T5-base is small (~250M) and fast to run on CPU
+generator = pipeline(
+    "text2text-generation",
+    model="google/flan-t5-base",
+    tokenizer="google/flan-t5-base",
+    max_new_tokens=150,
+    temperature=0.1,
+)
 
 # ------------------------------
-# Gradio interface
+# Combined function: retrieve → generate
 # ------------------------------
-def search(query, top_k):
-    return retriever.retrieve(query, k=top_k)
+retriever = Retriever(INDEX_DIR)
+
+def answer_query(query):
+    docs, scores = retriever.retrieve(query, k=1)
+    record = docs[0]
+    url = record["url"]
+    context = record["text"]
+
+    prompt = (
+        f"Answer the following question based on the context.\n\n"
+        f"Context:\n{context}\n\n"
+        f"Question: {query}\nAnswer:"
+    )
+    result = generator(prompt)[0]["generated_text"]
+
+    return f"""
+### Response
+{result}
 
+---
+
+**Context**  
+  
+{context}
+
+**[Source]({url})**
+"""
+
+# ------------------------------
+# Gradio UI
+# ------------------------------
 demo = gr.Interface(
-    fn=search,
-    inputs=[
-        gr.Textbox(label="Enter your query"),
-        gr.Slider(1, 10, value=5, step=1, label="Top K")
-    ],
-    outputs=gr.Markdown(label="Results"),
-    title="FAISS Retriever",
-    description="Semantic search using SentenceTransformers + FAISS.",
+    fn=answer_query,
+    inputs=gr.Textbox(label="Enter your question"),
+    outputs=gr.Markdown(label="Answer"),
+    title="RAG Demo",
+    description="Retrieves the top 1 passage and generates an answer using FLAN-T5.",
     theme="soft",
-    allow_flagging="never"
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()