ramouch/English-Encoder

1from sentence_transformers import SentenceTransformer2import numpy as np3from typing import List, Union4 5class EnglishSentenceEncoder:6    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):7        self.device = "cpu"8        self.model = SentenceTransformer(model_name, device=self.device)9 10    def encode(self, sentences: Union[str, List[str]]) -> np.ndarray:11        if isinstance(sentences, str):12            sentences = [sentences]13        return self.model.encode(sentences, convert_to_numpy=True)14 15    def __call__(self, sentences: Union[str, List[str]]) -> np.ndarray:16        return self.encode(sentences)