NightJar
NightJar は、ソースコードと、それを説明する自然言語の両方を理解することを目指して 事前学習した ModernBERT ベースのエンコーダーモデルです。
コード検索・分類・類似コード検出・リランキングなどの下流タスク向けの基盤モデルを 想定しています。Causal Language Model ではないため、コード補完や対話形式の コード生成を主目的としたモデルではありません。
モデル仕様
| 項目 | 設定 |
|---|---|
| モデル形式 | ModernBERT Masked Language Model |
| hidden size | 768 |
| Transformer layers | 12 |
| attention heads | 12 |
| intermediate size | 1,536 |
| vocabulary size | 50,368 |
| 学習時の最大系列長 | 1,024 tokens |
| mask probability | 0.3 |
学習データ
役割の異なる 3 系統のコーパスを組み合わせています。
| データ系統 | Phase 1 | Phase 2 | 主な役割 |
|---|---|---|---|
| Owl code/docstring corpus | 使用 | 使用 | 関数コードと自然言語による説明のペア |
| GitHub file programs | 使用 | 使用 | 実際のプログラムファイルに現れる長めのコード |
| StarCoder2 Extras | 全 subset を使用 | code-only 設定の subset を使用 | 技術文書、Q&A、issue、数式、IR などの周辺### Owl code/docstring corpus |
対象は次の 16 言語です。
| 既存の 8 言語 | 追加した 8 言語 |
|---|---|
| JavaScript | C |
| Ruby | C++ |
| Go | Kotlin |
| Java | Scala |
| Rust | Swift |
| PHP | C# |
| Python | Dart |
| TypeScript | Lua |
GitHub file programs
GitHub 上のプログラムファイルを言語別に集めたデータセットです。Python、JavaScript、 TypeScript、Java、Go、Rust、Ruby、PHP の 8 言語を対象に、ファイル単位での入力を扱えるようにします。
StarCoder2 Extras
bigcode/starcoder2data-extras から 12 subset（kaggle, stackoverflow, issues,
owm, lhq, wikipedia, arxiv, documentation, ir_cpp, ir_low_resource,
ir_python, ir_rust）を使用します。Phase 1 ではすべて、Phase 2 では
documentation, kaggle, stackoverflow, issues に絞ります。
前処理
- 空文字列を除外する
- Phase 1 では、長い入力を既定で 10,000 文字までに制限する
- 最大系列長を超えた文書は overflow chunk として分割する
- 学習データの 2% を検証用に分割する
- source ごとの上限件数は YAML で個別に指定
トークナイザ
NightOwl の設計を引き継いだ ByteLevel BPE トークナイザです。記号が多く、 空白にも意味があるソースコードを壊さずに扱うことを重視しています。
- 語彙サイズは 50,368
[NL]、[Code1]、[AST]などのドメイントークンを予約- StarCoder2 Extras の文書構造トークンも予約
- 1〜32 個の連続スペースを、それぞれ 1 つの予約トークンとして登録
- 1〜4 個の連続タブ / 連続改行も、同様に 1 トークンとして登録
これらの空白系トークンは special token ではないため、
skip_special_tokens=True でデコードしても取り除かれません。
学習方法
2 段階構成です。どちらも MLM probability は 0.3 ですが、マスクの選び方と データ構成が異なります。
Phase 1: 通常の事前学習
3 系統のデータをすべて混合し、標準的な Masked Language Modeling で学習します。
Phase 2: 継続事前学習
Phase 1 の最終チェックポイントから、line_no_space collator を使って継続学習します。
この collator はコードの行構造を利用しつつ、空白だけのトークンを無視してマスクします。
これによりモデルはよりコードの文脈を学習できると考えています。
長いコードは文字数で切り捨てず、トークナイズ時に複数の chunk へ分割します。
| 項目 | Phase 1 | Phase 2 |
|---|---|---|
| collator | Standard MLM | line_no_space |
| sequence length | 1,024 | 1,024 |
| epochs | 1 | 1 |
| per-device batch size | 8 | 8 |
| gradient accumulation steps | 32 | 32 |
| learning rate | 5e-5 | 5e-5 |
| scheduler | cosine | cosine |
| warmup ratio | 0.05 | 0.05 |
| weight decay | 0.01 | 0.01 |
| precision | FP16 | FP16 |
使い方
Mask されたトークンを予測する
from transformers import pipeline
fill_mask = pipeline(
"fill-mask",
model="Shuu12121/NightJar",
tokenizer="Shuu12121/NightJar",
)
results = fill_mask("def add(a, b):\n return a [MASK] b", top_k=5)
for result in results:
print(result["token_str"], result["score"])
エンコーダーとして読み込む
import torch
from transformers import AutoModel, AutoTokenizer
model_id = "Shuu12121/NightJar"
tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
model = AutoModel.from_pretrained(model_id)
model.eval()
texts = [
"[NL]\nAdd two values\n[Code1]\ndef add(a, b):\n return a + b",
"[NL]\nJoin two strings\n[Code1]\ndef join(a, b):\n return a + b",
]
inputs = tokenizer(
texts,
padding=True,
truncation=True,
max_length=1024,
return_tensors="pt",
)
with torch.no_grad():
outputs = model(**inputs)
# 動作確認用のmean poolingです。
attention_mask = inputs["attention_mask"].unsqueeze(-1)
embeddings = (
(outputs.last_hidden_state * attention_mask).sum(dim=1)
/ attention_mask.sum(dim=1).clamp(min=1)
)
print(embeddings.shape)
運用の検索用途には、SentenceTransformer や ColBERT などの構成で追加学習したチェックポイントを使ってください。 また私がコード検索向けにファインチューニングしたモデルShuu12121/NightJar-CodeSearch-Embeddingもあるので活用してください。