File size: 502 Bytes
f4c70c8
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from sentence_transformers import SentenceTransformer
import torch

class Embedder:
    def __init__(self, model_name: str = "all-MiniLM-L6-v2", device: str = None):
        if device is None:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device
            
        self.model = SentenceTransformer(model_name, device=self.device)

    def embed(self, texts: list[str]):
        return self.model.encode(texts, convert_to_numpy=True)