Spaces:
Sleeping
Sleeping
File size: 1,156 Bytes
84e25a6 144bde6 84e25a6 fe8b339 84e25a6 |
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
import numpy as np
from sentence_transformers import SentenceTransformer, CrossEncoder
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"))
reranker = CrossEncoder(os.getenv("RERANK_MODEL"), max_length=512)
def retrieve(query, n):
query_vec = retriever.encode(query)
try:
documents = TABLE.search(query_vec, vector_column_name=VECTOR_COLUMN).limit(n).to_list()
documents = [doc[TEXT_COLUMN] for doc in documents]
return documents
except Exception as e:
raise gr.Error(str(e))
def rerank(query, documents, k):
query_doc_pairs = [[query, doc] for doc in documents]
similarity_scores = reranker.predict(query_doc_pairs)
sim_scores_argsort = np.argsort(similarity_scores)[::-1]
rerank_documents = []
for idx in sim_scores_argsort[:k]:
rerank_documents.append(documents[idx])
return rerank_documents
|