Spaces:
Sleeping
Sleeping
File size: 4,473 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 | from __future__ import annotations
from pathlib import Path
import torch
from tqdm import trange
from tiny_transformer.config import ModelConfig, TrainConfig
from tiny_transformer.data import TextDataset, split_tokens
from tiny_transformer.model import TinyTransformer
from tiny_transformer.tokenizer import BytePairTokenizer, CharTokenizer, Tokenizer, tokenizer_from_dict
@torch.no_grad()
def estimate_loss(
model: TinyTransformer,
train_data: TextDataset,
val_data: TextDataset,
batch_size: int,
eval_batches: int,
) -> dict[str, float]:
model.eval()
losses: dict[str, float] = {}
for split, dataset in {"train": train_data, "val": val_data}.items():
split_losses = []
for _ in range(eval_batches):
x, y = dataset.get_batch(batch_size)
_, loss = model(x, y)
if loss is None:
raise RuntimeError("Expected a loss during evaluation")
split_losses.append(loss.item())
losses[split] = sum(split_losses) / len(split_losses)
model.train()
return losses
def train_from_text(
text: str,
model_config: ModelConfig | None = None,
train_config: TrainConfig | None = None,
device: str = "cpu",
tokenizer_name: str = "char",
bpe_vocab_size: int = 256,
) -> TinyTransformer:
train_config = train_config or TrainConfig()
torch.manual_seed(train_config.seed)
tokenizer = train_tokenizer(text, tokenizer_name, bpe_vocab_size)
token_ids = tokenizer.encode(text)
train_ids, val_ids = split_tokens(token_ids)
if model_config is None:
model_config = ModelConfig(vocab_size=tokenizer.vocab_size)
else:
model_config = ModelConfig(**{**model_config.to_dict(), "vocab_size": tokenizer.vocab_size})
train_data = TextDataset(train_ids, model_config.block_size, device)
val_data = TextDataset(val_ids, model_config.block_size, device)
model = TinyTransformer(model_config).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=train_config.learning_rate)
device_type = "cuda" if device.startswith("cuda") else "mps" if device == "mps" else "cpu"
amp_enabled = train_config.use_amp and device_type in {"cuda", "mps"}
scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled and device_type == "cuda")
progress = trange(train_config.max_steps, desc="training", leave=True)
for step in progress:
if step % train_config.eval_interval == 0 or step == train_config.max_steps - 1:
losses = estimate_loss(
model, train_data, val_data, train_config.batch_size, train_config.eval_batches
)
progress.set_postfix(train=f"{losses['train']:.3f}", val=f"{losses['val']:.3f}")
optimizer.zero_grad(set_to_none=True)
for _ in range(train_config.grad_accum_steps):
x, y = train_data.get_batch(train_config.batch_size)
with torch.autocast(device_type=device_type, enabled=amp_enabled):
_, loss = model(x, y)
if loss is None:
raise RuntimeError("Expected a loss during training")
loss = loss / train_config.grad_accum_steps
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
save_checkpoint(model, tokenizer, train_config.output_path)
return model
def train_tokenizer(text: str, tokenizer_name: str, bpe_vocab_size: int) -> Tokenizer:
if tokenizer_name == "char":
return CharTokenizer.train(text)
if tokenizer_name == "bpe":
return BytePairTokenizer.train(text, vocab_size=bpe_vocab_size)
raise ValueError("tokenizer_name must be 'char' or 'bpe'")
def save_checkpoint(model: TinyTransformer, tokenizer: Tokenizer, path: str) -> None:
output = Path(path)
output.parent.mkdir(parents=True, exist_ok=True)
torch.save(
{
"model_config": model.config.to_dict(),
"model_state": model.state_dict(),
"tokenizer": tokenizer.to_dict(),
},
output,
)
def load_checkpoint(path: str, device: str = "cpu") -> tuple[TinyTransformer, Tokenizer]:
payload = torch.load(path, map_location=device)
tokenizer = tokenizer_from_dict(payload["tokenizer"])
config = ModelConfig(**payload["model_config"])
model = TinyTransformer(config).to(device)
model.load_state_dict(payload["model_state"])
model.eval()
return model, tokenizer
|