Spaces:
Sleeping
Sleeping
| from sentence_transformers import SentenceTransformer | |
| from typing import List | |
| import numpy as np | |
| class EmbeddingModel: | |
| """Lightweight sentence-transformer for semantic similarity""" | |
| def __init__(self, model_name: str = "paraphrase-MiniLM-L3-v2"): | |
| # Use smallest model (3 layers, ~60MB, ~100MB RAM) | |
| self.model = SentenceTransformer(model_name) | |
| self.model_name = model_name | |
| self.dimension = self.model.get_sentence_embedding_dimension() | |
| def encode(self, texts: List[str], batch_size: int = 32) -> np.ndarray: | |
| """Generate embeddings for a list of texts""" | |
| if isinstance(texts, str): | |
| texts = [texts] | |
| cleaned_texts = [] | |
| for text in texts: | |
| if text and isinstance(text, str): | |
| cleaned = ' '.join(text.split()) | |
| if cleaned.strip(): | |
| cleaned_texts.append(cleaned) | |
| if not cleaned_texts: | |
| return np.array([]) | |
| embeddings = self.model.encode(cleaned_texts, batch_size=batch_size, convert_to_numpy=True) | |
| return embeddings | |
| def encode_single(self, text: str) -> np.ndarray: | |
| """Generate embedding for a single text""" | |
| cleaned = ' '.join(text.split()) if text else "" | |
| if not cleaned.strip(): | |
| return np.array([]) | |
| embedding = self.model.encode([cleaned], convert_to_numpy=True) | |
| return embedding[0] | |