Upload folder using huggingface_hub
Browse files- README.md +120 -1
- bigram.py +266 -0
- config.json +20 -0
- load_model.py +64 -0
- model.pt +3 -0
- tokenizer.py +79 -0
- tokenizer_char.json +69 -0
README.md
CHANGED
|
@@ -1,3 +1,122 @@
|
|
| 1 |
---
|
| 2 |
-
license:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 3 |
---
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
---
|
| 2 |
+
license: mit
|
| 3 |
+
tags:
|
| 4 |
+
- educational
|
| 5 |
+
- from-scratch
|
| 6 |
+
- character-level
|
| 7 |
+
- transformer
|
| 8 |
+
- pytorch
|
| 9 |
---
|
| 10 |
+
|
| 11 |
+
# SIL-v01 — Tiny From-Scratch Language Model
|
| 12 |
+
|
| 13 |
+
**This is an educational project, not a practical language model.** It is a
|
| 14 |
+
from-scratch, char-RNN-scale transformer built to learn the mechanics of
|
| 15 |
+
language modeling end-to-end: tokenization → embeddings → the
|
| 16 |
+
language-modeling objective → training loop → self-attention → subword
|
| 17 |
+
tokenization. It is not intended for any downstream or production use —
|
| 18 |
+
treat it as a worked example, not a tool.
|
| 19 |
+
|
| 20 |
+
The published checkpoint here is the **best-performing configuration found
|
| 21 |
+
in a controlled scaling sweep**: a 4-layer, 4-head, causal self-attention
|
| 22 |
+
transformer with `n_embd=128`, trained on a character-level tokenizer.
|
| 23 |
+
|
| 24 |
+
## Model architecture
|
| 25 |
+
|
| 26 |
+
Token + position embeddings → 4 stacked pre-norm transformer blocks
|
| 27 |
+
(multi-head causal self-attention + feedforward, residual connections) →
|
| 28 |
+
final LayerNorm → linear head to vocab logits. Same family as a minimal
|
| 29 |
+
GPT, at a scale that trains on a CPU in minutes. Full hyperparameters are
|
| 30 |
+
in `config.json`; nothing about the architecture is hardcoded in the loader
|
| 31 |
+
— see `load_model.py`.
|
| 32 |
+
|
| 33 |
+
| | |
|
| 34 |
+
|---|---|
|
| 35 |
+
| Parameters | 812,609 |
|
| 36 |
+
| n_embd | 128 |
|
| 37 |
+
| n_head | 4 |
|
| 38 |
+
| n_layer | 4 |
|
| 39 |
+
| block_size | 32 |
|
| 40 |
+
| Tokenizer | character-level, vocab_size=65 |
|
| 41 |
+
|
| 42 |
+
## Training corpus
|
| 43 |
+
|
| 44 |
+
Trained on **TinyShakespeare** (~1.1M characters), the classic small corpus
|
| 45 |
+
for character-level language modeling. It is fetched directly from Andrej
|
| 46 |
+
Karpathy's `char-rnn` GitHub repository
|
| 47 |
+
(`raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt`)
|
| 48 |
+
— a concatenation of Shakespeare's plays, ultimately public-domain text, but
|
| 49 |
+
sourced from that repo directly rather than downloaded from Project
|
| 50 |
+
Gutenberg. (This project's data pipeline also includes a Project Gutenberg
|
| 51 |
+
boilerplate-stripper for *other* corpora it experiments with — Dostoevsky,
|
| 52 |
+
Nietzsche, Suetonius — but that path was not used for this checkpoint.)
|
| 53 |
+
|
| 54 |
+
90/10 train/val split, held fixed across every run below.
|
| 55 |
+
|
| 56 |
+
## Results
|
| 57 |
+
|
| 58 |
+
Four runs, all on TinyShakespeare, same architecture family (`n_layer=4`,
|
| 59 |
+
`n_head=4`, `block_size=32`, `lr=1e-3`, `max_iters=8000`, `batch_size=32`,
|
| 60 |
+
seed=1337) — a controlled sweep over embedding width, plus one run swapping
|
| 61 |
+
in a trained BPE subword tokenizer instead of character-level.
|
| 62 |
+
|
| 63 |
+
**Raw cross-entropy (nats/token) is not comparable across tokenizers** —
|
| 64 |
+
a 1000-token BPE vocabulary and a 65-token char vocabulary have different
|
| 65 |
+
random-guessing floors (ln(1000) ≈ 6.9 vs ln(65) ≈ 4.2), so a higher raw
|
| 66 |
+
loss on BPE does not mean a worse model. **Bits-per-character (BPC)**
|
| 67 |
+
normalizes both onto the same unit — bits of model surprise per character
|
| 68 |
+
of the *original* text — and is the number to compare across the table.
|
| 69 |
+
|
| 70 |
+
| Tokenizer | n_embd | Params | Train loss (nats) | Train BPC | Val loss (nats) | Val BPC | Train/val gap (BPC) | Wall clock |
|
| 71 |
+
|---|---|---|---|---|---|---|---|---|
|
| 72 |
+
| char | 32 | 55,745 | 1.7159 | 2.4755 | 1.8849 | 2.7193 | 0.2438 | 4m48.8s |
|
| 73 |
+
| char | 64 | 209,729 | 1.5247 | 2.1997 | 1.7220 | 2.4843 | 0.2846 | 7m19.4s |
|
| 74 |
+
| char | **128** | **812,609** | **1.4043** | **2.0260** | **1.6257** | **2.3454** | 0.3194 | 16m37.6s |
|
| 75 |
+
| bpe (vocab=1000) | 32 | 116,520 | 3.5683 | 2.1272 | 3.8843 | 2.4189 | 0.2917 | 6m29.5s |
|
| 76 |
+
|
| 77 |
+
**Published checkpoint: char, n_embd=128** (bolded row) — lowest val BPC of
|
| 78 |
+
the char-level sweep.
|
| 79 |
+
|
| 80 |
+
Scaling observation: doubling `n_embd` costs a roughly constant ~3.8×
|
| 81 |
+
parameters each step (32→64, 64→128), but the val-BPC improvement shrinks
|
| 82 |
+
each time (0.235 → 0.139 in BPC terms), and the train/val gap widens —
|
| 83 |
+
diminishing, saturating returns from embedding width alone once depth
|
| 84 |
+
(`n_layer=4`) and context (`block_size=32`) are held fixed.
|
| 85 |
+
|
| 86 |
+
Tokenizer observation: char-level (BPC 2.35) still edges out this one BPE
|
| 87 |
+
run (BPC 2.42) at comparable parameter count, but BPE achieves that with
|
| 88 |
+
~7x fewer parameters than the 812K char model and visibly more coherent
|
| 89 |
+
generated chunks — a fair head-to-head at matched parameter count wasn't
|
| 90 |
+
run here.
|
| 91 |
+
|
| 92 |
+
## Files
|
| 93 |
+
|
| 94 |
+
| File | Purpose |
|
| 95 |
+
|---|---|
|
| 96 |
+
| `model.pt` | `state_dict()` of the trained `BigramLanguageModel` |
|
| 97 |
+
| `config.json` | Every hyperparameter needed to reconstruct the architecture and tokenizer |
|
| 98 |
+
| `tokenizer_char.json` | The character-level tokenizer's `stoi` vocabulary |
|
| 99 |
+
| `bigram.py` | Model architecture source (vendored so this repo is self-sufficient) |
|
| 100 |
+
| `tokenizer.py` | Character tokenizer source (vendored, same reason) |
|
| 101 |
+
| `load_model.py` | Reconstructs the model + tokenizer from `config.json` and generates a sample |
|
| 102 |
+
|
| 103 |
+
## Usage
|
| 104 |
+
|
| 105 |
+
```bash
|
| 106 |
+
pip install torch
|
| 107 |
+
python load_model.py
|
| 108 |
+
```
|
| 109 |
+
|
| 110 |
+
`load_model.py` reads `config.json` for every architectural parameter
|
| 111 |
+
(`vocab_size`, `n_embd`, `n_head`, `n_layer`, `block_size`, `tokenizer_type`)
|
| 112 |
+
— it does not assume or hardcode them — loads `model.pt` into a freshly
|
| 113 |
+
constructed model, and samples 400 characters to prove the checkpoint and
|
| 114 |
+
its config agree.
|
| 115 |
+
|
| 116 |
+
## Limitations
|
| 117 |
+
|
| 118 |
+
This is a ~65-character-vocabulary, 32-token-context toy model trained for
|
| 119 |
+
8,000 steps on 1MB of text. It reproduces surface Shakespeare-ish texture
|
| 120 |
+
(names, verse-like line breaks, archaic diction) but not coherent meaning,
|
| 121 |
+
plot, or factual content. **Do not use this for anything beyond studying
|
| 122 |
+
how these mechanics fit together.**
|
bigram.py
ADDED
|
@@ -0,0 +1,266 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Stage 2: The bigram language model — the simplest possible neural LM.
|
| 3 |
+
|
| 4 |
+
Originally this was ONE embedding table of shape (vocab_size, vocab_size):
|
| 5 |
+
row i a vector of raw next-token scores for "current token is i." Then a
|
| 6 |
+
second table added position information. Then a single self-attention Head.
|
| 7 |
+
Now the head is split into several narrower heads that run in parallel.
|
| 8 |
+
|
| 9 |
+
- token_embedding_table: one learned vector per vocab entry — the
|
| 10 |
+
token's identity, regardless of where it sits.
|
| 11 |
+
- position_embedding_table: one learned vector per slot in the context
|
| 12 |
+
window — "this is position 3," regardless of
|
| 13 |
+
which token is there.
|
| 14 |
+
- MultiHeadAttention: n_head independent Heads, each working in a
|
| 15 |
+
narrower head_size = n_embd // n_head space,
|
| 16 |
+
run in parallel and concatenated back to
|
| 17 |
+
n_embd. Same total width as one big head, but
|
| 18 |
+
each head can specialize on a different kind
|
| 19 |
+
of relationship between positions instead of
|
| 20 |
+
averaging all of them into one.
|
| 21 |
+
|
| 22 |
+
predict the next token; measure error with cross-entropy loss.
|
| 23 |
+
|
| 24 |
+
Everything a GPT does is this same objective with a smarter architecture.
|
| 25 |
+
Note the connection to the Word2vec slide from your lecture: like CBOW /
|
| 26 |
+
Skip-gram, this model learns embeddings. Position embeddings were a crude,
|
| 27 |
+
static fix for word order; self-attention is the real one — instead of a
|
| 28 |
+
fixed "position 3" vector, each token computes a data-dependent, weighted
|
| 29 |
+
average of what came before it.
|
| 30 |
+
"""
|
| 31 |
+
|
| 32 |
+
from __future__ import annotations
|
| 33 |
+
|
| 34 |
+
import torch
|
| 35 |
+
import torch.nn as nn
|
| 36 |
+
from torch.nn import functional as F
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
class Head(nn.Module):
|
| 40 |
+
"""One head of causal self-attention."""
|
| 41 |
+
|
| 42 |
+
def __init__(self, n_embd: int, head_size: int, block_size: int) -> None:
|
| 43 |
+
super().__init__()
|
| 44 |
+
# Three separate linear projections of the same input x, each
|
| 45 |
+
# (n_embd, head_size). No bias — these are pure projections, like
|
| 46 |
+
# rotating/rescaling x into a new space, not adding an offset.
|
| 47 |
+
self.key = nn.Linear(n_embd, head_size, bias=False)
|
| 48 |
+
self.query = nn.Linear(n_embd, head_size, bias=False)
|
| 49 |
+
self.value = nn.Linear(n_embd, head_size, bias=False)
|
| 50 |
+
# A (block_size, block_size) lower-triangular matrix of 1s, e.g. for
|
| 51 |
+
# block_size=4:
|
| 52 |
+
# 1 0 0 0
|
| 53 |
+
# 1 1 0 0
|
| 54 |
+
# 1 1 1 0
|
| 55 |
+
# 1 1 1 1
|
| 56 |
+
# Row t has 1s in columns 0..t: "position t may attend to positions
|
| 57 |
+
# 0..t." Registered as a buffer (not a parameter) so it moves with
|
| 58 |
+
# .to(device) and is saved in state_dict, but is never trained.
|
| 59 |
+
self.register_buffer("tril", torch.tril(torch.ones(block_size, block_size)))
|
| 60 |
+
self.head_size = head_size
|
| 61 |
+
|
| 62 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 63 |
+
"""
|
| 64 |
+
x: (B, T, n_embd) — the token+position embeddings for one batch.
|
| 65 |
+
Returns (B, T, head_size): each position's attention-weighted
|
| 66 |
+
summary of itself and everything before it.
|
| 67 |
+
"""
|
| 68 |
+
B, T, _ = x.shape
|
| 69 |
+
|
| 70 |
+
# Every position asks a question (query) and advertises an answer
|
| 71 |
+
# (key). Both (B, T, n_embd) -> (B, T, head_size).
|
| 72 |
+
q = self.query(x)
|
| 73 |
+
k = self.key(x)
|
| 74 |
+
|
| 75 |
+
# Compare every query to every key via dot product:
|
| 76 |
+
# (B, T, head_size) @ (B, head_size, T) -> (B, T, T).
|
| 77 |
+
# wei[b, i, j] = "how much position i's query matches position j's
|
| 78 |
+
# key" — raw, unnormalized affinity, before masking or softmax.
|
| 79 |
+
wei = q @ k.transpose(-2, -1)
|
| 80 |
+
# Divide by sqrt(head_size) before softmax — see the explanation
|
| 81 |
+
# printed by the demo script; in short, it keeps the variance of
|
| 82 |
+
# these dot products at ~1 instead of growing with head_size, so
|
| 83 |
+
# softmax doesn't saturate into a near one-hot distribution.
|
| 84 |
+
wei = wei * self.head_size**-0.5
|
| 85 |
+
|
| 86 |
+
# Causal mask: position i must not see position j > i (the future).
|
| 87 |
+
# tril[:T, :T] == 0 marks the upper triangle (j > i); those entries
|
| 88 |
+
# become -inf so softmax turns them into exactly 0.
|
| 89 |
+
wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf"))
|
| 90 |
+
# Softmax over the last dim (T, the "which key" axis): every row i
|
| 91 |
+
# becomes a probability distribution over positions 0..i that sums
|
| 92 |
+
# to 1. Still (B, T, T).
|
| 93 |
+
wei = F.softmax(wei, dim=-1)
|
| 94 |
+
|
| 95 |
+
# Every position also advertises a value: (B, T, n_embd) ->
|
| 96 |
+
# (B, T, head_size). This is what actually gets mixed together,
|
| 97 |
+
# as opposed to key/query which only decide the mixing weights.
|
| 98 |
+
v = self.value(x)
|
| 99 |
+
# (B, T, T) @ (B, T, head_size) -> (B, T, head_size): position i's
|
| 100 |
+
# output is the wei[i, :]-weighted average of every value vector at
|
| 101 |
+
# positions 0..i.
|
| 102 |
+
out = wei @ v
|
| 103 |
+
return out
|
| 104 |
+
|
| 105 |
+
|
| 106 |
+
class MultiHeadAttention(nn.Module):
|
| 107 |
+
"""Several causal self-attention heads run in parallel, then combined."""
|
| 108 |
+
|
| 109 |
+
def __init__(self, n_embd: int, n_head: int, block_size: int) -> None:
|
| 110 |
+
super().__init__()
|
| 111 |
+
# n_embd must split evenly across n_head so the concatenated output
|
| 112 |
+
# lands back at exactly n_embd (e.g. 32 = 4 heads * 8).
|
| 113 |
+
assert n_embd % n_head == 0, "n_embd must be divisible by n_head"
|
| 114 |
+
head_size = n_embd // n_head
|
| 115 |
+
self.heads = nn.ModuleList(
|
| 116 |
+
[Head(n_embd, head_size, block_size) for _ in range(n_head)]
|
| 117 |
+
)
|
| 118 |
+
# Each head reads the same x (n_embd-wide) and produces its own
|
| 119 |
+
# head_size-wide output. Concatenating n_head of those gets back to
|
| 120 |
+
# n_embd; this projection lets the heads' outputs mix before moving
|
| 121 |
+
# on, rather than just being stapled together.
|
| 122 |
+
self.proj = nn.Linear(n_embd, n_embd)
|
| 123 |
+
|
| 124 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 125 |
+
# Each h(x): (B, T, head_size). Concatenate along the last dim ->
|
| 126 |
+
# (B, T, n_head * head_size) == (B, T, n_embd).
|
| 127 |
+
out = torch.cat([h(x) for h in self.heads], dim=-1)
|
| 128 |
+
return self.proj(out)
|
| 129 |
+
|
| 130 |
+
|
| 131 |
+
class FeedForward(nn.Module):
|
| 132 |
+
"""Per-position MLP: the "thinking" step after attention's "looking."
|
| 133 |
+
|
| 134 |
+
Attention only moves information between positions via weighted
|
| 135 |
+
averages of value vectors — it can't recombine features nonlinearly.
|
| 136 |
+
This module runs independently on every position (no mixing across T)
|
| 137 |
+
and gives the model a place to do that recombination. The inner
|
| 138 |
+
dimension is widened 4x, as in the original Transformer paper, then
|
| 139 |
+
projected back to n_embd so it can be added into the residual stream.
|
| 140 |
+
"""
|
| 141 |
+
|
| 142 |
+
def __init__(self, n_embd: int) -> None:
|
| 143 |
+
super().__init__()
|
| 144 |
+
self.net = nn.Sequential(
|
| 145 |
+
nn.Linear(n_embd, 4 * n_embd),
|
| 146 |
+
nn.ReLU(),
|
| 147 |
+
nn.Linear(4 * n_embd, n_embd),
|
| 148 |
+
)
|
| 149 |
+
|
| 150 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 151 |
+
return self.net(x)
|
| 152 |
+
|
| 153 |
+
|
| 154 |
+
class Block(nn.Module):
|
| 155 |
+
"""One transformer block: communicate (attention), then compute (FFN).
|
| 156 |
+
|
| 157 |
+
Pre-norm residual design: x = x + sublayer(LN(x)) for both sublayers.
|
| 158 |
+
The "+" is what makes this residual — each sublayer adds a correction
|
| 159 |
+
to x rather than replacing it, so during backprop gradients have a
|
| 160 |
+
direct addition path (derivative 1) back through every block, no matter
|
| 161 |
+
how deep the stack gets. Without it, gradients must flow through each
|
| 162 |
+
block's full nonlinear transform in series, and repeated multiplication
|
| 163 |
+
by small Jacobians is exactly the vanishing-gradient failure mode from
|
| 164 |
+
the RNN slide. LayerNorm before each sublayer keeps that sublayer's
|
| 165 |
+
input at a stable, zero-mean/unit-variance scale regardless of how
|
| 166 |
+
large the residual stream has grown by this depth.
|
| 167 |
+
"""
|
| 168 |
+
|
| 169 |
+
def __init__(self, n_embd: int, n_head: int, block_size: int) -> None:
|
| 170 |
+
super().__init__()
|
| 171 |
+
self.sa = MultiHeadAttention(n_embd, n_head, block_size)
|
| 172 |
+
self.ffwd = FeedForward(n_embd)
|
| 173 |
+
self.ln1 = nn.LayerNorm(n_embd)
|
| 174 |
+
self.ln2 = nn.LayerNorm(n_embd)
|
| 175 |
+
|
| 176 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 177 |
+
x = x + self.sa(self.ln1(x))
|
| 178 |
+
x = x + self.ffwd(self.ln2(x))
|
| 179 |
+
return x
|
| 180 |
+
|
| 181 |
+
|
| 182 |
+
class BigramLanguageModel(nn.Module):
|
| 183 |
+
def __init__(
|
| 184 |
+
self,
|
| 185 |
+
vocab_size: int,
|
| 186 |
+
block_size: int,
|
| 187 |
+
n_embd: int,
|
| 188 |
+
n_head: int,
|
| 189 |
+
n_layer: int,
|
| 190 |
+
) -> None:
|
| 191 |
+
super().__init__()
|
| 192 |
+
self.block_size = block_size
|
| 193 |
+
# (vocab_size, n_embd): row i is token i's identity vector.
|
| 194 |
+
self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
|
| 195 |
+
# (block_size, n_embd): row t is position t's vector. There are
|
| 196 |
+
# only block_size rows because that's the longest context this
|
| 197 |
+
# model will ever be asked about — see the crop in generate().
|
| 198 |
+
self.position_embedding_table = nn.Embedding(block_size, n_embd)
|
| 199 |
+
# n_layer stacked (attention, feedforward) blocks with residual
|
| 200 |
+
# connections and layernorm — see Block's docstring for why those
|
| 201 |
+
# two aren't decoration.
|
| 202 |
+
self.blocks = nn.Sequential(
|
| 203 |
+
*[Block(n_embd, n_head, block_size) for _ in range(n_layer)]
|
| 204 |
+
)
|
| 205 |
+
# One more LayerNorm after the last block. The residual stream's
|
| 206 |
+
# scale is unconstrained as it accumulates additions across n_layer
|
| 207 |
+
# blocks; this normalizes it once before the final projection reads
|
| 208 |
+
# it, matching the scale lm_head was initialized to expect.
|
| 209 |
+
self.ln_f = nn.LayerNorm(n_embd)
|
| 210 |
+
# Project the final n_embd representation to vocab-sized logits.
|
| 211 |
+
self.lm_head = nn.Linear(n_embd, vocab_size)
|
| 212 |
+
|
| 213 |
+
def forward(
|
| 214 |
+
self, idx: torch.Tensor, targets: torch.Tensor | None = None
|
| 215 |
+
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
| 216 |
+
"""
|
| 217 |
+
idx: (B, T) batch of token-id sequences
|
| 218 |
+
targets: (B, T) the same sequences shifted one position left,
|
| 219 |
+
i.e. "the correct next token at every position"
|
| 220 |
+
|
| 221 |
+
Returns (logits, loss). Loss is None during generation.
|
| 222 |
+
"""
|
| 223 |
+
B, T = idx.shape
|
| 224 |
+
|
| 225 |
+
# (B, T) -> (B, T, n_embd): each token id looks up its identity vector.
|
| 226 |
+
tok_emb = self.token_embedding_table(idx)
|
| 227 |
+
# (T,) -> (T, n_embd): one vector per position, shared across the batch.
|
| 228 |
+
pos_emb = self.position_embedding_table(torch.arange(T, device=idx.device))
|
| 229 |
+
# (B, T, n_embd) + (T, n_embd): the position vectors broadcast over
|
| 230 |
+
# the batch dim, so every sequence gets its own token vectors plus
|
| 231 |
+
# the same T position vectors. Result: (B, T, n_embd).
|
| 232 |
+
x = tok_emb + pos_emb
|
| 233 |
+
# (B, T, n_embd) -> (B, T, n_embd): n_layer rounds of (attend, then
|
| 234 |
+
# think), each wrapped in a residual connection.
|
| 235 |
+
x = self.blocks(x)
|
| 236 |
+
# Normalize the accumulated residual stream once before reading it.
|
| 237 |
+
x = self.ln_f(x)
|
| 238 |
+
# (B, T, n_embd) -> (B, T, vocab_size): project back to next-token scores.
|
| 239 |
+
logits = self.lm_head(x)
|
| 240 |
+
|
| 241 |
+
loss = None
|
| 242 |
+
if targets is not None:
|
| 243 |
+
B, T, C = logits.shape
|
| 244 |
+
# cross_entropy wants (N, C) vs (N,), so flatten batch+time.
|
| 245 |
+
loss = F.cross_entropy(logits.view(B * T, C), targets.view(B * T))
|
| 246 |
+
return logits, loss
|
| 247 |
+
|
| 248 |
+
@torch.no_grad()
|
| 249 |
+
def generate(self, idx: torch.Tensor, max_new_tokens: int) -> torch.Tensor:
|
| 250 |
+
"""
|
| 251 |
+
Autoregressive sampling: feed the sequence in, take the logits at
|
| 252 |
+
the LAST position, turn them into probabilities, sample one token,
|
| 253 |
+
append it, repeat. This loop is identical in shape to how a full
|
| 254 |
+
GPT generates text — only the model in the middle changes.
|
| 255 |
+
"""
|
| 256 |
+
for _ in range(max_new_tokens):
|
| 257 |
+
# position_embedding_table has only block_size rows, so a
|
| 258 |
+
# sequence longer than that would index out of range — crop to
|
| 259 |
+
# the last block_size tokens before every forward pass.
|
| 260 |
+
idx_cond = idx[:, -self.block_size :]
|
| 261 |
+
logits, _ = self(idx_cond) # (B, T, C)
|
| 262 |
+
logits = logits[:, -1, :] # last time step: (B, C)
|
| 263 |
+
probs = F.softmax(logits, dim=-1) # scores -> probabilities
|
| 264 |
+
next_id = torch.multinomial(probs, num_samples=1) # sample
|
| 265 |
+
idx = torch.cat((idx, next_id), dim=1)
|
| 266 |
+
return idx
|
config.json
ADDED
|
@@ -0,0 +1,20 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architecture": "BigramLanguageModel",
|
| 3 |
+
"architecture_description": "Token + position embeddings -> n_layer stacked transformer blocks (pre-norm multi-head causal self-attention + feedforward, residual connections) -> final LayerNorm -> linear head to vocab logits.",
|
| 4 |
+
"vocab_size": 65,
|
| 5 |
+
"n_embd": 128,
|
| 6 |
+
"n_head": 4,
|
| 7 |
+
"n_layer": 4,
|
| 8 |
+
"block_size": 32,
|
| 9 |
+
"tokenizer_type": "char",
|
| 10 |
+
"tokenizer_file": "tokenizer_char.json",
|
| 11 |
+
"weights_file": "model.pt",
|
| 12 |
+
"training": {
|
| 13 |
+
"dataset": "tinyshakespeare.txt",
|
| 14 |
+
"batch_size": 32,
|
| 15 |
+
"learning_rate": 1e-3,
|
| 16 |
+
"max_iters": 8000,
|
| 17 |
+
"seed": 1337
|
| 18 |
+
},
|
| 19 |
+
"parameter_count": 812609
|
| 20 |
+
}
|
load_model.py
ADDED
|
@@ -0,0 +1,64 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Self-contained loader for this checkpoint.
|
| 3 |
+
|
| 4 |
+
Reconstructs the model architecture and tokenizer entirely from config.json
|
| 5 |
+
(no hardcoded hyperparameters here), loads the trained weights, and
|
| 6 |
+
generates a sample -- proving the artifact set (config.json + model.pt +
|
| 7 |
+
tokenizer_char.json + bigram.py + tokenizer.py) is sufficient on its own,
|
| 8 |
+
with no dependency on the original training repo or its train.py.
|
| 9 |
+
|
| 10 |
+
Usage:
|
| 11 |
+
python load_model.py
|
| 12 |
+
"""
|
| 13 |
+
|
| 14 |
+
from __future__ import annotations
|
| 15 |
+
|
| 16 |
+
import json
|
| 17 |
+
from pathlib import Path
|
| 18 |
+
|
| 19 |
+
import torch
|
| 20 |
+
|
| 21 |
+
from bigram import BigramLanguageModel
|
| 22 |
+
from tokenizer import CharTokenizer
|
| 23 |
+
|
| 24 |
+
HERE = Path(__file__).resolve().parent
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def load_model_and_tokenizer(dir_path: Path = HERE):
|
| 28 |
+
config = json.loads((dir_path / "config.json").read_text())
|
| 29 |
+
|
| 30 |
+
if config["tokenizer_type"] != "char":
|
| 31 |
+
raise ValueError(
|
| 32 |
+
f"This loader only wires up 'char'; config says {config['tokenizer_type']!r}"
|
| 33 |
+
)
|
| 34 |
+
tokenizer = CharTokenizer.load(dir_path / config["tokenizer_file"])
|
| 35 |
+
|
| 36 |
+
model = BigramLanguageModel(
|
| 37 |
+
vocab_size=config["vocab_size"],
|
| 38 |
+
block_size=config["block_size"],
|
| 39 |
+
n_embd=config["n_embd"],
|
| 40 |
+
n_head=config["n_head"],
|
| 41 |
+
n_layer=config["n_layer"],
|
| 42 |
+
)
|
| 43 |
+
state_dict = torch.load(dir_path / config["weights_file"], map_location="cpu")
|
| 44 |
+
model.load_state_dict(state_dict)
|
| 45 |
+
model.eval()
|
| 46 |
+
|
| 47 |
+
return model, tokenizer, config
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def main() -> None:
|
| 51 |
+
model, tokenizer, config = load_model_and_tokenizer()
|
| 52 |
+
|
| 53 |
+
n_params = sum(p.numel() for p in model.parameters())
|
| 54 |
+
print(f"Loaded model: {n_params:,} parameters (config.json says {config['parameter_count']:,})")
|
| 55 |
+
assert n_params == config["parameter_count"], "Reconstructed model doesn't match config!"
|
| 56 |
+
|
| 57 |
+
context = torch.zeros((1, 1), dtype=torch.long)
|
| 58 |
+
print("\n----- generated sample -----")
|
| 59 |
+
sample = tokenizer.decode(model.generate(context, max_new_tokens=400)[0].tolist())
|
| 60 |
+
print(sample)
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
if __name__ == "__main__":
|
| 64 |
+
main()
|
model.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:c13668656eb1fc5c87780a22b34dcb0ca4fdfe63e6e6c87a81235a95e5ffdc78
|
| 3 |
+
size 3348079
|
tokenizer.py
ADDED
|
@@ -0,0 +1,79 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Stage 1: Tokenization (character-level).
|
| 3 |
+
|
| 4 |
+
This maps directly onto the "Tokenization summary" slide you saw:
|
| 5 |
+
we're starting at the CHARACTER level because it's the simplest possible
|
| 6 |
+
scheme — near-zero risk of out-of-vocabulary tokens, trivially
|
| 7 |
+
interpretable, and the vocab is tiny (~65 symbols for Shakespeare).
|
| 8 |
+
The cost is longer sequences and less meaningful individual tokens.
|
| 9 |
+
|
| 10 |
+
Later (Stage 4) you'll swap this out for a subword tokenizer (BPE) and
|
| 11 |
+
watch how vocab size, sequence length, and sample quality change.
|
| 12 |
+
That contrast is one of the best lessons in this whole exercise.
|
| 13 |
+
|
| 14 |
+
Key idea: a tokenizer is just two lookup tables.
|
| 15 |
+
encode: string -> list of integers
|
| 16 |
+
decode: list of integers -> string
|
| 17 |
+
Everything downstream (embeddings, attention) operates on the integers.
|
| 18 |
+
"""
|
| 19 |
+
|
| 20 |
+
from __future__ import annotations
|
| 21 |
+
|
| 22 |
+
import json
|
| 23 |
+
from pathlib import Path
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
class CharTokenizer:
|
| 27 |
+
"""A character-level tokenizer built from a training corpus."""
|
| 28 |
+
|
| 29 |
+
def __init__(self, text: str) -> None:
|
| 30 |
+
# The vocabulary is simply every unique character, sorted so the
|
| 31 |
+
# mapping is deterministic across runs.
|
| 32 |
+
chars = sorted(set(text))
|
| 33 |
+
self.vocab_size = len(chars)
|
| 34 |
+
|
| 35 |
+
# stoi = "string to integer", itos = "integer to string".
|
| 36 |
+
# These two dicts ARE the tokenizer. That's it.
|
| 37 |
+
self.stoi: dict[str, int] = {ch: i for i, ch in enumerate(chars)}
|
| 38 |
+
self.itos: dict[int, str] = {i: ch for i, ch in enumerate(chars)}
|
| 39 |
+
|
| 40 |
+
def encode(self, text: str) -> list[int]:
|
| 41 |
+
"""Convert a string into a list of token ids."""
|
| 42 |
+
return [self.stoi[ch] for ch in text]
|
| 43 |
+
|
| 44 |
+
def decode(self, ids: list[int]) -> str:
|
| 45 |
+
"""Convert a list of token ids back into a string."""
|
| 46 |
+
return "".join(self.itos[i] for i in ids)
|
| 47 |
+
|
| 48 |
+
# --- persistence, so a trained model ships with its tokenizer ---
|
| 49 |
+
|
| 50 |
+
def save(self, path: str | Path) -> None:
|
| 51 |
+
Path(path).write_text(
|
| 52 |
+
json.dumps({"stoi": self.stoi}, indent=2), encoding="utf-8"
|
| 53 |
+
)
|
| 54 |
+
|
| 55 |
+
@classmethod
|
| 56 |
+
def load(cls, path: str | Path) -> "CharTokenizer":
|
| 57 |
+
stoi = json.loads(Path(path).read_text(encoding="utf-8"))["stoi"]
|
| 58 |
+
tok = cls.__new__(cls)
|
| 59 |
+
tok.stoi = {k: int(v) for k, v in stoi.items()}
|
| 60 |
+
tok.itos = {int(v): k for k, v in stoi.items()}
|
| 61 |
+
tok.vocab_size = len(tok.stoi)
|
| 62 |
+
return tok
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
if __name__ == "__main__":
|
| 66 |
+
# Quick self-test / demo
|
| 67 |
+
data_path = Path(__file__).resolve().parent.parent / "data" / "input.txt"
|
| 68 |
+
text = data_path.read_text(encoding="utf-8")
|
| 69 |
+
|
| 70 |
+
tok = CharTokenizer(text)
|
| 71 |
+
print(f"Vocab size: {tok.vocab_size}")
|
| 72 |
+
print(f"Vocabulary: {''.join(tok.itos[i] for i in range(tok.vocab_size))!r}")
|
| 73 |
+
|
| 74 |
+
sample = "To be, or not to be"
|
| 75 |
+
ids = tok.encode(sample)
|
| 76 |
+
print(f"\nencode({sample!r})\n -> {ids}")
|
| 77 |
+
print(f"decode(...)\n -> {tok.decode(ids)!r}")
|
| 78 |
+
assert tok.decode(ids) == sample, "Round-trip failed!"
|
| 79 |
+
print("\nRound-trip OK ✔")
|
tokenizer_char.json
ADDED
|
@@ -0,0 +1,69 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"stoi": {
|
| 3 |
+
"\n": 0,
|
| 4 |
+
" ": 1,
|
| 5 |
+
"!": 2,
|
| 6 |
+
"$": 3,
|
| 7 |
+
"&": 4,
|
| 8 |
+
"'": 5,
|
| 9 |
+
",": 6,
|
| 10 |
+
"-": 7,
|
| 11 |
+
".": 8,
|
| 12 |
+
"3": 9,
|
| 13 |
+
":": 10,
|
| 14 |
+
";": 11,
|
| 15 |
+
"?": 12,
|
| 16 |
+
"A": 13,
|
| 17 |
+
"B": 14,
|
| 18 |
+
"C": 15,
|
| 19 |
+
"D": 16,
|
| 20 |
+
"E": 17,
|
| 21 |
+
"F": 18,
|
| 22 |
+
"G": 19,
|
| 23 |
+
"H": 20,
|
| 24 |
+
"I": 21,
|
| 25 |
+
"J": 22,
|
| 26 |
+
"K": 23,
|
| 27 |
+
"L": 24,
|
| 28 |
+
"M": 25,
|
| 29 |
+
"N": 26,
|
| 30 |
+
"O": 27,
|
| 31 |
+
"P": 28,
|
| 32 |
+
"Q": 29,
|
| 33 |
+
"R": 30,
|
| 34 |
+
"S": 31,
|
| 35 |
+
"T": 32,
|
| 36 |
+
"U": 33,
|
| 37 |
+
"V": 34,
|
| 38 |
+
"W": 35,
|
| 39 |
+
"X": 36,
|
| 40 |
+
"Y": 37,
|
| 41 |
+
"Z": 38,
|
| 42 |
+
"a": 39,
|
| 43 |
+
"b": 40,
|
| 44 |
+
"c": 41,
|
| 45 |
+
"d": 42,
|
| 46 |
+
"e": 43,
|
| 47 |
+
"f": 44,
|
| 48 |
+
"g": 45,
|
| 49 |
+
"h": 46,
|
| 50 |
+
"i": 47,
|
| 51 |
+
"j": 48,
|
| 52 |
+
"k": 49,
|
| 53 |
+
"l": 50,
|
| 54 |
+
"m": 51,
|
| 55 |
+
"n": 52,
|
| 56 |
+
"o": 53,
|
| 57 |
+
"p": 54,
|
| 58 |
+
"q": 55,
|
| 59 |
+
"r": 56,
|
| 60 |
+
"s": 57,
|
| 61 |
+
"t": 58,
|
| 62 |
+
"u": 59,
|
| 63 |
+
"v": 60,
|
| 64 |
+
"w": 61,
|
| 65 |
+
"x": 62,
|
| 66 |
+
"y": 63,
|
| 67 |
+
"z": 64
|
| 68 |
+
}
|
| 69 |
+
}
|