OptimumTextEmbedder

使用 Hugging Face Optimum 库加载的模型嵌入文本。


pipeline 中的最常见位置	在查询/RAG 管道中的嵌入检索器之前
强制运行变量	“text”: 一个字符串
输出变量	“embedding”：浮点数列表（向量）
API 参考	Optimum
GitHub 链接	https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/optimum

概述

OptimumTextEmbedder 使用 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_integrations.components.embedders.optimum import OptimumTextEmbedder

text_to_embed = "I love pizza!"

text_embedder = OptimumTextEmbedder(model="sentence-transformers/all-mpnet-base-v2")
text_embedder.warm_up()

print(text_embedder.run(text_to_embed))

# {'embedding': [-0.07804739475250244, 0.1498992145061493,, ...]}

在 pipeline 中

请注意，此示例需要 GPU 支持才能执行。

from haystack import Pipeline

from haystack_integrations.components.embedders.optimum import (
    OptimumTextEmbedder,
    OptimumEmbedderPooling,
    OptimumEmbedderOptimizationConfig,
    OptimumEmbedderOptimizationMode,
)

pipeline = Pipeline()
embedder = OptimumTextEmbedder(
    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.add_component("embedder", embedder)

results = pipeline.run(
    {
        "embedder": {
            "text": "Ex profunditate antiquae doctrinae, Ad caelos supra semper, Hoc incantamentum evoco, draco apparet, Incantamentum iam transactum est"
        },
    }
)

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

更新于 9 个月前