ImportError: This modeling file requires the following packages that were not found in your environment: einops. Run `pip install einops`

app.py

@@ -102,7 +102,43 @@ 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})"
+
+
+# reranker 
+from sentence_transformers import CrossEncoder
+import numpy as np
+import torch
+
+class ProductionReranker:
+    def __init__(self, model_name="jinaai/jina-reranker-v2-base-multilingual"):
+        self.model = CrossEncoder(
+            model_name,
+            max_length=512,
+            device='cuda' if torch.cuda.is_available() else 'cpu',
+            trust_remote_code=True
+        )
     
+    def rerank(self, query, documents, k=5):
+        # Extract text
+        doc_texts = [
+            doc.page_content if hasattr(doc, 'page_content') else str(doc) 
+            for doc in documents
+        ]
+        
+        # Score in batches for efficiency
+        pairs = [[query, doc] for doc in doc_texts]
+        scores = self.model.predict(pairs, batch_size=32)
+        
+        # Get top-k
+        top_indices = np.argsort(scores)[::-1][:k]
+        
+        # Return with scores
+        reranked = [(documents[i], float(scores[i])) for i in top_indices]
+        return [doc for doc, score in reranked]
+
+
+
+
 # setup chatbot
 from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
 from langchain.chat_models import init_chat_model
@@ -151,16 +187,22 @@ def predict(message, history, request: gr.Request):
     
 
     # Retrieve relevant documents for the current message
+    print("retreive docs ...")
     relevant_docs = vectorstore.similarity_search(message,k=20)  #  retriever
     
+    # reank docs
+    print("reranking ...")
+    RERANKER = ProductionReranker()
+    relevant_docs = RERANKER.rerank(message, relevant_docs, k=10)
+
     # Build context from retrieved documents
+    print("build context ...")
     context = "\nExtracted documents:\n" + "\n".join([
         f"Content document {i+1}: {doc.page_content}\n\n---"
         for i, doc in enumerate(relevant_docs)
     ])
 
     
-
     # 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

@@ -5,6 +5,8 @@ torchaudio
 
 sentence-transformers
 faiss-cpu
+sentence-transformers>=2.5.0
+einops
 
 langchain-core==0.3.21
 langchain==0.3.8
@@ -12,6 +14,7 @@ langchain-community==0.3.8
 langchain-openai==0.2.9
 langchain-huggingface==0.1.0
 
+
 gradio
 
 python-dotenv