codekingpro/portable-devtools

1import os2import sys3from typing import Any, List4 5from langchain_core.embeddings import Embeddings6from pydantic import BaseModel, ConfigDict7 8 9class JohnSnowLabsEmbeddings(BaseModel, Embeddings):10    """JohnSnowLabs embedding models11 12    To use, you should have the ``johnsnowlabs`` python package installed.13    Example:14        .. code-block:: python15 16            from langchain_community.embeddings.johnsnowlabs import JohnSnowLabsEmbeddings17 18            embedding = JohnSnowLabsEmbeddings(model='embed_sentence.bert')19            output = embedding.embed_query("foo bar")20    """  # noqa: E50121 22    model: Any = "embed_sentence.bert"23 24    def __init__(25        self,26        model: Any = "embed_sentence.bert",27        hardware_target: str = "cpu",28        **kwargs: Any,29    ):30        """Initialize the johnsnowlabs model."""31        super().__init__(**kwargs)32        # 1) Check imports33        try:34            from johnsnowlabs import nlp35            from nlu.pipe.pipeline import NLUPipeline36        except ImportError as exc:37            raise ImportError(38                "Could not import johnsnowlabs python package. "39                "Please install it with `pip install johnsnowlabs`."40            ) from exc41 42        # 2) Start a Spark Session43        try:44            os.environ["PYSPARK_PYTHON"] = sys.executable45            os.environ["PYSPARK_DRIVER_PYTHON"] = sys.executable46            nlp.start(hardware_target=hardware_target)47        except Exception as exc:48            raise Exception("Failure starting Spark Session") from exc49 50        # 3) Load the model51        try:52            if isinstance(model, str):53                self.model = nlp.load(model)54            elif isinstance(model, NLUPipeline):55                self.model = model56            else:57                self.model = nlp.to_nlu_pipe(model)58        except Exception as exc:59            raise Exception("Failure loading model") from exc60 61    model_config = ConfigDict(62        extra="forbid",63    )64 65    def embed_documents(self, texts: List[str]) -> List[List[float]]:66        """Compute doc embeddings using a JohnSnowLabs transformer model.67 68        Args:69            texts: The list of texts to embed.70 71        Returns:72            List of embeddings, one for each text.73        """74 75        df = self.model.predict(texts, output_level="document")76        emb_col = None77        for c in df.columns:78            if "embedding" in c:79                emb_col = c80        return [vec.tolist() for vec in df[emb_col].tolist()]81 82    def embed_query(self, text: str) -> List[float]:83        """Compute query embeddings using a JohnSnowLabs transformer model.84 85        Args:86            text: The text to embed.87 88        Returns:89            Embeddings for the text.90        """91        return self.embed_documents([text])[0]92