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.
The section below explain how the model weights can be used.
Or copy the and paste the
GPTModel class and dependencies from
GitHub.
1import torch
2from llms_from_scratch.ch04 import GPTModel
3
4GPT_CONFIG_BASE = {
5 "vocab_size": 50257, # Vocabulary size
6 "context_length": 1024, # Original context length
7 "emb_dim": 768, # Embedding dimension
8 "n_heads": 12, # Number of attention heads
9 "n_layers": 12, # Number of layers
10 "drop_rate": 0.0, # Dropout rate
11 "qkv_bias": True # Query-key-value bias
12}
13
14model_configs = {
15 "gpt2-small (124M)": {"emb_dim": 768, "n_layers": 12, "n_heads": 12},
16 "gpt2-medium (355M)": {"emb_dim": 1024, "n_layers": 24, "n_heads": 16},
17 "gpt2-large (774M)": {"emb_dim": 1280, "n_layers": 36, "n_heads": 20},
18 "gpt2-xl (1558M)": {"emb_dim": 1600, "n_layers": 48, "n_heads": 25},
19}
20
21model_name = "gpt2-medium (355M)" # Example model name
22NEW_CONFIG = GPT_CONFIG_BASE.copy()
23NEW_CONFIG.update(model_configs[model_name])
24
25model = GPTModel(NEW_CONFIG)
26
27# Option A: state dict
28model.load_state_dict(torch.load("gpt2-medium-355M.pth", weights_only=True));
29model.eval();
30
31# Option B: safetensors
32# from safetensors.torch import load_file
33# model.load_state_dict(load_file("gpt2-medium-355M.safetensors"))
34
35model.eval();
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"))
The following showcases how the model can then be used to generate text.
1import tiktoken
2from llms_from_scratch.ch04 import generate_text_simple
3
4tokenizer = tiktoken.get_encoding("gpt2")
5
6prompt = "Ever effort moves"
7enc_prompt = tokenizer.encode(prompt)
8enc_prompt = torch.tensor([enc_prompt])
9
10token_ids = generate_text_simple(
11 model=model,
12 idx=enc_prompt,
13 max_new_tokens=25,
14 context_size=NEW_CONFIG["context_length"]
15)
16
17output = tokenizer.decode(token_ids.squeeze().tolist())
18print(output)
Ever effort moves the needle.
The first step is to understand the difference between a "good" and a "bad" goal.