🚀 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}, 
}