Spaces:
Running
Running
back to default values
Browse files
backend/semantic_search.py
CHANGED
@@ -6,17 +6,19 @@ from sentence_transformers import SentenceTransformer
|
|
6 |
|
7 |
db = lancedb.connect(".lancedb")
|
8 |
|
9 |
-
|
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))
|
13 |
|
14 |
-
|
15 |
|
16 |
|
17 |
def retrieve(query, k, table_name, emb_name):
|
18 |
-
|
19 |
-
|
|
|
|
|
20 |
query_vec = retriever.encode(query)
|
21 |
try:
|
22 |
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))
|
13 |
|
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()
|