跳转到内容
LlamaIndex Python Documentation
TypeScript

NVIDIA NIMs

llama-index-llms-nvidia 软件包包含与 NVIDIA NIM 推理微服务模型构建应用的 LlamaIndex 集成。NIM 支持来自社区及 NVIDIA 的跨领域模型，如聊天、嵌入和重排序模型。这些模型经过 NVIDIA 优化，可在 NVIDIA 加速基础设施上实现最佳性能，并部署为 NIM——一种易于使用的预构建容器，只需在 NVIDIA 加速基础设施上执行单一命令即可随处部署。

NVIDIA托管的NIM部署可在NVIDIA API目录上进行测试。测试完成后， 企业可通过NVIDIA AI Enterprise许可证从NVIDIA API目录导出NIM，并在本地或云端运行， 从而获得对其知识产权和AI应用程序的所有权及完全控制权。

NIMs 以容器镜像的形式按模型进行打包，并通过 NVIDIA NGC 目录作为 NGC 容器镜像分发。 NIMs 的核心功能是为 AI 模型推理运行提供简单、一致且熟悉的 API。

NVIDIA的LLM连接器

Section titled “NVIDIA’s LLM connector”

本示例介绍如何使用 LlamaIndex 通过公开可访问的 AI Foundation 端点与基于大语言模型的系统进行交互和开发。

通过此连接器，您将能够连接并生成来自兼容模型的输出，这些模型以托管的NVIDIA NIM形式提供，例如：

安装

Section titled “Installation”
%pip install --upgrade --quiet llama-index-llms-nvidia llama-index-embeddings-nvidia llama-index-readers-file

设置

Section titled “Setup”

开始使用：

  1. 使用托管 NVIDIA AI Foundation 模型的 NVIDIA 创建免费账户。

  2. 点击您选择的模型。

  3. 在输入下选择 Python 选项卡，并点击 Get API Key。然后点击 Generate Key。

  4. 复制并保存生成的密钥作为 NVIDIA_API_KEY。之后，您应该就能访问这些端点了。

import getpass
import os


# del os.environ['NVIDIA_API_KEY']  ## delete key and reset
if os.environ.get("NVIDIA_API_KEY", "").startswith("nvapi-"):
    print("Valid NVIDIA_API_KEY already in environment. Delete to reset")
else:
    nvapi_key = getpass.getpass("NVAPI Key (starts with nvapi-): ")
    assert nvapi_key.startswith(
        "nvapi-"
    ), f"{nvapi_key[:5]}... is not a valid key"
    os.environ["NVIDIA_API_KEY"] = nvapi_key
# llama-parse is async-first, running the async code in a notebook requires the use of nest_asyncio
import nest_asyncio


nest_asyncio.apply()

使用NVIDIA API目录

Section titled “Working with NVIDIA API Catalog”
from llama_index.llms.nvidia import NVIDIA
from llama_index.core.llms import ChatMessage, MessageRole


llm = NVIDIA()


messages = [
    ChatMessage(
        role=MessageRole.SYSTEM, content=("You are a helpful assistant.")
    ),
    ChatMessage(
        role=MessageRole.USER,
        content=("What are the most popular house pets in North America?"),
    ),
]


llm.chat(messages)

使用 NVIDIA NIMs

Section titled “Working with NVIDIA NIMs”

除了连接到托管的NVIDIA NIM之外，此连接器还可用于连接本地微服务实例。这有助于您在必要时将应用程序本地化。

有关如何设置本地微服务实例的说明，请参阅 https://developer.nvidia.com/blog/nvidia-nim-offers-optimized-inference-microservices-for-deploying-ai-models-at-scale/

from llama_index.llms.nvidia import NVIDIA


# connect to an chat NIM running at localhost:8080, spcecifying a specific model
llm = NVIDIA(
    base_url="http://localhost:8080/v1", model="meta/llama3-8b-instruct"
)

加载特定模型

Section titled “Loading a specific model”

现在我们可以通过传入模型名称来加载我们的 NVIDIA LLM，如文档中所述 - 位置在此处

注意：默认模型为 meta/llama3-8b-instruct。

# default model
llm = NVIDIA()
llm.model

我们可以通过 .model 属性观察当前与我们的 llm 对象关联的是哪个模型。

llm = NVIDIA(model="mistralai/mistral-7b-instruct-v0.2")
llm.model

基础功能

Section titled “Basic Functionality”

现在我们可以探索在LlamaIndex生态系统中使用连接器的不同方式！

在开始之前，让我们先设置一个 ChatMessage 对象列表 - 这是某些方法的预期输入。

我们将为每个示例遵循相同的基本模式：

  1. 我们将把我们的 NVIDIA LLM 指向我们所需的模型
  2. 我们将探讨如何使用该端点来实现所需任务！

完成：.complete().complete()

Section titled “Complete: .complete()”

我们可以使用.complete()/.acomplete()（接受字符串参数）来提示所选模型生成响应。

让我们使用默认模型来完成此任务。

completion_llm = NVIDIA()

我们可以通过检查 .model 属性来验证这是否是预期的默认值。

completion_llm.model

让我们在模型上调用 .complete() 并传入一个字符串，本例中使用 "Hello!"，然后观察响应结果。

completion_llm.complete("Hello!")

正如LlamaIndex所预期的那样 - 我们在响应中得到了一个CompletionResponse。

异步完成：.acomplete().acomplete()

Section titled “Async Complete: .acomplete()”

还有一个异步实现，可以以同样的方式利用！

await completion_llm.acomplete("Hello!")

聊天：.chat()

Section titled “Chat: .chat()”

现在我们可以尝试使用 .chat() 方法实现同样的功能。该方法需要传入一个聊天消息列表——因此我们将使用之前创建的那个列表。

我们将使用 mistralai/mixtral-8x7b-instruct-v0.1 模型作为示例。

chat_llm = NVIDIA(model="mistralai/mixtral-8x7b-instruct-v0.1")

现在我们需要做的就是在我们的 ChatMessages 列表上调用 .chat() 并观察响应结果。

您还会注意到我们可以传入一些额外的关键字参数来影响生成过程 - 在这种情况下，我们使用了 seed 参数来影响生成结果，并使用 stop 参数来指示模型在达到特定标记时停止生成！

注意：您可以通过参考所选模型的API文档来了解模型端点支持哪些额外的kwargs参数。例如，Mixtral的文档位于此处！

chat_llm.chat(messages, seed=4, stop=["cat", "cats", "Cat", "Cats"])

正如预期的那样，我们收到了一个 ChatResponse 作为响应。

异步聊天：(achat)

Section titled “Async Chat: (achat)”

我们还提供了 .chat() 方法的异步实现，可以通过以下方式调用。

await chat_llm.achat(messages)

流：.stream_chat()

Section titled “Stream: .stream_chat()”

我们也可以使用在 build.nvidia.com 上找到的模型来处理流式用例！

让我们选择另一个模型并观察这一行为。我们将使用谷歌的 gemma-7b 模型来完成此任务。

stream_llm = NVIDIA(model="google/gemma-7b")

让我们使用 .stream_chat() 调用模型，它同样期望接收一个 ChatMessage 对象列表，并捕获响应。

streamed_response = stream_llm.stream_chat(messages)
streamed_response

我们可以看到，响应是一个包含流式响应的生成器。

让我们在生成完成后查看最终响应。

last_element = None
for last_element in streamed_response:
    pass


print(last_element)

异步流：.astream_chat()

Section titled “Async Stream: .astream_chat()”

我们也有对应的异步方法用于流式处理，其使用方式与同步实现类似。

streamed_response = await stream_llm.astream_chat(messages)
streamed_response
last_element = None
async for last_element in streamed_response:
    pass


print(last_element)

流式查询引擎响应

Section titled “Streaming Query Engine Responses”

让我们来看一个使用查询引擎的稍微复杂一些的例子！

我们将从加载一些数据开始（我们将使用银河系漫游指南）。

加载数据

Section titled “Loading Data”

首先创建一个用于存放数据的目录。

!mkdir -p 'data/hhgttg'

我们将从上述来源下载我们的数据。

!wget 'https://web.eecs.utk.edu/~hqi/deeplearning/project/hhgttg.txt' -O 'data/hhgttg/hhgttg.txt'

这一步我们需要一个嵌入模型！我们将使用 NVIDIA NV-Embed-QA 模型来实现这一目标，并将其保存在我们的 Settings 中。

from llama_index.embeddings.nvidia import NVIDIAEmbedding
from llama_index.core import Settings


embedder = NVIDIAEmbedding(model="NV-Embed-QA", truncate="END")
Settings.embed_model = embedder

现在我们可以加载文档并利用上述内容创建索引

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader


documents = SimpleDirectoryReader("data/hhgttg").load_data()
index = VectorStoreIndex.from_documents(documents)

现在我们可以创建一个简单的查询引擎，并将我们的 streaming 参数设置为 True。

streaming_qe = index.as_query_engine(streaming=True)

让我们向我们的查询引擎发送一个查询，然后流式传输响应。

streaming_response = streaming_qe.query(
    "What is the significance of the number 42?",
)
streaming_response.print_response_stream()

工具调用

Section titled “Tool calling”

从 v0.2.1 版本开始，NVIDIA 支持工具调用功能。

NVIDIA 提供了与 build.nvidia.com 上的各类模型以及本地 NIM 的集成。并非所有这些模型都针对工具调用进行过训练。请务必选择确实具备工具调用功能的模型用于您的实验和应用。

你可以获取已知支持工具调用的模型列表，

NOTE: 更多示例请参考：nvidia_agent.ipynb

tool_models = [
    model
    for model in NVIDIA().available_models
    if model.is_function_calling_model
]

借助一个支持工具调用的模型，

from llama_index.core.tools import FunctionTool




def multiply(a: int, b: int) -> int:
    """Multiple two integers and returns the result integer"""
    return a * b




multiply_tool = FunctionTool.from_defaults(fn=multiply)




def add(a: int, b: int) -> int:
    """Add two integers and returns the result integer"""
    return a + b




add_tool = FunctionTool.from_defaults(fn=add)


llm = NVIDIA("meta/llama-3.1-70b-instruct")
from llama_index.core.agent import FunctionAgent


agent_worker = FunctionAgent(
    tools=[multiply_tool, add_tool],
    llm=llm,
)


response = await agent.run("What is (121 * 3) + 42?")
print(str(response))