Skip to content
BranchesTags

Repository files navigation

original_LLM

PyTorchのみで、外部の事前学習モデルやFTに頼らず、最小の独自LLMをゼロから作るための実験リポジトリです。

これは知的な汎用AIや高性能アシスタントを作る試みではありません。太宰っぽいキャラと口調を持ち、友達感覚で気軽にしゃべれる小さなLLMを目指します。

現時点では性能を追いません。まずは「自分で実装した学習器が動き、太宰っぽい返答を少しでも返せる」状態を最優先にします。

Goal

  • PyTorchだけで学習と推論を実装する
  • 外部の事前学習モデルは使わない
  • Fine-tuningではなく、ゼロから学習する
  • ローカル環境で完結する
  • まずは小さくても動くものを作る
  • 太宰っぽいキャラで、友達感覚の会話が少しずつ成立するところまで持っていく

Non-Goal

  • 知的な汎用AIや有能なアシスタントを作ること
  • 高性能な日本語LLMを作ること
  • 商用レベルの品質を狙うこと
  • 大規模分散学習
  • 重い学習基盤の導入
  • 複雑な最適化を最初から盛り込むこと

Current Assumptions

  • 開発マシン: Mac mini M4
  • メモリ: 32GB
  • フレームワーク: PyTorch
  • まずはMPSで動かす
  • Python と依存管理は uv を使う

v0 Plan

最初のバージョンでは、以下のような最小構成を目指します。

  1. char-level tokenizer
  2. decoder-only Transformer
  3. next-token predictionで学習
  4. checkpoint保存
  5. 文章生成スクリプト

Why Char-Level

byte-level は実装が単純だが、日本語では UTF-8 の断片を追うことになるため、自然な文章生成がかなり弱くなりやすい。

このリポジトリでは、まず corpus から登場文字を集めて、1文字を1トークンとして扱う。

利点:

  • 日本語の自然さを出しやすい
  • tokenizer 実装がまだ単純
  • 外部 tokenizer に依存しない

欠点:

  • 語彙サイズが corpus 依存で増える
  • 未知文字対応が必要
  • byte-level より実装は少し増える

First Model Target

まずは以下のような小型モデルから始めます。

  • n_layer=4
  • d_model=256
  • n_head=4
  • ffn_hidden=1024
  • context_length=256
  • vocab_size は corpus から動的に決まる

この段階では、性能よりも以下を重視します。

  • 学習が最後まで走る
  • lossが下がる
  • 生成が完全なランダム文字列ではなくなる

Planned Layout

original_LLM/
  README.md
  claude.md
  agent.md
  experiments.md
  data_policy.md
  scripts/
    download_aozora_dazai.py
  src/
    config.py
    data.py
    generate.py
    model.py
    train.py
    utils.py
  data/
    raw/
  checkpoints/
  logs/
  samples/

Milestones

Milestone 1

  • 学習用テキストを1ファイル読み込める
  • 文字列を token 列に変換できる
  • ミニバッチを作れる
  • 青空文庫の太宰作品を取得できる

Milestone 2

  • Transformer本体を実装する
  • 学習ループを回す
  • 小さなデータで過学習を確認する

Milestone 3

  • checkpoint save/load
  • 推論スクリプト
  • temperature / top-k sampling
  • train.py で run_name ごとの checkpoint / sample / log を分離

Milestone 4

  • 複数テキスト対応
  • validation loss
  • 学習再開

Development Principles

  • まずは小さく作る
  • 動くものを優先する
  • 余計な抽象化は後回しにする
  • 1ファイル過学習を最初の成功条件にする
  • 依存関係は増やしすぎない

Experiment Log

試した設定と結果は experiments.md に残す。

最低限、以下を記録する。

  • 何を変えたか
  • best / last の数値
  • 生成の使用感
  • 次に何を試すか

Environment Setup

このリポジトリは uv で Python バージョンと依存関係を管理する。

初回セットアップ:

uv sync

コマンド実行例:

uv run python scripts/download_aozora_dazai.py --limit 5
uv run python src/data.py --limit 2
uv run python src/model.py --context-length 256

Python バージョンは .python-version と pyproject.toml で管理する。 依存更新時は uv add または uv remove を使い、uv lock を更新する。

Training Operations

このリポジトリの学習運用は、最初から以下を前提にする。

  • caffeinate を付けてスリープで学習を止めない
  • tmux 上で学習を実行し、切断や端末終了に備える
  • checkpoint を定期保存し、resume で再開できるようにする
  • 複数実験を並行で回せるように、run ごとに出力先を分離する

想定コマンド例:

caffeinate -dimsu tmux new -s llm-train

tmux の中で学習:

uv run python src/train.py --run-name dazai-debug

再開:

uv run python src/train.py --run-name dazai-debug --resume checkpoints/dazai-debug/latest.pt

並行実験:

uv run python src/train.py --run-name dazai-a --limit 8
uv run python src/train.py --run-name dazai-b --limit 32

並行で回す場合は、checkpoint、sample、log の出力先を run_name 単位で必ず分ける。 同一ファイルへ複数プロセスが書く構成は避ける。

M4 Mac mini 32GB では、重い MPS 学習を同時に複数本走らせると競合しやすい。 そのため、実運用では「重い本学習を1本 + 軽い検証を1本」程度から始める。

現在の train.py は以下を持つ。

  • --run-name
  • --resume
  • --save-every
  • --sample-every
  • --eval-every
  • checkpoints/<run_name>/
  • samples/<run_name>/
  • logs/<run_name>/

Debug Generation

学習済み checkpoint を手で試すには generate.py を使う。

単発生成:

uv run python src/generate.py --checkpoint checkpoints/dazai-long/best.pt --prompt "私は"

対話モード:

uv run python src/generate.py --checkpoint checkpoints/dazai-long/best.pt --interactive

私: / 相手: 形式の会話モデルを対話モードで使う例:

uv run python src/generate.py \
  --checkpoint checkpoints/dazai-friend-simple/best.pt \
  --interactive \
  --user-label 私 \
  --reply-label 相手

開発者向けの会話 CLI 入口:

uv run dazai-chat

使い方はまずこれで確認できる:

uv run dazai-chat --help

これは現在の会話向け checkpoint を自動で拾い、次を既定で有効にする。

  • --interactive
  • --carry-context
  • --max-history-turns 1
  • --user-label 私
  • --reply-label 相手
  • --retrieval-examples 0
  • --normalize-chat-input
  • --temperature 0.2
  • --top-k 8
  • --repetition-penalty 1.1
  • --max-new-tokens 48

interactive では既定で 相手: の返答だけを表示する。 会話 prompt / 生の出力を見たいときだけ --show-prompt-output を付ける。

dazai-chat は会話の質を少し安定させるために、既定で次も行う。

  • 会話履歴は全保持ではなく直近 1 ターンだけを持つ
  • 近い 私: / 相手: 例を 1 件だけ prompt の前に差す
    • 既定では、手書きの data/chat_seed_simple/ があればそれを優先して使う
    • 短い雑談や口語を増やしたい時は、このディレクトリに .txt を足すと retrieval の候補に入る
  • 近い短文 seed がそのまま見つかった時は、その返答を優先して返す
  • 短い入力では、前の一手より現在の入力を優先して解釈する
  • どゆこと？ のような短い口語入力を軽く正規化してから対話する

調整したい時の例:

uv run dazai-chat --max-history-turns 3
uv run dazai-chat --retrieval-examples 0
uv run dazai-chat --no-normalize-chat-input

疑似的に文脈を持たせたい場合:

uv run python src/generate.py --checkpoint checkpoints/dazai-long/best.pt --interactive --carry-context

現在の generate.py は、デバッグ用途を優先して以下を既定で持つ。

  • max_new_tokens=64
  • 。, ！, ？, 」 または空行での簡易 stop
  • interactive 時の prompt / output 表示
  • 必要なら --repetition-penalty 1.1 のような軽い繰り返し抑制

短めに試す例:

uv run python src/generate.py --checkpoint checkpoints/dazai-char/best.pt --interactive --max-new-tokens 40

Data Acquisition

学習データは当面、青空文庫の太宰治作品を使います。

取得スクリプト:

uv run python scripts/download_aozora_dazai.py --limit 5

全件取得:

uv run python scripts/download_aozora_dazai.py

出力先:

  • data/raw/aozora/dazai/zips/
  • data/raw/aozora/dazai/txt/
  • data/raw/aozora/dazai/manifest.jsonl

最初の実験では --limit 付きで数作品だけ取得し、導線が固まったら全件に切り替える。

太宰に近い作家をまとめて取得する mixed corpus 用スクリプト:

uv run python scripts/download_aozora_authors.py

既定の対象:

  • 太宰治
  • 芥川龍之介
  • 坂口安吾
  • 織田作之助
  • 梶井基次郎
  • 中島敦

出力先:

  • data/raw/aozora/dazai_peers/zips/
  • data/raw/aozora/dazai_peers/txt/
  • data/raw/aozora/dazai_peers/manifest.jsonl
  • data/raw/aozora/dazai_peers/summary.json

既存チェックポイントのトークナイザで未知文字率を確認する例:

uv run python scripts/check_tokenizer_coverage.py \
  --checkpoint checkpoints/dazai-char-512x8/best.pt \
  --data-dir data/raw/aozora/dazai_peers/txt \
  --manifest-path data/raw/aozora/dazai_peers/manifest.jsonl

Dialogue Candidate Extraction

青空文庫の太宰テキストから、「...」 に包まれた発話候補を JSONL に抜き出すスクリプトを用意している。

これはそのまま学習に入れるためではなく、後段で AI エージェントに

  • 会話ペアかどうかの判定
  • 酒 / 文学 / 官能 / 日常などの分類
  • 私: / 相手: 形式への軽整形

をさせるための中間データである。

実行例:

uv run python scripts/extract_aozora_dialogue_candidates.py

まず一部だけ試す例:

uv run python scripts/extract_aozora_dialogue_candidates.py --limit 5

主な出力先:

  • data/intermediate/aozora/dazai/dialogue_candidates.jsonl
  • data/intermediate/aozora/dazai/dialogue_summary.json

各候補には、作品ID、タイトル、発話本文、前後文脈、前後の鉤括弧発話、近接する発話クラスタIDを入れている。

AI に渡しやすいように、高スコア候補だけを score 順でバッチ化する例:

uv run python scripts/prepare_dialogue_labeling_batches.py --min-heuristic-score 4 --batch-size 200

出力先:

  • data/intermediate/aozora/dazai/dialogue_batches/

分類ルールは dialogue_labeling.md にまとめている。

ラベル済み batch の検証:

uv run python scripts/validate_dialogue_labels.py

ラベル済み batch から自動 seed を書き出す例:

uv run python scripts/build_chat_seed_from_labels.py

first-pass の自動ラベルを作る例:

uv run python scripts/bootstrap_dialogue_labels.py --batches 0001 0002

Status

データ収集導線の実装を開始。学習系の実装は caffeinate + tmux + resume + run_name 分離 を前提に進める。

About

A small Japanese decoder-only Transformer trained from scratch in PyTorch without pretrained models.

Resources

Readme
Activity

Stars

1 star

Watchers

0 watching

Forks

0 forks
Report repository

Releases

Packages

Contributors

Languages