Update src/vector_db.py

src/vector_db.py

@@ -17,9 +17,7 @@ class VectorDB:
     db_location = ''
 
     def __init__(self, emb_model, db_location, actions_list_file_path, num_sub_vectors, batch_size):
-      self.emb_model = emb_model
-      self.db_location = db_location
-
+      self.retriever = SentenceTransformer(emb_model)
       emb_config = AutoConfig.from_pretrained(emb_model)
       emb_dimension = emb_config.hidden_size
 
@@ -50,7 +48,7 @@ class VectorDB:
             pa.field(self.name_column, pa.string())
         ]
       )
-      tbl = db.create_table(self.table_name, schema=schema, mode="overwrite")
+      tbl = db.create_table(table_name, schema=schema, mode="overwrite")
 
 
       df = pd.read_csv(actions_list_file_path)
@@ -76,23 +74,23 @@ class VectorDB:
               tbl.add(df)
           except:
               print(f"batch {i} was skipped")
+      
+      self.db = db
+      self.table = tbl
       print("Vector generation done.")
 
 
-    def get_embedding_db_as_pandas(self):
-        db = lancedb.connect(self.db_location)
-        tbl = db.open_table(self.table_name)
-        return tbl.to_pandas()
+    # def get_embedding_db_as_pandas(self):
+    #     db = lancedb.connect(self.db_location)
+    #     tbl = db.open_table(self.table_name)
+    #     return tbl.to_pandas()
 
 
 
     def retrieve_prefiltered_hits(self, query, k):
-        db = lancedb.connect(".lancedb")
-        table = db.open_table(self.table_name)
-        retriever = SentenceTransformer(self.emb_model)
 
-        query_vec = retriever.encode(query)
-        documents = table.search(query_vec, vector_column_name=self.vector_column).limit(k).to_list()
+        query_vec = self.retriever.encode(query)
+        documents = self.table.search(query_vec, vector_column_name=self.vector_column).limit(k).to_list()
         names = [doc[self.name_column] for doc in documents]
         descriptions = [doc[self.description_column] for doc in documents]