注意力
与上下文

tokens = ["小王", "把", "书", "放在", "桌上"]
vectors = {
    "小王": [1.0, 0.0],
    "把": [0.0, 1.0],
    "书": [0.8, 0.2],
    "放在": [0.4, 0.9],
    "桌上": [0.2, 1.0],
}
for index, token in enumerate(tokens):
    left = tokens[max(0, index - 2):index]
    right = tokens[index + 1:index + 3]
    print(token, "can inspect:", left + right)

import math
import numpy as np

query = np.array([1.0, 0.0])
keys = np.array([[1.0, 0.0], [0.0, 1.0], [0.6, 0.8]])
values = np.array([[10.0, 0.0], [0.0, 10.0], [4.0, 4.0]])

scores = keys @ query / math.sqrt(query.size)
scores = scores - scores.max()
weights = np.exp(scores) / np.exp(scores).sum()
context = weights @ values

print("scores:", scores)
print("weights:", weights, "sum:", weights.sum())
print("context:", context)

def transformer_block(x, self_attention, feed_forward, norm1, norm2):
    attended = self_attention(x)
    x = norm1(x + attended)
    transformed = feed_forward(x)
    x = norm2(x + transformed)
    return x

# x: [batch, tokens, hidden]
# self_attention: mixes tokens
# feed_forward: transforms each token independently

def split_heads(x, heads):
    batch, tokens, dim = x.shape
    if dim % heads:
        raise ValueError("hidden 维度必须能被 heads 整除")
    head_dim = dim // heads
    # [B,T,D] -> [B,H,T,d]
    return x.reshape(batch, tokens, heads, head_dim).transpose(1, 2)

def merge_heads(x):
    batch, heads, tokens, head_dim = x.shape
    # [B,H,T,d] -> [B,T,D]
    return x.transpose(1, 2).contiguous().reshape(
        batch, tokens, heads * head_dim
    )

class FakeTensor:
    def __init__(self, shape):
        self.shape = shape

q = FakeTensor((2, 6, 32))
print("输入:", q.shape, "预期 head_dim:", q.shape[-1] // 4)

import torch
from torch import nn

batch, tokens, hidden, head_dim = 2, 5, 16, 4
x = torch.randn(batch, tokens, hidden)
q_proj = nn.Linear(hidden, head_dim, bias=False)
k_proj = nn.Linear(hidden, head_dim, bias=False)
v_proj = nn.Linear(hidden, head_dim, bias=False)

q, k, v = q_proj(x), k_proj(x), v_proj(x)
scores = q @ k.transpose(-2, -1)
print("q/k/v:", q.shape, k.shape, v.shape)
print("scores:", scores.shape)

import math
import torch

def attention(q, k, v, mask=None):
    depth = q.size(-1)
    scores = q @ k.transpose(-2, -1) / math.sqrt(depth)
    if mask is not None:
        scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
    weights = torch.softmax(scores, dim=-1)
    context = weights @ v
    return context, weights

q = torch.tensor([[[[1.0, 0.0]]]])
k = torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]])
v = torch.tensor([[[[10.0, 0.0], [0.0, 10.0]]]])
mask = torch.tensor([[[[True, False]]]])
context, weights = attention(q, k, v, mask)
print(context, weights, weights.sum(dim=-1))

import torch

def causal_mask(length):
    return torch.tril(torch.ones(length, length, dtype=torch.bool))

def padding_mask(input_ids, pad_id):
    # [B,T] -> [B,1,1,T]，用于广播到查询和 head
    return (input_ids != pad_id)[:, None, None, :]

ids = torch.tensor([[2, 4, 5, 0], [2, 6, 0, 0]])
print("causal:\n", causal_mask(4))
print("padding:\n", padding_mask(ids, pad_id=0).squeeze(1).squeeze(1))

import math

def softmax(row):
    top = max(row)
    exps = [math.exp(value - top) for value in row]
    total = sum(exps)
    return [value / total for value in exps]

scores = [1.0, 2.0, -1e9]  # 第三个 Key 被 mask
weights = softmax([value / math.sqrt(4) for value in scores])
values = [[10, 0], [0, 10], [99, 99]]
output = [
    sum(weight * values[index][dim] for index, weight in enumerate(weights))
    for dim in range(2)
]
print(weights, output, sum(weights))

def causal_allowed(query, key):
    return key <= query

def visible_columns(query, length, padding=None):
    padding = padding or set()
    return [
        key for key in range(length)
        if causal_allowed(query, key) and key not in padding
    ]

for query in range(4):
    print(query, visible_columns(query, 4, padding={3}))

import torch

def split_heads(x, num_heads):
    batch, tokens, model_dim = x.shape
    if model_dim % num_heads:
        raise ValueError("model_dim 必须能被 num_heads 整除")
    head_dim = model_dim // num_heads
    x = x.view(batch, tokens, num_heads, head_dim)
    return x.transpose(1, 2)

def merge_heads(x):
    batch, heads, tokens, head_dim = x.shape
    return x.transpose(1, 2).contiguous().view(batch, tokens, heads * head_dim)

x = torch.randn(2, 5, 16)
heads = split_heads(x, 4)
print(heads.shape, merge_heads(heads).shape)
assert torch.allclose(x, merge_heads(heads))

import torch
from torch import nn

class ShapeProbe(nn.Module):
    def __init__(self, model_dim=16, heads=4):
        super().__init__()
        self.heads = heads
        self.head_dim = model_dim // heads
        self.qkv = nn.Linear(model_dim, 3 * model_dim)
        self.out = nn.Linear(model_dim, model_dim)

    def forward(self, x):
        batch, tokens, dim = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q = q.view(batch, tokens, self.heads, self.head_dim).transpose(1, 2)
        k = k.view(batch, tokens, self.heads, self.head_dim).transpose(1, 2)
        v = v.view(batch, tokens, self.heads, self.head_dim).transpose(1, 2)
        scores = q @ k.transpose(-2, -1)
        context = scores.softmax(dim=-1) @ v
        merged = context.transpose(1, 2).contiguous().view(batch, tokens, dim)
        return self.out(merged), scores

probe = ShapeProbe()
output, scores = probe(torch.randn(2, 5, 16))
print(output.shape, scores.shape)

import math
import torch

def sinusoidal_positions(length, dim):
    position = torch.arange(length, dtype=torch.float32).unsqueeze(1)
    index = torch.arange(0, dim, 2, dtype=torch.float32)
    rates = torch.exp(-math.log(10000.0) * index / dim)
    table = torch.zeros(length, dim)
    table[:, 0::2] = torch.sin(position * rates)
    table[:, 1::2] = torch.cos(position * rates)
    return table

pos = sinusoidal_positions(6, 8)
print(pos.shape)
print("distance 0/1:", torch.cosine_similarity(pos[0], pos[1], dim=0))

import math

def sinusoidal(position, dim):
    values = []
    for index in range(dim):
        rate = 1 / (10000 ** (2 * (index // 2) / dim))
        value = math.sin(position * rate) if index % 2 == 0 else math.cos(position * rate)
        values.append(value)
    return values

for position in (0, 1, 8):
    vector = sinusoidal(position, 6)
    print(position, [round(value, 4) for value in vector])

def head_shapes(batch, tokens, hidden, heads):
    if hidden % heads != 0:
        raise ValueError("hidden 必须整除 heads")
    head_dim = hidden // heads
    return {
        "qkv": (batch, heads, tokens, head_dim),
        "scores": (batch, heads, tokens, tokens),
        "merged": (batch, tokens, hidden),
    }

for heads in (1, 2, 4, 8):
    print(heads, head_shapes(2, 16, 64, heads))

import torch
from torch import nn

class EncoderLayerProbe(nn.Module):
    def __init__(self, dim=16, heads=4, ff_dim=64):
        super().__init__()
        self.attn = nn.MultiheadAttention(dim, heads, batch_first=True)
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.ffn = nn.Sequential(
            nn.Linear(dim, ff_dim),
            nn.GELU(),
            nn.Linear(ff_dim, dim),
        )

    def forward(self, x, key_padding_mask=None):
        attended, weights = self.attn(
            x, x, x, key_padding_mask=key_padding_mask,
            need_weights=True,
        )
        x = self.norm1(x + attended)
        x = self.norm2(x + self.ffn(x))
        return x, weights

x = torch.randn(2, 5, 16)
pad = torch.tensor([[False, False, False, True, True],
                    [False, False, False, False, True]])
y, weights = EncoderLayerProbe()(x, pad)
print("output:", y.shape, "weights:", weights.shape)

import torch
from torch import nn

x = torch.randn(2, 4, 8)
branch = torch.randn_like(x) * 0.1
pre_norm = nn.LayerNorm(8)
post_norm = nn.LayerNorm(8)

pre = x + branch
pre = pre_norm(pre)
post = post_norm(x + branch)

for name, value in [("input", x), ("pre", pre), ("post", post)]:
    print(name, "mean", value.mean().item(),
          "std", value.std(unbiased=False).item())
assert pre.shape == post.shape == x.shape

import torch
from torch import nn

class PositionwiseFFN(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.0):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
        )

    def forward(self, x):
        return self.net(x)

x = torch.randn(2, 6, 16)
ffn = PositionwiseFFN(16, 64)
y = ffn(x)
print(x.shape, y.shape)
assert y.shape == x.shape

import torch

def masked_mean(hidden, attention_mask):
    # hidden: [B,T,D], mask: [B,T]
    weights = attention_mask.to(hidden.dtype).unsqueeze(-1)
    total = (hidden * weights).sum(dim=1)
    count = weights.sum(dim=1).clamp_min(1)
    return total / count

hidden = torch.arange(2 * 4 * 3, dtype=torch.float32).reshape(2, 4, 3)
mask = torch.tensor([[1, 1, 0, 0], [1, 1, 1, 0]])
print(masked_mean(hidden, mask).shape)

def check_encoder_contract(hidden, mask, num_classes):
    if hidden.ndim != 3:
        raise ValueError("Encoder 输出必须是 [B,T,D]")
    batch, tokens, dim = hidden.shape
    if mask.shape != (batch, tokens):
        raise ValueError("attention_mask 必须是 [B,T]")
    if not hidden.isfinite().all():
        raise ValueError("Encoder 输出含 NaN/Inf")
    logits = hidden.new_zeros(batch, tokens, num_classes)
    return logits.shape

# 真实调用时把 hidden 换成 Encoder 的 Tensor。
print("contract: [B,T,D] -> [B,T,C]")

import torch

def make_causal(length, device=None):
    rows = torch.arange(length, device=device)[:, None]
    cols = torch.arange(length, device=device)[None, :]
    return cols <= rows

mask = make_causal(5)
print(mask.to(torch.int8))
assert mask[3, 0] and mask[3, 3]
assert not mask[3, 4]

import torch

def choose_token(logits, temperature=1.0, top_k=None):
    logits = logits.float()
    if temperature <= 0:
        raise ValueError("temperature 必须为正")
    scores = logits / temperature
    if top_k is not None:
        values, indices = torch.topk(scores, min(top_k, scores.numel()))
        filtered = torch.full_like(scores, torch.finfo(scores.dtype).min)
        filtered.scatter_(0, indices, values)
        scores = filtered
    probs = torch.softmax(scores, dim=-1)
    return torch.multinomial(probs, num_samples=1).item()

torch.manual_seed(7)
logits = torch.tensor([2.0, 1.5, 0.2, -1.0])
print(choose_token(logits, temperature=0.8, top_k=3))

import torch

def make_teacher_forcing(batch_ids, pad_id, ignore_index=-100):
    decoder_input = batch_ids[:, :-1]
    labels = batch_ids[:, 1:].clone()
    labels[labels == pad_id] = ignore_index
    loss_mask = labels != ignore_index
    return decoder_input, labels, loss_mask

ids = torch.tensor([[2, 8, 9, 3, 0],
                    [2, 7, 3, 0, 0]])
decoder_input, labels, mask = make_teacher_forcing(ids, pad_id=0)
print(decoder_input)
print(labels, mask)
assert decoder_input.shape == labels.shape == mask.shape

import torch

def greedy_decode(model, input_ids, eos_id, max_new_tokens):
    model.eval()
    generated = input_ids
    for _ in range(max_new_tokens):
        logits = model(generated)
        next_id = logits[:, -1, :].argmax(dim=-1, keepdim=True)
        generated = torch.cat([generated, next_id], dim=1)
        if bool((next_id == eos_id).all()):
            break
    return generated

def make_lm_batch(ids, pad_id=0):
    input_ids = ids[:, :-1]
    labels = ids[:, 1:].clone()
    labels[labels == pad_id] = -100
    loss_mask = labels.ne(-100)
    if not bool(loss_mask.any()):
        raise ValueError("batch 没有有效目标")
    return input_ids, labels, loss_mask

# 真实 Tensor 示例：最后一列 PAD 不应产生损失。
print("labels 的位置必须与 logits[:, :-1] 对齐")

import torch

def append_kv(past_key, past_value, new_key, new_value):
    if past_key is None:
        return new_key, new_value
    if past_key.shape[:2] != new_key.shape[:2]:
        raise ValueError("batch/head 维度不一致")
    key = torch.cat([past_key, new_key], dim=-2)
    value = torch.cat([past_value, new_value], dim=-2)
    return key, value

# key/value 约定为 [B,H,T,d]，沿 T 轴增长。

import torch
from torch import nn

batch, length, vocab = 2, 4, 10
logits = torch.randn(batch, length, vocab)
labels = torch.tensor([[2, 4, 5, -100],
                       [3, 1, 7, 8]])
loss_fn = nn.CrossEntropyLoss(ignore_index=-100, reduction="sum")
loss_sum = loss_fn(logits.reshape(-1, vocab), labels.reshape(-1))
valid = (labels != -100).sum()
loss = loss_sum / valid.clamp_min(1)
perplexity = loss.detach().exp()

print("valid tokens:", valid.item())
print("loss:", loss.item(), "ppl:", perplexity.item())

import torch
from torch import nn

model = nn.Linear(8, 3)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
warmup = torch.optim.lr_scheduler.LinearLR(
    optimizer, start_factor=0.1, end_factor=1.0, total_iters=10
)
decay = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=90)
scheduler = torch.optim.lr_scheduler.SequentialLR(
    optimizer, schedulers=[warmup, decay], milestones=[10]
)

for step in range(15):
    optimizer.zero_grad(set_to_none=True)
    loss = model(torch.randn(4, 8)).square().mean()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()
    scheduler.step()
    print(step, scheduler.get_last_lr()[0], loss.item())

def batch_stats(attention_mask, elapsed_seconds):
    valid_tokens = sum(sum(row) for row in attention_mask)
    total_tokens = sum(len(row) for row in attention_mask)
    padding_ratio = 1 - valid_tokens / max(total_tokens, 1)
    return {
        "valid_tokens": valid_tokens,
        "total_tokens": total_tokens,
        "padding_ratio": padding_ratio,
        "tokens_per_second": valid_tokens / max(elapsed_seconds, 1e-9),
        "max_length": max((len(row) for row in attention_mask), default=0),
    }

mask = [[1, 1, 1, 0], [1, 1, 0, 0]]
print(batch_stats(mask, elapsed_seconds=0.02))

import json
from pathlib import Path

class JsonlLogger:
    def __init__(self, path):
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)

    def write(self, record):
        required = {"step", "loss", "valid_tokens", "learning_rate"}
        missing = required - record.keys()
        if missing:
            raise ValueError(f"日志缺少字段: {sorted(missing)}")
        with self.path.open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(record, ensure_ascii=False) + "\n")

logger = JsonlLogger("artifacts/metrics.jsonl")
logger.write({"step": 1, "loss": 3.2, "valid_tokens": 120,
              "learning_rate": 0.001, "grad_norm": 1.8})

def train_step(model, batch, optimizer, loss_fn, update=True):
    input_ids = batch["input_ids"]
    labels = batch["labels"]
    logits = model(input_ids)
    if logits.shape[:2] != labels.shape:
        raise ValueError("logits 与 labels 的 Token 轴不一致")
    loss = loss_fn(logits.reshape(-1, logits.shape[-1]), labels.reshape(-1))
    if not loss.isfinite():
        raise FloatingPointError("loss 为 NaN/Inf")
    loss.backward()
    if update:
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)
    valid_tokens = int(labels.ne(-100).sum())
    return {"loss": float(loss.detach()), "valid_tokens": valid_tokens}

import torch
from torch import nn

class TinyLanguageModel(nn.Module):
    def __init__(self, vocab_size, dim=32, heads=4, layers=2, max_length=64):
        super().__init__()
        self.token = nn.Embedding(vocab_size, dim)
        self.position = nn.Embedding(max_length, dim)
        block = nn.TransformerEncoderLayer(
            d_model=dim, nhead=heads, dim_feedforward=4 * dim,
            batch_first=True, norm_first=True,
        )
        self.blocks = nn.TransformerEncoder(block, num_layers=layers)
        self.norm = nn.LayerNorm(dim)
        self.lm_head = nn.Linear(dim, vocab_size, bias=False)
        self.lm_head.weight = self.token.weight

    def forward(self, input_ids, padding_mask=None):
        batch, length = input_ids.shape
        positions = torch.arange(length, device=input_ids.device)
        x = self.token(input_ids) + self.position(positions)[None, :, :]
        causal = torch.triu(torch.ones(length, length, device=x.device), 1).bool()
        x = self.blocks(x, mask=causal, src_key_padding_mask=padding_mask)
        return self.lm_head(self.norm(x))

model = TinyLanguageModel(vocab_size=50)
logits = model(torch.randint(0, 50, (2, 8)))
print(logits.shape)

import torch

def check_batch(input_ids, mask, vocab_size, name="batch"):
    if input_ids.dtype != torch.long:
        raise TypeError(f"{name}: ids 必须是 long")
    if input_ids.ndim != 2 or mask.shape != input_ids.shape:
        raise ValueError(f"{name}: ids/mask shape 不匹配")
    if input_ids.numel() and (input_ids.min() < 0 or input_ids.max() >= vocab_size):
        raise ValueError(f"{name}: ID 越界")
    valid = int(mask.sum())
    if valid == 0:
        raise ValueError(f"{name}: 没有有效 Token")
    return {"shape": tuple(input_ids.shape), "valid_tokens": valid}

# 实际使用时把 dtype 检查改为 input_ids.dtype == torch.long
print(check_batch(torch.tensor([[1, 2, 0]]),
                  torch.tensor([[1, 1, 0]], dtype=torch.int64), 10))

import torch
from torch import nn

torch.manual_seed(7)
model = TinyLanguageModel(vocab_size=20, dim=32, heads=4, layers=2, max_length=12)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-3)
loss_fn = nn.CrossEntropyLoss()
ids = torch.tensor([[2, 5, 6, 7, 3, 0],
                    [2, 5, 6, 8, 3, 0]])
inputs, labels = ids[:, :-1], ids[:, 1:]
for step in range(200):
    optimizer.zero_grad(set_to_none=True)
    logits = model(inputs)
    loss = loss_fn(logits.reshape(-1, 20), labels.reshape(-1))
    loss.backward()
    optimizer.step()
    if step % 40 == 0:
        print(step, loss.item())
print("final:", loss.item())

import torch

def assert_model_contract(model, vocab_size, batch=2, tokens=5):
    ids = torch.randint(0, vocab_size, (batch, tokens), dtype=torch.long)
    with torch.inference_mode():
        logits = model(ids)
    assert logits.shape == (batch, tokens, vocab_size), logits.shape
    assert torch.isfinite(logits).all()
    return {"input": tuple(ids.shape), "output": tuple(logits.shape)}

def assert_weight_tying(model):
    tied = model.lm_head.weight.data_ptr() == model.token.weight.data_ptr()
    print("weight_tying:", tied)
    return tied

# 使用 TinyLanguageModel 实例调用：
# print(assert_model_contract(model, vocab_size=50))
# assert_weight_tying(model)

def throughput(valid_tokens, elapsed_seconds):
    if elapsed_seconds <= 0:
        raise ValueError("耗时必须为正")
    return valid_tokens / elapsed_seconds

runs = [
    {"name": "fp32", "tokens": 4096, "seconds": 0.42, "peak_mb": 812},
    {"name": "amp", "tokens": 4096, "seconds": 0.24, "peak_mb": 506},
]
for run in runs:
    print(run["name"], {
        "tokens_per_second": round(throughput(run["tokens"], run["seconds"])),
        "peak_mb": run["peak_mb"],
    })

from pathlib import Path
import hashlib
import json
import torch

def sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()

def load_manifest(root):
    root = Path(root)
    manifest = json.loads((root / "config.json").read_text(encoding="utf-8"))
    required = {"vocab_size", "hidden_size", "num_hidden_layers"}
    missing = required - manifest.keys()
    if missing:
        raise ValueError(f"config 缺少字段: {sorted(missing)}")
    weight_file = root / "model.pt"
    if not weight_file.exists():
        raise FileNotFoundError(weight_file)
    print({"sha256": sha256(weight_file), "config": manifest})
    # 仅加载可信来源；torch.load 可能反序列化对象，不要接收未知文件。
    state = torch.load(weight_file, map_location="cpu", weights_only=True)
    return manifest, state

config, state_dict = load_manifest("checkpoints/demo")
print(len(state_dict), config["hidden_size"])

def build_example(tokenizer, prompt, answer, max_length=512):
    prompt_text = f"用户：{prompt}\n助手："
    full_text = prompt_text + answer + tokenizer.eos_token
    prompt_ids = tokenizer(prompt_text, add_special_tokens=False)["input_ids"]
    full = tokenizer(
        full_text, add_special_tokens=True, truncation=True,
        max_length=max_length, return_attention_mask=True,
    )
    input_ids = full["input_ids"]
    labels = [-100] * len(input_ids)
    answer_start = min(len(prompt_ids), len(input_ids))
    for index in range(answer_start, len(input_ids)):
        labels[index] = input_ids[index]
    # 只有回答区域产生损失；PAD 在 batch padding 后仍设为 -100。
    return {
        "input_ids": input_ids,
        "attention_mask": full["attention_mask"],
        "labels": labels,
    }

sample = build_example(tokenizer, "解释什么是向量", "向量是有序的数值集合。")
print(len(sample["input_ids"]), sample["labels"])

def lora_params(in_features, out_features, rank, targets):
    base = in_features * out_features
    adapter = rank * (in_features + out_features)
    return {"base": base * targets, "lora": adapter * targets}

for rank in (4, 8, 16):
    result = lora_params(4096, 4096, rank, targets=8)
    ratio = result["lora"] / (result["base"] + result["lora"])
    print(rank, result["lora"], f"{ratio:.3%}")

from dataclasses import dataclass, asdict
import json

@dataclass
class CaseResult:
    case_id: str
    category: str
    expected: str
    actual: str
    passed: bool
    note: str = ""

def summarize(results):
    total = len(results)
    passed = sum(item.passed for item in results)
    by_category = {}
    for item in results:
        bucket = by_category.setdefault(item.category, [0, 0])
        bucket[0] += 1
        bucket[1] += int(item.passed)
    return {"total": total, "passed": passed, "rate": passed / total if total else 0,
            "by_category": by_category}

results = [
    CaseResult("safe-01", "privacy", "拒绝泄露", "拒绝并说明原因", True),
    CaseResult("safe-02", "injection", "忽略文档指令", "执行了文档指令", False),
]
report = {"results": [asdict(item) for item in results],
          "summary": summarize(results)}
print(json.dumps(report, ensure_ascii=False, indent=2))

from dataclasses import asdict, dataclass
from pathlib import Path
import json

@dataclass
class Config:
    seed: int = 7
    vocab_size: int = 128
    max_length: int = 64
    dim: int = 64
    heads: int = 4
    layers: int = 2
    learning_rate: float = 2e-3
    train_steps: int = 800

def save_manifest(path, config, files, metrics):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    record = {
        "config": asdict(config),
        "files": files,
        "metrics": metrics,
        "status": "candidate",
    }
    path.write_text(
        json.dumps(record, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    return record

cfg = Config()
manifest = save_manifest(
    "artifacts/manifest.json",
    cfg,
    files={"data": "data/clean.jsonl", "tokenizer": "artifacts/tokenizer.json"},
    metrics={"shape_probe": True, "overfit": False},
)
print(manifest["config"]["heads"])

def causal_example(ids, pad_id=0):
    if len(ids) < 2:
        raise ValueError("序列至少需要两个 Token")
    inputs = ids[:-1]
    labels = ids[1:]
    attention = [int(token != pad_id) for token in inputs]
    loss_mask = [int(token != pad_id) for token in labels]
    return {
        "input_ids": inputs,
        "labels": labels,
        "attention_mask": attention,
        "loss_mask": loss_mask,
    }

examples = [causal_example([2, 5, 6, 3, 0])]
print(examples[0])

from copy import deepcopy

base = {
    "seed": 7, "dim": 64, "heads": 4, "layers": 2,
    "position": "learned", "causal": True,
}
variants = [
    ("baseline", base),
    ("one_head", {**base, "heads": 1}),
    ("sinusoidal", {**base, "position": "sinusoidal"}),
    ("no_position", {**base, "position": "none"}),
]

def command(name, config):
    return {
        "name": name,
        "config": deepcopy(config),
        "metrics_file": f"artifacts/{name}/metrics.jsonl",
    }

for name, config in variants:
    print(command(name, config))

def relative_change(baseline, value):
    if baseline == 0:
        return None
    return (value - baseline) / abs(baseline)

baseline = {"val_loss": 2.31, "ppl": 10.08, "tokens_per_second": 520}
runs = {
    "one_head": {"val_loss": 2.45, "ppl": 11.58, "tokens_per_second": 610},
    "sinusoidal": {"val_loss": 2.28, "ppl": 9.78, "tokens_per_second": 505},
}
for name, metrics in runs.items():
    print(name, {
        key: relative_change(baseline[key], value)
        for key, value in metrics.items()
    })

from pathlib import Path
import json

REQUIRED = [
    "config.json", "token_probe.json", "shape_probe.txt",
    "metrics.jsonl", "report.md",
]

def verify(root):
    root = Path(root)
    missing = [name for name in REQUIRED if not (root / name).exists()]
    if missing:
        raise FileNotFoundError(f"缺少交付物: {missing}")
    manifest = json.loads((root / "config.json").read_text(encoding="utf-8"))
    for key in ("seed", "vocab_size", "max_length"):
        if key not in manifest:
            raise ValueError(f"配置缺少 {key}")
    return {"status": "ready", "files": len(REQUIRED)}

print(verify("artifacts/final"))

checklist = {
    "data_source_and_license": True,
    "tokenizer_round_trip": True,
    "shape_and_mask_probe": True,
    "single_batch_overfit": True,
    "validation_report": True,
    "ablation_report": True,
    "safety_and_limitations": True,
    "reproducible_readme": True,
}
failed = [name for name, passed in checklist.items() if not passed]
print("READY" if not failed else {"BLOCKED": failed})