はじめに

　前回の記事(コード特化型のModernBERTモデル(CodeModernBERT-Owl)を開発した話)で自作した事前学習済みのコード特化型ModernBERTモデルについて解説しました.

　今回はそのCodeModernBERT-Owlと新たに作成した軽量かつより長いシーケンスに対応したCodeModernBERT-Snakeに対してコード検索に特化するようファインチューニングしたモデルである

について,これまでのモデルとの比較や有名なベンチマークであるMTEBのCodeSearchNetを用いた結果を紹介させていただきます

そもそもコードサーチタスクとは (具体例を用いた解説)

コードサーチタスク、特にCodeSearchNetのようなベンチマークで評価されるのは、「自然言語のクエリで、意図する機能を持つコードを探し出す」能力です。実際に CodeSearch-ModernBERT-Owl モデルを使って、このタスクがどのように実行されるかを見てみましょう。

以下のPythonコードは、sentence-transformers ライブラリと今回紹介する CodeSearch-ModernBERT-Owl モデルを使って、簡単なコード検索を行う例です。

from sentence_transformers import SentenceTransformer, util

# 1. モデルのロード
# 事前学習およびファインチューニング済みのCodeSearch-ModernBERT-Owlモデルを読み込む
model = SentenceTransformer("Shuu12121/CodeSearch-ModernBERT-Owl")

def search_code(query, code_snippets):
    """
    クエリとコードスニペットのリストから類似度が高い順に結果を出力する。
    
    Parameters:
        query (str): 検索クエリ (自然言語)
        code_snippets (list of str): 検索対象のコードスニペットのリスト
    
    Returns:
        list of tuple: (コードスニペット, 類似度)のリスト (類似度降順)
    """
    # 2. ベクトル化
    # モデルを使って、クエリと各コードスニペットを数値ベクトル (embedding) に変換する
    query_embedding = model.encode(query, convert_to_tensor=True)
    code_embeddings = model.encode(code_snippets, convert_to_tensor=True)

    # 3. 類似度計算
    # クエリのベクトルと各コードのベクトルの間のコサイン類似度を計算する
    # 類似度が高いほど、意味的に近いとモデルが判断していることを示す
    similarities = util.pytorch_cos_sim(query_embedding, code_embeddings)[0]

    # 4. 結果のソート
    # 計算された類似度スコアが高い順に、(コードスニペット, 類似度) のペアを並べ替える
    sorted_results = sorted(
        zip(code_snippets, similarities.tolist()), # .tolist() を追加してPythonのfloatに変換
        key=lambda x: x[1],
        reverse=True
    )

    return sorted_results

# 検索対象のコードスニペット (コードの候補群)
code_snippets = [
    "def add(a, b): return a + b", # 足し算
    "def factorial(n): return 1 if n == 0 else n * factorial(n - 1)", # 階乗
    "def is_even(n): return n % 2 == 0", # 偶数判定
    "def multiply(x, y): return x * y", # 掛け算
    "def divide(a, b): return a / b if b != 0 else None", # 割り算
]

# 検索クエリ (探したい機能の説明)
query = "add two numbers" # 「2つの数を足す」

# コード検索の実行
results = search_code(query, code_snippets)

# 検索結果の表示
print("検索クエリ:", query)
print("\n検索結果:")
for snippet, similarity in results:
    print(f"類似度: {similarity:.4f} | コード: {snippet}")

処理の流れ:

モデルのロード: まず、コード検索用にファインチューニングされた CodeSearch-ModernBERT-Owl モデルを読み込みます
ベクトル化: モデルは、入力された自然言語クエリ ("add two numbers") と、検索対象の各コードスニペット (code_snippets) を、それぞれ高次元の数値ベクトル（埋め込みベクトル）に変換します。このベクトルは、テキストやコードの意味的な特徴量を捉えた表現となっています
類似度計算: 次に、クエリのベクトルと各コードスニペットのベクトルの「近さ」を計算します。ここではコサイン類似度が用いられており、ベクトル同士の向きがどれだけ似ているかを示します。値が1に近いほど類似度が高い（意味的に近い）と判断されます
結果のソート: 計算された類似度スコアに基づき、コードスニペットをスコアの高い順（= クエリに最も関連性が高いとモデルが判断した順）に並べ替えます

CodeSearchNetのクエリについて
CodeSearchNetや私の作成したデータセットではGitHubから大量にデータを取得してくるため、すべて手作業でクエリを作成するのは非常に困難です。
そこでDocstringと呼ばれる、コードの入出力や関数の振る舞いについて自然言語で記載されたコメントをクエリとし,手作業で作成することなく大量のデータセットを作成しています.

def search_code(query, code_snippets):
    """
    クエリとコードスニペットのリストから類似度が高い順に結果を出力する。
    
    Parameters:
        query (str): 検索クエリ (自然言語)
        code_snippets (list of str): 検索対象のコードスニペットのリスト
    
    Returns:
        list of tuple: (コードスニペット, 類似度)のリスト (類似度降順)
    """
    # 2. ベクトル化
    # モデルを使って、クエリと各コードスニペットを数値ベクトル (embedding) に変換する
    query_embedding = model.encode(query, convert_to_tensor=True)
    code_embeddings = model.encode(code_snippets, convert_to_tensor=True)

    # 3. 類似度計算
    # クエリのベクトルと各コードのベクトルの間のコサイン類似度を計算する
    # 類似度が高いほど、意味的に近いとモデルが判断していることを示す
    similarities = util.pytorch_cos_sim(query_embedding, code_embeddings)[0]

    # 4. 結果のソート
    # 計算された類似度スコアが高い順に、(コードスニペット, 類似度) のペアを並べ替える
    sorted_results = sorted(
        zip(code_snippets, similarities.tolist()), # .tolist() を追加してPythonのfloatに変換
        key=lambda x: x[1],
        reverse=True
    )

    return sorted_results

の場合は

"""
    クエリとコードスニペットのリストから類似度が高い順に結果を出力する。
    
    Parameters:
        query (str): 検索クエリ (自然言語)
        code_snippets (list of str): 検索対象のコードスニペットのリスト
    
    Returns:
        list of tuple: (コードスニペット, 類似度)のリスト (類似度降順)
    """

の部分がDocstringに該当します.

実行結果の解説:

上記のコードを実行すると、以下のような結果が得られます。

検索クエリ: add two numbers

検索結果:
類似度: 0.6868 | コード: def add(a, b): return a + b
類似度: 0.4275 | コード: def multiply(x, y): return x * y
類似度: 0.3940 | コード: def divide(a, b): return a / b if b != 0 else None   
類似度: 0.2802 | コード: def factorial(n): return 1 if n == 0 else n * factorial(n - 1)
類似度: 0.2292 | コード: def is_even(n): return n % 2 == 0

この結果からわかるように:

クエリ "add two numbers" (2つの数を足す) に対して、まさにその機能を持つコード def add(a, b): return a + b が最も高い類似度スコア (0.6868) を獲得し、検索結果の1位にランク付けされています
multiply (掛け算) や divide (割り算) も数値演算に関連するため、ある程度の類似度スコアを持っていますが、add よりは低くなっています
factorial (階乗) や is_even (偶数判定) は、クエリとの関連性が低いため、類似度スコアも低くなっています

このように、コードサーチモデルは、自然言語で表現された意図とコードスニペットの意味的な関連性を捉え、関連性の高いものを上位にランク付けすることができます。CodeSearchNetのようなベンチマークは、この「正解のコードをどれだけ上位にランク付けできるか」という能力を、より大規模なデータセットと厳密な指標（MRRなど）を用いて評価するものです。

言語について

現在公開している CodeSearch-ModernBERT-Owl および CodeSearch-ModernBERT-Snake は、主にCodeSearchNetデータセットや類するデータセットに含まれるプログラミング言語と英語のドキュメンテーション文字列で学習・ファインチューニングされています。そのため、現状では日本語のような英語以外の自然言語クエリに対しては、期待通りの検索性能を発揮できないです。
そもそもトークナイザが日本語に対応していないため日本語の入力はほぼUnknownTokenに置き換わってしまいます.
将来的に日本語対応も考えていますが,CodeSearchNetのような自然言語とコードのデータセットで日本語に対応しているものはあまりなく対応が難しいのが現状です
そこで既存のデータセットをLLM等で翻訳し,そのデータを使うといった方法で対応できると考えています.
しかし個人開発であるため潤沢な計算機資源を用意できないため,現状対応できていません.

この具体例を通して、コードサーチタスクがどのようなものであり、モデルがどのように機能して検索を実現しているか、そしてベンチマークが何を評価しているのか、ご理解いただけたかと思います。

MTEB と CoIR について

モデルの性能を客観的に評価し比較するためには、標準化されたベンチマークが不可欠です。今回、モデル評価に利用した CodeSearchNet データセットは、CoIR (Code Information Retrieval) というベンチマークの一部であり、さらにその CoIR は MTEB (Massive Text Embedding Benchmark) という、より大規模なベンチマークフレームワークに統合されています。

MTEB (Massive Text Embedding Benchmark):
- テキスト埋め込みモデルの性能を、検索、分類、クラスタリング、類似度計算など、多岐にわたるタスクで網羅的に評価するための、現在最も広く利用されているベンチマークの一つです。様々なモデルの性能を統一的な基準で比較することができます
- Webサイト: https://huggingface.co/spaces/mteb/leaderboard
CoIR (Code Information Retrieval):
- MTEB の中でも、特にコード検索（自然言語によるクエリでのコード検索や、コード間の類似検索など）の能力を評価するために設計された専門的なベンチマークです
- 8つの検索タスク、7つのドメインにわたる10個の厳選されたデータセット（合計200万ドキュメント）を含んでおり、コード検索モデルの性能を多角的に評価できます。今回使用する CodeSearchNet も CoIR に含まれる代表的なデータセットの一つです
- 近年、CodeSearchNet 単体よりもダウンロード数が多くなるなど、コード検索分野における重要なベンチマークとしてコミュニティでの存在感を増しています
- MTEB に統合されたことで、MTEB の評価フレームワークを用いて CoIR のデータセットを簡単に評価できるようになっています
- GitHub: https://github.com/CoIR-team/coir
- リーダーボード: https://archersama.github.io/coir/

この記事では、作成した CodeSearch-ModernBERT-Owl/Snake のコード検索性能を、CoIR ベンチマークに含まれる CodeSearchNet データセットを用いて評価し、その結果を他の既存モデルと比較しながら紹介します。

CodeModernBERT-Snakeの紹介

次に、前回の記事以降に作成した事前学習済みモデルであるCodeModernBERT-Snake 🐍 を簡単にご紹介します。

このモデルは、CodeSearchNet データセット (全言語)のみ を用いて、まず MLM (マスク付き言語モデル) 学習を行い、その後コード理解に特化した継続学習を実施しました。

CodeModernBERT-Snake の主な特徴**

軽量化: 従来のモデル (CodeModernBERT-Owl, 152M) と比較して、パラメータ数を約半分の 76M に削減。学習や推論のコストを低減しました
長シーケンス対応: 最大入力トークン長を 8192 に拡張。より長いコードの文脈を捉えやすくなりました

特別な学習方法を施したわけではないので8192トークンのようなデータで精度を保てるかどうかは不明です.
しかしモデルとして8192トークンを受け付けるようにしているのでファインチューニングにより対応可能の可能性があります.

CodeModernBERT-Snake 継続学習後の評価結果

次に評価結果を示します(継続学習前の結果が知りたい方はモデルカードをご確認ください!!)
データセット: code_x_glue_ct_code_to_text (Test), 評価指標: MRR, R-Precision, プールサイズ: 100

言語	MRR	R-Precision
Python	0.7553	0.6730
Java	0.7660	0.6890
JavaScript	0.5934	0.4870
PHP	0.7766	0.7030
Ruby	0.6056	0.5020
Go	0.8323	0.7550

継続学習の結果、code_x_glue_ct_code_to_text データセットでのコード検索において,CodeModernBERT-Owlほどの性能は出ていませんが,そこそこの結果を出すことができていました。
またモデルは Apache 2.0 ライセンスで公開しています。

CodeSearch-ModernBERT-Owl/Snake

　ここから本題です。前回の記事でCodeXEmbed(Salesforce/SFR-Embedding-Code-400M_R)というモデルについて触れました.
このモデルはSalesForce社が開発しているコード関連の検索タスクに特化しているモデルで,コード検索においても高い性能を出しており,CodeModernBERT-Owlでも太刀打ちできないレベルでした.
またCoIRでの順位も良好であり,特によほどモデルの大きさなどで困る理由がなければ,400Mとモデルサイズは大きくないのでCodeModernBERT-Owlはわざわざ選ぶほどの優位性を提供できていない状態でした.

以下の表は、各言語（Python, Java, JavaScript, PHP, Ruby, Go）におけるMRR（Mean Reciprocal Rank）の評価結果をまとめたものです。各値はそれぞれの言語に対するコード検索評価のMRRです。
またデータセットはgoogle/code_x_glue_ct_code_to_textを用いており,プールサイズ1000で扱っています

実験結果の違いについて
これまで記載されていた実験結果はプールサイズ100行われたコードサーチでのMRRです。
　つまり100個のランダムなコードからクエリに合うコードを探せるかというものでしたが,今回の評価はその10倍の1000個のコードから適切なコードを探さなければならないので難易度がかなり上がっています.

Model	Python	Java	JavaScript	PHP	Ruby	Go
codet5p-220m-bimodal	0.63	0.72	0.54	0.74	0.53	0.62
CodeModernBERT-Snake	0.51	0.54	0.37	0.61	0.35	0.64
CodeModernBERT-Owl	0.77	0.78	0.66	0.75	0.60	0.82
SFR-Embedding-Code-400M_R	0.92	0.91	0.84	0.95	0.81	0.96

　御覧の通りSFR-Embedding-Code-400M_Rが圧倒的な精度を誇っています。
特にプールサイズを増やしてもプールサイズが100の時のCodeModernBERT-Owlと同じくらい,むしろそれ以上の結果を出しています😢

SentenceTransformersによるファインチューニング

SFR-Embedding-Code-400M_Rとの性能差を埋めるべく、モデルサイズの拡大や学習データ量の増加など、様々なアプローチを試みました。しかし、期待したほどの改善は見られませんでした。そこで、次なる一手として、埋め込みベクトル表現の学習に特化したライブラリである Sentence Transformers を用いたファインチューニングを採用することにしました。

このファインチューニングでは、損失関数として MultipleNegativesRankingLoss を使用しました。この損失関数の大きな利点は、特別なラベリング作業を必要とせず、既存の「クエリ（例: docstring）」と「対応するコード」のペアデータセットをそのまま活用できる点にあります。

MultipleNegativesRankingLoss の仕組み

MultipleNegativesRankingLossは、対照学習（Contrastive Learning）と呼ばれるアプローチに基づいています。基本的な考え方は、「似ているペア（正例）の埋め込みベクトルは近くに、似ていないペア（負例）の埋め込みベクトルは遠くに」配置するようにモデルを学習させることです。MultipleNegativesRankingLossでは、これを効率的に行うためにバッチ内ネガティブサンプリングという手法を利用します。

具体的には、以下のように学習が進みます。

バッチの準備: まず、学習データから複数の「クエリ」と「対応する正解コード」のペアをランダムに選び出し、一つのミニバッチを作成します。（例: [(クエリ1, コード1), (クエリ2, コード2), ..., (クエリN, コードN)]）
ポジティブペアとネガティブペアの定義: バッチ内の各ペア (クエリ_i, コード_i) に注目します
- (クエリ_i, コード_i) はポジティブペア（正例）となります。モデルは、このペアの埋め込みベクトル間の類似度（通常はコサイン類似度）が高くなるように学習します
- クエリ_i に対して、バッチ内の他のコード コード_j (ただし j ≠ i) は全てネガティブサンプル（負例）とみなされます。つまり、(クエリ_i, コード_j) はネガティブペアとなります。モデルは、これらのネガティブペアの埋め込みベクトル間の類似度が低くなるように学習します
損失の計算: モデルは、クエリ_i の埋め込みベクトルと、バッチ内の全てのコード（コード_1 から コード_N まで）の埋め込みベクトルとの類似度を計算します。損失関数は、ポジティブペア (クエリ_i, コード_i) の類似度スコアが、全てのネガティブペア (クエリ_i, コード_j) の類似度スコアよりも十分に高くなるよう設計されています
パラメータ更新: 計算された損失に基づき、モデルのパラメータ（重み）が更新されます。このプロセスを繰り返すことで、モデルはクエリに対して関連性の高いコードを検索結果の上位に表示する能力（＝類似度スコアを高くする能力）を獲得していきます

このように、MultipleNegativesRankingLossは、ミニバッチ内のデータ構造を巧みに利用することで、「どのコードがクエリに対して不正解か」という情報を明示的にラベリングすることなく、ランキング学習（検索に適した埋め込み空間の学習）を実現します。これにより、CodeSearchNetのような既存のペアデータセットを効率的に活用して、検索性能の高い埋め込みモデルを学習させることが可能になります。

この学習方法により,自然言語のクエリとそれに対応するコードをしっかり学習することができました.

CodeSearch-ModernBERT-Owl/Snakeの評価

Sentence Transformersによるファインチューニングが、コード検索性能にどの程度貢献したかを見ていきましょう。

以下の表は、ファインチューニング前の CodeModernBERT-Owl/Snake、比較対象の codet5p-220m-bimodal、SFR-Embedding-Code-400M_R、そしてファインチューニング後の CodeSearch-ModernBERT-Owl/Snake の性能を、google/code_x_glue_ct_code_to_text の各言語テストセットで評価した結果（MRR）をまとめたものです。
(こちらもプールサイズ1000で実行しています)
ほかのデータセット(CodeSearchNetやcode_x_glue_tc_nl_code_search_adv)こちらでご確認していただけます

Model	Python	Java	JavaScript	PHP	Ruby	Go
Salesforce/codet5p-220m-bimodal	0.6323	0.7223	0.5423	0.7445	0.5300	0.6201
CodeModernBERT-Snake (Before)	0.5183	0.5403	0.3721	0.6193	0.3576	0.6409
CodeModernBERT-Owl (Before)	0.7767	0.7817	0.6638	0.7513	0.6097	0.8231
Salesforce/SFR-Embedding-Code-400M_R	0.9246	0.9197	0.8453	0.9554	0.8147	0.9619
CodeSearch-ModernBERT-Snake	0.9278 (+0.4095)	0.9328 (+0.3925)	0.8572 (+0.4851)	0.9644 (+0.3451)	0.8121 (+0.4545)	0.9627 (+0.3218)
CodeSearch-ModernBERT-Owl	0.9442 (+0.1675)	0.9519 (+0.1702)	0.8828 (+0.2190)	0.9744 (+0.2231)	0.8461 (+0.2364)	0.9736 (+0.1505)

(括弧内はファインチューニング前 (Before) からのMRR向上値)

結果の考察:

大幅な性能向上: Sentence Transformers (MultipleNegativesRankingLoss) を用いたファインチューニングにより、CodeModernBERT-Owl と CodeModernBERT-Snake の両モデルで、全ての言語においてMRRが劇的に向上 しました。特に CodeModernBERT-Snake は、元々の性能が高くなかった分、向上幅が非常に大きくなっています
SFR-Embedding-Code-400M_R との比較: ファインチューニング後の CodeSearch-ModernBERT-Owl/Snake は、多くの言語で強力なベースラインである SFR-Embedding-Code-400M_R と同等、あるいはそれを上回る性能 を達成しました
- CodeSearch-ModernBERT-Owl (152Mパラメータ) は、より大きな SFR-Embedding-Code-400M_R (400Mパラメータ) を全ての言語で上回っています
- 軽量な CodeSearch-ModernBERT-Snake (76Mパラメータ) も、多くの言語で SFR-Embedding-Code-400M_R に匹敵、あるいは上回る性能を示しており、特に PHP と Go ではわずかに高いスコアを記録しています
軽量モデルの可能性: パラメータ数が約半分の CodeSearch-ModernBERT-Snake が、より大きな CodeSearch-ModernBERT-Owl や SFR-Embedding-Code-400M_R に迫る、あるいは一部で凌駕する性能を示したことは特筆すべき点です。これは、リソースが限られた環境でも高性能なコード検索が可能であることを示唆しています

これらの結果から、Sentence Transformersを用いたファインチューニングがコード検索タスクにおいて極めて有効であり、自作の事前学習モデルである CodeModernBERT-Owl/Snake のポテンシャルを最大限に引き出すことに成功したと言えます。

MTEB (CodeSearchNet) での評価

MTEB (Massive Text Embedding Benchmark) に含まれるCodeSearchNetタスク（具体的には COIRCodeSearchNetRetrieval）において評価されました。

主な評価結果:

Shuu12121/CodeSearch-ModernBERT-Snake (75M パラメータ):
- CodeSearchNet スコア: 72.12
- このモデルは、約7500万という非常に少ないパラメータ数にも関わらず、72.12という高いスコアを達成しました
- これは、より大きなモデルであるSalesforce/SFR-Embedding-Code-400M_R (4億パラメータ) のスコア (72.53) に匹敵する性能であり、モデルサイズに対する効率の良さを示しています
Shuu12121/CodeSearch-ModernBERT-Owl (150M パラメータ):
- CodeSearchNet スコア: 76.89
- Snakeモデルよりパラメータ数が多く（約1億5000万）、最大シーケンス長は短い（2048）このモデルは、さらに高い76.89というスコアを達成しました
- このスコアは、OpenAI-Ada-002 (74.21) や Salesforce/SFR-Embedding-Code-2B_R (20億パラメータ, 73.5) といった著名な大規模モデルをも上回る結果であり、比較的小さなモデルサイズで優れたコード検索性能を実現していることを示しています

他モデルとの比較概要:

以下の表は、両モデルと他の主要なコード埋め込みモデルとのCodeSearchNetスコア比較を示しています。
またスコアは6言語(Python,PHP,Java,JavaScript,Ruby,Go)の平均値です。

モデル名	CodeSearchNet	パラメータ規模 (概算)
CodeSage-large-v2	94.26	1.3B
CodeSage-large	90.58	1.3B
Voyage-Code-002	81.79	非公開
Shuu12121/CodeSearch-ModernBERT-Owl	76.89	150M
OpenAI-Ada-002	74.21	非公開
Salesforce/SFR-Embedding-Code-2B_R	73.5	2B
Salesforce/SFR-Embedding-Code-400M_R	72.53	400M
Shuu12121/CodeSearch-ModernBERT-Snake	72.12	75M
BGE-Base-en-v1.5	69.6	110M
E5-Base-v2	67.99	110M
UniXcoder	60.2	123M
E5-Mistral	54.25	7B
GTE-Base-en-v1.5	43.35	110M
BGE-M3	43.23	567M
Contriever	34.72	110M

結論:

CodeSearch-ModernBERT-SnakeとCodeSearch-ModernBERT-Owlは、それぞれのモデルサイズ（75M、150M）において、CodeSearchNetベンチマークで非常に競争力のある性能を示しています。特にOwlモデルは、より大きなパラメータを持ついくつかの有名モデルを凌駕する結果を出しており、効率的なコード検索モデルとしての可能性を示唆しています。
　これらのことからコードサーチタスクで,特に小さいモデルが必要な状況ではCodeSearch-ModernBERT-Owl/Snakeも候補に挙がってくることと思います.

補足説明

CodeModernBERT-Owl🦉の重要性

CodeSearch-ModernBERT-Owlは、強力な事前学習済みモデルであるCodeModernBERT-Owlを、コード検索タスクに特化してファインチューニングした結果、これほどの高い性能を達成しました。これは、基盤となるCodeModernBERT-Owlの高いポテンシャルを示しています。
さらにSentenceTransformersを用いてファインチューニングをすれば思いつくだけでも

コードクローン検出
別言語のコード推薦
コード片から類似する完成済みコードの推薦
のようなタスクに応用することが可能であると考えられます.
さらに事前学習済みモデルの性能がファインチューニング後のモデルの性能にも影響を及ぼします.今回のCodeSearch-ModernBERT-Owlが高い性能を実現できたのも事前学習済みモデルのCodeModernBERT-Owlがコードの特徴をしっかりとらえることができるからだと推測されます。

ベンチマーク測定用コード

import mteb
import logging
from sentence_transformers import SentenceTransformer
from mteb import MTEB

# ログ設定
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# モデルのロード
model_name = "Shuu12121/CodeSearch-ModernBERT-Owl"
model = SentenceTransformer(model_name)

# 関係があるタスクのみを設定
tasks = mteb.get_tasks(tasks=[
    "COIRCodeSearchNetRetrieval",
])

# MTEB オブジェクトの初期化
evaluation = MTEB(tasks=tasks)

# 評価の実行
results = evaluation.run(
    model=model,
    encode_kwargs={'batch_size': 32},  # ここを変更
    overwrite_results=True
)

参考文献