最近、KuaishouのKwaipilotチームは、OASIS（Optimized Augmentation Strategy for Improved code Search）コード埋め込みモデルのオープンソース化を正式に発表しました。Kwaipilotチームはコード表現の分野において画期的な進歩を遂げ、わずか500万トークンのデータを用いてOpenAIを大幅に上回り、複数のコード検索ベンチマークにおいて既存のベストプラクティスを大幅に上回りました。

コード表現とは何ですか?

コードベースが拡大するにつれ、開発者は開発効率を向上させるために、効率的なコード検索システムにますます依存するようになっています。従来のキーワードマッチング手法は、もはや現代のソフトウェア開発のニーズを満たすことができません。コード埋め込み技術は、コードスニペットをベクトル表現に変換することで、機械がコードのセマンティクスを深く理解し、よりインテリジェントなコード検索を実現できるようにします。コード埋め込みモデルは、コード検索、リポジトリレベルのコードに関する質問応答、コード補完といったシナリオにおいて不可欠な役割を果たします。

図1. コード表現検索 [1]

OASIS はどのような最先端技術を活用していますか?

OASISは、最先端（SOTA）モデルのトレーニングサンプル数よりもはるかに少ない、わずか500万トークンのデータを用いて学習されました。手続き型分析手法と革新的に設計された損失関数を組み合わせることで、OASISは複数のテストセットにおいて、非常に低コストで既存のSOTAモデルを上回るパフォーマンスを発揮しました。

具体的には、OASIS モデルではいくつかの革新的なテクノロジーが採用されています。
1. リポジトリレベルのプログラム分析：従来の手法は個々のコードスニペットのみに焦点を当てており、コンテキスト情報が欠けていました。OASISは、南方科技大学のArise Labの技術を活用し、リポジトリレベルのプログラム分析を導入しています。関数呼び出しの関係と依存関係構造を分析することで、モデルはより広いコンテキストにおけるコードのセマンティクスをより深く理解できます。関数のドキュメント文字列を単独で使用する場合と比較して、プログラム分析は追加のコンテキスト情報を導入し、高品質なサンプルペアを構築できます。
2. OASIS-instructデータ合成アルゴリズム：独自のデータ拡張戦略により、高品質なトレーニングサンプルが自動生成されます。これらのサンプルにはコードと自然言語の対応関係が含まれており、モデルはより微妙な意味の違いを学習できます。
3. 融合損失関数: 革新的な多目的最適化戦略がモデルトレーニングに適用され、従来のコード埋め込みモデルの機能を維持しながら、モデルが類似サンプルを正確に区別し、微妙な意味の違いを識別できるようにし、最適化目標の観点から新しいトレーニング次元を追加します。

OASIS はどれほど強力なのでしょうか?

権威あるコード検索ベンチマークテストにおいて、OASIS は優れたパフォーマンスを示しました。

OASISは、CSN、CoSQA、AdvTestといった主流の評価セットにおいて、トレーニング中にテストセットから提供されたトレーニングセットがない場合でも、優れたパフォーマンスを発揮します。平均検索精度は、同規模の既存のコード埋め込みモデルをすべて上回ります。具体的には、500万データでトレーニングした結果、 3つの異なるコード検索データセットにおいてOpenAI-Ada-002を総合的に上回りました。モデルサイズはOASISのわずか3分の1であるにもかかわらず、平均検索成功率でCodeFuse-CGE-Smallを上回り、複数のデータセットにおいて顕著なリーダーシップを発揮しています。

CodeSearchNet (CSN) データセット

このデータセットには、Python、Java、JavaScriptを含む6つの主要プログラミング言語を網羅した200万以上のコードとドキュメントのペアが含まれています。データは主に関数レベルのコードとそれに対応するドキュメント文字列から生成され、自然言語のドキュメント文字列をクエリとして使用しています。

CoSQAデータセット

CoSQAデータセットは、20,000以上の自然言語クエリとコードスニペットのペアを含むアノテーション付きデータセットです。各クエリは少なくとも3人のアノテーターによってアノテーションされています。CoSQAデータセットのクエリは実際のウェブ検索から抽出されており、開発者が日常業務で直面する具体的なニーズを反映しています。一方、コードスニペットは多様性と品質を確保するために、様々なデータセットから慎重に選定されています。

AdvTest データセット

AdvTestデータセットは、コード検索タスクの難易度をテストするために特別に設計されています。関数名と変数名を正規化し、生データから高品質なサンプルをフィルタリングすることで、難易度を高めています。このデータセットには、コード理解と検索におけるモデルのパフォーマンスを評価するために設計された約20,000個のテストサンプルが含まれています。

OASIS モデルは、わずか 500 万のトレーニング データと 13 億のパラメータを使用して、CSN、CoSQA、および AdvTest データセットの最先端 (SOTA) モデルを上回り、そのパフォーマンスは OpenAI-Embedding-Ada-002 をはるかに上回ります。

OASIS のアプリケーション シナリオにはどのようなものがありますか?

OASIS には幅広い用途があります。

コード検索に関しては、OASISは開発者のクエリ意図を正確に理解し、膨大なコードベースからベストプラクティスのコードスニペットを正確に見つけ出すことができます。開発者が特定の要件記述を入力すると、モデルは膨大なコードベースからベストプラクティスのコードスニペットを正確に見つけ出し、現在のプロジェクトのテクノロジースタックに適合するサンプルを優先的に推奨します。

図2. IDEコード検索

コード推奨の分野において、OASISは従来のコード補完ツールの限界を打ち破ります。コードセマンティクスの深い理解に基づくこのモデルは、開発者のコ​​ーディング意図を予測し、API呼び出しシーケンスと機能実装スキームの補完をプロアクティブに推奨することで、プライベートドメインのコード方言におけるKwaipilotコード補完の補完品質を向上させます。

インテリジェントなコードレビューシナリオにおいて、OASISは類似した機能を持ちながらも実装の詳細が異なるコードスニペットを正確に識別できます。この機能は特にコードレビュー機能において有用であり、モデルがコード内の潜在的なエラーを特定するのに役立ちます。

図 3. Kwaipilot のコードレビュー

コードセマンティクス理解の観点では、このモデルはKwaipilot RepoChat機能に適用されています。このモデルは、レガシーシステムやサードパーティライブラリから主要なプログラムロジックを自動的に抽出し、正確な関数記述と呼び出し関係グラフを生成することで、開発者がリポジトリに素早く慣れ、チームのコード理解コストを大幅に削減するのに役立ちます。

オープンソースと将来の展望

コードインテリジェンスの分野を発展させるため、OASISを完全にオープンソース化することにしました。開発者はHugging Faceを通じてモデルを直接利用したり、私たちのコードに基づいてさらに改良・カスタマイズしたりすることができます。

ハギングフェイス URL: https://huggingface.co/Kwaipilot/OASIS-code-1.3B

今後も、Kwaipilot チームはコード インテリジェンスの分野での研究に投資を続け、以下の計画を立てます。

• より強力なコード埋め込みモデルをリリース
• 詳細なオープンソースの技術レポートと研究結果
• より多くのシナリオへのモデルの適用を拡大

OASISは単なるコード表現モデルではなく、Kwaipilotによるコードインテリジェンスの未来への探求でもあります。開発者コミュニティと協力し、この技術の進化とソフトウェア開発効率の向上に貢献できることを楽しみにしています。

コードインテリジェンスのオアシスで、一緒に新たな章を始めましょう。

参照

[1] Gu X、Zhang H、Kim S. ディープコードサーチ[C]//第40回国際ソフトウェア工学会議論文集。2018:933-944。