最近,Kwaipilot 团队正式宣布开源 OASIS(Optimized Augmentation Strategy for Improved code Search)Code Embedding 模型。Kwaipilot 与南方科技大学 Arise 实验室进行合作共研,在代码检索领域实现了突破性进展,仅仅使用 5M Tokens 数据大幅领先 OpenAI,在多个 Code Search benchmark 中显著超越现有最佳水平。
什么是代码表征 (Code Embedding)?
随着代码库规模的持续扩大,开发者越来越依赖高效的代码检索系统来提升开发效率。传统的关键词匹配方法已经无法满足现代软件开发的需求。代码表征(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。在使用的数据量显著少于 Codesage-Large 的情况下,在 CoSQA 数据集上超过 7.8%。在模型尺寸仅为三分之一的情况下在平均检索成功率全面超越 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 代码补全在私域代码方言中的补全质量。
针对代码复用分析,OASIS 能够精确识别功能相似但实现细节不同的代码片段。这一能力在 Code Review 功能中作用显著,可以帮助模型识别代码中的可能错误。
▲ 图 3. Kwaipilot 代码审查在代码理解辅助方面,模型应用于 Kwaipilot RepoChat 功能,模型能够自动提取遗留系统或第三方库中的关键程序逻辑,生成准确的功能描述和调用关系图谱,帮助开发者快速熟悉上手仓库,显著降低团队的代码理解成本。
同时,为推动代码智能领域的发展,OASIS 将完整开源。开发者可以直接通过 Hugging Face 使用该模型,也可以进行进一步的改进和定制。Huggingface 地址 https://huggingface.co/Kwaipilot/OASIS-code-1.3B
广告声明:文内含有的对外跳转链接(包括不限于超链接、二维码、口令等形式),用于传递更多信息,节省甄选时间,结果仅供参考,IT之家所有文章均包含本声明。