rag_project / backend /semantic_search.py
sashamn's picture
refactor && more choices
a76726b
raw
history blame contribute delete
726 Bytes
import json
import logging
import lancedb
import os
from pathlib import Path
from sentence_transformers import SentenceTransformer, CrossEncoder
with open(os.path.join(os.path.dirname(__file__), 'params.json')) as f_in:
params = json.load(f_in)
# Setting up the logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
retrievers = {
model_name: SentenceTransformer(model_name)
for model_name in params['model_names']
}
# db
db_uri = os.path.join(Path(__file__).parents[1], ".lancedb")
db = lancedb.connect(db_uri)
tables = {
table_name: db.open_table(table_name)
for table_name in params['table_names']
}
cross_model = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-12-v2')