File size: 313 Bytes
adad4ac
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
from sentence_transformers import SentenceTransformer
import numpy as np
from typing import Sequence

class Vectorizer:
    def __init__(self, model) -> None:
        self.model = SentenceTransformer(model)

    def transform(self, prompts: Sequence[str]) -> np.ndarray:
        return self.model.encode(prompts)