Build a JavaScript Code Generator

Build a GPT-style transformer from scratch that generates JavaScript code.

Introduction

30 Second Summary

Every time an AI assistant writes code, answers a question, or finishes your sentence, it is predicting the next piece of text one token at a time. The technology powering all of these tools is the same: a transformer model.

In this project, you will build your own miniature GPT model from scratch in PyTorch and train it on real JavaScript code from GitHub. You will give the model a prompt like function add( and watch it generate valid JavaScript completions.

What You'll Build

A ~10 million parameter GPT-style transformer that learns JavaScript syntax and patterns from real GitHub repositories, then generates code completions from any prompt you provide.

By the end of this project, you'll have:

  • A working GPT-style transformer model built entirely from scratch, with self-attention, embeddings, and residual connections you implemented yourself.
  • A trained model checkpoint that generates syntactically valid JavaScript code completions from prompts.
  • Hands-on understanding of how models like GPT-4 and Copilot work at the implementation level, from tokenization to autoregressive generation.
  • Secret Mission: Add temperature and top-k sampling to control how creative or predictable your model's output is.

Do I need machine learning experience?

No. This project assumes basic Python knowledge (variables, functions, loops). You will learn every machine learning concept as you build, starting from first principles.

Before We Start

Tools like GitHub Copilot and ChatGPT generate code by predicting the next token in a sequence. They use a transformer architecture at their core.

In this project, you will build a miniature version of that same architecture from scratch, train it on real JavaScript from GitHub, and watch it generate valid code completions. First, you need to set up your environment with the right tools.

In this step, get ready to:
  • Create a project folder and Python virtual environment.
  • Install PyTorch and the Hugging Face datasets library.
  • Verify your M1 GPU acceleration is working.
Create your project folder and install dependencies

Your project needs PyTorch (the deep learning framework) and the Hugging Face datasets library (for downloading training data). You will install both inside a virtual environment to keep things isolated.

  • Press Cmd+Space (macOS) to open your search bar.
  • Type Terminal and press Enter to open it.
  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir jsgpt
cd jsgpt
  • Create a Python virtual environment by running:
python3 -m venv venv
  • Activate the virtual environment by running:
source venv/bin/activate

You should see (venv) appear at the beginning of your terminal prompt. This confirms you are inside the virtual environment.

What is a virtual environment?

A virtual environment is an isolated Python installation inside your project folder. Packages you install here stay separate from your system Python, so different projects can use different library versions without conflicts.

  • Install PyTorch by running:
pip3 install torch torchvision torchaudio
  • Install the Hugging Face datasets library by running:
pip install datasets
  • Verify both packages installed correctly by running:
python3 -c "import torch; print('PyTorch:', torch.__version__)" && python3 -c "import datasets; print('Datasets:', datasets.__version__)"

You should see version numbers printed for both PyTorch and Datasets. No errors means everything installed successfully.

Seeing an error during install?

Make sure you see (venv) in your terminal prompt. If not, run source venv/bin/activate again before installing.

If pip3 is not found, try pip instead. Inside a virtual environment both commands work the same way.

Help me troubleshoot my PyTorch or datasets installation error.

Verify MPS GPU acceleration

Your MacBook's M1 chip has a built-in GPU that PyTorch can use through Apple's Metal Performance Shaders (MPS) backend. This makes training significantly faster than running on CPU alone.

  • Confirm MPS is available by running this test script:
python3 -c "
import torch
print('MPS available:', torch.backends.mps.is_available())
if torch.backends.mps.is_available():
    x = torch.ones(3, 3, device='mps')
    print('Tensor on MPS device:')
    print(x)
    print('MPS is working!')
else:
    print('MPS is NOT available - training will use CPU')
"

You should see MPS available: True followed by a 3x3 tensor of ones and the message MPS is working! printed.

What is MPS?

MPS stands for Metal Performance Shaders. It is Apple's GPU computing framework that lets PyTorch run tensor operations on your M1 chip's GPU instead of the CPU. This is the same acceleration that powers machine learning features built into macOS.

When you call .to("mps") on a tensor or model, PyTorch moves it from system RAM to GPU memory where math operations run in parallel across GPU cores.

Seeing MPS available: False?

Make sure you are running macOS 14.0 or later. Check by clicking the Apple menu and selecting About This Mac.

Ensure Xcode command-line tools are installed by running xcode-select --install in your terminal.

Help me figure out why MPS is not available on my Mac.

Run a GPU speed benchmark

Now that MPS is confirmed working, run a quick benchmark to see how much faster the GPU is compared to CPU for the kind of matrix math your transformer model will do.

  • Run this benchmark script to compare MPS vs CPU speed:
python3 -c "
import torch
import time

size = 2000
a = torch.randn(size, size)
b = torch.randn(size, size)

# CPU benchmark
start = time.time()
for _ in range(10):
    _ = a @ b
cpu_time = time.time() - start

# MPS benchmark
a_mps = a.to('mps')
b_mps = b.to('mps')
_ = a_mps @ b_mps  # warm-up run
torch.mps.synchronize()

start = time.time()
for _ in range(10):
    _ = a_mps @ b_mps
torch.mps.synchronize()
mps_time = time.time() - start

print(f'CPU time:  {cpu_time:.3f}s')
print(f'MPS time:  {mps_time:.3f}s')
print(f'Speedup:   {cpu_time / mps_time:.1f}x faster on GPU')
"

You should see MPS completing the same computation several times faster than CPU. The exact speedup varies, but expect somewhere between 2x and 10x for this matrix size.

What does this benchmark do?

It multiplies two 2000x2000 matrices ten times on both CPU and MPS, then compares the total time. Matrix multiplication is the core operation inside every transformer layer, so this benchmark directly predicts how much faster your training loop will run on GPU.

The torch.mps.synchronize() call ensures all GPU operations finish before measuring time. Without it, the timer would only measure how long it takes to queue the work, not actually complete it.

MPS is slower than CPU?

For very small matrices, the overhead of copying data to the GPU can make MPS slower. Make sure you are using size = 2000 as shown above.

If you have many apps open, the GPU may be under memory pressure. Try closing other applications and re-running the benchmark.

Help me understand my MPS benchmark results.

Your environment is ready. Next up, you will download real JavaScript code from GitHub and prepare it as training data for your model.

Prepare JavaScript Training Data

Your environment is ready with PyTorch and MPS acceleration confirmed. Now you need data to train your model on.

Every language model learns by studying massive amounts of text. In this step, you will download real JavaScript source code from GitHub and prepare it for training. You will also build a character-level tokenizer that converts raw text into numbers your model can process.

In this step, get ready to:
  • Download approximately 5MB of real JavaScript code from GitHub.
  • Examine the dataset and understand its vocabulary.
  • Build a character-level tokenizer and split the data for training.
Download JavaScript code from GitHub

Your model needs real-world code to learn from. The Hugging Face datasets library gives you access to massive public datasets, including one called codeparrot/github-code that contains over 11 million JavaScript files scraped from GitHub repositories.

You will stream a small portion of this dataset (about 5MB) and save it locally as a training file.

  • In VS Code, create a new file called prepare_data.py in your project folder.
  • Paste the following code into prepare_data.py:
from datasets import load_dataset

print("Streaming JavaScript code from GitHub...")
ds = load_dataset(
    "codeparrot/github-code",
    streaming=True,
    split="train",
    languages=["JavaScript"],
    trust_remote_code=True
)

texts = []
total_chars = 0
target_chars = 5_000_000

for i, example in enumerate(ds):
    code = example["code"]
    texts.append(code)
    total_chars += len(code)
    if total_chars >= target_chars:
        break
    if (i + 1) % 200 == 0:
        print(f"  Collected {i + 1} files ({total_chars:,} characters)...")

print(f"\nDone! Collected {len(texts)} files ({total_chars:,} characters)")

with open("javascript_data.txt", "w", encoding="utf-8") as f:
    f.write("\n".join(texts))

print("Saved to javascript_data.txt")

What does this code do?

  • load_dataset connects to the codeparrot/github-code dataset on Hugging Face. The streaming=True flag means you download files one at a time instead of pulling the entire 1TB dataset.
  • languages=["JavaScript"] filters the stream server-side so you only receive JavaScript files.
  • The loop collects files until you hit 5 million characters, printing progress every 200 files.
  • Finally, all collected code is joined with newlines and saved to javascript_data.txt.

Why trust_remote_code=True?

This dataset uses a custom loading script hosted in its repository. Hugging Face requires you to explicitly opt in to running that script by passing trust_remote_code=True. The codeparrot/github-code dataset is a well-known, widely-used dataset maintained by Hugging Face staff.

  • Save prepare_data.py.
  • Run the script by entering this command in your terminal:
python3 prepare_data.py

The script will print progress updates as it collects files. It typically takes 1-2 minutes depending on your internet connection.

  • Confirm javascript_data.txt was created by running:
ls -lh javascript_data.txt

You should see the file listed at approximately 5MB in size.

Script stuck or showing an error?

  • Make sure your virtual environment is activated. You should see the environment name in your terminal prompt.
  • Check your internet connection. The script streams data from Hugging Face servers.
  • If you see a message about trust_remote_code, make sure you included trust_remote_code=True in your load_dataset call.

Having trouble with the download? Help me debug why my prepare_data.py script is failing

✔️ Awesome, I've got everything!

Great. Make sure you have saved prepare_data.py and that javascript_data.txt exists in your project folder.

ⓧ I'd like to double check the full code

Here is the complete prepare_data.py file:

from datasets import load_dataset

print("Streaming JavaScript code from GitHub...")
ds = load_dataset(
    "codeparrot/github-code",
    streaming=True,
    split="train",
    languages=["JavaScript"],
    trust_remote_code=True
)

texts = []
total_chars = 0
target_chars = 5_000_000

for i, example in enumerate(ds):
    code = example["code"]
    texts.append(code)
    total_chars += len(code)
    if total_chars >= target_chars:
        break
    if (i + 1) % 200 == 0:
        print(f"  Collected {i + 1} files ({total_chars:,} characters)...")

print(f"\nDone! Collected {len(texts)} files ({total_chars:,} characters)")

with open("javascript_data.txt", "w", encoding="utf-8") as f:
    f.write("\n".join(texts))

print("Saved to javascript_data.txt")
Examine the training data

Before feeding data to a model, you need to understand what it looks like. Knowing the character set (vocabulary) tells you how many unique symbols the model needs to learn.

  • Start an interactive Python session by running:
python3
  • Load the data and print basic statistics by entering these lines:
# Load the JavaScript data you just downloaded
with open("javascript_data.txt", "r", encoding="utf-8") as f:
    text = f.read()

print(f"Total characters: {len(text):,}")
print(f"\nFirst 200 characters:")
print(text[:200])

You should see a character count around 5 million and a snippet of real JavaScript code.

  • Print the unique character set (the vocabulary) by entering:
# Find all unique characters in the dataset
chars = sorted(list(set(text)))
print(f"Vocabulary size: {len(chars)} unique characters")
print(f"Characters: {''.join(chars)}")

Why character-level tokenization?

You will likely see around 90-100 unique characters. This is your model's vocabulary. Every bracket, semicolon, space, and newline is a separate token.

Character-level tokenization works especially well for code because individual characters carry syntactic meaning. A missing } or an extra ; completely changes what code does. Your model needs to learn every character matters.

Build the character-level tokenizer

A tokenizer converts text into numbers that a neural network can process. You need two functions: encode (text to integers) and decode (integers back to text).

  • Still in your Python session, create the tokenizer mappings and functions:
# Create character-to-index and index-to-character mappings
char_to_idx = {ch: i for i, ch in enumerate(chars)}
idx_to_char = {i: ch for i, ch in enumerate(chars)}

# Encode: convert a string to a list of integers
def encode(s):
    return [char_to_idx[c] for c in s]

# Decode: convert a list of integers back to a string
def decode(indices):
    return "".join([idx_to_char[i] for i in indices])

# Test it
test = "function add("
encoded = encode(test)
print(f"Text:    {test}")
print(f"Encoded: {encoded}")
print(f"Decoded: {decode(encoded)}")

What does this code do?

  • char_to_idx maps each character to a unique integer. For example, a might map to 65 and { might map to 93.
  • idx_to_char is the reverse mapping, turning integers back into characters.
  • encode looks up each character's integer. decode looks up each integer's character. Together they convert losslessly between text and numbers.

You should see the test string converted to a list of integers and then perfectly reconstructed.

  • Now encode the full dataset into a PyTorch tensor and split it for training:
import torch

# Encode the entire dataset into a tensor of integers
data = torch.tensor(encode(text), dtype=torch.long)
print(f"Dataset tensor shape: {data.shape}")
print(f"Data type: {data.dtype}")

# Split 90% training, 10% validation
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]
print(f"\nTraining tokens: {len(train_data):,}")
print(f"Validation tokens: {len(val_data):,}")

Why split into training and validation?

The training set is what the model learns from. The validation set is held back to check whether the model is actually learning general patterns or just memorizing the training data.

A 90/10 split gives the model plenty of data to learn from while keeping enough aside to measure real performance.

  • Exit the Python session by typing:
exit()

You now have a working tokenizer and understand your data's structure. This code will become part of your model file in the next step.

Seeing a KeyError when encoding?

  • Make sure you built chars from the same text variable you are encoding. If you reloaded the file, rebuild chars and the mappings first.
  • If encoding the full dataset is slow (it processes millions of characters), give it 10-20 seconds to complete.

Need help with the tokenizer? Help me understand why my encode or decode function is not working

You have real JavaScript training data ready and a tokenizer that can convert it into numbers. Next up, you will build the transformer model that learns patterns from this data.

Build the Transformer Model

Your training data is ready and your tokenizer can convert JavaScript into numbers. Now you need something that can actually learn from those numbers.

In this step, you will build a transformer model from scratch. This is the same architecture behind GPT-4, Copilot, and every modern LLM. You will implement each component yourself: the attention mechanism, feed-forward networks, residual connections, and the full model assembly.

In this step, get ready to:
  • Define the model configuration and implement causal self-attention.
  • Build the feed-forward network, transformer blocks, and the full GPT model class.
  • Instantiate the model on GPU and verify it has approximately 10 million parameters.
Define the GPT configuration and implement CausalSelfAttention

The transformer architecture has several hyperparameters that control its size and behavior. You will define these first, then build the core attention mechanism that lets the model look at previous tokens when predicting the next one.

  • Create a new file called jsgpt.py in your project folder.
  • Add the imports and model configuration by copying this code:
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass

# ============================================================
# Configuration
# ============================================================

@dataclass
class GPTConfig:
    block_size: int = 256
    vocab_size: int = 0
    n_layer: int = 6
    n_head: int = 6
    n_embd: int = 384
    dropout: float = 0.2

BATCH_SIZE = 32
MAX_ITERS = 3000
EVAL_INTERVAL = 300
EVAL_ITERS = 100
LEARNING_RATE = 3e-4
DEVICE = "mps" if torch.backends.mps.is_available() else "cpu"

print(f"Using device: {DEVICE}")

What does this code do?

GPTConfig is a dataclass that holds all the model's architectural choices. block_size (256) is how many characters the model can see at once. n_layer (6) is how many transformer blocks are stacked. n_head (6) is how many attention heads run in parallel. n_embd (384) is the size of each token's internal representation.

The training constants below control how long and how fast the model learns. DEVICE automatically selects your M1 GPU (MPS) for faster training.

  • Save jsgpt.py and run it to verify the device detection:
python jsgpt.py
  • You should see Using device: mps printed in your terminal.

Seeing "cpu" instead of "mps"?

Make sure your virtual environment is activated (you should see (venv) in your terminal prompt). If MPS still shows as unavailable, confirm you ran the MPS verification script successfully in Step 1.

help me debug why MPS is not available on my Mac

  • Next, add the data preparation section below the print line. This loads your JavaScript data and builds the character-level tokenizer:
# ============================================================
# Data Preparation
# ============================================================

print("Loading JavaScript training data...")
with open("javascript_data.txt", "r", encoding="utf-8") as f:
    text = f.read()

print(f"Dataset size: {len(text):,} characters")

chars = sorted(list(set(text)))
vocab_size = len(chars)
print(f"Vocabulary size: {vocab_size} unique characters")

char_to_idx = {ch: i for i, ch in enumerate(chars)}
idx_to_char = {i: ch for i, ch in enumerate(chars)}


def encode(s):
    return [char_to_idx[c] for c in s]


def decode(indices):
    return "".join([idx_to_char[i] for i in indices])


data = torch.tensor(encode(text), dtype=torch.long)

n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]
print(f"Train: {len(train_data):,} tokens | Val: {len(val_data):,} tokens")

What does this code do?

This section loads your javascript_data.txt file and builds a character-level tokenizer. It maps every unique character in the dataset to an integer (char_to_idx) and back (idx_to_char).

The encode function converts a string to a list of integers. The decode function reverses it. The full dataset gets encoded into a single PyTorch tensor, then split 90/10 into training and validation sets.

  • Save jsgpt.py and run it again:
python jsgpt.py
  • You should see the device line followed by dataset statistics: character count, vocabulary size (around 96 unique characters), and the train/val split sizes.

Seeing FileNotFoundError?

This means javascript_data.txt is not in the same folder as jsgpt.py. Make sure both files are in your project directory and that you ran prepare_data.py successfully in the previous step.

help me fix a FileNotFoundError for javascript_data.txt

  • Now add the data loading utilities below the train/val split. These functions will supply batches of data during training:
# ============================================================
# Data Loading
# ============================================================

config = GPTConfig(vocab_size=vocab_size)


def get_batch(split):
    data_split = train_data if split == "train" else val_data
    ix = torch.randint(len(data_split) - config.block_size, (BATCH_SIZE,))
    x = torch.stack([data_split[i : i + config.block_size] for i in ix])
    y = torch.stack([data_split[i + 1 : i + config.block_size + 1] for i in ix])
    x, y = x.to(DEVICE), y.to(DEVICE)
    return x, y

What does this code do?

This creates the GPTConfig instance with the actual vocabulary size from your data. The get_batch function randomly samples sequences of length block_size (256 characters) from the data.

For each input sequence x, the target y is the same sequence shifted one position forward. This teaches the model to predict the next character given the previous ones.

You will not see new output from this chunk yet. These functions are called during training in the next step.

  • Now for the most important piece. Add the CausalSelfAttention class below the get_batch function. This is the core mechanism that lets the model attend to previous tokens:
# ============================================================
# Model Architecture
# ============================================================


class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.n_head = config.n_head
        self.head_dim = config.n_embd // config.n_head
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.register_buffer(
            "mask",
            torch.tril(torch.ones(config.block_size, config.block_size)).view(
                1, 1, config.block_size, config.block_size
            ),
        )

    def forward(self, x):
        B, T, C = x.shape
        qkv = self.c_attn(x)
        q, k, v = qkv.split(C, dim=2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        att = (q @ k.transpose(-2, -1)) * (self.head_dim**-0.5)
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)
        y = att @ v

        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.resid_dropout(self.c_proj(y))
        return y

What is self-attention?

Self-attention is how the model decides which earlier tokens matter when predicting the next one. For each position, it computes three vectors: a query ("what am I looking for?"), a key ("what do I contain?"), and a value ("what information do I carry?").

The "causal" part means the model can only look backward. The mask is a lower-triangular matrix that blocks the model from peeking at future tokens. This is what makes it autoregressive: it generates one token at a time, only using past context.

Multiple attention heads (n_head=6) run in parallel, each learning to focus on different patterns. One head might track bracket matching while another tracks variable names.

You will not see new output from this class definition yet. It is used when the full model is assembled in the next substep.

Implement FeedForward, TransformerBlock, and the full GPT model

Attention alone is not enough. Each transformer block also needs a feed-forward network that processes the attended information. Then you will stack multiple blocks together and wrap them in the full GPT model class.

  • Add the FeedForward class below CausalSelfAttention:
class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(config.n_embd, 4 * config.n_embd),
            nn.GELU(),
            nn.Linear(4 * config.n_embd, config.n_embd),
            nn.Dropout(config.dropout),
        )

    def forward(self, x):
        return self.net(x)

What does FeedForward do?

After attention decides which tokens to focus on, the feed-forward network processes that information. It expands the representation to 4x its size (384 to 1536), applies a GELU activation, then projects back down. This expansion gives the model room to learn complex transformations.

  • Add the TransformerBlock class below FeedForward:
class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln2 = nn.LayerNorm(config.n_embd)
        self.ff = FeedForward(config)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ff(self.ln2(x))
        return x

What are residual connections?

Notice the x = x + ... pattern. These are residual connections. Instead of replacing x entirely, each sub-layer adds its output to the original input. This makes deep networks trainable by giving gradients a "shortcut" path backward through the network.

Layer normalization (ln1, ln2) stabilizes training by normalizing the values before each sub-layer. The full pattern is: normalize, attend, add back, normalize, feed-forward, add back.

  • Now add the full GPT model class below TransformerBlock. This assembles all the components into a complete language model:
class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
        self.position_embedding = nn.Embedding(config.block_size, config.n_embd)
        self.drop = nn.Dropout(config.dropout)
        self.blocks = nn.Sequential(
            *[TransformerBlock(config) for _ in range(config.n_layer)]
        )
        self.ln_f = nn.LayerNorm(config.n_embd)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.token_embedding.weight = self.lm_head.weight
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok_emb = self.token_embedding(idx)
        pos_emb = self.position_embedding(torch.arange(T, device=idx.device))
        x = self.drop(tok_emb + pos_emb)
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)), targets.view(-1)
            )
        return logits, loss

    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.config.block_size :]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

What does the GPT class do?

This is the full model. It has two embedding layers: token_embedding converts each character ID into a 384-dimensional vector, and position_embedding adds positional information so the model knows where each token sits in the sequence.

The forward method runs the full prediction pipeline: embed tokens, add positions, pass through 6 transformer blocks, normalize, and project back to vocabulary-sized logits. If targets are provided, it also computes the cross-entropy loss for training.

The generate method is for inference. It feeds the model its own output one token at a time (autoregressive generation). Weight tying between token_embedding and lm_head means both layers share parameters. This is standard practice since GPT-2.

You will not see new output from these class definitions yet. The next substep instantiates the model and confirms everything connects.

Instantiate the model and verify the parameter count

All the architecture is defined. Now you will create an actual model instance, move it to your GPU, and confirm it has the expected number of parameters.

  • Add the model instantiation code at the very bottom of jsgpt.py, below the GPT class:
model = GPT(config).to(DEVICE)
param_count = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Model parameters: {param_count:,}")

What does this code do?

GPT(config) creates the model with your configuration. .to(DEVICE) moves all the model's weights to your M1 GPU for faster computation.

The parameter count sums up every trainable number in the model. With 6 layers, 6 heads, and 384 embedding dimensions, you should see approximately 10 million parameters.

  • Save jsgpt.py and run the full file:
python jsgpt.py
  • You should see output that includes Using device: mps, your dataset statistics, and a line showing Model parameters: followed by a number around 10 million.

Seeing an error when instantiating the model?

If you see an AssertionError, check that n_embd (384) is evenly divisible by n_head (6). If you see a RuntimeError about tensor sizes, verify your GPTConfig values match the code above exactly.

If the parameter count is wildly different from 10 million, double check that n_layer is 6 and n_embd is 384 in your GPTConfig.

help me debug my GPT model instantiation error

✔️ Awesome, I've got everything!

Great. Double check you have saved jsgpt.py and the script runs without errors.

ⓧ I'd like to double check the full code

Here is the complete jsgpt.py file as it should look at the end of this step:

import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass

# ============================================================
# Configuration
# ============================================================

@dataclass
class GPTConfig:
    block_size: int = 256
    vocab_size: int = 0
    n_layer: int = 6
    n_head: int = 6
    n_embd: int = 384
    dropout: float = 0.2

BATCH_SIZE = 32
MAX_ITERS = 3000
EVAL_INTERVAL = 300
EVAL_ITERS = 100
LEARNING_RATE = 3e-4
DEVICE = "mps" if torch.backends.mps.is_available() else "cpu"

print(f"Using device: {DEVICE}")

# ============================================================
# Data Preparation
# ============================================================

print("Loading JavaScript training data...")
with open("javascript_data.txt", "r", encoding="utf-8") as f:
    text = f.read()

print(f"Dataset size: {len(text):,} characters")

chars = sorted(list(set(text)))
vocab_size = len(chars)
print(f"Vocabulary size: {vocab_size} unique characters")

char_to_idx = {ch: i for i, ch in enumerate(chars)}
idx_to_char = {i: ch for i, ch in enumerate(chars)}


def encode(s):
    return [char_to_idx[c] for c in s]


def decode(indices):
    return "".join([idx_to_char[i] for i in indices])


data = torch.tensor(encode(text), dtype=torch.long)

n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]
print(f"Train: {len(train_data):,} tokens | Val: {len(val_data):,} tokens")

# ============================================================
# Data Loading
# ============================================================

config = GPTConfig(vocab_size=vocab_size)


def get_batch(split):
    data_split = train_data if split == "train" else val_data
    ix = torch.randint(len(data_split) - config.block_size, (BATCH_SIZE,))
    x = torch.stack([data_split[i : i + config.block_size] for i in ix])
    y = torch.stack([data_split[i + 1 : i + config.block_size + 1] for i in ix])
    x, y = x.to(DEVICE), y.to(DEVICE)
    return x, y


# ============================================================
# Model Architecture
# ============================================================


class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.n_head = config.n_head
        self.head_dim = config.n_embd // config.n_head
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.register_buffer(
            "mask",
            torch.tril(torch.ones(config.block_size, config.block_size)).view(
                1, 1, config.block_size, config.block_size
            ),
        )

    def forward(self, x):
        B, T, C = x.shape
        qkv = self.c_attn(x)
        q, k, v = qkv.split(C, dim=2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        att = (q @ k.transpose(-2, -1)) * (self.head_dim**-0.5)
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)
        y = att @ v

        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.resid_dropout(self.c_proj(y))
        return y


class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(config.n_embd, 4 * config.n_embd),
            nn.GELU(),
            nn.Linear(4 * config.n_embd, config.n_embd),
            nn.Dropout(config.dropout),
        )

    def forward(self, x):
        return self.net(x)


class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln2 = nn.LayerNorm(config.n_embd)
        self.ff = FeedForward(config)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ff(self.ln2(x))
        return x


class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
        self.position_embedding = nn.Embedding(config.block_size, config.n_embd)
        self.drop = nn.Dropout(config.dropout)
        self.blocks = nn.Sequential(
            *[TransformerBlock(config) for _ in range(config.n_layer)]
        )
        self.ln_f = nn.LayerNorm(config.n_embd)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.token_embedding.weight = self.lm_head.weight
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok_emb = self.token_embedding(idx)
        pos_emb = self.position_embedding(torch.arange(T, device=idx.device))
        x = self.drop(tok_emb + pos_emb)
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)), targets.view(-1)
            )
        return logits, loss

    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.config.block_size :]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx


model = GPT(config).to(DEVICE)
param_count = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Model parameters: {param_count:,}")

Your transformer is built and verified. Next up, you will write the training loop and watch this model learn JavaScript patterns from real code.

Train the Model

You have a 10-million parameter transformer model sitting on your GPU, ready to learn. Right now it predicts characters randomly. Training is the process that turns those random guesses into meaningful JavaScript.

In this step, you will write the training loop that feeds thousands of code batches through the model, computes how wrong each prediction is, and nudges the weights in a better direction. After 3000 iterations, your model will have absorbed the patterns of real JavaScript.

In this step, get ready to:
  • Write an evaluation function that measures loss on training and validation data.
  • Build the full training loop with gradient clipping and periodic loss reporting.
  • Train the model and save the checkpoint to disk.
Write the evaluation function

Before you start training, you need a way to measure how well the model is doing. The estimate_loss function runs the model on random batches from both splits and averages the loss, giving you a stable number to track progress.

  • In jsgpt.py, find this line at the bottom of the file:
print(f"Model parameters: {param_count:,}")
  • Add the following code below that line:
optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)


@torch.no_grad()
def estimate_loss():
    model.eval()
    losses = {}
    for split in ["train", "val"]:
        batch_losses = torch.zeros(EVAL_ITERS)
        for k in range(EVAL_ITERS):
            X, Y = get_batch(split)
            _, loss = model(X, Y)
            batch_losses[k] = loss.item()
        losses[split] = batch_losses.mean()
    model.train()
    return losses

What does this code do?

The first line creates an AdamW optimizer. This is the algorithm that updates the model's weights based on computed gradients. AdamW adapts the learning rate for each parameter individually, which works well for transformers.

The estimate_loss function switches the model into evaluation mode (disabling dropout), runs 100 random batches through each split, and averages the loss. The @torch.no_grad() decorator tells PyTorch not to track gradients during evaluation, which saves memory.

Why evaluate on two splits?

The training loss tells you how well the model memorizes its training data. The validation loss tells you how well it generalizes to code it has never seen. If train loss keeps dropping but val loss rises, the model is overfitting.

  • Save jsgpt.py.

Seeing a syntax error?

Make sure the @torch.no_grad() decorator is on the line directly above def estimate_loss(): with no blank lines between them.

Also check that the optimizer line appears after the print(f"Model parameters: {param_count:,}") line, not inside any class or function.

Help me debug a syntax error in my estimate_loss function.

Write the training loop

The training loop is where the model actually learns. Each iteration, it grabs a batch of sequences, predicts the next character for each position, measures how wrong it was, and adjusts its weights to be less wrong next time.

  • In jsgpt.py, add the following code below the estimate_loss function:
print("\nStarting training...")
for iter_num in range(MAX_ITERS):
    if iter_num % EVAL_INTERVAL == 0:
        losses = estimate_loss()
        print(
            f"Step {iter_num:5d} | train loss: {losses['train']:.4f} | val loss: {losses['val']:.4f}"
        )

    xb, yb = get_batch("train")
    _, loss = model(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()

losses = estimate_loss()
print(
    f"Final    | train loss: {losses['train']:.4f} | val loss: {losses['val']:.4f}"
)

torch.save(
    {"model_state_dict": model.state_dict(), "config": config, "chars": chars},
    "jsgpt_model.pt",
)
print("Model saved to jsgpt_model.pt")

What does this code do?

The loop runs for 3000 steps. Every 300 steps, it pauses to evaluate train and validation loss so you can watch the model improve.

Inside each step, get_batch("train") grabs 32 random sequences of 256 characters. The model predicts the next character at every position and computes cross-entropy loss against the true characters.

Then loss.backward() computes gradients, clip_grad_norm_ caps their magnitude at 1.0 to prevent training instability, and optimizer.step() updates the weights.

After training finishes, torch.save writes the model weights, config, and character list to jsgpt_model.pt so you can reload the model later without retraining.

Why clip gradients?

During training, gradients can occasionally spike to very large values. Without clipping, a single bad batch could blow up your model's weights and ruin hours of training. Gradient clipping is standard practice for transformer models.

  • Save jsgpt.py.

Seeing an indentation error in the loop?

Make sure the for iter_num in range(MAX_ITERS): line is at the module level (no indentation). Everything inside the loop should be indented by 4 spaces.

Also check that the losses = estimate_loss() and torch.save(...) lines after the loop are back at the module level, not inside the for loop.

Help me fix my training loop indentation.

Run training and observe results

Your code is ready. Time to let the model learn JavaScript. Training takes approximately 3-5 minutes on an M1 GPU.

  • Run the training script in your terminal:
python jsgpt.py

You should see output like this as training progresses:

Step     0 | train loss: 4.5XXX | val loss: 4.5XXX
Step   300 | train loss: 2.8XXX | val loss: 2.9XXX
Step   600 | train loss: 2.2XXX | val loss: 2.3XXX
...
Step  2700 | train loss: 1.6XXX | val loss: 1.8XXX
Final    | train loss: 1.5XXX | val loss: 1.7XXX
Model saved to jsgpt_model.pt

What do these numbers mean?

Loss starts around 4.5 because log(vocab_size) for a ~96 character vocabulary is about 4.6. This is the loss you get from random guessing across all possible characters.

As loss drops to 1.5-2.0, the model has learned common JavaScript patterns: indentation, bracket matching, keyword usage, and variable naming conventions. The exact numbers will vary slightly each run.

  • Confirm the model checkpoint was saved by running:
ls jsgpt_model.pt

You should see jsgpt_model.pt listed in the output.

Training is extremely slow?

Check that the output says Using device: mps at the top. If it says cpu, the MPS device is not being detected. Make sure you verified MPS in Step 1.

If you have many apps open and see memory warnings, close other applications. The model uses approximately 200-300MB of memory for weights and optimizer states.

Help me figure out why training is running on CPU instead of MPS.

✔️ Awesome, I've got everything!

Great. Double check you have saved jsgpt.py and that jsgpt_model.pt exists in your project folder.

ⓧ I'd like to double check the full code

Here is the complete jsgpt.py file up to this point:

import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass

# ============================================================
# Configuration
# ============================================================

@dataclass
class GPTConfig:
    block_size: int = 256
    vocab_size: int = 0
    n_layer: int = 6
    n_head: int = 6
    n_embd: int = 384
    dropout: float = 0.2

BATCH_SIZE = 32
MAX_ITERS = 3000
EVAL_INTERVAL = 300
EVAL_ITERS = 100
LEARNING_RATE = 3e-4
DEVICE = "mps" if torch.backends.mps.is_available() else "cpu"

print(f"Using device: {DEVICE}")

# ============================================================
# Data Preparation
# ============================================================

print("Loading JavaScript training data...")
with open("javascript_data.txt", "r", encoding="utf-8") as f:
    text = f.read()

print(f"Dataset size: {len(text):,} characters")

chars = sorted(list(set(text)))
vocab_size = len(chars)
print(f"Vocabulary size: {vocab_size} unique characters")

char_to_idx = {ch: i for i, ch in enumerate(chars)}
idx_to_char = {i: ch for i, ch in enumerate(chars)}


def encode(s):
    return [char_to_idx[c] for c in s]


def decode(indices):
    return "".join([idx_to_char[i] for i in indices])


data = torch.tensor(encode(text), dtype=torch.long)

n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]
print(f"Train: {len(train_data):,} tokens | Val: {len(val_data):,} tokens")

# ============================================================
# Data Loading
# ============================================================

config = GPTConfig(vocab_size=vocab_size)


def get_batch(split):
    data_split = train_data if split == "train" else val_data
    ix = torch.randint(len(data_split) - config.block_size, (BATCH_SIZE,))
    x = torch.stack([data_split[i : i + config.block_size] for i in ix])
    y = torch.stack([data_split[i + 1 : i + config.block_size + 1] for i in ix])
    x, y = x.to(DEVICE), y.to(DEVICE)
    return x, y


# ============================================================
# Model Architecture
# ============================================================


class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.n_head = config.n_head
        self.head_dim = config.n_embd // config.n_head
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.register_buffer(
            "mask",
            torch.tril(torch.ones(config.block_size, config.block_size)).view(
                1, 1, config.block_size, config.block_size
            ),
        )

    def forward(self, x):
        B, T, C = x.shape
        qkv = self.c_attn(x)
        q, k, v = qkv.split(C, dim=2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        att = (q @ k.transpose(-2, -1)) * (self.head_dim**-0.5)
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)
        y = att @ v

        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.resid_dropout(self.c_proj(y))
        return y


class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(config.n_embd, 4 * config.n_embd),
            nn.GELU(),
            nn.Linear(4 * config.n_embd, config.n_embd),
            nn.Dropout(config.dropout),
        )

    def forward(self, x):
        return self.net(x)


class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln2 = nn.LayerNorm(config.n_embd)
        self.ff = FeedForward(config)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ff(self.ln2(x))
        return x


class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
        self.position_embedding = nn.Embedding(config.block_size, config.n_embd)
        self.drop = nn.Dropout(config.dropout)
        self.blocks = nn.Sequential(
            *[TransformerBlock(config) for _ in range(config.n_layer)]
        )
        self.ln_f = nn.LayerNorm(config.n_embd)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.token_embedding.weight = self.lm_head.weight
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok_emb = self.token_embedding(idx)
        pos_emb = self.position_embedding(torch.arange(T, device=idx.device))
        x = self.drop(tok_emb + pos_emb)
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)), targets.view(-1)
            )
        return logits, loss

    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.config.block_size :]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx


# ============================================================
# Training
# ============================================================

model = GPT(config).to(DEVICE)
param_count = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Model parameters: {param_count:,}")

optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)


@torch.no_grad()
def estimate_loss():
    model.eval()
    losses = {}
    for split in ["train", "val"]:
        batch_losses = torch.zeros(EVAL_ITERS)
        for k in range(EVAL_ITERS):
            X, Y = get_batch(split)
            _, loss = model(X, Y)
            batch_losses[k] = loss.item()
        losses[split] = batch_losses.mean()
    model.train()
    return losses


print("\nStarting training...")
for iter_num in range(MAX_ITERS):
    if iter_num % EVAL_INTERVAL == 0:
        losses = estimate_loss()
        print(
            f"Step {iter_num:5d} | train loss: {losses['train']:.4f} | val loss: {losses['val']:.4f}"
        )

    xb, yb = get_batch("train")
    _, loss = model(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()

losses = estimate_loss()
print(
    f"Final    | train loss: {losses['train']:.4f} | val loss: {losses['val']:.4f}"
)

torch.save(
    {"model_state_dict": model.state_dict(), "config": config, "chars": chars},
    "jsgpt_model.pt",
)
print("Model saved to jsgpt_model.pt")

Your model has learned JavaScript. Next up, you will put it to work by feeding it prompts and watching it generate real code completions.

Generate JavaScript Code

Your model has finished training and learned real patterns from JavaScript code. Now it's time for the payoff.

In this step, you'll use your trained model to generate brand-new JavaScript. You'll give it a few starting characters and watch it predict one token at a time, building up syntactically plausible code from scratch.

In this step, get ready to:
  • Write the generation pipeline that feeds prompts to your trained model.
  • Test generation with JavaScript prompts and observe the output.
  • Experiment with different prompt lengths to see how context affects quality.
Write the generation pipeline

Your GPT class already has a generate method that handles autoregressive generation. It takes a starting sequence, predicts the next token, appends it, and repeats.

Now you need the code that actually uses it. This generation section will encode text prompts, feed them to the model, and decode the output back into readable JavaScript.

  • In jsgpt.py, scroll to the very bottom of the file (after the line print("Model saved to jsgpt_model.pt")).
  • Add the generation section by pasting the following code below it:
# ============================================================
# Generation
# ============================================================

print("\n" + "=" * 50)
print("GENERATING JAVASCRIPT")
print("=" * 50 + "\n")

# Switch from training mode to evaluation mode
model.eval()
prompts = ["function add(", "const app = ", "if (x > ", "class User {"]

for prompt in prompts:
    # Encode the prompt string into token indices
    context = torch.tensor([encode(prompt)], dtype=torch.long, device=DEVICE)
    # Generate 200 new tokens autoregressively
    generated = model.generate(context, max_new_tokens=200)
    # Decode the full sequence (prompt + generated) back to text
    output = decode(generated[0].tolist())
    print(f"--- Prompt: {prompt!r} ---")
    print(output[:300])
    print()

What does this code do?

model.eval() switches off dropout so the model uses all its learned knowledge during generation. During training, dropout randomly silences neurons to prevent overfitting. During generation, you want the full model.

For each prompt, the pipeline encodes the string into integer tokens, passes them to model.generate, and decodes the result back to text. The generate method works by feeding the model's own output back as input, one token at a time. This is exactly how ChatGPT and GitHub Copilot produce text.

How does autoregressive generation work?

The model only ever predicts one token at a time. Given the sequence so far, it outputs a probability distribution over the entire vocabulary (every possible next character). Then torch.multinomial randomly samples from that distribution to pick the next character.

That chosen character gets appended to the sequence, and the whole thing repeats. Each new prediction sees all previous tokens (up to the 256-character context window), so the model can maintain coherent structure across the output.

  • Save jsgpt.py.

You won't see output yet. The next substep runs the full script and shows the generated JavaScript.

Run generation and observe the output

Your script is now complete. When you run it, it will train the model from scratch (this takes a few minutes), then immediately generate JavaScript from each of the four prompts.

  • Run the full script by entering this command in your terminal:
python jsgpt.py

What should I see?

First you'll see the training output (loss values dropping from ~4.5 to ~1.5-2.0 over 3000 steps). This takes 2-5 minutes on MPS. After training completes and the model is saved, you'll see the generation header followed by four generated code blocks.

Each block starts with the prompt you gave and continues with 200 characters of model-generated JavaScript. Look for proper indentation, matching brackets, semicolons at line endings, and recognizable JavaScript keywords.

  • Wait for the script to finish. After the training loss output and the "Model saved" message, you should see the generation section print four code blocks under the GENERATING JAVASCRIPT banner.

Why does the output look imperfect?

Your model has ~10 million parameters and trained on ~5MB of code. GPT-4 has hundreds of billions of parameters and trained on trillions of tokens. Your model learns the structure of JavaScript (brackets, indentation, keywords) but won't produce logically correct programs.

The impressive part is that it learned these patterns at all from raw characters. It discovered that parentheses close, that semicolons end statements, and that indentation increases after opening braces. All from seeing character sequences.

Not seeing generated output?

  • If the script stops after printing "Model saved to jsgpt_model.pt" with no generation output, check that you pasted the generation section AFTER the save line, not inside the training loop.
  • If you see a KeyError, one of your prompts contains a character that wasn't in the training data. Stick with basic ASCII characters for prompts.
  • If the script seems frozen during training, it's working. Training 3000 steps takes 2-5 minutes on MPS.

Help me debug my generation output.

Experiment with different prompt lengths

The amount of context you give the model affects how good its completions are. A longer prompt gives the model more information to work with. A single character forces it to guess from almost nothing.

  • In jsgpt.py, find the prompts list in the Generation section:
prompts = ["function add(", "const app = ", "if (x > ", "class User {"]
  • Try replacing it with prompts of varying lengths to compare output quality:
prompts = ["f", "function ", "function calculateTotal(items) {\n  let total = "]
  • Save jsgpt.py and run the script again:
python jsgpt.py
  • Compare the three outputs. The single-character prompt f produces scattered, incoherent code. The medium prompt produces recognizable structure. The long prompt produces the most coherent continuation.

Why does more context help?

The self-attention mechanism attends to all previous tokens when predicting the next one. With a single character, the model has almost no signal. With a full function signature, it has strong evidence about what kind of code should follow (variable declarations, return statements, closing braces).

This is the same reason GitHub Copilot gives better suggestions when your file already has context. More input means more attention signal.

  • Once you've finished experimenting, change the prompts list back to the original four prompts:
prompts = ["function add(", "const app = ", "if (x > ", "class User {"]
  • Save jsgpt.py.

✔️ Awesome, I've got everything!

Great. Double check you've saved your file with the original four prompts restored.

ⓧ I'd like to double check the full code

Here is the complete jsgpt.py file including the generation section at the bottom:

import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass

# ============================================================
# Configuration
# ============================================================

@dataclass
class GPTConfig:
    block_size: int = 256
    vocab_size: int = 0
    n_layer: int = 6
    n_head: int = 6
    n_embd: int = 384
    dropout: float = 0.2

BATCH_SIZE = 32
MAX_ITERS = 3000
EVAL_INTERVAL = 300
EVAL_ITERS = 100
LEARNING_RATE = 3e-4
DEVICE = "mps" if torch.backends.mps.is_available() else "cpu"

print(f"Using device: {DEVICE}")

# ============================================================
# Data Preparation
# ============================================================

print("Loading JavaScript training data...")
with open("javascript_data.txt", "r", encoding="utf-8") as f:
    text = f.read()

print(f"Dataset size: {len(text):,} characters")

chars = sorted(list(set(text)))
vocab_size = len(chars)
print(f"Vocabulary size: {vocab_size} unique characters")

char_to_idx = {ch: i for i, ch in enumerate(chars)}
idx_to_char = {i: ch for i, ch in enumerate(chars)}


def encode(s):
    return [char_to_idx[c] for c in s]


def decode(indices):
    return "".join([idx_to_char[i] for i in indices])


data = torch.tensor(encode(text), dtype=torch.long)

n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]
print(f"Train: {len(train_data):,} tokens | Val: {len(val_data):,} tokens")

# ============================================================
# Data Loading
# ============================================================

config = GPTConfig(vocab_size=vocab_size)


def get_batch(split):
    data_split = train_data if split == "train" else val_data
    ix = torch.randint(len(data_split) - config.block_size, (BATCH_SIZE,))
    x = torch.stack([data_split[i : i + config.block_size] for i in ix])
    y = torch.stack([data_split[i + 1 : i + config.block_size + 1] for i in ix])
    x, y = x.to(DEVICE), y.to(DEVICE)
    return x, y


# ============================================================
# Model Architecture
# ============================================================


class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.n_head = config.n_head
        self.head_dim = config.n_embd // config.n_head
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.register_buffer(
            "mask",
            torch.tril(torch.ones(config.block_size, config.block_size)).view(
                1, 1, config.block_size, config.block_size
            ),
        )

    def forward(self, x):
        B, T, C = x.shape
        qkv = self.c_attn(x)
        q, k, v = qkv.split(C, dim=2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        att = (q @ k.transpose(-2, -1)) * (self.head_dim**-0.5)
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)
        y = att @ v

        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.resid_dropout(self.c_proj(y))
        return y


class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(config.n_embd, 4 * config.n_embd),
            nn.GELU(),
            nn.Linear(4 * config.n_embd, config.n_embd),
            nn.Dropout(config.dropout),
        )

    def forward(self, x):
        return self.net(x)


class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln2 = nn.LayerNorm(config.n_embd)
        self.ff = FeedForward(config)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ff(self.ln2(x))
        return x


class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
        self.position_embedding = nn.Embedding(config.block_size, config.n_embd)
        self.drop = nn.Dropout(config.dropout)
        self.blocks = nn.Sequential(
            *[TransformerBlock(config) for _ in range(config.n_layer)]
        )
        self.ln_f = nn.LayerNorm(config.n_embd)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.token_embedding.weight = self.lm_head.weight
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok_emb = self.token_embedding(idx)
        pos_emb = self.position_embedding(torch.arange(T, device=idx.device))
        x = self.drop(tok_emb + pos_emb)
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)), targets.view(-1)
            )
        return logits, loss

    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.config.block_size :]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx


# ============================================================
# Training
# ============================================================

model = GPT(config).to(DEVICE)
param_count = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Model parameters: {param_count:,}")

optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)


@torch.no_grad()
def estimate_loss():
    model.eval()
    losses = {}
    for split in ["train", "val"]:
        batch_losses = torch.zeros(EVAL_ITERS)
        for k in range(EVAL_ITERS):
            X, Y = get_batch(split)
            _, loss = model(X, Y)
            batch_losses[k] = loss.item()
        losses[split] = batch_losses.mean()
    model.train()
    return losses


print("\nStarting training...")
for iter_num in range(MAX_ITERS):
    if iter_num % EVAL_INTERVAL == 0:
        losses = estimate_loss()
        print(
            f"Step {iter_num:5d} | train loss: {losses['train']:.4f} | val loss: {losses['val']:.4f}"
        )

    xb, yb = get_batch("train")
    _, loss = model(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()

losses = estimate_loss()
print(
    f"Final    | train loss: {losses['train']:.4f} | val loss: {losses['val']:.4f}"
)

torch.save(
    {"model_state_dict": model.state_dict(), "config": config, "chars": chars},
    "jsgpt_model.pt",
)
print("Model saved to jsgpt_model.pt")

# ============================================================
# Generation
# ============================================================

print("\n" + "=" * 50)
print("GENERATING JAVASCRIPT")
print("=" * 50 + "\n")

model.eval()
prompts = ["function add(", "const app = ", "if (x > ", "class User {"]

for prompt in prompts:
    context = torch.tensor([encode(prompt)], dtype=torch.long, device=DEVICE)
    generated = model.generate(context, max_new_tokens=200)
    output = decode(generated[0].tolist())
    print(f"--- Prompt: {prompt!r} ---")
    print(output[:300])
    print()

Secret mission

Temperature and Top-k Sampling

Your model generates JavaScript, but you have no control over how creative or predictable its output is. Right now every token is sampled with equal randomness. What if you could turn a dial between safe, repetitive code and wild, experimental completions?

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally on your machine, so there are no ongoing costs. Decide whether to keep your files to experiment further, or delete them to free up disk space.

Resources you used:

  • Python virtual environment (with PyTorch and datasets library installed).
  • Training data file (javascript_data.txt, ~5MB).
  • Trained model checkpoint (jsgpt_model.pt).
  • Python scripts (jsgpt.py and prepare_data.py).

Keep everything running

No action needed. Choose this if you want to keep experimenting with your model, trying different prompts, adjusting hyperparameters, or retraining on more data.

  • Your project folder stays on your Desktop with all scripts, data, and the trained checkpoint intact.
  • You can reactivate the virtual environment any time to run jsgpt.py again.

Pause - I'll come back to this later

Deactivate your virtual environment to free up your terminal session. Your files stay on disk so you can pick up where you left off.

  • Deactivate the virtual environment by running:
deactivate
  • Close VS Code when you are done.
  • When you return, reactivate the virtual environment with source venv/bin/activate from inside your project folder.

Delete - I don't want to use this again

Remove all project files, the virtual environment, and the trained model to reclaim disk space.

  • Deactivate the virtual environment first by running:
deactivate
  • Navigate to your Desktop and remove the project folder by running:
cd ~/Desktop
rm -rf [[PROJECT_FOLDER="your project folder name"]]
  • Confirm deletion by running:
ls
  • Your project folder should no longer appear in the output.

That's a Wrap!

That's a Wrap!

Nice work! You just built a working GPT-style transformer model entirely from scratch and trained it to generate real JavaScript code.

You've learned how to:

  • Build a ~10M parameter transformer model from scratch in PyTorch, implementing causal self-attention, embeddings, residual connections, and layer normalization.
  • Train the model on real JavaScript code from GitHub using MPS GPU acceleration on your M1, with a complete training loop including loss tracking, gradient clipping, and checkpoint saving.
  • Generate syntactically valid JavaScript code completions using autoregressive generation, understanding how models like GPT-4 and Copilot produce text one token at a time.
  • Secret Mission: Implemented temperature and top-k sampling to control how creative vs. predictable your model's output is.

Ready to quiz yourself?