File size: 440 Bytes
28b14ff |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 |
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") |