File size: 437 Bytes
de292ea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 |
"""Text encoder."""
from sentence_transformers import SentenceTransformer
from torch import Tensor
class TextEncoder:
"""Text encoder."""
def __init__(self) -> None:
"""Initialize the text encoder."""
self.model = SentenceTransformer("sentence-transformers/all-MiniLM-L12-v1")
def encode(self, text: str) -> Tensor:
"""Encode a string into an embedding."""
return self.model.encode(text)
|