Update app.py

app.py

@@ -8,7 +8,7 @@ import faiss
 import numpy as np
 import gradio as gr
 from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer, AutoModelForQuestionAnswering, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from PyPDF2 import PdfReader
 import docx
 
@@ -40,7 +40,7 @@ def get_default_config():
         "models": {
             # Embedding model for FAISS
             "embedding": "sentence-transformers/all-MiniLM-L6-v2",
-            # Abstractive generation model (can upgrade to flan-t5-base if resources allow)
+            # Abstractive generation model
             "qa": "google/flan-t5-small",
         },
         "chunking": {
@@ -181,7 +181,8 @@ def load_kb_documents(kb_dir: str) -> List[Tuple[str, str]]:
 class RAGIndex:
     def __init__(self):
         self.embedder = None
-        self.qa_pipeline = None  # now a generative pipeline
+        self.qa_tokenizer = None
+        self.qa_model = None
         self.chunks: List[str] = []
         self.chunk_sources: List[str] = []
         self.index = None
@@ -203,13 +204,9 @@ class RAGIndex:
             print(f"Loading embedding model: {EMBEDDING_MODEL_NAME}")
             self.embedder = SentenceTransformer(EMBEDDING_MODEL_NAME)
 
-            print(f"Loading QA (generation) model: {QA_MODEL_NAME}")
-            # Abstractive generation pipeline (Flan-T5)
-            self.qa_pipeline = pipeline(
-                "text2text-generation",
-                model=QA_MODEL_NAME,
-                tokenizer=QA_MODEL_NAME,
-            )
+            print(f"Loading QA (seq2seq) model: {QA_MODEL_NAME}")
+            self.qa_tokenizer = AutoTokenizer.from_pretrained(QA_MODEL_NAME)
+            self.qa_model = AutoModelForSeq2SeqLM.from_pretrained(QA_MODEL_NAME)
         except Exception as e:
             print(f"Error loading models: {e}")
             raise
@@ -329,6 +326,31 @@ class RAGIndex:
             print(f"Retrieval error: {e}")
             return []
 
+    def _generate_from_context(self, prompt: str) -> str:
+        """Run Flan-T5 on the given prompt and return the decoded answer."""
+        if self.qa_model is None or self.qa_tokenizer is None:
+            return "Model not loaded."
+
+        inputs = self.qa_tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=768,
+        )
+
+        output_ids = self.qa_model.generate(
+            **inputs,
+            max_new_tokens=256,
+            do_sample=False,
+        )
+
+        answer = self.qa_tokenizer.decode(
+            output_ids[0],
+            skip_special_tokens=True,
+        ).strip()
+
+        return answer
+
     def answer(self, question: str) -> str:
         """Answer a question using RAG + abstractive generation"""
         if not self.initialized:
@@ -363,7 +385,7 @@ class RAGIndex:
 
         combined_text = "\n\n".join(combined_context)
 
-        # Limit context length to keep it manageable for the model
+        # Limit context length to keep it manageable
         max_context_chars = 4000
         if len(combined_text) > max_context_chars:
             combined_text = combined_text[:max_context_chars]
@@ -379,13 +401,7 @@ class RAGIndex:
         )
 
         try:
-            result = self.qa_pipeline(
-                prompt,
-                max_new_tokens=256,
-                do_sample=False,
-            )
-            # text2text-generation returns list of dicts with 'generated_text'
-            answer_text = result[0]["generated_text"].strip()
+            answer_text = self._generate_from_context(prompt)
         except Exception as e:
             print(f"Generation error: {e}")
             return (