Implemented new rerank

rag.py

@@ -1,7 +1,8 @@
 import logging
 import os
 import requests
-
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
 
 
 from langchain_community.vectorstores import FAISS
@@ -15,11 +16,13 @@ class RAG:
     #vectorstore = "vectorestore" # CA only
     vectorstore = "index-BAAI_bge-m3-1500-200-recursive_splitter-CA_ES_UE"
 
-    def __init__(self, hf_token, embeddings_model, model_name):
+    def __init__(self, hf_token, embeddings_model, model_name, rerank_model, rerank_number_contexts):
 
 
         self.model_name = model_name
         self.hf_token = hf_token
+        self.rerank_model = rerank_model
+        self.rerank_number_contexts = rerank_number_contexts
         
         # load vectore store
         embeddings = HuggingFaceEmbeddings(model_name=embeddings_model, model_kwargs={'device': 'cpu'})
@@ -27,10 +30,50 @@ class RAG:
 
         logging.info("RAG loaded!")
     
+    def rerank_contexts(self, instruction, contexts, number_of_contexts=1):
+        """
+        Rerank the contexts based on their relevance to the given instruction.
+        """
+
+        rerank_model = self.rerank_model
+        
+
+        tokenizer = AutoTokenizer.from_pretrained(rerank_model)
+        model = AutoModelForSequenceClassification.from_pretrained(rerank_model)
+
+        def get_score(query, passage):
+            """Calculate the relevance score of a passage with respect to a query."""
+
+
+            inputs = tokenizer(query, passage, return_tensors='pt', truncation=True, padding=True, max_length=512)
+            
+
+            with torch.no_grad():
+                outputs = model(**inputs)
+            
+
+            logits = outputs.logits
+            
+
+            score = logits.view(-1, ).float()  
+            
+
+            return score
+
+        scores = [get_score(instruction, c[0].page_content) for c in contexts]
+        combined = list(zip(contexts, scores))
+        sorted_combined = sorted(combined, key=lambda x: x[1], reverse=True)
+        sorted_texts, _ = zip(*sorted_combined)
+
+        return sorted_texts[:number_of_contexts]
+
     def get_context(self, instruction, number_of_contexts=2):
+        """Retrieve the most relevant contexts for a given instruction."""
+        documentos = self.vectore_store.similarity_search_with_score(instruction, k=self.rerank_number_contexts)
 
-        documentos = self.vectore_store.similarity_search_with_score(instruction, k=number_of_contexts)
+        documentos = self.rerank_contexts(instruction, documentos, number_of_contexts=number_of_contexts)
 
+        print("Reranked documents")
         return documentos
         
     def predict(self, instruction, context, model_parameters):