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

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

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

一、什么是代码表征？

随着代码库规模的持续扩大，开发者越来越依赖高效的代码检索系统来提升开发效率。传统的关键词匹配方法已经无法满足现代软件开发的需求。代码表征（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有哪些应用场景？

OASISI的广泛应用场景

在智能代码检索方面，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.