Update app.py

app.py

@@ -3,7 +3,9 @@ from langchain_huggingface import HuggingFaceEndpoint, HuggingFaceEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
 from langchain_community.document_loaders import PyPDFLoader
-from langchain.chains.retrieval_qa.base import RetrievalQA
+from langchain_core.prompts import PromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
 
 import gradio as gr
 import warnings
@@ -149,14 +151,30 @@ def retriever_qa(file, query, model_choice, max_tokens, temperature, embedding_m
         selected_model = model_choice or MODEL_OPTIONS[0]
         llm = get_llm(selected_model, int(max_tokens), float(temperature))
         retriever_obj = retriever(file, int(chunk_size), int(chunk_overlap), embedding_model)
-        qa = RetrievalQA.from_chain_type(
-            llm=llm,
-            chain_type="stuff",
-            retriever=retriever_obj,
-            return_source_documents=True,
+        
+        # Create a simple RAG chain
+        template = """Answer the question based only on the following context:
+{context}
+
+Question: {question}
+
+Answer:"""
+        
+        prompt = PromptTemplate.from_template(template)
+        
+        def format_docs(docs):
+            return "\n\n".join(doc.page_content for doc in docs)
+        
+        # Build the chain
+        rag_chain = (
+            {"context": retriever_obj | format_docs, "question": RunnablePassthrough()}
+            | prompt
+            | llm
+            | StrOutputParser()
         )
-        response = qa.invoke({"query": query})
-        return response['result']
+        
+        response = rag_chain.invoke(query)
+        return response
     except Exception as e:
         return f"Error: {str(e)}"