Модуль 20. LLM целиком: от токенов до обучения и генерации

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="Позиция ключа", ylabel="Позиция запроса", title="Причинное внимание")
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="Шаг обучения", ylabel="Средняя потеря, нат", title="Обучение на одной последовательности")
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("Температура должна быть положительным конечным числом")
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("Температура:", experiment_temperature, "top-k:", experiment_top_k)
print("Токен   p модели   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" / "ru", search_root / "minillm" / "ru"):
        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 "пакет проекта: " + found

    try:
        import minillm  # noqa: F401
        return "minillm уже доступен"
    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("Загрузите course-materials.zip со страницы ноутбуков курса. Структура каталогов будет восстановлена автоматически.")
            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 "архив проекта: " + 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 "архив проекта: " + str(archive_path)

    raise RuntimeError("minillm не найден. Скачайте course-materials.zip со страницы ноутбуков курса и либо распакуйте его целиком, либо загрузите ZIP в текущую сессию Colab.")

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())

# ---------------------------------------------------------------- профиль забега
PROFILE = "SMOKE"          # "SMOKE" | "FULL"
USE_REAL_CORPUS = False    # сетевая загрузка настоящих данных — см. §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, "на", DEVICE, "с вычислениями", COMPUTE_DTYPE, "->", 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]:
    """Двуязычный корпус со слотовой грамматикой и намеренно внесённым мусором.

    Никакой это не претрейн-корпус: грамматика одна, словарь узкий. Задача —
    дать конвейеру и токенизатору воспроизводимый, заведомо лицензионно чистый
    материал, на котором виден каждый механизм. Слоты дают комбинаторное
    разнообразие, иначе near-dup дедуп справедливо схлопнул бы весь корпус.
    """
    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))))
    # Мусор для воронки: точные копии, почти-копии, короткие и символьные.
    docs += docs[:n_docs // 20]                                   # точные дубли
    docs += [d + " Дополнение." for d in docs[:n_docs // 40]]     # почти-дубли
    docs += ["ок"] * (n_docs // 40)                               # короткие
    docs += ["!@#$%^&*()" * 8] * (n_docs // 40)                   # символьный мусор
    rng.shuffle(docs)
    return docs

raw_docs = synthetic_corpus(SETTINGS["docs"], seed=SEED)
print(f"сырых документов: {len(raw_docs):,}")

clean_docs, report = pipeline(raw_docs, FilterConfig(min_chars=60),
                              near_dup_threshold=0.85)
print(report)
print(f"\nитог: {len(clean_docs):,} документов, "
      f"{sum(len(d) for d in clean_docs):,} символов")

# Разделяем ДО обучения tokenizer: validation не участвует ни в словаре,
# ни в merge table. Глобальная очистка выполнена до split, поэтому точные и
# near-дубли не расходятся между train и 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)
    # Защита от точной утечки между источниками/splits после нормализации.
    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("valid:", valid_report)
    print("revisions:", DATA_REVISIONS)
    print("\nАтрибуция: FineWeb2 (ODC-By 1.0 + Common Crawl Terms), "
          "TinyStories (CDLA-Sharing-1.0). Это техническое описание, не юридическая консультация.")
else:
    DATA_REVISIONS = {"synthetic": f"generator-v1-seed-{SEED}"}
    print("USE_REAL_CORPUS = False — используем document-level split синтетического корпуса")

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

# Управляющие токены чата резервируются в словаре СРАЗУ, хотя претрейну они не
# нужны. Причина ровно та же, по которой в o200k_harmony лежат 1 075
# зарезервированных идентификаторов (M2): добавить их потом — значит менять
# vocab_size и пересобирать эмбеддинги. Ноутбук 2 берёт этот же чекпойнт и
# рассчитывает найти в нём `<|eot_id|>`.
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"обучено за {time.time() - t0:.1f} с: словарь {tokenizer.vocab_size}, "
      f"merge'ей {len(tokenizer.merges)}")
if tokenizer.vocab_size < SETTINGS["vocab_size"]:
    print('Доступные пары для слияний закончились до заданного размера словаря.')

sample = "Языковая модель предсказывает следующий токен."
assert tokenizer.decode(tokenizer.encode(sample)) == sample, "round-trip сломан"
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"{'язык':<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 для кириллицы даёт ~2 токена на символ — это и есть цена "
      "отсутствия merge'ей для языка (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):,}; valid 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'])} шардов, {index['total_tokens']:,} токенов, "
      f"{len(train_stream):,} окон, checksum {train_stream.checksum()}")
print(f"valid: {len(valid_stream):,} окон — документно отделён до 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(), "формула и код разошлись"
print(f"параметров: {analytic['total']:,} (формула == факт)")
for key in ("embeddings", "attention", "ffn", "norms"):
    print(f"  {key:<12} {analytic[key]:>12,}  ({analytic[key] / analytic['total']:.1%})")
print(f"  блок стоит {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-кэш, BF16/FP16 reference: {reference_cache_bpt} байт/токен "
      f"→ {reference_cache_bpt * SETTINGS['seq_len'] / 1024:.1f} KiB на окно")
print(f"KV-кэш при текущем compute dtype ({COMPUTE_DTYPE}): {runtime_cache_bpt} байт/токен")
print("постоянная память обучения:", 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"плановых next-token целей за шаг: {tcfg.tokens_per_step:,}")
print(f"запланировано:  {planned_tokens:,} обучаемых целей за {tcfg.steps:,} шагов")
print(f"reference 20×: {chinchilla_tokens:,} токенов "
      f"→ {chinchilla_tokens / tcfg.tokens_per_step:,.0f} шагов")
print(f"доля бюджета:   {planned_tokens / chinchilla_tokens:.1%}")
print(f"эпох по корпусу: {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"\nготово: {state.step} шагов, "
      f"{state.tokens:,} входных позиций, {state.supervised_tokens:,} целей, "
      f"финальный лосс {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) — равномерная модель")
axes[0].set_xlabel("шаг"); axes[0].set_ylabel("кросс-энтропия"); axes[0].legend()
axes[0].set_title("лосс")

axes[1].plot(steps, [h["lr_scale"] for h in state.history], color="#A3BE8C")
axes[1].set_xlabel("шаг"); axes[1].set_ylabel("множитель lr")
axes[1].set_title("WSD-расписание")

axes[2].plot(steps, [h["tokens_per_s"] for h in state.history], color="#D08770")
axes[2].set_xlabel("шаг"); axes[2].set_ylabel("входных позиций/с")
axes[2].set_title("пропускная способность")

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"финальная кросс-энтропия {state.history[-1]['loss']:.4f} нат "
      f"= {bpc:.4f} бит на токен")
print('Кривая относится к данному бюджету и корпусу; её плато само по себе не определяет причину ограничения.')

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"{'вариант':<8} {'total':>14} {'active':>12} {'разреж.':>8} "
      f"{'KV Б/ток':>9} {'FLOP/ток':>12}")
for name, c, m in rows:
    assert c.parameter_report()["total"] == m.num_parameters(), "формула разошлась с моделью"
    total, active = c.parameter_report()["total"], c.active_parameters()
    print(f"{name:<8} {total:>14,} {active:>12,} {total / active:>7.2f}× "
          f"{kv_bytes_per_token(c):>9,} {c.flops_per_token():>12,}")

print(f"\nэкспертов {moe_cfg.num_local_experts}, top-{moe_cfg.num_experts_per_tok}, "
      f"shared {moe_cfg.num_shared_experts}, первый слой плотный: "
      f"{not moe_cfg.moe_layer(0)}")
print(f"активных у MoE к плотному: "
      f"×{moe_cfg.active_parameters() / dense_cfg.active_parameters():.3f}")
print(f"KV-кэш совпадает до байта: "
      f"{kv_bytes_per_token(dense_cfg) == kv_bytes_per_token(moe_cfg)} — "
      f"внимание не трогали, сравнение чистое")

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{'вариант':<8} {'train':>9} {'valid':>9} {'позиций/с':>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"\nнагрузка экспертов последнего замера: min {load.min():.3f}, "
      f"max {load.max():.3f}, std {load.std():.3f}")
print(f"идеально равномерная была бы "
      f"{moe_cfg.num_experts_per_tok / moe_cfg.num_local_experts:.3f} на эксперта")
print(f"смещений aux-loss-free: 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("Языковая модель")
    out, stats = generate(model, prompt, 80,
                          SamplingConfig(temperature=temperature, top_p=top_p, seed=SEED),
                          stop_tokens={EOS})
    label = "жадно" if temperature == 0 else f"T={temperature}, top_p={top_p}"
    print(f"--- {label}\n{tokenizer.decode(out)!r}\n")
print("статистика последнего запуска:", 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("сохранено:", sorted(p.name for p in run_dir.iterdir()))

import json

def read_passport(path: Path) -> None:
    """Паспорт модели по config.json — тот же приём, что в модуле 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("Ось I  · размер     :", f"d={c['hidden_size']}, L={c['num_hidden_layers']}, "
          f"V={c['vocab_size']}, эмбеддинги "
          f"{'связаны' if c['tie_word_embeddings'] else 'раздельны'}")
    print("Ось II · контекст    :", f"{c['max_position_embeddings']} токенов, "
          f"RoPE base {c['rope_theta']:g}, растяжка "
          f"{c['rope_scaling'] or 'нет'}")
    print("Ось III· внимание    :", attention,
          f"· head_dim {c['head_dim']} · sinks {'да' if c['attention_sinks'] else 'нет'}"
          f" · окно {c['sliding_window'] or 'нет'}")
    print("Ось IV · блок        :", f"норма {c['norm_placement']}, {c['hidden_act']}, "
          f"d_ffn/d = {ffn_ratio:.2f}"
          + (f", кламп {c['swiglu_limit']}" if c["swiglu_limit"] else ""))
    print("Ось V  · разреженность:", "dense" if not c["num_local_experts"]
          else f"MoE {c['num_local_experts']} экспертов, top-{c['num_experts_per_tok']}")
    print("Производные          :", f"KV-кэш {kv_bytes} Б/токен в измеренном {stats.cache_dtype}, "
          f"{kv_bytes * c['max_position_embeddings'] / 1024 ** 2:.1f} МиБ на полный контекст")

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" / "ru", search_root / "minillm" / "ru"):
        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 "пакет проекта: " + found

    try:
        import minillm  # noqa: F401
        return "minillm уже доступен"
    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("Загрузите course-materials.zip со страницы ноутбуков курса. Структура каталогов будет восстановлена автоматически.")
            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 "архив проекта: " + 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 "архив проекта: " + str(archive_path)

    raise RuntimeError("minillm не найден. Скачайте course-materials.zip со страницы ноутбуков курса и либо распакуйте его целиком, либо загрузите ZIP в текущую сессию Colab.")

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

# Профиль забега. SMOKE рассчитан на CPU и укладывается в минуты; FULL имеет
# смысл только на GPU и даёт заметно более связные ответы.
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)

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

RU_SLOTS = dict(
    subject=["языковая модель", "трансформер", "механизм внимания", "токенизатор",
             "оптимизатор", "нормализация", "остаточная связь", "роутер экспертов"],
    verb=["предсказывает", "обрабатывает", "сокращает", "ускоряет", "стабилизирует"],
    obj=["следующий токен", "длинный контекст", "стоимость шага", "объём памяти",
         "масштаб активаций", "нагрузку экспертов"],
)

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, RU_SLOTS) for _ in range(4))
                      for _ in range(4000))

def usable_checkpoint(path: Path):
    """Чекпойнт годится, только если его словарь знает управляющие токены чата.

    Словарь без `<|eot_id|>` уронил бы ноутбук на первой же маске лосса, причём
    с невнятным KeyError. Ноутбук 1 резервирует эти идентификаторы заранее, но
    чекпойнт мог быть собран старой версией — поэтому проверяем явно.
    """
    if not (path / "config.json").exists():
        return None, None, "чекпойнта ноутбука 1 рядом нет"
    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"в словаре чекпойнта нет токенов {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("загружен чекпойнт ноутбука 1:", base_model.num_parameters(), "параметров")
else:
    print(f"учим базовую модель сами: {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"параметров: {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("словарь:", tokenizer.vocab_size)

from minillm.chat import ChatTemplate, Message, SFTDataset

# Формулировки подобраны так, чтобы согласование работало для подлежащих
# любого рода: «языковая модель», «трансформер», «остаточная связь».
QUESTIONS = [
    ("Что делает {s}?", "{s} {v} {o}."),
    ("Для чего служит {s}?", "{s} {v} {o}. Это и есть основная функция."),
    ("Как влияет {s} на модель?", "{s} {v} {o} и тем самым меняет поведение модели."),
]

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(RU_SLOTS["subject"])
        v, o = rng.choice(RU_SLOTS["verb"]), rng.choice(RU_SLOTS["obj"])
        out.append([
            Message("system", "Ты отвечаешь кратко и по делу."),
            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"токенов: {len(ids)}, обучаемых: {sum(mask)} "
      f"({template.trainable_fraction(mask):.1%})\n")

print(f"{'токен':<22} {'id':>6}  учим")
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}  {'ДА' if flag else '—'}")

dataset, stats = SFTDataset.from_dialogues(dialogues, tokenizer, template,
                                           seq_len=BUDGET["seq_len"], pad_id=PAD_ID)
print("упаковка:", stats)
print(f"\nбез упаковки понадобилось бы {stats['dialogues']} строк по "
      f"{BUDGET['seq_len']} токенов,")
print(f"с упаковкой — {stats['bins']}: экономия ×{stats['packing_ratio']}, "
      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", "Ты отвечаешь кратко и по делу."),
         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 = "Что делает токенизатор?"
for label, model in (("до SFT", base_model), ("после SFT", sft_model)):
    answer, n = ask(model, question)
    print(f"--- {label} ({n} токенов)\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"последовательности идентичны: {a_ids == b_ids}")
print(f"маски различаются в позициях {diff} — это {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"{'модель':<16} {'p(<|eot_id|>) после ответа':>28} {'длина ответа':>14}")
for label, model in (("маска верная", sft_model), ("EOT забыт", broken_model)):
    ids, _ = template.encode_with_mask(dialogues[0], tokenizer)
    prefix = ids[:len(ids) - 1]        # всё, кроме закрывающего 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, "Что делает токенизатор?", 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, "Что делает токенизатор?", max_new_tokens=120)
print(f"\nответ модели с забытым EOT ({n} токенов):\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 = " Это важно и нужно учитывать при работе с моделью."

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(RU_SLOTS["subject"])
        v, o = rng.choice(RU_SLOTS["verb"]), rng.choice(RU_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"пар: {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 — старт")
axes[0].set_title("DPO-лосс"); 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: разность неявных наград")
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("неявные награды"); axes[2].legend()
for ax in axes:
    ax.set_xlabel("шаг"); 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 = "Для чего служит оптимизатор?"
for label, model in (("после SFT", sft_model), ("после DPO", policy)):
    answer, n = ask(model, question, max_new_tokens=80, temperature=0.7)
    print(f"--- {label} ({n} токенов)\n{answer!r}\n")
print('Два ответа — иллюстрация. Для вывода об изменении поведения сравните несколько фиксированных запросов и сохраните исходы.')

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("сохранено:", 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" / "ru", search_root / "minillm" / "ru"):
        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 "пакет проекта: " + found

    try:
        import minillm  # noqa: F401
        return "minillm уже доступен"
    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("Загрузите course-materials.zip со страницы ноутбуков курса. Структура каталогов будет восстановлена автоматически.")
            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 "архив проекта: " + 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 "архив проекта: " + str(archive_path)

    raise RuntimeError("minillm не найден. Скачайте course-materials.zip со страницы ноутбуков курса и либо распакуйте его целиком, либо загрузите ZIP в текущую сессию Colab.")

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, "| вычисления:", 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", "Ты отвечаешь кратко и по делу."),
         Message("user", "Что делает токенизатор?")],
        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("ответ:", repr(handoff_tokenizer.decode(handoff_out[len(handoff_ids):])))
    print("KV bytes:", handoff_stats.kv_bytes)
else:
    print("post-training артефакт не найден — self-contained маршрут продолжается")

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:
    """Псевдоязык: большой лексикон, зипфовское распределение частот.

    Неравномерные частоты задают контролируемую задачу предсказания.
    Loss измеряется на BPE-токен, поэтому энтропия распределения слов
    не является готовым эталоном для этой величины.
    """
    rng = random.Random(seed)
    alphabet = "абвгдеиклмнопрстуфхцчшы"
    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)]   # зипфовский хвост
    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"корпус {len(corpus):,} символов · словарь {tokenizer.vocab_size} · "
      f"{len(stream):,} окон · sha {stream.checksum()}")
print(f"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)

# Бюджет рассчитан на CPU. Увеличение числа шагов — отдельный опыт:
# после него заново измеряются loss и чувствительность к квантизации.
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():,} параметров")
    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"{'промпт':>7} {'prefill, мс':>12} {'мс/токен decode':>17} {'KV-кэш, KiB':>13}")
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:>13.1f}")

print('Сравните измеренные фазы; для определения узкого места нужны повторности и профилирование.')

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"{'схема':>6} {'H_kv':>5} {('байт/токен ' + BASELINE_COMPUTE):>20} {'на 128K контекста':>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} ГиБ")

# Сверка формулы с фактом на нашей модели.
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"\nфактический dtype кэша: {str(cache.dtype()).replace('torch.', '')}")
print(f"факт {cache.nbytes():,} Б == формула {expected:,} Б: {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"результаты совпадают: {identical}")
print(f"prefill без кэша:  {cold_total * 1000:6.1f} мс на {len(questions)} запросов")
print(f"prefill с кэшем:   {hot_total * 1000:6.1f} мс (+{warmup_ms:.1f} мс прогрева один раз)")
print(f"экономия при {len(questions)} запросах: "
      f"×{cold_total / max(hot_total, 1e-9):.1f} на фазе prefill")

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"{'формат':>7} {'эфф. бит/вес':>12} {'сжатие':>8} {'лосс':>8} {'Δлосс':>8} "
      f"{'KL, нат':>9} {'top-1 совп.':>12}")
print(f"{BASELINE_LABEL:>7} {BASELINE_STORAGE_BITS:>12.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']:>12.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("лосс (обрезанная ось!)"); axes[0].set_ylabel("нат")

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("нат")

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 с 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} {'приём':>8} {'токенов/проход target':>24} {'проходов draft':>16}")
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:>24.2f} {stats.draft_forwards:>16}")

# Черновик, равный target: правило min(1, p/q) обязано принимать всё.
_, twin = speculative_generate(target, copy.deepcopy(target), prompt, 32, lookahead=4,
                               sampling=SamplingConfig(temperature=1.0, seed=SEED))
print(f"черновик == target: приём {twin.acceptance_rate:.1%}, "
      f"{twin.tokens_per_target_forward:.2f} токена на проход")

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("глубина черновика k"); axes[0].set_ylabel("доля приёма")
axes[0].set_title("измеренная доля принятия")
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="без спекулятива")
axes[1].set_xlabel("глубина черновика k")
axes[1].set_ylabel("токенов на проход target")
axes[1].set_title("выдача на проверочный проход"); 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" / "ru", search_root / "minillm" / "ru"):
        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 "пакет проекта: " + found

    try:
        import minillm  # noqa: F401
        return "minillm уже доступен"
    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("Загрузите course-materials.zip со страницы ноутбуков курса. Структура каталогов будет восстановлена автоматически.")
            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 "архив проекта: " + 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 "архив проекта: " + str(archive_path)

    raise RuntimeError("minillm не найден. Скачайте course-materials.zip со страницы ноутбуков курса и либо распакуйте его целиком, либо загрузите ZIP в текущую сессию Colab.")

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     # на GPU разница между режимами чётче
print("устройство:", DEVICE, "| шагов на прогон:", 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:
    """Тот же псевдоязык, что в ноутбуке 3: зипфовский лексикон оставляет
    моделям запас, в котором различимы архитектурные варианты."""
    rng = random.Random(seed)
    alphabet = "абвгдеиклмнопрстуфхцчшы"
    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"словарь {tokenizer.vocab_size} · train {len(train_stream):,} окон · "
      f"valid {len(valid_stream):,} окон · равномерный ориентир 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):
    """Один прогон абляции: фиксированный seed, данные и расписание."""
    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} параметров {total:>9,} · активных "
          f"{cfg.active_parameters():>9,} · valid-лосс {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("слой"); plt.ylabel("rms потока")
plt.title("Нелинейная остаточная цепочка: 36 слоёв")
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}: слой 1 = {trace[0]:8.3f} → слой 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_runs = [run(f"norm_placement={p}", norm_placement=p)
             for p in ("pre", "post", "peri", "olmo2", "none")]

print("\nrms residual-потока по слоям обученных моделей:")
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("абляция внимания:")
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{'схема':<14} {'KV, Б/токен':>12} {'×к MHA':>8} {'valid-лосс':>11} {'Δлосс':>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']:>7.1f}× "
          f"{r['loss']:>11.4f} {r['loss'] - mha['loss']:>+8.4f}")

d = BASE["hidden_size"]
gated_ff = round(8 * d / 3 / 64) * 64          # правило 2/3, выровнено под 64
print(f"все варианты: 3 · {d} · {gated_ff} = {3 * d * gated_ff:,} параметров FFN\n")

print("абляция φ внутри GLU-семейства:")
ffn_runs = [run(f"{name} (3 матрицы)", 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, "бюджеты обязаны совпадать"
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"\nбюджеты совпадают до параметра: {ffn_runs[0]['params']:,}")
print(f"лучший на этом прогоне — {best['label']}, разброс по лоссу {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("абляция разреженности (бюджет шагов удвоен):")
dense = run("dense", steps=MOE_STEPS, intermediate_size=gated_ff)
moe = run(f"MoE {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{'вариант':<26} {'total':>12} {'active':>11} {'valid-лосс':>11}")
for r in (dense, moe):
    print(f"{r['label']:<26} {r['params']:>12,} {r['active']:>11,} {r['loss']:>11.4f}")
print(f"\nактивных у MoE к плотному: ×{parity:.3f} — активные параметры; время измеряется отдельно")
print(f"ёмкость: total ×{moe['params'] / dense['params']:.2f}, "
      f"разрежённость ×{sparsity:.2f}")
print(f"KV-кэш совпадает до байта: "
      f"{kv_bytes_per_token(dense['cfg']) == kv_bytes_per_token(moe['cfg'])}")

# Что с нагрузкой экспертов? Это первая вещь, которую надо смотреть у MoE.
def load_table(result, title):
    layers = [l.ffn for l in result["model"].layers if isinstance(l.ffn, MoELayer)]
    print(title)
    print(f"{'слой':>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"идеально равномерная нагрузка — {ideal:.3f} на эксперта\n")
std_default = load_table(moe, "смещение по умолчанию (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, "ускоренное смещение (router_bias_speed = 5e-2):")
print(f"\nмаксимальный разброс нагрузки: {std_default:.3f} → {std_fast:.3f}")
print(f"valid-лосс: {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} {'потолок':>9} {'отброшено':>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(probe):.4f}")
print("измеренный бюджет RoPE:", 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 поверх уже обученной модели: у RoPE нет параметров, поэтому достаточно
# заменить таблицу частот — веса остаются те же. Это и есть «растяжка без
# дообучения» из M4. Окно претрейна равно 96 токенам (две половины по 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"{'вариант':<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("\nобучение шло на расстояниях 12–48; всё правее 48 — экстраполяция")

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, "граница обучения", 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, "случайное угадывание",
         ha="right", fontsize=9, color="#4C566A")
plt.xlabel("расстояние зависимости (длина половины)")
plt.ylabel("доля верно восстановленных токенов")
plt.title('Точность по расстоянию зависимости')
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 = 'ru'
"""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)