reranker from langchain

app.py

@@ -4,7 +4,8 @@
 from dotenv import load_dotenv
 import os 
 load_dotenv()
-
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain_community.document_compressors import ColbertReranker
 
 from langchain.chat_models import init_chat_model
 
@@ -102,7 +103,7 @@ def format_source(doc):
         page_label = doc.metadata["pagpage_labele"]
         total_page = doc.metadata["total_page"]
         return f"{source.split('/')[-1]} page({page_label/total_page})"
-    
+
 # setup chatbot
 from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
 from langchain.chat_models import init_chat_model
@@ -153,14 +154,16 @@ def predict(message, history, request: gr.Request):
     # Retrieve relevant documents for the current message
     relevant_docs = vectorstore.similarity_search(message,k=20)  #  retriever
     
-    # reranker
-    from ragatouille import RAGPretrainedModel
-
-    RERANKER = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")
+    # reank docs
+    reranker = ColbertReranker(
+        model_name="colbert-ir/colbertv2.0",
+        top_n=10
+    )
+    relevant_docs = reranker.compress_documents(
+        documents=relevant_docs,
+        query=message
+    )
 
-    relevant_docs = RERANKER.rerank(message, relevant_docs, k=10)
-    relevant_docs = [doc["content"] for doc in relevant_docs]
-    
     # Build context from retrieved documents
     context = "\nExtracted documents:\n" + "\n".join([
         f"Content document {i+1}: {doc.page_content}\n\n---"

requirements.txt

@@ -11,7 +11,7 @@ langchain==0.3.8
 langchain-community==0.3.8
 langchain-openai==0.2.9
 langchain-huggingface==0.1.0
-RAGatouille
+
 
 gradio