wsman

NV-Embed-v2

NV-Embed-v2 本地 API 服务 - 基于 NVIDIA 嵌入模型的 OpenAI 兼容 API

0
0
0
Python
public
View on GitHub

NV-Embed-v2 本地 API 服务

基于 NVIDIA NV-Embed-v2 的本地嵌入向量 API 服务，提供与 OpenAI 兼容的接口格式。

📖 项目简介

NV-Embed-v2 是 NVIDIA 开发的通用文本嵌入模型，在 MTEB (Massive Text Embedding Benchmark) 基准测试中排名第一（截至 2024 年 8 月），特别适合检索任务和 RAG 应用。

本项目提供了一个开箱即用的本地 API 服务，让您可以轻松地将 NV-Embed-v2 集成到您的应用中。

特性

  • 🚀 开箱即用 - 简单的配置和启动流程
  • 🔌 OpenAI 兼容 - 提供与 OpenAI 相同的 API 格式
  • 💻 本地运行 - 数据不出本地，保护隐私
  • ⚡ GPU 加速 - 支持 CUDA 加速推理
  • 🛡️ 显存保护 - 智能的并发控制和批处理

🔧 系统要求

硬件要求

配置级别 GPU 显存 内存 说明
最低配置 11 GB 16 GB RTX 2080 Ti / 1080 Ti
推荐配置 24 GB 32 GB RTX 3090 / 4090
最佳配置 48 GB+ 64 GB A100 / A6000

软件要求

  • Python 3.10+
  • CUDA 11.8+ (如需 GPU 加速)
  • Windows / Linux / macOS

📦 安装指南

1. 克隆仓库

git clone https://github.com/wsman/NV-Embed-v2.git
cd NV-Embed-v2

2. 创建虚拟环境（推荐）

python -m venv venv

# Windows
venv\Scripts\activate

# Linux/macOS
source venv/bin/activate

3. 安装依赖

pip install -r requirements.txt

4. 下载模型

模型文件较大（约 15GB），需要单独下载：

# 安装下载工具依赖
pip install huggingface_hub

# 下载模型（需要 HuggingFace 账号和访问权限）
python scripts/download_model.py

注意：NV-Embed-v2 是受限访问模型，您需要：

  1. 在 HuggingFace 注册账号
  2. 访问 模型页面 申请访问权限
  3. 获取 Access Token：Settings → Access Tokens → New token
  4. 运行 huggingface-cli login 或使用 --token 参数

手动下载（可选）

如果自动下载失败，您可以手动下载：

# 使用 huggingface-cli
huggingface-cli download nvidia/NV-Embed-v2 --local-dir assets/nvembed-v2

🚀 快速开始

启动服务

python server.py

服务将在 http://0.0.0.0:14514 启动。

API 使用示例

cURL

curl http://localhost:14514/v1/embeddings \
  -H "Content-Type: application/json" \
  -d '{
    "input": "这是一段测试文本",
    "model": "nv-embed-v2"
  }'

Python (requests)

import requests

response = requests.post(
    "http://localhost:14514/v1/embeddings",
    json={
        "input": ["这是第一段文本", "这是第二段文本"],
        "model": "nv-embed-v2"
    }
)

embeddings = response.json()["data"]
print(f"向量维度: {len(embeddings[0]['embedding'])}")

OpenAI SDK

from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:14514/v1",
    api_key="not-needed"  # 本地服务不需要 API Key
)

response = client.embeddings.create(
    input="这是一段测试文本",
    model="nv-embed-v2"
)

embedding = response.data[0].embedding
print(f"向量维度: {len(embedding)}")

⚙️ 配置选项

在 server.py 中可以修改以下配置：

配置项 默认值 说明
MODEL_PATH assets/nvembed-v2 模型路径
MAX_CONCURRENT_REQUESTS 1 最大并发请求数
MAX_LENGTH 5120 最大文本长度（tokens）
INTERNAL_BATCH_SIZE 1 内部批处理大小
INSTRUCTION_TEMPLATE 见代码 指令模板

显存优化建议

对于显存较小的 GPU（如 11GB），建议：

  1. 降低 MAX_LENGTH（如 2048 或 4096）
  2. 保持 INTERNAL_BATCH_SIZE = 1
  3. 使用 torch.float16 而非 torch.float32

📁 项目结构

NV-Embed-v2/
├── server.py              # API 服务主程序
├── requirements.txt       # Python 依赖
├── start_server.bat       # Windows 快捷启动脚本
├── LICENSE               # 许可证
├── README.md             # 项目文档
├── .gitignore            # Git 忽略规则
├── assets/
│   └── nvembed-v2/       # 模型文件（需下载）
├── scripts/
│   ├── download_model.py # 模型下载工具
│   ├── diagnostic_check.py
│   ├── error_diagnostic.py
│   ├── test_load.py
│   └── test_sdpa.py
└── src/
    └── nvembed/          # 自定义模型实现
        ├── __init__.py
        ├── configuration_nvembed.py
        └── modeling_nvembed.py

🔌 API 接口

POST /v1/embeddings

创建文本嵌入向量。

请求体：

{
  "input": "string | string[]",
  "model": "nv-embed-v2"
}

响应：

{
  "object": "list",
  "data": [
    {
      "object": "embedding",
      "embedding": [0.1, 0.2, ...],
      "index": 0
    }
  ],
  "model": "nv-embed-v2",
  "usage": {
    "prompt_tokens": 10,
    "total_tokens": 10
  }
}

GET /v1/models

列出可用模型。

❓ 常见问题

1. 模型加载失败

错误：Access to model nvidia/NV-Embed-v2 is restricted

解决：确保您已登录 HuggingFace 并获得了模型访问权限。

huggingface-cli login

2. CUDA 内存不足 (OOM)

错误：torch.cuda.OutOfMemoryError

解决：

  • 降低 MAX_LENGTH 值
  • 确保 INTERNAL_BATCH_SIZE = 1
  • 检查是否有其他程序占用显存

3. 模型路径错误

错误：nvidia/NV-Embed-v2 is not the path to a directory

解决：检查 assets/nvembed-v2/config.json 中的 _name_or_path 是否正确指向本地路径。

📄 许可证

本项目采用 CC-BY-NC-4.0 许可证。

重要：本模型不可用于商业用途。如需商业使用，请访问 NVIDIA NeMo Retriever。

🙏 致谢

📚 参考文献

@article{lee2024nv,
  title={NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models},
  author={Lee, Chankyu and Roy, Rajarshi and Xu, Mengyao and Raiman, Jonathan and Shoeybi, Mohammad and Catanzaro, Bryan and Ping, Wei},
  journal={arXiv preprint arXiv:2405.17428},
  year={2024}
}

Star ⭐ 本项目以获取最新更新！

v0.3.3[beta]