Spaces:
Sleeping
Sleeping
File size: 849 Bytes
ed4afb1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | import numpy as np
from sentence_transformers import SentenceTransformer
class Embedder:
def __init__(self, model_name: str):
self.model = SentenceTransformer(model_name)
def embed_query(self, query: str) -> np.ndarray:
embedding = self.model.encode(query, normalize_embeddings=True)
return np.array(embedding, dtype=np.float32).reshape(1, -1)
def embed_datasets(self, datasets: list[dict]) -> np.ndarray:
if not datasets:
return np.array([]).reshape(0, self.model.get_sentence_embedding_dimension())
texts = [
f"{ds.get('name', '')}. {ds.get('description', '')}"
for ds in datasets
]
embeddings = self.model.encode(texts, normalize_embeddings=True)
return np.array(embeddings, dtype=np.float32) |