z00mP commited on
Commit
f46a74d
1 Parent(s): 98a9e92
Files changed (1) hide show
  1. backend/semantic_search.py +3 -3
backend/semantic_search.py CHANGED
@@ -6,7 +6,7 @@ from sentence_transformers import SentenceTransformer
6
 
7
  db = lancedb.connect(".lancedb")
8
 
9
- TABLE = db.open_table(os.getenv("TABLE_NAME"))
10
  VECTOR_COLUMN = os.getenv("VECTOR_COLUMN", "vector")
11
  TEXT_COLUMN = os.getenv("TEXT_COLUMN", "text")
12
  BATCH_SIZE = int(os.getenv("BATCH_SIZE", 32))
@@ -14,11 +14,11 @@ BATCH_SIZE = int(os.getenv("BATCH_SIZE", 32))
14
  retriever = SentenceTransformer(os.getenv("EMB_MODEL"))
15
 
16
 
17
- def retrieve(query, k, table_name, emb_name):
18
  #print(table_name)
19
  #print(emb_name)
20
  TABLE = db.open_table(table_name)
21
- retriever = SentenceTransformer(emb_name)
22
  query_vec = retriever.encode(query)
23
  try:
24
  documents = TABLE.search(query_vec, vector_column_name=VECTOR_COLUMN).limit(k).to_list()
 
6
 
7
  db = lancedb.connect(".lancedb")
8
 
9
+ #TABLE = db.open_table(os.getenv("TABLE_NAME"))
10
  VECTOR_COLUMN = os.getenv("VECTOR_COLUMN", "vector")
11
  TEXT_COLUMN = os.getenv("TEXT_COLUMN", "text")
12
  BATCH_SIZE = int(os.getenv("BATCH_SIZE", 32))
 
14
  retriever = SentenceTransformer(os.getenv("EMB_MODEL"))
15
 
16
 
17
+ def retrieve(query, k, table_name, embedding_model_name):
18
  #print(table_name)
19
  #print(emb_name)
20
  TABLE = db.open_table(table_name)
21
+ retriever = SentenceTransformer(embedding_model_name)
22
  query_vec = retriever.encode(query)
23
  try:
24
  documents = TABLE.search(query_vec, vector_column_name=VECTOR_COLUMN).limit(k).to_list()