rasbt / gpt2-from-scratch-pytorch

GPT-2 PyTorch

The original GPT-2 model weights from https://openaipublic.blob.core.windows.net/gpt-2/models converted from TensorFlow to PyTorch state dicts and PyTorch safetensors files.

Usage

The section below explain how the model weights can be used.

Setup

Install the llms-from-scratch PyPI package:

pip install llms_from_scratch

Or copy the and paste the GPTModel class and dependencies from GitHub .

Loading the model weights

The following shows how to load the weights into the 355M parameter model:

import torch
from llms_from_scratch.ch04 import GPTModel

GPT_CONFIG_BASE = {
    "vocab_size": 50257,    # Vocabulary size
    "context_length": 1024, # Original context length
    "emb_dim": 768,         # Embedding dimension
    "n_heads": 12,          # Number of attention heads
    "n_layers": 12,         # Number of layers
    "drop_rate": 0.0,       # Dropout rate
    "qkv_bias": True        # Query-key-value bias
}

model_configs = {
    "gpt2-small (124M)": {"emb_dim": 768, "n_layers": 12, "n_heads": 12},
    "gpt2-medium (355M)": {"emb_dim": 1024, "n_layers": 24, "n_heads": 16},
    "gpt2-large (774M)": {"emb_dim": 1280, "n_layers": 36, "n_heads": 20},
    "gpt2-xl (1558M)": {"emb_dim": 1600, "n_layers": 48, "n_heads": 25},
}

model_name = "gpt2-medium (355M)"  # Example model name
NEW_CONFIG = GPT_CONFIG_BASE.copy()
NEW_CONFIG.update(model_configs[model_name])

model = GPTModel(NEW_CONFIG)

# Option A: state dict
model.load_state_dict(torch.load("gpt2-medium-355M.pth", weights_only=True));
model.eval();

# Option B: safetensors
# from safetensors.torch import load_file
# model.load_state_dict(load_file("gpt2-medium-355M.safetensors"))

model.eval();

To use the other models, simply replace the model names:

model_name = "gpt2-medium (355M)"
...
model.load_state_dict(torch.load("gpt2-medium-355M.pth"))
# or
model.load_state_dict(load_file("gpt2-medium-355M.safetensors"))

with the desired model names. For example:

model_name = "gpt2-small (124M)"
...
model.load_state_dict(torch.load("gpt2-small-124M.pth"))
# or
model.load_state_dict(load_file("gpt2-small-124M.safetensors"))

Generating text

The following showcases how the model can then be used to generate text.

import tiktoken
from llms_from_scratch.ch04 import generate_text_simple

tokenizer = tiktoken.get_encoding("gpt2")

prompt = "Ever effort moves"
enc_prompt = tokenizer.encode(prompt)
enc_prompt = torch.tensor([enc_prompt])

token_ids = generate_text_simple(
    model=model,
    idx=enc_prompt, 
    max_new_tokens=25, 
    context_size=NEW_CONFIG["context_length"]
)

output = tokenizer.decode(token_ids.squeeze().tolist())
print(output)

Ever effort moves the needle.

The first step is to understand the difference between a "good" and a "bad" goal.