reading this record

code is folded so the argument and saved outputs come first. expand any cell to inspect the implementation. saved outputs are historical evidence from the source notebook, not a new run or independent verification.

Bigram language model¶

Prepare the Tiny Shakespeare token stream. The next step is to build context-window (x, y) samples.

inspect code · cell 01

In [1]:

import sys
from pathlib import Path

repo_root = Path.cwd()
if not (repo_root / "data").exists():
    repo_root = repo_root.parent
sys.path.insert(0, str(repo_root))

from src.dataset import get_batch, load_tiny_shakespeare_tokens, split_token_stream
from src.tokenizer import bpe_decode

Load and encode Tiny Shakespeare¶

inspect code · cell 02

In [2]:

tokens, vocab, _ = load_tiny_shakespeare_tokens(repo_root / "data")

Split the token stream¶

inspect code · cell 03

In [3]:

train_tokens, validation_tokens = split_token_stream(tokens)

Sanity check¶

inspect code · cell 04

In [4]:

print(f"Total tokens: {len(tokens):,}")
print(f"Training tokens: {len(train_tokens):,}")
print(f"Validation tokens: {len(validation_tokens):,}")
# print(bpe_decode(train_tokens[:100], vocab))

Total tokens: 590,498
Training tokens: 531,448
Validation tokens: 59,050

inspect code · cell 05

In [5]:

import random

# Set parameters for batching
seed = 42
random.seed(seed)
block_size = 8
batch_size = 32


def sample_batch(split):
    return get_batch(
        split,
        train_tokens,
        validation_tokens,
        block_size,
        batch_size,
    )

inspect code · cell 06

In [6]:

import torch
from torch import nn

vocab_size = len(vocab)
n_embd = 32


class BigramLanguageModel(nn.Module):
    def __init__(self, vocab_size, n_embd):
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        x = self.token_embedding_table(idx)
        logits = self.lm_head(x)

        if targets is None:
            loss = None
        else:
            B, T, V = logits.shape
            logits = logits.view(B * T, V)
            targets = targets.view(B * T)
            loss = nn.functional.cross_entropy(logits, targets)

        return logits, loss

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

        return idx

inspect code · cell 07

In [7]:

model = BigramLanguageModel(vocab_size, n_embd)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

for step in range(1000):
    x_batch, y_batch = sample_batch("train")
    x_batch = torch.tensor(x_batch, dtype=torch.long)
    y_batch = torch.tensor(y_batch, dtype=torch.long)

    logits, loss = model(x_batch, y_batch)

    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

    if step % 100 == 0:
        print(f"Step {step}: loss = {loss.item():.4f}")

Step 0: loss = 6.4571
Step 100: loss = 6.1014
Step 200: loss = 5.7985
Step 300: loss = 5.5250
Step 400: loss = 5.2603
Step 500: loss = 5.1149
Step 600: loss = 4.8774
Step 700: loss = 4.6288
Step 800: loss = 4.5432
Step 900: loss = 4.4973

inspect code · cell 08

In [8]:

import math

model.eval()

with torch.no_grad():
    xb, yb = sample_batch("validation")

    xb = torch.tensor(xb, dtype=torch.long)
    yb = torch.tensor(yb, dtype=torch.long)

    _, val_loss = model(xb, yb)

print("validation loss:", val_loss.item())
print("log vocab size:", math.log(len(vocab)))

model.train()

validation loss: 4.3758955001831055
log vocab size: 6.238324625039508

Out[8]:

BigramLanguageModel(
  (token_embedding_table): Embedding(512, 32)
  (lm_head): Linear(in_features=32, out_features=512, bias=True)
)

inspect code · cell 09

In [9]:

context = torch.tensor([[0]], dtype=torch.long)
out = model.generate(context, max_new_tokens=100)
print(bpe_decode(out[0].tolist(), vocab, errors="replace"))

�d
ookPAE:
Oerhour,
R:
Tll on , and the e to to youPant shs ow, behoughghton neygmeweie emcyour aanour  and such'd I git my atOr ce, syspbuINGAORove:
ut aupnoRLence:
NG that sece �y bebla er