Spaces:
Sleeping
Sleeping
File size: 645 Bytes
d0abef8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
# src/embed_service/embedder.py
from sentence_transformers import SentenceTransformer
import numpy as np
class Embedder:
def __init__(self, model_name="all-MiniLM-L6-v2"):
print(f"Loading embedding model: {model_name}")
self.model = SentenceTransformer(model_name)
def embed_text(self, text: str):
emb = self.model.encode(text, convert_to_numpy=True)
return emb.astype("float32")
def embed_batch(self, texts: list):
embs = self.model.encode(texts, convert_to_numpy=True)
return embs.astype("float32")
def dim(self):
return self.model.get_sentence_embedding_dimension()
|