Nomic Embed Code

🚀 Nomic Embed Code：一款先進的代碼檢索器

nomic-embed-code 是一款先進的代碼嵌入模型，在代碼檢索任務中表現出色。它能有效解決代碼檢索難題，為開發者提供高效、準確的代碼搜索體驗，助力提升開發效率。

博客 | 技術報告 | AWS SageMaker | Atlas 嵌入和非結構化數據分析平臺

🚀 快速開始

nomic-embed-code 是一款先進的代碼嵌入模型，在代碼檢索任務中表現卓越：

• 高性能：在 CodeSearchNet 上的表現優於 Voyage Code 3 和 OpenAI Embed 3 Large。
• 多語言代碼支持：針對多種編程語言（Python、Java、Ruby、PHP、JavaScript、Go）進行了訓練。
• 先進架構：擁有 70 億參數的代碼嵌入模型。
• 完全開源：模型權重、訓練數據和評估代碼均已發佈。

✨ 主要特性

模型架構

• 總參數：70 億
• 訓練方法：在 CoRNStack 數據集上進行訓練，採用雙一致性過濾和漸進式難負樣本挖掘。
• 支持語言：Python、Java、Ruby、PHP、JavaScript 和 Go

📦 安裝指南

你可以使用以下命令安裝必要的依賴項：

pip install transformers sentence-transformers torch

💻 使用示例

基礎用法

Transformers 庫使用示例

import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("nomic-ai/nomic-embed-code")
model = AutoModel.from_pretrained("nomic-ai/nomic-embed-code")

def last_token_pooling(hidden_states, attention_mask):
    sequence_lengths = attention_mask.sum(-1) - 1
    return hidden_states[torch.arange(hidden_states.shape[0]), sequence_lengths]

queries = ['Represent this query for searching relevant code: Calculate the n-th factorial']
codes = ['def fact(n):\n if n < 0:\n  raise ValueError\n return 1 if n == 0 else n * fact(n - 1)']
code_snippets = queries + codes

encoded_input = tokenizer(code_snippets, padding=True, truncation=True, return_tensors='pt')
model.eval()
with torch.no_grad():
    model_output = model(**encoded_input)[0]

embeddings = last_token_pooling(model_output, encoded_input['attention_mask'])
embeddings = F.normalize(embeddings, p=2, dim=1)
print(embeddings.shape)

similarity = F.cosine_similarity(embeddings[0], embeddings[1], dim=0)
print(similarity)

SentenceTransformers 庫使用示例

from sentence_transformers import SentenceTransformer

queries = ['Calculate the n-th factorial']
code_snippets = ['def fact(n):\n if n < 0:\n  raise ValueError\n return 1 if n == 0 else n * fact(n - 1)']

model = SentenceTransformer("nomic-ai/nomic-embed-code")
query_emb = model.encode(queries, prompt_name="query")
code_emb = model.encode(code_snippets)

similarity = model.similarity(query_emb[0], code_emb[0])
print(similarity)

高級用法

CoRNStack 數據集整理

從去重後的 Stackv2 開始，我們從函數文檔字符串和相應的代碼中創建文本 - 代碼對。我們過濾掉了低質量的對，例如文檔字符串不是英文、太短或包含 URL、HTML 標籤或無效字符的對。此外，我們保留了文本長度為 256 個標記或更長的文檔字符串，以幫助模型學習長距離依賴關係。

在初始過濾之後，我們使用雙一致性過濾來去除可能有噪聲的示例。我們對每個文檔字符串和代碼對進行嵌入，並計算每個文檔字符串與每個代碼示例之間的相似度。如果給定文檔字符串的前 2 個最相似示例中沒有找到相應的代碼示例，我們就從數據集中刪除該對。

在訓練期間，我們採用了一種基於課程的新型難負樣本挖掘策略，以確保模型從具有挑戰性的示例中學習。我們使用基於 softmax 的採樣策略，隨著時間的推移逐步採樣難度不斷增加的難負樣本。

📚 詳細文檔

加入 Nomic 社區

• Nomic 嵌入生態系統：https://www.nomic.ai/embed
• 網站：https://nomic.ai
• Twitter：https://twitter.com/nomic_ai
• Discord：https://discord.gg/myY5YDR8z8

📄 許可證

本項目採用 Apache - 2.0 許可證。

📚 引用

如果你發現該模型、數據集或訓練代碼有用，請引用我們的工作：

@misc{suresh2025cornstackhighqualitycontrastivedata,
      title={CoRNStack: High-Quality Contrastive Data for Better Code Retrieval and Reranking}, 
      author={Tarun Suresh and Revanth Gangi Reddy and Yifei Xu and Zach Nussbaum and Andriy Mulyar and Brandon Duderstadt and Heng Ji},
      year={2025},
      eprint={2412.01007},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2412.01007}, 
}