Spaces:
Sleeping
Sleeping
| from sentence_transformers import SentenceTransformer | |
| from typing import List | |
| import numpy as np | |
| from src.config import config | |
| class EmbeddingService: | |
| def __init__(self, model_name: str = None): | |
| self.model_name = model_name or config.EMBEDDING_MODEL | |
| self.model = None | |
| def load_model(self): | |
| if self.model is None: | |
| print(f"Loading embedding model: {self.model_name}") | |
| self.model = SentenceTransformer(self.model_name) | |
| print(f"Model loaded successfully") | |
| def embed_text(self, text: str) -> List[float]: | |
| self.load_model() | |
| embedding = self.model.encode(text, convert_to_numpy=True) | |
| return embedding.tolist() | |
| def embed_batch(self, texts: List[str]) -> List[List[float]]: | |
| self.load_model() | |
| embeddings = self.model.encode(texts, convert_to_numpy=True, show_progress_bar=True) | |
| return embeddings.tolist() | |
| def get_embedding_dimension(self) -> int: | |
| self.load_model() | |
| return self.model.get_sentence_embedding_dimension() | |
| embedding_service = EmbeddingService() | |