Update app.py

app.py

@@ -13,6 +13,7 @@ import requests
 from functools import lru_cache
 import torch
 from sentence_transformers import SentenceTransformer
+from sentence_transformers.cross_encoder import CrossEncoder
 import threading
 from queue import Queue
 import concurrent.futures
@@ -46,6 +47,7 @@ class OptimizedRAGLoader:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.encoder = SentenceTransformer("intfloat/multilingual-e5-large")
         self.encoder.to(self.device)
+        self.reranker = model = CrossEncoder("cross-encoder/mmarco-mMiniLMv2-L12-H384-v1",trust_remote_code=True)
         
         # Initialize thread pool
         self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=4)
@@ -212,7 +214,7 @@ llm = ChatMistralAI(
 )
 
 rag_loader = OptimizedRAGLoader()
-retriever = rag_loader.get_retriever(k=4)  # Reduced k for faster retrieval
+retriever = rag_loader.get_retriever(k=30)  # Reduced k for faster retrieval
 
 # Cache for processed questions
 question_cache = {}
@@ -335,9 +337,21 @@ def process_question(question: str) -> Iterator[str]:
         return
         
     relevant_docs = retriever(question)
-    context = "\n".join([doc.page_content for doc in relevant_docs])
+    # context = "\n".join([doc.page_content for doc in relevant_docs])
+
+
+    context = [doc.page_content for doc in relevant_docs]
+    text_pairs = [[question, text] for text in context]
+    scores = rag_loader.reranker.predict(text_pairs)
+    
+    scored_docs = list(zip(scores, context, relevant_docs))
+    # scored_docs.sort(reverse=True)
+    scored_docs.sort(key=lambda x: x[0], reverse=True)
+    reranked_docs = [d[2].page_content for d in scored_docs][:6]
+
+    
     prompt = prompt_template.format_messages(
-        context=context,
+        context=reranked_docs,
         question=question
     )
     full_response = ""
@@ -353,6 +367,12 @@ def process_question(question: str) -> Iterator[str]:
             # yield full_response + "\n\n\nالمصادر المحتملة :\n" + "".join(sources)
             sources = [doc.metadata.get("source") for doc in relevant_docs]
             sources = list(set([os.path.splitext(source)[0] for source in sources]))
+
+
+            sources = [d[2].metadata['source'] for d in scored_docs][:6]
+            sources = list(set([os.path.splitext(source)[0] for source in sources]))
+
+            
             yield full_response + "\n\n\nالمصادر المحتملة :\n" + "\n".join(sources)
             # yield full_response + "\n\n\nالمصادر المحتملة:\n" + "\n".join([doc.metadata.get("source") for doc in relevant_docs])
         question_cache[question] = (full_response, relevant_docs)