Spaces:
Sleeping
Sleeping
File size: 1,690 Bytes
7e92170 b3c39b6 7e92170 b3c39b6 7e92170 b3c39b6 7e92170 b3c39b6 7e92170 4f4776c 7e92170 b3c39b6 7e92170 |
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 41 42 43 44 45 46 |
import lancedb
import gradio as gr
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import os
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))
CROSS_ENCODER = os.getenv("CROSS_ENCODER")
retriever = SentenceTransformer(os.getenv("EMB_MODEL"))
cross_encoder = AutoModelForSequenceClassification.from_pretrained(CROSS_ENCODER)
cross_encoder.eval()
cross_encoder_tokenizer = AutoTokenizer.from_pretrained(CROSS_ENCODER)
def reranking(query, list_of_documents, k):
received_tokens = cross_encoder_tokenizer([query] * len(list_of_documents), list_of_documents, padding=True, truncation=True, return_tensors="pt")
with torch.no_grad():
logits_on_tokens = cross_encoder(**received_tokens).logits
probabilities = logits_on_tokens.reshape(-1).tolist()
documents = sorted(zip(list_of_documents, probabilities), key=lambda x: x[1], reverse=True)
result = [document[0] for document in documents[:k]]
return result
def retrieve(query, top_k_retriever=30, use_reranking=True, top_k_reranker=5):
query_vec = retriever.encode(query)
try:
documents = TABLE.search(query_vec, vector_column_name=VECTOR_COLUMN).limit(top_k_retriever).to_list()
documents = [doc[TEXT_COLUMN] for doc in documents]
if use_reranking:
documents = reranking(query, documents, top_k_reranker)
return documents
except Exception as e:
raise gr.Error(str(e))
|