Visualisation, text processing and PyTorch

import matplotlib.pyplot as plt, numpy as np

x = np.linspace(0, 10, 200)

# Always fig, ax. The implicit plt.plot interface relies on a hidden 'current
# axes' and becomes unmanageable the moment you have subplots.
fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(x, np.sin(x), label="sin", lw=2)
ax.plot(x, np.cos(x), label="cos", ls="--", lw=2)
ax.scatter(x[::20], np.sin(x[::20]), s=30, zorder=3, label="samples")

ax.set(xlabel="x", ylabel="value", title="Two curves")
ax.legend(frameon=False)
ax.grid(alpha=0.3)
ax.spines[["top", "right"]].set_visible(False)   # remove chart junk
fig.tight_layout()
fig.savefig("figure.png", dpi=150, bbox_inches="tight")

import matplotlib.pyplot as plt, numpy as np

fig, axes = plt.subplots(2, 3, figsize=(14, 7), sharex=True, sharey=True)
for i, ax in enumerate(axes.flat):          # .flat iterates the 2-D grid
    ax.plot(np.random.randn(50).cumsum())
    ax.set_title(f"layer {i}")
fig.supxlabel("step"); fig.supylabel("value")
fig.tight_layout()

# Uneven layouts:
fig = plt.figure(figsize=(12, 5))
gs = fig.add_gridspec(2, 3)
big = fig.add_subplot(gs[:, :2])            # spans both rows, first two columns
top = fig.add_subplot(gs[0, 2])
bot = fig.add_subplot(gs[1, 2])

# The standard interpretability heatmap: layers x heads.
fig, ax = plt.subplots(figsize=(8, 6))
limit = np.abs(effects).max()
im = ax.imshow(effects, cmap="RdBu_r", vmin=-limit, vmax=limit, aspect="auto")
ax.set(xlabel="head", ylabel="layer", title="Ablation effect")
fig.colorbar(im, ax=ax, label="Δ logit difference")
# Symmetric limits and a diverging map, so zero is the visual midpoint. With a
# sequential map a reader cannot tell positive from negative effects.

name, score, n = "gpt2", 0.8734, 1234567

print(f"{name}: {score:.2%}")          # gpt2: 87.34%
print(f"{score:.3f} {score:8.3f} {score:<8.3f}|")   # precision and alignment
print(f"{n:,}")                        # 1,234,567
print(f"{n:.2e}")                      # 1.23e+06
print(f"{name!r}")                     # 'gpt2'  — repr, shows quotes
print(f"{score=}")                     # score=0.8734  — debugging shortcut

# Multi-line and aligned table output:
for layer, value in enumerate(profile):
    print(f"layer {layer:>2}  {value:+7.3f}  {'#' * int(abs(value) * 20)}")

# Common string operations
s = "  Tokenization Is Fun!  "
print(s.strip().lower())
print(s.replace("Fun", "Subtle"))
print("a,b,,c".split(","))             # ['a', 'b', '', 'c'] — keeps empties
print(" ".join(["a", "b", "c"]))

# Build strings with join, never with += in a loop: += is O(n^2) because
# strings are immutable and each step copies everything so far.
parts = [f"line {i}" for i in range(1000)]
text = "\n".join(parts)

import re, requests

# A public-domain text — the standard small corpus for these exercises.
url = "https://www.gutenberg.org/files/35/35-0.txt"
raw = requests.get(url, timeout=30).text

# Gutenberg wraps the work in licence headers; strip them or you train on them.
start = raw.find("*** START OF")
end = raw.find("*** END OF")
body = raw[raw.index("\n", start) + 1: end]

def clean(text):
    text = text.replace("\r\n", "\n")
    text = re.sub(r"\n{3,}", "\n\n", text)     # collapse blank runs
    text = re.sub(r"[ \t]+", " ", text)        # collapse spaces, keep newlines
    # Normalise typographic quotes and dashes — they tokenize differently.
    for a, b in [("\u201c", '"'), ("\u201d", '"'), ("\u2018", "'"),
                 ("\u2019", "'"), ("\u2014", "--")]:
        text = text.replace(a, b)
    return text.strip()

body = clean(body)
print(f"{len(body):,} characters, {len(body.split()):,} words")

# Word frequencies
from collections import Counter
words = re.findall(r"[a-z']+", body.lower())
print(Counter(words).most_common(10))

# Encoding: always specify it. The default differs by platform, which is why
# code that works on Linux fails on Windows.
with open("corpus.txt", "w", encoding="utf-8") as f:
    f.write(body)

class Tokenizer:
    """Plain Python class: state plus methods over it."""
    def __init__(self, texts):
        self.vocab = sorted({w for t in texts for w in t.split()})
        self.stoi = {s: i for i, s in enumerate(self.vocab)}
        self.itos = {i: s for s, i in self.stoi.items()}

    def encode(self, text):
        return [self.stoi[w] for w in text.split() if w in self.stoi]

    def decode(self, ids):
        return " ".join(self.itos[i] for i in ids)

    def __len__(self):                    # enables len(tokenizer)
        return len(self.vocab)

    def __repr__(self):                   # what the REPL prints
        return f"Tokenizer(vocab={len(self)})"


import torch.nn as nn

class TinyModel(nn.Module):
    def __init__(self, vocab, dim=64):
        super().__init__()                # MUST come first
        self.emb = nn.Embedding(vocab, dim)     # auto-registered as a submodule
        self.out = nn.Linear(dim, vocab)
        # A tensor that should move with .to(device) but is not trained:
        self.register_buffer("scale", torch.tensor(dim ** -0.5))

    def forward(self, x):                 # called by model(x), never directly
        return self.out(self.emb(x) * self.scale)

model = TinyModel(1000)
print(sum(p.numel() for p in model.parameters()))   # 0 if super() was omitted
print(model)                                        # prints the module tree

import torch

x = torch.tensor([[1., 2., 3.], [4., 5., 6.]])
print(x.shape, x.dtype, x.device)        # torch.Size([2, 3]) torch.float32 cpu

# dtype matters. float32 is the default; integer tensors cannot hold gradients.
torch.zeros(2, 3); torch.ones(4); torch.eye(3)
torch.arange(6); torch.linspace(0, 1, 5)
torch.randn(2, 3)                        # standard normal
torch.randint(0, 10, (2, 3))             # int64 — correct for token ids

# Device: a mismatch is the most common runtime error in PyTorch.
device = "cuda" if torch.cuda.is_available() else "cpu"
x = x.to(device)
model = model.to(device)                 # both, or it fails

# Conversions
x.float(); x.long(); x.half()
x.cpu().numpy()                          # must be on CPU and detached first
torch.from_numpy(np_array)               # shares memory with the array

# Autograd
w = torch.randn(3, requires_grad=True)
loss = (w ** 2).sum()
loss.backward()
print(w.grad)                            # 2w

with torch.no_grad():                    # inference: no graph, much less memory
    y = model(x)

# .item() extracts a Python number from a 1-element tensor. Accumulating
# `total += loss` instead of `total += loss.item()` retains the whole graph
# and leaks memory until the process dies.

import torch

x = torch.arange(24).reshape(2, 3, 4)      # -1 infers a dimension
print(x.shape)

print(x.view(2, 12).shape)                 # shares storage; needs contiguity
print(x.reshape(6, 4).shape)               # copies if it has to
print(x.permute(1, 0, 2).shape)            # (3, 2, 4) — strides only
print(x.transpose(0, 1).shape)             # same for two dimensions

p = x.permute(1, 0, 2)
print(p.is_contiguous())                   # False
# p.view(6, 4)                             # RuntimeError
print(p.contiguous().view(6, 4).shape)     # fine

print(torch.randn(1, 3, 1).squeeze().shape)        # (3,)
print(torch.randn(3).unsqueeze(0).shape)           # (1, 3) — add batch dim

# The pattern from the GPT module, spelled out:
B, T, C, n_heads = 2, 5, 64, 8
q = torch.randn(B, T, C)
q = q.view(B, T, n_heads, C // n_heads).transpose(1, 2)   # [B, heads, T, hd]
out = q.transpose(1, 2).contiguous().view(B, T, C)        # and back

# Combining and splitting
torch.cat([x, x], dim=0).shape             # concatenate along an existing dim
torch.stack([x, x], dim=0).shape           # new dimension
a, b = torch.randn(4, 6).split(3, dim=1)   # how Q, K, V are separated

import torch, numpy as np, random

def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)       # harmless if no GPU

set_seed(0)
print(torch.randn(3))
set_seed(0)
print(torch.randn(3))                      # identical

# A local generator, so one experiment's sampling cannot disturb another's.
g = torch.Generator().manual_seed(123)
print(torch.randn(3, generator=g))

# Sampling functions you will use for generation:
probs = torch.tensor([0.1, 0.7, 0.2])
torch.multinomial(probs, num_samples=1)    # proportional sampling
torch.bernoulli(torch.full((5,), 0.3))     # coin flips — dropout's primitive
torch.randperm(10)                         # shuffling

# Full determinism costs speed and is only sometimes achievable:
# torch.use_deterministic_algorithms(True)
# torch.backends.cudnn.deterministic = True
# Some GPU kernels have no deterministic implementation, so seed-identical
# results are not always available. Report seeds and repeat runs instead of
# assuming bit-identical reproducibility.