gradio_attempt_3 / backend /semantic_search.py
ivan
initial add
2febf52
raw
history blame contribute delete
714 Bytes
import logging
import lancedb
import os
from pathlib import Path
from sentence_transformers import SentenceTransformer
from sentence_transformers import CrossEncoder
#EMB_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
EMB_MODEL_NAME = "thenlper/gte-large"
RANKER_MODEL_NAME ='cross-encoder/ms-marco-TinyBERT-L-2-v2'
DB_TABLE_NAME = "db_table"
# Setting up the logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
retriever = SentenceTransformer(EMB_MODEL_NAME)
ranker = CrossEncoder(RANKER_MODEL_NAME, max_length=512)
# db
db_uri = os.path.join(Path(__file__).parents[1], ".lancedb_large_model_256_32")
db = lancedb.connect(db_uri)
table = db.open_table(DB_TABLE_NAME)