Source code for langchain_community.embeddings.text2vec

"""封装了text2vec嵌入模型。"""

from typing import Any, List, Optional

from langchain_core.embeddings import Embeddings
from langchain_core.pydantic_v1 import BaseModel


class Text2vecEmbeddings(Embeddings, BaseModel):
    """text2vec嵌入模型。

    首先安装text2vec，运行'pip install -U text2vec'。
    text2vec的gitbub存储库为：https://github.com/shibing624/text2vec

    示例：
        .. code-block:: python

            from langchain_community.embeddings.text2vec import Text2vecEmbeddings

            embedding = Text2vecEmbeddings()
            embedding.embed_documents([
                "This is a CoSENT(Cosine Sentence) model.",
                "It maps sentences to a 768 dimensional dense vector space.",
            ])
            embedding.embed_query(
                "It can be used for text matching or semantic search."
            )
"""

    model_name_or_path: Optional[str] = None
    encoder_type: Any = "MEAN"
    max_seq_length: int = 256
    device: Optional[str] = None
    model: Any = None

    def __init__(
        self,
        *,
        model: Any = None,
        model_name_or_path: Optional[str] = None,
        **kwargs: Any,
    ):
        try:
            from text2vec import SentenceModel
        except ImportError as e:
            raise ImportError(
                "Unable to import text2vec, please install with "
                "`pip install -U text2vec`."
            ) from e

        model_kwargs = {}
        if model_name_or_path is not None:
            model_kwargs["model_name_or_path"] = model_name_or_path
        model = model or SentenceModel(**model_kwargs, **kwargs)
        super().__init__(model=model, model_name_or_path=model_name_or_path, **kwargs)

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """使用text2vec嵌入模型嵌入文档。

参数：
    texts：要嵌入的文本列表。

返回：
    嵌入列表，每个文本对应一个嵌入。
"""

        return self.model.encode(texts)

    def embed_query(self, text: str) -> List[float]:
        """嵌入一个使用text2vec嵌入模型的查询。

参数：
    text: 要嵌入的文本。

返回：
    文本的嵌入。
"""

        return self.model.encode(text)