LLM_Model / embedding_model_instance.py
Shreekant Kalwar (Nokia)
major changess
28b14ff
raw
history blame
440 Bytes
import torch
from sentence_transformers import SentenceTransformer, CrossEncoder
# --- Embedding model
EMBEDDING_MODEL = "BAAI/bge-base-en-v1.5"
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
embedding_model = SentenceTransformer(EMBEDDING_MODEL, device=device)
embedding_dim = embedding_model.get_sentence_embedding_dimension()
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")