llama-index-embeddings-nvidia 包包含 LlamaIndex 集成，用于在 NVIDIA NIM 推理微服务上使用模型构建应用。NIM 支持社区以及 NVIDIA 提供的聊天、嵌入和重排序等领域的模型。这些模型经过 NVIDIA 优化，可在 NVIDIA 加速基础设施上提供最佳性能，并以 NIM 形式部署——这是一种易于使用的预构建容器，只需一条命令即可部署到任何 NVIDIA 加速基础设施上。

NVIDIA 托管的 NIM 部署可在 NVIDIA API 目录 中试用。测试完成后，可以使用 NVIDIA AI Enterprise 许可证从 NVIDIA API 目录导出 NIM，并在本地或云端运行，让企业拥有并完全控制其 IP 和 AI 应用。

NIM 按模型分别打包为容器镜像，并通过 NVIDIA NGC Catalog 以 NGC 容器镜像形式分发。NIM 的核心是提供简单、一致且熟悉的 API，用于在 AI 模型上运行推理。

安装

%pip install --upgrade --quiet llama-index-embeddings-nvidia

设置

开始使用：

1. 在 NVIDIA 创建免费账户，该平台托管 NVIDIA AI Foundation 模型。

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 目录

初始化嵌入模型时，你可以通过传入模型名称（例如 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 NIM

除了连接托管的 NVIDIA NIM，该连接器还可用于连接本地 NIM 实例。这在需要时可以帮助你将应用本地化。

有关如何设置本地 NIM 实例的说明，请参阅 NVIDIA NIM。

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)]

相似度

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

查询：

• 堪察加的天气怎么样？

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

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

• 人生的意义到底是什么？

• 人生的意义就是享受乐趣 😄

文本：

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

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

• 我记不住个人姓名，只能提供信息。

• 人生的意义因人而异，通常被视为个人价值的实现。

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

嵌入查询

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])))

嵌入文档

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()

提醒一下，发送到我们系统的查询和文档是：

查询：

• 堪察加的天气怎么样？

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

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

• 人生的意义到底是什么？

• 人生的意义就是享受乐趣 😄

文本：

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

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

• 我记不住个人姓名，只能提供信息。

• 人生的意义因人而异，通常被视为个人价值的实现。

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

截断

嵌入模型通常有一个固定的上下文窗口，它决定了可以嵌入的最大输入 token 数。这个限制可能是硬限制，等于模型的最大输入 token 长度；也可能是有效限制，超过它之后嵌入的准确率会下降。

由于模型基于 token 运行，而应用通常处理文本，应用可能很难确保其输入保持在模型的 token 限制内。默认情况下，如果输入过大，就会抛出异常。

为帮助解决这个问题，NVIDIA NIM 提供了一个 truncate 参数，可在服务器端对过大的输入进行截断。

truncate 参数有三个选项：

• “NONE”：默认选项。输入过大时会抛出异常。
• “START”：服务器从开头（左侧）截断输入，按需丢弃 token。
• “END”：服务器从末尾（右侧）截断输入，按需丢弃 token。

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]