rag-transformers / backend /semantic_search.py
waleko's picture
Update db; Add Reranking
322806a
raw
history blame
1.28 kB
import lancedb
import os
import gradio as gr
from sentence_transformers import SentenceTransformer
from FlagEmbedding import FlagReranker
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 = FlagReranker(os.getenv("RERANKER_MODEL", 'BAAI/bge-reranker-large'), use_fp16=True)
def retrieve(query, k):
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))
def rerank(documents, query, k):
try:
query_pairs = [[query, doc] for doc in documents]
scores = reranker.compute_score(query_pairs)
scored_documents = list(zip(documents, scores))
scored_documents.sort(key=lambda x: x[1], reverse=True)
top_k_documents = [doc for doc, _ in scored_documents[:k]]
return top_k_documents
except Exception as e:
raise gr.Error(str(e))