Module 20. The complete LLM: from tokens to training and generation

import importlib.util
import subprocess
import sys
missing = [name for name in ("numpy", "torch", "matplotlib") if importlib.util.find_spec(name) is None]
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", *missing],
                          creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
import numpy as np
import torch
import matplotlib
from IPython import get_ipython
ipython = get_ipython()
if ipython is not None and getattr(ipython, "kernel", None) is not None:
    ipython.run_line_magic("matplotlib", "inline")
import matplotlib.pyplot as plt
torch.set_num_threads(1)
torch.manual_seed(20)
np.set_printoptions(precision=6, suppress=True)
print("NumPy", np.__version__, "PyTorch", torch.__version__, "device: CPU")

"""M20: one complete pre-norm decoder block, with inspectable intermediates.

Rows are token representations; linear weights have shape [input, output].
This synthetic fixture illustrates arithmetic, not language competence.
"""
import numpy as np

EPS = 1e-6
INPUT = np.array([0, 1, 2], dtype=np.int64)
TARGET = np.array([1, 2, 3], dtype=np.int64)
TOKENS = ("BOS", "A", "B", "EOS")


def fixture():
    """Return fresh arrays: vocabulary 4, width 2, one head, FFN width 2."""
    return {
        "E": np.array([[1., 0.], [0., 1.], [1., 1.], [-1., 0.]]),
        "Q": .5 * np.eye(2), "K": .5 * np.eye(2),
        "V": np.diag([.5, 1.]), "O": .25 * np.eye(2),
        "gate": np.array([[.5, -.25], [.25, .5]]),
        "up": np.array([[1., .5], [-.5, 1.]]),
        "down": .25 * np.eye(2),
        "na": np.ones(2), "nf": np.ones(2), "nout": np.ones(2),
        "U": np.array([[0., .5, -.5, .25], [0., -.25, .5, -.5]]),
    }


def rms(x, gain):
    return x / np.sqrt(np.mean(x * x, axis=-1, keepdims=True) + EPS) * gain


def softmax(x):
    e = np.exp(x - np.max(x, axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)


def silu(x):
    return x / (1 + np.exp(-x))


def rope(x, positions):
    """[heads, time, even head_dim]; pair first and second halves."""
    dh = x.shape[-1]
    if dh % 2:
        raise ValueError("RoPE requires an even head dimension")
    angle = np.asarray(positions)[:, None] * 10000. ** (-np.arange(0, dh, 2) / dh)
    angle = np.concatenate([angle, angle], axis=-1)
    rotated = np.concatenate([-x[..., dh // 2:], x[..., :dh // 2]], axis=-1)
    return x * np.cos(angle) + rotated * np.sin(angle)


def forward(w, ids, hq=1, hkv=1, cache=None):
    """A single layer; cache contains already rotated K and unrotated V.

    Cache is valid only for these weights and this uninterrupted prefix.
    Tied embeddings are selected by omitting U from w.
    """
    ids = np.asarray(ids, dtype=np.int64)
    if ids.ndim != 1 or len(ids) == 0:
        raise ValueError("ids must be a nonempty one-dimensional sequence")
    if np.any(ids < 0) or np.any(ids >= len(w["E"])):
        raise ValueError("token id outside the vocabulary")
    h = w["E"][ids]
    t, d = h.shape
    if hq < 1 or hkv < 1 or hq % hkv or d % hq:
        raise ValueError("head counts must divide the query count and model width")
    dh = d // hq
    offset = 0 if cache is None else cache[0].shape[1]
    positions = np.arange(offset, offset + t)
    x = rms(h, w["na"])
    q = rope((x @ w["Q"]).reshape(t, hq, dh).transpose(1, 0, 2), positions)
    k = rope((x @ w["K"]).reshape(t, hkv, dh).transpose(1, 0, 2), positions)
    v = (x @ w["V"]).reshape(t, hkv, dh).transpose(1, 0, 2)
    if cache is not None:
        k = np.concatenate([cache[0], k], axis=1)
        v = np.concatenate([cache[1], v], axis=1)
    kr, vr = (np.repeat(a, hq // hkv, axis=0) for a in (k, v))
    scores = q @ kr.transpose(0, 2, 1) / np.sqrt(dh)
    mask = np.arange(k.shape[1])[None, :] <= positions[:, None]
    a = softmax(np.where(mask[None], scores, -np.inf))
    mixed = (a @ vr).transpose(1, 0, 2).reshape(t, d)
    z = h + mixed @ w["O"]
    n = rms(z, w["nf"])
    gate, up = n @ w["gate"], n @ w["up"]
    f = (silu(gate) * up) @ w["down"]
    h1 = z + f
    final = rms(h1, w["nout"])
    logits = final @ (w["U"] if "U" in w else w["E"].T)
    return dict(H0=h, X=x, Q=q, K=k, V=v, scores=scores, A=a,
                mixed=mixed, Z=z, N=n, gate=gate, up=up, F=f, H1=h1,
                final=final, logits=logits, probabilities=softmax(logits), cache=(k, v))


def cross_entropy(logits, targets, mask=None):
    """Mean over explicitly selected targets; never mean over padding."""
    targets = np.asarray(targets, dtype=np.int64)
    if logits.ndim != 2 or targets.shape != (len(logits),):
        raise ValueError("one target is required for every logit row")
    selected = np.ones(len(targets), dtype=bool) if mask is None else np.asarray(mask, dtype=bool)
    if selected.shape != targets.shape or not selected.any():
        raise ValueError("the target mask must select at least one row")
    if np.any(targets[selected] < 0) or np.any(targets[selected] >= logits.shape[1]):
        raise ValueError("selected target outside the vocabulary")
    x = logits[selected]
    m = x.max(axis=-1)
    return np.mean(m + np.log(np.exp(x - m[:, None]).sum(axis=-1))
                   - x[np.arange(len(x)), targets[selected]])


def finite_gradients(w, ids=INPUT, targets=TARGET, step=1e-5):
    """Central differences for every scalar; restore each weight after use."""
    result = {}
    for name, value in w.items():
        g = np.empty_like(value)
        for index in np.ndindex(value.shape):
            old = value[index]
            try:
                value[index] = old + step
                plus = cross_entropy(forward(w, ids)["logits"], targets)
                value[index] = old - step
                minus = cross_entropy(forward(w, ids)["logits"], targets)
            finally:
                value[index] = old
            g[index] = (plus - minus) / (2 * step)
        result[name] = g
    return result


def adamw_first_step(w, grads, lr=.01, decay=.1, clip=1., eps=1e-8):
    """Zero initial moments, bias correction: m_hat=g, v_hat=g^2.

    Only matrices receive weight decay. Clipping precedes moment updates.
    The first bias-corrected step is independent of beta1 and beta2.
    """
    norm = np.sqrt(sum(np.sum(g * g) for g in grads.values()))
    scale = min(1., clip / max(norm, 1e-300))
    updated = {}
    for name, value in w.items():
        g = grads[name] * scale
        wd = decay if value.ndim == 2 else 0.
        updated[name] = (1 - lr * wd) * value - lr * g / (np.abs(g) + eps)
    return updated, float(norm), float(scale)


def torch_forward(w, ids, hq=1, hkv=1):
    """Independent float64 autograd path; PyTorch is imported only when used."""
    import torch
    import torch.nn.functional as F
    ids = torch.as_tensor(ids, dtype=torch.long)
    h = w["E"][ids]
    t, d = h.shape
    dh = d // hq
    def norm(a, gain):
        return a * torch.rsqrt(a.square().mean(-1, keepdim=True) + EPS) * gain
    def rotate(a):
        freq = 10000. ** (-torch.arange(0, dh, 2, dtype=h.dtype) / dh)
        angle = torch.arange(t, dtype=h.dtype)[:, None] * freq
        angle = torch.cat([angle, angle], dim=-1)
        return a * angle.cos() + torch.cat([-a[..., dh//2:], a[..., :dh//2]], -1) * angle.sin()
    n = norm(h, w["na"])
    q = rotate((n @ w["Q"]).view(t, hq, dh).transpose(0, 1))
    k = rotate((n @ w["K"]).view(t, hkv, dh).transpose(0, 1))
    v = (n @ w["V"]).view(t, hkv, dh).transpose(0, 1)
    k, v = [a.repeat_interleave(hq // hkv, dim=0) for a in (k, v)]
    scores = q @ k.transpose(-2, -1) / dh ** .5
    scores = scores.masked_fill(torch.ones(t, t, dtype=torch.bool).triu(1), -torch.inf)
    mixed = (scores.softmax(-1) @ v).transpose(0, 1).reshape(t, d)
    z = h + mixed @ w["O"]
    n = norm(z, w["nf"])
    h1 = z + (F.silu(n @ w["gate"]) * (n @ w["up"])) @ w["down"]
    return norm(h1, w["nout"]) @ (w["U"] if "U" in w else w["E"].T)


def autograd_gradients(w, ids=INPUT, targets=TARGET):
    import torch
    params = {k: torch.tensor(v, dtype=torch.float64, requires_grad=True) for k, v in w.items()}
    logits = torch_forward(params, ids)
    loss = torch.nn.functional.cross_entropy(logits, torch.tensor(targets, dtype=torch.long))
    loss.backward()
    return {k: v.grad.detach().numpy().copy() for k, v in params.items()}


def sampling_distribution(logits, temperature=1., top_k=None):
    """One row; stable token-index tie rule. Distribution q, not base p."""
    if temperature <= 0 or np.asarray(logits).ndim != 1:
        raise ValueError("use a positive temperature and one logit row")
    z = np.asarray(logits) / temperature
    if top_k is not None:
        if not isinstance(top_k, int) or not 1 <= top_k <= len(z):
            raise ValueError("top_k must be an integer in [1, vocabulary_size]")
        keep = np.argsort(-z, kind="stable")[:top_k]
        z = np.where(np.isin(np.arange(len(z)), keep), z, -np.inf)
    return softmax(z)


def moe_weights(logits, selection_bias, k=2, eps=1e-9):
    """MiniLLM's sigmoid routing: bias changes selection, not mixture weights."""
    p = 1 / (1 + np.exp(-np.asarray(logits)))
    selected = np.argsort(-(p + selection_bias), kind="stable")[:k]
    return selected, p[selected] / (p[selected].sum() + eps)


def architecture_examples():
    """Small exact SSM recurrence and a two-step discrete reverse chain."""
    state, states = 0., []
    for u in (1., 0., -1.):
        state = .5 * state + u
        states.append(state)
    # Four binary sequences: 00, 01, 10, 11. At s=2 both positions are masked.
    first = np.array([.6*.3, .6*.7, .4*.3, .4*.7])
    # At the next refinement each bit is independently kept with probability .8.
    bits = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])
    distance = (bits[:, None, :] != bits[None, :, :]).sum(-1)
    kernel = .8 ** (2-distance) * .2 ** distance
    return {"ssm_states": states, "diffusion_first": first.tolist(),
            "diffusion_kernel": kernel.tolist(), "diffusion_final": (first @ kernel).tolist()}


def reference_metrics():
    w = fixture()
    trace = forward(w, INPUT)
    grads = autograd_gradients(w)
    changed, norm, scale = adamw_first_step(w, grads)
    out = {k: v.tolist() for k, v in trace.items() if isinstance(v, np.ndarray)}
    out.update(parameter_count=sum(v.size for v in w.values()),
               loss=float(cross_entropy(trace["logits"], TARGET)),
               target_losses=(-np.log(trace["probabilities"][np.arange(3), TARGET])).tolist(),
               gradient_norm=norm, clip_scale=scale,
               gradients={k: v.tolist() for k, v in grads.items()},
               loss_after_step=float(cross_entropy(forward(changed, INPUT)["logits"], TARGET)),
               probabilities_after_step=forward(changed, INPUT)["probabilities"].tolist(),
               q_temperature_half_top2=sampling_distribution(trace["logits"][-1], .5, 2).tolist(),
               architecture_examples=architecture_examples())
    return out

baseline = fixture()
trace = forward(baseline, INPUT)
assert sum(v.size for v in baseline.values()) == 50
for name in ("H0", "X", "Q", "K", "V", "A", "Z", "N", "gate", "up", "F", "H1", "final"):
    print(name, trace[name].shape, "\n", trace[name])

np.testing.assert_allclose(trace["A"].sum(-1), 1.)
assert np.all(np.triu(trace["A"][0], 1) == 0)
print("logits\n", trace["logits"])
print("p\n", trace["probabilities"])
fig, ax = plt.subplots(figsize=(6, 4))
ax.imshow(trace["A"][0], vmin=0, vmax=1, cmap="Blues")
for i in range(3):
    for j in range(3):
        ax.text(j, i, f"{trace['A'][0,i,j]:.3f}", ha="center", va="center",
                color="white" if trace["A"][0,i,j] > .6 else "#2E3440")
ax.set(xticks=range(3), yticks=range(3), xlabel="Key position", ylabel="Query position", title="Causal attention")
fig.tight_layout()
plt.show()

loss_before = cross_entropy(trace["logits"], TARGET)
loss_eos = cross_entropy(trace["logits"], TARGET, [False, False, True])
np.testing.assert_allclose(loss_before, 1.1093739242148548, atol=1e-10)
np.testing.assert_allclose(loss_eos, 1.6955297196508368, atol=1e-10)
print("all targets:", loss_before, "EOS only:", loss_eos)
try:
    cross_entropy(trace["logits"], TARGET, [False]*3)
except ValueError as error:
    print(type(error).__name__, str(error))
else:
    raise AssertionError("empty target mask was accepted")

gradients = autograd_gradients(baseline)
finite = finite_gradients(baseline)
max_error = max(np.max(np.abs(gradients[k] - finite[k])) for k in baseline)
assert max_error < 2e-8
print("50 parameters; max gradient error:", max_error)
tied = fixture()
del tied["U"]
tied_auto, tied_fd = autograd_gradients(tied), finite_gradients(tied)
tied_error = max(np.max(np.abs(tied_auto[k] - tied_fd[k])) for k in tied)
assert tied_error < 2e-8
print("42 tied parameters; max gradient error:", tied_error)
logit_gradient = trace["probabilities"].copy()
logit_gradient[np.arange(len(TARGET)), TARGET] -= 1
logit_gradient /= len(TARGET)
np.testing.assert_allclose(logit_gradient.sum(-1), 0., atol=1e-15)
print("gradient of loss w.r.t. logits:\n", logit_gradient)

updated, gradient_norm, clip_scale = adamw_first_step(baseline, gradients)
params = {k: torch.tensor(v, dtype=torch.float64, requires_grad=True) for k, v in baseline.items()}
optimizer = torch.optim.AdamW([
    {"params": [p for p in params.values() if p.ndim == 2], "weight_decay": .1},
    {"params": [p for p in params.values() if p.ndim == 1], "weight_decay": 0.}],
    lr=.01, betas=(.9, .95), eps=1e-8)
loss_torch = torch.nn.functional.cross_entropy(torch_forward(params, INPUT), torch.tensor(TARGET))
optimizer.zero_grad()
loss_torch.backward()
for p in params.values():
    p.grad.mul_(clip_scale)
optimizer.step()
for k, p in params.items():
    np.testing.assert_allclose(p.detach().numpy(), updated[k], atol=1e-10)
loss_after = cross_entropy(forward(updated, INPUT)["logits"], TARGET)
np.testing.assert_allclose(loss_after, 1.0848315484318878, atol=1e-10)
print("gradient norm:", gradient_norm, "scale:", clip_scale)
print("loss:", loss_before, "->", loss_after)

training_params = {k: torch.tensor(v, dtype=torch.float64, requires_grad=True) for k, v in fixture().items()}
training_optimizer = torch.optim.AdamW(list(training_params.values()), lr=.03, weight_decay=0.)
training_losses = []
for step in range(80):
    training_optimizer.zero_grad()
    loss = torch.nn.functional.cross_entropy(torch_forward(training_params, INPUT), torch.tensor(TARGET))
    training_losses.append(loss.item())
    loss.backward()
    torch.nn.utils.clip_grad_norm_(list(training_params.values()), 1.)
    training_optimizer.step()
with torch.no_grad():
    final_training_loss = torch.nn.functional.cross_entropy(torch_forward(training_params, INPUT), torch.tensor(TARGET)).item()
training_losses.append(final_training_loss)
assert np.isfinite(training_losses).all() and final_training_loss < training_losses[0]
fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(training_losses, color="#5E81AC")
ax.set(xlabel="Training step", ylabel="Mean loss, nats", title="Fitting one sequence")
fig.tight_layout()
plt.show()
print("training loss:", training_losses[0], "->", final_training_loss)

changed_input = forward(baseline, [0, 1, 3])
np.testing.assert_allclose(changed_input["logits"][:2], trace["logits"][:2])
cache, cached_logits = None, []
for token in INPUT:
    current = forward(baseline, [token], cache=cache)
    cached_logits.append(current["logits"])
    cache = current["cache"]
np.testing.assert_allclose(np.concatenate(cached_logits), trace["logits"], atol=1e-12)
prefix = forward(baseline, INPUT[:1])
suffix = forward(baseline, INPUT[1:], cache=prefix["cache"])
np.testing.assert_allclose(suffix["logits"], trace["logits"][1:], atol=1e-12)
wrong_q = rope((trace["X"][2:] @ baseline["Q"])[None], [0])
wrong_scores = wrong_q @ trace["K"].transpose(0, 2, 1) / np.sqrt(2)
offset_error = np.max(np.abs(wrong_scores[0, 0] - trace["scores"][0, 2]))
assert offset_error > .1
print("cache scalars:", sum(a.size for a in cache), "wrong-offset score difference:", offset_error)

q = sampling_distribution(trace["logits"][-1], temperature=.5, top_k=2)
np.testing.assert_allclose(q, [0., .5631163013756344, .4368836986243656, 0.])
print("p:", trace["probabilities"][-1], "q:", q)
trained = {k: p.detach().numpy().copy() for k, p in training_params.items()}
generated = [0]
generation_cache = None
for _ in range(12):
    output = forward(trained, [generated[-1]], cache=generation_cache)
    generation_cache = output["cache"]
    next_token = int(np.argmax(output["logits"][-1]))
    generated.append(next_token)
    if next_token == 3:
        break
print(" ".join(TOKENS[i] for i in generated))

experiment_temperature = 1.0
experiment_top_k = 2
if not np.isfinite(experiment_temperature) or experiment_temperature <= 0:
    raise ValueError("Temperature must be positive and finite")
experiment_q = sampling_distribution(
    trace["logits"][-1], temperature=experiment_temperature, top_k=experiment_top_k)
assert np.isfinite(experiment_q).all() and (experiment_q >= 0).all()
np.testing.assert_allclose(experiment_q.sum(), 1., rtol=1e-12, atol=1e-12)
print("Temperature:", experiment_temperature, "top-k:", experiment_top_k)
print("Token   model p    selection q")
for token, model_p, selection_q in zip(TOKENS, trace["probabilities"][-1], experiment_q):
    print(f"{token:<5}   {model_p:.6f}   {selection_q:.6f}")

selected, routing_weights = moe_weights([0., 1., -1.], np.array([0., -.5, .5]), k=2)
assert selected.tolist() == [2, 0]
np.testing.assert_allclose(routing_weights, [.3497554090542189, .6502445909457811])
print("selected:", selected, "weights:", routing_weights)

examples = architecture_examples()
np.testing.assert_allclose(examples["ssm_states"], [1., .5, -.75])
kernel = np.array(examples["diffusion_kernel"])
first = np.array(examples["diffusion_first"])
np.testing.assert_allclose(kernel.sum(-1), 1.)
final = first @ kernel
np.testing.assert_allclose(final, [.2128, .3472, .1672, .2728])
manual_00 = .18*.64 + .42*.16 + .12*.16 + .28*.04
np.testing.assert_allclose(manual_00, final[0])
print("SSM:", examples["ssm_states"])
print("reverse chain:", first, "->", final, "sum:", final.sum())

experiment_ids = np.array([0, 2, 1], dtype=np.int64)
experiment_targets = np.array([2, 1, 3], dtype=np.int64)
experiment_weights = fixture()
experiment_trace = forward(experiment_weights, experiment_ids)
experiment_loss = cross_entropy(experiment_trace["logits"], experiment_targets)
experiment_gradients = autograd_gradients(experiment_weights, experiment_ids, experiment_targets)
print("inputs:", experiment_ids, "targets:", experiment_targets)
print("p:\n", experiment_trace["probabilities"])
print("loss:", experiment_loss)
print("gradient norm:", np.sqrt(sum(np.sum(g*g) for g in experiment_gradients.values())))

import importlib
import sys
import subprocess
import zipfile
from pathlib import Path

ARCHIVE_GLOB = "modern-llms-capstone-*.zip"

def _activate_capstone_root(candidate: Path):
    package_init = candidate / "minillm" / "__init__.py"
    if not package_init.exists():
        return None
    root = str(candidate.resolve())
    if root not in sys.path:
        sys.path.insert(0, root)
    return root

def _find_capstone_root(search_root: Path):
    direct = _activate_capstone_root(search_root)
    if direct:
        return direct
    for package_root in (search_root / "capstone" / "minillm" / "en", search_root / "minillm" / "en"):
        found = _activate_capstone_root(package_root)
        if found:
            return found
    for child in sorted(search_root.glob("modern-llms-capstone-*")):
        found = _activate_capstone_root(child)
        if found:
            return found
    return None

def _safe_extract(archive_path: Path, destination: Path):
    destination = destination.resolve()
    with zipfile.ZipFile(archive_path) as archive:
        for member in archive.infolist():
            member_path = Path(member.filename)
            if member_path.is_absolute() or ".." in member_path.parts:
                raise RuntimeError(f"unsafe path in archive: {member.filename}")
        archive.extractall(destination)

def ensure_minillm():
    for candidate in (Path.cwd(), *Path.cwd().parents):
        found = _find_capstone_root(candidate)
        if found:
            return "project package: " + found

    try:
        import minillm  # noqa: F401
        return "minillm is already importable"
    except ImportError:
        pass

    archives = sorted(Path.cwd().glob(ARCHIVE_GLOB)) + sorted(Path.cwd().glob("course-materials.zip"))
    if not archives:
        try:
            from google.colab import files
        except ImportError:
            files = None
        if files is not None:
            print("Upload course-materials.zip from the course notebook page. The directory structure will be restored automatically.")
            files.upload()
            archives = sorted(Path.cwd().glob(ARCHIVE_GLOB)) + sorted(Path.cwd().glob("course-materials.zip"))

    if archives:
        archive_path = archives[-1]
        destination = Path.cwd() / ".capstone_runtime"
        destination.mkdir(exist_ok=True)
        _safe_extract(archive_path, destination)
        root = _find_capstone_root(destination)
        if root:
            return "project archive: " + str(archive_path)
        matches = sorted(destination.glob("*/minillm/__init__.py"))
        if len(matches) != 1:
            raise RuntimeError(
                f"expected one minillm package in {archive_path.name}, "
                f"found {len(matches)}"
            )
        root = _activate_capstone_root(matches[0].parent.parent)
        if root:
            return "project archive: " + str(archive_path)

    raise RuntimeError("minillm was not found. Download course-materials.zip from the course notebook page and either extract it in full or upload the ZIP to the current Colab session.")

print(ensure_minillm())

for package in ("torch", "numpy", "matplotlib", "regex"):
    try:
        importlib.import_module(package)
    except ImportError:
        subprocess.run(
            [sys.executable, "-m", "pip", "install", "-q", package],
            check=True,
        )

import os, torch, numpy as np, hashlib
print("torch", torch.__version__, "| CUDA:", torch.cuda.is_available())

# ------------------------------------------------------------------ run profile
PROFILE = "SMOKE"          # "SMOKE" | "FULL"
USE_REAL_CORPUS = False    # network download of real data — see §3
SEED = 0

torch.manual_seed(SEED)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
from minillm.precision import resolve_compute_dtype
COMPUTE_DTYPE = resolve_compute_dtype(DEVICE, "auto")

SETTINGS = {
    "SMOKE": dict(docs=8_000, vocab_size=2_048, seq_len=128,
                  model_profile="S", steps=250, micro_batch_size=8, grad_accum=1),
    "FULL":  dict(docs=200_000, vocab_size=16_384, seq_len=512,
                  model_profile="M", steps=18_886, micro_batch_size=16, grad_accum=4),
}[PROFILE]
if os.getenv("CAPSTONE_FAST_VALIDATION") == "1":
    SETTINGS = {**SETTINGS, "docs": min(SETTINGS["docs"], 800),
                "vocab_size": min(SETTINGS["vocab_size"], 512),
                "steps": min(SETTINGS["steps"], 12),
                "micro_batch_size": min(SETTINGS["micro_batch_size"], 2),
                "grad_accum": 1}
print(PROFILE, "on", DEVICE, "with", COMPUTE_DTYPE, "compute ->", SETTINGS)

import random
from minillm.corpus import FilterConfig, pipeline

RU_SLOTS = dict(
    subject=["языковая модель", "трансформер", "механизм внимания", "токенизатор",
             "оптимизатор", "нормализация", "остаточная связь", "роутер экспертов",
             "кэш ключей", "эмбеддинг", "спекулятивный декод", "квантизация"],
    verb=["предсказывает", "обрабатывает", "сокращает", "ускоряет", "стабилизирует",
          "ограничивает", "восстанавливает", "перераспределяет", "измеряет"],
    obj=["следующий токен", "длинный контекст", "стоимость шага", "объём памяти",
         "масштаб активаций", "разброс градиентов", "поток данных", "качество ответа",
         "нагрузку экспертов", "пропускную способность"],
    tail=["на длинных последовательностях", "при обучении в низкой точности",
          "без потери качества", "за счёт дополнительной матрицы",
          "ценой роста числа параметров", "в фазе декодирования",
          "по сравнению с базовой схемой", "при неизменном бюджете вычислений"],
)
EN_SLOTS = dict(
    subject=["a language model", "the transformer", "the attention mechanism",
             "the tokenizer", "the optimizer", "normalization", "the residual stream",
             "the expert router", "the key cache", "the embedding table",
             "speculative decoding", "quantization"],
    verb=["predicts", "processes", "reduces", "accelerates", "stabilizes", "bounds",
          "restores", "redistributes", "measures"],
    obj=["the next token", "a long context", "the cost per step", "the memory footprint",
         "the scale of activations", "the spread of gradients", "the data stream",
         "the quality of the answer", "the load across experts", "the throughput"],
    tail=["on long sequences", "when training in low precision", "without losing quality",
          "at the price of one more matrix", "at the cost of more parameters",
          "during the decode phase", "compared with the baseline scheme",
          "under an unchanged compute budget"],
)


def sentence(rng: random.Random, slots: dict) -> str:
    parts = [rng.choice(slots["subject"]), rng.choice(slots["verb"]),
             rng.choice(slots["obj"]), rng.choice(slots["tail"])]
    return " ".join(parts).capitalize() + "."


def synthetic_corpus(n_docs: int, seed: int = 0) -> list[str]:
    """A bilingual corpus from a slot grammar, with junk mixed in on purpose.

    This is no pretraining corpus: one grammar, a narrow lexicon. Its job is to
    give the pipeline and the tokenizer reproducible, unambiguously licence-clean
    material where every mechanism is visible. The slots supply combinatorial
    variety — without it near-dup dedup would rightly collapse the whole corpus.
    """
    rng = random.Random(seed)
    docs = []
    for i in range(n_docs):
        slots = RU_SLOTS if i % 2 else EN_SLOTS
        docs.append(" ".join(sentence(rng, slots) for _ in range(rng.randint(3, 6))))
    # Junk for the funnel: exact copies, near-copies, short and symbol-only.
    docs += docs[:n_docs // 20]                                   # exact duplicates
    docs += [d + " Addendum." for d in docs[:n_docs // 40]]       # near-duplicates
    docs += ["ok"] * (n_docs // 40)                               # too short
    docs += ["!@#$%^&*()" * 8] * (n_docs // 40)                   # symbol junk
    rng.shuffle(docs)
    return docs

raw_docs = synthetic_corpus(SETTINGS["docs"], seed=SEED)
print(f"raw documents: {len(raw_docs):,}")

clean_docs, report = pipeline(raw_docs, FilterConfig(min_chars=60),
                              near_dup_threshold=0.85)
print(report)
print(f"\nresult: {len(clean_docs):,} documents, "
      f"{sum(len(d) for d in clean_docs):,} characters")

# Split documents BEFORE tokenizer training: validation contributes neither
# vocabulary entries nor BPE merges. Global cleaning happened before the split,
# so exact and near duplicates cannot straddle train and validation.
split_rng = random.Random(SEED + 17)
split_rng.shuffle(clean_docs)
valid_count = max(1, round(0.05 * len(clean_docs)))
valid_docs = clean_docs[:valid_count]
train_docs = clean_docs[valid_count:]
print(f"document split: train={len(train_docs):,}, valid={len(valid_docs):,}")

if USE_REAL_CORPUS:
    subprocess.run([sys.executable, "-m", "pip", "install", "-q", "datasets"], check=True)
    from datasets import load_dataset

    DATA_REVISIONS = {
        "fineweb2": "v2.1.1",
        "tinystories": "f54c09fd23315a6f9c86f9dc80f725de7d8f9c64",
    }
    train_target = SETTINGS["docs"] // 2
    eval_target = max(100, SETTINGS["docs"] // 40)

    ru_train = load_dataset("HuggingFaceFW/fineweb-2", name="rus_Cyrl",
                            split="train", streaming=True,
                            revision=DATA_REVISIONS["fineweb2"])
    ru_eval = load_dataset("HuggingFaceFW/fineweb-2", name="rus_Cyrl",
                           split="test", streaming=True,
                           revision=DATA_REVISIONS["fineweb2"])
    en_train = load_dataset("roneneldan/TinyStories", split="train", streaming=True,
                            revision=DATA_REVISIONS["tinystories"])
    en_eval = load_dataset("roneneldan/TinyStories", split="validation", streaming=True,
                           revision=DATA_REVISIONS["tinystories"])

    def take_text(dataset, limit):
        rows = []
        for row in dataset:
            rows.append(row["text"])
            if len(rows) >= limit:
                break
        return rows

    raw_train = take_text(ru_train, train_target) + take_text(en_train, train_target)
    raw_valid = take_text(ru_eval, eval_target) + take_text(en_eval, eval_target)
    train_docs, train_report = pipeline(raw_train, FilterConfig(), near_dup_threshold=0.85)
    valid_docs, valid_report = pipeline(raw_valid, FilterConfig(), near_dup_threshold=0.85)
    train_fingerprints = {hashlib.sha256(d.encode("utf-8")).digest() for d in train_docs}
    valid_docs = [d for d in valid_docs
                  if hashlib.sha256(d.encode("utf-8")).digest() not in train_fingerprints]
    print("train:", train_report)
    print("validation:", valid_report)
    print("revisions:", DATA_REVISIONS)
    print("\nAttribution: FineWeb2 (ODC-By 1.0 + Common Crawl Terms), "
          "TinyStories (CDLA-Sharing-1.0). This is technical documentation, not legal advice.")
else:
    DATA_REVISIONS = {"synthetic": f"generator-v1-seed-{SEED}"}
    print("USE_REAL_CORPUS = False — using the synthetic document-level split")

import time
from minillm.tokenizer import BBPETokenizer
from minillm.data import ByteTokenizer
from minillm.chat import CHAT_SPECIALS

# The chat control tokens are reserved in the vocabulary right away, even though
# pretraining does not need them. The reason is exactly the one behind the 1,075
# reserved identifiers in o200k_harmony (M2): adding them later means changing
# vocab_size and rebuilding the embeddings. Notebook 2 picks up this very
# checkpoint and expects to find `<|eot_id|>` inside it.
train_text = "\n".join(train_docs)
t0 = time.time()
tokenizer = BBPETokenizer.train(train_text, vocab_size=SETTINGS["vocab_size"],
                                specials=["<|bos|>", "<|eos|>"] + CHAT_SPECIALS)
print(f"trained in {time.time() - t0:.1f} s: vocabulary {tokenizer.vocab_size}, "
      f"{len(tokenizer.merges)} merges")
if tokenizer.vocab_size < SETTINGS["vocab_size"]:
    print('Available merge pairs ran out before the requested vocabulary size.')

sample = "The tokenizer predicts nothing; it only decides the price."
assert tokenizer.decode(tokenizer.encode(sample)) == sample, "round trip is broken"
print("round trip:", repr(tokenizer.decode(tokenizer.encode(sample))))

probe_rng = random.Random(123)
ru_probe = " ".join(sentence(probe_rng, RU_SLOTS) for _ in range(40))
en_probe = " ".join(sentence(probe_rng, EN_SLOTS) for _ in range(40))
plain = ByteTokenizer()

rows = []
for name, probe in (("RU", ru_probe), ("EN", en_probe)):
    trained = tokenizer.fertility(probe)
    baseline = len(plain.encode(probe)) / len(probe)
    rows.append((name, trained["tokens_per_char"], baseline,
                 trained["bytes_per_token"]))

print(f"{'lang':<5} {'BBPE tok/char':>14} {'byte-level tok/char':>21} {'bytes/token':>12}")
for name, trained, baseline, bpt in rows:
    print(f"{name:<5} {trained:>14.3f} {baseline:>21.3f} {bpt:>12.2f}")
print("\nByte-level costs about two tokens per Cyrillic character — that is the "
      "price of having no merges for a language (M2).")

from minillm.data import ShardStream, TokenStream, write_shards

EOS = tokenizer.specials["<|eos|>"]

def encode_documents(documents):
    stream = []
    for doc in documents:
        stream.extend(tokenizer.encode(doc))
        stream.append(EOS)
    return stream

train_token_ids = encode_documents(train_docs)
valid_token_ids = encode_documents(valid_docs)
print(f"train tokens: {len(train_token_ids):,}; validation tokens: {len(valid_token_ids):,}")

shard_dir = Path("data_shards")
index = write_shards(train_token_ids, shard_dir / "train", shard_tokens=250_000)
write_shards(valid_token_ids, shard_dir / "valid", shard_tokens=250_000)

train_stream = ShardStream(shard_dir / "train", seq_len=SETTINGS["seq_len"], verify=True)
valid_stream = ShardStream(shard_dir / "valid", seq_len=SETTINGS["seq_len"], verify=True)
print(f"train: {len(index['shards'])} shards, {index['total_tokens']:,} tokens, "
      f"{len(train_stream):,} windows, checksum {train_stream.checksum()}")
print(f"validation: {len(valid_stream):,} windows — document-separated before tokenizer training")

from minillm.config import MiniLLMConfig, kv_bytes_per_token
from minillm.model import MiniLLM, MoELayer
from minillm.train import TrainConfig, memory_report

cfg = MiniLLMConfig.profile(
    SETTINGS["model_profile"],
    vocab_size=tokenizer.vocab_size,
    max_position_embeddings=max(SETTINGS["seq_len"], 256),
)
model = MiniLLM(cfg).to(DEVICE)

analytic = cfg.parameter_report()
assert analytic["total"] == model.num_parameters(), "formula and code disagree"
print(f"parameters: {analytic['total']:,} (formula == actual)")
for key in ("embeddings", "attention", "ffn", "norms"):
    print(f"  {key:<12} {analytic[key]:>12,}  ({analytic[key] / analytic['total']:.1%})")
print(f"  the block costs {analytic['block_over_d2']}·d²")
reference_cache_bpt = kv_bytes_per_token(cfg, dtype_bytes=2)
runtime_cache_bpt = kv_bytes_per_token(cfg, dtype_bytes={"fp32": 4, "bf16": 2, "fp16": 2}[COMPUTE_DTYPE])
print(f"\nKV cache, BF16/FP16 reference: {reference_cache_bpt} bytes/token "
      f"→ {reference_cache_bpt * SETTINGS['seq_len'] / 1024:.1f} KiB per window")
print(f"KV cache at the current compute dtype ({COMPUTE_DTYPE}): {runtime_cache_bpt} bytes/token")
print("persistent training memory:", memory_report(cfg, "adamw"))

tcfg = TrainConfig(
    steps=SETTINGS["steps"],
    seq_len=SETTINGS["seq_len"],
    micro_batch_size=SETTINGS["micro_batch_size"],
    grad_accum=SETTINGS["grad_accum"],
    optimizer="muon",
    compute_dtype="auto",
    gradient_checkpointing=DEVICE.startswith("cuda"),
    seed=SEED,
    device=DEVICE,
    log_every=max(1, SETTINGS["steps"] // 10),
    out_dir="runs/pretrain",
)

chinchilla_tokens = 20 * analytic["total"]
planned_tokens = tcfg.steps * tcfg.tokens_per_step
print(f"planned next-token targets per step: {tcfg.tokens_per_step:,}")
print(f"planned:         {planned_tokens:,} supervised targets over {tcfg.steps:,} steps")
print(f"reference 20×: {chinchilla_tokens:,} tokens "
      f"→ {chinchilla_tokens / tcfg.tokens_per_step:,.0f} steps")
print(f"share of budget: {planned_tokens / chinchilla_tokens:.1%}")
print(f"epochs over the corpus: {planned_tokens / len(train_token_ids):.2f}")

from minillm.train import train, save_checkpoint

state = train(model, train_stream, tcfg, eval_stream=valid_stream)
print(f"\ndone: {state.step} steps, "
      f"{state.tokens:,} input positions, {state.supervised_tokens:,} targets, "
      f"final loss {state.history[-1]['loss']:.4f}")

import matplotlib.pyplot as plt

steps = [h["step"] for h in state.history]
fig, axes = plt.subplots(1, 3, figsize=(15, 4))

axes[0].plot(steps, [h["loss"] for h in state.history], label="train", color="#5E81AC")
if "eval_loss" in state.history[0]:
    axes[0].plot(steps, [h["eval_loss"] for h in state.history], label="valid",
                 color="#BF616A")
axes[0].axhline(np.log(cfg.vocab_size), ls="--", lw=1, color="#4C566A",
                label="ln(V) — a uniform model")
axes[0].set_xlabel("step"); axes[0].set_ylabel("cross-entropy"); axes[0].legend()
axes[0].set_title("loss")

axes[1].plot(steps, [h["lr_scale"] for h in state.history], color="#A3BE8C")
axes[1].set_xlabel("step"); axes[1].set_ylabel("lr multiplier")
axes[1].set_title("WSD schedule")

axes[2].plot(steps, [h["tokens_per_s"] for h in state.history], color="#D08770")
axes[2].set_xlabel("step"); axes[2].set_ylabel("input positions/s")
axes[2].set_title("throughput")

for ax in axes:
    ax.spines[["top", "right"]].set_visible(False)
plt.tight_layout(); plt.show()

bpc = state.history[-1]["loss"] / np.log(2)
print(f"final cross-entropy {state.history[-1]['loss']:.4f} nats "
      f"= {bpc:.4f} bits per token")
print('The curve describes this budget and corpus; a plateau alone does not identify the limiting cause.')

dense_cfg = cfg
moe_cfg = MiniLLMConfig(**{**MiniLLMConfig.twin(SETTINGS["model_profile"]).to_dict(),
                           "vocab_size": tokenizer.vocab_size,
                           "max_position_embeddings": cfg.max_position_embeddings})
moe_model = MiniLLM(moe_cfg).to(DEVICE)

rows = [("dense", dense_cfg, model), ("MoE", moe_cfg, moe_model)]
print(f"{'variant':<8} {'total':>14} {'active':>12} {'sparsity':>9} "
      f"{'KV B/tok':>9} {'FLOP/tok':>12}")
for name, c, m in rows:
    assert c.parameter_report()["total"] == m.num_parameters(), "formula diverged"
    total, active = c.parameter_report()["total"], c.active_parameters()
    print(f"{name:<8} {total:>14,} {active:>12,} {total / active:>8.2f}× "
          f"{kv_bytes_per_token(c):>9,} {c.flops_per_token():>12,}")

print(f"\n{moe_cfg.num_local_experts} experts, top-{moe_cfg.num_experts_per_tok}, "
      f"{moe_cfg.num_shared_experts} shared, first layer dense: "
      f"{not moe_cfg.moe_layer(0)}")
print(f"MoE active relative to dense: "
      f"×{moe_cfg.active_parameters() / dense_cfg.active_parameters():.3f}")
print(f"KV cache matches to the byte: "
      f"{kv_bytes_per_token(dense_cfg) == kv_bytes_per_token(moe_cfg)} — "
      f"attention untouched, so the comparison is clean")

moe_state = train(moe_model, train_stream, tcfg, eval_stream=valid_stream)

dense_last, moe_last = state.history[-1], moe_state.history[-1]
print(f"\n{'variant':<8} {'train':>9} {'valid':>9} {'input/s':>9}")
for name, h in (("dense", dense_last), ("MoE", moe_last)):
    print(f"{name:<8} {h['loss']:>9.4f} {h.get('eval_loss', float('nan')):>9.4f} "
          f"{h['tokens_per_s']:>9.0f}")

moe_layer = next(l.ffn for l in moe_model.layers if isinstance(l.ffn, MoELayer))
load = moe_layer.expert_load
print(f"\nexpert load at the last measurement: min {load.min():.3f}, "
      f"max {load.max():.3f}, std {load.std():.3f}")
print(f"a perfectly even load would be "
      f"{moe_cfg.num_experts_per_tok / moe_cfg.num_local_experts:.3f} per expert")
print(f"aux-loss-free bias: max|b| = {moe_layer.expert_bias.abs().max():.4f}")

from minillm.generate import SamplingConfig, generate

for temperature, top_p in ((0.0, 1.0), (0.8, 0.9)):
    prompt = tokenizer.encode("a language model")
    out, stats = generate(model, prompt, 80,
                          SamplingConfig(temperature=temperature, top_p=top_p, seed=SEED),
                          stop_tokens={EOS})
    label = "greedy" if temperature == 0 else f"T={temperature}, top_p={top_p}"
    print(f"--- {label}\n{tokenizer.decode(out)!r}\n")
print("stats of the last run:", stats.as_dict())

run_dir = Path("runs/pretrain/final")
save_checkpoint(run_dir, model, state.optimizer, tcfg, state)
tokenizer.save(run_dir / "tokenizer.json")
print("saved:", sorted(p.name for p in run_dir.iterdir()))

import json

def read_passport(path: Path) -> None:
    """A model passport from config.json — the same move as in module 1."""
    c = json.loads((path / "config.json").read_text(encoding="utf-8"))
    heads, kv = c["num_attention_heads"], c["num_key_value_heads"]
    attention = "MHA" if heads == kv else ("MQA" if kv == 1 else f"GQA {heads}/{kv}")
    ffn_ratio = c["intermediate_size"] / c["hidden_size"]
    cache_bytes = {"fp32": 4, "bf16": 2, "fp16": 2}[stats.cache_dtype]
    kv_bytes = 2 * c["num_hidden_layers"] * kv * c["head_dim"] * cache_bytes

    print("Axis I   · size      :", f"d={c['hidden_size']}, L={c['num_hidden_layers']}, "
          f"V={c['vocab_size']}, embeddings "
          f"{'tied' if c['tie_word_embeddings'] else 'untied'}")
    print("Axis II  · context   :", f"{c['max_position_embeddings']} tokens, "
          f"RoPE base {c['rope_theta']:g}, scaling "
          f"{c['rope_scaling'] or 'none'}")
    print("Axis III · attention :", attention,
          f"· head_dim {c['head_dim']} · sinks {'yes' if c['attention_sinks'] else 'no'}"
          f" · window {c['sliding_window'] or 'none'}")
    print("Axis IV  · block     :", f"{c['norm_placement']}-norm, {c['hidden_act']}, "
          f"d_ffn/d = {ffn_ratio:.2f}"
          + (f", clamp {c['swiglu_limit']}" if c["swiglu_limit"] else ""))
    print("Axis V   · sparsity  :", "dense" if not c["num_local_experts"]
          else f"MoE, {c['num_local_experts']} experts, top-{c['num_experts_per_tok']}")
    print("Derived              :", f"KV cache {kv_bytes} B/token in measured {stats.cache_dtype}, "
          f"{kv_bytes * c['max_position_embeddings'] / 1024 ** 2:.1f} MiB at full context")

read_passport(run_dir)

import importlib
import sys
import subprocess
import zipfile
from pathlib import Path

ARCHIVE_GLOB = "modern-llms-capstone-*.zip"

def _activate_capstone_root(candidate: Path):
    package_init = candidate / "minillm" / "__init__.py"
    if not package_init.exists():
        return None
    root = str(candidate.resolve())
    if root not in sys.path:
        sys.path.insert(0, root)
    return root

def _find_capstone_root(search_root: Path):
    direct = _activate_capstone_root(search_root)
    if direct:
        return direct
    for package_root in (search_root / "capstone" / "minillm" / "en", search_root / "minillm" / "en"):
        found = _activate_capstone_root(package_root)
        if found:
            return found
    for child in sorted(search_root.glob("modern-llms-capstone-*")):
        found = _activate_capstone_root(child)
        if found:
            return found
    return None

def _safe_extract(archive_path: Path, destination: Path):
    destination = destination.resolve()
    with zipfile.ZipFile(archive_path) as archive:
        for member in archive.infolist():
            member_path = Path(member.filename)
            if member_path.is_absolute() or ".." in member_path.parts:
                raise RuntimeError(f"unsafe path in archive: {member.filename}")
        archive.extractall(destination)

def ensure_minillm():
    for candidate in (Path.cwd(), *Path.cwd().parents):
        found = _find_capstone_root(candidate)
        if found:
            return "project package: " + found

    try:
        import minillm  # noqa: F401
        return "minillm is already importable"
    except ImportError:
        pass

    archives = sorted(Path.cwd().glob(ARCHIVE_GLOB)) + sorted(Path.cwd().glob("course-materials.zip"))
    if not archives:
        try:
            from google.colab import files
        except ImportError:
            files = None
        if files is not None:
            print("Upload course-materials.zip from the course notebook page. The directory structure will be restored automatically.")
            files.upload()
            archives = sorted(Path.cwd().glob(ARCHIVE_GLOB)) + sorted(Path.cwd().glob("course-materials.zip"))

    if archives:
        archive_path = archives[-1]
        destination = Path.cwd() / ".capstone_runtime"
        destination.mkdir(exist_ok=True)
        _safe_extract(archive_path, destination)
        root = _find_capstone_root(destination)
        if root:
            return "project archive: " + str(archive_path)
        matches = sorted(destination.glob("*/minillm/__init__.py"))
        if len(matches) != 1:
            raise RuntimeError(
                f"expected one minillm package in {archive_path.name}, "
                f"found {len(matches)}"
            )
        root = _activate_capstone_root(matches[0].parent.parent)
        if root:
            return "project archive: " + str(archive_path)

    raise RuntimeError("minillm was not found. Download course-materials.zip from the course notebook page and either extract it in full or upload the ZIP to the current Colab session.")

print(ensure_minillm())

for package in ("torch", "numpy", "matplotlib", "regex"):
    try:
        importlib.import_module(package)
    except ImportError:
        subprocess.run(
            [sys.executable, "-m", "pip", "install", "-q", package],
            check=True,
        )

import os, torch, numpy as np, random

# Run profile. SMOKE is sized for CPU and finishes in minutes; FULL only makes
# sense on a GPU and produces noticeably more coherent answers.
PROFILE = "SMOKE"
BUDGET = {
    "SMOKE": dict(base_steps=250, sft_steps=200, sft_batch=4, seq_len=192,
                  dpo_steps=100, dpo_batch=4, hidden=128, layers=3, dialogues=800),
    "FULL":  dict(base_steps=2000, sft_steps=1500, sft_batch=8, seq_len=256,
                  dpo_steps=600, dpo_batch=8, hidden=384, layers=6, dialogues=6000),
}[PROFILE]
if os.getenv("CAPSTONE_FAST_VALIDATION") == "1":
    BUDGET = {**BUDGET, "base_steps": min(BUDGET["base_steps"], 10),
              "sft_steps": min(BUDGET["sft_steps"], 12),
              "dpo_steps": min(BUDGET["dpo_steps"], 8),
              "dialogues": min(BUDGET["dialogues"], 80),
              "sft_batch": min(BUDGET["sft_batch"], 2),
              "dpo_batch": min(BUDGET["dpo_batch"], 2)}
SEED = 0
torch.manual_seed(SEED)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print("torch", torch.__version__, "| device:", DEVICE)

from minillm.chat import CHAT_SPECIALS
from minillm.config import MiniLLMConfig
from minillm.data import TokenStream
from minillm.model import MiniLLM
from minillm.tokenizer import BBPETokenizer
from minillm.train import TrainConfig, train

EN_SLOTS = dict(
    subject=["the language model", "the transformer", "the attention mechanism",
             "the tokenizer", "the optimizer", "normalization", "the residual stream",
             "the expert router"],
    verb=["predicts", "processes", "reduces", "accelerates", "stabilizes"],
    obj=["the next token", "a long context", "the cost per step", "the memory footprint",
         "the scale of activations", "the load across experts"],
)

def sentence(rng, slots):
    return " ".join([rng.choice(slots["subject"]), rng.choice(slots["verb"]),
                     rng.choice(slots["obj"])]).capitalize() + "."

rng = random.Random(SEED)
base_text = "\n".join(" ".join(sentence(rng, EN_SLOTS) for _ in range(4))
                      for _ in range(4000))

def usable_checkpoint(path: Path):
    """A checkpoint is usable only if its vocabulary knows the chat control tokens.

    A vocabulary without `<|eot_id|>` would crash the notebook on the very first
    loss mask, and with an unhelpful KeyError at that. Notebook 1 reserves those
    identifiers up front, but a checkpoint may have been built by an older
    version — so we check explicitly.
    """
    if not (path / "config.json").exists():
        return None, None, "no notebook-1 checkpoint nearby"
    tok = BBPETokenizer.load(path / "tokenizer.json")
    missing = [s for s in CHAT_SPECIALS if s not in tok.specials]
    if missing:
        return None, None, f"the checkpoint vocabulary lacks {missing}"
    from minillm.train import load_checkpoint
    model, _, _, _ = load_checkpoint(path, device=DEVICE)
    return model, tok, None

CHECKPOINT = Path("runs/pretrain/final")
base_model, tokenizer, reason = usable_checkpoint(CHECKPOINT)
if base_model is not None:
    print("loaded the checkpoint from notebook 1:", base_model.num_parameters(), "parameters")
else:
    print(f"training a base model ourselves: {reason}")
    tokenizer = BBPETokenizer.train(base_text, vocab_size=1024, specials=CHAT_SPECIALS)
    cfg = MiniLLMConfig.profile("S", vocab_size=tokenizer.vocab_size,
                                hidden_size=BUDGET["hidden"],
                                num_hidden_layers=BUDGET["layers"],
                                num_attention_heads=4, num_key_value_heads=2,
                                head_dim=32, max_position_embeddings=BUDGET["seq_len"])
    base_model = MiniLLM(cfg).to(DEVICE)
    stream = TokenStream.from_text(base_text, tokenizer, seq_len=128)
    print(f"parameters: {base_model.num_parameters():,}")
    train(base_model, stream, TrainConfig(steps=BUDGET["base_steps"], seq_len=128,
                                          micro_batch_size=8, optimizer="muon",
                                          seed=SEED, device=DEVICE, log_every=125))
print("vocabulary:", tokenizer.vocab_size)

from minillm.chat import ChatTemplate, Message, SFTDataset

QUESTIONS = [
    ("What does {s} do?", "{s} {v} {o}."),
    ("What is {s} used for?", "{s} {v} {o}. That is its main role here."),
    ("How does {s} affect the model?", "{s} {v} {o}, changing the model behaviour."),
]

def make_dialogues(n: int, seed: int = 0) -> list[list[Message]]:
    rng = random.Random(seed)
    out = []
    for _ in range(n):
        question, answer = rng.choice(QUESTIONS)
        s = rng.choice(EN_SLOTS["subject"])
        v, o = rng.choice(EN_SLOTS["verb"]), rng.choice(EN_SLOTS["obj"])
        out.append([
            Message("system", "Answer briefly and directly."),
            Message("user", question.format(s=s)),
            Message("assistant", answer.format(s=s.capitalize(), v=v, o=o)),
        ])
    return out

template = ChatTemplate()
dialogues = make_dialogues(BUDGET["dialogues"], seed=SEED)
print(template.render(dialogues[0], add_generation_prompt=False))

EOT_ID = tokenizer.specials["<|eot_id|>"]
PAD_ID = tokenizer.specials["<|begin_of_text|>"]

ids, mask = template.encode_with_mask(dialogues[0], tokenizer)
print(f"tokens: {len(ids)}, trainable: {sum(mask)} "
      f"({template.trainable_fraction(mask):.1%})\n")

print(f"{'token':<22} {'id':>6}  train")
for token_id, flag in list(zip(ids, mask)):
    piece = tokenizer.decode([token_id]).replace("\n", "\\n")
    print(f"{piece!r:<22} {token_id:>6}  {'YES' if flag else '—'}")

dataset, stats = SFTDataset.from_dialogues(dialogues, tokenizer, template,
                                           seq_len=BUDGET["seq_len"], pad_id=PAD_ID)
print("packing:", stats)
print(f"\nwithout packing this would need {stats['dialogues']} rows of "
      f"{BUDGET['seq_len']} tokens,")
print(f"with packing — {stats['bins']}: a ×{stats['packing_ratio']} saving, "
      f"padding {stats['padding_fraction']:.1%}")

import copy

sft_model = copy.deepcopy(base_model)
sft_cfg = TrainConfig(steps=BUDGET["sft_steps"], seq_len=BUDGET["seq_len"],
                      micro_batch_size=BUDGET["sft_batch"], optimizer="muon",
                      seed=SEED, device=DEVICE, log_every=BUDGET["sft_steps"] // 4,
                      out_dir="runs/sft")
sft_state = train(sft_model, dataset, sft_cfg)

from minillm.generate import SamplingConfig, generate

def ask(model, question: str, max_new_tokens: int = 60, temperature: float = 0.0) -> tuple[str, int]:
    ids, _ = template.encode_with_mask(
        [Message("system", "Answer briefly and directly."),
         Message("user", question)], tokenizer, add_generation_prompt=True)
    out, stats = generate(model, ids, max_new_tokens,
                          SamplingConfig(temperature=temperature, seed=SEED),
                          stop_tokens={EOT_ID})
    answer = tokenizer.decode(out[len(ids):])
    return answer, stats.generated_tokens

question = "What does the tokenizer do?"
for label, model in (("before SFT", base_model), ("after SFT", sft_model)):
    answer, n = ask(model, question)
    print(f"--- {label} ({n} tokens)\n{answer!r}\n")

broken_dataset, _ = SFTDataset.from_dialogues(dialogues, tokenizer, template,
                                              seq_len=BUDGET["seq_len"], pad_id=PAD_ID,
                                              train_eot=False)
a_ids, a_mask = template.encode_with_mask(dialogues[0], tokenizer, train_eot=True)
b_ids, b_mask = template.encode_with_mask(dialogues[0], tokenizer, train_eot=False)
diff = [i for i, (x, y) in enumerate(zip(a_mask, b_mask)) if x != y]
print(f"sequences identical: {a_ids == b_ids}")
print(f"masks differ at positions {diff} — that is {tokenizer.decode([a_ids[diff[0]]])!r}")

broken_model = copy.deepcopy(base_model)
train(broken_model, broken_dataset,
      TrainConfig(steps=BUDGET["sft_steps"], seq_len=BUDGET["seq_len"],
                  micro_batch_size=BUDGET["sft_batch"], optimizer="muon", seed=SEED,
                  device=DEVICE, log_every=BUDGET["sft_steps"], out_dir="runs/sft_broken"))

eot_rows = []

print(f"{'model':<16} {'p(<|eot_id|>) after reply':>28} {'reply length':>14}")
for label, model in (("mask correct", sft_model), ("EOT forgotten", broken_model)):
    ids, _ = template.encode_with_mask(dialogues[0], tokenizer)
    prefix = ids[:len(ids) - 1]        # everything but the closing EOT
    with torch.no_grad():
        logits, _ = model(torch.tensor([prefix], device=DEVICE))
        p_eot = torch.softmax(logits[0, -1].float(), dim=-1)[EOT_ID].item()
    _, generated = ask(model, "What does the tokenizer do?", max_new_tokens=120)
    eot_rows.append(dict(label=label, probability=p_eot, generated_tokens=generated))
    print(f"{label:<16} {p_eot:>28.4f} {generated:>14}")

answer, n = ask(broken_model, "What does the tokenizer do?", max_new_tokens=120)
print(f"\nreply from the model with the forgotten EOT ({n} tokens):\n{answer!r}")

# Use one correct mask for both models; the generator is a diagnostic fixture.
from minillm.train import evaluate
check_dialogues = make_dialogues(32, seed=SEED + 991)
check_dataset, _ = SFTDataset.from_dialogues(check_dialogues, tokenizer, template,
    seq_len=BUDGET["seq_len"], pad_id=PAD_ID, train_eot=True)
shared_scores = {name: evaluate(model, check_dataset, sft_cfg, batches=4)
                 for name, model in (("correct_mask", sft_model), ("omitted_eot", broken_model))}
assert sum(a_mask) == sum(b_mask) + len(diff)
print("EOT supervised positions:", sum(a_mask), "vs", sum(b_mask))
print("Shared correct-mask evaluation:", shared_scores)

from minillm.preference import (DPOConfig, PreferencePair, freeze_reference, train_dpo)

FILLER = " This matters and should be considered when working with the model."

def build_pairs(n: int, seed: int = 0) -> list[PreferencePair]:
    rng = random.Random(seed)
    pairs = []
    for _ in range(n):
        question, answer = rng.choice(QUESTIONS)
        s = rng.choice(EN_SLOTS["subject"])
        v, o = rng.choice(EN_SLOTS["verb"]), rng.choice(EN_SLOTS["obj"])
        prompt_ids, _ = template.encode_with_mask(
            [Message("user", question.format(s=s))], tokenizer,
            add_generation_prompt=True)
        short = answer.format(s=s.capitalize(), v=v, o=o)
        verbose = short + FILLER * 2
        pairs.append(PreferencePair(
            prompt_ids=prompt_ids,
            chosen_ids=tokenizer.encode(short) + [EOT_ID],
            rejected_ids=tokenizer.encode(verbose) + [EOT_ID],
        ))
    return pairs

pairs = build_pairs(192, seed=SEED)
print(f"pairs: {len(pairs)}")
print("chosen  :", tokenizer.decode(pairs[0].chosen_ids)[:90])
print("rejected:", tokenizer.decode(pairs[0].rejected_ids)[:90])

policy = copy.deepcopy(sft_model)
reference = freeze_reference(sft_model)
dpo_state = train_dpo(policy, pairs,
                      DPOConfig(beta=0.1, lr=1e-4, steps=BUDGET["dpo_steps"],
                                batch_size=BUDGET["dpo_batch"], seed=SEED,
                                device=DEVICE, log_every=BUDGET["dpo_steps"] // 8),
                      pad_id=PAD_ID, reference=reference)

import matplotlib.pyplot as plt

steps = [h["step"] for h in dpo_state.history]
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
axes[0].plot(steps, [h["loss"] for h in dpo_state.history], color="#5E81AC")
axes[0].axhline(np.log(2), ls="--", lw=1, color="#4C566A", label="ln 2 — the start")
axes[0].set_title("DPO loss"); axes[0].legend()
axes[1].plot(steps, [h["margin"] for h in dpo_state.history], color="#A3BE8C")
axes[1].axhline(0, ls="--", lw=1, color="#4C566A")
axes[1].set_title("margin: difference of implicit rewards")
axes[2].plot(steps, [h["chosen_reward"] for h in dpo_state.history],
             label="chosen", color="#A3BE8C")
axes[2].plot(steps, [h["rejected_reward"] for h in dpo_state.history],
             label="rejected", color="#BF616A")
axes[2].set_title("implicit rewards"); axes[2].legend()
for ax in axes:
    ax.set_xlabel("step"); ax.spines[["top", "right"]].set_visible(False)
plt.tight_layout(); plt.show()

last = dpo_state.history[-1]
print(f"margin {dpo_state.history[0]['margin']:+.3f} → {last['margin']:+.3f}, "
      f"accuracy {last['accuracy']:.2f}")

question = "What is the optimizer used for?"
for label, model in (("after SFT", sft_model), ("after DPO", policy)):
    answer, n = ask(model, question, max_new_tokens=80, temperature=0.7)
    print(f"--- {label} ({n} tokens)\n{answer!r}\n")
print('Two responses are illustrations. Compare several fixed prompts and save outcomes before inferring a behavioral change.')

from minillm.train import save_model_artifact

out_dir = Path("runs/posttrain/final")
save_model_artifact(out_dir, policy, {
    "stage": "dpo",
    "sft_steps": sft_state.step,
    "dpo_steps": dpo_state.step,
    "note": "inference artifact; DPO optimizer state is intentionally not stored",
})
tokenizer.save(out_dir / "tokenizer.json")
print("saved:", sorted(p.name for p in out_dir.iterdir()))

import json
metrics = {"seed": SEED, "budget": BUDGET, "eot": eot_rows,
           "shared_mask_scores": shared_scores, "dpo_history": dpo_state.history}
(out_dir / "teaching_metrics.json").write_text(
    json.dumps(metrics, ensure_ascii=False, indent=2, allow_nan=False), encoding="utf-8")

import importlib
import sys
import subprocess
import zipfile
from pathlib import Path

ARCHIVE_GLOB = "modern-llms-capstone-*.zip"

def _activate_capstone_root(candidate: Path):
    package_init = candidate / "minillm" / "__init__.py"
    if not package_init.exists():
        return None
    root = str(candidate.resolve())
    if root not in sys.path:
        sys.path.insert(0, root)
    return root

def _find_capstone_root(search_root: Path):
    direct = _activate_capstone_root(search_root)
    if direct:
        return direct
    for package_root in (search_root / "capstone" / "minillm" / "en", search_root / "minillm" / "en"):
        found = _activate_capstone_root(package_root)
        if found:
            return found
    for child in sorted(search_root.glob("modern-llms-capstone-*")):
        found = _activate_capstone_root(child)
        if found:
            return found
    return None

def _safe_extract(archive_path: Path, destination: Path):
    destination = destination.resolve()
    with zipfile.ZipFile(archive_path) as archive:
        for member in archive.infolist():
            member_path = Path(member.filename)
            if member_path.is_absolute() or ".." in member_path.parts:
                raise RuntimeError(f"unsafe path in archive: {member.filename}")
        archive.extractall(destination)

def ensure_minillm():
    for candidate in (Path.cwd(), *Path.cwd().parents):
        found = _find_capstone_root(candidate)
        if found:
            return "project package: " + found

    try:
        import minillm  # noqa: F401
        return "minillm is already importable"
    except ImportError:
        pass

    archives = sorted(Path.cwd().glob(ARCHIVE_GLOB)) + sorted(Path.cwd().glob("course-materials.zip"))
    if not archives:
        try:
            from google.colab import files
        except ImportError:
            files = None
        if files is not None:
            print("Upload course-materials.zip from the course notebook page. The directory structure will be restored automatically.")
            files.upload()
            archives = sorted(Path.cwd().glob(ARCHIVE_GLOB)) + sorted(Path.cwd().glob("course-materials.zip"))

    if archives:
        archive_path = archives[-1]
        destination = Path.cwd() / ".capstone_runtime"
        destination.mkdir(exist_ok=True)
        _safe_extract(archive_path, destination)
        root = _find_capstone_root(destination)
        if root:
            return "project archive: " + str(archive_path)
        matches = sorted(destination.glob("*/minillm/__init__.py"))
        if len(matches) != 1:
            raise RuntimeError(
                f"expected one minillm package in {archive_path.name}, "
                f"found {len(matches)}"
            )
        root = _activate_capstone_root(matches[0].parent.parent)
        if root:
            return "project archive: " + str(archive_path)

    raise RuntimeError("minillm was not found. Download course-materials.zip from the course notebook page and either extract it in full or upload the ZIP to the current Colab session.")

print(ensure_minillm())

for package in ("torch", "numpy", "matplotlib", "regex"):
    try:
        importlib.import_module(package)
    except ImportError:
        subprocess.run(
            [sys.executable, "-m", "pip", "install", "-q", package],
            check=True,
        )

import copy, math, os, random, time
import torch, numpy as np
import matplotlib.pyplot as plt

SEED = 0
torch.manual_seed(SEED)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
from minillm.precision import autocast_context, model_parameter_dtype, resolve_compute_dtype
COMPUTE_DTYPE = resolve_compute_dtype(DEVICE, "auto")
print("torch", torch.__version__, "| device:", DEVICE, "| compute:", COMPUTE_DTYPE)

from minillm.chat import ChatTemplate, Message
from minillm.generate import SamplingConfig, generate
from minillm.tokenizer import BBPETokenizer
from minillm.train import load_model_artifact

POSTTRAIN_ARTIFACT = Path("runs/posttrain/final")
if (POSTTRAIN_ARTIFACT / "model.pt").exists():
    handoff_model, handoff_meta = load_model_artifact(POSTTRAIN_ARTIFACT, DEVICE)
    handoff_tokenizer = BBPETokenizer.load(POSTTRAIN_ARTIFACT / "tokenizer.json")
    handoff_template = ChatTemplate()
    handoff_ids, _ = handoff_template.encode_with_mask(
        [Message("system", "Answer briefly and directly."),
         Message("user", "What does a tokenizer do?")],
        handoff_tokenizer, add_generation_prompt=True)
    handoff_out, handoff_stats = generate(
        handoff_model, handoff_ids, 40, SamplingConfig(temperature=0.0),
        stop_tokens={handoff_tokenizer.specials["<|eot_id|>"]})
    print("post-training handoff:", handoff_meta)
    print("reply:", repr(handoff_tokenizer.decode(handoff_out[len(handoff_ids):])))
    print("KV bytes:", handoff_stats.kv_bytes)
else:
    print("post-training artifact not found — continuing via the self-contained route")

from minillm.tokenizer import BBPETokenizer
from minillm.data import TokenStream

def high_entropy_corpus(n_lines: int = 12_000, lexicon: int = 600,
                        seed: int = 0) -> str:
    """A pseudo-language: large lexicon, Zipfian frequencies.

    Unequal frequencies define a controlled prediction task. Loss is measured
    per BPE token, so the entropy of the word distribution is not a direct
    reference value for this metric.
    """
    rng = random.Random(seed)
    alphabet = "abcdefghijklmnopqrstuvwxyz"
    words = ["".join(rng.choice(alphabet) for _ in range(rng.randint(3, 8)))
             for _ in range(lexicon)]
    weights = [1.0 / (i + 1) ** 0.9 for i in range(lexicon)]   # Zipfian tail
    lines = []
    for _ in range(n_lines):
        k = rng.randint(6, 14)
        lines.append(" ".join(rng.choices(words, weights=weights, k=k)) + ".\n")
    return "".join(lines)

corpus = high_entropy_corpus(n_lines=8_000, seed=SEED)
lines = corpus.splitlines(keepends=True)
cut = int(len(lines) * 0.9)
train_text, eval_text = "".join(lines[:cut]), "".join(lines[cut:])
tokenizer = BBPETokenizer.train(train_text, vocab_size=1024, specials=["<|eos|>"])
stream = TokenStream.from_text(train_text, tokenizer, seq_len=128)
eval_stream = TokenStream.from_text(eval_text, tokenizer, seq_len=128)
print(f"corpus {len(corpus):,} characters · vocabulary {tokenizer.vocab_size} · "
      f"{len(stream):,} windows · sha {stream.checksum()}")
print(f"uniform-model reference loss: ln(V) = "
      f"{math.log(tokenizer.vocab_size):.3f}")

from minillm.config import MiniLLMConfig
from minillm.model import MiniLLM
from minillm.train import TrainConfig, evaluate, train

def build(hidden: int, layers: int, heads: int, kv: int) -> MiniLLM:
    cfg = MiniLLMConfig.profile("S", vocab_size=tokenizer.vocab_size,
                                hidden_size=hidden, num_hidden_layers=layers,
                                num_attention_heads=heads, num_key_value_heads=kv,
                                head_dim=hidden // heads, max_position_embeddings=256)
    return MiniLLM(cfg).to(DEVICE)

# This budget targets CPU execution. Longer training is a separate experiment;
# remeasure both loss and sensitivity to quantization afterwards.
models = {}
for name, (hidden, layers, heads, kv, steps) in {
    "target": (192, 4, 6, 2, 250),
    "draft":  (96, 2, 4, 2, 250),
}.items():
    if os.getenv("CAPSTONE_FAST_VALIDATION") == "1":
        steps = 12
    torch.manual_seed(SEED)
    model = build(hidden, layers, heads, kv)
    print(f"--- {name}: {model.num_parameters():,} parameters")
    train(model, stream, TrainConfig(steps=steps, seq_len=128, micro_batch_size=8,
                                     optimizer="muon", seed=SEED, device=DEVICE, compute_dtype="auto",
                                     log_every=steps // 2))
    models[name] = model

target, draft = models["target"], models["draft"]
eval_cfg = TrainConfig(seq_len=128, micro_batch_size=8, device=DEVICE, compute_dtype="auto")
base_loss = evaluate(target, eval_stream, eval_cfg, batches=(2 if os.getenv("CAPSTONE_FAST_VALIDATION") == "1" else 16))
BASELINE_COMPUTE = resolve_compute_dtype(DEVICE, "auto")
BASELINE_PARAMETER = model_parameter_dtype(target)
BASELINE_LABEL = f"{BASELINE_PARAMETER}/{BASELINE_COMPUTE}"
BASELINE_STORAGE_BITS = 32.0 if BASELINE_PARAMETER == "fp32" else 16.0
print("target eval loss:", base_loss, "uniform reference:", math.log(tokenizer.vocab_size))

from minillm.generate import SamplingConfig, generate

rows = []
for prompt_len in (16, 32, 64, 128):
    prompt = tokenizer.encode(corpus[:4000])[:prompt_len]
    _, stats = generate(target, prompt, 32, SamplingConfig(temperature=0.0))
    rows.append((prompt_len, stats.prefill_seconds * 1000,
                 stats.decode_seconds / stats.generated_tokens * 1000,
                 stats.kv_bytes / 1024))

print(f"{'prompt':>7} {'prefill, ms':>12} {'ms/token decode':>17} {'KV cache, KiB':>15}")
for prompt_len, prefill_ms, decode_ms, kv_kib in rows:
    print(f"{prompt_len:>7} {prefill_ms:>12.2f} {decode_ms:>17.2f} {kv_kib:>15.1f}")

print('Compare measured phases; diagnosing a bottleneck needs repeated measurements and profiling.')

from minillm.config import gqa_variants, kv_bytes_per_token

variants = gqa_variants(target.cfg)
runtime_bytes = 4 if BASELINE_COMPUTE == "fp32" else 2
print(f"{'scheme':>6} {'H_kv':>5} {('bytes/token ' + BASELINE_COMPUTE):>20} {'at 128K context':>20}")
for name, cfg in variants.items():
    per_token = kv_bytes_per_token(cfg, dtype_bytes=runtime_bytes)
    print(f"{name:>6} {cfg.num_key_value_heads:>5} {per_token:>20,} "
          f"{per_token * 131072 / 1024 ** 3:>18.2f} GiB")

# Checking the formula against the actual cache allocation.
prompt = tokenizer.encode(corpus[:2000])[:48]
cache = target.new_cache()
with torch.no_grad():
    with autocast_context(DEVICE, BASELINE_COMPUTE):
        target(torch.tensor([prompt], device=DEVICE), cache=cache)
actual_bytes = cache.nbytes() // (len(prompt) * 2 * target.cfg.num_hidden_layers * target.cfg.kv_dim)
expected = kv_bytes_per_token(target.cfg, dtype_bytes=actual_bytes) * len(prompt)
print(f"\nactual cache dtype: {str(cache.dtype()).replace('torch.', '')}")
print(f"actual {cache.nbytes():,} B == formula {expected:,} B: {cache.nbytes() == expected}")

assert cache.nbytes() == expected

system_prompt = tokenizer.encode(corpus[:1200])[:96]
questions = [tokenizer.encode(corpus[5000 + 500 * i:5000 + 500 * i + 200])[:16]
             for i in range(4)]

warm = target.new_cache()
if DEVICE.startswith("cuda"):
    torch.cuda.synchronize()
t0 = time.perf_counter()
with torch.no_grad():
    with autocast_context(DEVICE, BASELINE_COMPUTE):
        target(torch.tensor([system_prompt], device=DEVICE), cache=warm)
if DEVICE.startswith("cuda"):
    torch.cuda.synchronize()
warmup_ms = (time.perf_counter() - t0) * 1000

cold_total = hot_total = 0.0
identical = True
for tail in questions:
    cold, cold_stats = generate(target, system_prompt + tail, 16,
                                SamplingConfig(temperature=0.0))
    hot, hot_stats = generate(target, system_prompt + tail, 16,
                              SamplingConfig(temperature=0.0), cache=warm.clone())
    identical &= cold == hot
    cold_total += cold_stats.prefill_seconds
    hot_total += hot_stats.prefill_seconds

print(f"results identical: {identical}")
print(f"prefill without cache: {cold_total * 1000:6.1f} ms over {len(questions)} requests")
print(f"prefill with cache:    {hot_total * 1000:6.1f} ms (+{warmup_ms:.1f} ms warmup, once)")
print(f"saving over {len(questions)} requests: "
      f"×{cold_total / max(hot_total, 1e-9):.1f} on the prefill phase")

from minillm.generate import quantize_model

probe = eval_stream.batch(torch.arange(16)).to(DEVICE)
with torch.no_grad():
    with autocast_context(DEVICE, BASELINE_COMPUTE):
        reference_logits, _ = target(probe)
reference_probs = torch.softmax(reference_logits.float(), dim=-1)
reference_top1 = reference_logits.argmax(-1)

print(f"{'format':>7} {'eff bits/wt':>11} {'compr.':>8} {'loss':>8} {'Δloss':>8} "
      f"{'KL, nats':>9} {'top-1 agree':>12}")
print(f"{BASELINE_LABEL:>7} {BASELINE_STORAGE_BITS:>11.2f} {'×1.00':>8} {base_loss:>8.4f} {0.0:>8.4f} "
      f"{0.0:>9.5f} {1.0:>11.1%}")

quant_rows = []
for method in ("int8", "nvfp4", "mxfp4"):
    clone = copy.deepcopy(target)
    report = quantize_model(clone, method)
    loss = evaluate(clone, eval_stream, eval_cfg, batches=(2 if os.getenv("CAPSTONE_FAST_VALIDATION") == "1" else 16))
    with torch.no_grad():
        with autocast_context(DEVICE, BASELINE_COMPUTE):
            logits, _ = clone(probe)
    probs = torch.softmax(logits.float(), dim=-1)
    kl = torch.sum(reference_probs * (torch.log(reference_probs + 1e-12)
                                      - torch.log(probs + 1e-12)), dim=-1).mean().item()
    agreement = (logits.argmax(-1) == reference_top1).float().mean().item()
    quant_rows.append((method, report, loss, kl, agreement))
    print(f"{method:>7} {report['effective_bits_per_weight']:>11.2f} "
          f"{'×' + format(report['compression_vs_runtime_fp32'], '.2f'):>8} {loss:>8.4f} "
          f"{loss - base_loss:>+8.4f} {kl:>9.5f} {agreement:>11.1%}")

fig, axes = plt.subplots(1, 3, figsize=(15, 4))
names = [BASELINE_LABEL] + [r[0] for r in quant_rows]
colors = ["#4C566A", "#A3BE8C", "#EBCB8B", "#BF616A"]

axes[0].bar(names, [base_loss] + [r[2] for r in quant_rows], color=colors)
axes[0].set_ylim(min([base_loss] + [r[2] for r in quant_rows]) * 0.98, max([base_loss] + [r[2] for r in quant_rows]) * 1.02)
axes[0].set_title("loss (note the clipped axis!)"); axes[0].set_ylabel("nats")

axes[1].bar(names, [0.0] + [r[3] for r in quant_rows], color=colors)
axes[1].set_title("KL(baseline ‖ fake quant)"); axes[1].set_ylabel("nats")

axes[2].bar(names, [1.0] + [r[4] for r in quant_rows], color=colors)
axes[2].set_ylim(0, 1.05); axes[2].set_title("top-1 agreement with baseline")

for ax in axes:
    ax.spines[["top", "right"]].set_visible(False)
plt.tight_layout(); plt.show()

from minillm.generate import speculative_generate

prompt = tokenizer.encode(corpus[:2000])[:32]
print(f"{'k':>3} {'accept':>8} {'tokens/target pass':>21} {'draft passes':>14}")
lookahead_rows = []
for k in (1, 2, 4, 8):
    _, stats = speculative_generate(target, draft, prompt, 48, lookahead=k,
                                    sampling=SamplingConfig(temperature=1.0, seed=SEED))
    lookahead_rows.append((k, stats.acceptance_rate, stats.tokens_per_target_forward,
                           stats.draft_decode_forwards))
    print(f"{k:>3} {stats.acceptance_rate:>7.1%} "
          f"{stats.tokens_per_target_forward:>21.2f} {stats.draft_decode_forwards:>14}")

# A draft equal to the target: the min(1, p/q) rule must accept everything.
_, twin = speculative_generate(target, copy.deepcopy(target), prompt, 32, lookahead=4,
                               sampling=SamplingConfig(temperature=1.0, seed=SEED))
print(f"draft == target: acceptance {twin.acceptance_rate:.1%}, "
      f"{twin.tokens_per_target_forward:.2f} tokens per pass")

fig, axes = plt.subplots(1, 2, figsize=(11, 4))
ks = [r[0] for r in lookahead_rows]
axes[0].plot(ks, [r[1] for r in lookahead_rows], "o-", color="#5E81AC")
axes[0].set_xlabel("draft depth k"); axes[0].set_ylabel("acceptance rate")
axes[0].set_title("observed acceptance")
axes[1].plot(ks, [r[2] for r in lookahead_rows], "o-", color="#A3BE8C")
axes[1].axhline(1.0, ls="--", lw=1, color="#4C566A", label="no speculation")
axes[1].set_xlabel("draft depth k")
axes[1].set_ylabel("tokens per target pass")
axes[1].set_title("observed yield per verification"); axes[1].legend()
for ax in axes:
    ax.spines[["top", "right"]].set_visible(False)
plt.tight_layout(); plt.show()

import json
inference_record = {"seed": SEED, "device": DEVICE,
    "compute_dtype": BASELINE_COMPUTE, "parameter_dtype": BASELINE_PARAMETER,
    "evaluation_sha": eval_stream.checksum(), "base_loss": base_loss,
    "quantization": [dict(method=r[0], report=r[1], loss=r[2], kl=r[3], top1=r[4]) for r in quant_rows],
    "speculative": [dict(k=r[0], acceptance=r[1], tokens_per_verification=r[2], draft_decode_passes=r[3]) for r in lookahead_rows],
    "identical_draft_acceptance": twin.acceptance_rate,
    "prefix_greedy_identical": identical}
Path("runs/inference").mkdir(parents=True, exist_ok=True)
Path("runs/inference/measurements.json").write_text(
    json.dumps(inference_record, ensure_ascii=False, indent=2, allow_nan=False), encoding="utf-8")

import importlib
import sys
import subprocess
import zipfile
from pathlib import Path

ARCHIVE_GLOB = "modern-llms-capstone-*.zip"

def _activate_capstone_root(candidate: Path):
    package_init = candidate / "minillm" / "__init__.py"
    if not package_init.exists():
        return None
    root = str(candidate.resolve())
    if root not in sys.path:
        sys.path.insert(0, root)
    return root

def _find_capstone_root(search_root: Path):
    direct = _activate_capstone_root(search_root)
    if direct:
        return direct
    for package_root in (search_root / "capstone" / "minillm" / "en", search_root / "minillm" / "en"):
        found = _activate_capstone_root(package_root)
        if found:
            return found
    for child in sorted(search_root.glob("modern-llms-capstone-*")):
        found = _activate_capstone_root(child)
        if found:
            return found
    return None

def _safe_extract(archive_path: Path, destination: Path):
    destination = destination.resolve()
    with zipfile.ZipFile(archive_path) as archive:
        for member in archive.infolist():
            member_path = Path(member.filename)
            if member_path.is_absolute() or ".." in member_path.parts:
                raise RuntimeError(f"unsafe path in archive: {member.filename}")
        archive.extractall(destination)

def ensure_minillm():
    for candidate in (Path.cwd(), *Path.cwd().parents):
        found = _find_capstone_root(candidate)
        if found:
            return "project package: " + found

    try:
        import minillm  # noqa: F401
        return "minillm is already importable"
    except ImportError:
        pass

    archives = sorted(Path.cwd().glob(ARCHIVE_GLOB)) + sorted(Path.cwd().glob("course-materials.zip"))
    if not archives:
        try:
            from google.colab import files
        except ImportError:
            files = None
        if files is not None:
            print("Upload course-materials.zip from the course notebook page. The directory structure will be restored automatically.")
            files.upload()
            archives = sorted(Path.cwd().glob(ARCHIVE_GLOB)) + sorted(Path.cwd().glob("course-materials.zip"))

    if archives:
        archive_path = archives[-1]
        destination = Path.cwd() / ".capstone_runtime"
        destination.mkdir(exist_ok=True)
        _safe_extract(archive_path, destination)
        root = _find_capstone_root(destination)
        if root:
            return "project archive: " + str(archive_path)
        matches = sorted(destination.glob("*/minillm/__init__.py"))
        if len(matches) != 1:
            raise RuntimeError(
                f"expected one minillm package in {archive_path.name}, "
                f"found {len(matches)}"
            )
        root = _activate_capstone_root(matches[0].parent.parent)
        if root:
            return "project archive: " + str(archive_path)

    raise RuntimeError("minillm was not found. Download course-materials.zip from the course notebook page and either extract it in full or upload the ZIP to the current Colab session.")

print(ensure_minillm())

for package in ("torch", "numpy", "matplotlib", "regex"):
    try:
        importlib.import_module(package)
    except ImportError:
        subprocess.run(
            [sys.executable, "-m", "pip", "install", "-q", package],
            check=True,
        )

import math, os, random
import torch, numpy as np
import matplotlib.pyplot as plt

SEED = 0
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
STEPS = 60 if DEVICE == "cpu" else 600
if os.getenv("CAPSTONE_FAST_VALIDATION") == "1":
    STEPS = 5     # on a GPU the contrasts are sharper
print("device:", DEVICE, "| steps per run:", STEPS)

from minillm.config import MiniLLMConfig, kv_bytes_per_token
from minillm.data import TokenStream
from minillm.model import MiniLLM, RMSNorm
from minillm.tokenizer import BBPETokenizer
from minillm.train import TrainConfig, evaluate, train

def high_entropy_corpus(n_lines=8_000, lexicon=600, seed=0) -> str:
    """The same pseudo-language as notebook 3: a Zipfian lexicon leaves models
    the headroom in which architectural variants are distinguishable."""
    rng = random.Random(seed)
    alphabet = "abcdefghijklmnopqrstuvwxyz"
    words = ["".join(rng.choice(alphabet) for _ in range(rng.randint(3, 8)))
             for _ in range(lexicon)]
    weights = [1.0 / (i + 1) ** 0.9 for i in range(lexicon)]
    return "".join(" ".join(rng.choices(words, weights=weights, k=rng.randint(6, 14)))
                   + ".\n" for _ in range(n_lines))

corpus = high_entropy_corpus(seed=SEED)
lines = corpus.splitlines(keepends=True)
a, b = int(len(lines) * .8), int(len(lines) * .9)
train_text, valid_text, test_text = map("".join, (lines[:a], lines[a:b], lines[b:]))
tokenizer = BBPETokenizer.train(train_text, vocab_size=1024, specials=["<|eos|>"])
train_stream = TokenStream.from_text(train_text, tokenizer, 128)
valid_stream = TokenStream.from_text(valid_text, tokenizer, 128)
test_stream = TokenStream.from_text(test_text, tokenizer, 128)
CEILING = math.log(tokenizer.vocab_size)
print(f"vocabulary {tokenizer.vocab_size} · train {len(train_stream):,} windows · "
      f"valid {len(valid_stream):,} windows · uniform reference ln(V) = {CEILING:.3f}")

BASE = dict(vocab_size=tokenizer.vocab_size, hidden_size=128, num_hidden_layers=3,
            num_attention_heads=4, num_key_value_heads=2, head_dim=32,
            max_position_embeddings=256)

def run(label: str, steps: int = None, **overrides):
    """One ablation run: fixed seed, fixed data, fixed schedule."""
    cfg = MiniLLMConfig.profile("S", **{**BASE, **overrides})
    torch.manual_seed(SEED)
    model = MiniLLM(cfg).to(DEVICE)
    tcfg = TrainConfig(steps=steps or STEPS, seq_len=128, micro_batch_size=8,
                       optimizer="muon", seed=SEED, device=DEVICE, compute_dtype="auto", log_every=10_000)
    train(model, train_stream, tcfg)
    loss = evaluate(model, valid_stream, tcfg, batches=8)
    total = cfg.parameter_report()["total"]
    print(f"  {label:<26} parameters {total:>9,} · active "
          f"{cfg.active_parameters():>9,} · valid loss {loss:.4f}")
    return dict(label=label, cfg=cfg, model=model, loss=loss, params=total,
                active=cfg.active_parameters())

torch.manual_seed(1)
d, layers = 256, 36
h0 = torch.randn(1, d)
weights = [torch.randn(d, d) / math.sqrt(d) for _ in range(layers)]
norm = RMSNorm(d, 1e-6)

traces = {}
with torch.no_grad():
    for mode in ("none", "pre", "post"):
        h, trace = h0.clone(), []
        for W in weights:
            if mode == "none":
                h = h + torch.tanh(h) @ W
            elif mode == "pre":
                h = h + torch.tanh(norm(h)) @ W
            else:
                h = norm(h + torch.tanh(h) @ W)
            trace.append(h.pow(2).mean().sqrt().item())
        traces[mode] = trace

plt.figure(figsize=(7, 4))
for mode, color in (("none", "#BF616A"), ("pre", "#5E81AC"), ("post", "#A3BE8C")):
    plt.plot(range(1, layers + 1), traces[mode], label=mode, color=color)
plt.yscale("log"); plt.xlabel("layer"); plt.ylabel("stream rms")
plt.title("Nonlinear residual chain: 36 layers")
plt.legend(); plt.gca().spines[["top", "right"]].set_visible(False)
plt.tight_layout(); plt.show()

for mode, trace in traces.items():
    print(f"{mode:>5}: layer 1 = {trace[0]:8.3f} → layer 36 = {trace[-1]:10.3f}")

def stream_rms(model: MiniLLM) -> list[float]:
    ids = valid_stream.batch(torch.arange(4)).to(DEVICE)
    values = []
    with torch.no_grad():
        x = model.embed_tokens(ids)
        cos, sin = model.rope(torch.arange(ids.shape[1], device=DEVICE))
        for layer in model.layers:
            x = layer(x, cos, sin)
            values.append(x.float().pow(2).mean().sqrt().item())
    return values

print("norm placement ablation:")
norm_runs = [run(f"norm_placement={p}", norm_placement=p)
             for p in ("pre", "post", "peri", "olmo2", "none")]

print("\nresidual stream rms by layer of the trained models:")
for r in norm_runs:
    r["rms"] = stream_rms(r["model"])
    print(f"  {r['label']:<26} " +
          " → ".join(f"{v:.2f}" for v in r["rms"]))

print("attention ablation:")
attn_runs = []
for name, kv in (("MHA (H_kv=4)", 4), ("GQA (H_kv=2)", 2), ("MQA (H_kv=1)", 1)):
    r = run(name, num_key_value_heads=kv)
    r["kv_per_token"] = kv_bytes_per_token(r["cfg"])
    attn_runs.append(r)

print(f"\n{'scheme':<14} {'KV, B/token':>12} {'× vs MHA':>9} {'valid loss':>11} {'Δloss':>8}")
mha = attn_runs[0]
for r in attn_runs:
    print(f"{r['label']:<14} {r['kv_per_token']:>12,} "
          f"{mha['kv_per_token'] / r['kv_per_token']:>8.1f}× "
          f"{r['loss']:>11.4f} {r['loss'] - mha['loss']:>+8.4f}")

d = BASE["hidden_size"]
gated_ff = round(8 * d / 3 / 64) * 64          # the 2/3 rule, aligned to 64
print(f"all variants: 3 · {d} · {gated_ff} = {3 * d * gated_ff:,} FFN parameters\n")

print("ablation of φ inside the GLU family:")
ffn_runs = [run(f"{name} (3 matrices)", hidden_act=act, intermediate_size=gated_ff)
            for name, act in (("SwiGLU", "swiglu"), ("GeGLU", "geglu"),
                              ("ReGLU", "relu"))]

assert len({r["params"] for r in ffn_runs}) == 1, "budgets must match"
best = min(ffn_runs, key=lambda r: r["loss"])
spread = max(r["loss"] for r in ffn_runs) - min(r["loss"] for r in ffn_runs)
print(f"\nbudgets match to the parameter: {ffn_runs[0]['params']:,}")
print(f"best on this run — {best['label']}, loss spread {spread:.4f}")

from minillm.model import MoELayer

MOE_STEPS = STEPS * 2
EXPERTS, TOPK, SHARED = 8, 2, 1
expert_ff = max(64, round(gated_ff / (TOPK + SHARED) / 64) * 64)

print("sparsity ablation (step budget doubled):")
dense = run("dense", steps=MOE_STEPS, intermediate_size=gated_ff)
moe = run(f"MoE {EXPERTS} experts, top-{TOPK}", steps=MOE_STEPS,
          intermediate_size=gated_ff, num_local_experts=EXPERTS,
          num_experts_per_tok=TOPK, num_shared_experts=SHARED,
          moe_intermediate_size=expert_ff, first_k_dense_replace=1)

parity = moe["active"] / dense["active"]
sparsity = moe["params"] / moe["active"]
print(f"\n{'variant':<26} {'total':>12} {'active':>11} {'valid loss':>11}")
for r in (dense, moe):
    print(f"{r['label']:<26} {r['params']:>12,} {r['active']:>11,} {r['loss']:>11.4f}")
print(f"\nMoE active relative to dense: ×{parity:.3f} — active counts; runtime is measured separately")
print(f"capacity: total ×{moe['params'] / dense['params']:.2f}, "
      f"sparsity ×{sparsity:.2f}")
print(f"KV cache matches to the byte: "
      f"{kv_bytes_per_token(dense['cfg']) == kv_bytes_per_token(moe['cfg'])}")

# What about expert load? This is the first thing to look at in an MoE.
def load_table(result, title):
    layers = [l.ffn for l in result["model"].layers if isinstance(l.ffn, MoELayer)]
    print(title)
    print(f"{'layer':>5} {'min':>7} {'max':>7} {'std':>7}  {'max|bias|':>9}")
    for i, layer in enumerate(layers, start=1):
        load = layer.expert_load
        print(f"{i:>5} {load.min():>7.3f} {load.max():>7.3f} {load.std():>7.3f}"
              f"  {layer.expert_bias.abs().max():>9.4f}")
    return max(float(l.expert_load.std()) for l in layers)

ideal = TOPK / EXPERTS
print(f"a perfectly even load would be {ideal:.3f} per expert\n")
std_default = load_table(moe, "default bias step (router_bias_speed = 1e-3):")

moe_fast = run(f"MoE, balance ×50", steps=MOE_STEPS, intermediate_size=gated_ff,
               num_local_experts=EXPERTS, num_experts_per_tok=TOPK,
               num_shared_experts=SHARED, moe_intermediate_size=expert_ff,
               first_k_dense_replace=1, router_bias_speed=5e-2)
print()
std_fast = load_table(moe_fast, "accelerated bias (router_bias_speed = 5e-2):")
print(f"\nmaximum load spread: {std_default:.3f} → {std_fast:.3f}")
print(f"valid loss: {moe['loss']:.4f} → {moe_fast['loss']:.4f} "
      f"(dense {dense['loss']:.4f})")

probe = valid_stream.batch(torch.arange(8)).to(DEVICE)
print(f"{'capacity_factor':>16} {'ceiling':>9} {'dropped':>11}")
capacity_rows = []
for factor in (None, 2.0, 1.5, 1.25, 1.0, 0.75):
    cfg_c = MiniLLMConfig(**{**moe["cfg"].to_dict(), "capacity_factor": factor})
    probe_model = MiniLLM(cfg_c).to(DEVICE)
    probe_model.load_state_dict(moe["model"].state_dict())
    with torch.no_grad():
        probe_model(probe)
    layer = next(l.ffn for l in probe_model.layers if isinstance(l.ffn, MoELayer))
    cap = cfg_c.expert_capacity(probe.numel())
    capacity_rows.append(dict(factor=factor, capacity=cap, dropped=float(layer.dropped_fraction)))
    print(f"{str(factor):>16} {str(cap):>9} {layer.dropped_fraction.item():>10.1%}")

from minillm.probes import RepeatTask, chance_level, repeat_accuracy
from minillm.model import RotaryEmbedding

PROBE_V = 64
probe = RepeatTask(vocab_size=PROBE_V, train_range=(12, 48), seed=SEED)
PROBE_BASE = dict(vocab_size=PROBE_V, hidden_size=128, num_hidden_layers=2,
                  num_attention_heads=4, num_key_value_heads=2, head_dim=32,
                  max_position_embeddings=512)
PROBE_STEPS = 1200 if DEVICE == "cpu" else 3000
if os.getenv("CAPSTONE_FAST_VALIDATION") == "1":
    PROBE_STEPS = 10

probe_budgets = {}

def train_probe(theta: float) -> MiniLLM:
    cfg = MiniLLMConfig.profile("S", **{**PROBE_BASE, "rope_theta": theta})
    torch.manual_seed(SEED)
    model = MiniLLM(cfg).to(DEVICE)
    probe_state = train(model, probe, TrainConfig(steps=PROBE_STEPS, seq_len=64,
                                    micro_batch_size=16, optimizer="muon",
                                    seed=SEED, device=DEVICE, log_every=10_000))
    probe_budgets[str(theta)] = {
        "input_positions": probe_state.tokens,
        "supervised_targets": probe_state.supervised_tokens,
    }
    return model

probe_models = {theta: train_probe(theta) for theta in (10_000.0, 500_000.0)}
print(f"chance level: {chance_level(probe):.4f}")
print("measured RoPE budgets:", probe_budgets)
probe_in_range = {str(theta): repeat_accuracy(model, probe, 32, device=DEVICE)
                  for theta, model in probe_models.items()}
for theta, accuracy in probe_in_range.items():
    print("rope_theta:", theta, "in-range accuracy at 32:", accuracy)

LENGTHS = (48, 64, 80, 96, 128)
rope_rows = {f"base {theta:g}": [repeat_accuracy(m, probe, n, device=DEVICE)
                                 for n in LENGTHS]
             for theta, m in probe_models.items()}

# YaRN on top of an already trained model: RoPE has no parameters, so swapping
# the frequency table is enough — the weights stay as they are. That is the
# "stretching without fine-tuning" of M4. The pretraining window is 96 tokens
# (two halves of 48).
for factor in (2.0, 4.0):
    model = probe_models[10_000.0]
    saved = model.rope
    yarn_cfg = MiniLLMConfig.profile("S", **{**PROBE_BASE, "rope_theta": 10_000.0,
        "rope_scaling": {"rope_type": "yarn", "factor": factor,
                         "original_max_position_embeddings": 96}})
    model.rope = RotaryEmbedding(yarn_cfg).to(DEVICE)
    rope_rows[f"10000 + YaRN ×{factor:g}"] = [
        repeat_accuracy(model, probe, n, device=DEVICE) for n in LENGTHS]
    model.rope = saved

print(f"{'variant':<22}" + " ".join(f"{n:>8}" for n in LENGTHS))
print("-" * 64)
for name, values in rope_rows.items():
    print(f"{name:<22}" + " ".join(f"{v:>8.3f}" for v in values))
print("\ntraining used distances 12–48; everything right of 48 is extrapolation")

plt.figure(figsize=(7.5, 4.5))
palette = ("#BF616A", "#D08770", "#5E81AC", "#A3BE8C")
for (name, values), color in zip(rope_rows.items(), palette):
    plt.plot(LENGTHS, values, "o-", label=name, color=color)
plt.axvline(48, ls="--", lw=1, color="#4C566A")
plt.text(49, 0.5, "training boundary", rotation=90, va="center",
         fontsize=9, color="#4C566A")
plt.axhline(chance_level(probe), ls=":", lw=1, color="#4C566A")
plt.text(120, chance_level(probe) + 0.03, "chance level",
         ha="right", fontsize=9, color="#4C566A")
plt.xlabel("dependency distance (half length)")
plt.ylabel("share of correctly restored tokens")
plt.title('Accuracy across dependency distances')
plt.legend(); plt.ylim(-0.03, 1.05)
plt.gca().spines[["top", "right"]].set_visible(False)
plt.tight_layout(); plt.show()

def run_optimizer(optimizer: str, lr: float):
    cfg = MiniLLMConfig.profile("S", **BASE)
    torch.manual_seed(SEED)
    model = MiniLLM(cfg).to(DEVICE)
    tcfg = TrainConfig(steps=STEPS, seq_len=128, micro_batch_size=8,
                       optimizer=optimizer, lr=lr if optimizer == "adamw" else 3e-3,
                       muon_lr=lr if optimizer == "muon" else 0.02,
                       seed=SEED, device=DEVICE, log_every=10_000)
    train(model, train_stream, tcfg)
    return dict(model=model, config=tcfg, valid_loss=evaluate(model, valid_stream, tcfg, batches=8))

grid = {
    "adamw": [3e-3, 1e-2],
    "muon":  [2e-2, 6e-2],
}
optimizer_report = {}
for optimizer, lrs in grid.items():
    runs = {lr: run_optimizer(optimizer, lr) for lr in lrs}
    chosen_lr, selected = min(runs.items(), key=lambda kv: kv[1]["valid_loss"])
    test_loss = evaluate(selected["model"], test_stream, selected["config"], batches=8)
    optimizer_report[optimizer] = dict(selected_lr=chosen_lr,
        validation_grid={str(lr): r["valid_loss"] for lr,r in runs.items()},
        selected_valid_loss=selected["valid_loss"], selected_test_loss=test_loss)
    print(optimizer, optimizer_report[optimizer])

CAPSTONE_LANGUAGE = 'en'
"""Embedded into both ablation notebooks; reads only this run's measurements."""
import json
from pathlib import Path

ablation_record = {
    "schema": 2,
    "language": CAPSTONE_LANGUAGE,
    "mode": "fast_execution_check" if os.getenv("CAPSTONE_FAST_VALIDATION") == "1" else "teaching_experiment",
    "seed": SEED, "device": DEVICE, "torch": torch.__version__,
    "steps": STEPS, "moe_steps": MOE_STEPS, "probe_steps": PROBE_STEPS,
    "data": {"train_sha": train_stream.checksum(), "valid_sha": valid_stream.checksum(),
             "test_sha": test_stream.checksum()},
    "chain_rms": traces,
    "norm": [{"label": r["label"], "valid_loss": r["loss"], "rms": r["rms"]} for r in norm_runs],
    "attention": [{"label": r["label"], "valid_loss": r["loss"],
                   "calculated_kv_bytes_per_token_bf16": r["kv_per_token"]} for r in attn_runs],
    "glu": [{"label": r["label"], "valid_loss": r["loss"], "parameters": r["params"]} for r in ffn_runs],
    "moe": {"dense_total": dense["params"], "moe_total": moe["params"],
            "dense_active": dense["active"], "moe_active": moe["active"],
            "active_ratio": parity, "dense_loss": dense["loss"], "moe_loss": moe["loss"],
            "faster_balance_loss": moe_fast["loss"],
            "last_batch_load_std": std_default, "faster_balance_load_std": std_fast},
    "capacity": capacity_rows,
    "rope": {"distances": list(LENGTHS), "accuracy": rope_rows,
             "training_budgets": probe_budgets,
             "in_range_32": probe_in_range, "chance": chance_level(probe)},
    "optimizers": optimizer_report,
    "not_tested": ["ungated FFN", "packed quantization speed", "RL answer quality",
                   "scaling-law optimum", "QK-norm ablation",
                   "speculative distribution test in this notebook", "EOT ablation in this notebook"],
}
report_dir = Path("runs/ablations")
report_dir.mkdir(parents=True, exist_ok=True)
report_path = report_dir / "measurements.json"
report_path.write_text(json.dumps(ablation_record, ensure_ascii=False, indent=2, allow_nan=False), encoding="utf-8")
saved = json.loads(report_path.read_text(encoding="utf-8"))
assert saved == ablation_record
ru = CAPSTONE_LANGUAGE == "ru"
mode_ru = {"fast_execution_check": "быстрая проверка исполнения", "teaching_experiment": "учебный эксперимент"}
print("Режим:" if ru else "Mode:", mode_ru[saved["mode"]] if ru else saved["mode"])
print("Измерения сохранены:" if ru else "Measurements saved:", report_path)
labels_ru = {"chain_rms": "RMS остаточной цепочки", "norm": "Нормализация",
             "attention": "Внимание", "glu": "Варианты GLU", "moe": "Эксперты",
             "capacity": "Ограничение нагрузки", "rope": "Позиции", "optimizers": "Оптимизаторы"}
for key in ("chain_rms", "norm", "attention", "glu", "moe", "capacity", "rope", "optimizers"):
    print("\n" + (labels_ru[key] if ru else key))
    print(json.dumps(saved[key], ensure_ascii=False, indent=2))
print("\nЗдесь не проверялось:" if ru else "\nNot tested here:")
not_tested_ru = {"ungated FFN": "FFN без гейта",
    "packed quantization speed": "скорость упакованной квантизации",
    "RL answer quality": "качество ответов после RL",
    "scaling-law optimum": "оптимум закона масштабирования",
    "QK-norm ablation": "абляция QK-нормировки",
    "speculative distribution test in this notebook": "распределение спекулятивного семплирования в этой тетради",
    "EOT ablation in this notebook": "абляция EOT в этой тетради"}
for item in saved["not_tested"]:
    print("-", not_tested_ru[item] if ru else item)