
pipeline 中的最常见位置	在 `PromptBuilder` 之后
必需的初始化变量	"api_key": NVIDIA NIM 的 API 密钥。可以通过以下方式设置环境变量 `NVIDIA_API_KEY`。
强制运行变量	“prompt”：一个包含 LLM 提示的字符串
输出变量	“replies”：一个包含 LLM 生成的所有回复的字符串列表 “meta”: 一个包含与每个回复相关的元数据的字典列表，例如 token 数量等。
API 参考	Nvidia
GitHub 链接	https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/nvidia

概述

该NvidiaGenerator 提供了一个接口，用于使用 NVIDIA NIM 自行托管的 LLM 或托管在 NVIDIA API catalog 上的模型来生成文本。

用法

要开始使用NvidiaGenerator，首先，安装nvidia-haystack 包

pip install nvidia-haystack

您可以使用NvidiaGenerator 即可使用 NVIDIA API catalog 中所有可用的 LLM，或使用 NVIDIA NIM 部署的模型。请遵循 LLMs 的 NVIDIA NIM Playbook 来了解如何在您的基础设施上部署所需模型。

单独使用

要使用 NVIDIA API catalog 中的 LLM，您需要指定正确的api_url 和您的 API 密钥。您可以直接从 catalog 网站获取您的 API 密钥。

该NvidiaGenerator 需要 Nvidia API 密钥才能工作。它使用了NVIDIA_API_KEY 环境变量。否则，您可以在初始化时通过api_key 传递 API 密钥，如下例所示。

from haystack.utils.auth import Secret
from haystack_integrations.components.generators.nvidia import NvidiaGenerator

generator = NvidiaGenerator(
    model="meta/llama3-70b-instruct",
    api_url="https://integrate.api.nvidia.com/v1",
    api_key=Secret.from_token("<your-api-key>"),
    model_arguments={
        "temperature": 0.2,
        "top_p": 0.7,
        "max_tokens": 1024,
    },
)
generator.warm_up()

result = generator.run(prompt="What is the answer?")
print(result["replies"])
print(result["meta"])

要使用本地部署的模型，您需要将api_url 设置为您的 localhost 并取消设置您的api_key.

from haystack_integrations.components.generators.nvidia import NvidiaGenerator

generator = NvidiaGenerator(
    model="llama-2-7b",
    api_url="http://0.0.0.0:9999/v1",
    api_key=None,
    model_arguments={
        "temperature": 0.2,
    },
)
generator.warm_up()

result = generator.run(prompt="What is the answer?")
print(result["replies"])
print(result["meta"])

在 Pipeline 中

以下是 RAG 管道的示例

from haystack import Pipeline, Document
from haystack.utils.auth import Secret
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
from haystack.components.builders.prompt_builder import PromptBuilder
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack_integrations.components.generators.nvidia import NvidiaGenerator

docstore = InMemoryDocumentStore()
docstore.write_documents([Document(content="Rome is the capital of Italy"), Document(content="Paris is the capital of France")])

query = "What is the capital of France?"

template = """
Given the following information, answer the question.

Context: 
{% for document in documents %}
    {{ document.content }}
{% endfor %}

Question: {{ query }}?
"""
pipe = Pipeline()

pipe.add_component("retriever", InMemoryBM25Retriever(document_store=docstore))
pipe.add_component("prompt_builder", PromptBuilder(template=template))
pipe.add_component("llm", NvidiaGenerator(
    model="meta/llama3-70b-instruct",
    api_url="https://integrate.api.nvidia.com/v1",
    api_key=Secret.from_token("<your-api-key>"),
    model_arguments={
        "temperature": 0.2,
        "top_p": 0.7,
        "max_tokens": 1024,
    },
))
pipe.connect("retriever", "prompt_builder.documents")
pipe.connect("prompt_builder", "llm")

res=pipe.run({
    "prompt_builder": {
        "query": query
    },
    "retriever": {
        "query": query
    }
})

print(res)

其他参考资料

🧑‍🍳 食谱：使用 NVIDIA NIMs 的自部署 AI 模型的 Haystack RAG 管道