Spaces:
Running
Running
File size: 1,418 Bytes
58f4a9c 16530ae 967868b 58f4a9c 967868b 58f4a9c 16530ae 4f9495d 58f4a9c 16530ae 58f4a9c 16530ae 58f4a9c 967868b 16530ae 58f4a9c 16530ae 967868b 16530ae 58f4a9c 967868b 4f9495d 16530ae 4f9495d 58f4a9c | 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 30 31 32 33 34 35 36 37 38 39 40 41 42 | import logging
import torch
from sentence_transformers import SentenceTransformer
logger = logging.getLogger("EmbedService")
class MultiEmbeddingService:
def __init__(self):
self.models = {}
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.model_map = {
384: "./models/bge-384",
768: "./models/bge-768",
1024: "./models/bge-1024"
}
def load_all_models(self):
"""Loads all defined models into memory."""
logger.info(f"🚀 Acceleration Device: {self.device.upper()}")
for dim, path in self.model_map.items():
try:
logger.info(f"Loading {dim}-dimension model...")
model = SentenceTransformer(path, device=self.device)
model.eval()
self.models[dim] = model
except Exception as e:
logger.error(f"❌ Failed to load {dim}-dim model: {e}")
def generate_embedding(self, text, dimension):
if dimension not in self.models:
raise ValueError(f"Dimension {dimension} not supported.")
# show_progress_bar=False stops the spam
return self.models[dimension].encode(
text,
normalize_embeddings=True,
convert_to_numpy=True,
show_progress_bar=False,
batch_size=32
).tolist() |