from dataclasses import dataclass

import numpy as np
from sentence_transformers import SentenceTransformer


@dataclass(frozen=True)
class Encoder:
    st: SentenceTransformer
    truncation_ratio: float

    @classmethod
    def from_model_name(
        cls, model_name: str, truncation_ratio: float = 0.0
    ) -> "Encoder":
        st = SentenceTransformer(model_name)
        return cls(st, truncation_ratio)

    @property
    def truncate_dim(self) -> int:
        return int(self.st.get_sentence_embedding_dimension() * self.truncation_ratio)

    def encode(self, docs: list[str], is_query: bool) -> np.ndarray:
        prompt_name = "query" if is_query else None
        embs = self.st.encode(
            docs, prompt_name=prompt_name, truncate_dim=self.truncate_dim
        )
        return embs
