AI技术文章

Jina Code Embeddings: 为高质量代码搜索而生的0.5B/1.5B向量模型

作者:Jina AI0 次浏览 来源:Jina AI
Jina Code Embeddings: 为高质量代码搜索而生的0.5B/1.5B向量模型

我们今天正式开源 jina-code-embeddings,一套全新的代码向量模型。包含 0.5B 和 1.5B 两种参数规模,并同步推出了 1-4 bit 的 GGUF 量化版本,方便在各类端侧硬件上部署。

Jina Code Embeddings: 为高质量代码搜索而生的0.5B/1.5B向量模型

我们今天正式开源jina-code-embeddings,一套全新的代码向量模型。包含 0.5B 和 1.5B 两种参数规模,并同步推出了 1-4 bit 的 GGUF 量化版本,方便在各类端侧硬件上部署。

技术报告: https://arxiv.org/abs/2508.21290v1

模型介绍:https://jina.ai/models/jina-code-embeddings-1.5b/

🤗:https://huggingface.co/jinaai/jina-code-embeddings-1.5b

我们选择代码生成大语言模型(LLM)作为骨干网络进行训练,这使得模型在保持紧凑体积的同时,其代码检索性能达到了领域顶尖水平。

模型能支持 5 种核心检索任务,包括自然语言搜代码(nl2code)、代码间相似性搜索(code2code)、代码生成文档(code2nl)、代码补全(code2completions)以及技术问答(qa)。

同时,模型原生支持超过 15 种主流编程语言,覆盖了 Python, JavaScript, Java, C++, C#, Go, Rust, TypeScript, SQL, MATLAB, R, Swift, Kotlin, HTML/CSS, PHP, Ruby, Scala, Perl, 以及 Shell 脚本。

在 25 个代码检索基准测试的综合评估中,jina-code-embeddings的 0.5B 模型取得了 78.41% 的平均分,1.5B 模型则达到了 79.04%。为了更清晰地展示其性能优势,我们进行了横向对比:

  • 我们的 0.5B 模型,参数量比Qwen3-Embedding-0.6B模型少 20%,但其性能反而高出 5 个百分点。
  • 我们的 1.5B 模型,其性能不仅与voyage-code-3(79.23%) 基本持平,还明显超过了顶级的闭源模型gemini-embedding-001(77.38%)。后两者均为架构未公开的专有模型。

下表汇总了详细的性能对比数据:

Jina Code Embeddings: 为高质量代码搜索而生的0.5B/1.5B向量模型

注:标星号的模型为闭源模型,其具体架构未公开。

Jina Code Embeddings: 为高质量代码搜索而生的0.5B/1.5B向量模型

为了引导模型精准地处理不同类型的检索任务,我们在训练流程中设计了五种任务特定的指令前缀(instruction prefixes)。请注意:以下表格内容仅为便于阅读而翻译为中文。实际使用时,请将英文原文插入至代码中,并参考下方链接获取对应的英文前缀格式。

📎: https://huggingface.co/jinaai/jina-code-embeddings-1.5b/blob/main/config_sentence_transformers.json

每一种指令都区分为查询(query)和文档(document)两种角色,以支持非对称检索。例如,在处理自然语言到代码的检索任务时,可以用nl2code_query指令来编码用户查询,并用nl2code_document指令来编码待检索的代码库。

核心训练方案

我们直接选用预训练的代码生成模型作为向量模型的骨干网络(backbone)。我们的模型分别基于Qwen2.5-Coder-0.5B1.5B构建,其核心技术规格如下:

Jina Code Embeddings: 为高质量代码搜索而生的0.5B/1.5B向量模型

传统代码向量模型的训练,长期受制于一个核心瓶颈:高质量的监督训练数据,即“注释-代码”配对,极度稀缺。为此,我们选择了一条不同的技术路径。

我们选用Qwen2.5-Coder作为预训练基座,该模型已在覆盖 92 种以上编程语言、总计 5.5 万亿 token 的海量数据上完成了预训练。使我们的模型得以直接继承其对编程结构的深层语义理解、跨语言的模式识别能力,以及对语法范式的内化知识。

在此基础上,我们仅需通过对比学习(contrastive fine-tuning)进行微调,就能将这些既有知识高效地迁移到代码检索任务上。这种方法只需极少的对齐数据,从而成功规避了传统 Encoder-only 模型因数据稀缺而普遍面临的性能瓶颈。

针对部分训练数据不足的任务,例如跨框架代码翻译,我们利用大语言模型生成了合成数据,并对每一条合成样本都进行了严格的人工校验,确保其质量。最终,我们的最终训练数据集,是 MTEB 代码任务现有的训练集,与我们从 CommitPackFT、SWE-Bench、Spider、MBPP 和 CodeSearchNet 等多个公开数据集中筛选并适配数据的有机结合。

与我们早前的jina-embeddings-v3v4模型不同,这次我们放弃了 LoRA,转而采用完整的后训练(full post-training)。

我们做出此项决策,主要基于以下考量:对于 494M 和 1.54B 这样的小参数模型,LoRA 的参数效率优势并不明显。相反,适配器(adapter)的额外开销在模型容量有限时,反而可能拖累性能。我们需要调动模型的每一个参数,使其完全服务于向量化任务。

即便是处理多任务场景,我们发现采用任务特定的指令前缀,也比挂载多个 LoRA 适配器来得更简洁高效。无需切换权重配置,仅需在输入前添加不同指令。这种方式不仅更轻量,也更契合大语言模型处理条件化信息的内在机制。

整个训练过程极为高效。我们基于对比学习和 InfoNCE 损失函数,在 4 块 A100 80GB GPU 上完成了两个模型的训练。其中,0.5B 模型仅耗时 8.3 小时,1.5B 模型也只用了 12 小时。

最后,我们系统性地评测了不同的池化(pooling)策略。实验数据表明,last-token pooling(取最后一词元)策略取得了 78.41% 的综合平均分,在所有基准测试类别中都稳定优于 mean pooling(平均池化,77.20%)和 latent attention pooling(潜注意力池化,78.27%)。

这 1.2 个百分点的显著优势,促使我们放弃了在jina-embeddings-v2v3v4中沿用的 mean pooling 传统。

我们认为,随着越来越多的检索模型转向基于 Decoder-only 架构的大语言模型构建,last-token pooling 正成为更自然的选择。因为它与单向注意力机制(unidirectional attention)的原理天然契合。虽然 mean pooling 也能工作,且在训练初期往往更容易收敛(这可能得益于其更平滑的优化曲面),但我们的实验反复证明,它的性能最终会停滞在某个天花板之下,而 last-token pooling 则能达到更高的性能上限。

快速上手

我们已将两个模型无缝集成到 Search Foundation API 中,并确保它们与sentence-transformerstransformersllama.cpp等主流框架完全兼容,方便快速上手和集成。

通过 API 直接调用

您可以直接通过 Jina AI 的 Search Foundation API 端点调用模型。在请求中,只需指定模型名称(例如jina-code-embeddings-1.5b)、输入文本以及相应的任务类型(例如nl2code.passage),即可获取向量。

curl http://api.jina.ai/v1/embeddings \ -H"Content-Type: application/json"\ -H"Authorization: Bearer$JINA_API_KEY"\ -d @- <<EOF {"model":"jina-code-embeddings-1.5b","input": ["print hello world in python"],"task":"nl2code.passage"}EOF

通过sentence-transformers

模型已原生集成至sentence-transformers库,调用过程十分便捷。以下代码演示了如何加载模型,并利用prompt_name参数自动添加任务指令前缀,从而为查询(query)和文档(document)生成针对性的向量。

fromsentence_transformersimportSentenceTransformer# 加载模型 (可选择 0.5b 或 1.5b 版本)model = SentenceTransformer("jinaai/jina-code-embeddings-1.5b", model_kwargs={"torch_dtype":"bfloat16"}, tokenizer_kwargs={"padding_side":"left"})# 定义自然语言查询与待检索的代码文档queries = ["print hello world in python","initialize array of 5 zeros in c++"]documents = ["print('Hello World!')","int arr[5] = {0, 0, 0, 0, 0};"]# 使用任务指令前缀生成向量query_embeddings = model.encode(queries, prompt_name="nl2code_query")document_embeddings = model.encode(documents, prompt_name="nl2code_document")# 计算相似度similarity = model.similarity(query_embeddings, document_embeddings)

通过transformers

如果需要更底层的控制,您也可以直接使用transformers库。由于我们的模型采用last-token pooling策略,因此需要自定义一个池化函数来准确提取向量。我们提供的以下函数已妥善处理了分词器(tokenizer)在处理批量数据时可能产生的左填充(left-padding)问题。

在代码实现中,你需要手动将任务指令前缀拼接到输入文本的前面,然后依次执行分词、模型推理和池化操作。

fromtransformersimportAutoModel, AutoTokenizerimporttorchimporttorch.nn.functionalasFdeflast_token_pool(last_hidden_states, attention_mask): left_padding = (attention_mask[:,-1].sum() == attention_mask.shape[0])ifleft_padding:returnlast_hidden_states[:,-1]else: sequence_lengths = attention_mask.sum(dim=1) -1batch_size = last_hidden_states.shape[0]returnlast_hidden_states[torch.arange(batch_size), sequence_lengths]tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-code-embeddings-1.5b')model = AutoModel.from_pretrained('jinaai/jina-code-embeddings-1.5b')# 手动添加任务指令前缀query ="Find the most relevant code snippet given the following query:\nprint hello world"code ="Candidate code snippet:\nprint('Hello World!')"# 执行分词和模型推理batch_dict = tokenizer([query, code], padding=True, truncation=True, return_tensors="pt")outputs = model(**batch_dict)# 应用 last-token pooling 提取向量embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])

Matryoshka 动态截断

我们的两个模型在训练时都集成了 Matryoshka 表示学习(MRL)技术。模型生成的完整向量,其前缀本身就是低维度的、经过优化的有效表示。

您在获得完整的向量后,无需重新计算,可以直接截取其前缀,从而获得一个维度更低但依然高质量的向量。例如,从 1536 维的向量中直接截取前 256 维,就能得到一个可直接用于下游任务的 256 维向量。0.5B 模型支持的最低维度为 64,1.5B 模型则为 128。

# 生成完整的向量,维度为 896 (0.5B) 或 1536 (1.5B)full_embedding = model.encode(text)# 无需重计算,直接截断以获得更小的维度,提升效率small_embedding = full_embedding[:256]# 两个模型均支持tiny_embedding = full_embedding[:128]# 两个模型均支持

这样,您可以根据具体的应用场景,在模型性能、内存占用和检索效率之间做出最优的权衡。

结论

jina-code-embeddings的实践证明:构建顶尖的代码向量模型,并不需要依赖巨大的参数规模。我们选择以强大的代码生成模型为基座,再施以目标明确的微调,最终在 1.5B 以下的参数规模上,成功实现了业界顶尖的性能。

这次的成果,也清晰地验证了我们的核心技术理念:正确的模型基座远比参数量更为关键。代码生成模型在其海量数据的预训练阶段,已经内化了对代码的深层语义理解。我们的工作表明,这种理解能力能够别直接、高效地迁移至代码表示任务中。

这一技术路径与 Jina AI 的长远愿景完全契合。我们致力于构建统一的基座模型,让强大的向量(Embedding)能力与生成(Generation)能力,最终能够源自同一个基础模型。我们相信,这种融合将持续推动搜索基础模型的技术边界,开创新的可能。

文章来自于微信公众号 “Jina AI”,作者 “Jina AI”

相关主题

读完之后

加入 AINRK 交流群

了解加入交流群的方式,与更多 AI 从业者交流。

查看交流群说明(在新窗口打开)