NightJar-large

NightJar-large is a 346M-parameter ModernBERT encoder pre-trained from scratch on source code and the natural language around it: docstrings, documentation, Q&A and issues.

It is meant as a base model to fine-tune for code retrieval, code classification, clone detection and similar code-understanding tasks. It is not a generative model.

Highlights

• Whitespace is kept as tokens. Runs of 1–32 spaces, 1–4 tabs and 1–4 newlines are each a single token, and none of them are special tokens. Indentation survives tokenization and decoding.
• Docstring and code in the same sequence. A docstring and its code are packed into one training document, so the model learns how each description relates to its implementation:

  [NL]
  Add two numbers.
  [Code1]
  def add(a, b):
      return a + b
  

• Two-stage training. Stage 1 is standard token-level MLM on a broad mix of data. Stage 2 continues with line-level masking: whole lines are hidden and the model has to reconstruct them from the surrounding code.
• Strong after fine-tuning. With the same fine-tuning recipe, NightJar-large outperforms both the smaller NightJar and NightOwl on every language of CodeSearchNetRetrieval and CodeEditSearchRetrieval (see Evaluation).

Technical details

Architecture

NightJar-large uses the standard ModernBERT design, trained from random initialization with its own tokenizer. It does not start from any ModernBERT checkpoint.

The alternating local/global attention keeps the cost of long inputs low. Only 1 in 3 layers attends over the full sequence, while the others model local syntax such as nearby tokens on the same line or block. The architecture is ModernBERT's, so it supports unpadding and FlashAttention 2 when they are available.

Tokenizer

The tokenizer is a ByteLevel BPE with 50,368 tokens (49,766 merges). It was trained on samples from the same data sources as the model, following the NightOwl tokenizer design.

Whitespace tokens

Most BPE tokenizers split indentation into many small pieces or merge it with the next word. NightJar adds 40 dedicated whitespace tokens instead:

These are normal (non-special) tokens. They stay in the output when you decode with skip_special_tokens=True, so decoding round-trips the original code layout exactly. A line indented by 8 spaces costs one token, and blank lines are explicit.

Structure tokens

Role markers were reserved in the vocabulary from the start:

Only [NL] and [Code1] (and the StarCoder2 markers already in the raw data) appear during pre-training. The other tokens are reserved: their embeddings are untrained, and they are there so fine-tuning tasks can use them without resizing the vocabulary.

Line-level masking (Stage 2)

The Stage 2 objective is a custom collator (line_no_space):

for each sequence:
    split into segments at every token whose text contains "\n" or "\r"
    for each segment (= one line):
        with probability 0.3:
            replace every token in the line with [MASK]
            except tokens that are only spaces/tabs, and the first/last position
    loss is computed on masked positions only

This differs from token-level MLM in three ways:

• Span-like difficulty. A masked line usually hides a whole statement, including identifiers, operators and literals. Most of the masked tokens cannot be copied from neighbours in the same line.
• Structure stays visible. The indentation tokens and the newlines are never masked, so the model still sees where a line sits in the block structure: its depth, and whether it follows an if or a for.
• Masking follows lines, not tokens. The mask boundaries match how code is written, which pushes the model toward statement-level representations that are useful for retrieval and clone detection.

A newline is detected by decoding each vocabulary entry. This catches both ByteLevel Ċ pieces and the added literal \n tokens described above.

How it was built

Tokenizer (ByteLevel BPE, 50k)
        │
        ▼
Stage 1: token-level MLM from scratch   ── 159k steps ──▶  NightJar-large-pretrain
        │
        ▼
Stage 2: line-level MLM (continued)     ── 191k steps ──▶  NightJar-large  (this model)

Training data

1. Docstring–code pairs (17 datasets)

Each row's docstring and code fields are joined into one document: [NL]\n{docstring}\n[Code1]\n{code}. A row with only code starts at [Code1], and a row with only a docstring starts at [NL].

Up to 1M rows per dataset were used in Stage 1, and up to 2M in Stage 2.

2. GitHub source files (8 datasets)

The content field is used as-is, one file per document. These files add longer, whole-file context to the mostly function-level pairs above.

Up to 1M files per language were used in Stage 1, and up to 2M in Stage 2.

3. StarCoder2 Extras (12 subsets)

These subsets come from bigcode/starcoder2data-extras.

Stage 2 applies no character cap. Long documents are chunked after tokenization instead.

Preprocessing

• Empty docstrings, code or texts are dropped.
• Documents longer than 1,024 tokens are split into consecutive chunks rather than truncated, so every token is used for training.
• Datasets are streamed and tokenized one at a time, and each result is cached to disk.
• 2% of the data is held out for validation.

Stage 2 narrows the mix toward code but keeps the docstrings, docs, Q&A and issues, so the model still learns the language used to describe code.

Hyperparameters

Token counts are upper bounds (steps × 256 × 1,024) because short documents are padded.

The Stage 2 loss is higher than Stage 1 because reconstructing whole lines is a harder task. It does not mean the model got worse. During Stage 2 the eval loss fell steadily from 1.85 at step 10k to 1.41 at step 190k.

Evaluation after fine-tuning

NightJar-large is a base model, so we measure it by fine-tuning it for code retrieval and evaluating the result with MTEB.

Fine-tuning setup

Results (nDCG@10)

CodeSearchNetRetrieval

CodeEditSearchRetrieval

NightJar-large (fine-tuned) is the final checkpoint after the full epoch, and it scores highest on every language of both benchmarks.

"NightJar-CodeSearch" is the smaller Shuu12121/NightJar fine-tuned with exactly this setup for the full epoch (4,073,472 examples). "NightOwl (same recipe)" is Shuu12121/NightOwl fine-tuned with the same losses and hyperparameters (without the additional-languages dataset). Scores were computed with MTEB 2.5.1.

The fine-tuning data is the same decontaminated set used for NightOwl-CodeEmbedding: overlaps with the CodeSearchNet test splits, and between the commitpackft-derived code-edit data and the CodeEditSearchRetrieval evaluation examples, were removed before training. MTEB names the CodeEditSearchRetrieval evaluation split train, but those examples were not used for fine-tuning.

Usage

Fill-mask

from transformers import pipeline

fill = pipeline("fill-mask", model="Shuu12121/NightJar-large")
for r in fill("def add(a, b):\n    return a [MASK] b", top_k=5):
    print(repr(r["token_str"]), round(r["score"], 3))

Embeddings (as a starting point for fine-tuning)

import torch
from transformers import AutoModel, AutoTokenizer

model_id = "Shuu12121/NightJar-large"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id).eval()

texts = [
    "[NL]\nAdd two values\n[Code1]\ndef add(a, b):\n    return a + b",
    "def join(a, b):\n    return a + b",
]
batch = tok(texts, padding=True, truncation=True, max_length=1024, return_tensors="pt")

with torch.no_grad():
    hidden = model(**batch).last_hidden_state

mask = batch["attention_mask"].unsqueeze(-1)
emb = (hidden * mask).sum(1) / mask.sum(1)   # mean pooling -> (2, 1024)

Raw embeddings from a pre-trained MLM are not tuned for similarity. For search, fine-tune first, for example with Sentence Transformers (contrastive loss) or in a ColBERT setup.

Input format tips

• Code alone can be passed as a plain string.
• For a description paired with code, use the training format: f"[NL]\n{description}\n[Code1]\n{code}".
• Always pass max_length=1024 (or smaller). The tokenizer config sets no default maximum length.
• The model was trained on sequences of up to 1,024 tokens. Longer inputs are accepted but have not been tested.

Limitations

• Downstream results so far cover code retrieval only (CodeSearchNet and CodeEditSearch, after fine-tuning). Classification and clone detection have not been evaluated.
• Quality varies by language with the amount of data. Languages outside the 16 listed above are tokenizable, but they are not well covered.
• The model may reproduce bugs, outdated APIs, insecure patterns or memorized snippets from its training data.
• Fill-mask predictions are not guaranteed to be correct, safe or compilable.

License

The model weights are released under Apache 2.0. The training datasets carry their own licenses and terms of use. Check them if your use case depends on the provenance of the training data.

---

日本語版

概要

NightJar-largeは、ソースコードと、コードを説明する文章（docstring、技術文書、Q&A、issueなど）を使ってゼロから事前学習した、約3.46億パラメータのModernBERTエンコーダーです。

コード検索、コード分類、クローン検出など、コードの内容を理解する必要があるタスク向けに、追加学習（ファインチューニング）の土台として作りました。文章やコードを生成するモデルではありません。

特徴

1. 空白・インデントをトークンとして保持する

一般的なBPEトークナイザーは、インデントを細かく分割したり、直後の単語とくっつけたりします。そのため、コードのブロック構造がトークン列からは読み取りにくくなります。

NightJarでは、次の40個を専用トークンとして語彙に登録しています。

• 1〜32個の連続スペース
• 1〜4個の連続タブ
• 1〜4個の連続改行

これらは特殊トークンではなく、通常のトークンです。そのため、skip_special_tokens=Trueでデコードしても消えず、インデントや空行を含めて元のコードをそのまま復元できます。8スペースのインデントも1トークンで表せるため、系列長も節約できます。

2. 説明文とコードを1つの文書として学習する

docstringと実装コードを別々の文書にせず、次の形式で1つの系列にまとめています。

[NL]
2つの数値を足し合わせる。
[Code1]
def add(a, b):
    return a + b

説明文とコードが同じ文脈に入るため、「この説明はどんな処理を指しているのか」「このコードは何をしているのか」という対応関係を、事前学習の段階から学ばせることができます。docstringがない行は[Code1]から、コードがない行は[NL]から始まる文書として扱います。

3. 2段階の事前学習

• 1段階目: 全データを使って、通常のトークン単位MLMを行い、コード・自然言語・技術文書に共通する幅広い知識を学ばせました。
• 2段階目: 1段階目のモデルから、行単位でマスクするMLMを続け、文単位の構造やコードの意味を学ばせました。

4. ファインチューニング後の検索性能

同じ手順で追加学習した場合、NightJar-largeは、小さいNightJarおよびNightOwlを、CodeSearchNetRetrievalとCodeEditSearchRetrievalのすべての言語で上回りました（詳細は「ファインチューニング後の評価」を参照）。

技術的な詳細

アーキテクチャ

標準的なModernBERTの構成ですが、既存のModernBERTの重みは使わず、独自のトークナイザーを使ってランダムな初期値から学習しています。

• Pre-Norm型のTransformer（28層、hidden 1024、16ヘッド、ヘッドあたり64次元）
• MLPはGeGLUで、中間層の大きさは2048
• LayerNormはバイアスなし（eps 1e-5）。最終層の後にもNormを配置
• Attention・MLP・Normにはバイアスなし（MLMのデコーダーのみバイアスあり）
• 位置表現はRoPE。θはグローバル層が160,000、ローカル層が10,000
• 3層ごとにグローバルアテンション。残りは128トークン（前後64）のローカルアテンション
• Dropoutはすべて0
• 分類ヘッドの既定値は[CLS]プーリング＋SiLU

ローカルアテンションとグローバルアテンションを交互に使うため、長い入力でも計算量を抑えられます。ローカル層では同じ行やブロック内の構文を、グローバル層では離れた位置の依存関係（定義と呼び出しの関係など）を捉えることを想定しています。ModernBERTの実装なので、利用可能な環境ではunpaddingとFlashAttention 2も使えます。

トークナイザー

• ByteLevel BPE（語彙数50,368、マージ数49,766）
• NFC正規化
• 入力の先頭にスペースを自動で追加しない（add_prefix_space=False）
• 256バイトすべてを初期語彙に含めるため、どんな入力も[UNK]にならない
• 入力は[CLS] A [SEP]、ペアの場合は[CLS] A [SEP] B [SEP]の形式
• 基本の特殊トークン: [UNK]=0、[PAD]=1、[CLS]=2、[SEP]=3、[MASK]=4

用途別のトークンも、最初から語彙に登録しています。

事前学習で実際に使ったのは、[NL]、[Code1]と、元データに含まれるStarCoder2のマーカーだけです。それ以外のトークンは予約枠で、埋め込みは学習されていません。語彙サイズを変えずにファインチューニングで使えるよう、あらかじめ用意しています。

行単位マスク（2段階目）

2段階目では、次の手順でマスクを作ります。

1. 改行（\nまたは\rを含むトークン）の位置で系列を「行」に区切る
2. 各行を30%の確率で選ぶ
3. 選んだ行のトークンをまとめて[MASK]に置き換える。ただし、スペースやタブだけのトークン（インデント）は隠さずに残す
4. 損失は、マスクした位置だけで計算する

通常のトークン単位MLMと比べて、次の点が違います。

• 難しさ: 1行分の文（識別子、演算子、リテラルなど）をまとめて隠すため、同じ行の周りのトークンから答えを写すことができません。前後の行の文脈から、文全体を推測する必要があります。
• 構造は見える: インデントと改行は隠さないので、その行がブロックのどの深さにあるか、ifやforの直後かどうかといった構造は、手がかりとして残ります。
• コードの書き方に沿った単位: 行というコードの自然な単位で隠すため、検索やクローン検出で役立つ、文単位の表現を学ぶことが期待できます。

改行の判定では、語彙の各トークンをデコードして調べています。そのため、ByteLevelのĊを含むトークンと、追加した\nトークンの両方を正しく改行として扱えます。

学習データ

合計37のデータソースを使っています（各データセットへのリンクは英語版の表にあります）。

① docstringとコードのペア（17データセット）

• 16言語分: Python、JavaScript、TypeScript、Java、Go、Rust、Ruby、PHP、C、C++、C#、Kotlin、Scala、Swift、Dart、Lua（Shuu12121/*-treesitter-*）
• クラス単位の多言語データ: Fsoft-AIC/the-vault-class
• 使用件数: 1段階目は各最大100万件、2段階目は各最大200万件

② GitHubのソースファイル（8データセット）

• 8言語分: Python、JavaScript、TypeScript、Java、Go、Rust、Ruby、PHP（Shuu12121/github-file-programs-dataset-*）
• 1ファイルを1文書として使い、関数単位のペアデータだけでは得にくい、ファイル全体の長い流れを学ばせる
• 使用件数: 1段階目は各最大100万件、2段階目は各最大200万件

③ StarCoder2 Extras（12サブセット）

• 1段階目: kaggle、stackoverflow、documentation、issues、lhq、owm、wikipedia、arxiv、ir_cpp、ir_python、ir_rust、ir_low_resource
• 2段階目: documentation、kaggle、stackoverflow、issuesのみ（各最大100万件）

前処理

• 空のデータは除外
• 1,024トークンを超える文書は切り捨てず、連続したチャンクに分割してすべてのトークンを学習に使用
• 1段階目のみ、原則10,000文字を上限とする（arxivは1万文字、IR系は5千文字）
• 全体の2%を検証用として取り分け

2段階目はコード寄りの構成ですが、説明文、技術文書、Q&A、issueも引き続き含めています。コードを説明する言葉も、忘れずに学び続けられるようにするためです。

学習設定

2段階目の損失が1段階目より高いのは、行全体を復元するタスクの方が難しいためで、モデルの性能が落ちたわけではありません。2段階目の中では、検証損失が1.85（1万ステップ）から1.41（19万ステップ）まで順調に下がっています。

ファインチューニング後の評価

NightJar-largeはファインチューニングの土台となるモデルなので、コード検索向けに追加学習したうえで、MTEBで評価しました。

追加学習の設定

• Sentence Transformers、[CLS]プーリング、コサイン類似度
• 損失関数はCached MultipleNegativesRankingLoss（scale 100）。1クエリあたりハードネガティブ15件、クエリ→文書と文書→クエリの双方向で学習
• データ: owl_code_search_hard_negative_datasets_V2_kd（8言語）、owl_code_search_hard_negative_datasets_additional_languages_V2_kd（追加8言語）、codeedit_hard_negative_datasets_kd（11言語）を 8 : 8 : 1 の比率でサンプリング
• バッチサイズ1,024、学習率6e-5（コサイン減衰、ウォームアップなし）、weight decay 0.01、bf16
• タスクプレフィックスなし
• 評価したのは最終チェックポイント（全3,978ステップ中3,978ステップ目、1エポック、4,073,472件）

結果（nDCG@10、平均）

NightJar-CodeSearchは、小さいNightJarをまったく同じ設定で1エポック（4,073,472件）学習したモデルです。NightJar-largeは、CodeSearchNetRetrievalの6言語、CodeEditSearchRetrievalの13言語のすべてで、NightJar-CodeSearchと、同じ手順で追加学習したNightOwlの両方を上回りました。

言語別のスコアは英語版の表を参照してください。

追加学習のデータは、NightOwl-CodeEmbeddingと同じ重複除去済みのデータセットです。学習前に、コード検索データとCodeSearchNetのテストデータの重複、commitpackft由来のコード編集データとCodeEditSearchRetrievalの評価データの重複を取り除いています。MTEBではCodeEditSearchRetrievalの評価データがtrainという名前になっていますが、この評価データは追加学習には使っていません。

使い方

• 穴埋め・埋め込み: 英語版の「Usage」にあるコードがそのまま使えます。
• コードだけを入力する場合: 文字列をそのまま渡してください。
• 説明文とコードを組にする場合: 学習時と同じf"[NL]\n{説明文}\n[Code1]\n{コード}"の形式を推奨します。
• 最大長: トークナイザーの設定に最大長の既定値がないため、max_length=1024以下を必ず指定してください。
• 検索に使う場合: 事前学習しただけのモデルの埋め込みは、類似度の計算に最適化されていません。Sentence Transformers（対照学習）やColBERTの形式で追加学習してから使ってください。

制限事項

• 下流タスクの評価は、追加学習後のコード検索（CodeSearchNet、CodeEditSearch）のみです。分類やクローン検出は、まだ評価していません。
• 言語ごとにデータ量が違うため、性能にも差があります。上記の16言語以外もトークン化はできますが、十分には学習されていません。
• 学習データに含まれるバグ、古いAPI、安全でない実装、記憶したコード片を再現する可能性があります。
• 穴埋めの結果が正しいこと、安全であること、コンパイルできることは保証しません。

ライセンス

モデルの重みはApache 2.0で公開しています。学習データセットには、それぞれ独自のライセンスと利用条件があります。用途によっては、各データセットの条件も確認してください。