from typing import List
from langchain.embeddings.base import Embeddings
from sentence_transformers import SentenceTransformer


class PingChatEmbeddings(Embeddings):
    def __init__(self, model_path: str):
        self.embedding = SentenceTransformer(model_path)

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """Embed search docs."""
        return self.embedding.encode(['[DOC]' + t for t in texts]).tolist()

    def embed_query(self, text: str) -> List[float]:
        """Embed query text."""
        return self.embedding.encode('[QRY]' + text).tolist()