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()