挑战项目 C · 迷你生成式推荐：在笔记本上复现 2026 前沿

   你要亲手实现的三个前沿组件：

   ① RQ-VAE：把电影的内容向量量化成 3 层语义码字
      《星际穿越》 → <c₁=12, c₂=47, c₃=201>
      《盗梦空间》 → <c₁=12, c₂=47, c₃=88>     ← 前两层相同=模型知道它们相似！

   ② 生成式召回：一个小型 GPT，输入用户看过的语义 ID 序列，
      逐 token"写出"下一部电影的语义 ID

   ③ Trie 约束解码：保证生成的 ID 是真实存在的电影
      （否则模型会"幻觉"出不存在的电影）

   最终考试：
   🥊 和你自己训的 SASRec 打一架（Recall@10）
   ❄️ 冷启动挑战赛：故意藏起 5% 的电影不参与训练，
      看谁能把它们推出来 —— 生成式的招牌优势在这里兑现

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

class VQLayer(nn.Module):
    """一层向量量化：找最近的码字，EMA 更新码本"""
    def __init__(self, n_codes=256, dim=64, decay=0.99, eps=1e-5):
        super().__init__()
        self.decay, self.eps = decay, eps
        self.register_buffer("codebook", torch.randn(n_codes, dim) * 0.1)
        self.register_buffer("cluster_size", torch.zeros(n_codes))
        self.register_buffer("embed_avg", self.codebook.clone())

    def forward(self, x):                       # x: (B, D)
        d = (x.pow(2).sum(1, keepdim=True)
             - 2 * x @ self.codebook.T
             + self.codebook.pow(2).sum(1))     # (B, n_codes) 欧氏距离
        idx = d.argmin(1)                       # 每个样本最近的码字
        quant = self.codebook[idx]

        if self.training:
            onehot = F.one_hot(idx, self.codebook.size(0)).float()
            # EMA 统计
            self.cluster_size.mul_(self.decay).add_(onehot.sum(0), alpha=1-self.decay)
            self.embed_avg.mul_(self.decay).add_(onehot.T @ x, alpha=1-self.decay)
            n = self.cluster_size.sum()
            smoothed = (self.cluster_size + self.eps) / (n + self.codebook.size(0)*self.eps) * n
            self.codebook.copy_(self.embed_avg / smoothed.unsqueeze(1))
            # ⭐ 死码字重启：长期没人用的码字，随机跳到一个真实样本上
            dead = self.cluster_size < 0.01
            if dead.any():
                rand = x[torch.randint(0, x.size(0), (int(dead.sum()),))]
                self.codebook[dead] = rand

        # straight-through：前向用量化值，反向梯度直通
        return x + (quant - x).detach(), idx

class RQVAE(nn.Module):
    def __init__(self, in_dim=128, latent=64, n_layers=3, n_codes=256):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(in_dim, 256), nn.ReLU(), nn.Linear(256, latent))
        self.dec = nn.Sequential(nn.Linear(latent, 256), nn.ReLU(), nn.Linear(256, in_dim))
        self.vqs = nn.ModuleList([VQLayer(n_codes, latent) for _ in range(n_layers)])

    def forward(self, x):
        z = self.enc(x)
        residual, quant_sum, codes, commit = z, 0, [], 0.0
        for vq in self.vqs:                     # ⭐ 残差量化：逐层逼近
            q, idx = vq(residual)
            quant_sum = quant_sum + q
            commit = commit + F.mse_loss(residual, q.detach())
            residual = residual - q.detach()
            codes.append(idx)
        recon = self.dec(quant_sum)
        loss = F.mse_loss(recon, x) + 0.25 * commit
        return loss, torch.stack(codes, dim=1)  # (B, n_layers)

# 训练循环：普通 Adam，几千步就收敛（数据只有几千部电影）
# 训完：_, codes = model(content_vecs)  → 每部电影 3 个码字

class Trie:
    """所有合法电影的语义 ID 前缀树"""
    def __init__(self):
        self.root = {}
    def insert(self, tokens, movie_id):
        node = self.root
        for t in tokens:
            node = node.setdefault(t, {})
        node["_end"] = movie_id
    def valid_next(self, prefix):
        node = self.root
        for t in prefix:
            node = node.get(t)
            if node is None: return []
        return [k for k in node if k != "_end"]

@torch.no_grad()
def constrained_beam_search(model, user_tokens, trie, n_steps=4, beam=20, topk=10):
    """从用户历史 token 流出发，生成 top-k 合法语义 ID"""
    beams = [(user_tokens, 0.0)]                    # (序列, 累计 log 概率)
    for step in range(n_steps):
        cand = []
        for seq, lp in beams:
            logits = model(torch.tensor([seq]))[0, -1]        # 最后位置的预测
            prefix = seq[len(user_tokens):]                   # 已生成的部分
            allowed = trie.valid_next(prefix)                 # ⭐ 只许走合法分支
            if not allowed: continue
            logp = torch.log_softmax(logits, -1)
            for t in allowed:
                cand.append((seq + [t], lp + logp[t].item()))
        beams = sorted(cand, key=lambda x: -x[1])[:beam]
    results = []
    for seq, lp in beams:
        node = trie.root
        for t in seq[len(user_tokens):]: node = node[t]
        results.append((node["_end"], lp))                    # 反查电影 ID
    return [m for m, _ in results[:topk]]