NightJar

NightJar は、ソースコードと、それを説明する自然言語の両方を理解することを目指して 事前学習した ModernBERT ベースのエンコーダーモデルです。

コード検索・分類・類似コード検出・リランキングなどの下流タスク向けの基盤モデルを 想定しています。Causal Language Model ではないため、コード補完や対話形式の コード生成を主目的としたモデルではありません。

モデル仕様

学習データ

役割の異なる 3 系統のコーパスを組み合わせています。

対象は次の 16 言語です。

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 へ分割します。

使い方

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