from typing import List

from sentence_transformers import SentenceTransformer

from base import BaseEmbedding


class LocalEmbeddings(BaseEmbedding):
    def __init__(self, model_path: str):
        self.embedding = SentenceTransformer(model_path)
        self.description = "Finetuned on `paraphrase-multilingual-mpnet-base-v2` with PingCAP docs"

    def embed(
        self, texts: List[str], prefix: str = "[QRY]", suffix: str = ""
    ) -> List[List[float]]:
        """Embed search docs."""
        return self.embedding.encode([prefix + t + suffix for t in texts]).tolist()