SkyOSFullCore / core /embeddings.py
sky-meilin's picture
Update core/embeddings.py
7f37031 verified
Raw
History Blame Contribute Delete
1.07 kB
from sentence_transformers import SentenceTransformer
from typing import List, Union
import numpy as np
class Embedder:
def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
# Nutzt verfügbare Hardware (Cuda/CPU) automatisch
self.model = SentenceTransformer(model_name)
def encode(self, texts: Union[str, List[str]], normalize: bool = True) -> np.ndarray:
"""
Encodes text(s) into embeddings.
Args:
texts: single string or list of strings
normalize: optional normalization for cosine similarity use
Returns:
numpy.ndarray embeddings
"""
if not texts:
return np.array([])
if isinstance(texts, str):
if not texts.strip():
return np.array([])
texts = [texts]
# Konvertiert leere Strings in der Liste, um Crashs zu vermeiden
texts = [t if t.strip() else " " for t in texts]
embeddings = self.model.encode(texts, normalize_embeddings=normalize)
return embeddings