codekingpro/portable-devtools

1"""Wrapper around text2vec embedding models."""2 3from typing import Any, List, Optional4 5from langchain_core.embeddings import Embeddings6from pydantic import BaseModel, ConfigDict7 8 9class Text2vecEmbeddings(Embeddings, BaseModel):10    """text2vec embedding models.11 12    Install text2vec first, run 'pip install -U text2vec'.13    The github repository for text2vec is : https://github.com/shibing624/text2vec14 15    Example:16        .. code-block:: python17 18            from langchain_community.embeddings.text2vec import Text2vecEmbeddings19 20            embedding = Text2vecEmbeddings()21            embedding.embed_documents([22                "This is a CoSENT(Cosine Sentence) model.",23                "It maps sentences to a 768 dimensional dense vector space.",24            ])25            embedding.embed_query(26                "It can be used for text matching or semantic search."27            )28    """29 30    model_name_or_path: Optional[str] = None31    encoder_type: Any = "MEAN"32    max_seq_length: int = 25633    device: Optional[str] = None34    model: Any = None35 36    model_config = ConfigDict(protected_namespaces=())37 38    def __init__(39        self,40        *,41        model: Any = None,42        model_name_or_path: Optional[str] = None,43        **kwargs: Any,44    ):45        try:46            from text2vec import SentenceModel47        except ImportError as e:48            raise ImportError(49                "Unable to import text2vec, please install with "50                "`pip install -U text2vec`."51            ) from e52 53        model_kwargs = {}54        if model_name_or_path is not None:55            model_kwargs["model_name_or_path"] = model_name_or_path56        model = model or SentenceModel(**model_kwargs, **kwargs)57        super().__init__(model=model, model_name_or_path=model_name_or_path, **kwargs)58 59    def embed_documents(self, texts: List[str]) -> List[List[float]]:60        """Embed documents using the text2vec embeddings model.61 62        Args:63            texts: The list of texts to embed.64 65        Returns:66            List of embeddings, one for each text.67        """68 69        return self.model.encode(texts)70 71    def embed_query(self, text: str) -> List[float]:72        """Embed a query using the text2vec embeddings model.73 74        Args:75            text: The text to embed.76 77        Returns:78            Embeddings for the text.79        """80 81        return self.model.encode(text)82