Source code for langchain_community.embeddings.clarifai

import logging
from typing import Any, Dict, List, Optional

from langchain_core.embeddings import Embeddings
from langchain_core.pydantic_v1 import BaseModel, Extra, Field, root_validator

logger = logging.getLogger(__name__)


class ClarifaiEmbeddings(BaseModel, Embeddings):
    """Clarifai嵌入模型。

    要使用，您应该安装``clarifai`` python包，并设置环境变量``CLARIFAI_PAT``为您的个人访问令牌，或将其作为命名参数传递给构造函数。

    示例:
        .. code-block:: python

            from langchain_community.embeddings import ClarifaiEmbeddings
            clarifai = ClarifaiEmbeddings(user_id=USER_ID,
                                          app_id=APP_ID,
                                          model_id=MODEL_ID)
                             (或)
            Example_URL = "https://clarifai.com/clarifai/main/models/BAAI-bge-base-en-v15"
            clarifai = ClarifaiEmbeddings(model_url=EXAMPLE_URL)"""

    model_url: Optional[str] = None
    """要使用的模型URL。"""
    model_id: Optional[str] = None
    """要使用的模型ID。"""
    model_version_id: Optional[str] = None
    """要使用的模型版本ID。"""
    app_id: Optional[str] = None
    """用于Clarifai应用程序的应用程序ID。"""
    user_id: Optional[str] = None
    """用于Clarifai的用户ID。"""
    pat: Optional[str] = Field(default=None, exclude=True)
    """用于使用Clarifai的个人访问令牌。"""
    token: Optional[str] = Field(default=None, exclude=True)
    """用于使用Clarifai会话令牌。"""
    model: Any = Field(default=None, exclude=True)  #: :meta private:
    api_base: str = "https://api.clarifai.com"

    class Config:
        """此pydantic对象的配置。"""

        extra = Extra.forbid

    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """验证我们是否具有访问Clarifai平台所需的所有必要信息，并且Python包存在于环境中。
"""

        try:
            from clarifai.client.model import Model
        except ImportError:
            raise ImportError(
                "Could not import clarifai python package. "
                "Please install it with `pip install clarifai`."
            )
        user_id = values.get("user_id")
        app_id = values.get("app_id")
        model_id = values.get("model_id")
        model_version_id = values.get("model_version_id")
        model_url = values.get("model_url")
        api_base = values.get("api_base")
        pat = values.get("pat")
        token = values.get("token")

        values["model"] = Model(
            url=model_url,
            app_id=app_id,
            user_id=user_id,
            model_version=dict(id=model_version_id),
            pat=pat,
            token=token,
            model_id=model_id,
            base_url=api_base,
        )

        return values

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """调用Clarifai的嵌入模型。

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

返回：
    每个文本的嵌入列表。
"""
        from clarifai.client.input import Inputs

        input_obj = Inputs.from_auth_helper(self.model.auth_helper)
        batch_size = 32
        embeddings = []

        try:
            for i in range(0, len(texts), batch_size):
                batch = texts[i : i + batch_size]
                input_batch = [
                    input_obj.get_text_input(input_id=str(id), raw_text=inp)
                    for id, inp in enumerate(batch)
                ]
                predict_response = self.model.predict(input_batch)
                embeddings.extend(
                    [
                        list(output.data.embeddings[0].vector)
                        for output in predict_response.outputs
                    ]
                )

        except Exception as e:
            logger.error(f"Predict failed, exception: {e}")

        return embeddings

    def embed_query(self, text: str) -> List[float]:
        """调用Clarifai的嵌入模型。

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

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

        try:
            predict_response = self.model.predict_by_bytes(
                bytes(text, "utf-8"), input_type="text"
            )
            embeddings = [
                list(op.data.embeddings[0].vector) for op in predict_response.outputs
            ]

        except Exception as e:
            logger.error(f"Predict failed, exception: {e}")

        return embeddings[0]