Update app.py

app.py

@@ -6,19 +6,24 @@ from langchain_openai import OpenAIEmbeddings
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnablePassthrough
 from langchain_core.prompts import ChatPromptTemplate
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import FlashrankRerank
+from langchain_openai import ChatOpenAI
+
 import gradio as gr
 
 
 def format_docs(docs):
-  print(docs)
   return "\n\n".join(doc.page_content for doc in docs)
 
 
-
 embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
 db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
-retriever = db.as_retriever(k=5)
-
+retriever = db.as_retriever(k=10)
+compressor = FlashrankRerank()
+compression_retriever = ContextualCompressionRetriever(
+    base_compressor=compressor, base_retriever=retriever
+)
 
 llm = ChatOpenAI(model="gpt-4o")
 
@@ -33,7 +38,7 @@ prompt = ChatPromptTemplate.from_messages([
 ])
 
 rag_chain = (
-    {"context": retriever | format_docs, "question": RunnablePassthrough()}
+    {"context": compression_retriever | format_docs, "question": RunnablePassthrough()}
     | prompt
     | llm
     | StrOutputParser()
@@ -48,7 +53,6 @@ def chat_gen(message, history):
 
   partial_message=""
   for chunk in rag_chain.stream(message):
-    # if chunk.choices[0].delta.content is not None:
     partial_message = partial_message + chunk
     yield partial_message
 
@@ -59,7 +63,7 @@ demo = gr.ChatInterface(chat_gen, chatbot=chatbot).queue()
 
 
 try:
-    demo.launch(debug=True, share=True, show_api=False)
+    demo.launch(debug=True, share=False, show_api=False)
     demo.close()
 except Exception as e:
     demo.close()