Add eldamar-tokenizer project

c60ec61 verified 2 months ago

6.16 kB

"""Train ONE (architecture, vocab size) pair and save its checkpoint.
This is the worker invoked once per model by train_all.py. It lives in its own
process on purpose: every architecture folder uses flat module names
(config.py, model.py, block.py, ...), so importing two of them into a single
interpreter would collide in sys.modules. One process = one architecture keeps
each import clean (the repo's "one architecture per kernel" rule).
The training recipe is the repo's own -- a plain windowed next-token loop -- on
the Middle-earth corpus. The tokeniser is selectable: our character-level BPE at
vocab 256 or 512, or the repo's plain CharTokenizer ("char") as a baseline.
Run:  python src/train_one.py qwen3 256
      python src/train_one.py qwen3 char
      python src/train_one.py deepseek3 512 --steps 200   # quick smoke test
"""

from __future__ import annotations

import argparse
import importlib
import sys
from pathlib import Path

import torch

ROOT = Path(__file__).resolve().parents[1]
REPO = ROOT / "single_letter_transformers"
DATA_FILE = ROOT / "data" / "middle_earth_names.txt"
BPE_DIR = ROOT / "bpe"
CKPT_DIR = ROOT / "checkpoints"

# architecture -> (folder, model class, default training steps).
# DeepSeek's sparse routing needs a little longer to settle, as in the repo.
ARCHITECTURES = {
    "qwen3": ("qwen3", "TinyQwen", 5000),
    "qwen3_5": ("qwen3_5", "TinyQwen35", 5000),
    "gemma4": ("gemma4", "TinyGemma", 5000),
    "deepseek3": ("deepseek3", "TinyDeepSeek", 5000),
}

# Shared hyperparameters (identical to the repo's train.py scripts).
BATCH_SIZE = 64
BLOCK_SIZE = 16
LEARNING_RATE = 3e-3
EVAL_EVERY = 200
SEED = 1337


def load_architecture(folder: str, class_name: str):
    """Import one architecture's ModelConfig + model class from its own folder."""
    sys.path.insert(0, str(REPO / folder))
    model_config = importlib.import_module("config").ModelConfig
    model_class = getattr(importlib.import_module("model"), class_name)
    return model_config, model_class


def build_tokenizer(kind: str):
    """Return (tokenizer, label, checkpoint-metadata) for 'char', '256' or '512'.
    Call this only after load_architecture(), because the 'char' baseline imports
    the CharTokenizer from the architecture folder now on sys.path. Both tokenisers
    expose the same interface (encode/decode/vocab_size/newline_id/eos_id).
    """
    if kind == "char":
        # The repo's plain character tokeniser; its vocabulary is built directly
        # from the corpus (42 letters + newline).
        char_tokenizer = importlib.import_module("tokenizer").CharTokenizer
        tokenizer = char_tokenizer.from_file(str(DATA_FILE))
        return tokenizer, "char", {"tokenizer_kind": "char", "chars": tokenizer.chars}

    from bpe_tokenizer import BpeTokenizer

    vocab = int(kind)
    tokenizer = BpeTokenizer.from_file(BPE_DIR / f"bpe_{vocab}.json")
    return tokenizer, f"bpe{vocab}", {"tokenizer_kind": "bpe", "tokenizer": f"bpe/bpe_{vocab}.json"}


def main() -> None:
    parser = argparse.ArgumentParser(description="Train one tiny model on the names corpus.")
    parser.add_argument("arch", choices=sorted(ARCHITECTURES))
    parser.add_argument("tokenizer", choices=("char", "256", "512"),
                        help="'char' = plain CharTokenizer baseline; 256/512 = BPE vocab size")
    parser.add_argument("--steps", type=int, default=None, help="override the default step count")
    args = parser.parse_args()

    folder, class_name, default_steps = ARCHITECTURES[args.arch]
    steps = args.steps if args.steps is not None else default_steps

    model_config, model_class = load_architecture(folder, class_name)

    torch.manual_seed(SEED)
    device = "cuda" if torch.cuda.is_available() else "cpu"

    tokenizer, label, tok_meta = build_tokenizer(args.tokenizer)
    text = DATA_FILE.read_text(encoding="utf-8")
    data = torch.tensor(tokenizer.encode(text), dtype=torch.long)

    cfg = model_config(vocab_size=tokenizer.vocab_size)
    model = model_class(cfg).to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"[{args.arch} · {label}] device={device} "
          f"params={n_params:,} steps={steps} corpus_ids={len(data)}")

    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)

    def get_batch():
        """Sample random windows; targets are inputs shifted by one token."""
        ix = torch.randint(len(data) - BLOCK_SIZE - 1, (BATCH_SIZE,))
        x = torch.stack([data[i:i + BLOCK_SIZE] for i in ix])
        y = torch.stack([data[i + 1:i + 1 + BLOCK_SIZE] for i in ix])
        return x.to(device), y.to(device)

    def sample_names(n: int = 10, max_new_tokens: int = 20) -> list[str]:
        """Generate a few names, each starting from the newline (EOS) token."""
        model.eval()
        start = torch.full((n, 1), tokenizer.newline_id, dtype=torch.long, device=device)
        out = model.generate(start, max_new_tokens=max_new_tokens, temperature=1.0,
                             top_k=None, eos_id=tokenizer.eos_id)
        model.train()
        return [tokenizer.decode(row[1:]).split("\n")[0] for row in out.tolist()]

    final_loss = float("nan")
    for step in range(1, steps + 1):
        x, y = get_batch()
        _, loss = model(x, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        final_loss = loss.item()
        if step % EVAL_EVERY == 0 or step == 1:
            print(f"  step {step:5d}  loss {final_loss:.4f}")

    baseline = torch.log(torch.tensor(float(tokenizer.vocab_size))).item()
    print(f"  baseline loss (uniform guessing): {baseline:.4f}")
    print("  samples: " + ", ".join(sample_names(10)))

    CKPT_DIR.mkdir(parents=True, exist_ok=True)
    out_path = CKPT_DIR / f"{args.arch}_{label}.pt"
    torch.save(
        {
            "model": model.state_dict(),
            "cfg": cfg,
            "arch": args.arch,
            "vocab_size": tokenizer.vocab_size,
            **tok_meta,
        },
        out_path,
    )
    print(f"  saved {out_path.relative_to(ROOT)}  (final loss {final_loss:.4f})")


if __name__ == "__main__":
    main()

	"""Train ONE (architecture, vocab size) pair and save its checkpoint.

	This is the worker invoked once per model by train_all.py. It lives in its own
	process on purpose: every architecture folder uses flat module names
	(config.py, model.py, block.py, ...), so importing two of them into a single
	interpreter would collide in sys.modules. One process = one architecture keeps
	each import clean (the repo's "one architecture per kernel" rule).

	The training recipe is the repo's own -- a plain windowed next-token loop -- on
	the Middle-earth corpus. The tokeniser is selectable: our character-level BPE at
	vocab 256 or 512, or the repo's plain CharTokenizer ("char") as a baseline.

	Run: python src/train_one.py qwen3 256
	python src/train_one.py qwen3 char
	python src/train_one.py deepseek3 512 --steps 200 # quick smoke test
	"""

	from __future__ import annotations

	import argparse
	import importlib
	import sys
	from pathlib import Path

	import torch

	ROOT = Path(__file__).resolve().parents[1]
	REPO = ROOT / "single_letter_transformers"
	DATA_FILE = ROOT / "data" / "middle_earth_names.txt"
	BPE_DIR = ROOT / "bpe"
	CKPT_DIR = ROOT / "checkpoints"

	# architecture -> (folder, model class, default training steps).
	# DeepSeek's sparse routing needs a little longer to settle, as in the repo.
	ARCHITECTURES = {
	"qwen3": ("qwen3", "TinyQwen", 5000),
	"qwen3_5": ("qwen3_5", "TinyQwen35", 5000),
	"gemma4": ("gemma4", "TinyGemma", 5000),
	"deepseek3": ("deepseek3", "TinyDeepSeek", 5000),
	}

	# Shared hyperparameters (identical to the repo's train.py scripts).
	BATCH_SIZE = 64
	BLOCK_SIZE = 16
	LEARNING_RATE = 3e-3
	EVAL_EVERY = 200
	SEED = 1337


	def load_architecture(folder: str, class_name: str):
	"""Import one architecture's ModelConfig + model class from its own folder."""
	sys.path.insert(0, str(REPO / folder))
	model_config = importlib.import_module("config").ModelConfig
	model_class = getattr(importlib.import_module("model"), class_name)
	return model_config, model_class


	def build_tokenizer(kind: str):
	"""Return (tokenizer, label, checkpoint-metadata) for 'char', '256' or '512'.

	Call this only after load_architecture(), because the 'char' baseline imports
	the CharTokenizer from the architecture folder now on sys.path. Both tokenisers
	expose the same interface (encode/decode/vocab_size/newline_id/eos_id).
	"""
	if kind == "char":
	# The repo's plain character tokeniser; its vocabulary is built directly
	# from the corpus (42 letters + newline).
	char_tokenizer = importlib.import_module("tokenizer").CharTokenizer
	tokenizer = char_tokenizer.from_file(str(DATA_FILE))
	return tokenizer, "char", {"tokenizer_kind": "char", "chars": tokenizer.chars}

	from bpe_tokenizer import BpeTokenizer

	vocab = int(kind)
	tokenizer = BpeTokenizer.from_file(BPE_DIR / f"bpe_{vocab}.json")
	return tokenizer, f"bpe{vocab}", {"tokenizer_kind": "bpe", "tokenizer": f"bpe/bpe_{vocab}.json"}


	def main() -> None:
	parser = argparse.ArgumentParser(description="Train one tiny model on the names corpus.")
	parser.add_argument("arch", choices=sorted(ARCHITECTURES))
	parser.add_argument("tokenizer", choices=("char", "256", "512"),
	help="'char' = plain CharTokenizer baseline; 256/512 = BPE vocab size")
	parser.add_argument("--steps", type=int, default=None, help="override the default step count")
	args = parser.parse_args()

	folder, class_name, default_steps = ARCHITECTURES[args.arch]
	steps = args.steps if args.steps is not None else default_steps

	model_config, model_class = load_architecture(folder, class_name)

	torch.manual_seed(SEED)
	device = "cuda" if torch.cuda.is_available() else "cpu"

	tokenizer, label, tok_meta = build_tokenizer(args.tokenizer)
	text = DATA_FILE.read_text(encoding="utf-8")
	data = torch.tensor(tokenizer.encode(text), dtype=torch.long)

	cfg = model_config(vocab_size=tokenizer.vocab_size)
	model = model_class(cfg).to(device)
	n_params = sum(p.numel() for p in model.parameters())
	print(f"[{args.arch} · {label}] device={device} "
	f"params={n_params:,} steps={steps} corpus_ids={len(data)}")

	optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)

	def get_batch():
	"""Sample random windows; targets are inputs shifted by one token."""
	ix = torch.randint(len(data) - BLOCK_SIZE - 1, (BATCH_SIZE,))
	x = torch.stack([data[i:i + BLOCK_SIZE] for i in ix])
	y = torch.stack([data[i + 1:i + 1 + BLOCK_SIZE] for i in ix])
	return x.to(device), y.to(device)

	def sample_names(n: int = 10, max_new_tokens: int = 20) -> list[str]:
	"""Generate a few names, each starting from the newline (EOS) token."""
	model.eval()
	start = torch.full((n, 1), tokenizer.newline_id, dtype=torch.long, device=device)
	out = model.generate(start, max_new_tokens=max_new_tokens, temperature=1.0,
	top_k=None, eos_id=tokenizer.eos_id)
	model.train()
	return [tokenizer.decode(row[1:]).split("\n")[0] for row in out.tolist()]

	final_loss = float("nan")
	for step in range(1, steps + 1):
	x, y = get_batch()
	_, loss = model(x, y)
	optimizer.zero_grad()
	loss.backward()
	optimizer.step()
	final_loss = loss.item()
	if step % EVAL_EVERY == 0 or step == 1:
	print(f" step {step:5d} loss {final_loss:.4f}")

	baseline = torch.log(torch.tensor(float(tokenizer.vocab_size))).item()
	print(f" baseline loss (uniform guessing): {baseline:.4f}")
	print(" samples: " + ", ".join(sample_names(10)))

	CKPT_DIR.mkdir(parents=True, exist_ok=True)
	out_path = CKPT_DIR / f"{args.arch}_{label}.pt"
	torch.save(
	{
	"model": model.state_dict(),
	"cfg": cfg,
	"arch": args.arch,
	"vocab_size": tokenizer.vocab_size,
	**tok_meta,
	},
	out_path,
	)
	print(f" saved {out_path.relative_to(ROOT)} (final loss {final_loss:.4f})")


	if __name__ == "__main__":
	main()