12. ML Coding 专题

import math
from typing import Optional

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


def scaled_dot_product_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
                                 mask: Optional[torch.Tensor] = None) -> torch.Tensor:
    """
    Args:
        q: (B, H, T_q, d_k)；k: (B, H, T_k, d_k)；v: (B, H, T_k, d_v)
        mask: 可广播到 (B, H, T_q, T_k)，1 = 可见，0 = 屏蔽
    Returns:
        (B, H, T_q, d_v)
    """
    # 1. 打分并缩放：(B, H, T_q, d_k) @ (B, H, d_k, T_k) -> (B, H, T_q, T_k)
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(q.size(-1))
    # 2. 屏蔽必须在 softmax 之前：填 -1e9，softmax 后权重为 0
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    # 3. 沿 T_k 归一化，再加权求和 V
    weights = F.softmax(scores, dim=-1)                    # (B, H, T_q, T_k)
    return torch.matmul(weights, v)                        # (B, H, T_q, d_v)


class MultiHeadAttention(nn.Module):
    """self-attention 传 (x, x, x)；cross-attention 传 (decoder 状态, memory, memory)。"""

    def __init__(self, d_model: int, num_heads: int):
        super().__init__()
        assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除"
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)

    def split_heads(self, x: torch.Tensor) -> torch.Tensor:
        """(B, T, C) -> (B, H, T, d_k)"""
        batch_size, seq_len, _ = x.size()
        return x.reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)

    def combine_heads(self, x: torch.Tensor) -> torch.Tensor:
        """(B, H, T, d_k) -> (B, T, C)；transpose 后内存不连续，reshape 会自动复制（view 必须先 .contiguous()）"""
        batch_size, _, seq_len, _ = x.size()
        return x.transpose(1, 2).reshape(batch_size, seq_len, self.d_model)

    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
                mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        """query: (B, T_q, C)；key、value: (B, T_k, C)；返回 (B, T_q, C)"""
        # 1. 线性投影后切头
        q = self.split_heads(self.w_q(query))              # (B, H, T_q, d_k)
        k = self.split_heads(self.w_k(key))                # (B, H, T_k, d_k)
        v = self.split_heads(self.w_v(value))              # (B, H, T_k, d_k)
        # 2. H 个头并行算 attention（H 维和 B 维一样当 batch 维）
        out = scaled_dot_product_attention(q, k, v, mask)  # (B, H, T_q, d_k)
        # 3. 拼回 (B, T_q, C)，再过 w_o 融合各头
        return self.w_o(self.combine_heads(out))

class PositionwiseFeedForward(nn.Module):
    """逐 token 的前馈网络：变换特征，不混合不同位置的信息，也不添加位置信息。"""

    def __init__(self, d_model: int, d_ff: int):
        super().__init__()
        self.fc1 = nn.Linear(d_model, d_ff)  # C -> d_ff：投影到中间特征空间，通常 d_ff > C
        self.fc2 = nn.Linear(d_ff, d_model)  # d_ff -> C：恢复宽度，才能和输入做残差相加
        self.relu = nn.ReLU()               # 引入非线性，否则两次线性变换仍等价于一次

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """(B, T, C) -> (B, T, C)"""
        # nn.Linear 只作用于最后一维：B、T 不变，所有位置共用同一套权重
        # 1. fc1：(B, T, C) -> (B, T, d_ff)
        # 2. ReLU：负值置零，形状不变
        # 3. fc2：(B, T, d_ff) -> (B, T, C)
        return self.fc2(self.relu(self.fc1(x)))


class PositionalEncoding(nn.Module):
    """
    固定正弦位置编码：给每个位置一个 C 维向量，加到 token embedding 上。
    PE[pos, 2i] = sin(pos / 10000^(2i/C))，PE[pos, 2i+1] = cos(同上)。
    本示例要求 d_model 为偶数，输入长度 T <= max_len。
    """

    def __init__(self, d_model: int, max_len: int = 512, dropout: float = 0.1):
        super().__init__()
        # dropout 是相加后的正则化，不负责生成位置编码；model.eval() 时不丢弃元素
        self.dropout = nn.Dropout(dropout)
        # 1. 列出位置 pos = 0, 1, ..., max_len-1；加一维，便于和各频率广播相乘
        position = torch.arange(max_len).unsqueeze(1).float()                    # (max_len, 1)
        # 2. 每一对通道共用一个频率：exp(-log(10000) * 2i/C) = 1 / 10000^(2i/C)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))  # (C/2,)
        # 3. 广播得到 (max_len, C/2) 的相位；偶数通道填 sin，奇数通道填 cos
        pe = torch.zeros(max_len, d_model)                                       # (max_len, C)
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        # 4. 加 batch 维便于复用；buffer 不是可训练参数，随 .to(device) 移动并存入 state_dict
        self.register_buffer("pe", pe.unsqueeze(0))   # (1, max_len, C)：不训练，但跟着 .to(device) 走、存进 state_dict

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """(B, T, C) -> (B, T, C)"""
        # 1. 按实际长度 T 取前 T 个位置：(1, T, C)，同一套位置向量广播到 B 个样本
        # 2. 与 embedding 逐元素相加，形状仍为 (B, T, C)
        # 3. 对相加结果做 dropout；若只想加位置编码，可以省略这一步
        return self.dropout(x + self.pe[:, :x.size(1)])


class EncoderLayer(nn.Module):
    """self-attention -> FFN，每个子层 Post-LN：LN(x + Dropout(Sublayer(x)))"""

    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.ffn = PositionwiseFeedForward(d_model, d_ff)
        # 每个子层各有一套可学习的 LayerNorm 参数；沿最后的 C 维归一化
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        # 对子层输出做 dropout，再和子层输入相加；不是对残差相加后的整体做 dropout
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, src_mask: torch.Tensor) -> torch.Tensor:
        """x: (B, T_src, C)；src_mask: (B, 1, 1, T_src)。返回 (B, T_src, C)"""
        # 1. self-attention：Q = K = V = x，读取源句各位置；src_mask 只屏蔽 padding
        #    执行顺序：Attention -> Dropout -> 加原输入 x（残差）-> LayerNorm
        x = self.norm1(x + self.dropout(self.self_attn(x, x, x, src_mask)))
        # 2. FFN：每个 token 独立变换特征，再做 Dropout -> 残差相加 -> LayerNorm
        #    两个子层的输入输出均为 (B, T_src, C)，可以直接堆叠下一层
        return self.norm2(x + self.dropout(self.ffn(x)))


class DecoderLayer(nn.Module):
    """masked self-attention -> cross-attention -> FFN，每个子层 Post-LN"""

    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.cross_attn = MultiHeadAttention(d_model, num_heads)
        self.ffn = PositionwiseFeedForward(d_model, d_ff)
        # 三个子层分别归一化，参数互不共享；都只沿最后的 C 维计算
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, memory: torch.Tensor,
                src_mask: torch.Tensor, tgt_mask: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, T_tgt, C)，目标端当前表示
            memory: (B, T_src, C)，Encoder 输出，在各 Decoder 层中复用
            src_mask: (B, 1, 1, T_src)，屏蔽源端 padding
            tgt_mask: (B, 1, T_tgt, T_tgt)，屏蔽目标端未来位置和 padding
        Returns:
            (B, T_tgt, C)，每个目标位置融合前缀和源句信息后的表示
        """
        # 每步都按 Sublayer -> Dropout -> 残差相加 -> LayerNorm 执行（Post-LN）
        # 1. masked self-attention：Q = K = V = decoder，tgt_mask 挡住未来和 padding
        x = self.norm1(x + self.dropout(self.self_attn(x, x, x, tgt_mask)))
        # 2. cross-attention：Q 来自 decoder，K、V 来自 memory，注意力矩阵 (B, H, T_tgt, T_src)
        #    每个目标位置查询完整源句；输出长度由 Q 决定，仍为 T_tgt，memory 不被修改
        x = self.norm2(x + self.dropout(self.cross_attn(x, memory, memory, src_mask)))
        # 3. FFN：逐 token 做 C -> d_ff -> C，返回 (B, T_tgt, C)
        return self.norm3(x + self.dropout(self.ffn(x)))

def make_src_mask(src: torch.Tensor, pad_id: int = 0) -> torch.Tensor:
    """(B, T_src) -> (B, 1, 1, T_src)，True = 真实 token；广播到每个头、每个 query"""
    # 只限制哪些 key 可被读取：两个单例维分别对应 H 和 T_q
    # Encoder self-attention 和 Decoder cross-attention 都可以复用这个源端 mask
    return (src != pad_id).unsqueeze(1).unsqueeze(2)


def make_tgt_mask(tgt: torch.Tensor, pad_id: int = 0) -> torch.Tensor:
    """(B, T_tgt) -> (B, 1, T_tgt, T_tgt)：padding mask 与因果 mask 同时为 True 才可见"""
    seq_len = tgt.size(1)
    # 1. padding mask：每个 query 都不能读取 PAD key
    pad_mask = (tgt != pad_id).unsqueeze(1).unsqueeze(2)                                # (B, 1, 1, T)
    # 2. 因果 mask：行是 query 位置 i，列是 key 位置 j，只保留 j <= i（包括自己）
    causal = torch.tril(torch.ones(seq_len, seq_len, device=tgt.device)).bool()      # (T, T)，下三角
    # 3. 广播后取交集，得到每个样本的可见位置；H 维仍为 1，供所有头共用
    return pad_mask & causal                                                          # (B, 1, T, T)


class Transformer(nn.Module):
    """Encoder-Decoder Transformer：forward(src, tgt) -> logits (B, T_tgt, V)"""

    def __init__(self, src_vocab_size: int, tgt_vocab_size: int, d_model: int = 512,
                 num_heads: int = 8, num_layers: int = 6, d_ff: int = 2048,
                 max_len: int = 512, dropout: float = 0.1, pad_id: int = 0):
        super().__init__()
        self.d_model = d_model
        self.pad_id = pad_id
        # 源语言和目标语言分别查自己的词表：token id (B, T) -> 向量 (B, T, C)
        self.src_embed = nn.Embedding(src_vocab_size, d_model)
        self.tgt_embed = nn.Embedding(tgt_vocab_size, d_model)
        # 正弦编码没有可训练参数，两端共用同一张位置表；各自在入栈前调用一次
        self.pos_enc = PositionalEncoding(d_model, max_len, dropout)
        # 每次构造一个新层，各层参数独立；ModuleList 让 PyTorch 注册并管理这些子模块
        self.encoder_layers = nn.ModuleList(
            [EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])
        self.decoder_layers = nn.ModuleList(
            [DecoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])
        # 输出投影与 FFN 是不同用途：C -> V，为目标词表中每个 token 产生一个分数
        self.fc_out = nn.Linear(d_model, tgt_vocab_size)

    def encode(self, src: torch.Tensor, src_mask: torch.Tensor) -> torch.Tensor:
        """src: (B, T_src) -> memory: (B, T_src, C)"""
        # 1. 查 embedding -> 乘 sqrt(C) 调整量级 -> 加位置编码并 dropout
        x = self.pos_enc(self.src_embed(src) * math.sqrt(self.d_model))
        # 2. 依次经过 N 个 Encoder 层，形状始终为 (B, T_src, C)
        for layer in self.encoder_layers:
            x = layer(x, src_mask)
        # 3. 最后一层表示作为 memory，供每个 Decoder 层的 cross-attention 查询
        return x

    def decode(self, tgt: torch.Tensor, memory: torch.Tensor,
               src_mask: torch.Tensor, tgt_mask: torch.Tensor) -> torch.Tensor:
        """tgt: (B, T_tgt)，memory: (B, T_src, C) -> logits: (B, T_tgt, V)"""
        # 1. 目标前缀也要查自己的 embedding，并在入栈前加位置编码
        x = self.pos_enc(self.tgt_embed(tgt) * math.sqrt(self.d_model))
        # 2. 每层先读目标前缀，再查相同的源端 memory，形状始终为 (B, T_tgt, C)
        for layer in self.decoder_layers:
            x = layer(x, memory, src_mask, tgt_mask)
        # 3. 投影到词表：(B, T_tgt, C) -> (B, T_tgt, V)
        #    返回未归一化的 logits；训练时 CrossEntropyLoss 内部处理，无须先 softmax
        return self.fc_out(x)

    def forward(self, src: torch.Tensor, tgt: torch.Tensor) -> torch.Tensor:
        """src: (B, T_src)；tgt: (B, T_tgt)，以 BOS 开头的目标前缀（teacher forcing）"""
        # 1. 根据 token id 创建 mask：源端屏蔽 PAD，目标端同时屏蔽 PAD 和未来位置
        src_mask = make_src_mask(src, self.pad_id)
        tgt_mask = make_tgt_mask(tgt, self.pad_id)
        # 2. 源句编码成 memory，再结合目标前缀解码；训练时一次计算全部目标位置的 logits
        return self.decode(tgt, self.encode(src, src_mask), src_mask, tgt_mask)

def test_transformer() -> None:
    torch.manual_seed(42)
    model = Transformer(11, 13, d_model=32, num_heads=4, num_layers=2, d_ff=64).eval()
    src = torch.tensor([[5, 6, 7, 8, 9], [5, 6, 7, 0, 0]])    # (B=2, T_src=5)，第 2 句后两个是 PAD
    tgt = torch.tensor([[1, 3, 4, 5], [1, 6, 7, 8]])           # (B=2, T_tgt=4)，1 是 BOS
    logits = model(src, tgt)
    assert logits.shape == (2, 4, 13)                          # 1. (B, T_tgt, V)
    logits2 = model(src, torch.tensor([[1, 3, 9, 5], [1, 6, 9, 8]]))   # 2. 只改位置 2
    assert torch.allclose(logits[:, :2], logits2[:, :2], atol=1e-6)
    assert not torch.allclose(logits[:, 2], logits2[:, 2])
    assert torch.allclose(model(src[1:, :3], tgt[1:]), logits[1:], atol=1e-5)   # 3. padding 无影响
    print("all tests passed")


if __name__ == "__main__":
    test_transformer()   # all tests passed

import numpy as np


def apply_rope(x: np.ndarray, freq_base: float = 10000.0) -> np.ndarray:
    """x: (B, T, H, d_k)，d_k 为偶数（相邻两维一对）。Returns: 同形状，位置 m 的第 i 对通道转 m * theta_i。"""
    _, seq_len, _, d_k = x.shape
    assert d_k % 2 == 0, "d_k 必须是偶数才能两两旋转"
    # 1. 每对通道的频率 theta_i = base^(-2i/d_k)
    theta = 1.0 / (freq_base ** (np.arange(0, d_k, 2) / d_k))          # (d_k/2,)
    # 2. 角度 m * theta_i：位置和频率的外积
    freqs = np.outer(np.arange(seq_len), theta)                       # (T, d_k/2)
    # 3. cos/sin 每个复制两次，对齐交错排列的 d_k 维
    cos_val = np.repeat(np.cos(freqs), 2, axis=-1)[None, :, None, :]  # (1, T, 1, d_k)
    sin_val = np.repeat(np.sin(freqs), 2, axis=-1)[None, :, None, :]  # (1, T, 1, d_k)
    # 4. 旋转：每对 (x_even, x_odd) 变成 x * cos + (-x_odd, x_even) * sin
    swapped = np.stack([-x[..., 1::2], x[..., ::2]], axis=-1).reshape(x.shape)  # (B, T, H, d_k)
    return x * cos_val + swapped * sin_val


def test_rope() -> None:
    rng = np.random.default_rng(0)
    x = rng.normal(size=(2, 8, 2, 16))                                # (B, T, H, d_k)
    # 1. 旋转是正交变换，保范数
    assert np.allclose(np.linalg.norm(apply_rope(x), axis=-1), np.linalg.norm(x, axis=-1))
    # 2. 同一对 q、k 放到不同位置：点积只依赖 m - n
    q, k = rng.normal(size=16), rng.normal(size=16)
    rq = apply_rope(np.tile(q, (1, 8, 1, 1)))[0, :, 0]                # (T, d_k)，第 m 行是 R_m q
    rk = apply_rope(np.tile(k, (1, 8, 1, 1)))[0, :, 0]                # (T, d_k)
    assert np.isclose(rq[3] @ rk[1], rq[7] @ rk[5])                   # m - n 都是 2
    assert not np.isclose(rq[3] @ rk[1], rq[3] @ rk[2])               # 距离变了，点积跟着变
    print("all tests passed")


if __name__ == "__main__":
    test_rope()

from typing import Callable, Dict, List, Optional

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


@torch.no_grad()
def greedy_decode(model: nn.Module, src: torch.Tensor, max_len: int, bos_id: int, eos_id: int,
                  pad_id: int = 0, pick_fn: Optional[Callable[[torch.Tensor], torch.Tensor]] = None
                  ) -> torch.Tensor:
    """
    Args:
        src: (B, T_src)；pick_fn: logits (B, V) -> token (B,)，默认 argmax（greedy）
    Returns:
        (B, 1 + n)，以 BOS 开头，n <= max_len；每行 EOS 之后补 pad_id
    """
    model.eval()                                                          # 关掉 dropout
    src_mask = (src != pad_id).unsqueeze(1).unsqueeze(2)                  # (B, 1, 1, T_src)
    memory = model.encode(src, src_mask)                                  # 1. 源句只编码一次，(B, T_src, C)
    ys = torch.full((src.size(0), 1), bos_id, dtype=torch.long, device=src.device)   # (B, 1)
    finished = torch.zeros(src.size(0), dtype=torch.bool, device=src.device)        # (B,)
    for _ in range(max_len):
        tgt_mask = torch.tril(torch.ones(ys.size(1), ys.size(1), device=src.device))  # (T, T) 因果
        logits = model.decode(ys, memory, src_mask, tgt_mask)[:, -1, :]   # 2. 只取最后一个位置，(B, V)
        next_token = logits.argmax(dim=-1) if pick_fn is None else pick_fn(logits)   # (B,)
        next_token = next_token.masked_fill(finished, pad_id)             # 3. 已结束的行只补 PAD
        ys = torch.cat([ys, next_token.unsqueeze(1)], dim=1)              # (B, T + 1)
        finished = finished | (next_token == eos_id)                      # 4. 用 | 累积，PAD 不会把它改回去
        if finished.all():
            break
    return ys


def filter_logits(logits: torch.Tensor, temperature: float = 1.0, top_k: int = 0,
                  top_p: float = 1.0) -> torch.Tensor:
    """logits: (B, V) -> (B, V)；被过滤的位置置 -inf，softmax 后概率正好是 0，剩下的自动重新归一化"""
    logits = logits / temperature                                         # 1. 温度
    if top_k > 0:                                                         # 2. top-k：比每行第 k 大还小的屏蔽
        kth = torch.topk(logits, min(top_k, logits.size(-1)), dim=-1).values[:, -1:]   # (B, 1)
        logits = logits.masked_fill(logits < kth, float("-inf"))
    if top_p < 1.0:                                                       # 3. top-p：降序排序后算累计概率
        sorted_logits, sorted_idx = torch.sort(logits, dim=-1, descending=True)        # (B, V)
        probs = F.softmax(sorted_logits, dim=-1)
        sorted_remove = (probs.cumsum(dim=-1) - probs) >= top_p           # 前面已凑够 p 的都多余；第 1 个永远保留
        remove = sorted_remove.scatter(1, sorted_idx, sorted_remove)      # 排序后的顺序 -> 原顺序
        logits = logits.masked_fill(remove, float("-inf"))
    return logits


def sample_next(logits: torch.Tensor, temperature: float = 1.0, top_k: int = 0, top_p: float = 1.0,
                generator: Optional[torch.Generator] = None) -> torch.Tensor:
    """logits: (B, V) -> (B,)；temperature = 0 是 greedy 的极限，单独处理避免除零"""
    if temperature == 0:
        return logits.argmax(dim=-1)
    probs = F.softmax(filter_logits(logits, temperature, top_k, top_p), dim=-1)    # (B, V)
    return torch.multinomial(probs, num_samples=1, generator=generator).squeeze(-1)


@torch.no_grad()
def beam_search(model: nn.Module, src: torch.Tensor, max_len: int, bos_id: int, eos_id: int,
                beam_width: int = 3, alpha: float = 0.0, pad_id: int = 0) -> Dict:
    """
    Args:
        src: (1, T_src)，一条源句；alpha: 长度惩罚，排名用 score / L ** alpha
    Returns:
        {"tokens": [BOS, ...], "score": 累计 log-prob}
    """
    model.eval()
    src_mask = (src != pad_id).unsqueeze(1).unsqueeze(2)                  # (1, 1, 1, T_src)
    memory = model.encode(src, src_mask)                                  # (1, T_src, C)

    def rank(beam: Dict) -> float:
        return beam["score"] / max(len(beam["tokens"]) - 1, 1) ** alpha

    beams = [{"tokens": [bos_id], "score": 0.0}]
    for _ in range(max_len):
        candidates = []
        for beam in beams:
            if beam["tokens"][-1] == eos_id:      # 1. 已结束：原样带到下一轮，用自己的分数继续竞争
                candidates.append(beam)
                continue
            tgt = torch.tensor([beam["tokens"]], device=src.device)                       # (1, T)
            tgt_mask = torch.tril(torch.ones(tgt.size(1), tgt.size(1), device=src.device))
            log_probs = F.log_softmax(model.decode(tgt, memory, src_mask, tgt_mask)[0, -1], dim=-1)  # (V,)
            top_lp, top_ids = torch.topk(log_probs, min(beam_width, log_probs.numel()))  # 2. 每条只展开前 k 个
            for lp, tok in zip(top_lp.tolist(), top_ids.tolist()):
                candidates.append({"tokens": beam["tokens"] + [tok], "score": beam["score"] + lp})
        beams = sorted(candidates, key=rank, reverse=True)[:beam_width]  # 3. 全局剪枝回 k 条
        if all(b["tokens"][-1] == eos_id for b in beams):                 # 4. 全部结束就停
            break
    return beams[0]                                                       # 已按 rank 降序

import itertools


class StubSeq2Seq(nn.Module):
    """只为测试解码：随机初始化，接口同第 1 节的 Transformer，mask 参数接收但不用"""

    def __init__(self, vocab_size: int = 8, d_model: int = 16):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.fc_out = nn.Linear(d_model, vocab_size)

    def encode(self, src, src_mask):
        return self.embed(src)                                                    # (B, T_src, C)

    def decode(self, tgt, memory, src_mask, tgt_mask):
        h = self.embed(tgt).cumsum(dim=1) + memory.mean(dim=1, keepdim=True)     # (B, T_tgt, C)
        logits = self.fc_out(torch.tanh(h))                                       # (B, T_tgt, V)
        return logits.index_fill(-1, torch.tensor([0, 1]), -1e9)                  # 不生成 PAD、BOS


def test_decoding() -> None:
    torch.manual_seed(45)
    model, bos, eos = StubSeq2Seq(), 1, 2
    src = torch.tensor([[3, 4, 5], [6, 7, 0]])                                    # (B=2, T_src=3)
    out = greedy_decode(model, src, 6, bos, eos)
    print(out.tolist())
    # 1. 采样时 top_k = 1 等于 greedy；top-p 小例子：p = 0.7 只留 0.5 和 0.3
    assert torch.equal(greedy_decode(model, src, 6, bos, eos, pick_fn=lambda z: sample_next(z, top_k=1)), out)
    toy = torch.log(torch.tensor([[0.15, 0.5, 0.05, 0.3]]))
    assert torch.isfinite(filter_logits(toy, top_p=0.7)).tolist() == [[False, True, False, True]]
    # 2. beam：宽度 1 等于 greedy；宽度 64 时前两步一条都不剪（8 条、7 * 8 + 1 条），结果等于暴力枚举
    src1, memory = src[:1], model.encode(src[:1], None)
    beams = {w: beam_search(model, src1, 3, bos, eos, beam_width=w) for w in [1, 3, 64]}
    for w, beam in beams.items():
        print(w, beam["tokens"], round(beam["score"], 4))
    assert beams[1]["tokens"] == greedy_decode(model, src1, 3, bos, eos)[0].tolist()

    def full_score(tokens: List[int]) -> float:   # 一次完整前向给整条序列打分，和 beam 的逐步累加相互独立
        lp = F.log_softmax(model.decode(torch.tensor([tokens]), memory, None, None)[0], dim=-1)   # (T, V)
        return sum(lp[t, tokens[t + 1]].item() for t in range(len(tokens) - 1))
    conts = {c[:c.index(eos) + 1] if eos in c else c for c in itertools.product(range(8), repeat=3)}
    best = max(([bos] + list(c) for c in conts), key=full_score)
    assert beams[64]["tokens"] == best and abs(beams[64]["score"] - full_score(best)) < 1e-5
    print(beam_search(model, src1, 3, bos, eos, beam_width=64, alpha=1.0)["tokens"])
    print("all tests passed")


if __name__ == "__main__":
    test_decoding()
    # [[1, 6, 6, 6, 2, 0, 0], [1, 6, 6, 6, 6, 4, 2]]
    # 1 [1, 6, 6, 6] -4.4276
    # 3 [1, 7, 4, 4] -3.3934
    # 64 [1, 2] -2.0994
    # [1, 7, 4, 4]
    # all tests passed

import numpy as np


class KMeans:
    def __init__(self, n_clusters: int, max_iter: int = 100, tol: float = 1e-4):
        self.n_clusters = n_clusters
        self.max_iter = max_iter
        self.tol = tol

    def fit(self, X: np.ndarray) -> "KMeans":
        """X: (n, d)。分配到最近中心，再用簇内均值更新中心。"""
        X = np.asarray(X, dtype=float)  # 中心要存均值，使用浮点数
        # 1. 随机选 k 个不同样本作为初始中心：(k, d)
        indices = np.random.choice(len(X), self.n_clusters, replace=False)
        self.centroids = X[indices].copy()

        for _ in range(self.max_iter):
            # 2. 算距离，找归属
            distances = np.sum((X[:, None, :] - self.centroids[None, :, :]) ** 2, axis=2)
            labels = np.argmin(distances, axis=1)

            # 3. 更新：对每个簇沿样本维求均值，得到新的中心 (k, d)
            # 3. 更新中心
            new_centroids = np.array([X[labels == i].mean(axis=0) for i in range(self.n_clusters)])

            # 4. 如果中心不动了，就停止
            if np.all(centroids == new_centroids):
                break
            self.centroids = new_centroids

def squared_distance(p: list, q: list) -> float:
    return sum((a - b) ** 2 for a, b in zip(p, q))


def solution(data: list, k: int, max_iter: int = 100) -> list:
    """data: n 个 d 维点。返回每个点的簇下标。"""
    centroids = [list(p) for p in data[:k]]
    for _ in range(max_iter):
        # 1. 分配：min 返回距离最小的中心下标
        labels = [min(range(k), key=lambda j: squared_distance(p, centroids[j])) for p in data]

        # 2. 更新：逐簇计算每一维的均值；空簇保留旧中心
        new_centroids = []
        for j in range(k):
            points = [p for p, label in zip(data, labels) if label == j]
            if points:
                mean = [sum(col) / len(points) for col in zip(*points)]
                new_centroids.append(mean)
            else:
                new_centroids.append(centroids[j])

        # 相同的分组会算出相同的均值；这里比较列表，不引入额外的容差参数
        if new_centroids == centroids:
            break
        centroids = new_centroids

    return [min(range(k), key=lambda j: squared_distance(p, centroids[j])) for p in data]

np.random.seed(0)  # 固定随机种子，让例子可复现
X = np.array([[0, 0], [0, 1], [1, 0], [9, 9], [9, 10], [10, 9]])
model = KMeans(n_clusters=2).fit(X)
print(model.centroids)  # 两个中心约为 (0.33, 0.33) 和 (9.33, 9.33)，编号可能相反
print(model.labels_)   # 前三个点一簇，后三个点一簇

# 簇编号本身没有意义：检查两组点各自同簇、两组之间不同簇
labels = model.labels_
assert np.all(labels[:3] == labels[0])
assert np.all(labels[3:] == labels[3])
assert labels[0] != labels[3]
assert np.allclose(model.centroids[labels[0]], [1 / 3, 1 / 3])
assert np.allclose(model.centroids[labels[3]], [28 / 3, 28 / 3])
assert np.array_equal(model.predict(np.array([[0, 0], [10, 10]])), labels[[0, 3]])
assert solution([[0, 0], [10, 10], [0, 1], [10, 9]], k=2) == [0, 1, 0, 1]
print("all tests passed")

from collections import Counter

import numpy as np


class KNN:
    def __init__(self, k: int = 3):
        self.k = k

    def fit(self, X: np.ndarray, y: np.ndarray) -> "KNN":
        """X: (n, d)，y: (n,)。保存数据，没有参数训练。"""
        self.X_train = X
        self.y_train = y
        return self

    def predict(self, X: np.ndarray) -> np.ndarray:
        """X: (m, d) -> 预测类别 (m,)"""
        predictions = []
        for x in X:
            # 1. 当前点 (d,) 与全部训练点 (n, d) 相减，得到 n 个平方距离
            distances = ((self.X_train - x) ** 2).sum(axis=1)  # (n,)，不用开根号
            # 2. 按距离从小到大排序，取前 k 个训练样本的下标
            indices = np.argsort(distances)[:self.k]         # (k,)
            # 3. 取邻居标签，多数投票得到当前点的预测类别
            neighbor_labels = self.y_train[indices]          # (k,)
            prediction = Counter(neighbor_labels).most_common(1)[0][0]
            predictions.append(prediction)
        return np.array(predictions)

X_train = np.array([[0, 0], [0, 1], [1, 0], [9, 9], [9, 10], [10, 9]])
y_train = np.array([0, 0, 0, 1, 1, 1])
X_test = np.array([[0.2, 0.2], [9.2, 9.2]])

model = KNN(k=3).fit(X_train, y_train)
predictions = model.predict(X_test)
print(predictions)  # [0 1]
assert np.array_equal(predictions, [0, 1])

# 用 k=1 检查预测确实跟随最近样本的标签；标签也可以是字符串
model = KNN(k=1).fit(X_train, np.array(["a", "b", "c", "d", "e", "f"]))
assert np.array_equal(model.predict(np.array([[0, 0.9], [9.9, 9]])), ["b", "f"])
print("all tests passed")

from typing import List, Optional

import numpy as np


def sigmoid(z: np.ndarray) -> np.ndarray:
    """数值稳定的 sigmoid。z: 任意形状，返回同形状"""
    z = np.asarray(z, dtype=float)
    out = np.empty_like(z)
    pos = z >= 0
    out[pos] = 1.0 / (1.0 + np.exp(-z[pos]))      # 1. z >= 0：exp(-z) <= 1
    exp_z = np.exp(z[~pos])                        # 2. z < 0：exp(z) < 1
    out[~pos] = exp_z / (1.0 + exp_z)
    return out


def bce_with_logits(logits: np.ndarray, y: np.ndarray) -> float:
    """logits, y: (n,)。返回平均 loss：max(z, 0) - z*y + log(1 + exp(-|z|))"""
    logits = np.asarray(logits, dtype=float)
    y = np.asarray(y, dtype=float)
    loss = np.maximum(logits, 0) - logits * y + np.log1p(np.exp(-np.abs(logits)))
    return float(loss.mean())


class LogisticRegression:
    def __init__(self, lr: float = 0.1, max_iter: int = 1000, l2: float = 0.0):
        self.lr = lr
        self.max_iter = max_iter
        self.l2 = l2                                       # L2 正则强度，截距不参与
        self.w: Optional[np.ndarray] = None                # (d,)
        self.b = 0.0
        self.loss_history: List[float] = []

    def fit(self, X: np.ndarray, y: np.ndarray) -> "LogisticRegression":
        """X: (n, d)，y: (n,)，取值 0/1"""
        X = np.asarray(X, dtype=float)
        y = np.asarray(y, dtype=float).reshape(-1)         # 拉平成 (n,)，防止 (n, 1) 广播成 (n, n)
        n, d = X.shape
        self.w, self.b = np.zeros(d), 0.0                  # 凸问题，零初始化即可
        self.loss_history = []
        for _ in range(self.max_iter):
            logits = X @ self.w + self.b                   # 1. 前向只算到 logits (n,)
            reg = 0.5 * self.l2 * float(self.w @ self.w)
            self.loss_history.append(bce_with_logits(logits, y) + reg)
            err = sigmoid(logits) - y                      # 2. dL/dlogits = p - y，(n,)
            self.w -= self.lr * (X.T @ err / n + self.l2 * self.w)   # 3. 加上 L2 的梯度 λw
            self.b -= self.lr * float(err.mean())          # 截距只管整体基准，不正则
        return self

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        """X: (m, d)。返回正类概率 (m,)"""
        if self.w is None:
            raise RuntimeError("call fit before predict")
        return sigmoid(np.asarray(X, dtype=float) @ self.w + self.b)

    def predict(self, X: np.ndarray, threshold: float = 0.5) -> np.ndarray:
        """返回 0/1 标签 (m,)"""
        return (self.predict_proba(X) >= threshold).astype(int)

from sklearn.linear_model import LogisticRegression as SkLogisticRegression


def test_logistic_regression() -> None:
    z = np.linspace(-30, 30, 61)
    assert np.allclose(sigmoid(z), 1 / (1 + np.exp(-z)))                     # 1. 正常范围和原公式一致
    assert sigmoid(np.array([-1000.0, 1000.0])).tolist() == [0.0, 1.0]        # 2. 极端值无溢出、无 warning
    assert np.isclose(bce_with_logits(np.array([40.0]), np.array([0.0])), 40.0)   # 原公式这里是 inf
    rng = np.random.default_rng(0)
    X = rng.normal(size=(200, 3))                                             # (200, 3)
    y = (X @ np.array([1.0, -2.0, 0.5]) + 0.3 + rng.normal(size=200) > 0).astype(int)
    model = LogisticRegression(lr=0.5, max_iter=2000, l2=0.01).fit(X, y)
    ref = SkLogisticRegression(C=1 / (200 * 0.01), tol=1e-8).fit(X, y)       # 3. 和 sklearn 对拍
    assert np.allclose(model.w, ref.coef_[0], atol=1e-4)
    assert np.isclose(model.b, ref.intercept_[0], atol=1e-4)
    assert np.all(np.diff(model.loss_history) <= 1e-12)                       # 4. 凸问题，loss 单调下降
    print("all tests passed")


if __name__ == "__main__":
    test_logistic_regression()

import numpy as np


def entropy(y: np.ndarray) -> float:
    """y: (n,) 标签。返回以 2 为底的熵"""
    if len(y) == 0:
        return 0.0
    _, counts = np.unique(y, return_counts=True)       # 各类别频次
    p = counts / len(y)                                 # np.unique 保证 p > 0，不会 log(0)
    return float(-np.sum(p * np.log2(p)))


def information_gain(y_parent: np.ndarray, y_left: np.ndarray, y_right: np.ndarray) -> float:
    """父节点熵减去子节点熵的加权平均；不加权的话，只有 2 个样本的纯净子节点会把分数骗高"""
    n = len(y_parent)
    child = len(y_left) / n * entropy(y_left) + len(y_right) / n * entropy(y_right)
    return entropy(y_parent) - child


if __name__ == "__main__":
    y = np.array([0, 0, 1, 1])
    assert entropy(y) == 1.0 and entropy(np.array([1, 1, 1])) == 0.0
    assert information_gain(y, y[:2], y[2:]) == 1.0                    # 完美切分，增益等于父节点的熵
    assert np.isclose(information_gain(y, y[[0, 2]], y[[1, 3]]), 0.0)  # 两边仍是一半一半，没有增益
    print("all tests passed")

import numpy as np


def cosine_similarity_matrix(A: np.ndarray, B: np.ndarray, eps: float = 1e-8) -> np.ndarray:
    """A: (n, d)，B: (m, d)。返回 (n, m)，第 i 行第 j 列是 cos(A[i], B[j])；单个向量传 a[None]"""
    # 1. 每行除以自己的 L2 模长，(n, d) 和 (m, d)；eps 防零向量除出 nan
    A_unit = A / (np.linalg.norm(A, axis=1, keepdims=True) + eps)
    B_unit = B / (np.linalg.norm(B, axis=1, keepdims=True) + eps)
    # 2. 归一化之后内积就是余弦，一次矩阵乘拿到全部 n * m 对
    return A_unit @ B_unit.T

if __name__ == "__main__":
    from sklearn.metrics.pairwise import cosine_similarity

    rng = np.random.default_rng(0)
    A, B = rng.normal(size=(4, 8)), rng.normal(size=(3, 8))
    S = cosine_similarity_matrix(A, B)
    assert S.shape == (4, 3) and np.allclose(S, cosine_similarity(A, B))
    assert np.allclose(np.diag(cosine_similarity_matrix(A, 5 * A)), 1.0)    # 只看方向，不看长度
    assert np.allclose(cosine_similarity_matrix(np.zeros((1, 8)), B), 0.0)  # 零向量得到 0，不出 nan
    print("all tests passed")

from typing import List, Tuple


def precision_recall_f1(y_true: List[int], y_pred: List[int]) -> Tuple[float, float, float]:
    """二分类，y_true / y_pred 是长度 n 的 0/1 列表。分母为 0 时返回 0（同 sklearn 的 zero_division=0）。"""
    # 1. 数混淆矩阵里的 TP、FP、FN
    tp = sum(1 for t, p in zip(y_true, y_pred) if t == 1 and p == 1)
    fp = sum(1 for t, p in zip(y_true, y_pred) if t == 0 and p == 1)
    fn = sum(1 for t, p in zip(y_true, y_pred) if t == 1 and p == 0)
    # 2. 一个正类都没预测出来，或者没有正样本时，分母为 0，显式返回 0
    precision = tp / (tp + fp) if tp + fp > 0 else 0.0
    recall = tp / (tp + fn) if tp + fn > 0 else 0.0
    f1 = 2 * precision * recall / (precision + recall) if precision + recall > 0 else 0.0
    return precision, recall, f1

def solution(labels: list, scores: list) -> float:
    """labels: 长度 n，1/True 为正类；scores: 长度 n，概率或 logits 都行。返回 [0, 1] 的 AUC。"""
    # 1. 按标签分成两堆；if y 同时兼容 1/0 和 True/False
    pos = [s for y, s in zip(labels, scores) if y]
    neg = [s for y, s in zip(labels, scores) if not y]
    if not pos or not neg:      # 只有一类时 AUC 无定义：报错，别偷偷返回 0.5
        raise ValueError("only one class present, AUC is undefined")
    # 2. 每个（正, 负）对：正样本分高记 1，相等记 0.5
    wins = 0.0
    for p in pos:
        for q in neg:
            if p > q:
                wins += 1.0
            elif p == q:
                wins += 0.5
    # 3. 除以总对数 P * N
    return wins / (len(pos) * len(neg))


print(solution([1, 1, 0, 1, 0, 0], [0.9, 0.8, 0.8, 0.4, 0.3, 0.1]))   # 0.8333333333333334

def auc_rank(labels: list, scores: list) -> float:
    """排名公式，纯 Python。labels / scores 同 solution。时间 O(n log n)。"""
    n, n_pos = len(scores), sum(1 for y in labels if y)
    n_neg = n - n_pos
    if n_pos == 0 or n_neg == 0:
        raise ValueError("only one class present, AUC is undefined")
    order = sorted(range(n), key=lambda i: scores[i])     # 1. 按分数升序排下标
    # 2. 同分组 order[i..j] 占名次 i+1 .. j+1，统一给平均名次
    ranks = [0.0] * n
    i = 0
    while i < n:
        j = i
        while j + 1 < n and scores[order[j + 1]] == scores[order[i]]:
            j += 1
        for k in range(i, j + 1):
            ranks[order[k]] = (i + j) / 2 + 1
        i = j + 1
    # 3. 正样本名次和减去 P(P+1)/2 得到赢的对数，再除以 P * N
    rank_sum = sum(r for r, y in zip(ranks, labels) if y)
    return (rank_sum - n_pos * (n_pos + 1) / 2) / (n_pos * n_neg)

def gauc(user_ids: list, labels: List[int], scores: List[float]) -> float:
    """按曝光数加权；用上面的 auc_rank。"""
    groups = {}                                         # 1. 用户 -> 样本下标
    for i, u in enumerate(user_ids):
        groups.setdefault(u, []).append(i)
    num, den = 0.0, 0.0
    for idx in groups.values():
        y = [labels[i] for i in idx]
        if 0 < sum(y) < len(y):                         # 2. 两类都有才算
            num += len(idx) * auc_rank(y, [scores[i] for i in idx])   # 3. 权重 = 曝光数
            den += len(idx)
    return num / den


users = ["a"] * 4 + ["b"] * 4 + ["c"] * 2
labels = [1, 0, 1, 0, 0, 1, 0, 0, 0, 0]                # 用户 c 没有点击，跳过
scores = [0.8, 0.9, 0.7, 0.6, 0.3, 0.2, 0.1, 0.4, 0.5, 0.35]
print(auc_rank(labels, scores), gauc(users, labels, scores))   # 0.6190476190476191 0.41666666666666663

if __name__ == "__main__":
    import random
    from sklearn.metrics import precision_recall_fscore_support, roc_auc_score
    random.seed(0)
    for _ in range(200):
        n = random.randint(2, 30)
        y = [0, 1] + [random.randint(0, 1) for _ in range(n - 2)]   # 保证两类都在
        s = [random.randint(0, 4) / 4 for _ in range(n)]             # 只有 5 种分数，平局很多
        pred = [int(v >= 0.5) for v in s]
        ref = precision_recall_fscore_support(y, pred, average="binary", zero_division=0)[:3]
        assert all(abs(a - b) < 1e-12 for a, b in zip(precision_recall_f1(y, pred), ref))
        assert all(abs(f(y, s) - roc_auc_score(y, s)) < 1e-12 for f in (solution, auc_rank))
    print("all tests passed")

import numpy as np


def softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
    """x: 任意形状，例如 (B, C)；沿 axis 归一化，返回同形状，每行和为 1"""
    # 1. 减每行最大值：数学上等价，最大项变成 e^0 = 1，不会上溢
    e = np.exp(x - x.max(axis=axis, keepdims=True))
    # 2. keepdims=True 让 (B, 1) 和 (B, C) 广播，每行各自归一化
    return e / e.sum(axis=axis, keepdims=True)


def cross_entropy_from_logits(logits: np.ndarray, y: np.ndarray) -> float:
    """logits: (B, C) 未经 softmax 的分数；y: (B,) 整数类别（不是 one-hot）。返回 batch 平均的交叉熵"""
    # 1. log-softmax = (z - max) - log Σ e^(z - max)，(B, C)；没有除法，也不会 log(0)
    shifted = logits - logits.max(axis=1, keepdims=True)
    log_probs = shifted - np.log(np.exp(shifted).sum(axis=1, keepdims=True))
    # 2. 花式索引：第 i 行取第 y[i] 列，(B,)
    return float(-log_probs[np.arange(len(y)), y].mean())


z = np.array([[1000.0, 0.0], [1.0, 2.0]])            # (2, 2)，第 0 行直接 np.exp 会溢出成 inf
print(softmax(z).round(4).tolist())                  # [[1.0, 0.0], [0.2689, 0.7311]]
print(f"{cross_entropy_from_logits(z, np.array([1, 1])):.4f}")   # 500.1566：p 下溢成 0 也不会得到 inf

import math
from typing import Callable, Optional

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


def reduce_loss(loss: torch.Tensor, reduction: str) -> torch.Tensor:
    """loss: 逐元素 loss；reduction: 'none' | 'sum' | 'mean'"""
    if reduction == "none":
        return loss
    return loss.sum() if reduction == "sum" else loss.mean()


def check(my_fn: Callable, ref_fn: Callable, x: torch.Tensor, atol: float = 1e-6) -> bool:
    """对拍：my_fn、ref_fn 输入 x、返回标量 loss；loss 值和对 x 的梯度都一致才返回 True"""
    x1, x2 = x.detach().clone().requires_grad_(True), x.detach().clone().requires_grad_(True)
    loss1, loss2 = my_fn(x1), ref_fn(x2)
    g1, g2 = torch.autograd.grad(loss1, x1)[0], torch.autograd.grad(loss2, x2)[0]
    return torch.allclose(loss1, loss2, atol=atol) and torch.allclose(g1, g2, atol=atol)


def cross_entropy_loss(logits: torch.Tensor, target: torch.Tensor, ignore_index: int = -100,
                       label_smoothing: float = 0.0, reduction: str = "mean") -> torch.Tensor:
    """logits: (N, C) 未经 softmax 的分数；target: (N,) 整数类别，等于 ignore_index 的位置不计入 loss。
    返回：'none' 时 (N,)，忽略位置为 0；否则标量（全部被忽略时 mean 是 nan，和内置一致）"""
    log_probs = F.log_softmax(logits, dim=1)                         # 1. (N, C)
    valid = target != ignore_index                                    # 2. (N,) True 表示要算
    safe_target = target.masked_fill(~valid, 0)                       # 3. 忽略位置换成合法下标 0
    nll = -log_probs.gather(1, safe_target.unsqueeze(1)).squeeze(1)   # 4. (N,) 真实类别的 -log p
    smooth = -log_probs.mean(dim=1)                                   # 5. (N,) 对均匀分布的交叉熵
    loss = (1 - label_smoothing) * nll + label_smoothing * smooth     # 6. (N,)
    loss = loss.masked_fill(~valid, 0.0)                              # 7. 忽略位置清零
    if reduction == "mean":
        return loss.sum() / valid.sum()                               # 分母 = 有效位置数
    return reduce_loss(loss, reduction)

def bce_with_logits(logits: torch.Tensor, target: torch.Tensor,
                    pos_weight: Optional[torch.Tensor] = None, reduction: str = "mean") -> torch.Tensor:
    """logits: (N,) 或 (N, C)；target: 同形状的 0/1（或软标签）；pos_weight: 标量或 (C,)"""
    w_p = 1.0 if pos_weight is None else pos_weight
    pos_term = w_p * target * F.softplus(-logits)    # 1. -w_p · y · log σ(x)
    neg_term = (1 - target) * F.softplus(logits)     # 2. -(1 - y) · log(1 - σ(x))
    return reduce_loss(pos_term + neg_term, reduction)