File size: 1,097 Bytes
2aeeaf6 |
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 |
import logging
import lancedb
import os
from pathlib import Path
from sentence_transformers import SentenceTransformer
#from FlagEmbedding import LLMEmbedder, FlagReranker # Al document present here https://github.com/FlagOpen/FlagEmbedding/tree/master
#EMB_MODEL_NAME = "thenlper/gte-base"
EMB_MODEL_NAME = 'BAAI/llm-embedder'
task = "qa" # Encode for a specific task (qa, icl, chat, lrlm, tool, convsearch)
#EMB_MODEL_NAME = LLMEmbedder('BAAI/llm-embedder', use_fp16=False) # Load model (automatically use GPUs)
#reranker_model = FlagReranker('BAAI/bge-reranker-base', use_fp16=True) # use_fp16 speeds up computation with a slight performance degradation
#EMB_MODEL_NAME = "thenlper/gte-base"
#DB_TABLE_NAME = "Huggingface_docs"
DB_TABLE_NAME = "cbse"
# Setting up the logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
retriever = SentenceTransformer(EMB_MODEL_NAME)
# db
db_uri = os.path.join(Path(__file__).parents[1], "lancedb")
print(f'DB URL is {db_uri}')
db = lancedb.connect(db_uri)
table = db.open_table(DB_TABLE_NAME)
|