mirror of
https://github.com/rasbt/LLMs-from-scratch.git
synced 2026-04-10 12:33:42 +00:00
Ch05 supplementary code (#81)
This commit is contained in:
Sebastian Raschka
committed by
GitHub
GitHub
parent
861a2788f3
commit
a2cd8436cb
@@ -1,5 +1,16 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "45398736-7e89-4263-89c8-92153baff553",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"<font size=\"1\">\n",
|
||||
"Supplementary code for \"Build a Large Language Model From Scratch\": <a href=\"https://www.manning.com/books/build-a-large-language-model-from-scratch\">https://www.manning.com/books/build-a-large-language-model-from-scratch</a> by <a href=\"https://sebastianraschka.com\">Sebastian Raschka</a><br>\n",
|
||||
"Code repository: <a href=\"https://github.com/rasbt/LLMs-from-scratch\">https://github.com/rasbt/LLMs-from-scratch</a>\n",
|
||||
"</font>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "66dd524e-864c-4012-b0a2-ccfc56e80024",
|
||||
@@ -2161,7 +2172,7 @@
|
||||
" response = requests.get(url, stream=True)\n",
|
||||
"\n",
|
||||
" # Get the total file size from headers, defaulting to 0 if not present\n",
|
||||
" file_size = int(response.headers.get('content-length', 0))\n",
|
||||
" file_size = int(response.headers.get(\"content-length\", 0))\n",
|
||||
"\n",
|
||||
" # Check if file exists and has the same size\n",
|
||||
" if os.path.exists(destination):\n",
|
||||
@@ -2174,10 +2185,10 @@
|
||||
" block_size = 1024 # 1 Kilobyte\n",
|
||||
"\n",
|
||||
" # Initialize the progress bar with total file size\n",
|
||||
" progress_bar_description = url.split('/')[-1] # Extract filename from URL\n",
|
||||
" with tqdm(total=file_size, unit='iB', unit_scale=True, desc=progress_bar_description) as progress_bar:\n",
|
||||
" progress_bar_description = url.split(\"/\")[-1] # Extract filename from URL\n",
|
||||
" with tqdm(total=file_size, unit=\"iB\", unit_scale=True, desc=progress_bar_description) as progress_bar:\n",
|
||||
" # Open the destination file in binary write mode\n",
|
||||
" with open(destination, 'wb') as file:\n",
|
||||
" with open(destination, \"wb\") as file:\n",
|
||||
" # Iterate over the file data in chunks\n",
|
||||
" for chunk in response.iter_content(block_size):\n",
|
||||
" progress_bar.update(len(chunk)) # Update progress bar\n",
|
||||
|
||||
247
ch05/01_main-chapter-code/generate.py
Normal file
247
ch05/01_main-chapter-code/generate.py
Normal file
@@ -0,0 +1,247 @@
|
||||
# Copyright (c) Sebastian Raschka under Apache License 2.0 (see LICENSE.txt).
|
||||
# Source for "Build a Large Language Model From Scratch"
|
||||
# - https://www.manning.com/books/build-a-large-language-model-from-scratch
|
||||
# Code: https://github.com/rasbt/LLMs-from-scratch
|
||||
|
||||
import json
|
||||
import numpy as np
|
||||
import os
|
||||
import requests
|
||||
import tensorflow as tf
|
||||
import tiktoken
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
# Import from local files
|
||||
from previous_chapters import GPTModel
|
||||
|
||||
|
||||
def text_to_token_ids(text, tokenizer):
|
||||
encoded = tokenizer.encode(text)
|
||||
encoded_tensor = torch.tensor(encoded).unsqueeze(0) # add batch dimension
|
||||
return encoded_tensor
|
||||
|
||||
|
||||
def token_ids_to_text(token_ids, tokenizer):
|
||||
flat = token_ids.squeeze(0) # remove batch dimension
|
||||
return tokenizer.decode(flat.tolist())
|
||||
|
||||
|
||||
def download_and_load_gpt2(model_size, models_dir):
|
||||
# Validate model size
|
||||
allowed_sizes = ("124M", "355M", "774M", "1558M")
|
||||
if model_size not in allowed_sizes:
|
||||
raise ValueError(f"Model size not in {allowed_sizes}")
|
||||
|
||||
# Define paths
|
||||
model_dir = os.path.join(models_dir, model_size)
|
||||
base_url = "https://openaipublic.blob.core.windows.net/gpt-2/models"
|
||||
filenames = [
|
||||
"checkpoint", "encoder.json", "hparams.json",
|
||||
"model.ckpt.data-00000-of-00001", "model.ckpt.index",
|
||||
"model.ckpt.meta", "vocab.bpe"
|
||||
]
|
||||
|
||||
# Download files
|
||||
os.makedirs(model_dir, exist_ok=True)
|
||||
for filename in filenames:
|
||||
file_url = os.path.join(base_url, model_size, filename)
|
||||
file_path = os.path.join(model_dir, filename)
|
||||
download_file(file_url, file_path)
|
||||
|
||||
# Load hparams and params
|
||||
tf_ckpt_path = tf.train.latest_checkpoint(model_dir)
|
||||
hparams = json.load(open(os.path.join(model_dir, "hparams.json")))
|
||||
params = load_gpt2_params_from_tf_ckpt(tf_ckpt_path, hparams)
|
||||
|
||||
return hparams, params
|
||||
|
||||
|
||||
def download_file(url, destination):
|
||||
# Send a GET request to download the file in streaming mode
|
||||
response = requests.get(url, stream=True)
|
||||
|
||||
# Get the total file size from headers, defaulting to 0 if not present
|
||||
file_size = int(response.headers.get("content-length", 0))
|
||||
|
||||
# Check if file exists and has the same size
|
||||
if os.path.exists(destination):
|
||||
file_size_local = os.path.getsize(destination)
|
||||
if file_size == file_size_local:
|
||||
print(f"File already exists and is up-to-date: {destination}")
|
||||
return
|
||||
|
||||
# Define the block size for reading the file
|
||||
block_size = 1024 # 1 Kilobyte
|
||||
|
||||
# Initialize the progress bar with total file size
|
||||
progress_bar_description = url.split("/")[-1] # Extract filename from URL
|
||||
with tqdm(total=file_size, unit="iB", unit_scale=True, desc=progress_bar_description) as progress_bar:
|
||||
# Open the destination file in binary write mode
|
||||
with open(destination, "wb") as file:
|
||||
# Iterate over the file data in chunks
|
||||
for chunk in response.iter_content(block_size):
|
||||
progress_bar.update(len(chunk)) # Update progress bar
|
||||
file.write(chunk) # Write the chunk to the file
|
||||
|
||||
|
||||
def load_gpt2_params_from_tf_ckpt(ckpt_path, hparams):
|
||||
# Initialize parameters dictionary with empty blocks for each layer
|
||||
params = {"blocks": [{} for _ in range(hparams["n_layer"])]}
|
||||
|
||||
# Iterate over each variable in the checkpoint
|
||||
for name, _ in tf.train.list_variables(ckpt_path):
|
||||
# Load the variable and remove singleton dimensions
|
||||
variable_array = np.squeeze(tf.train.load_variable(ckpt_path, name))
|
||||
|
||||
# Process the variable name to extract relevant parts
|
||||
variable_name_parts = name.split("/")[1:] # Skip the 'model/' prefix
|
||||
|
||||
# Identify the target dictionary for the variable
|
||||
target_dict = params
|
||||
if variable_name_parts[0].startswith("h"):
|
||||
layer_number = int(variable_name_parts[0][1:])
|
||||
target_dict = params["blocks"][layer_number]
|
||||
|
||||
# Recursively access or create nested dictionaries
|
||||
for key in variable_name_parts[1:-1]:
|
||||
target_dict = target_dict.setdefault(key, {})
|
||||
|
||||
# Assign the variable array to the last key
|
||||
last_key = variable_name_parts[-1]
|
||||
target_dict[last_key] = variable_array
|
||||
|
||||
return params
|
||||
|
||||
|
||||
def assign(left, right):
|
||||
if left.shape != right.shape:
|
||||
raise ValueError(f"Shape mismatch. Left: {left.shape}, Right: {right.shape}")
|
||||
return torch.nn.Parameter(torch.tensor(right))
|
||||
|
||||
|
||||
def load_weights_into_gpt(gpt, params):
|
||||
# Weight tying
|
||||
gpt.pos_emb.weight = assign(gpt.pos_emb.weight, params['wpe'])
|
||||
gpt.tok_emb.weight = assign(gpt.tok_emb.weight, params['wte'])
|
||||
|
||||
for b in range(len(params["blocks"])):
|
||||
q_w, k_w, v_w = np.split((params["blocks"][b]["attn"]["c_attn"])["w"], 3, axis=-1)
|
||||
gpt.trf_blocks[b].att.W_query.weight = assign(gpt.trf_blocks[b].att.W_query.weight, q_w.T)
|
||||
gpt.trf_blocks[b].att.W_key.weight = assign(gpt.trf_blocks[b].att.W_key.weight, k_w.T)
|
||||
gpt.trf_blocks[b].att.W_value.weight = assign(gpt.trf_blocks[b].att.W_value.weight, v_w.T)
|
||||
|
||||
q_b, k_b, v_b = np.split((params["blocks"][b]["attn"]["c_attn"])["b"], 3, axis=-1)
|
||||
gpt.trf_blocks[b].att.W_query.bias = assign(gpt.trf_blocks[b].att.W_query.bias, q_b)
|
||||
gpt.trf_blocks[b].att.W_key.bias = assign(gpt.trf_blocks[b].att.W_key.bias, k_b)
|
||||
gpt.trf_blocks[b].att.W_value.bias = assign(gpt.trf_blocks[b].att.W_value.bias, v_b)
|
||||
|
||||
gpt.trf_blocks[b].att.out_proj.weight = assign(gpt.trf_blocks[b].att.out_proj.weight, params["blocks"][b]["attn"]["c_proj"]["w"].T)
|
||||
gpt.trf_blocks[b].att.out_proj.bias = assign(gpt.trf_blocks[b].att.out_proj.bias, params["blocks"][b]["attn"]["c_proj"]["b"])
|
||||
|
||||
gpt.trf_blocks[b].ff.layers[0].weight = assign(gpt.trf_blocks[b].ff.layers[0].weight, params["blocks"][b]["mlp"]["c_fc"]["w"].T)
|
||||
gpt.trf_blocks[b].ff.layers[0].bias = assign(gpt.trf_blocks[b].ff.layers[0].bias, params["blocks"][b]["mlp"]["c_fc"]["b"])
|
||||
gpt.trf_blocks[b].ff.layers[2].weight = assign(gpt.trf_blocks[b].ff.layers[2].weight, params["blocks"][b]["mlp"]["c_proj"]["w"].T)
|
||||
gpt.trf_blocks[b].ff.layers[2].bias = assign(gpt.trf_blocks[b].ff.layers[2].bias, params["blocks"][b]["mlp"]["c_proj"]["b"])
|
||||
|
||||
gpt.trf_blocks[b].norm1.scale = assign(gpt.trf_blocks[b].norm1.scale, params["blocks"][b]["ln_1"]["g"])
|
||||
gpt.trf_blocks[b].norm1.shift = assign(gpt.trf_blocks[b].norm1.shift, params["blocks"][b]["ln_1"]["b"])
|
||||
gpt.trf_blocks[b].norm2.scale = assign(gpt.trf_blocks[b].norm2.scale, params["blocks"][b]["ln_2"]["g"])
|
||||
gpt.trf_blocks[b].norm2.shift = assign(gpt.trf_blocks[b].norm2.shift, params["blocks"][b]["ln_2"]["b"])
|
||||
|
||||
gpt.final_norm.scale = assign(gpt.final_norm.scale, params["g"])
|
||||
gpt.final_norm.shift = assign(gpt.final_norm.shift, params["b"])
|
||||
gpt.out_head.weight = assign(gpt.out_head.weight, params["wte"])
|
||||
|
||||
|
||||
def generate(model, idx, max_new_tokens, context_size, temperature, top_k=None):
|
||||
|
||||
# For-loop is the same as before: Get logits, and only focus on last time step
|
||||
for _ in range(max_new_tokens):
|
||||
idx_cond = idx[:, -context_size:]
|
||||
with torch.no_grad():
|
||||
logits = model(idx_cond)
|
||||
logits = logits[:, -1, :]
|
||||
|
||||
# New: Filter logits with top_k sampling
|
||||
if top_k is not None:
|
||||
# Keep only top_k values
|
||||
top_logits, _ = torch.topk(logits, top_k)
|
||||
min_val = top_logits[:, -1]
|
||||
logits = torch.where(logits < min_val, torch.tensor(float('-inf')).to(logits.device), logits)
|
||||
|
||||
# New: Apply temperature scaling
|
||||
if temperature > 0.0:
|
||||
logits = logits / temperature
|
||||
|
||||
# Apply softmax to get probabilities
|
||||
probs = torch.softmax(logits, dim=-1) # (batch_size, context_len)
|
||||
|
||||
# Sample from the distribution
|
||||
idx_next = torch.multinomial(probs, num_samples=1) # (batch_size, 1)
|
||||
|
||||
# Otherwise same as before: get idx of the vocab entry with the highest logits value
|
||||
else:
|
||||
idx_next = torch.argmax(logits, dim=-1, keepdim=True) # (batch_size, 1)
|
||||
|
||||
# Same as before: append sampled index to the running sequence
|
||||
idx = torch.cat((idx, idx_next), dim=1) # (batch_size, num_tokens+1)
|
||||
|
||||
return idx
|
||||
|
||||
|
||||
def main(gpt_config, input_prompt, model_size):
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
hparams, params = download_and_load_gpt2(model_size=model_size, models_dir="gpt2")
|
||||
|
||||
gpt = GPTModel(gpt_config)
|
||||
load_weights_into_gpt(gpt, params)
|
||||
gpt.to(device)
|
||||
|
||||
tokenizer = tiktoken.get_encoding("gpt2")
|
||||
|
||||
token_ids = generate(
|
||||
model=gpt,
|
||||
idx=text_to_token_ids(input_prompt, tokenizer),
|
||||
max_new_tokens=65,
|
||||
context_size=gpt_config["ctx_len"],
|
||||
top_k=50,
|
||||
temperature=1.5
|
||||
)
|
||||
|
||||
print("Output text:\n", token_ids_to_text(token_ids, tokenizer))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
torch.manual_seed(123)
|
||||
|
||||
CHOOSE_MODEL = "gpt2-small"
|
||||
INPUT_PROMPT = "Every effort moves you"
|
||||
|
||||
BASE_CONFIG = {
|
||||
"vocab_size": 50257, # Vocabulary size
|
||||
"ctx_len": 1024, # Context length
|
||||
"drop_rate": 0.0, # Dropout rate
|
||||
"qkv_bias": True # Query-key-value bias
|
||||
}
|
||||
|
||||
model_configs = {
|
||||
"gpt2-small": {"emb_dim": 768, "n_layers": 12, "n_heads": 12},
|
||||
"gpt2-medium": {"emb_dim": 1024, "n_layers": 24, "n_heads": 16},
|
||||
"gpt2-large": {"emb_dim": 1280, "n_layers": 36, "n_heads": 20},
|
||||
"gpt2-xl": {"emb_dim": 1600, "n_layers": 48, "n_heads": 25},
|
||||
}
|
||||
|
||||
model_sizes = {
|
||||
"gpt2-small": "124M",
|
||||
"gpt2-medium": "355M",
|
||||
"gpt2-large": "774M",
|
||||
"gpt2-xl": "1558"
|
||||
}
|
||||
|
||||
BASE_CONFIG.update(model_configs[CHOOSE_MODEL])
|
||||
|
||||
main(BASE_CONFIG, INPUT_PROMPT, model_sizes[CHOOSE_MODEL])
|
||||
40
ch05/01_main-chapter-code/tests.py
Normal file
40
ch05/01_main-chapter-code/tests.py
Normal file
@@ -0,0 +1,40 @@
|
||||
# Copyright (c) Sebastian Raschka under Apache License 2.0 (see LICENSE.txt).
|
||||
# Source for "Build a Large Language Model From Scratch"
|
||||
# - https://www.manning.com/books/build-a-large-language-model-from-scratch
|
||||
# Code: https://github.com/rasbt/LLMs-from-scratch
|
||||
|
||||
# File for internal use (unit tests)
|
||||
|
||||
import pytest
|
||||
from train import main
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def gpt_config():
|
||||
return {
|
||||
"vocab_size": 50257,
|
||||
"ctx_len": 12, # small for testing efficiency
|
||||
"emb_dim": 32, # small for testing efficiency
|
||||
"n_heads": 4, # small for testing efficiency
|
||||
"n_layers": 2, # small for testing efficiency
|
||||
"drop_rate": 0.1,
|
||||
"qkv_bias": False
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def other_hparams():
|
||||
return {
|
||||
"learning_rate": 5e-4,
|
||||
"num_epochs": 1, # small for testing efficiency
|
||||
"batch_size": 2,
|
||||
"weight_decay": 0.1
|
||||
}
|
||||
|
||||
|
||||
def test_main(gpt_config, other_hparams):
|
||||
train_losses, val_losses, tokens_seen, model = main(gpt_config, other_hparams)
|
||||
|
||||
assert len(train_losses) == 39, "Unexpected number of training losses"
|
||||
assert len(val_losses) == 39, "Unexpected number of validation losses"
|
||||
assert len(tokens_seen) == 39, "Unexpected number of tokens seen"
|
||||
235
ch05/01_main-chapter-code/train.py
Normal file
235
ch05/01_main-chapter-code/train.py
Normal file
@@ -0,0 +1,235 @@
|
||||
# Copyright (c) Sebastian Raschka under Apache License 2.0 (see LICENSE.txt).
|
||||
# Source for "Build a Large Language Model From Scratch"
|
||||
# - https://www.manning.com/books/build-a-large-language-model-from-scratch
|
||||
# Code: https://github.com/rasbt/LLMs-from-scratch
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import os
|
||||
import torch
|
||||
import urllib.request
|
||||
|
||||
# Import from local files
|
||||
from previous_chapters import GPTModel, create_dataloader_v1, generate_text_simple
|
||||
|
||||
|
||||
def text_to_token_ids(text, tokenizer):
|
||||
encoded = tokenizer.encode(text)
|
||||
encoded_tensor = torch.tensor(encoded).unsqueeze(0) # add batch dimension
|
||||
return encoded_tensor
|
||||
|
||||
|
||||
def token_ids_to_text(token_ids, tokenizer):
|
||||
flat = token_ids.squeeze(0) # remove batch dimension
|
||||
return tokenizer.decode(flat.tolist())
|
||||
|
||||
|
||||
def calc_loss_batch(input_batch, target_batch, model, device):
|
||||
input_batch, target_batch = input_batch.to(device), target_batch.to(device)
|
||||
|
||||
logits = model(input_batch)
|
||||
logits = logits.view(-1, logits.size(-1))
|
||||
loss = torch.nn.functional.cross_entropy(logits, target_batch.view(-1))
|
||||
return loss
|
||||
|
||||
|
||||
def calc_loss_loader(data_loader, model, device, num_batches=None):
|
||||
total_loss, batches_seen = 0., 0.
|
||||
if num_batches is None:
|
||||
num_batches = len(data_loader)
|
||||
for i, (input_batch, target_batch) in enumerate(data_loader):
|
||||
if i < num_batches:
|
||||
loss = calc_loss_batch(input_batch, target_batch, model, device)
|
||||
total_loss += loss.item()
|
||||
batches_seen += 1
|
||||
else:
|
||||
break
|
||||
return total_loss / batches_seen
|
||||
|
||||
|
||||
def evaluate_model(model, train_loader, val_loader, device, eval_iter):
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
train_loss = calc_loss_loader(train_loader, model, device, num_batches=eval_iter)
|
||||
val_loss = calc_loss_loader(val_loader, model, device, num_batches=eval_iter)
|
||||
model.train()
|
||||
return train_loss, val_loss
|
||||
|
||||
|
||||
def generate_and_print_sample(model, tokenizer, device, start_context):
|
||||
model.eval()
|
||||
context_size = model.pos_emb.weight.shape[0]
|
||||
encoded = text_to_token_ids(start_context, tokenizer).to(device)
|
||||
with torch.no_grad():
|
||||
token_ids = generate_text_simple(
|
||||
model=model, idx=encoded,
|
||||
max_new_tokens=50, context_size=context_size
|
||||
)
|
||||
decoded_text = token_ids_to_text(token_ids, tokenizer)
|
||||
print(decoded_text.replace("\n", " ")) # Compact print format
|
||||
model.train()
|
||||
|
||||
|
||||
def train_model_simple(model, train_loader, val_loader, optimizer, device, num_epochs,
|
||||
eval_freq, eval_iter, start_context):
|
||||
# Initialize lists to track losses and tokens seen
|
||||
train_losses, val_losses, track_tokens_seen = [], [], []
|
||||
tokens_seen = 0
|
||||
global_step = -1
|
||||
|
||||
# Main training loop
|
||||
for epoch in range(num_epochs):
|
||||
model.train() # Set model to training mode
|
||||
|
||||
for input_batch, target_batch in train_loader:
|
||||
optimizer.zero_grad() # Reset loss gradients from previous epoch
|
||||
loss = calc_loss_batch(input_batch, target_batch, model, device)
|
||||
loss.backward() # Calculate loss gradients
|
||||
optimizer.step() # Update model weights using loss gradients
|
||||
tokens_seen += input_batch.numel()
|
||||
global_step += 1
|
||||
|
||||
# Optional evaluation step
|
||||
if global_step % eval_freq == 0:
|
||||
train_loss, val_loss = evaluate_model(
|
||||
model, train_loader, val_loader, device, eval_iter)
|
||||
train_losses.append(train_loss)
|
||||
val_losses.append(val_loss)
|
||||
track_tokens_seen.append(tokens_seen)
|
||||
print(f"Ep {epoch+1} (Step {global_step:06d}): "
|
||||
f"Train loss {train_loss:.3f}, Val loss {val_loss:.3f}")
|
||||
|
||||
# Print a sample text after each epoch
|
||||
generate_and_print_sample(
|
||||
model, train_loader.dataset.tokenizer, device, start_context
|
||||
)
|
||||
|
||||
return train_losses, val_losses, track_tokens_seen
|
||||
|
||||
|
||||
def plot_losses(epochs_seen, tokens_seen, train_losses, val_losses):
|
||||
fig, ax1 = plt.subplots()
|
||||
|
||||
# Plot training and validation loss against epochs
|
||||
ax1.plot(epochs_seen, train_losses, label="Training loss")
|
||||
ax1.plot(epochs_seen, val_losses, linestyle="-.", label="Validation loss")
|
||||
ax1.set_xlabel("Epochs")
|
||||
ax1.set_ylabel("Loss")
|
||||
ax1.legend(loc="upper right")
|
||||
|
||||
# Create a second x-axis for tokens seen
|
||||
ax2 = ax1.twiny() # Create a second x-axis that shares the same y-axis
|
||||
ax2.plot(tokens_seen, train_losses, alpha=0) # Invisible plot for aligning ticks
|
||||
ax2.set_xlabel("Tokens seen")
|
||||
|
||||
fig.tight_layout() # Adjust layout to make room
|
||||
# plt.show()
|
||||
|
||||
|
||||
def main(gpt_config, hparams):
|
||||
|
||||
torch.manual_seed(123)
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
##############################
|
||||
# Download data if necessary
|
||||
##############################
|
||||
|
||||
file_path = "the-verdict.txt"
|
||||
url = "https://raw.githubusercontent.com/rasbt/LLMs-from-scratch/main/ch02/01_main-chapter-code/the-verdict.txt"
|
||||
|
||||
if not os.path.exists(file_path):
|
||||
with urllib.request.urlopen(url) as response:
|
||||
text_data = response.read().decode('utf-8')
|
||||
with open(file_path, "w", encoding="utf-8") as file:
|
||||
file.write(text_data)
|
||||
else:
|
||||
with open(file_path, "r", encoding="utf-8") as file:
|
||||
text_data = file.read()
|
||||
|
||||
##############################
|
||||
# Initialize model
|
||||
##############################
|
||||
|
||||
model = GPTModel(gpt_config)
|
||||
model.to(device) # no assignment model = model.to(device) necessary for nn.Module classes
|
||||
optimizer = torch.optim.AdamW(
|
||||
model.parameters(), lr=hparams["learning_rate"], weight_decay=hparams["weight_decay"]
|
||||
)
|
||||
|
||||
##############################
|
||||
# Set up dataloaders
|
||||
##############################
|
||||
|
||||
# Train/validation ratio
|
||||
train_ratio = 0.90
|
||||
split_idx = int(train_ratio * len(text_data))
|
||||
|
||||
train_loader = create_dataloader_v1(
|
||||
text_data[:split_idx],
|
||||
batch_size=hparams["batch_size"],
|
||||
max_length=gpt_config["ctx_len"],
|
||||
stride=gpt_config["ctx_len"],
|
||||
drop_last=True,
|
||||
shuffle=True
|
||||
)
|
||||
|
||||
val_loader = create_dataloader_v1(
|
||||
text_data[split_idx:],
|
||||
batch_size=hparams["batch_size"],
|
||||
max_length=gpt_config["ctx_len"],
|
||||
stride=gpt_config["ctx_len"],
|
||||
drop_last=False,
|
||||
shuffle=False
|
||||
)
|
||||
|
||||
##############################
|
||||
# Train model
|
||||
##############################
|
||||
|
||||
train_losses, val_losses, tokens_seen = train_model_simple(
|
||||
model, train_loader, val_loader, optimizer, device,
|
||||
num_epochs=hparams["num_epochs"], eval_freq=5, eval_iter=1,
|
||||
start_context="Every effort moves you",
|
||||
)
|
||||
|
||||
return train_losses, val_losses, tokens_seen, model
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
GPT_CONFIG_124M = {
|
||||
"vocab_size": 50257, # Vocabulary size
|
||||
"ctx_len": 256, # Shortened context length (orig: 1024)
|
||||
"emb_dim": 768, # Embedding dimension
|
||||
"n_heads": 12, # Number of attention heads
|
||||
"n_layers": 12, # Number of layers
|
||||
"drop_rate": 0.1, # Dropout rate
|
||||
"qkv_bias": False # Query-key-value bias
|
||||
}
|
||||
|
||||
OTHER_HPARAMS = {
|
||||
"learning_rate": 5e-4,
|
||||
"num_epochs": 10,
|
||||
"batch_size": 2,
|
||||
"weight_decay": 0.1
|
||||
}
|
||||
|
||||
###########################
|
||||
# Initiate training
|
||||
###########################
|
||||
|
||||
train_losses, val_losses, tokens_seen, model = main(GPT_CONFIG_124M, OTHER_HPARAMS)
|
||||
|
||||
###########################
|
||||
# After training
|
||||
###########################
|
||||
|
||||
# Plot results
|
||||
epochs_tensor = torch.linspace(0, OTHER_HPARAMS["num_epochs"], len(train_losses))
|
||||
plot_losses(epochs_tensor, tokens_seen, train_losses, val_losses)
|
||||
plt.savefig("loss.pdf")
|
||||
|
||||
# Save and load model
|
||||
torch.save(model.state_dict(), "model.pth")
|
||||
model = GPTModel(GPT_CONFIG_124M)
|
||||
model.load_state_dict(torch.load("model.pth"))
|
||||
Reference in New Issue
Block a user