Update utils_st.py

utils_st.py

@@ -30,18 +30,11 @@ class SentenceTransformerRerank():
         self.top_n=top_n
         
  
-    def predict(
-        self,
-        nodes,
-        query_bundle = None,
+    def predict(self,nodes,query = None,
     ) :
-
         query_and_nodes = [
-            (
-                query_bundle,
-                node.text,
-            )
-            for node in nodes
+            (str(query),str(nodes[i].text))
+            for i in range(len(nodes))
         ]
         def predict_score(pair):
             return self.model.predict([pair])[0]
@@ -68,18 +61,19 @@ class SentenceTransformerRerank():
         return new_nodes
 
 @st.cache_resource(show_spinner=False)
-def load_data(_embed_model):   
+def load_data():   
     with open('nodes_clean.pkl', 'rb') as file:
+        embed_model, reranker=load_models()
         nodes=pickle.load( file)
         d = 768
         faiss_index = faiss.IndexFlatL2(d)
         vector_store = FaissVectorStore(faiss_index=faiss_index )
         storage_context = StorageContext.from_defaults(vector_store=vector_store)
         # use later nodes_clean
-        index = VectorStoreIndex(nodes,embed_model=_embed_model,storage_context=storage_context)
-        retriever_dense = index.as_retriever(similarity_top_k=20,embedding=True)
-        retrieverBM25 = BM25Retriever.from_defaults(nodes=nodes, similarity_top_k=5)
-        hybrid_retriever = HybridRetriever(retriever_dense, retrieverBM25)
+        index = VectorStoreIndex(nodes,embed_model=embed_model,storage_context=storage_context)
+        retriever_dense = index.as_retriever(similarity_top_k=40,embedding=True)
+        retrieverBM25 = BM25Retriever.from_defaults(nodes=nodes, similarity_top_k=15)
+        hybrid_retriever = HybridRetriever(retriever_dense, retrieverBM25,reranker)
         return hybrid_retriever
 
 
@@ -93,9 +87,10 @@ def load_models():
 
 
 class HybridRetriever(BaseRetriever):
-    def __init__(self, vector_retriever, bm25_retriever):
+    def __init__(self, vector_retriever, bm25_retriever,reranker):
         self.vector_retriever = vector_retriever
         self.bm25_retriever = bm25_retriever
+        self.reranker = reranker
         super().__init__()
     def _retrieve(self, query, **kwargs):
         with ThreadPoolExecutor() as executor:
@@ -105,13 +100,22 @@ class HybridRetriever(BaseRetriever):
         bm25_nodes = bm25_future.result()
         vector_nodes = vector_future.result()
         # combine the two lists of nodes
+        dense_n=20
+        bm25_n=5
+        combined_nodes = vector_nodes[dense_n:] + bm25_nodes[bm25_n:]
+
         all_nodes = []
         node_ids = set()
-        for n in bm25_nodes + vector_nodes:
+        for n in bm25_nodes.copy()[:bm25_n] + vector_nodes[:dense_n]:
             if n.node.node_id not in node_ids:
                 all_nodes.append(n)
                 node_ids.add(n.node.node_id)
-        return all_nodes
+        #reRank only best of retrieved_nodes
+        reranked_nodes = self.reranker.predict(
+            all_nodes,query
+            )
+        
+        return reranked_nodes+combined_nodes
 
 
 import re