git clone https://github.com/workofart/ml-by-hand.git
cd ml-by-hand
./bootstrap.sh
source .venv/bin/activate

RELEASE=https://github.com/workofart/ml-by-hand/releases/download/gpt2-124m-openwebtext-56000
CKPT=gpt2_124m_openwebtext_56000_inference
VOCAB=openwebtext_vocab_49990.pkl
curl -L --create-dirs -o "checkpoints/$CKPT.json" "$RELEASE/$CKPT.json"
curl -L --create-dirs -o "checkpoints/$CKPT.npz" "$RELEASE/$CKPT.npz"
curl -L --create-dirs -o "training_data/$VOCAB" "$RELEASE/$VOCAB"

Then, from the repo root (after the Environment Setup below):

from autograd.text.tokenizer import BytePairEncoder
from autograd.text.utils import generate_text
from autograd.tools.model import load_checkpoint
from examples.gpt_2 import GPT2, GPT2ForwardFn

ckpt = load_checkpoint(
    "checkpoints/gpt2_124m_openwebtext_56000_inference.json",
    "checkpoints/gpt2_124m_openwebtext_56000_inference.npz",
)
model = GPT2(**ckpt["model_init_kwargs"])
model.load_state_dict(ckpt["model_state_dict"])

bpe = BytePairEncoder(
    num_merges=49990,
    vocab_file_path="training_data/openwebtext_vocab_49990.pkl",
)

generate_text(
    model=model,
    prediction_func=GPT2ForwardFn(),
    bpe=bpe,
    start_tokens="The meaning of life is",
    max_length=100,
    temperature=0.8,
    top_k=50,
)

> Inference: 100%|██████████| 95/95 [00:02<00:00, 31.90it/s]
> [prompt 5 tokens + 95 new tokens in 2.98s, 31.9 tok/s]
> 'The meaning of life is based on the ability to feel and remember the things that ...'

• Regression (Code)

• Binary Classification:

  • MNIST (One vs Rest) (Code)
  • Breast Cancer (Code)

• Multi-Class Classification:

  • MNIST (Code)
  • CIFAR-10/CIFAR-100 (Code)

• Convolutional Neural Networks:

  • MNIST (Code)
  • CIFAR-10/CIFAR-100 (Code)

• Residual Neural Networks:

  • MNIST (Code)
  • CIFAR-10/CIFAR-100 (Code)

• RNN + LSTM:

  • Movie Sentiment Analysis (Code)

• Neural Turing Machine (LSTM Controller):

  • Copy Tasks (Code)

• Sequence-to-Sequence:

  • WikiSum (Code)

from autograd.tensor import Tensor
from autograd.nn import Linear, Module
from autograd.optim import SGD
from autograd.backend import xp

class SimpleNN(Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        # A single linear layer (input_dim -> output_dim).
        # Mathematically: fc(x) = xW^T + b
        # where W is weight and b is bias.
        self.fc = Linear(input_dim, output_dim)

    def forward(self, x):
        # Simply compute xW^T + b without any additional activation.
        return self.fc(x)

# Create a sample input tensor x with shape (1, 3).
# 'requires_grad=True' means we want to track gradients for x.
x = Tensor([[-1.0, 0.0, 2.0]], requires_grad=True)

# We want the output to get close to 1.0 over time.
y_true = 1.0

# Initialize the simple neural network.
# This layer has a weight matrix W of shape (3, 1) and a bias of shape (1,).
model = SimpleNN(input_dim=3, output_dim=1)

# Use SGD with a learning rate of 0.03
optimizer = SGD(model.parameters, lr=0.03)

for epoch in range(20):
    # Reset (zero out) all accumulated gradients before each update.
    optimizer.zero_grad()

    # --- Forward pass ---
    # prediction = xW^T + b
    y_pred = model(x)
    print(f"Epoch {epoch}: {y_pred}")

    # Define a simple mean squared error function
    loss = ((y_pred - y_true) ** 2).mean()

    # --- Backward pass ---
    # Ultimately we need to compute the gradient of the loss with respect to the weights
    # Specifically, if Loss = (pred - 1)^2, then:
    #   dL/d(pred) = 2 * (pred - 1)
    #   d(pred)/dW = d(xW^T + b) / dW = x^T
    # By chain rule, dL/dW = dL/d(pred) * d(pred)/dW = [2 * (pred - 1)] * x^T
    loss.backward()

    # --- Update weights ---
    optimizer.step()

# See the computed gradients for the linear layer’s weight matrix:
weights = model.fc.parameters["weight"].data
bias = model.fc.parameters["bias"].data
gradient = model.fc.parameters["weight"].grad
print("[After Training] Gradients for fc weights:", gradient)
print("[After Training] layer weights:", weights)
print("[After Training] layer bias:", bias)
assert xp.to_scalar(xp.allclose(x.data @ weights + bias, y_true))