Source code for langchain_community.embeddings.huggingface_hub

import json
import os
from typing import Any, Dict, List, Optional

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

DEFAULT_MODEL = "sentence-transformers/all-mpnet-base-v2"
VALID_TASKS = ("feature-extraction",)


class HuggingFaceHubEmbeddings(BaseModel, Embeddings):
    """HuggingFaceHub嵌入模型。

    要使用，您应该已安装``huggingface_hub`` python包，并且
    将环境变量``HUGGINGFACEHUB_API_TOKEN``设置为您的API令牌，或者传递
    它作为构造函数的一个命名参数。

    示例:
        .. code-block:: python

            from langchain_community.embeddings import HuggingFaceHubEmbeddings
            model = "sentence-transformers/all-mpnet-base-v2"
            hf = HuggingFaceHubEmbeddings(
                model=model,
                task="feature-extraction",
                huggingfacehub_api_token="my-api-key",
            )
    """

    client: Any  #: :meta private:
    async_client: Any  #: :meta private:
    model: Optional[str] = None
    """要使用的模型名称。"""
    repo_id: Optional[str] = None
    """Huggingfacehub存储库ID，用于向后兼容。"""
    task: Optional[str] = "feature-extraction"
    """调用模型的任务。"""
    model_kwargs: Optional[dict] = None
    """传递给模型的关键字参数。"""

    huggingfacehub_api_token: Optional[str] = None

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

        extra = Extra.forbid

    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """验证环境中是否存在API密钥和Python包。"""
        huggingfacehub_api_token = values["huggingfacehub_api_token"] or os.getenv(
            "HUGGINGFACEHUB_API_TOKEN"
        )

        try:
            from huggingface_hub import AsyncInferenceClient, InferenceClient

            if values["model"]:
                values["repo_id"] = values["model"]
            elif values["repo_id"]:
                values["model"] = values["repo_id"]
            else:
                values["model"] = DEFAULT_MODEL
                values["repo_id"] = DEFAULT_MODEL

            client = InferenceClient(
                model=values["model"],
                token=huggingfacehub_api_token,
            )

            async_client = AsyncInferenceClient(
                model=values["model"],
                token=huggingfacehub_api_token,
            )

            if values["task"] not in VALID_TASKS:
                raise ValueError(
                    f"Got invalid task {values['task']}, "
                    f"currently only {VALID_TASKS} are supported"
                )
            values["client"] = client
            values["async_client"] = async_client

        except ImportError:
            raise ImportError(
                "Could not import huggingface_hub python package. "
                "Please install it with `pip install huggingface_hub`."
            )
        return values

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """调用HuggingFaceHub的嵌入端点来进行嵌入搜索文档。

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

返回：
    每个文本的嵌入列表。
"""
        # replace newlines, which can negatively affect performance.
        texts = [text.replace("\n", " ") for text in texts]
        _model_kwargs = self.model_kwargs or {}
        responses = self.client.post(
            json={"inputs": texts, "parameters": _model_kwargs}, task=self.task
        )
        return json.loads(responses.decode())

    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
        """异步调用HuggingFaceHub的嵌入端点以进行嵌入搜索文档。

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

返回：
    嵌入列表，每个文本对应一个嵌入。
"""
        # replace newlines, which can negatively affect performance.
        texts = [text.replace("\n", " ") for text in texts]
        _model_kwargs = self.model_kwargs or {}
        responses = await self.async_client.post(
            json={"inputs": texts, "parameters": _model_kwargs}, task=self.task
        )
        return json.loads(responses.decode())

    def embed_query(self, text: str) -> List[float]:
        """调用HuggingFaceHub的嵌入端点来嵌入查询文本。

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

返回：
    文本的嵌入。
"""
        response = self.embed_documents([text])[0]
        return response

    async def aembed_query(self, text: str) -> List[float]:
        """异步调用HuggingFaceHub的嵌入端点，用于嵌入查询文本。

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

返回：
    文本的嵌入。
"""
        response = (await self.aembed_documents([text]))[0]
        return response