AI教我做事之RAG开发-25 如何調用 NV-Embed-v2

關於如何調用 NV-Embed-v2 的詳細分析，涵蓋了所有相關步驟和考慮因素，旨在為用戶提供全面的指導。

背景與模型概述

NV-Embed-v2 是由 NVIDIA 開發的一個通用嵌入模型，於 2024 年 8 月 30 日在 Massive Text Embedding Benchmark (MTEB) 上排名第一，得分為 72.31，涵蓋 56 個文本嵌入任務。它特別在檢索子類別中表現出色，得分為 62.65，涵蓋 15 個任務，這對檢索增強生成（RAG）技術的發展至關重要。該模型基於 Mistral-7B-v0.1，採用解碼器-only 大語言模型（LLM），並引入了潛在注意力層和兩階段指令調優方法以提升性能。

模型的詳細技術信息可參考論文 NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models。它在 Hugging Face 上提供，允許通過 Python 訪問，特別是通過 HuggingFace Transformers 和 Sentence-Transformers 庫。

使用方法與代碼實現

有兩種主要方法可以調用 NV-Embed-v2：使用 HuggingFace Transformers 和 Sentence-Transformers。基於簡單性和文檔提供的示例，HuggingFace Transformers 方法被選為主要推薦。

所需包

在使用前，需確保安裝以下包：

• torch==2.2.0
• transformers==4.42.0
• flash-attn==2.2.0

安裝命令為：

pip uninstall -y transformer-enginepip install torch==2.2.0 pip install transformers==4.42.4 pip install flash-attn==2.2.0

這些版本信息來自模型的 Hugging Face 頁面 NV-Embed-v2 on Hugging Face。

HuggingFace Transformers 方法

以下是完整的 Python 代碼，展示如何調用 NV-Embed-v2：

from transformers import AutoModel
import torch
import torch.nn.functional as F
# 加載模型
model = AutoModel.from_pretrain('nvidia/NV-Embed-v2', trust_remote_code=True)
# 定義查詢前綴
query_prefix = "Instruct: Given a question, retrieve passages that answer the question\nQuery: "
# 定義原始查詢和通道
raw_queries = ["are judo throws allowed in wrestling?", "how to become a radiology technician in michigan?"]
passages = ["Some passage 1", "Some passage 2"]  # 用實際通道替換# 格式化查詢，添加前綴
queries = [query_prefix + q for q in raw_queries]
# 編碼查詢和通道
query_embs = model.encode(queries)
passage_embs = model.encode(passages)

# 標準化嵌入
query_embs = F.normalize(query_embs, dim=1)
passage_embs = F.normalize(passage_embs, dim=1)
# 計算得分（點積）
scores = torch.mm(query_embs, passage_embs.T)

# 打印得分
print(scores)

關鍵步驟解釋

1. 模型加載：使用 AutoModel.from_pretrain 加載模型，trust_remote_code=True 允許執行模型自帶的遠程代碼，這對於自定義模型（如 NV-Embed-v2）是必要的。
2. 查詢格式化：查詢需添加特定前綴 “Instruct: Given a question, retrieve passages that answer the question\nQuery: “，這是模型訓練時使用的指令模板，確保與訓練數據一致。這個細節可能對用戶來說是意想不到的，因為通常嵌入模型不要求這樣的格式。
3. 編碼：使用 model.encode() 方法編碼查詢和通道。文檔顯示該方法是模型自帶的，允許直接處理文本輸入。
4. 標準化和得分：使用 PyTorch 的 F.normalize 對嵌入進行 L2 標準化，然後通過矩陣乘法 torch.mm 計算查詢和通道之間的相似度得分。

Sentence-Transformers 方法（備選）

雖然主要推薦 HuggingFace Transformers 方法，但 Sentence-Transformers 也是一個可選方案。以下是相關步驟：

• 安裝 sentence-transformers==2.7.0 和 torch==2.2.0。
• 加載模型：model = SentenceTransformer(‘nvidia/NV-Embed-v2’, trust_remote_code=True)。
• 設置 model.max_seq_length = 32768 和 model.tokenizer.pad_side = “right”。
• 添加 EOS 標記：定義 add_eos 函數，將 EOS 標記添加到輸入文本。
• 編碼和標準化步驟類似，但需注意可能的結果不匹配問題，可能需要修改 Sentence-Transformers 源代碼（具體在 SentenceTransformer.py 第 353 行）。

由於 Sentence-Transformers 方法涉及更多手動設置和潛在的代碼修改，HuggingFace Transformers 方法被認為更適合一般用戶。

性能與應用

NV-Embed-v2 在 MTEB 基準測試中表現出色，特別是在檢索任務中，得分為 62.65，涵蓋 15 個任務。這使其成為 RAG 技術和文本檢索應用的理想選擇。模型的嵌入維度為 4096，支持長文檔（最大長度 32768），適合處理大規模文本數據。

注意事項與疑慮

• Hugging Face 訪問：如果遇到訪問問題，可能需要登錄 Hugging Face 帳戶，使用 huggingface-cli login 命令設置令牌。
• 多 GPU 支持：對於 HuggingFace Transformers 方法，若需多 GPU 支持，可使用 DataParallel，但需參考文檔的故障排除部分。
• 查詢前綴的重要性：文檔強調使用特定前綴以確保最佳性能，這可能對用戶來說是新信息，需特別注意。

表格：使用方法對比

以下表格比較了兩種方法的關鍵差異：