pick_vocab.py

"""語彙サイズを決めるための実測ツール.

サブワード化の効果は「1トークンあたり何文字を運べるか」で決まるが、
語彙を大きくするほど埋め込みのパラメータが増え、1トークンあたりの
出現回数が減る。手元のコーパス量に対して割に合う点を、実測して選ぶ。

    python tools/pick_vocab.py
    python tools/pick_vocab.py --sizes 4000,8000,32000 --sample-lines 4000

Mac 版は 24,000 で止めた。5090 なら埋め込みが増えても VRAM に余裕があるので
32,000 まで伸ばして、伸びが止まる場所を見る。

既定では **コーパス全体を実際に符号化して** 総トークン数を数える。
標本から外挿すると 1% 前後ずれ、学習に使ったトークン数と突き合わせられない。
--sample-lines に正の数を渡すと標本からの外挿に切り替わる (速いが不正確)。
"""

from __future__ import annotations

import argparse
import json
import sys
import tempfile
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from src.tokenizer import SubwordTokenizer  # noqa: E402


def measure(text: str, size: int, sample: str | None) -> tuple[float, float]:
    """語彙 size で学習し、(文字/トークン, 総トークン数) を返す."""
    with tempfile.TemporaryDirectory() as tmp:
        tokenizer = SubwordTokenizer.train(text, vocab_size=size, model_dir=Path(tmp))
        target = sample if sample is not None else text
        ids = tokenizer.encode(target)
    chars_per_token = len(target) / len(ids)
    total_tokens = len(ids) if sample is None else len(text) / chars_per_token
    return chars_per_token, total_tokens


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--corpus", default="data/corpus.txt")
    ap.add_argument("--sizes", default="4000,8000,12000,16000,24000,32000")
    ap.add_argument("--n-embd", type=int, default=384)
    ap.add_argument(
        "--sample-lines",
        type=int,
        default=0,
        help="0 ならコーパス全体を符号化する (既定)。正の数なら標本から外挿する",
    )
    ap.add_argument("--out", default="runs/vocab_table.json")
    args = ap.parse_args()

    text = Path(args.corpus).read_text(encoding="utf-8")
    lines = text.splitlines()
    sample = "\n".join(lines[: args.sample_lines]) if args.sample_lines > 0 else None
    n_chars_total = len(text)

    print(f"コーパス: {args.corpus}  ({n_chars_total:,} 文字 / {len(lines):,} 会話)")
    if sample is None:
        print("測定方法: コーパス全体を符号化 (外挿なし)\n")
    else:
        print(f"測定方法: 先頭 {args.sample_lines:,} 会話 ({len(sample):,} 文字) から外挿\n")

    header = (
        f"{'語彙':>7}  {'文字/トークン':>12}  {'総トークン数':>13}"
        f"  {'埋め込み':>9}  {'トークン/語彙':>12}  {'学習秒':>7}"
    )
    print(header)
    print("-" * len(header))

    rows = []
    for size in [int(s) for s in args.sizes.split(",")]:
        started = time.time()
        chars_per_token, total_tokens = measure(text, size, sample)
        took = time.time() - started
        embedding_m = size * args.n_embd / 1e6
        print(
            f"{size:>7,}  {chars_per_token:>12.3f}  {total_tokens:>13,.0f}"
            f"  {embedding_m:>8.2f}M  {total_tokens / size:>12,.0f}  {took:>7.1f}"
        )
        rows.append(
            {
                "vocab_size": size,
                "chars_per_token": round(chars_per_token, 3),
                "total_tokens": round(total_tokens),
                "embedding_params_m": round(embedding_m, 2),
                "tokens_per_vocab": round(total_tokens / size),
                "train_sec": round(took, 1),
            }
        )

    print(
        "\n文字/トークン が大きいほど文脈に多くの文章が入る。"
        "\nトークン/語彙 は1つの語彙あたりの平均学習回数。小さすぎると覚えきれない。"
    )

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(
        json.dumps(
            {
                "corpus": args.corpus,
                "chars": n_chars_total,
                "conversations": len(lines),
                "sample_lines": args.sample_lines,
                "n_embd": args.n_embd,
                "rows": rows,
            },
            ensure_ascii=False,
            indent=2,
        ),
        encoding="utf-8",
        newline="\n",
    )
    print(f"\n保存: {out}")


if __name__ == "__main__":
    main()