Spaces:
Running
Running
| import lancedb | |
| import os | |
| import gradio as gr | |
| from sentence_transformers import SentenceTransformer | |
| db = lancedb.connect(".lancedb") | |
| #TABLE = db.open_table(os.getenv("TABLE_NAME")) | |
| 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 retrieve(query, k, table_name, embedding_model_name): | |
| #print(table_name) | |
| #print(emb_name) | |
| TABLE = db.open_table(table_name) | |
| retriever = SentenceTransformer(embedding_model_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)) | |