codekingpro/portable-devtools

1from typing import Any, List, Optional2 3from langchain_core.embeddings import Embeddings4from pydantic import BaseModel, ConfigDict5 6 7class ModelScopeEmbeddings(BaseModel, Embeddings):8    """ModelScopeHub embedding models.9 10    To use, you should have the ``modelscope`` python package installed.11 12    Example:13        .. code-block:: python14 15            from langchain_community.embeddings import ModelScopeEmbeddings16            model_id = "damo/nlp_corom_sentence-embedding_english-base"17            embed = ModelScopeEmbeddings(model_id=model_id, model_revision="v1.0.0")18    """19 20    embed: Any = None21    model_id: str = "damo/nlp_corom_sentence-embedding_english-base"22    """Model name to use."""23    model_revision: Optional[str] = None24 25    def __init__(self, **kwargs: Any):26        """Initialize the modelscope"""27        super().__init__(**kwargs)28        try:29            from modelscope.pipelines import pipeline30            from modelscope.utils.constant import Tasks31        except ImportError as e:32            raise ImportError(33                "Could not import some python packages."34                "Please install it with `pip install modelscope`."35            ) from e36        self.embed = pipeline(37            Tasks.sentence_embedding,38            model=self.model_id,39            model_revision=self.model_revision,40        )41 42    model_config = ConfigDict(extra="forbid", protected_namespaces=())43 44    def embed_documents(self, texts: List[str]) -> List[List[float]]:45        """Compute doc embeddings using a modelscope embedding model.46 47        Args:48            texts: The list of texts to embed.49 50        Returns:51            List of embeddings, one for each text.52        """53        texts = list(map(lambda x: x.replace("\n", " "), texts))54        inputs = {"source_sentence": texts}55        embeddings = self.embed(input=inputs)["text_embedding"]56        return embeddings.tolist()57 58    def embed_query(self, text: str) -> List[float]:59        """Compute query embeddings using a modelscope embedding model.60 61        Args:62            text: The text to embed.63 64        Returns:65            Embeddings for the text.66        """67        text = text.replace("\n", " ")68        inputs = {"source_sentence": [text]}69        embedding = self.embed(input=inputs)["text_embedding"][0]70        return embedding.tolist()71