update

rag.py

@@ -1,6 +1,6 @@
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.prompts import PromptTemplate
-from utility import load_data, process_data, CustomRetriever
+from utility import load_data, process_data, CustomRetriever, CustomRetriever1
 
 
 data1 = load_data('raw_data/sv')
@@ -137,9 +137,9 @@ ensemble_retriever3 = EnsembleRetriever(retrievers=[bm25_retriever3, retriever3]
 
 #########################################################################################
 
-custom_retriever1 = CustomRetriever(retriever = ensemble_retriever1)
-custom_retriever2 = CustomRetriever(retriever = ensemble_retriever2)
-custom_retriever3 = CustomRetriever(retriever = ensemble_retriever3)
+custom_retriever1 = CustomRetriever1(retriever = ensemble_retriever1)
+custom_retriever2 = CustomRetriever1(retriever = ensemble_retriever2)
+custom_retriever3 = CustomRetriever1(retriever = ensemble_retriever3)
 
 multiq_chain1 = generate_queries | custom_retriever1
 multiq_chain2 = generate_queries | custom_retriever2

utility.py

@@ -144,4 +144,70 @@ class CustomRetriever(BaseRetriever):
 
         docs_top_10 = docs[0:10]
 
-        return docs_top_10
+        return docs_top_10
+    
+
+import cohere
+COHERE_API_KEY = 'axMzubIv9l3UTObYnIaHuZhE6tR3Nj8eGReXTws9'
+
+class CustomRetriever1(BaseRetriever):
+    # vectorstores:Chroma
+    retriever:Any
+
+    def reciprocal_rank_fusion(self, results: list[list], k=60):
+        """ Reciprocal_rank_fusion that takes multiple lists of ranked documents
+            and an optional parameter k used in the RRF formula """
+
+        # Initialize a dictionary to hold fused scores for each unique document
+        fused_scores = {}
+
+        # Iterate through each list of ranked documents
+        for docs in results:
+            # Iterate through each document in the list, with its rank (position in the list)
+            for rank, doc in enumerate(docs):
+                # Convert the document to a string format to use as a key (assumes documents can be serialized to JSON)
+                doc_str = dumps(doc)
+                # If the document is not yet in the fused_scores dictionary, add it with an initial score of 0
+                if doc_str not in fused_scores:
+                    fused_scores[doc_str] = 0
+                # Retrieve the current score of the document, if any
+                previous_score = fused_scores[doc_str]
+                # Update the score of the document using the RRF formula: 1 / (rank + k)
+                fused_scores[doc_str] += 1 / (rank + k)
+
+        # Sort the documents based on their fused scores in descending order to get the final reranked results
+        reranked_results = [
+            (loads(doc), score)
+            for doc, score in sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)   #[:10] #Top 10
+        ]
+
+        # Return the reranked results as a list of tuples, each containing the document and its fused score
+        rr_list=[]
+        for doc in reranked_results:
+          rr_list.append(doc[0])
+        return rr_list[:30]
+
+    def _get_relevant_documents(
+        self, queries: list, *, run_manager: CallbackManagerForRetrieverRun
+    ) -> List[Document]:
+        # Use your existing retriever to get the documents
+        documents=[]
+        for i in range(len(queries)):
+          document = self.retriever.get_relevant_documents(queries[i], callbacks=run_manager.get_child())
+          documents.append(document)
+
+        unique_documents = self.reciprocal_rank_fusion(documents)
+
+        # Get page content
+        docs_content = []
+        for i in range(len(unique_documents)):
+          docs_content.append(unique_documents[i].page_content)
+
+        co = cohere.Client(COHERE_API_KEY)
+        results = co.rerank(query=queries[0], documents=docs_content, top_n=10, model='rerank-multilingual-v3.0', return_documents=True)
+
+        reranked_indices = [result.index for result in results.results]
+
+        sorted_documents = [unique_documents[idx] for idx in reranked_indices]
+
+        return sorted_documents