integrate reranker

app.py

@@ -153,6 +153,14 @@ 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")
+
+    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---"
@@ -160,7 +168,6 @@ def predict(message, history, request: gr.Request):
     ])
 
     
-
     # RAG tool
     RAG_PROMPT_TEMPLATE="""You will be asked information related to Rémi Cazelles's specific projects, work and education.
                         Using the information contained in the context, provide a structured answer to the question.

requirements.txt

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