库名称: sentence-transformers
流水线标签: 句子相似度
标签:

• sentence-transformers
• 句子相似度
• 特征提取
  许可证: apache-2.0
  数据集:
• nomic-ai/cornstack-python-v1
• nomic-ai/cornstack-javascript-v1
• nomic-ai/cornstack-java-v1
• nomic-ai/cornstack-go-v1
• nomic-ai/cornstack-php-v1
• nomic-ai/cornstack-ruby-v1
  基础模型:
• Qwen/Qwen2.5-Coder-7B-Instruct

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 = ['为搜索相关代码表示此查询：计算第n个阶乘']  
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 = ['计算第n个阶乘']  
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

引用

如果模型、数据集或训练代码对您有帮助，请引用我们的工作：

@misc{suresh2025cornstackhighqualitycontrastivedata,  
      title={CoRNStack: 高质量对比数据提升代码检索与重排序},  
      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},  
}