最近，快手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.