出品 | 51CTO技術棧（微信號：blog51cto）

最近，快手Kwaipilot 團隊正式宣布開源 OASIS（Optimized Augmentation Strategy for Improved code Search）Code Embedding模型?？焓?Kwaipilot 團隊在代碼表征領域?qū)崿F(xiàn)了突破性進展，僅僅使用 5M Tokens數(shù)據(jù)大幅領先 OpenAI，在多個 Code Search benchmark 中顯著超越現(xiàn)有最佳水平。

1.什么是代碼表征?

隨著代碼庫規(guī)模的持續(xù)擴大，開發(fā)者越來越依賴高效的代碼檢索系統(tǒng)來提升開發(fā)效率。傳統(tǒng)的關鍵詞匹配方法已經(jīng)無法滿足現(xiàn)代軟件開發(fā)的需求。代碼表征（Code Embedding）技術通過將代碼片段轉化為向量表示，使機器能夠深入理解代碼語義，從而實現(xiàn)更智能的代碼檢索。Code Embedding模型在代碼檢索、倉庫級代碼問答，代碼補全等場景中發(fā)揮不可或缺的作用。

圖 1. 代碼表征搜索[1]

2.OASIS用了哪些黑科技？

OASIS 的訓練只使用了 5M Tokens的數(shù)據(jù)，遠低于 SOTA 模型的訓練樣本數(shù)，通過結合程序分析技術與創(chuàng)新設計的損失函數(shù)，以極低的成本就在多個測試集上超越了現(xiàn)有的 SOTA 模型。

具體來說，OASIS 模型采用了多項創(chuàng)新技術：

• 倉庫級程序分析：傳統(tǒng)方法僅關注單個代碼片段，缺少代碼的上下文信息，結合南方科技大學Arise實驗室技術，OASIS 引入了倉庫級別的程序分析技術。通過分析函數(shù)調(diào)用關系和依賴結構，模型能夠更好地理解代碼在更大上下文中的語義，相比于孤立地使用函數(shù)的 docstring，程序分析技術可以引入額外的上下文信息構建高質(zhì)量樣本對。
• OASIS-instruct 數(shù)據(jù)合成算法：我們開發(fā)的專有數(shù)據(jù)增強策略能夠自動生成高質(zhì)量的訓練樣本。這些樣本包含代碼和自然語言的對應關系，使模型能夠?qū)W習到更細膩的語義差異。
• 融合式損失函數(shù)：創(chuàng)新性地將多目標優(yōu)化策略應用于模型訓練，確保模型在保持傳統(tǒng)Code Embedding模型能力的情況下，既能準確區(qū)分相似樣本，又能識別細微的語義差異，在優(yōu)化目標的角度添加了新的訓練維度。

3.OASIS到底有多強？

在權威的代碼檢索基準測試中，OASIS 展現(xiàn)出了令人矚目的性能：

OASIS在訓練過程中并不包含測試集提供的訓練集的情況下，在 CSN、CoSQA、AdvTest 等主流評測集上，性能優(yōu)越。平均檢索準確率優(yōu)于現(xiàn)有所有同尺寸的Code Embedding模型，具體來說，以5M數(shù)據(jù)訓練在三個不同的代碼檢索數(shù)據(jù)集全面超越OpenAI-Ada-002，在模型尺寸僅為三分之一的情況下在平均檢索成功率超越CodeFuse-CGE-Small，多個數(shù)據(jù)集領先明顯。

CodeSearchNet (CSN)數(shù)據(jù)集

該數(shù)據(jù)集包含超過200萬個代碼-文檔對，涵蓋Python、Java、JavaScript等六種主流編程語言。數(shù)據(jù)主要來源于函數(shù)級別的代碼及其對應的文檔字符串（docstring），使用自然語言的docstring作為查詢（query）。

CoSQA 數(shù)據(jù)集

CoSQA數(shù)據(jù)集是一個包含20,000+對自然語言查詢和代碼的標注數(shù)據(jù)集，每個數(shù)據(jù)對都至少由3名人類標注者進行標注。CoSQA數(shù)據(jù)集中的查詢來源于真實的網(wǎng)絡搜索，反映了開發(fā)者在日常工作中遇到的具體需求，而代碼片段則從不同的數(shù)據(jù)集中精心篩選，以確保多樣性和質(zhì)量。

AdvTest 數(shù)據(jù)集

AdvTest數(shù)據(jù)集專門用于測試代碼搜索任務的難度。它通過規(guī)范化函數(shù)名和變量名來增加挑戰(zhàn)性，并從原始數(shù)據(jù)中過濾出高質(zhì)量的樣本。該數(shù)據(jù)集包含近2萬的測試樣本，旨在評估模型在代碼理解和搜索方面的性能。

OASIS模型僅使用了5M的訓練數(shù)據(jù)，以1.3B的參數(shù)量在CSN，CoSQA，AdvTest數(shù)據(jù)集了超過SOTA模型，性能表現(xiàn)遠超OpenAI-Embedding-Ada-002。

4.OASIS有哪些應用場景呢？

OASIS 的應用場景廣泛：

在智能代碼檢索方面，OASIS能準確理解開發(fā)者的查詢意圖，從海量代碼庫中精準定位最佳實踐代碼片段。當開發(fā)者輸入具體需求描述時，模型能從海量代碼庫中精準定位最佳實踐代碼片段，并優(yōu)先推薦與當前項目技術棧相匹配的示例。

圖 2. IDE代碼檢索

在代碼推薦領域，OASIS突破了傳統(tǒng)代碼補全工具的局限。基于對代碼語義的深度理解，模型能預測開發(fā)者的編碼意圖，主動推薦API調(diào)用序列和完整的功能實現(xiàn)方案，提升Kwaipilot代碼補全在私域代碼方言中的補全質(zhì)量。

在智能CR場景，OASIS能夠精確識別功能相似但實現(xiàn)細節(jié)不同的代碼片段。這一能力在Code Review 功能中作用顯著，可以幫助模型識別代碼中的可能錯誤。

圖 3. Kwaipilot代碼審查

在代碼語義理解方面，模型應用于Kwaipilot RepoChat功能，模型能夠自動提取遺留系統(tǒng)或第三方庫中的關鍵程序邏輯，生成準確的功能描述和調(diào)用關系圖譜，幫助開發(fā)者快速熟悉上手倉庫，顯著降低團隊的代碼理解成本。

5.開源與未來展望

為推動代碼智能領域的發(fā)展，我們決定將 OASIS 完整開源。開發(fā)者可以直接通過 Hugging Face 使用模型，也可以基于我們的代碼進行進一步的改進和定制。

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

未來，Kwaipilot 團隊將持續(xù)投入代碼智能領域的研究，計劃：

• 發(fā)布性能更強的Code Embedding模型
• 開源詳細的技術報告和研究成果
• 拓展模型在更多場景下的應用

OASIS 不僅是一個代碼表征模型，更是 Kwaipilot 團隊對代碼智能未來的探索。我們期待與開發(fā)者社區(qū)一起，繼續(xù)推動這項技術的發(fā)展，為軟件開發(fā)效率的提升貢獻力量。

讓我們一起，在代碼智能的綠洲中開啟新的篇章。

Reference

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

本文轉載自??51CTO技術棧??