Code Walkthrough
このページは、既存の解説ページとハンズオンで実際にどのコードが動いているのかを読むための補助線です。 特に Hands-on は、コマンドを実行するだけだと「裏で何が起きたか」が見えにくいため、主要な処理を GitHub permanent link ベースのコード抜粋として埋め込んでいます。
解説ページとコードの対応
Foundations は token、logits、causal mask、attention などの概念を説明します。 Base Hands-on は tokenizer、next-token 学習、生成パラメータの差を動かします。 Reasoning と Reasoning Hands-on は hidden scratchpad、reasoning effort、verification、self-consistency を小さな制御ループとして観察します。
Runnable now の範囲は、Transformer 本体を小さく実装し、データを token ID にして、1つ先の token を予測する学習を回し、 保存した checkpoint から続きを生成するところまでです。SFT、preference / alignment、RAG、tool use、agent loop は Implementation Plan 側の Planned 項目です。
1. Tokenizer
Tokenizer のハンズオンで実際に動いているコードは、文字単位の tokenizer と、識別子・空白・数字・記号を分ける 正規表現 tokenizer の比較です。`devbox run compare` はこの差を使って、自然文とコードで token の切れ方が変わることを見せます。
@dataclassclass CharTokenizer(_VocabularyMixin): """One token per Unicode character."""
kind: ClassVar[str] = "char"
@classmethod def train(cls, text: str) -> "CharTokenizer": return cls(vocab=list(dict.fromkeys(text)))
def encode(self, text: str) -> list[int]: return self._encode_tokens(list(text))
def decode(self, token_ids: list[int]) -> str: return "".join(self.vocab[token_id] for token_id in token_ids)
@dataclassclass RegexTokenizer(_VocabularyMixin): """Small regex tokenizer that keeps code identifiers and whitespace chunks intact."""
kind: ClassVar[str] = "regex" pattern: ClassVar[re.Pattern[str]] = re.compile( r"\s+|[A-Za-z_][A-Za-z_0-9]*|\d+|[^\sA-Za-z_0-9]", flags=re.UNICODE, )
@classmethod def train(cls, text: str) -> "RegexTokenizer": return cls(vocab=list(dict.fromkeys(cls._tokenize(text))))
@classmethod def _tokenize(cls, text: str) -> list[str]: return cls.pattern.findall(text)
def encode(self, text: str) -> list[int]: return self._encode_tokens(self._tokenize(text))
def decode(self, token_ids: list[int]) -> str: return "".join(self.vocab[token_id] for token_id in token_ids)CharTokenizer は1文字を1 token にし、RegexTokenizer はコードの識別子や空白をまとまりとして残します。モデルに入る前に、文字列は必ず整数 ID の列へ変換されます。
2. Next-token batch
Foundations で説明している「次の token を予測する」学習は、教師データの作り方としてはかなり単純です。 `inputs` はある長さの token window、`targets` は同じ window を1つ右にずらしたものです。 つまり next-token 学習の正解を1つ右にずらす処理が、この lab の学習問題そのものです。
def make_lm_batch( token_ids: list[int], *, context_length: int, batch_size: int, device: str | torch.device = "cpu", generator: torch.Generator | None = None,) -> LanguageModelBatch: if context_length < 2: raise ValueError("context_length must be at least 2") if batch_size < 1: raise ValueError("batch_size must be at least 1") if len(token_ids) < 2: raise ValueError("token_ids must contain at least 2 tokens")
needed = context_length + 1 if len(token_ids) < needed: repeat_count = (needed // len(token_ids)) + 1 token_ids = token_ids * repeat_count
data = torch.tensor(token_ids, dtype=torch.long, device=device) max_start = len(data) - needed starts = torch.randint(0, max_start + 1, (batch_size,), generator=generator, device=device) windows = torch.stack([data[start : start + needed] for start in starts]) return LanguageModelBatch(inputs=windows[:, :-1], targets=windows[:, 1:])短すぎるデータは繰り返して context_length + 1 の window を作り、最後に inputs と targets に分けています。
3. Attention と causal mask
Attention の章で出てくる query、key、value、head、causal mask は、ここで実際に tensor 操作として実装されています。 causal mask は未来位置の score を `-inf` にして、softmax 後に未来 token を参照できないようにします。
class CausalSelfAttention(nn.Module): def __init__(self, config: TransformerConfig) -> None: super().__init__() self.n_heads = config.n_heads self.head_dim = config.d_model // config.n_heads self.qkv = nn.Linear(config.d_model, 3 * config.d_model) self.output = nn.Linear(config.d_model, config.d_model) self.attn_dropout = nn.Dropout(config.dropout) self.resid_dropout = nn.Dropout(config.dropout) mask = torch.tril(torch.ones(config.context_length, config.context_length)) self.register_buffer("causal_mask", mask.view(1, 1, config.context_length, config.context_length))
def forward(self, x: torch.Tensor) -> torch.Tensor: batch_size, sequence_length, d_model = x.shape qkv = self.qkv(x) query, key, value = qkv.split(d_model, dim=2)
query = query.view(batch_size, sequence_length, self.n_heads, self.head_dim).transpose(1, 2) key = key.view(batch_size, sequence_length, self.n_heads, self.head_dim).transpose(1, 2) value = value.view(batch_size, sequence_length, self.n_heads, self.head_dim).transpose(1, 2)
scores = query @ key.transpose(-2, -1) scores = scores / math.sqrt(self.head_dim) mask = self.causal_mask[:, :, :sequence_length, :sequence_length] scores = scores.masked_fill(mask == 0, float("-inf")) weights = F.softmax(scores, dim=-1) weights = self.attn_dropout(weights)
attended = weights @ value attended = attended.transpose(1, 2).contiguous().view(batch_size, sequence_length, d_model) return self.resid_dropout(self.output(attended))`qkv` を3分割し、head ごとに score を計算し、mask をかけてから value を混ぜます。ここが Transformer の中心です。
4. Transformer 本体と生成
TransformerLanguageModel は token embedding と position embedding を足し、複数の block を通して logits を出します。 `targets` が渡されたときは cross entropy loss を返し、生成時は最後の位置の logits から次 token を選んで列に追加します。
def forward( self, inputs: torch.Tensor, targets: torch.Tensor | None = None ) -> tuple[torch.Tensor, torch.Tensor | None]: _, sequence_length = inputs.shape if sequence_length > self.config.context_length: raise ValueError("input sequence is longer than context_length")
positions = torch.arange(sequence_length, device=inputs.device) x = self.token_embedding(inputs) + self.position_embedding(positions) for block in self.blocks: x = block(x) x = self.final_norm(x) logits = self.output(x)
loss = None if targets is not None: loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.reshape(-1)) return logits, loss
@torch.no_grad() def generate( self, token_ids: list[int], *, max_new_tokens: int = 80, temperature: float = 0.8, top_k: int | None = None, ) -> list[int]: device = next(self.parameters()).device generated = torch.tensor([token_ids], dtype=torch.long, device=device)
for _ in range(max_new_tokens): context = generated[:, -self.config.context_length :] logits, _ = self(context) next_logits = logits[:, -1, :] if next_logits.size(-1) > 1: next_logits[:, 0] = float("-inf")
if temperature <= 0: next_token = torch.argmax(next_logits, dim=-1, keepdim=True) else: next_logits = next_logits / temperature if top_k is not None: values, _ = torch.topk(next_logits, min(top_k, next_logits.size(-1))) next_logits[next_logits < values[:, [-1]]] = float("-inf") probs = F.softmax(next_logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) generated = torch.cat([generated, next_token], dim=1)
return generated[0].tolist()temperature が0以下なら argmax、0より大きければ softmax から sampling します。top-k は候補を上位 k 個に絞るための設定です。
5. Training loop
Hands-on の `train:natural` と `train:code` は、dataset を読み、tokenizer を作り、model を初期化し、 `make_lm_batch` で作った batch に対して loss を計算して optimizer を進めます。 最後に checkpoint、tokenizer、metrics、sample を保存するので、生成コマンドが同じ run directory を再利用できます。
text = read_text(args.dataset) tokenizer = build_tokenizer(args.tokenizer, text) token_ids = tokenizer.encode(text) train_ids, valid_ids = train_valid_split(token_ids) if len(valid_ids) < 2: valid_ids = train_ids
config = TransformerConfig( vocab_size=tokenizer.vocab_size, context_length=args.context_length, d_model=args.d_model, n_heads=args.n_heads, n_layers=args.n_layers, d_ff=args.d_ff, dropout=args.dropout, ) model = TransformerLanguageModel(config).to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr) metrics: list[dict[str, float | int]] = []
print(json.dumps({ "dataset": args.dataset, "tokenizer": args.tokenizer, "tokens": len(token_ids), "vocab_size": tokenizer.vocab_size, "parameters": count_parameters(model), "device": device, }, ensure_ascii=False, indent=2))
for step in range(1, args.steps + 1): batch = make_lm_batch( train_ids, context_length=args.context_length, batch_size=args.batch_size, device=device, ) _, loss = model(batch.inputs, batch.targets) assert loss is not None
optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step()
if step == 1 or step % args.eval_interval == 0 or step == args.steps: valid_loss = estimate_loss( model, valid_ids, context_length=args.context_length, batch_size=args.batch_size, device=device, ) row = {"step": step, "train_loss": float(loss.item()), "valid_loss": valid_loss} metrics.append(row) print(json.dumps(row, ensure_ascii=False))
output_dir = Path(args.output_dir) save_experiment(output_dir, model=model, tokenizer=tokenizer, metrics=metrics)
prompt = args.sample_prompt or text[: min(20, len(text))] prompt_ids = tokenizer.encode(prompt) generated_ids = model.generate(prompt_ids, max_new_tokens=args.sample_tokens, temperature=0.8) sample = tokenizer.decode(generated_ids) (output_dir / "sample.txt").write_text(sample, encoding="utf-8") print(f"saved: {output_dir}") print("--- sample ---") print(sample)学習の最小単位は forward、loss.backward、optimizer.step です。評価 loss は一定間隔で別 batch から測ります。
6. Reasoning loop
Reasoning Hands-on は、Transformer 本体を賢くした実装ではなく、推論らしい制御の形を外側の loop として観察する教材です。 `run_reasoning` は prompt から数式を取り出し、candidate を作り、投票し、verification を行い、hidden scratchpad 相当の内部手順と visible summary を分けて返します。
def run_reasoning(prompt: str, effort: ReasoningEffort = "medium", samples: int | None = None) -> ReasoningRun: expression = extract_arithmetic_expression(prompt) sample_count = samples if samples is not None else _default_samples(effort) candidates = [_solve_once(expression, index) for index in range(sample_count)] votes = Counter(str(candidate) for candidate in candidates) final_answer, _ = votes.most_common(1)[0] hidden_steps = _build_hidden_steps(expression, effort, candidates, final_answer) verification = "passed" if safe_evaluate(expression) == int(final_answer) else "failed"
return ReasoningRun( prompt=prompt, expression=expression, effort=effort, final_answer=final_answer, hidden_steps=tuple(hidden_steps), visible_summary=( f"Used {effort} effort with {sample_count} candidate path(s); " f"verification {verification}." ), reasoning_token_count=sum(_rough_token_count(step) for step in hidden_steps), verification=verification, self_consistency_votes=dict(votes), )ここでは候補生成を決定的にしてあります。狙いは性能ではなく、reasoning effort と self-consistency の制御面を観察することです。