Spaces:
Running
Running
| from typing import List | |
| from sentence_transformers import SentenceTransformer | |
| class Encoder(): | |
| def __init__(self): | |
| print("Loading embedding model...") | |
| self.model = SentenceTransformer( | |
| "KaLM-Embedding/KaLM-embedding-multilingual-mini-instruct-v2.5", | |
| model_kwargs={"attn_implementation": "sdpa"} | |
| ) | |
| self.model = self.model.half() | |
| def encode( | |
| self, | |
| texts: List[str], | |
| batch_size: int = 8, | |
| show_progress_bar: bool = False, | |
| save_path: str = None): | |
| embeddings = self.model.encode(texts, convert_to_tensor=True, show_progress_bar=show_progress_bar, batch_size=batch_size) | |
| # if save_path: | |
| # torch.save(embeddings, save_path) | |
| return embeddings |