sergeipetrov commited on
Commit
192f6d9
·
verified ·
1 Parent(s): 973b824

Update src/vector_db.py

Browse files
Files changed (1) hide show
  1. src/vector_db.py +7 -0
src/vector_db.py CHANGED
@@ -6,6 +6,11 @@ import pyarrow as pa
6
  import pandas as pd
7
  import numpy as np
8
  import tqdm
 
 
 
 
 
9
 
10
  class VectorDB:
11
 
@@ -90,8 +95,10 @@ class VectorDB:
90
  def retrieve_prefiltered_hits(self, query, k):
91
 
92
  query_vec = self.retriever.encode(query)
 
93
  documents = self.table.search(query_vec, vector_column_name=self.vector_column).limit(k).to_list()
94
  names = [doc[self.name_column] for doc in documents]
95
  descriptions = [doc[self.description_column] for doc in documents]
 
96
 
97
  return names, descriptions
 
6
  import pandas as pd
7
  import numpy as np
8
  import tqdm
9
+ import logging
10
+
11
+ logger = logging.getLogger(__name__)
12
+ logger.setLevel(logging.DEBUG)
13
+
14
 
15
  class VectorDB:
16
 
 
95
  def retrieve_prefiltered_hits(self, query, k):
96
 
97
  query_vec = self.retriever.encode(query)
98
+ logger.info('encoded')
99
  documents = self.table.search(query_vec, vector_column_name=self.vector_column).limit(k).to_list()
100
  names = [doc[self.name_column] for doc in documents]
101
  descriptions = [doc[self.description_column] for doc in documents]
102
+ logger.info('done topK lookup')
103
 
104
  return names, descriptions