updated description

app.py

@@ -32,7 +32,6 @@ def download_file(url, dest_path):
         f.write(r.content)
         print(f"Saved to {dest_path}")
 
-
 # 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"))
@@ -41,12 +40,7 @@ download_file(DOCSTORE_URL, os.path.join(INDEX_DIR, "docstore.pkl"))
 # Retriever
 # ------------------------------
 class Retriever:
-
-    def __init__(
-            self, 
-            index_dir, 
-            cross_encoder_model="cross-encoder/ms-marco-MiniLM-L-6-v2"
-            ):
+    def __init__(self, index_dir, cross_encoder_model="cross-encoder/ms-marco-MiniLM-L-6-v2"):
         index, segments = self._load_index(index_dir)
         self.index = index
         self.segments = segments
@@ -59,7 +53,7 @@ class Retriever:
 
     def _load_index(self, index_dir):
         index = faiss.read_index(os.path.join(index_dir, "index.faiss"))
-        with open(os.path.join(index_dir, "docstore.pkl") , "rb") as f:
+        with open(os.path.join(index_dir, "docstore.pkl"), "rb") as f:
             segments = pickle.load(f)
         return index, segments
     
@@ -68,29 +62,20 @@ class Retriever:
         faiss.normalize_L2(embedding)
         return embedding
 
-    def _cosine_similarity(self, a, b):
-        return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
-
     def retrieve(self, query, k=50):
-        """
-        1. Retrieve top-k segments using bi-encoder (FAISS)
-        2. Re-rank segments using cross-encoder on segment['text']
-        3. Re-score each sentence inside chosen segment using cross-encoder
-        4. Highlight the best sentence
-        """
-        # ---------- Stage 1: Bi-Encoder Retrieval ----------
+        # ---------- Stage 1: Bi-Encoder ----------
         embedding = self.preprocess_query(query)
         D, I = self.index.search(embedding, k)
 
         candidates = []
-        ce_pairs_segments = []   # (query, segment_text)
+        ce_pairs_segments = []
 
         for idx in I[0]:
             seg = self.segments[idx]
             candidates.append(seg)
             ce_pairs_segments.append([query, seg["text"]])
 
-        # ---------- Stage 2: Cross-Encoder Re-Rank Segments ----------
+        # ---------- Stage 2: Cross-Encoder Re-Rank ----------
         segment_scores = self.cross.predict(ce_pairs_segments)
         best_seg_idx = int(np.argmax(segment_scores))
         best_segment = candidates[best_seg_idx]
@@ -98,21 +83,18 @@ class Retriever:
         # ---------- Stage 3: Cross-Encoder Sentence Ranking ----------
         sentences = best_segment["sentences"]
         ce_pairs_sentences = [[query, s] for s in sentences]
-
         sentence_scores = self.cross.predict(ce_pairs_sentences)
-        best_sent_idx = int(np.argmax(sentence_scores))
 
+        best_sent_idx = int(np.argmax(sentence_scores))
         best_sentence = sentences[best_sent_idx].strip()
 
-        # Highlight within full segment
         highlighted_text = (
             best_segment["text"]
             .replace(best_sentence, f"**{best_sentence}**")
             .replace("\n", " ")
         )
 
-        # ---------- Result ----------
-        result = {
+        return {
             "text": highlighted_text,
             "url": best_segment.get("url"),
             "document_id": best_segment.get("document_id"),
@@ -120,46 +102,51 @@ class Retriever:
             "sentence_score": float(sentence_scores[best_sent_idx]),
         }
 
-        return result
-
-
 # ------------------------------
-# Lightweight Generator
+# Generators (loaded once)
 # ------------------------------
-# Finetuned TinyLlama
-generator = pipeline(
-    "text-generation",
-    model="LoneWolfgang/tinyllama-for-abalone-RAG",
-    max_new_tokens=150,
-    temperature=0.1,
-)
+generators = {
+    "TinyLlama": pipeline(
+        "text-generation",
+        model="LoneWolfgang/tinyllama-for-abalone-RAG",
+        max_new_tokens=150,
+        temperature=0.1,
+    ),
+    "FLAN-T5": pipeline(
+        "text2text-generation",
+        model="google/flan-t5-base",
+        max_length=200,
+    )
+}
 
+retriever = Retriever(INDEX_DIR)
 
 # ------------------------------
 # Combined function: retrieve → generate
 # ------------------------------
-retriever = Retriever(INDEX_DIR)
-
-def answer_query(query):
+def answer_query(query, model_choice):
     doc = retriever.retrieve(query)
 
     url = doc["url"]
     context = doc["text"].replace("\n", " ")
 
     prompt = f"""
-    <|system|>
     You answer questions strictly using the provided context.
-    <|user|>
+
     Context: {context}
 
     Question: {query}
-    <|assistant|>
     """
 
-    result = generator(prompt)[0]["generated_text"]
+    # Choose generator
+    gen = generators[model_choice]
 
-    # Keep only model completion after the assistant token
-    result = result.split("<|assistant|>")[-1].strip()
+    if model_choice == "TinyLlama":
+        out = gen(f"<|system|>{prompt}<|assistant|>")[0]["generated_text"]
+        result = out.split("<|assistant|>")[-1].strip()
+    else:
+        # FLAN-T5 returns text in "generated_text"
+        result = gen(prompt)[0]["generated_text"]
 
     return (
         f"#### Response\n\n"
@@ -176,12 +163,26 @@ def answer_query(query):
 # ------------------------------
 demo = gr.Interface(
     fn=answer_query,
-    inputs=gr.Textbox(label="Enter your question"),
+    inputs=[
+        gr.Textbox(label="Enter your question"),
+        gr.Radio(
+            ["TinyLlama", "FLAN-T5"],
+            label="Choose Model",
+            value="FLAN-T5"
+        )
+    ],
     outputs=gr.Markdown(label="Answer"),
     title="Abalone RAG Demo",
-    description="This RAG system uses SBERT + Cross-Encoders for Retrieval with TinyLlama finetuned on responses from GPT5."
+    description="""This RAG system uses [SBERT](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2) for initial retrieval and a [Cross Encoder](https://huggingface.co/cross-encoder/ms-marco-MiniLM-L6-v2) for re-ranking and highlighting.
+
+Sentence embeddings are computed and [indexed](https://huggingface.co/LoneWolfgang/abalone-index) using FAISS.
+
+For generation, you can choose between:
+
+- [FLAN-T5](https://huggingface.co/google/flan-t5-base) — fast, reliable, and ideal for exploring retrieval quality.
+- [Finetuned TinyLlama](https://huggingface.co/LoneWolfgang/tinyllama-for-abalone-RAG) — slower, but more expressive.
+"""
 )
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
-