Spaces:
Sleeping
Sleeping
File size: 1,232 Bytes
b7f4e8c 285b311 b7f4e8c c1d292c 9751345 b7f4e8c 9751345 b7f4e8c 9751345 b7f4e8c 9751345 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
import lancedb
import os
import gradio as gr
from sentence_transformers import SentenceTransformer
db = lancedb.connect("./.lancedb")
MODEL1_STRATEGY1 = "model1_fixed"
MODEL2_STRATEGY1 = "model2_fixed"
MODEL3_STRATEGY1 = "model3_fixed"
VECTOR_COLUMN = os.getenv("VECTOR_COLUMN", "vector")
TEXT_COLUMN = os.getenv("TEXT_COLUMN", "text")
BATCH_SIZE = int(os.getenv("BATCH_SIZE", 32))
retriever = SentenceTransformer(os.getenv("EMB_MODEL"))
def get_table_name():
emb_model = os.getenv("EMB_MODEL")
if emb_model == "sentence-transformers/all-MiniLM-L6-v2":
return MODEL1_STRATEGY1
elif emb_model == "BAAI/bge-large-en-v1.5":
return MODEL2_STRATEGY1
elif emb_model == "openai/text-embedding-ada-002":
return MODEL3_STRATEGY1
else:
raise ValueError(f"Unsupported embedding model: {emb_model}")
def retrieve(query, k):
table_name = get_table_name()
TABLE = db.open_table(table_name)
query_vec = retriever.encode(query)
try:
documents = TABLE.search(query_vec, vector_column_name=VECTOR_COLUMN).limit(k).to_list()
documents = [doc[TEXT_COLUMN] for doc in documents]
return documents
except Exception as e:
raise gr.Error(str(e)) |