from typing import List

from pydantic import BaseModel
from sentence_transformers import SentenceTransformer

class EmbeddingModelConfig(BaseModel):
    model_name: str
    batch_size: int


class EmbeddingModel:
    def __init__(self, config: EmbeddingModelConfig):
        self.config = config
        self._model = SentenceTransformer(self.config.model_name)

    def encode(self, texts: List[str], convert_to_tensor: bool = False):
        return self._model.encode(texts, convert_to_tensor=convert_to_tensor, batch_size=self.config.batch_size)