Shuu12121 / CodeHawks-ModernBERT

CodeHawks‑ModernBERT

概要

CodeHawks‑ModernBERT は、コード検索およびコード理解のタスク向けに１からトレーニングしたCodeMorph-ModernBERT-BPE-1.3に追加学習を行った事前学習済みモデルです. 従来のCodeMorph-ModernBERTシリーズとは異なる,より構造を意識した継続学習を行い,これまでのCodeMorphシリーズとは 一線を画す性能 を実現しました.

最大シーケンス長2048トークン （従来のMicrosoftモデルは512トークン）

• アーキテクチャ : ModernBERT ベース
• 目的 : コード検索 / コード理解 /コードクローン発見
• トレーニングデータ : CodeSearchNet (全言語)　code_x_glue_tc_text_to_code
• ライセンス : Apache 2.0

主な特徴

• 長いシーケンス対応
  最大2048トークンのシーケンス処理が可能。長いコードや複雑な関数にも対応します。

• 高いコード検索,コード理解,コードクローン発見性能
  Pythonをはじめとする6言語対応のBPEレベルのトークナイザとコード全体を把握させるための事前学習によりこれまで苦手だったJava,JavaScript,RubyやGoでの性能が改善し,すべての言語でこれまでのシリーズの中で最高性能を達成しました.

• 専用にトレーニングされたモデル
  CodeSearchNet,CodeXGLUEのデータセットを用いて1から事前学習したCodeMorph-ModernBERT-BPE-1.3に対して,同じデータセットで異なる継続学習を行い,より洗練されたコード理解を実現しました.

パラメータについて

以下のパラメータで設計しています。

モデルの使用方法

Hugging Face Transformers ライブラリを利用して、本モデルを簡単にロードできます。（※ Transformers のバージョンは 4.48.0 以上のみ動作します）

• 簡単な動作例はこちらです(使用例はCodeMorph-ModernBERTになっているので名前は置き換えてください)

モデルのロード

from transformers import AutoModelForMaskedLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Shuu12121/CodeHawks-ModernBERT")
model = AutoModelForMaskedLM.from_pretrained("Shuu12121/CodeHawks-ModernBERT")

コード埋め込みの取得

import torch

def get_embedding(text, model, tokenizer, device="cuda"):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=256)
    if "token_type_ids" in inputs:
        inputs.pop("token_type_ids")
    inputs = {k: v.to(device) for k, v in inputs.items()}
    outputs = model.model(**inputs)
    embedding = outputs.last_hidden_state[:, 0, :]
    return embedding

embedding = get_embedding("def my_function(): pass", model, tokenizer)
print(embedding.shape)

評価結果

code_x_glue_ct_code_to_text データセットのTestを用いて評価を行いました.(プールサイズ100) 実験結果は こちらのリンクからご確認 できます

Python

Java

JavaScript

PHP

Ruby

Go

code_x_glue_tc_nl_code_search_adv (Test)　で評価した結果はこちらです

結果の考察と工夫点

初期のCodeMorph-ModernBERTではSentencePieceを使って作成したトークナイザを利用しモデルを作成していましたが,pythonやphpでは高い性能を出せたものの,ほかの言語ではあまり実用的なレベルではありませんでした. そこでよりプログラムの構造をつかんだトークナイザが必要だと感じ,BPEで作成したトークナイザを作成し,それを用いて事前学習を行いました(CodeMorp-ModernBERT-BPE-1.0)が、初期のモデルと比べてもいい性能が出せなかったので,より全体像を意識した学習方法を試してみた結果これまでのモデルよりも高い性能かつ.これまで苦手だったJavaやGoでの性能を上げることができました. また実験結果から考察できることは,Python,GOにおいては実験したすべてのモデルの中でCodeHawks‑ModernBERTモデルが一番優れており,ほかの言語でもcodet5p‑220m‑bimodalと競合できるレベルまで性能を上げることができています. さらに同じアーキテクチャ、同じデータセットで事前学習しているCodeMorph-ModernBERT‑BPE‑1.0 よりも圧倒的に高い性能を出していることから、 この継続学習がコード検索において有効 であったことが推測できます.

ライセンス

本モデルは Apache-2.0 ライセンスのもとで提供されます。

連絡先

このモデルで何か質問等がございましたらこちらのメールアドレスまでお願いします [email protected]