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もあるので活用してください。