Update app.py

app.py

@@ -36,9 +36,15 @@ def init_llm():
         raise ValueError("HUGGINGFACEHUB_API_TOKEN is not set in environment variables.")
     
     model_id = "tiiuae/falcon-rw-1b"  # ✅ Can switch to "tiiuae/falcon-rw-1b" for lighter model
-    hf_pipeline = pipeline("text-generation", model=model_id, device=DEVICE)
+    hf_pipeline = pipeline(
+    "text-generation",
+    model=model_id,
+    device=DEVICE,
+    max_new_tokens=512  # Increase this as needed
+    )
+
     llm_pipeline = HuggingFacePipeline(pipeline=hf_pipeline)
-    
+
     embeddings = HuggingFaceEmbeddings(
         model_name="sentence-transformers/all-MiniLM-L6-v2",
         model_kwargs={"device": DEVICE}
@@ -63,7 +69,7 @@ def process_document(file):
 
         # Load or create ChromaDB
         db = Chroma.from_documents(texts, embedding=embeddings, persist_directory=persist_directory)
-        retriever = db.as_retriever(search_type="similarity", search_kwargs={'k': 6})
+        retriever = db.as_retriever(search_type="mmr", search_kwargs={'k': 6})
 
         conversation_retrieval_chain = ConversationalRetrievalChain.from_llm(
             llm=llm_pipeline, retriever=retriever
@@ -82,11 +88,11 @@ def process_prompt(prompt, chat_history_display):
     if not conversation_retrieval_chain:
         return chat_history_display + [("❌ No document uploaded.", "Please upload a PDF first.")]
 
-    output = conversation_retrieval_chain({"question": prompt, "chat_history": chat_history})
+    output = conversation_retrieval_chain.invoke({"question": prompt, "chat_history": chat_history})
     answer = output["answer"]
 
     chat_history.append((prompt, answer))
-    return chat_history
+    return answer
 
 # Define Gradio UI
 with gr.Blocks(theme=gr.themes.Soft()) as demo: