Create HybridRetriever.py

HybridRetriever.py

@@ -0,0 +1,62 @@
+from llama_index.retrievers.bm25 import BM25Retriever 
+from llama_index.core.retrievers import VectorIndexRetriever
+from llama_index.core import Document
+
+class HybridRetriever:
+    def __init__(self, bm25_retriever: BM25Retriever, vector_retriever: VectorIndexRetriever):
+        """
+        Inıtializes a Hybrid Retriever with BM25Retriever and VectorIndexRetriever.
+
+        Args:
+            bm25_retriever (BM25Retriever): An instance of BM25Retriever for keyword-based retrieval.
+            vector_retriever (VectorIndexRetriever): An instance of VectorIndexRetriever for vector-based retrieval.
+        """
+        
+        self.bm25_retriever = bm25_retriever
+        self.vector_retriever = vector_retriever
+        self.top_k = vector_retriever._similarity_top_k + bm25_retriever._similarity_top_k
+
+    def retrieve(self, query: str):
+        """
+        Retrieves documents relevant to the query using both BM25 and vector retrieval methods.
+        
+        Args:
+            query (str): The query string for which relevant documents are to be retrieved.
+        
+        Returns:
+            list: A list of tuples, each containing the document text and its combined score.
+        """
+        query = "[INST] " + " [/INST]"
+        # Perform keyword search using BM25 retriever
+        bm25_results = self.bm25_retriever.retrieve(query)
+        
+        # Perform vector search using VectorIndexRetriever
+        vector_results = self.vector_retriever.retrieve(query)
+        
+        # Combine results, filter duplicates, and calculate combined scores
+        combined_results = {}
+        for result in bm25_results:
+            combined_results[result.node.text] = {'score': result.score}
+        
+        for result in vector_results:
+            if result.node.text in combined_results:
+                combined_results[result.node.text]['score'] += result.score
+            else:
+                combined_results[result.node.text] = {'score': result.score}
+
+        # Convert combined results to a list of tuples and sort by score
+        combined_results_list = sorted(combined_results.items(), key=lambda item: item[1]['score'], reverse=True)
+        return combined_results_list # {score, document}
+        
+    def best_docs(self, query: str):
+        """
+        Retrieves the most relevant documents to the query as Document objects with their scores.
+
+        Args:
+            query (str): The query string for which the most relevant documents are to be retrieved.
+
+        Returns:
+            list: A list of tuples, each containing a Document object and its score.
+        """
+        top_results = self.retrieve(query)
+        return [(Document(text=text), score) for text, score in top_results]