Source code for langchain_community.embeddings.self_hosted_hugging_face

import importlib
import logging
from typing import Any, Callable, List, Optional

from langchain_community.embeddings.self_hosted import SelfHostedEmbeddings

DEFAULT_MODEL_NAME = "sentence-transformers/all-mpnet-base-v2"
DEFAULT_INSTRUCT_MODEL = "hkunlp/instructor-large"
DEFAULT_EMBED_INSTRUCTION = "Represent the document for retrieval: "
DEFAULT_QUERY_INSTRUCTION = (
    "Represent the question for retrieving supporting documents: "
)

logger = logging.getLogger(__name__)


def _embed_documents(client: Any, *args: Any, **kwargs: Any) -> List[List[float]]:
    """推断函数，用于发送到远程硬件。

接受一个 sentence_transformer 模型标识，并返回批处理中每个文档的嵌入列表。
"""
    return client.encode(*args, **kwargs)


def load_embedding_model(model_id: str, instruct: bool = False, device: int = 0) -> Any:
    """加载嵌入模型。"""
    if not instruct:
        import sentence_transformers

        client = sentence_transformers.SentenceTransformer(model_id)
    else:
        from InstructorEmbedding import INSTRUCTOR

        client = INSTRUCTOR(model_id)

    if importlib.util.find_spec("torch") is not None:
        import torch

        cuda_device_count = torch.cuda.device_count()
        if device < -1 or (device >= cuda_device_count):
            raise ValueError(
                f"Got device=={device}, "
                f"device is required to be within [-1, {cuda_device_count})"
            )
        if device < 0 and cuda_device_count > 0:
            logger.warning(
                "Device has %d GPUs available. "
                "Provide device={deviceId} to `from_model_id` to use available"
                "GPUs for execution. deviceId is -1 for CPU and "
                "can be a positive integer associated with CUDA device id.",
                cuda_device_count,
            )

        client = client.to(device)
    return client


class SelfHostedHuggingFaceEmbeddings(SelfHostedEmbeddings):
    """在自托管远程硬件上使用HuggingFace嵌入模型。

支持的硬件包括在AWS、GCP、Azure和Lambda上自动启动的实例，以及通过IP地址和SSH凭据指定的服务器（例如在本地，或者在Paperspace、Coreweave等其他云上）。

要使用，您应该已安装``runhouse`` python包。

示例：
```python
from langchain_community.embeddings import SelfHostedHuggingFaceEmbeddings
import runhouse as rh
model_id = "sentence-transformers/all-mpnet-base-v2"
gpu = rh.cluster(name="rh-a10x", instance_type="A100:1")
hf = SelfHostedHuggingFaceEmbeddings(model_id=model_id, hardware=gpu)
```"""

    client: Any  #: :meta private:
    model_id: str = DEFAULT_MODEL_NAME
    """要使用的模型名称。"""
    model_reqs: List[str] = ["./", "sentence_transformers", "torch"]
    """在硬件上安装推断模型所需的要求。"""
    hardware: Any
    """远程硬件发送推理函数。"""
    model_load_fn: Callable = load_embedding_model
    """在服务器上远程加载模型的函数。"""
    load_fn_kwargs: Optional[dict] = None
    """传递给模型加载函数的关键字参数。"""
    inference_fn: Callable = _embed_documents
    """用于提取嵌入向量的推理函数。"""

    def __init__(self, **kwargs: Any):
        """初始化远程推理函数。"""
        load_fn_kwargs = kwargs.pop("load_fn_kwargs", {})
        load_fn_kwargs["model_id"] = load_fn_kwargs.get("model_id", DEFAULT_MODEL_NAME)
        load_fn_kwargs["instruct"] = load_fn_kwargs.get("instruct", False)
        load_fn_kwargs["device"] = load_fn_kwargs.get("device", 0)
        super().__init__(load_fn_kwargs=load_fn_kwargs, **kwargs)


class SelfHostedHuggingFaceInstructEmbeddings(SelfHostedHuggingFaceEmbeddings):
    """HuggingFace在自托管远程硬件上的InstructEmbedding模型。

支持的硬件包括在AWS、GCP、Azure和Lambda上自动启动的实例，以及通过IP地址和SSH凭据指定的服务器（例如在本地，或者在Paperspace、Coreweave等其他云上）。

要使用，您应该已安装``runhouse`` python包。

示例：
    .. code-block:: python

        from langchain_community.embeddings import SelfHostedHuggingFaceInstructEmbeddings
        import runhouse as rh
        model_name = "hkunlp/instructor-large"
        gpu = rh.cluster(name='rh-a10x', instance_type='A100:1')
        hf = SelfHostedHuggingFaceInstructEmbeddings(
            model_name=model_name, hardware=gpu)"""  # noqa: E501

    model_id: str = DEFAULT_INSTRUCT_MODEL
    """要使用的模型名称。"""
    embed_instruction: str = DEFAULT_EMBED_INSTRUCTION
    """用于嵌入文档的指令。"""
    query_instruction: str = DEFAULT_QUERY_INSTRUCTION
    """用于嵌入查询的指令。"""
    model_reqs: List[str] = ["./", "InstructorEmbedding", "torch"]
    """在硬件上安装推断模型所需的要求。"""

    def __init__(self, **kwargs: Any):
        """初始化远程推理函数。"""
        load_fn_kwargs = kwargs.pop("load_fn_kwargs", {})
        load_fn_kwargs["model_id"] = load_fn_kwargs.get(
            "model_id", DEFAULT_INSTRUCT_MODEL
        )
        load_fn_kwargs["instruct"] = load_fn_kwargs.get("instruct", True)
        load_fn_kwargs["device"] = load_fn_kwargs.get("device", 0)
        super().__init__(load_fn_kwargs=load_fn_kwargs, **kwargs)

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """使用HuggingFace预训练模型计算文档嵌入。

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

返回：
    嵌入列表，每个文本对应一个嵌入。
"""
        instruction_pairs = []
        for text in texts:
            instruction_pairs.append([self.embed_instruction, text])
        embeddings = self.client(self.pipeline_ref, instruction_pairs)
        return embeddings.tolist()

    def embed_query(self, text: str) -> List[float]:
        """使用HuggingFace instruct模型计算查询嵌入。

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

返回：
    文本的嵌入。
"""
        instruction_pair = [self.query_instruction, text]
        embedding = self.client(self.pipeline_ref, [instruction_pair])[0]
        return embedding.tolist()