| --- |
| license: mit |
| language: |
| - en |
| tags: |
| - pytorch |
| - gpt |
| - gpt-2 |
| --- |
| # GPT-2 PyTorch |
|
|
|
|
|
|
| The original GPT-2 model weights from [https://openaipublic.blob.core.windows.net/gpt-2/models](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](https://pypi.org/project/llms-from-scratch/) PyPI package: |
|
|
| ```python |
| pip install llms_from_scratch |
| ``` |
|
|
| Or copy the and paste the `GPTModel` class and dependencies from [GitHub](https://github.com/rasbt/LLMs-from-scratch/blob/main/ch05/01_main-chapter-code/previous_chapters.py). |
|
|
|
|
| |
| ### Loading the model weights |
|
|
| The following shows how to load the weights into the 355M parameter model: |
|
|
| ```python |
| 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. |
|
|
| ```python |
| 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. |
| ``` |