Spaces:
Sleeping
Sleeping
File size: 7,014 Bytes
c407ed8 8098a3f c407ed8 8098a3f c407ed8 8098a3f c407ed8 8098a3f c407ed8 8098a3f c407ed8 8098a3f c407ed8 8098a3f c407ed8 8098a3f c407ed8 8098a3f 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 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | from __future__ import annotations
import argparse
from pathlib import Path
import torch
from tiny_transformer.config import ModelConfig, TrainConfig
from tiny_transformer.train import evaluate_checkpoint, 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("--eval-interval", type=int, default=100)
train.add_argument("--eval-batches", type=int, default=20)
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.")
train.add_argument("--loss-history", help="Optional CSV path for train/validation loss history.")
train.add_argument("--resume", help="Resume model and optimizer state from a checkpoint.")
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")
evaluate = subparsers.add_parser("evaluate", help="Report loss and perplexity on a corpus.")
evaluate.add_argument("--checkpoint", required=True)
evaluate.add_argument("--data", required=True)
evaluate.add_argument("--device", default="cpu")
evaluate.add_argument("--batch-size", type=int, default=32)
evaluate.add_argument("--eval-batches", type=int, default=20)
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,
eval_interval=args.eval_interval,
eval_batches=args.eval_batches,
grad_accum_steps=args.grad_accum_steps,
use_amp=args.amp,
output_path=args.output,
loss_history_path=args.loss_history,
)
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,
resume_path=args.resume,
)
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 evaluate_command(args: argparse.Namespace) -> None:
text = Path(args.data).read_text(encoding="utf-8")
metrics = evaluate_checkpoint(
args.checkpoint,
text,
batch_size=args.batch_size,
eval_batches=args.eval_batches,
device=args.device,
)
for name, value in metrics.items():
print(f"{name}: {value:.4f}" if isinstance(value, float) else f"{name}: {value}")
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)
elif args.command == "evaluate":
evaluate_command(args)
else:
parser.error(f"Unknown command: {args.command}")
if __name__ == "__main__":
main()
|