Spaces:
Running on Zero
Running on Zero
File size: 3,494 Bytes
98d1e65 5bdeead 98d1e65 5bdeead 6c01195 5bdeead 6c01195 98d1e65 137f7ef 98d1e65 5bdeead 972b71a 98d1e65 5bdeead 6c01195 98d1e65 6c01195 98d1e65 6c01195 98d1e65 6c01195 98d1e65 6c01195 98d1e65 6c01195 98d1e65 972b71a 98d1e65 6c01195 98d1e65 6c01195 98d1e65 | 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 | from pathlib import Path
import torch
from src.infra.dataset.util import split_tokens, get_batch
from src.infra.tokenizer import CharacterTokenizer
from src.infra.dataset.autoregressive import AutoRegressiveDataset
from src.task.config import ModelConfig,TrainingConfig
from src.task.model import TinyTransformerLM
from src.task.training.autoregressive import train, evaluate
from src.task.training.util import TrainingLogger, count_parameters, separator
# Custom Logger
logger = TrainingLogger(
"logs/text_generator.log"
)
# Runtime configs
corpus = "data/tiny_shakespear.txt"
checkpoint = "mh_tiny_text_generator.pt"
# Configs
separator("Configuration")
torch.manual_seed(42)
# Configs
model_config = ModelConfig(context_length=128, d_model=128, d_ff=512, num_layers=4, number_of_heads=4)
training_config = TrainingConfig()
# Overrides
training_config.num_steps = 10000
training_config.eval_interval = 1000
logger.log(model_config)
logger.log(training_config)
# Load the corpus
separator("Corpus")
data_path = Path(corpus)
text = data_path.read_text(
encoding="utf-8",
)
logger.log(f"Corpus length: {len(text):,} characters")
# Tokens
separator("Tokenization")
tokenizer = CharacterTokenizer(text)
TEXT_TOKENIZER_PATH = Path("checkpoints/shakespeare_tokenizer.json")
tokenizer.save(TEXT_TOKENIZER_PATH)
tokens = torch.tensor(
tokenizer.encode(text),
dtype=torch.long,
)
logger.log(f"Vocabulary size: {tokenizer.vocab_size}")
logger.log(f"Token count: {len(tokens):,}")
# Train/Val/Test split
separator("Data Split")
train_tokens, val_tokens, test_tokens = split_tokens(
tokens,
)
logger.log(f"Train tokens: {len(train_tokens):,}")
logger.log(f"Val tokens: {len(val_tokens):,}")
logger.log(f"Test tokens: {len(test_tokens):,}")
# Datasets
separator("Datasets")
train_dataset = AutoRegressiveDataset(
tokens=train_tokens,
context_length=model_config.context_length,
)
val_dataset = AutoRegressiveDataset(
tokens=val_tokens,
context_length=model_config.context_length,
)
test_dataset = AutoRegressiveDataset(
tokens=test_tokens,
context_length=model_config.context_length,
)
logger.log(f"Train examples: {len(train_dataset):,}")
logger.log(f"Val examples: {len(val_dataset):,}")
logger.log(f"Test examples: {len(test_dataset):,}")
# Models
separator("Model")
model = TinyTransformerLM(
config=model_config,
vocab_size=tokenizer.vocab_size,
)
logger.log(model)
# Train........
separator("Training")
# train(
# model=model,
# train_dataset=train_dataset,
# val_dataset=val_dataset,
# training_config=training_config,
# )
model, steps, train_history, val_history = train(
model=model,
train_dataset=train_dataset,
val_dataset=val_dataset,
training_config=training_config,
)
# Final test evaluation
separator("Test Evaluation")
test_loss = evaluate(
model=model,
dataset=test_dataset,
training_config=training_config,
)
logger.log(f"Test loss: {test_loss:.4f}")
logger.log(f"Parameters: {count_parameters(model):,}")
# --------------------------------------------------
# Save checkpoint
# --------------------------------------------------
separator("Saving")
checkpoint_dir = Path("checkpoints")
checkpoint_dir.mkdir(parents=True, exist_ok=True)
checkpoint_path = (
checkpoint_dir / checkpoint
)
torch.save(
model.state_dict(),
checkpoint_path,
)
print(f"Checkpoint saved to: {checkpoint_path}") |