GPT-2 From Scratch

A PyTorch implementation of the GPT-2 architecture built for experimentation. Use it to train models on custom text or run inference using the original OpenAI weights.

---

Table of Contents

Table of Contents

1. Project Structure
2. Model Architecture
3. Attention Tensor Shapes
4. Key Features
5. Setup & Installation
6. Usage: Training
7. Usage: Inference
8. Fine-Tuning
9. Interactive Web UI
10. Contributing
11. Acknowledgments

---

Project Structure

GPT2-From-Scratch/
├── src/                        # Core GPT-2 Source Code
│   ├── model.py                # The GPT-2 Architecture (Blocks, Attention)
│   ├── generate.py             # Text generation logic (Top-K, Temperature)
│   ├── dataset.py              # Sliding window dataset logic for pre-training
│   ├── train_utils.py          # General training utilities
│   ├── download_weights.py     # TF checkpoint converter
│   └── download_weights_hf.py  # Hugging Face weights loader
├── classification-finetuning/  # Standalone module for Sentiment Analysis
├── instruction-finetuning/     # Standalone module for Instruction Following
├── data/                       # Main Dataset storage (input.txt)
├── model_weights/              # Stores Pre-trained GPT-2 checkpoints
├── train.ipynb                 # Main Training Loop (Pre-training)
├── inference.ipynb             # Inference using Raw/TF converted weights
├── inference_hf.ipynb          # Inference using Hugging Face weights
├── requirements.txt            # Python dependencies
├── README.md                   # Project documentation
├── app.py                      # Chainlit Interactive Web UI application
├── chainlit.md                 # Welcome screen configuration for the Web UI
├── assets/                     # Images and diagrams (e.g., model_architecture.png)
├── LICENSE
├── .gitignore
└── .venv/                      # Virtual Environment

---

Model Architecture

The model follows the standard GPT-2 (Decoder-Only Transformer) design. It processes sequences of tokens to predict the next token in the sequence.

---

Attention Tensor Shapes

Understanding the dimensionality of tensors as they flow through the network is critical. This table breaks down the exact tensor shape transformations that occur during the Multi-Head Attention forward pass within the model configuration.

Step	Before Shape	Operation	After Shape	Major Change
1. Input X	(batch_size, num_tokens, d_in)	-	(batch_size, num_tokens, d_in)	No change
2. Q, K, V projections	(batch_size, num_tokens, d_in)	Linear (W_q, W_k, W_v)	(batch_size, num_tokens, d_out)	d_in → d_out
3. Split into heads	(batch_size, num_tokens, d_out)	Reshape into heads	(batch_size, num_tokens, num_heads, head_dim)	d_out → num_heads × head_dim
4. Transpose for attention	(batch_size, num_tokens, num_heads, head_dim)	Swap axes (1,2)	(batch_size, num_heads, num_tokens, head_dim)	num_tokens ↔ num_heads
5. Compute Attention Scores	(batch_size, num_heads, num_tokens, head_dim) and (batch_size, num_heads, head_dim, num_tokens)	Q @ Kᵀ	(batch_size, num_heads, num_tokens, num_tokens)	head_dim → num_tokens
6. Apply Mask & Softmax	(batch_size, num_heads, num_tokens, num_tokens)	Mask + Softmax	(batch_size, num_heads, num_tokens, num_tokens)	No shape change
7. Multiply with Values (V)	(batch_size, num_heads, num_tokens, num_tokens) and (batch_size, num_heads, num_tokens, head_dim)	Attention @ V	(batch_size, num_heads, num_tokens, head_dim)	num_tokens ↔ head_dim
8. Merge Heads	(batch_size, num_heads, num_tokens, head_dim)	Reshape back	(batch_size, num_tokens, d_out)	num_heads × head_dim → d_out
9. Final Linear Layer	(batch_size, num_tokens, d_out)	Linear layer	(batch_size, num_tokens, d_out)	No change

---

Key Features

• Dual Weight Support:
  • Native: Train from scratch and save as .pth.
  • Legacy: Download original OpenAI TensorFlow checkpoints and convert them to PyTorch (./src/download_weights.py).
  • Modern: Load weights directly from Hugging Face and map them to the custom architecture (./src/download_weights_hf.py).
• Gradient Accumulation: Training loop simulates larger batch sizes on limited hardware by accumulating gradients over multiple steps before updating weights.
• Hardware Acceleration: Automatically detects and utilizes CUDA (NVIDIA), MPS (Apple Silicon), or CPU.
• Custom Generation: Implements Top-K sampling and Temperature scaling to control the creativity and coherence of generated text.
• Auto-Splitting: The dataset handler automatically splits a single raw text file (./data/input.txt) into Train/Validation/Test sets.
• Adaptive Data Handling: Uses a sliding window approach with configurable stride to create efficient, overlapping input-target pairs, maximizing data usage from limited text sources.
• Weight Tying: Explains weight tying between the token embedding layer and the final output head (a standard GPT-2 optimization) to reduce parameter count (used in the official implementation). But we don't use it in our implementation (explantion provided in notebook).
• BPE Tokenization: Integrates OpenAI's tiktoken library for fast Byte Pair Encoding (BPE), enabling efficient subword processing that matches the official GPT-2 vocabulary schema.

---

Setup & Installation

This project requires Python 3.10 or higher.

Method 1: Using pip (Standard)

1. Create and activate a virtual environment:

   python -m venv .venv
   # Windows
   .venv\Scripts\activate
   # macOS/Linux
   source .venv/bin/activate

2. Install dependencies:

   pip install -r requirements.txt

Method 2: Using uv (Fast & Recommended)

1. Install uv:

   pip install uv

2. Setup environment and install:

   uv venv --python=python3.10
   source .venv/bin/activate
   uv pip install -r requirements.txt

---

Usage: Training

To train the model on your own text data:

1. Prepare Data:

   • Rename your text file to input.txt.
   • Place it inside the ./data/ folder.

2. Launch the Training Notebook:

   jupyter lab train.ipynb

3. Run the Pipeline:

   • Data Loading: The notebook will detect input.txt. If train.txt and valid.txt don't exist, it will prompt you to auto-split the file.
   • Configuration: You can adjust GPT_CONFIG_124M (layers, heads) and SETTINGS (batch size, learning rate) in the second cell.
   • Training: The loop runs for the specified epochs, printing loss and generating sample text at the end of every epoch.
   • Saving: The model is saved to ./model_weights/trained_gpt2.pth.

---

Usage: Inference

You have two options for generating text.

Option A: Running Your Custom Model

You can use either inference_hf.ipynb or inference.ipynb.

1. Open the notebook.
2. Set the selection variable:

   MODEL_SELECTION = "custom"

3. Ensure your trained_gpt2.pth is located in ./model_weights/.
4. Run the notebook to generate text.

Option B: Running Official GPT-2 Weights

If you want to run the pre-trained 124M, 355M, 774M, or 1558M models, choose one of the methods below based on your needs.

Method 1: Modern & Fast (Recommended) This uses the Hugging Face transformers library for faster downloads and caching.

1. Open inference_hf.ipynb.
2. Select a size:

   MODEL_SELECTION = "gpt2-small (124M)"
   # Or "gpt2-medium (355M)", "gpt2-large (774M)", etc.

3. Run the notebook. It will download the official weights via Hugging Face, map them layer-by-layer into your custom GPTModel class, and let you generate text.

Method 2: Legacy & Manual This downloads the original OpenAI checkpoints directly from their official Azure bucket (openaipublic.blob.core.windows.net) and manually converts them to PyTorch.

1. Open inference.ipynb.
2. Select a size:

   MODEL_SELECTION = "gpt2-small (124M)"
   # Or "gpt2-medium (355M)", "gpt2-large (774M)", etc.

3. Run the notebook. It will fetch the raw files from the OpenAI public server, convert the TensorFlow checkpoints to PyTorch .pth format, and load them into your model.

---

Fine-Tuning

Fine-tuning specializes the pre-trained GPT-2 model for specific tasks. This repository demonstrates two distinct approaches:

1. Classification: Adapting the model to output labels (e.g., Sentiment Analysis).
2. Instruction: Teaching the model to follow prompts and act as an assistant.

Classification

• This directory (classification-finetuning/) is a self-contained module for the Sentiment Analysis task.
• The Goal is to adapt a pre-trained GPT-2 Model — originally designed for text generation — to classify input text into distinct categories (Positive vs. Negative).
• Pre-trained Weights: The fine-tuned model file gpt2-small_classifier.pth is available in the GitHub Releases section. Please check the release notes for instructions on how to download and use it.

Directory Structure

classification-finetuning/
├── src/                    # Classification-specific source code
│   ├── dataset.py          # Handles IMDB tokenization & DataLoader creation
│   └── train_utils.py      # Helper functions for accuracy, loss tracking & plotting
├── data/                 # Auto-downloaded IMDB dataset and processed CSVs
├── model_weights/        # Stores the saved fine-tuned model (.pth)
└── main.ipynb            # The main notebook to Run Training, Evaluation & Inference

Instruction

• This directory (instruction-finetuning/) is a self-contained module designed to teach the model to follow natural language prompts and act as an assistant.
• The goal is to adapt a pre-trained GPT-2 model using the Stanford Alpaca dataset (52,000 instruction, input, and output pairs generated via Self-Instruct).
• Custom Collation: Implements a PyTorch collation strategy to mask out padding tokens (-100) during the forward pass, ensuring the model's loss is calculated purely on the actual response tokens.
• Automated Evaluation: Features an "LLM-as-a-Judge" pipeline. It connects to a local instance of Ollama (e.g., gemma4:e2b) to automatically evaluate and score the fine-tuned model's generated responses against ground-truth answers.
• Pre-trained Weights: The fine-tuned model file gpt2-small_instruction.pth is available in the GitHub Releases section. Please check the release notes for instructions on how to download and use it.

Directory Structure

instruction-finetuning/
├── src/                    # Instruction-specific source code
│   ├── dataset.py          # Handles Alpaca dataset downloading, formatting, and custom collation
│   ├── eval_utils.py       # LLM-as-a-Judge automated scoring via local Ollama
│   └── train_utils.py      # Helper functions for Causal LM loss and text generation
├── data/                   # Auto-downloaded Stanford Alpaca dataset (alpaca_data.json)
├── model_weights/          # Stores the saved fine-tuned model (e.g., gpt2-small_instruction.pth)
└── main.ipynb              # The main notebook to run Training, Automated Evaluation & Inference

---

Interactive Web UI

This repository includes a fully featured interactive web interface powered by Chainlit. It provides a unified chat environment to seamlessly switch between and evaluate the different models trained in this project.

Key Features

• Chat Profiles: Dynamically toggle between the Base Pretrained Model, the Instruction Assistant, and the Sentiment Analyzer using a built-in dropdown menu.
• Auto-Scaling Architecture: The application automatically inspects the loaded .pth file to dynamically adjust the context window (e.g., from 1024 down to 256 tokens) based on your specific training configuration.
• Hardware Auto-Detection: Automatically utilizes CUDA, MPS, or CPU to ensure optimal inference speed based on your machine.

How to Launch

1. Run the application from the root directory:

   chainlit run app.py -w

2. The interface will automatically open in your default web browser at http://localhost:8000.

Using Custom Models

To load your own fine-tuned weights into the UI:

1. Place your successfully trained .pth file into the corresponding model_weights/ directory.
2. Open app.py and update the MODEL_PATHS dictionary to point to your custom filename. The interface will automatically detect the architectural adjustments upon your next session initialization.

---

Contributing

Contributions are welcome! Whether it's optimizing the attention mechanism, adding new schedulers, or improving documentation.

How to Get Started

1. Fork the repository

   git clone https://github.com/yourusername/GPT2-From-Scratch.git
   cd GPT2-From-Scratch

2. Create a feature branch

   git checkout -b feature/flash-attention # An Example branch name

3. Make your changes

   • Ensure all notebooks run sequentially.
   • If modifying ./src/model.py, ensure the shapes align in ./src/download_weights.py to maintain compatibility with official weights.

4. Push and Open a PR

   • Provide a description of your changes.
   • If you improved training speed, please provide metrics.

---

Acknowledgements

• OpenAI for their paper on GPT-2, which defined the architecture used in this project:

  Language Models are Unsupervised Multitask Learners.

• Sebastian Raschka for his repository and book, which served as a primary reference for my "from scratch" implementation:

  Raschka, Sebastian. Build A Large Language Model (From Scratch). Manning, 2024. ISBN: 978-1633437166.
  Source Code: Github

• Stanford (tatsu-lab) for the Alpaca dataset used in the instruction fine-tuning module:

  Taori, Rohan et al. Stanford Alpaca: An Instruction-following LLaMA model. 2023.
  Repository: Github

• Stanford AI Lab for the Large Movie Review Dataset (IMDB) used in the sentiment classification module:

  Maas, Andrew L., et al. Learning Word Vectors for Sentiment Analysis. Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies. 2011.
  Dataset: Stanford IMDB

---

Star the repository if you found this useful! 🌟