仅用 5M 数据超过 OpenAI?快手最新 Code Embedding 模型 OASIS(绿洲)发布

图片

最近,快手 Kwaipilot 团队正式宣布开源 OASIS(Optimized Augmentation Strategy for Improved code Search)Code Embedding 模型。快手 Kwaipilot 团队在代码表征领域实现了突破性进展,仅仅使用 5M Tokens 数据大幅领先 OpenAI,在多个 Code Search benchmark 中显著超越现有最佳水平。

什么是代码表征?

随着代码库规模的持续扩大,开发者越来越依赖高效的代码检索系统来提升开发效率。传统的关键词匹配方法已经无法满足现代软件开发的需求。代码表征(Code Embedding)技术通过将代码片段转化为向量表示,使机器能够深入理解代码语义,从而实现更智能的代码检索。Code Embedding 模型在代码检索、仓库级代码问答,代码补全等场景中发挥不可或缺的作用。

图片

图 1. 代码表征搜索 [1]

OASIS 用了哪些黑科技?

OASIS 的训练只使用了 5M Tokens 的数据,远低于 SOTA 模型的训练样本数,通过结合程序分析技术与创新设计的损失函数,以极低的成本就在多个测试集上超越了现有的 SOTA 模型。

具体来说,OASIS 模型采用了多项创新技术:

  1. 仓库级程序分析:传统方法仅关注单个代码片段,缺少代码的上下文信息,结合南方科技大学 Arise 实验室技术,OASIS 引入了仓库级别的程序分析技术。通过分析函数调用关系和依赖结构,模型能够更好地理解代码在更大上下文中的语义,相比于孤立地使用函数的 docstring,程序分析技术可以引入额外的上下文信息构建高质量样本对。

  2. OASIS-instruct 数据合成算法:我们开发的专有数据增强策略能够自动生成高质量的训练样本。这些样本包含代码和自然语言的对应关系,使模型能够学习到更细腻的语义差异。

  3. 融合式损失函数:创新性地将多目标优化策略应用于模型训练,确保模型在保持传统 Code Embedding 模型能力的情况下,既能准确区分相似样本,又能识别细微的语义差异,在优化目标的角度添加了新的训练维度。

OASIS 到底有多强?

图片

在权威的代码检索基准测试中,OASIS 展现出了令人瞩目的性能:

OASIS 在训练过程中并不包含测试集提供的训练集的情况下,在 CSN、CoSQA、AdvTest 等主流评测集上,性能优越。平均检索准确率优于现有所有同尺寸的 Code Embedding 模型,具体来说,以 5M 数据训练在三个不同的代码检索数据集全面超越 OpenAI-Ada-002,在模型尺寸仅为三分之一的情况下在平均检索成功率超越 CodeFuse-CGE-Small,多个数据集领先明显。

CodeSearchNet (CSN) 数据集

该数据集包含超过 200 万个代码 - 文档对,涵盖 Python、Java、JavaScript 等六种主流编程语言。数据主要来源于函数级别的代码及其对应的文档字符串(docstring),使用自然语言的 docstring 作为查询(query)。

CoSQA 数据集

CoSQA 数据集是一个包含 20,000+ 对自然语言查询和代码的标注数据集,每个数据对都至少由 3 名人类标注者进行标注。CoSQA 数据集中的查询来源于真实的网络搜索,反映了开发者在日常工作中遇到的具体需求,而代码片段则从不同的数据集中精心筛选,以确保多样性和质量。

AdvTest 数据集

AdvTest 数据集专门用于测试代码搜索任务的难度。它通过规范化函数名和变量名来增加挑战性,并从原始数据中过滤出高质量的样本。该数据集包含近 2 万的测试样本,旨在评估模型在代码理解和搜索方面的性能。

OASIS 模型仅使用了 5M 的训练数据,以 1.3B 的参数量在 CSN,CoSQA,AdvTest 数据集了超过 SOTA 模型,性能表现远超 OpenAI-Embedding-Ada-002。

OASIS 有哪些应用场景呢?

OASIS 的应用场景广泛:

在智能代码检索方面,OASIS 能准确理解开发者的查询意图,从海量代码库中精准定位最佳实践代码片段。当开发者输入具体需求描述时,模型能从海量代码库中精准定位最佳实践代码片段,并优先推荐与当前项目技术栈相匹配的示例。

图片

图 2. IDE 代码检索

代码推荐领域,OASIS 突破了传统代码补全工具的局限。基于对代码语义的深度理解,模型能预测开发者的编码意图,主动推荐 API 调用序列和完整的功能实现方案,提升 Kwaipilot 代码补全在私域代码方言中的补全质量。

智能 CR 场景,OASIS 能够精确识别功能相似但实现细节不同的代码片段。这一能力在 Code Review 功能中作用显著,可以帮助模型识别代码中的可能错误。

图片

图 3. Kwaipilot 代码审查

在代码语义理解方面,模型应用于 Kwaipilot RepoChat 功能,模型能够自动提取遗留系统或第三方库中的关键程序逻辑,生成准确的功能描述和调用关系图谱,帮助开发者快速熟悉上手仓库,显著降低团队的代码理解成本。

开源与未来展望

为推动代码智能领域的发展,我们决定将 OASIS 完整开源。开发者可以直接通过 Hugging Face 使用模型,也可以基于我们的代码进行进一步的改进和定制。

Huggingface 地址:https://huggingface.co/Kwaipilot/OASIS-code-1.3B

未来,Kwaipilot 团队将持续投入代码智能领域的研究,计划:

  • 发布性能更强的 Code Embedding 模型

  • 开源详细的技术报告和研究成果

  • 拓展模型在更多场景下的应用

OASIS 不仅是一个代码表征模型,更是 Kwaipilot 团队对代码智能未来的探索。我们期待与开发者社区一起,继续推动这项技术的发展,为软件开发效率的提升贡献力量。

让我们一起,在代码智能的绿洲中开启新的篇章。

Reference

[1] Gu X, Zhang H, Kim S. Deep code search[C]//Proceedings of the 40th International Conference on Software Engineering. 2018: 933-944.

会议推荐

12 月 13 日至 14 日(下周五至周六),AICon 全球人工智能开发与应用大会将在北京盛大开幕!本次大会汇聚 70+ 位 AI 及技术领域的顶尖专家,深入探讨大模型与推理、AI Agent、多模态、具身智能等前沿话题。此外还有丰富的圆桌论坛、以及展区活动,带你深入探索大模型的最新实践与未来趋势。年度最后一次 AI 盛宴,让我们一起见证 AI 未来。