17 題 從 NLP 基礎到打造 GPT 的 Python 筆記

題目整理

建立最基本的 NLP 前處理管線：把原始文字正規化（轉小寫、去標點）、斷詞成 token，從語料建立詞彙表，再把任意句子轉成 bag-of-words 計數向量（每個維度是某個字出現幾次）。這是文字進到任何模型前的第一步。

解法說明

preprocess 用 lower() 統一大小寫，再用正規表示式 [a-z0-9']+ 抓出「連續字母數字」當作 word（順手把標點與空白當分隔）。BagOfWords 從語料收集排序後的 vocab 建 stoi，vectorize 掃過 token 累加對應維度；沒看過的字直接忽略。BoW 的特點是只看字頻、不看順序。

import re
from typing import Dict, List


def preprocess(text: str) -> List[str]:
    text = text.lower()
    return re.findall(r"[a-z0-9']+", text)   # Words = maximal alnum runs.


class BagOfWords:
    def __init__(self, corpus: List[str]):
        vocab = sorted({tok for doc in corpus for tok in preprocess(doc)})
        self.stoi: Dict[str, int] = {w: i for i, w in enumerate(vocab)}

    def vectorize(self, text: str) -> List[int]:
        vec = [0] * len(self.stoi)
        for tok in preprocess(text):
            if tok in self.stoi:             # Ignore out-of-vocab words.
                vec[self.stoi[tok]] += 1
        return vec

Interview Explanation Flow

corpus = ["The cat sat.", "The dog ran!"]
preprocess("The cat sat.") -> ["the", "cat", "sat"]
vocab = ["cat","dog","ran","sat","the"]  (sorted)
stoi  = {cat:0, dog:1, ran:2, sat:3, the:4}

vectorize("the cat the") -> [1, 0, 0, 0, 2]
#                            cat        the appears twice

延伸練習：Unicode-aware preprocessing（follow-up）

上面的 preprocess 只吃 ASCII：[a-z0-9'] 遇到 café、Zürich、北京、Москва 會把重音字、非拉丁字母切爛或丟掉。以下是不改動原題解、額外提供的 Unicode 版本；Python 3 的 \w 對 str pattern 預設就是 Unicode。

import re
import unicodedata


def preprocess_unicode(text: str) -> list:
    text = unicodedata.normalize("NFC", text)   # 1. Compose accents to canonical form.
    text = text.casefold()                       # 2. Unicode-aware lowercasing.
    # 3. Word = run of Unicode letters/digits (not underscore), keeping contractions.
    return re.findall(r"[^\W_]+(?:['’][^\W_]+)*", text)

三個步驟各修一個 Unicode 陷阱：

• normalize("NFC")：é 可能是單一 codepoint U+00E9，也可能是 e + 結合重音 U+0301；兩者外觀相同但是不同字串，會讓 vocab 對不上。NFC 先合成成同一種標準形式。
• casefold() 取代 lower()：casefold 是給「無視大小寫比對」用的正確版本。例如德文 "Straße".casefold() → 'strasse'、希臘文字尾 sigma 也處理正確，lower() 都會漏。
• [^\W_]+(?:['’][^\W_]+)*：[^\W_] 是「Unicode 詞字元但排除底線」（純 \w 會把 _ 也算進去）；後面允許詞內單引號，且同時吃 ASCII ' 與真實文本常見的 Unicode 右單引號 ’（U+2019），讓 it's / it’s 都保持完整。

行為對比：

s = "Café Zürich — it's naïve. 北京 Москва don’t"

# 原題 ASCII 版
re.findall(r"[a-z0-9']+", s.lower())
-> ['caf', 'z', 'rich', "it's", 'na', 've', 'don']              # 重音與非拉丁字母被切爛

preprocess_unicode(s)
-> ['café', 'zürich', "it's", 'naïve', '北京', 'москва', 'don’t']  # 保持完整

題目整理

把 bag-of-words 那種稀疏、彼此正交的表示，換成稠密向量：每個 token id 對應一列可學習的 embedding。實作 embedding 查表，並用 cosine similarity 找出與某個字最相近的 k 個字（word2vec / GloVe 的核心操作）。

解法說明

nn.Embedding 本質是一張 (vocab_size, dim) 的查表：給 id 就回傳對應的列向量。語意相近的字在訓練後會落在方向相近的位置，所以用 cosine similarity（只看方向、不看長度）衡量相似度。nearest 把 query 向量對整張表算相似度，排除自己後取 top-k。

import torch
import torch.nn as nn
from typing import List


class WordEmbeddings(nn.Module):
    def __init__(self, vocab_size: int, dim: int):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, dim)   # Lookup table.

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        return self.embedding(ids)                        # (..., dim)


def cosine_similarity(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    a = a / a.norm(dim=-1, keepdim=True)                  # Direction only.
    b = b / b.norm(dim=-1, keepdim=True)
    return a @ b.transpose(-2, -1)


def nearest(query_id: int, weight: torch.Tensor, k: int = 5) -> List[int]:
    q = weight[query_id].unsqueeze(0)                     # (1, dim)
    sims = cosine_similarity(q, weight).squeeze(0)        # (vocab,)
    sims[query_id] = float("-inf")                        # Exclude the word itself.
    return torch.topk(sims, k).indices.tolist()

Interview Explanation Flow

weight: (vocab=10000, dim=64)
q = weight[stoi["king"]]              # (64,)
sims = cosine(q, weight)             # (10000,)
sims[stoi["king"]] = -inf           # don't return itself
topk(sims, 5) -> ["queen","prince","monarch","throne","royal"]

題目整理

做一個文字情感分類器（正面／負面）。輸入是一批 token id (B, T)（不等長的句子用 <pad> 補齊），模型要輸出每個類別的 logits。做法：查 embedding → 對每句忽略 padding 做 mean-pooling得到句向量 → 線性層分類 → 用 cross-entropy 訓練。

解法說明

關鍵是 masked mean-pooling：直接對整個 (B, T, d) 取平均會把 pad token 算進去而稀釋語意。用 ids != pad_id 造 mask，只加總真實 token 再除以真實長度（clamp(min=1) 防止空句除以零）。padding_idx 讓 pad 的 embedding 固定為 0 且不更新。

import torch
import torch.nn as nn
from torch.nn import functional as F


class SentimentClassifier(nn.Module):
    def __init__(self, vocab_size: int, dim: int, num_classes: int = 2, pad_id: int = 0):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, dim, padding_idx=pad_id)
        self.fc = nn.Linear(dim, num_classes)
        self.pad_id = pad_id

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        emb = self.embedding(ids)                    # (B, T, dim)
        mask = (ids != self.pad_id).unsqueeze(-1)    # (B, T, 1)
        summed = (emb * mask).sum(dim=1)             # (B, dim), pads zeroed out
        counts = mask.sum(dim=1).clamp(min=1)        # (B, 1), avoid div-by-zero
        pooled = summed / counts                     # Masked mean over real tokens.
        return self.fc(pooled)                       # (B, num_classes)


def train_step(model, ids, labels, optimizer) -> float:
    logits = model(ids)
    loss = F.cross_entropy(logits, labels)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    return loss.item()


@torch.no_grad()
def predict(model, ids) -> torch.Tensor:
    return model(ids).argmax(dim=-1)                 # Class with highest logit.

Interview Explanation Flow

ids = [[The, movie, was, great],
       [Bad,  <pad>, <pad>, <pad>]]     # pad_id = 0
mask = [[1,1,1,1],
        [1,0,0,0]]
pooled row 1 = mean of 4 embeddings
pooled row 2 = embedding("Bad") / 1     # not / 4

題目整理

Self-attention 本身對位置無感（打亂 token 順序，輸出只是跟著換位置），所以要額外把「位置」資訊加進去。實作 Attention Is All You Need 的 sinusoidal positional encoding：產生一個 (seq_len, d_model) 矩陣，偶數維用 sin、奇數維用 cos，且不同維度用不同頻率，之後直接加到 token embedding 上。

解法說明

位置 pos 在維度 i 的角度是 pos / 10000^(2i/d)。用 div_term = exp(arange(0,d,2) · (-log(10000)/d)) 一次算好所有頻率（等價但數值更穩）。偶數維 0::2 填 sin、奇數維 1::2 填 cos。這種編碼固定不用學、對任意長度都定義得出來，而且相鄰位置的編碼有平滑關係。

import math
import torch


def sinusoidal_positional_encoding(seq_len: int, d_model: int) -> torch.Tensor:
    pe = torch.zeros(seq_len, d_model)
    position = torch.arange(seq_len).unsqueeze(1).float()          # (T, 1)
    # div_term[i] = 1 / 10000^(2i / d_model), computed in log-space for stability.
    div_term = torch.exp(
        torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
    )
    pe[:, 0::2] = torch.sin(position * div_term)                   # Even dims -> sin.
    pe[:, 1::2] = torch.cos(position * div_term)                   # Odd dims  -> cos.
    return pe                                                      # (seq_len, d_model)


# Usage: x = token_embedding(idx) + sinusoidal_positional_encoding(T, d_model)

Interview Explanation Flow

d_model = 4, positions 0..3
pos=0: [sin0, cos0, sin0,  cos0 ] = [0, 1, 0, 1]
pos=1: [sin1, cos1, sin(1/100), cos(1/100)]
       fast-freq dims          slow-freq dims
even index -> sin, odd index -> cos

題目整理

實作單頭的 scaled dot-product self-attention。輸入 x（shape (B, T, n_embd)），從 x 投影出 Query、Key、Value，用 softmax(Q Kᵀ / √d) · V 讓每個位置根據相關性聚合其他位置的資訊。支援可選的 mask（例如 causal 或 padding mask）。這是整個 Transformer 的最小積木。

解法說明

三個線性層各自把 x 投影成 Q、K、V。注意力分數 Q Kᵀ 是每個 query 對每個 key 的內積，除以 √head_size 避免維度變大時 softmax 過度飽和。若有 mask，把不該看的位置設 -inf，softmax 後歸零，再對 V 加權求和。

import torch
import torch.nn as nn
from torch.nn import functional as F


class SelfAttention(nn.Module):
    def __init__(self, n_embd: int, head_size: int):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)

    def forward(self, x, mask=None):
        k = self.key(x)      # (B, T, head_size)
        q = self.query(x)
        v = self.value(x)
        wei = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5   # (B, T, T)
        if mask is not None:
            wei = wei.masked_fill(mask == 0, float("-inf"))   # Block disallowed pairs.
        wei = F.softmax(wei, dim=-1)                          # Row-normalized weights.
        return wei @ v                                        # (B, T, head_size)

Interview Explanation Flow

x: (B, T, n_embd) = (1, 3, 8), head_size = 4
q,k,v: (1, 3, 4)
wei = q @ kᵀ / sqrt(4)  -> (1, 3, 3)   # token-to-token scores
softmax rows -> weights sum to 1
out = wei @ v -> (1, 3, 4)

題目整理

把單頭 attention 擴成 multi-head：把 n_embd 切成 n_head 份，每個 head 在自己的子空間獨立做 attention，讓模型能同時關注不同面向（語法、指代、距離…），最後把各 head 輸出 concat 再投影回 n_embd。要求用張量 reshape做平行運算，而不是 Python for-loop 逐 head 跑。

解法說明

用單一個 qkv 線性層一次算出 Q、K、V（再 split），比多個 Linear 更省。接著 view 成 (B, T, n_head, head_size) 再 transpose 到 (B, n_head, T, head_size)，讓 head 變成 batch 維度一起算 attention。算完 transpose 回來、contiguous().view 把各 head concat 成 (B, T, n_embd)，再過輸出投影 proj。

import torch
import torch.nn as nn
from torch.nn import functional as F


class MultiHeadSelfAttention(nn.Module):
    def __init__(self, n_embd: int, n_head: int):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.head_size = n_embd // n_head
        self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)   # Fused Q,K,V.
        self.proj = nn.Linear(n_embd, n_embd)                  # Mix heads back.

    def forward(self, x, mask=None):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        # (B, T, C) -> (B, n_head, T, head_size): heads become a batch dim.
        q = q.view(B, T, self.n_head, self.head_size).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_size).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_size).transpose(1, 2)

        att = q @ k.transpose(-2, -1) * self.head_size ** -0.5   # (B, nh, T, T)
        if mask is not None:
            att = att.masked_fill(mask == 0, float("-inf"))
        att = F.softmax(att, dim=-1)
        out = att @ v                                            # (B, nh, T, hs)
        out = out.transpose(1, 2).contiguous().view(B, T, C)     # Concat heads.
        return self.proj(out)

Interview Explanation Flow

n_embd = 64, n_head = 8 -> head_size = 8
qkv(x): (B, T, 192) -> split -> q,k,v each (B, T, 64)
view:      (B, T, 8, 8)
transpose: (B, 8, T, 8)   # (batch, heads, time, head_size)
att = q @ kᵀ : (B, 8, T, T)   # 8 heads computed together
out -> transpose+view -> (B, T, 64) -> proj

題目整理

把 multi-head attention 與 feed-forward 組成一個可堆疊的 Transformer block。用 pre-norm 殘差結構：x = x + attn(ln1(x))、x = x + ffn(ln2(x))。attention 負責跨位置「溝通」，feed-forward 負責每個位置獨立「思考」。堆 N 層就是 GPT / BERT 的主體。

解法說明

兩個 sub-layer 各包一層 LayerNorm 與 residual。Pre-norm（先 norm 再進 sub-layer）比原論文的 post-norm 更好訓練、梯度更穩，是現代做法。feed-forward 先擴張到 4 · n_embd、過非線性（GELU）、再投影回來，給模型逐位置的非線性容量。這正是第 13 題 Code GPT 裡 Block 的通用版（多接了一個 mask 參數，可當 encoder 或 decoder 用）。

import torch
import torch.nn as nn


class TransformerBlock(nn.Module):
    def __init__(self, n_embd: int, n_head: int, dropout: float = 0.1):
        super().__init__()
        self.attn = MultiHeadSelfAttention(n_embd, n_head)   # From problem 6.
        self.ff = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),   # Expand.
            nn.GELU(),                        # Per-position non-linearity.
            nn.Linear(4 * n_embd, n_embd),   # Project back.
            nn.Dropout(dropout),
        )
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x, mask=None):
        x = x + self.attn(self.ln1(x), mask)   # Communicate across positions.
        x = x + self.ff(self.ln2(x))           # Think at each position.
        return x

Interview Explanation Flow

x --> ln1 --> attn --+--> ln2 --> ff --+--> out
 \___________________/ \_______________/
      residual add         residual add

題目整理

實作一個 byte-level 的 BPE tokenizer。給一段訓練文字與目標 vocab_size，先把文字編成 UTF-8 bytes（0–255 共 256 個基礎 token），接著反覆找出出現頻率最高的相鄰 token pair，把它合併成一個新的 token id，直到達到目標詞彙量。之後要能用學到的 merges 對任意文字 encode，並能 decode 回原字串。

解法說明

三個核心函式：get_stats 統計相鄰 pair 次數、merge 把某個 pair 換成新 id、訓練時每輪選 max 頻率的 pair 合併並記錄到 merges。encode 時要依照學習順序套用 merge（每次挑 merge index 最小、也就是最早學到的 pair），這樣才能重現訓練時的合併結果。decode 則把每個 id 對應的 bytes 串起來再 UTF-8 解碼。

from collections import Counter
from typing import Dict, List, Tuple


def get_stats(ids: List[int]) -> Counter:
    counts = Counter()
    for pair in zip(ids, ids[1:]):   # Every adjacent pair.
        counts[pair] += 1
    return counts


def merge(ids: List[int], pair: Tuple[int, int], new_id: int) -> List[int]:
    result = []
    i = 0
    while i < len(ids):
        if i < len(ids) - 1 and ids[i] == pair[0] and ids[i + 1] == pair[1]:
            result.append(new_id)
            i += 2
        else:
            result.append(ids[i])
            i += 1
    return result


class BPETokenizer:
    def __init__(self):
        self.merges: Dict[Tuple[int, int], int] = {}
        self.vocab: Dict[int, bytes] = {i: bytes([i]) for i in range(256)}

    def train(self, text: str, vocab_size: int) -> None:
        assert vocab_size >= 256
        num_merges = vocab_size - 256
        ids = list(text.encode("utf-8"))

        for i in range(num_merges):
            stats = get_stats(ids)
            if not stats:
                break
            pair = max(stats, key=stats.get)   # Most frequent adjacent pair.
            new_id = 256 + i
            ids = merge(ids, pair, new_id)
            self.merges[pair] = new_id
            self.vocab[new_id] = self.vocab[pair[0]] + self.vocab[pair[1]]

    def encode(self, text: str) -> List[int]:
        ids = list(text.encode("utf-8"))
        while len(ids) >= 2:
            stats = get_stats(ids)
            # Apply the earliest-learned merge that still appears.
            pair = min(stats, key=lambda p: self.merges.get(p, float("inf")))
            if pair not in self.merges:
                break
            ids = merge(ids, pair, self.merges[pair])
        return ids

    def decode(self, ids: List[int]) -> str:
        raw = b"".join(self.vocab[i] for i in ids)
        return raw.decode("utf-8", errors="replace")

Interview Explanation Flow

text = "aaabdaaabac"   # treat as chars for illustration
ids  = [a, a, a, b, d, a, a, a, b, a, c]

round 1: most frequent pair = (a, a) -> Z
ids  = [Z, a, b, d, Z, a, b, a, c]      merges[(a,a)] = Z

round 2: most frequent pair = (Z, a) -> Y
ids  = [Y, b, d, Y, b, a, c]            merges[(Z,a)] = Y

round 3: most frequent pair = (Y, b) -> X
ids  = [X, d, X, a, c]                  merges[(Y,b)] = X

題目整理

建立字元級（character-level）詞彙表。給一段文字，取出所有不重複字元並排序，建立 stoi（字元 → 整數 id）與 itos（id → 字元）兩個對照表，並提供 encode(str) -> List[int] 與 decode(List[int]) -> str。這是最簡單的 tokenizer，也是 Karpathy「makemore / mini-GPT」教學的起點。

解法說明

sorted(set(text)) 一行取得穩定、可重現的字元順序（排序很重要，否則每次跑出來的 id 都不同，模型無法載入舊 checkpoint）。stoi 與 itos 用 enumerate 一次建好；encode/decode 就是 list comprehension 查表。

from typing import List


class CharVocab:
    def __init__(self, text: str):
        chars = sorted(set(text))   # Deterministic order for reproducibility.
        self.stoi = {ch: i for i, ch in enumerate(chars)}
        self.itos = {i: ch for i, ch in enumerate(chars)}

    @property
    def size(self) -> int:
        return len(self.stoi)

    def encode(self, s: str) -> List[int]:
        return [self.stoi[ch] for ch in s]

    def decode(self, ids: List[int]) -> str:
        return "".join(self.itos[i] for i in ids)

Interview Explanation Flow

text = "hello"
set(text)          -> {'h','e','l','o'}   (unordered)
sorted(set(text))  -> ['e','h','l','o']

stoi = {'e':0, 'h':1, 'l':2, 'o':3}
encode("hell") -> [1, 0, 2, 2]
decode([1,0,2,2]) -> "hell"

題目整理

把基礎 vocab 強化成能處理真實輸入的邊界情況：(1) special tokens 如 <pad>、<unk>、<|endoftext|> 要保留固定 id；(2) 遇到訓練時沒看過的未知字元要回退到 <unk> 而不是 crash；(3) 空字串要回傳空 list；(4) 文字中若出現 special token 字面字串，要當成單一 token 切出來，而不是逐字元拆開。

解法說明

把 special tokens 放在 id 空間最前面（固定、可預期），一般字元接在後面。encode 先用正規表示式把 special token 從文字切出來，落在中間的普通片段再逐字元查表，查不到就用 stoi.get(ch, unk_id) 回退。decode 對稱處理，未知 id 也回退到 <unk> 字面。

import re
from typing import List, Tuple


class RobustVocab:
    def __init__(self, text: str, specials: Tuple[str, ...] = ("<pad>", "<unk>", "<|endoftext|>")):
        self.unk = "<unk>"
        chars = sorted(set(text))
        # Special tokens occupy the lowest, fixed ids.
        tokens = list(specials) + [c for c in chars if c not in specials]
        self.stoi = {tok: i for i, tok in enumerate(tokens)}
        self.itos = {i: tok for tok, i in self.stoi.items()}
        # Match any special token, longest first, so "<|endoftext|>" wins.
        escaped = sorted(specials, key=len, reverse=True)
        self.pattern = re.compile("(" + "|".join(re.escape(s) for s in escaped) + ")")

    def encode(self, s: str) -> List[int]:
        if not s:                        # Empty input is valid -> empty output.
            return []
        unk_id = self.stoi[self.unk]
        ids = []
        for chunk in self.pattern.split(s):
            if chunk == "":
                continue
            if chunk in self.stoi and chunk in ("<pad>", "<unk>", "<|endoftext|>"):
                ids.append(self.stoi[chunk])     # Whole special token.
            else:
                ids.extend(self.stoi.get(ch, unk_id) for ch in chunk)
        return ids

    def decode(self, ids: List[int]) -> str:
        return "".join(self.itos.get(i, self.unk) for i in ids)

Interview Explanation Flow

specials = <pad>=0, <unk>=1, <|endoftext|>=2
chars:  a=3, b=4, i=5, h=6

encode("hi<|endoftext|>")
split -> ["hi", "<|endoftext|>", ""]
"hi"  -> [6, 5]
token -> [2]
result -> [6, 5, 2]

encode("bzi")   # 'z' never seen
-> [4, 1, 5]    # z falls back to <unk>=1

encode("")  -> []

題目整理

給一整串已編碼的 token（一維 tensor data）、batch_size 與 block_size（context 長度），寫一個 get_batch 隨機抽出訓練 batch。回傳 x（shape (B, T)）與 y；y 是 x 整體右移一格的下一個 token，也就是每個位置的預測目標。

解法說明

用 torch.randint 抽 B 個起點 ix，範圍是 len(data) - block_size（保證 i + block_size 不越界，因為 y 還要多取一格）。每個起點切出長度 block_size 的片段當 x，起點 +1 的片段當 y，最後 torch.stack 疊成 batch。

import torch
from typing import Tuple


def get_batch(data: torch.Tensor, batch_size: int, block_size: int) -> Tuple[torch.Tensor, torch.Tensor]:
    # Random start offsets; -block_size keeps y's last index in range.
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i + block_size] for i in ix])
    y = torch.stack([data[i + 1:i + 1 + block_size] for i in ix])
    return x, y

Interview Explanation Flow

data = [10, 11, 12, 13, 14, 15]   block_size = 3
i = 1
x = data[1:4] = [11, 12, 13]
y = data[2:5] = [12, 13, 14]

# position-by-position targets:
# see 11        -> predict 12
# see 11,12     -> predict 13
# see 11,12,13  -> predict 14

題目整理

把上一題的隨機抽樣，改寫成標準的 torch.utils.data.Dataset，好讓官方 DataLoader 負責 shuffle、batching、多進程載入。需要實作 __len__ 與 __getitem__，每個 index 用滑動視窗回傳一組 (x, y)，兩者都長 block_size 且 y 相對 x 右移一格。

解法說明

__getitem__ 一次切出長度 block_size + 1 的 chunk，前 T 個當 x、後 T 個當 y，這樣只需一次 slice。__len__ 是 len(data) - block_size，確保最後一個 index 也切得出完整的 block_size + 1。

import torch
from torch.utils.data import Dataset
from typing import Tuple


class GPTDataset(Dataset):
    def __init__(self, data: torch.Tensor, block_size: int):
        self.data = data
        self.block_size = block_size

    def __len__(self) -> int:
        # Last valid start needs block_size + 1 tokens available.
        return len(self.data) - self.block_size

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
        chunk = self.data[idx:idx + self.block_size + 1]
        x = chunk[:-1]   # tokens 0 .. T-1
        y = chunk[1:]    # tokens 1 .. T   (shifted by one)
        return x, y


# Usage:
# loader = DataLoader(GPTDataset(data, block_size=8), batch_size=32, shuffle=True)

Interview Explanation Flow

block_size = 4, idx = 0
chunk = data[0:5] = [t0, t1, t2, t3, t4]
x = chunk[:-1] = [t0, t1, t2, t3]
y = chunk[1:]  = [t1, t2, t3, t4]

題目整理

實作一個 decoder-only GPT。輸入 token id idx（shape (B, T)），輸出每個位置對整個 vocab 的 logits (B, T, vocab_size)；若提供 targets 也回傳 cross-entropy loss。要點：token embedding + position embedding、堆疊多個 transformer block（每個 block 含 causal multi-head self-attention 與 feed-forward，皆用 pre-LayerNorm + residual），最後接 LayerNorm 與線性 lm_head。

解法說明

單個 attention head 算 Q, K, V，注意力分數 Q Kᵀ / √d，再用下三角 tril mask 把「看未來」的位置設成 -inf（causal），softmax 後對 V 加權。多頭就是把數個 head 的輸出 concat 再投影。Block 用 pre-norm 殘差：x = x + attn(ln1(x))、x = x + ffn(ln2(x))。GPT 把 token 與 position embedding 相加後過所有 block，最後 lm_head 投影到 vocab。

import torch
import torch.nn as nn
from torch.nn import functional as F


class Head(nn.Module):
    """One head of causal self-attention."""

    def __init__(self, n_embd: int, head_size: int, block_size: int, dropout: float):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        # Lower-triangular mask, registered as a buffer (not a parameter).
        self.register_buffer("tril", torch.tril(torch.ones(block_size, block_size)))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)      # (B, T, head_size)
        q = self.query(x)
        wei = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5   # (B, T, T)
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf"))   # Causal.
        wei = F.softmax(wei, dim=-1)
        wei = self.dropout(wei)
        v = self.value(x)
        return wei @ v       # (B, T, head_size)


class MultiHeadAttention(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        head_size = n_embd // n_head
        self.heads = nn.ModuleList(
            [Head(n_embd, head_size, block_size, dropout) for _ in range(n_head)]
        )
        self.proj = nn.Linear(n_embd, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)   # (B, T, n_embd)
        return self.dropout(self.proj(out))


class FeedForward(nn.Module):
    def __init__(self, n_embd, dropout):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),   # Expand.
            nn.ReLU(),
            nn.Linear(4 * n_embd, n_embd),   # Project back.
            nn.Dropout(dropout),
        )

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


class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        self.sa = MultiHeadAttention(n_embd, n_head, block_size, dropout)
        self.ff = FeedForward(n_embd, dropout)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        x = x + self.sa(self.ln1(x))   # Pre-norm + residual.
        x = x + self.ff(self.ln2(x))
        return x


class GPT(nn.Module):
    def __init__(self, vocab_size, n_embd, n_head, n_layer, block_size, dropout=0.1):
        super().__init__()
        self.block_size = block_size
        self.token_embedding = nn.Embedding(vocab_size, n_embd)
        self.position_embedding = nn.Embedding(block_size, n_embd)
        self.blocks = nn.Sequential(
            *[Block(n_embd, n_head, block_size, dropout) for _ in range(n_layer)]
        )
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok = self.token_embedding(idx)                                   # (B, T, C)
        pos = self.position_embedding(torch.arange(T, device=idx.device)) # (T, C)
        x = tok + pos                                                     # Broadcast add.
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)                                          # (B, T, vocab)

        loss = None
        if targets is not None:
            B, T, V = logits.shape
            loss = F.cross_entropy(logits.view(B * T, V), targets.view(B * T))
        return logits, loss

Interview Explanation Flow

x: (B, T, C) = (32, 8, 64),  n_head = 4 -> head_size = 16
q,k,v: (B, T, 16)
wei = q @ k.transpose(-2,-1) / sqrt(16)  -> (B, T, T) = (32, 8, 8)
masked_fill upper triangle with -inf, softmax over last dim
out = wei @ v -> (B, T, 16); concat 4 heads -> (B, T, 64)

題目整理

寫出 GPT 的訓練迴圈：用 AdamW 當 optimizer，每個 step 抽一個 batch、forward 得到 loss、zero_grad → backward → step，並週期性用不算梯度的方式估計 train/val loss 監控是否過擬合。

解法說明

estimate_loss 加 @torch.no_grad() 與 model.eval()（關掉 dropout），多抽幾個 batch 取平均降低雜訊，結束再切回 model.train()。訓練迴圈標準四步：zero_grad(set_to_none=True) 較省記憶體、loss.backward() 反傳、optimizer.step() 更新。

import torch


@torch.no_grad()
def estimate_loss(model, get_batch, eval_iters: int = 200):
    out = {}
    model.eval()                       # Disable dropout for stable measurement.
    for split in ("train", "val"):
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            x, y = get_batch(split)
            _, loss = model(x, y)
            losses[k] = loss.item()
        out[split] = losses.mean().item()
    model.train()                      # Back to training mode.
    return out


def train(model, get_batch, max_iters=5000, lr=3e-4, eval_interval=500):
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)

    for step in range(max_iters):
        if step % eval_interval == 0:
            losses = estimate_loss(model, get_batch)
            print(f"step {step}: train {losses['train']:.4f}, val {losses['val']:.4f}")

        xb, yb = get_batch("train")
        _, loss = model(xb, yb)
        optimizer.zero_grad(set_to_none=True)   # Clear old gradients.
        loss.backward()                         # Backprop.
        optimizer.step()                        # Update parameters.

    return model

Interview Explanation Flow

step 0:    train 4.20, val 4.21
step 500:  train 2.55, val 2.58
step 1000: train 2.10, val 2.19   # val lagging -> mild overfit starting
...

題目整理

實作自回歸生成 generate。給一段起始 context idx（shape (B, T)）與要生成的 max_new_tokens，反覆：把序列餵進模型 → 取最後一個位置的 logits → 用 temperature（與可選 top_k）調整 → softmax 成機率 → multinomial 抽下一個 token → 接回序列。注意每步都要把 context 裁到最多 block_size 個 token。

解法說明

只取 logits[:, -1, :] 是因為生成只關心「下一個」token。temperature < 1 讓分布更尖銳（保守），> 1 更發散（有創意）。top_k 只保留機率最高的 k 個，其餘設 -inf，避免抽到長尾雜訊。multinomial 依機率抽樣（不是永遠取 argmax，才有多樣性）。

import torch
from torch.nn import functional as F


@torch.no_grad()
def generate(model, idx, max_new_tokens, block_size, temperature=1.0, top_k=None):
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -block_size:]          # Never exceed the context window.
        logits, _ = model(idx_cond)
        logits = logits[:, -1, :] / temperature  # Only the last step matters.

        if top_k is not None:
            v, _ = torch.topk(logits, top_k)
            logits[logits < v[:, [-1]]] = float("-inf")   # Keep top-k only.

        probs = F.softmax(logits, dim=-1)
        next_id = torch.multinomial(probs, num_samples=1)  # Sample, not argmax.
        idx = torch.cat((idx, next_id), dim=1)             # Append.
    return idx

Interview Explanation Flow

idx = [[The, cat]]           # (1, 2)
logits = model(idx)          # (1, 2, vocab)
last = logits[:, -1, :]      # (1, vocab)  distribution after "cat"
probs = softmax(last / temperature)
next = multinomial(probs)    # e.g. "sat"
idx = [[The, cat, sat]]      # append, repeat

題目整理

上一題的生成每步都對整段 context 重算 attention，複雜度 O(T²)。KV-Cache 的洞見是：過去 token 的 Key 與 Value 不會改變，所以可以快取起來；每生成一個新 token，只需算它自己的 Q/K/V，把新的 K、V append 到快取，再讓新 Q 對「快取 + 自己」的 K、V 做 attention。實作一個吃 cache、回傳更新後 cache 的 self-attention。

解法說明

一次算好 qkv 再 split，reshape 成 (B, n_head, T, head_size)。若有 cache=(past_k, past_v)，就沿時間軸（dim=2）把新的 k、v 接到後面；attention 分數是新 Q 對「全部歷史 K」。在解碼階段每步 T == 1，新 query 天生只看得到過去，所以不需要 causal mask。

import torch
import torch.nn as nn
from torch.nn import functional as F


class CachedSelfAttention(nn.Module):
    def __init__(self, n_embd: int, n_head: int):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.head_size = n_embd // n_head
        self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)
        self.proj = nn.Linear(n_embd, n_embd)

    def forward(self, x, cache=None):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        # (B, T, C) -> (B, n_head, T, head_size)
        q = q.view(B, T, self.n_head, self.head_size).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_size).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_size).transpose(1, 2)

        if cache is not None:
            past_k, past_v = cache
            k = torch.cat([past_k, k], dim=2)   # Append along the time axis.
            v = torch.cat([past_v, v], dim=2)
        new_cache = (k, v)                      # Hand back for the next step.

        att = q @ k.transpose(-2, -1) * self.head_size ** -0.5
        # In decode mode T == 1: the single new query only sees past keys,
        # so an explicit causal mask is unnecessary.
        att = F.softmax(att, dim=-1)
        out = att @ v
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(out), new_cache

Interview Explanation Flow

Without cache, generating token 4 recomputes K,V for tokens 1..4.
With cache:
  step t: q_t, k_t, v_t = project(x_t)          # 1 token only
  K = cat(cache_K, k_t); V = cat(cache_V, v_t)  # append
  out = softmax(q_t @ Kᵀ / sqrt(d)) @ V         # attend over history
  cache = (K, V)                                # carry forward

題目整理

標準 multi-head attention 每個 query head 都有專屬的 KV head，KV-Cache 因此很大。GQA 讓 n_kv_head < n_head：把 query head 分成 n_kv_head 組，同一組的多個 query head 共用一組 K、V。實作 GQA：Q 投影出 n_head 組、K/V 只投影出 n_kv_head 組，再用 repeat_kv 把 KV 複製到與 Q 對齊後做 attention。（n_kv_head = 1 就退化成 MQA，n_kv_head = n_head 就是原本的 MHA。）

解法說明

關鍵是 n_rep = n_head // n_kv_head：每個 KV head 要被 n_rep 個 query head 共用。repeat_kv 用 expand（不實際複製記憶體）再 reshape 把 KV 從 (B, n_kv_head, T, d) 展開成 (B, n_head, T, d)，之後就跟一般 MHA 一樣算。省下的是 K/V 投影參數與 KV-Cache 記憶體（縮小 n_rep 倍）。

import torch
import torch.nn as nn
from torch.nn import functional as F


def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    # (B, n_kv_head, T, d) -> (B, n_kv_head * n_rep, T, d)
    B, n_kv_head, T, d = x.shape
    if n_rep == 1:
        return x
    x = x[:, :, None, :, :].expand(B, n_kv_head, n_rep, T, d)  # No real copy.
    return x.reshape(B, n_kv_head * n_rep, T, d)


class GroupedQueryAttention(nn.Module):
    def __init__(self, n_embd: int, n_head: int, n_kv_head: int):
        super().__init__()
        assert n_head % n_kv_head == 0
        self.n_head = n_head
        self.n_kv_head = n_kv_head
        self.n_rep = n_head // n_kv_head          # Queries per KV group.
        self.head_size = n_embd // n_head

        self.q_proj = nn.Linear(n_embd, n_head * self.head_size, bias=False)
        self.k_proj = nn.Linear(n_embd, n_kv_head * self.head_size, bias=False)
        self.v_proj = nn.Linear(n_embd, n_kv_head * self.head_size, bias=False)
        self.o_proj = nn.Linear(n_head * self.head_size, n_embd, bias=False)

    def forward(self, x):
        B, T, C = x.shape
        q = self.q_proj(x).view(B, T, self.n_head, self.head_size).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_kv_head, self.head_size).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_kv_head, self.head_size).transpose(1, 2)

        k = repeat_kv(k, self.n_rep)   # Share each KV head across its query group.
        v = repeat_kv(v, self.n_rep)

        att = q @ k.transpose(-2, -1) * self.head_size ** -0.5
        mask = torch.tril(torch.ones(T, T, device=x.device)).bool()
        att = att.masked_fill(~mask, float("-inf"))    # Causal.
        att = F.softmax(att, dim=-1)
        out = att @ v
        out = out.transpose(1, 2).contiguous().view(B, T, self.n_head * self.head_size)
        return self.o_proj(out)

Interview Explanation Flow

n_head = 8, n_kv_head = 2  ->  n_rep = 4
Q heads:  q0 q1 q2 q3 | q4 q5 q6 q7
KV heads:      kv0     |     kv1
group 0: q0..q3 all attend using kv0
group 1: q4..q7 all attend using kv1

repeat_kv: (B, 2, T, d) -> expand -> (B, 8, T, d)   # aligns with Q