深入探索NVIDIA NIMs：优化AI推理服务的秘密武器

bhawfgrcbtwny

于 2024-10-07 01:51:58 发布

阅读量125

点赞数 3

文章标签：人工智能 python

本文链接：https://blog.csdn.net/bhawfgrcbtwny/article/details/142734389

版权

引言

在现代AI应用中，推理性能和部署灵活性是成功的关键。NVIDIA的NIM（NVIDIA Inference Microservice）通过提供优化的模型和轻松部署的容器化解决方案，成为开发者的强大工具。本篇文章将深入探讨如何使用LangChain库与NVIDIA NIM进行集成，从而构建高效的AI模型推理应用。

主要内容

什么是NVIDIA NIMs？

NVIDIA NIMs是一种预构建的容器，优化以在NVIDIA加速的基础设施上提供最佳性能。它支持来自社区和NVIDIA的多个领域模型，包括聊天、嵌入和重新排名模型。开发者可以通过NIMs在NVIDIA的API目录上测试模型，并通过NVIDIA AI企业许可导出到本地或云上运行，实现对知识产权和AI应用的完全控制。

如何使用LangChain与NIMs集成

首先，你需要在NVIDIA注册免费账户，并在API目录中选择合适的模型，获取API密钥。通过langchain-nvidia-ai-endpoints包，可以轻松与NVIDIA提供的模型交互，并用于复杂的任务如检索增强生成（RAG）。

安装和设置

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

import getpass
import os

# 示例代码，设置API密钥
if os.environ.get("NVIDIA_API_KEY", "").startswith("nvapi-"):
    print("有效的NVIDIA_API_KEY已存在。需要重置请删除当前密钥。")
else:
    nvapi_key = getpass.getpass("NVAPI Key (starts with nvapi-): ")
    assert nvapi_key.startswith("nvapi-"), f"{nvapi_key[:5]}... 不是有效的密钥"
    os.environ["NVIDIA_API_KEY"] = nvapi_key

使用NVIDIAEmbeddings类

你可以通过传递模型名称来初始化嵌入模型，或者使用默认设置。

from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings

# 使用API代理服务提高访问稳定性
embedder = NVIDIAEmbeddings(model="NV-Embed-QA")

代码示例

让我们通过一个简单的相似性检测任务展示NIM如何工作：

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

# 示例：嵌入问答并计算相似性
embedder = NVIDIAEmbeddings(model="NV-Embed-QA")

# 查询和文档嵌入
q_embeddings = [embedder.embed_query("What's the weather like in Komchatka?")]
d_embeddings = embedder.embed_documents(["Komchatka's weather is cold, with long, severe winters."])

# 计算相似性矩阵
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()