Spaces:
Sleeping
Sleeping
File size: 715 Bytes
4722db8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
from sentence_transformers import SentenceTransformer
from typing import List
class EmbeddingGenerator:
def __init__(self, model_name: str):
self.model = SentenceTransformer(model_name)
self.dimension = self.model.get_sentence_embedding_dimension()
def generate(self, texts: List[str]) -> List[List[float]]:
"""Generate embeddings for a list of texts"""
embeddings = self.model.encode(texts, convert_to_numpy=True)
return embeddings.tolist()
def generate_single(self, text: str) -> List[float]:
"""Generate embedding for a single text"""
embedding = self.model.encode([text], convert_to_numpy=True)
return embedding[0].tolist()
|