Update app.py

app.py

@@ -8,8 +8,9 @@ from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
 from langchain.llms import HuggingFacePipeline
 from langchain.chains import RetrievalQA
+from langchain.prompts import PromptTemplate
 
-from transformers.pipelines import pipeline   # ✅ FIXED IMPORT
+from transformers.pipelines import pipeline
 
 # -------------------------------
 # Page Config
@@ -19,7 +20,7 @@ st.title("📄 Chat with Your Documents (RAG)")
 st.write("🚀 App started successfully")
 
 # -------------------------------
-# Load Documents (FIXED)
+# Load Documents
 # -------------------------------
 def load_documents(uploaded_files):
     documents = []
@@ -27,7 +28,6 @@ def load_documents(uploaded_files):
     for file in uploaded_files:
         file_extension = os.path.splitext(file.name)[1]
 
-        # Save safely as temp file
         with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as tmp:
             tmp.write(file.getbuffer())
             temp_path = tmp.name
@@ -76,30 +76,56 @@ def create_vectorstore(chunks):
 
 
 # -------------------------------
-# Cached LLM (FIXED)
+# Cached LLM
 # -------------------------------
 @st.cache_resource
 def load_llm():
     pipe = pipeline(
-        "text2text-generation",   # ✅ CORRECT TASK
+        "text2text-generation",
         model="google/flan-t5-small",
         max_length=256
     )
     return HuggingFacePipeline(pipeline=pipe)
 
 
+# -------------------------------
+# Custom Prompt (IMPORTANT)
+# -------------------------------
+prompt_template = """
+Use the following context to answer the question clearly.
+
+Context:
+{context}
+
+Question:
+{question}
+
+Answer:
+"""
+
+PROMPT = PromptTemplate(
+    template=prompt_template,
+    input_variables=["context", "question"]
+)
+
+
 # -------------------------------
 # Build QA Chain
 # -------------------------------
 def build_qa(vectorstore):
     llm = load_llm()
-    retriever = vectorstore.as_retriever()
+
+    retriever = vectorstore.as_retriever(
+        search_kwargs={"k": 3}   # 🔥 improves answer quality
+    )
 
     qa = RetrievalQA.from_chain_type(
         llm=llm,
         retriever=retriever,
+        chain_type_kwargs={"prompt": PROMPT},
         return_source_documents=False
     )
+
     return qa
 
 
@@ -134,7 +160,9 @@ if uploaded_files:
         with st.spinner("🤖 Generating answer..."):
             try:
                 result = qa_chain.run(query)
+
                 st.markdown("### 🧠 Answer:")
                 st.write(result)
+
             except Exception as e:
                 st.error(f"❌ Error generating answer: {e}")