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

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 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 による事前学習です。

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

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

Hardware

Tokenizer

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

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