File size: 878 Bytes
8f17b3d
 
 
 
 
 
 
 
f46a74d
8f17b3d
 
 
 
97723c6
8f17b3d
 
f46a74d
9f3ddbf
 
 
f46a74d
8f17b3d
 
 
 
 
 
 
 
 
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
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))