使用NVIDIA NIMs构建高效的检索增强生成应用

原创于 2025-06-23 14:33:05 发布 · 154 阅读

CC 4.0 BY-SA版权

文章标签：

技术背景介绍

在AI应用开发中，检索增强生成（Retrieval-Augmented Generation，RAG）是一种强大技术，通过结合大语言模型（LLM）和嵌入模型，可以实现高效的信息检索与生成。NVIDIA的NIM（NVIDIA Inference Microservice）专为在NVIDIA加速基础设施上实现最佳性能而优化，支持包括嵌入、重排等多个领域的模型。

核心原理解析

NIM以容器镜像形式通过NVIDIA NGC Catalog分发，每个镜像对应一个模型。它提供了一致且易用的API，用于AI模型的推理任务。特别是，NIMs可以部署在本地或云端，企业可以完全掌控其知识产权和AI应用。

代码实现演示

下面我们将展示如何使用LangChain与NVIDIA NIMs交互，特别是使用支持检索增强生成的嵌入模型。

安装依赖

%pip install --upgrade --quiet langchain-nvidia-ai-endpoints

环境配置与初始化

在开始之前，请确保您有有效的NVIDIA API Key。

import getpass
import os

# 初始化NVIDIA API密钥
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

嵌入模型接口

通过NVIDIAEmbeddings类，您可以访问NIM支持的嵌入模型。

from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings

embedder = NVIDIAEmbeddings(model="NV-Embed-QA")  # 选择一个嵌入模型

执行查询嵌入与文档嵌入

# 生成查询的嵌入
q_embeddings = [
    embedder.embed_query("What's the weather like in Komchatka?"),
    embedder.embed_query("What kinds of food is Italy known for?"),
]

# 文档嵌入
d_embeddings = embedder.embed_documents([
    "Komchatka's weather is cold, with long, severe winters.",
    "Italy is famous for pasta, pizza, gelato, and espresso."
])

相似度计算

我们将计算查询和文档嵌入之间的相似度。

import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

# 计算相似度矩阵
cross_similarity_matrix = cosine_similarity(np.array(q_embeddings), np.array(d_embeddings))

# 绘制相似度矩阵
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()