takumiecd/kairo-rnnt-trf-v2

rnnt-trf-v2 (kairo-ai)

ローマ字入力列から日本語文（漢字かな交じり文）への変換を行う RNN-Transducer モデルです。 タイプミス（重複・脱字・隣接キー誤打・転置）を含む入力に対しても頑健に変換できるよう、ノイズを付与したデータで学習しています。

A RNN-Transducer model that converts romaji input sequences into Japanese text (kanji-kana mixed sentences), trained to be robust against typos.

rnnt-trf-v1 の改良版です（同一データセット）。主な変更点は以下のとおりです。

• 入出力で埋め込み次元を分離（入力 96 / 出力 256 を共有 model 次元に射影）。v1 は入出力共通の 256 次元でした。
• 学習基盤の効率化: 自動混合精度（AMP）、lattice セル数で平準化する bucket_lattice バッチング、バッチサイズ 128（v1 は 8）。
• 40 エポックを完走（v1 は計算環境の都合で epoch 20 で中断）。検証 CER は v1 の 5.87%（サンプル 100）から約 4.2%（サンプル 500）へ改善。

モデル概要

• タスク: ローマ字 → 日本語文変換（例: wagahaihanekodearu. → 吾輩は猫である。）
• アーキテクチャ: RNN-T（エンコーダ・予測ネットワークともに Transformer）
  • エンコーダ: Transformer 4層（hidden 256）
  • 予測ネットワーク (predictor): Transformer 2層（hidden 256）
  • 埋め込み次元: 入力 96 / 出力 256（共有 model 次元に射影）
  • Joint ネットワーク隠れ次元: 256、feedforward 次元: 1024
  • アテンションヘッド数: 4、dropout: 0.1
  • 最大位置数: 256、最大系列長フィルタ: 96（--max-len 96）
  • パラメータ数: 約 7.17M（7,170,438）
  • 語彙: 入力 420 トークン（文字単位）、出力 4,102 トークン（文字単位、最低出現頻度 2）
• 損失関数: RNN-T loss
• 最適化: Adam（cosine 学習率スケジュール、ピーク 1e-3 → epoch 20 で約 1e-4、weight decay 0.01、gradient clip 1.0、AMP 有効）。epoch 35 以降は学習率 1e-5 で継続学習し epoch 40 で終了。
• バッチング: bucket_lattice（lattice セル数 ≤ 160,000）、バッチサイズ 128
• 学習環境: NVIDIA RTX 6000 Ada（48GB）、40 エポック
• 公開チェックポイント: epoch 40 時点（epoch_040.pt）。decode のデフォルトは検証損失最小の best.pt。
• ソースコード: github.com/takumiecd/kairo-ai（MIT License）

評価結果

学習中の推移（検証データ 500 サンプル抽出、greedy デコード）

Epoch	train loss	valid loss	valid CER
5	9.01	7.93	13.13%
10	5.18	5.16	8.54%
15	2.83	3.66	6.56%
20	1.81	3.11	6.76%
25	1.51	2.98	5.02%
35	0.84	2.94	4.46%
40	0.65	2.98	4.36%

検証 CER の最良値は 4.23%（epoch 39）。epoch 34 前後を境に train loss は下がり続ける一方で valid loss は微増に転じており（過学習の兆候）、decode のデフォルトは valid loss が最小の best.pt を使用します。

Note: 上表の CER は検証データ 500 サンプルの抽出値です。v1 のような未使用テストセット全件・ノイズ種別ごとの評価は本リリースには含まれていません（eval.run_test で別途算出できます。後述の「再現手順」を参照）。

学習データ

rnnt-trf-v1 と同一のデータセットを使用しています。計 301,429 ペア（input: ローマ字列, target: 日本語文）。約半数にタイプミスを模したノイズを付与（train:valid:test = 8:1:1, seed 0 → 241,143 / 30,142 / 30,144 ペア）。

ソース構成

ソース	ペア数	割合	ライセンス
日本語 Wikipedia（記事ダンプ）	199,929	66.3%	CC BY-SA / GFDL
Tatoeba（日本語例文）	100,000	33.2%	CC BY 2.0 FR
青空文庫（夏目漱石『吾輩は猫である』）	1,500	0.5%	パブリックドメイン（確認済み）

ノイズ構成

ノイズ種別	ペア数	割合
none（ノイズなし）	150,500	49.9%
duplication（文字重複）	35,546	11.8%
deletion（脱字）	35,518	11.8%
keyboard_typo（隣接キー誤打）	35,313	11.7%
swap（隣接文字転置）	33,646	11.2%
romaji_variant（ローマ字表記ゆれ）	10,906	3.6%

ファイル構成

train.rnnt.train が出力する artifact ディレクトリと同じ構成です。hf download でそのまま --artifact-dir に使えます。

ファイル	内容
checkpoints/best.pt	検証損失が最小のチェックポイント（decode のデフォルト）
checkpoints/epoch_040.pt	epoch 40 のチェックポイント（model / optimizer / scheduler の state と学習 config を含む）
config.json	モデル・学習設定（チェックポイント内 config と同一）
input_vocab.json	入力（ローマ字）語彙 420 トークン
output_vocab.json	出力（日本語文字）語彙 4,102 トークン
metrics.jsonl	エポックごとの train/valid loss・CER・学習率
loss_curve.png / cer_curve.png	学習曲線

使い方

本モデルは独自アーキテクチャのため、transformers の from_pretrained には対応していません。 kairo-ai のコードと併せて使用します。

git clone https://github.com/takumiecd/kairo-ai
cd kairo-ai && uv sync

# モデル一式（checkpoint + config + vocab）をダウンロード
hf download takumiecd/kairo-rnnt-trf-v2 --local-dir artifacts/rnnt-trf-v2

# greedy デコードで推論（デフォルトで checkpoints/best.pt が使われる）
uv run python -m decode.greedy \
  --artifact-dir artifacts/rnnt-trf-v2 \
  --input "wagahaihanekodearu."
# => 吾輩は猫である。

再現手順

本モデルの構築に使用したコマンド。詳細なフラグの説明は kairo-ai リポジトリ の README を参照。データセット構築・分割は v1 と同一のため、rnnt-trf-v1 のモデルカードを参照してください。

学習（チェックポイント内 config より）

python -m train.rnnt.train \
  --data data/combined/all_sources/train.jsonl \
  --valid-data data/combined/all_sources/valid.jsonl \
  --output-dir artifacts/rnnt-trf-v2 \
  --encoder-type transformer --prediction-type transformer \
  --encoder-layers 4 --prediction-layers 2 \
  --input-embed-dim 96 --output-embed-dim 256 \
  --encoder-hidden-dim 256 --prediction-hidden-dim 256 --joint-hidden-dim 256 \
  --feedforward-dim 1024 --num-heads 4 \
  --epochs 40 --batch-size 128 \
  --learning-rate 1e-3 \
  --amp \
  --batch-order bucket_lattice --max-batch-lattice-cells 160000 \
  --max-len 96 --max-positions 256 \
  --device cuda \
  --valid-decode greedy --valid-cer-samples 500 --valid-cer-every 1

評価（テストセット全件）

python -m eval.run_test \
  --artifact-dir artifacts/rnnt-trf-v2 \
  --data data/combined/all_sources/test.jsonl \
  --device cuda

推論例（epoch 40, greedy, CPU で動作確認済み）

入力（ローマ字）	出力
wagahaihanekodearu.	吾輩は猫である。
wagaahaihanekodearu.（重複ノイズ）	吾輩は猫である。
namaehamdaanai.（転置ノイズ）	名前はまだ無い。
kyouhayoitenkidesune.	今日はよい天気ですね。
kikaigakushuunorepo-towokakimasu.	機械学習のレポートを書きます。

ライセンスと帰属

• 学習データに CC BY-SA の Wikipedia 由来テキストを含むため、本モデルは CC BY-SA 4.0 で公開しています。
• Tatoeba の例文は tatoeba.org より CC BY 2.0 FR で提供されています。
• 青空文庫のテキストは 青空文庫 より取得したパブリックドメイン作品です。

出典

• Wikipedia 日本語版ダンプ
• Tatoeba Project
• 青空文庫