mini-transformer/
  data/           # corpus + vocab (created by scripts)
  build_corpus.py
  train.py
  generate.py
  checkpoints/    # saved model weights

# build_corpus.py
from pathlib import Path
import re
 
# Path from mini-transformer/ → repo content/blog
root = Path(__file__).resolve().parent.parent / "content" / "blog"
out = Path(__file__).resolve().parent / "data"
out.mkdir(parents=True, exist_ok=True)
 
texts = []
for path in sorted(root.glob("*.mdx")):
    raw = path.read_text(encoding="utf-8")
 
    # Remove YAML frontmatter between first --- pair
    body = re.sub(r"^---.*?---\s*", "", raw, flags=re.S)
 
    # Strip images ![alt](url) and links [text](url) — keep readable words only
    body = re.sub(r"!\[[^\]]*\]\([^)]+\)", " ", body)
    body = re.sub(r"\[[^\]]+\]\([^)]+\)", " ", body)
 
    # Drop # headings markers but keep words; remove extra whitespace
    body = re.sub(r"#+\s*", "", body)
    body = re.sub(r"\s+", " ", body).strip()
 
    if body:
        texts.append(body)
 
corpus = "\n".join(texts).lower()
(out / "blog_corpus.txt").write_text(corpus, encoding="utf-8")
print(f"Wrote {len(corpus):,} characters from {len(texts)} posts")

# inside train.py later, or a small split helper
split_at = int(len(all_ids) * 0.9)
train_ids = all_ids[:split_at]
val_ids = all_ids[split_at:]

# tokenizer.py (or top of train.py)
import re
import json
from collections import Counter
from pathlib import Path
 
text = Path("data/blog_corpus.txt").read_text(encoding="utf-8")
 
# Split into words and punctuation tokens
tokens = re.findall(r"[a-z0-9]+|[^\s]", text)
counts = Counter(tokens)
 
vocab = {"<pad>": 0, "<unk>": 1}  # reserve IDs for padding and unknown words
for tok, _ in counts.most_common(8000):
    vocab[tok] = len(vocab)
 
id_to_tok = {i: t for t, i in vocab.items()}
 
def encode(s: str) -> list[int]:
    pieces = re.findall(r"[a-z0-9]+|[^\s]", s.lower())
    return [vocab.get(t, vocab["<unk>"]) for t in pieces]
 
def decode(ids: list[int]) -> str:
    return " ".join(id_to_tok[i] for i in ids)
 
Path("data/vocab.json").write_text(json.dumps(vocab), encoding="utf-8")
print("vocab size:", len(vocab))
all_ids = encode(text)

import torch
from torch.utils.data import Dataset
 
class NextTokenDataset(Dataset):
    def __init__(self, ids: list[int], block_size: int = 128):
        self.ids = ids
        self.block_size = block_size
 
    def __len__(self):
        # One sample per starting position where a full window fits
        return max(0, len(self.ids) - self.block_size - 1)
 
    def __getitem__(self, i):
        chunk = self.ids[i : i + self.block_size + 1]
        x = torch.tensor(chunk[:-1], dtype=torch.long)   # input:  positions 0..T-1
        y = torch.tensor(chunk[1:], dtype=torch.long)    # target: positions 1..T
        return x, y

from torch.utils.data import DataLoader
 
train_ds = NextTokenDataset(train_ids, block_size=128)
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)

import torch
import torch.nn as nn
 
class TinyGPT(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        d_model: int = 128,
        nhead: int = 4,
        num_layers: int = 3,
        block_size: int = 128,
    ):
        super().__init__()
        self.block_size = block_size
        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(block_size, d_model)
 
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=512,
            batch_first=True,
        )
        self.blocks = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.lm_head = nn.Linear(d_model, vocab_size)
 
    def forward(self, idx: torch.Tensor) -> torch.Tensor:
        B, T = idx.shape
        if T > self.block_size:
            raise ValueError(f"Sequence length {T} > block_size {self.block_size}")
 
        pos = torch.arange(T, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(pos)
 
        # Causal mask: position i cannot attend to j > i
        mask = nn.Transformer.generate_square_subsequent_mask(T, device=idx.device)
        x = self.blocks(x, mask=mask, is_causal=True)
 
        # Logits: one vocab-sized vector per position
        return self.lm_head(x)

sum(p.numel() for p in model.parameters())

import torch
from torch.utils.data import DataLoader
 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = TinyGPT(vocab_size=len(vocab), block_size=128).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
criterion = torch.nn.CrossEntropyLoss()
 
def evaluate(loader):
    model.eval()
    total_loss = 0.0
    n = 0
    with torch.no_grad():
        for x, y in loader:
            x, y = x.to(device), y.to(device)
            logits = model(x)
            # Flatten (B,T,C) and (B,T) for cross-entropy
            loss = criterion(logits.view(-1, logits.size(-1)), y.view(-1))
            total_loss += loss.item() * x.size(0)
            n += x.size(0)
    model.train()
    return total_loss / max(n, 1)
 
for epoch in range(10):
    model.train()
    running = 0.0
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
 
        optimizer.zero_grad()
        logits = model(x)
        loss = criterion(logits.view(-1, logits.size(-1)), y.view(-1))
        loss.backward()
 
        # Stabilize training on small models / noisy grads
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
 
        running += loss.item() * x.size(0)
 
    train_loss = running / len(train_ds)
    val_loss = evaluate(val_loader)
    print(f"epoch {epoch+1:02d}  train_loss={train_loss:.4f}  val_loss={val_loss:.4f}")

torch.save(model.state_dict(), "checkpoints/tiny_gpt.pt")

import torch
 
@torch.no_grad()
def generate(
    model,
    prompt_ids: list[int],
    id_to_tok: dict,
    max_new: int = 40,
    temperature: float = 0.8,
):
    model.eval()
    ids = list(prompt_ids)
 
    for _ in range(max_new):
        # Only feed the last block_size tokens (context window limit)
        window = ids[-model.block_size :]
        x = torch.tensor([window], device=next(model.parameters()).device)
 
        logits = model(x)[0, -1]  # logits for LAST position only → next token
 
        if temperature <= 0:
            next_id = logits.argmax().item()
        else:
            probs = torch.softmax(logits / temperature, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1).item()
 
        ids.append(next_id)
 
        # Optional early stop at sentence end
        if id_to_tok[next_id] in (".", "!", "?"):
            break
 
    return decode(ids)
 
prompt = "on device ai"
out = generate(model, encode(prompt), id_to_tok, temperature=0.8)
print(out)