Crow-mk2🐦‍⬛

Crow-mk2 は、関数コード・ソースコンテキスト・パースツリーを含む Crow-Corpus データセットで事前学習した ModernBERT ベースの Masked Language Model です。 コード検索・コード理解などの下流タスクへのファインチューニングを前提としたエンコーダモデルです。

Crow-mk2 is a ModernBERT-based Masked Language Model pre-trained on the Crow-Corpus datasets, which include function code with source context and parse trees (S-expressions). It is designed as an encoder model for fine-tuning on downstream tasks such as code search and code understanding.

Model Details

Property	Value
Architecture	ModernBERT (ModernBertForMaskedLM)
Parameters	152M
Hidden Size	768
Layers	19
Attention Heads	12
Intermediate Size (FFN)	1,536
Max Sequence Length	8,192 positions (trained with 1,024)
Vocab Size	50,368
Precision	FP32 (safetensors)

Architectural Features

• Local + Global Attention: Local sliding-window attention (window=128) with global attention every 3 layers
• RoPE: Rotary Position Embeddings (θ=160,000 global / θ=10,000 local)
• Flash Attention: Compatible with Flash Attention 2 for efficient inference
• Activation: GELU

Training

Training Data

以下の 2 つの Crow-Corpus データセットをストリーミングでインターリーブ（各 50%）して学習しています:

Dataset	Description
Shuu12121/crow-corpus-v2-func-with-context	関数コード + ソースコンテキスト + コールコンテキスト + 説明文
Shuu12121/crow-corpus-func-with-parse	関数コード + パースツリー (S式) + 説明文

• Dataset B (func-with-context): 関数の元のコードに加え、そのソースファイルのコンテキストとコールコンテキストを結合
• Dataset C (func-with-parse): 関数コードと tree-sitter による S式パースツリーを結合（S式は省略形に圧縮して約 63% トークン削減）

Training Procedure

2段階の学習パイプラインで事前学習しています:

Phase 1: Pre-training (MLM)

標準的な Masked Language Modeling による事前学習です。

Hyperparameter	Value
Collator	Standard MLM
MLM Probability	0.3 (30%)
Max Sequence Length	1,024
Max Steps	100,000
Batch Size (per device)	8
Gradient Accumulation	32
Effective Batch Size	256
Learning Rate	5.0 × 10⁻⁵
LR Scheduler	Cosine Annealing
Warmup Steps	10,000
Weight Decay	0.01
Precision	FP16 (mixed precision)

Phase 2: Continue-training (Line-level Masking)

コード構造を意識した行単位のマスキング戦略 (line_no_space) による継続学習です。 行全体をマスクする際にインデント（空白・タブ）はマスク対象から除外することで、 コードの構造情報を保持したまま学習します。

Hyperparameter	Value
Collator	line_no_space
MLM Probability	0.3 (30%)
Max Sequence Length	1,024
Max Steps	100,000
Batch Size (per device)	8
Gradient Accumulation	32
Effective Batch Size	256
Learning Rate	2.0 × 10⁻⁵
LR Scheduler	Cosine Annealing
Warmup Steps	5,000
Weight Decay	0.01
Precision	FP16 (mixed precision)

Hardware

Component	Spec
GPU	NVIDIA GeForce RTX 5090 (32 GB) × 1
CPU	24 cores
RAM	128 GB
CUDA	12.9
PyTorch	2.8.0
Transformers	4.56.2
Flash Attention	2.8.3

Tokenizer

answerdotai/ModernBERT-baseのトークナイザを使用しています。

Property	Value
Type	BPE (Byte-Level)
Vocab Size	50,368

Usage

Masked Language Modeling (Fill-Mask)

from transformers import AutoModelForMaskedLM, AutoTokenizer

model = AutoModelForMaskedLM.from_pretrained("Shuu12121/Crow-mk2")
tokenizer = AutoTokenizer.from_pretrained("Shuu12121/Crow-mk2")

text = "def [MASK](self, x, y):"
inputs = tokenizer(text, return_tensors="pt")
outputs = model(**inputs)

Feature Extraction (Embeddings)

import torch
from transformers import AutoModelForMaskedLM, AutoTokenizer

model = AutoModelForMaskedLM.from_pretrained("Shuu12121/Crow-mk2")
tokenizer = AutoTokenizer.from_pretrained("Shuu12121/Crow-mk2")

code = "def fibonacci(n):\n    if n <= 1:\n        return n\n    return fibonacci(n-1) + fibonacci(n-2)"
inputs = tokenizer(code, return_tensors="pt", padding=True, truncation=True, max_length=1024)

with torch.no_grad():
    outputs = model(**inputs, output_hidden_states=True)

# Use [CLS] token embedding from the last hidden state
cls_embedding = outputs.hidden_states[-1][:, 0, :]

Fine-tuning for Code Search

コード検索タスクへの fine-tuning には sentence-transformers を推奨します:

from sentence_transformers import SentenceTransformer

# Fine-tuning 済みモデルの場合
model = SentenceTransformer("Shuu12121/Crow-mk2")
embeddings = model.encode(["def hello(): pass", "function to greet"])

Intended Uses & Limitations

Intended Uses

• コード検索 (Code Search / Code Retrieval)
• コード理解・分類 (Code Understanding / Classification)
• コード類似度計算 (Code Similarity)
• 下流タスクへの転移学習のベースモデル

Limitations

• 事前学習モデル (MLM) のため、そのままでは生成タスクには使用できません
• 学習データは 8 言語に限定されており、それ以外の言語での性能は未検証です
• 学習時の最大系列長は 1,024 トークンです（RoPE により 8,192 まで外挿可能ですが性能低下の可能性あり）
• 自然言語のみのテキスト（コードを含まない文章）に対する性能は最適化されていません

Citation

@misc{crow-mk2,
  author       = {Shuu12121},
  title        = {Crow-mk2: A ModernBERT-based Code Pre-trained Model},
  year         = {2026},
  publisher    = {Hugging Face},
  url          = {https://huggingface.co/Shuu12121/Crow-mk2}
}

Acknowledgements

• ModernBERT
• Hugging Face Transformers