OptimumDocumentEmbedder

使用 Hugging Face Optimum 库加载的模型计算文档嵌入的组件。


pipeline 中的最常见位置	在索引管道中的 `DocumentWriter` 之前
强制运行变量	“documents”：文档列表
输出变量	“documents”: 一份包含嵌入的文档列表
API 参考	Optimum
GitHub 链接	https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/optimum

概述

OptimumDocumentEmbedder 使用 HuggingFace Optimum 库加载的模型嵌入文本字符串。它使用 ONNX Runtime 进行高速推理。

默认模型为sentence-transformers/all-mpnet-base-v2.

与其他 Embedder 类似，此组件允许添加前缀（和后缀）以包含说明。有关更多详细信息，请参阅组件的 API 参考。

Optimum Embedder 有三个有用的特定参数，您可以通过各种模式进行控制

Pooling：从可变长度的句子嵌入生成固定大小的句子嵌入
Optimization：对模型应用图优化并提高推理速度
Quantization：降低计算和内存成本

在我们的 Optimum API 参考中查找所有可用的模式详细信息。

身份验证

只有在通过 Serverless Inference API 或 Inference Endpoints 访问私有或受限模型时，才需要使用 Hugging Face API Token 进行身份验证。

该组件使用HF_API_TOKEN 或HF_TOKEN 环境变量，或者您可以在初始化时传递 Hugging Face API Token。有关更多信息，请参阅我们的 Secret Management 页面。

用法

要开始将此集成与 Haystack 一起使用，请使用以下命令安装它：

pip install optimum-haystack

单独使用

from haystack.dataclasses import Document
from haystack_integrations.components.embedders.optimum import OptimumDocumentEmbedder

doc = Document(content="I love pizza!")

document_embedder = OptimumDocumentEmbedder(model="sentence-transformers/all-mpnet-base-v2")
document_embedder.warm_up()

result = document_embedder.run([doc])
print(result["documents"][0].embedding)

# [0.017020374536514282, -0.023255806416273117, ...]

在 pipeline 中

from haystack import Pipeline
from haystack import Document
from haystack_integrations.components.embedders.optimum import (
    OptimumDocumentEmbedder,
    OptimumEmbedderPooling,
    OptimumEmbedderOptimizationConfig,
    OptimumEmbedderOptimizationMode,
)

documents = [
    Document(content="My name is Wolfgang and I live in Berlin"),
    Document(content="I saw a black horse running"),
    Document(content="Germany has many big cities"),
]

embedder = OptimumDocumentEmbedder(
    model="intfloat/e5-base-v2",
    normalize_embeddings=True,
    onnx_execution_provider="CUDAExecutionProvider",
    optimizer_settings=OptimumEmbedderOptimizationConfig(
        mode=OptimumEmbedderOptimizationMode.O4,
        for_gpu=True,
    ),
    working_dir="/tmp/optimum",
    pooling_mode=OptimumEmbedderPooling.MEAN,
)

pipeline = Pipeline()
pipeline.add_component("embedder", embedder)

pipeline.run({"embedder": {"documents": documents}})

print(results["embedder"]["embedding"])

更新于 9 个月前