跳转到内容
LlamaIndex Python Documentation
TypeScript

NVIDIA NIMs

llama-index-embeddings-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。

安装

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

注意：您可能需要重启内核才能使用更新后的包。

设置

Section titled “Setup”

开始使用：

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

  2. 选择 Retrieval 标签页，然后选择您所需的模型。

  3. 在 Input 下选择 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

使用NVIDIA API目录

Section titled “Working with NVIDIA API Catalog”

初始化嵌入模型时，您可以通过传递模型来选择模型，例如下方的NV-Embed-QA，或者不传递任何参数以使用默认模型。

from llama_index.embeddings.nvidia import NVIDIAEmbedding


embedder = NVIDIAEmbedding(model="NV-Embed-QA")

该模型是一个经过微调的 E5-large 模型，支持预期的 Embeddings 方法，包括：

  • get_query_embedding: 为查询样本生成查询嵌入。

  • get_text_embedding_batch: 为您希望搜索的文档列表生成文本嵌入。

  • 以及上述内容的异步版本。

使用 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.embeddings.nvidia import NVIDIAEmbedding


# connect to an embedding NIM running at localhost:8080
embedder = NVIDIAEmbedding(base_url="http://localhost:8080/v1")
embedder.available_models
/home/raspawar/Desktop/llama_index/llama-index-integrations/embeddings/llama-index-embeddings-nvidia/llama_index/embeddings/nvidia/base.py:161: UserWarning: Default model is set as: NV-Embed-QA.
Set model using model parameter.
To get available models use available_models property.
  warnings.warn(










[Model(id='NV-Embed-QA', base_model=None)]

相似度

Section titled “Similarity”

以下是对这些数据点相似性的快速测试：

查询：

  • 堪察加的天气怎么样？

  • 意大利以哪些种类的食物闻名？

  • 我的名字是什么？我打赌你不记得了……

  • 人生的意义究竟是什么呢？

  • 人生的意义在于享受乐趣 :D

文本：

  • 堪察加的天气寒冷，冬季漫长而严酷。

  • 意大利以意大利面、披萨、意式冰淇淋和浓缩咖啡而闻名。

  • 我无法回忆个人姓名，只能提供信息。

  • 人生的目的各不相同，通常被视为个人实现。

  • 享受生活中的时刻确实是一种美好的方式。

嵌入查询

Section titled “Embedding queries”
print("\nSequential Embedding: ")
q_embeddings = [
    embedder.get_query_embedding("What's the weather like in Komchatka?"),
    embedder.get_query_embedding("What kinds of food is Italy known for?"),
    embedder.get_query_embedding(
        "What's my name? I bet you don't remember..."
    ),
    embedder.get_query_embedding("What's the point of life anyways?"),
    embedder.get_query_embedding("The point of life is to have fun :D"),
]
print("Shape:", (len(q_embeddings), len(q_embeddings[0])))

文档嵌入

Section titled “Document Embedding”
print("\nBatch Document Embedding: ")
d_embeddings = embedder.get_text_embedding_batch(
    [
        "Komchatka's weather is cold, with long, severe winters.",
        "Italy is famous for pasta, pizza, gelato, and espresso.",
        "I can't recall personal names, only provide information.",
        "Life's purpose varies, often seen as personal fulfillment.",
        "Enjoying life's moments is indeed a wonderful approach.",
    ]
)
print("Shape:", (len(d_embeddings), len(d_embeddings[0])))

既然我们已经生成了嵌入向量，现在可以对结果进行简单的相似度检查，看看在检索任务中哪些文档会被触发作为合理答案：

%pip install --upgrade --quiet matplotlib scikit-learn
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity


# Compute the similarity matrix between q_embeddings and d_embeddings
cross_similarity_matrix = cosine_similarity(
    np.array(q_embeddings),
    np.array(d_embeddings),
)


# Plotting the cross-similarity matrix
plt.figure(figsize=(8, 6))
plt.imshow(cross_similarity_matrix, cmap="Greens", interpolation="nearest")
plt.colorbar()
plt.title("Cross-Similarity Matrix")
plt.xlabel("Query Embeddings")
plt.ylabel("Document Embeddings")
plt.grid(True)
plt.show()

作为提醒，发送到我们系统的查询和文档如下：

查询：

  • 堪察加的天气怎么样？

  • 意大利以哪些种类的食物闻名？

  • 我的名字是什么？我打赌你不记得了……

  • 人生的意义究竟是什么呢？

  • 人生的意义在于享受乐趣 :D

文本：

  • 堪察加的天气寒冷，冬季漫长而严酷。

  • 意大利以意大利面、披萨、意式冰淇淋和浓缩咖啡而闻名。

  • 我无法回忆个人姓名，只能提供信息。

  • 人生的目的各不相同，通常被视为个人实现。

  • 享受生活中的时刻确实是一种美好的方式。

截断

Section titled “Truncation”

嵌入模型通常具有固定的上下文窗口，用于确定可嵌入的输入令牌的最大数量。此限制可能是硬性限制，等于模型的最大输入令牌长度，也可能是有效限制，超过该限制后嵌入的准确性会降低。

由于模型基于令牌运行，而应用程序通常处理文本，应用程序确保其输入保持在模型令牌限制内可能具有挑战性。默认情况下，如果输入过大，会抛出异常。

为了协助这一点，NVIDIA NIMs 提供了一个 truncate 参数，如果输入过大，它会在服务器端截断输入。

truncate 参数有三个选项：

  • “NONE”：默认选项。如果输入过大，将抛出异常。
  • “START”：服务器从起始位置（左侧）截断输入，根据需要丢弃令牌。
  • “END”：服务器从末尾（右侧）截断输入，根据需要丢弃令牌。
long_text = "AI is amazing, amazing is " * 100
strict_embedder = NVIDIAEmbedding()
try:
    strict_embedder.get_query_embedding(long_text)
except Exception as e:
    print("Error:", e)
truncating_embedder = NVIDIAEmbedding(truncate="END")
truncating_embedder.get_query_embedding(long_text)[:5]