Update app.py

app.py

@@ -1,19 +1,20 @@
 import streamlit as st
 import os
 
-# ✅ Imports
 from langchain_community.document_loaders import PyPDFLoader, TextLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
+
 from langchain.chains import RetrievalQA
+from langchain.prompts import PromptTemplate
+from langchain_community.llms import HuggingFacePipeline
 
 from transformers import pipeline
-from langchain_community.llms import HuggingFacePipeline
 
 
 # -------------------------------
-# Load Documents (SAFE PATH)
+# Load Documents (SAFE)
 # -------------------------------
 def load_documents(uploaded_files):
     documents = []
@@ -35,48 +36,67 @@ def load_documents(uploaded_files):
 
 
 # -------------------------------
-# Split Documents
+# Split Documents (BETTER CHUNKS)
 # -------------------------------
 def split_documents(documents):
     splitter = RecursiveCharacterTextSplitter(
-        chunk_size=500,
-        chunk_overlap=50
+        chunk_size=800,
+        chunk_overlap=100
     )
     return splitter.split_documents(documents)
 
 
 # -------------------------------
-# Create Vector Store (LOCAL)
+# Embeddings
 # -------------------------------
 def create_vectorstore(chunks):
     embeddings = HuggingFaceEmbeddings(
-        model_name="sentence-transformers/all-MiniLM-L6-v2"  # works without token
+        model_name="sentence-transformers/all-MiniLM-L6-v2"
     )
     return FAISS.from_documents(chunks, embeddings)
 
 
 # -------------------------------
-# Load LOCAL LLM (VERY LIGHT)
+# LLM (Balanced quality + speed)
 # -------------------------------
 def load_llm():
     pipe = pipeline(
         "text2text-generation",
-        model="sshleifer/tiny-t5",   # 🔥 super light, no auth needed
-        max_length=256
+        model="google/flan-t5-small",   # BEST without token
+        max_length=512,
+        temperature=0.3
     )
     return HuggingFacePipeline(pipeline=pipe)
 
 
 # -------------------------------
-# Build QA Chain
+# Prompt (VERY IMPORTANT)
 # -------------------------------
 def build_qa(vectorstore):
     llm = load_llm()
-    retriever = vectorstore.as_retriever()
+
+    prompt_template = """
+    Use the following context to answer the question.
+    If the answer is not in the context, say "Answer not found in document".
+
+    Context:
+    {context}
+
+    Question:
+    {question}
+
+    Answer:
+    """
+
+    PROMPT = PromptTemplate(
+        template=prompt_template,
+        input_variables=["context", "question"]
+    )
 
     return RetrievalQA.from_chain_type(
         llm=llm,
-        retriever=retriever
+        retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
+        chain_type_kwargs={"prompt": PROMPT}
     )
 
 
@@ -103,10 +123,8 @@ if uploaded_files:
     query = st.text_input("Ask a question from your documents")
 
     if query:
-        with st.spinner("Generating answer..."):
-            try:
-                result = qa_chain.run(query)
-                st.write("### Answer:")
-                st.write(result)
-            except Exception as e:
-                st.error(str(e))
+        with st.spinner("Thinking..."):
+            result = qa_chain.run(query)
+
+            st.write("### 📌 Answer:")
+            st.write(result)