Source code for langchain_community.embeddings.self_hosted

from typing import Any, Callable, List

from langchain_core.embeddings import Embeddings
from langchain_core.pydantic_v1 import Extra

from langchain_community.llms.self_hosted import SelfHostedPipeline


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

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


class SelfHostedEmbeddings(SelfHostedPipeline, Embeddings):
    """自定义嵌入模型在自托管远程硬件上。

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

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

使用模型加载函数的示例：
```python
from langchain_community.embeddings import SelfHostedEmbeddings
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import runhouse as rh

gpu = rh.cluster(name="rh-a10x", instance_type="A100:1")
def get_pipeline():
    model_id = "facebook/bart-large"
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(model_id)
    return pipeline("feature-extraction", model=model, tokenizer=tokenizer)
embeddings = SelfHostedEmbeddings(
    model_load_fn=get_pipeline,
    hardware=gpu
    model_reqs=["./", "torch", "transformers"],
)
```

传递管道路径的示例：
```python
from langchain_community.embeddings import SelfHostedHFEmbeddings
import runhouse as rh
from transformers import pipeline

gpu = rh.cluster(name="rh-a10x", instance_type="A100:1")
pipeline = pipeline(model="bert-base-uncased", task="feature-extraction")
rh.blob(pickle.dumps(pipeline),
    path="models/pipeline.pkl").save().to(gpu, path="models")
embeddings = SelfHostedHFEmbeddings.from_pipeline(
    pipeline="models/pipeline.pkl",
    hardware=gpu,
    model_reqs=["./", "torch", "transformers"],
)
```"""

    inference_fn: Callable = _embed_documents
    """远程硬件上提取嵌入向量的推断函数。"""
    inference_kwargs: Any = None
    """要传递给模型推理函数的任何kwargs。"""

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

        extra = Extra.forbid

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

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

返回：
    嵌入列表，每个文本对应一个嵌入。
"""
        texts = list(map(lambda x: x.replace("\n", " "), texts))
        embeddings = self.client(self.pipeline_ref, texts)
        if not isinstance(embeddings, list):
            return embeddings.tolist()
        return embeddings

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

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

返回：
    文本的嵌入。
"""
        text = text.replace("\n", " ")
        embeddings = self.client(self.pipeline_ref, text)
        if not isinstance(embeddings, list):
            return embeddings.tolist()
        return embeddings