Spaces:
Sleeping
Sleeping
File size: 5,624 Bytes
c407ed8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | from __future__ import annotations
import argparse
from pathlib import Path
import torch
from tiny_transformer.config import ModelConfig, TrainConfig
from tiny_transformer.train import load_checkpoint, train_from_text
from tiny_transformer.visualize import save_attention_heatmap
from tiny_transformer.web import serve_playground
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Train and sample a tiny GPT-style Transformer.")
subparsers = parser.add_subparsers(dest="command", required=True)
train = subparsers.add_parser("train", help="Train a model on a plain-text corpus.")
train.add_argument("--data", required=True, help="Path to a UTF-8 text file.")
train.add_argument("--output", default="runs/tiny-transformer.pt", help="Checkpoint path.")
train.add_argument("--device", default="cpu", help="Device such as cpu, cuda, or mps.")
train.add_argument("--steps", type=int, default=1_000)
train.add_argument("--batch-size", type=int, default=32)
train.add_argument("--block-size", type=int, default=32)
train.add_argument("--layers", type=int, default=4)
train.add_argument("--heads", type=int, default=4)
train.add_argument("--embedding", type=int, default=128)
train.add_argument("--dropout", type=float, default=0.1)
train.add_argument("--learning-rate", type=float, default=3e-4)
train.add_argument("--tokenizer", choices=["char", "bpe"], default="char")
train.add_argument("--bpe-vocab-size", type=int, default=256)
train.add_argument("--grad-accum-steps", type=int, default=1)
train.add_argument("--amp", action="store_true", help="Use mixed precision on CUDA or MPS.")
generate = subparsers.add_parser("generate", help="Generate text from a trained checkpoint.")
generate.add_argument("--checkpoint", required=True, help="Path to a model checkpoint.")
generate.add_argument("--prompt", default="\n", help="Prompt text.")
generate.add_argument("--device", default="cpu")
generate.add_argument("--max-new-tokens", type=int, default=200)
generate.add_argument("--temperature", type=float, default=0.8)
generate.add_argument("--top-k", type=int, default=20)
attention = subparsers.add_parser("attention", help="Export an attention heatmap SVG.")
attention.add_argument("--checkpoint", required=True, help="Path to a model checkpoint.")
attention.add_argument("--prompt", required=True, help="Prompt text to inspect.")
attention.add_argument("--output", default="runs/attention.svg", help="SVG output path.")
attention.add_argument("--device", default="cpu")
attention.add_argument("--layer", type=int, default=-1, help="Layer index to visualize.")
attention.add_argument("--head", type=int, default=0, help="Attention head index to visualize.")
serve = subparsers.add_parser("serve", help="Launch a local text-generation playground.")
serve.add_argument("--checkpoint", required=True, help="Path to a model checkpoint.")
serve.add_argument("--host", default="127.0.0.1")
serve.add_argument("--port", type=int, default=8000)
serve.add_argument("--device", default="cpu")
return parser
def train_command(args: argparse.Namespace) -> None:
text = Path(args.data).read_text(encoding="utf-8")
train_config = TrainConfig(
batch_size=args.batch_size,
learning_rate=args.learning_rate,
max_steps=args.steps,
grad_accum_steps=args.grad_accum_steps,
use_amp=args.amp,
output_path=args.output,
)
model_config = ModelConfig(
vocab_size=1,
block_size=args.block_size,
n_layer=args.layers,
n_head=args.heads,
n_embd=args.embedding,
dropout=args.dropout,
)
train_from_text(
text,
model_config=model_config,
train_config=train_config,
device=args.device,
tokenizer_name=args.tokenizer,
bpe_vocab_size=args.bpe_vocab_size,
)
print(f"Saved checkpoint to {args.output}")
def generate_command(args: argparse.Namespace) -> None:
model, tokenizer = load_checkpoint(args.checkpoint, device=args.device)
encoded = tokenizer.encode(args.prompt)
idx = torch.tensor([encoded], dtype=torch.long, device=args.device)
out = model.generate(
idx,
max_new_tokens=args.max_new_tokens,
temperature=args.temperature,
top_k=args.top_k,
)
print(tokenizer.decode(out[0].tolist()))
def attention_command(args: argparse.Namespace) -> None:
model, tokenizer = load_checkpoint(args.checkpoint, device=args.device)
encoded = tokenizer.encode(args.prompt)
idx = torch.tensor([encoded], dtype=torch.long, device=args.device)
save_attention_heatmap(
model=model,
tokenizer=tokenizer,
idx=idx,
output_path=args.output,
layer=args.layer,
head=args.head,
)
print(f"Saved attention heatmap to {args.output}")
def serve_command(args: argparse.Namespace) -> None:
serve_playground(
checkpoint=args.checkpoint,
host=args.host,
port=args.port,
device=args.device,
)
def main() -> None:
parser = build_parser()
args = parser.parse_args()
if args.command == "train":
train_command(args)
elif args.command == "generate":
generate_command(args)
elif args.command == "attention":
attention_command(args)
elif args.command == "serve":
serve_command(args)
else:
parser.error(f"Unknown command: {args.command}")
if __name__ == "__main__":
main()
|