Updated main.py

app.py

@@ -10,9 +10,14 @@ from transformers import pipeline
 def load_models():
     # Embedding model (lightweight)
     embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
-    # QA model (distilled SQuAD)
-    qa = pipeline('question-answering', model='distilbert-base-uncased-distilled-squad')
-    return embedder, qa
+    # Generative QA model
+    qa_gen = pipeline(
+        'text2text-generation',
+        model='google/flan-t5-base',
+        tokenizer='google/flan-t5-base',
+        device=-1  # CPU
+    )
+    return embedder, qa_gen
 
 # Extract text from uploaded file
 def extract_text_from_file(uploaded_file):
@@ -40,8 +45,8 @@ def chunk_text(text, chunk_size=500, overlap=50):
 
 # Build FAISS index from chunks
 @st.cache_resource
-def build_faiss_index(chunks, _embedder):  # underscore to avoid hashing issues
-    embeddings = _embedder.encode(chunks)
+def build_faiss_index(chunks, _embedder):  # underscore avoids hashing
+    embeddings = _embedder.encode(chunks, convert_to_numpy=True)
     dim = embeddings.shape[1]
     index = faiss.IndexFlatL2(dim)
     index.add(embeddings)
@@ -52,48 +57,51 @@ def main():
     st.set_page_config(page_title='📄 RAGbot', layout='wide')
     st.title('🤖 RagBot')
     st.sidebar.header('Upload Documents')
-    
-    # Initialize chat history in session state
+
+    # Initialize chat history
     if 'history' not in st.session_state:
         st.session_state.history = []
 
     uploaded = st.sidebar.file_uploader('Upload PDF/DOCX/TXT', type=['pdf', 'docx', 'txt'])
-    if uploaded:
-        text = extract_text_from_file(uploaded)
-        chunks = chunk_text(text)
-        embedder, qa = load_models()
-        index = build_faiss_index(chunks, embedder)
-
-        # Display existing chat history
-        for chat in st.session_state.history:
-            with st.chat_message('user'):
-                st.markdown(f"**You:** {chat['question']}")
-            with st.chat_message('assistant'):
-                st.markdown(f"**RagBot:** {chat['answer']}")
+    if not uploaded:
+        st.info('Please upload a document in the sidebar to begin.')
+        return
 
-        # Chat input
-        question = st.chat_input('Ask a question about the document...')
-        if question:
-            # Retrieve top-k relevant chunks
-            q_emb = embedder.encode([question])
-            D, I = index.search(q_emb, k=3)
-            context = '\n\n'.join(chunks[i] for i in I[0])
+    # On first load of a doc, process and index
+    if 'chunks' not in st.session_state or st.session_state.uploaded_name != uploaded.name:
+        text = extract_text_from_file(uploaded)
+        st.session_state.chunks = chunk_text(text)
+        st.session_state.embedder, st.session_state.qa_gen = load_models()
+        st.session_state.index = build_faiss_index(st.session_state.chunks, st.session_state.embedder)
+        st.session_state.uploaded_name = uploaded.name
+        st.session_state.history = []  # reset history on new doc
 
-            # Get answer
-            result = qa({'question': question, 'context': context})
-            answer = result.get('answer', 'Sorry, could not find an answer.')
+    # Display existing chat history
+    for chat in st.session_state.history:
+        with st.chat_message('user'):
+            st.markdown(f"**You:** {chat['question']}")
+        with st.chat_message('assistant'):
+            st.markdown(f"**RagBot:** {chat['answer']}")
 
-            # Save to history
-            st.session_state.history.append({'question': question, 'answer': answer})
+    # Chat input
+    question = st.chat_input('Ask a question about the document...')
+    if question:
+        # Retrieve top-k relevant chunks
+        q_emb = st.session_state.embedder.encode([question], convert_to_numpy=True)
+        _, I = st.session_state.index.search(q_emb, k=3)
+        context = '\n\n'.join(st.session_state.chunks[i] for i in I[0])
 
-            # Display new messages
-            with st.chat_message('user'):
-                st.markdown(f"**You:** {question}")
-            with st.chat_message('assistant'):
-                st.markdown(f"**RagBot:** {answer}")
+        # Generate answer
+        prompt = f"Context:\n{context}\n\nQuestion: {question}\nAnswer in detail:"
+        response = st.session_state.qa_gen(prompt, max_new_tokens=200, do_sample=False)
+        answer = response[0]['generated_text'].strip()
 
-    else:
-        st.info('Please upload a document in the sidebar to begin.')
+        # Save & display new messages
+        st.session_state.history.append({'question': question, 'answer': answer})
+        with st.chat_message('user'):
+            st.markdown(f"**You:** {question}")
+        with st.chat_message('assistant'):
+            st.markdown(f"**RagBot:** {answer}")
 
 if __name__ == '__main__':
     main()