循环神经网络 RNN 与词嵌入

"""one-hot 与 embedding 的正面对照：为什么查表比稀疏编码划算。

结论可以纯算术推导，不需要跑起来也能验证：
    词表 50000 词、想要 128 维表示
    one-hot   : 每个词 50000 维，其中 49999 个是 0
    embedding : 每个词 128 维，全是有意义的小数
"""
import torch
import torch.nn as nn


def one_hot(idx, vocab_size):
    """手写 one-hot：开一排 0，只有第 idx 位是 1。"""
    vec = torch.zeros(vocab_size)
    vec[idx] = 1.0
    return vec


def main():
    torch.manual_seed(0)

    vocab = ['我', '爱', '你', '北京', '天安门']
    vocab_size, embed_dim = len(vocab), 4

    # ------------------------------------------------------------ one-hot
    print('=== one-hot ===')
    for i, w in enumerate(vocab):
        print('%6s' % w, one_hot(i, vocab_size).tolist())
    # 看得出来三件事：
    #   1) 维度 = 词表大小，词表一大就爆炸
    #   2) 任意两个词的向量点积都是 0，"北京" 和 "天安门" 一样不相关
    #   3) 全是 0 和 1，没有任何可学的东西

    # ------------------------------------------------------------ embedding
    print('=== embedding ===')
    embed = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embed_dim)
    print('嵌入矩阵形状', embed.weight.shape)       # (5, 4)

    ids = torch.tensor([0, 1, 2])                   # 我 爱 你
    vecs = embed(ids)
    print('三个词一次查表的输出形状', vecs.shape)     # (3, 4)

    # 查表的本质：embedding 的输出 = one-hot 矩阵乘嵌入矩阵，
    # 只是没人会真去乘一个 99.998% 都是 0 的矩阵，直接按下标取行就行
    manual = torch.stack([one_hot(i, vocab_size) for i in [0, 1, 2]]) @ embed.weight
    print('按下标取行 与 one-hot 乘矩阵 是否一致:',
          torch.allclose(vecs, manual))             # True

    # ------------------------------------------------------------ 规模账
    print('=== 规模账（50000 词，128 维）===')
    big_vocab, dim = 50000, 128
    print('one-hot   单词向量元素数 :', big_vocab)
    print('embedding 单词向量元素数 :', dim)
    print('embedding 矩阵总参数    :', big_vocab * dim)   # 6,400,000
    print('压缩比                  : %.1f 倍' % (big_vocab / dim))  # 390.6 倍

    # 注意：embedding 矩阵是网络参数，会被反向传播更新。
    # 一开始是随机数，训练之后语义相近的词，向量也会靠得更近。
    print('嵌入矩阵需要梯度吗:', embed.weight.requires_grad)   # True


if __name__ == '__main__':
    main()

"""手写一个 RNN 单元，再和 nn.RNN 对账。

单元内部只有一条公式：
    h_t = tanh(x_t @ W_ih^T + b_ih  +  h_{t-1} @ W_hh^T + b_hh)
把 nn.RNN 的权重搬过来自己算一遍，两边数值一致才算真看懂。
"""
import torch
import torch.nn as nn


def rnn_cell(x_t, h_prev, w_ih, b_ih, w_hh, b_hh):
    """一个时间步：吃「这一步的词向量」和「上一步的 h」，吐新的 h。"""
    from_input = x_t @ w_ih.T + b_ih        # 当前输入贡献的那一份
    from_state = h_prev @ w_hh.T + b_hh     # 历史记忆贡献的那一份
    return torch.tanh(from_input + from_state)   # 两份相加再压到 (-1, 1)


def main():
    torch.manual_seed(0)

    input_size, hidden_size, seq_len = 3, 4, 5

    rnn = nn.RNN(input_size=input_size, hidden_size=hidden_size, num_layers=1)

    # nn.RNN 把四个参数摊在 _l0 后缀的属性里，形状分别是：
    #   weight_ih_l0 (hidden_size, input_size)
    #   weight_hh_l0 (hidden_size, hidden_size)
    #   bias_ih_l0 / bias_hh_l0 (hidden_size,)
    w_ih, b_ih = rnn.weight_ih_l0.detach(), rnn.bias_ih_l0.detach()
    w_hh, b_hh = rnn.weight_hh_l0.detach(), rnn.bias_hh_l0.detach()
    print('w_ih', tuple(w_ih.shape), ' w_hh', tuple(w_hh.shape))

    # 一句话、batch=1，默认布局是 (seq_len, batch, input_size)
    x = torch.randn(seq_len, 1, input_size)
    h0 = torch.zeros(1, 1, hidden_size)      # 全 0 起步：还没有任何记忆

    # -------------------------------------------------- 官方实现
    output, hn = rnn(x, h0)

    # -------------------------------------------------- 自己按公式循环
    h = h0[0]                                # (batch, hidden_size)
    manual_outputs = []
    for t in range(seq_len):
        h = rnn_cell(x[t], h, w_ih, b_ih, w_hh, b_hh)
        manual_outputs.append(h)             # 每一步的 h 就是这一步的 output
    manual_output = torch.stack(manual_outputs)

    print('官方 output 形状', output.shape)           # (5, 1, 4)
    print('手写 output 形状', manual_output.shape)    # (5, 1, 4)
    print('两边是否一致 :', torch.allclose(output, manual_output, atol=1e-6))
    print('hn 就是最后一步的 h :',
          torch.allclose(hn[0], manual_output[-1], atol=1e-6))

    # 三个必须记住的结论：
    #   1) output 是「每一步的 h 摞起来」，hn 是「最后一步的 h」，
    #      单层网络里 output[-1] 和 hn[0] 是同一个东西
    #   2) 所有时间步共用同一套 w_ih / w_hh —— 这就是「同一个神经元」的含义
    #   3) 参数量与句子长度无关：句子再长，参数还是这四个张量
    total = sum(p.numel() for p in rnn.parameters())
    print('参数量 :', total,
          '= 4*3 + 4*4 + 4 + 4 =', hidden_size * input_size
          + hidden_size * hidden_size + 2 * hidden_size)


if __name__ == '__main__':
    main()

"""nn.RNN 的输入输出形状：搞错维度顺序是这一讲最常见的报错来源。"""
import torch
import torch.nn as nn

# input_size：每个词向量多少维；hidden_size：隐藏状态多少维；num_layers：堆几层
rnn = nn.RNN(input_size=128, hidden_size=256, num_layers=1)

# 默认布局是 (seq_len, batch, input_size)——句子长度在最前面，不是 batch
inputs = torch.randn(5, 32, 128)      # 5 个词、32 句话、每词 128 维
h0 = torch.zeros(1, 32, 256)          # (num_layers, batch, hidden_size)，一般全 0 起步

output, hn = rnn(inputs, h0)

print('output 形状', output.shape)    # (5, 32, 256)：每个时间步都有一个输出
print('hn 形状    ', hn.shape)        # (1, 32, 256)：只有最后一个时间步的隐藏状态

# 如果你的数据是 (batch, seq_len, input_size)，两条路二选一：
#   1) 建层时写 nn.RNN(..., batch_first=True)
#   2) 送进去之前 x.transpose(0, 1)
rnn_bf = nn.RNN(128, 256, 1, batch_first=True)
out_bf, hn_bf = rnn_bf(torch.randn(32, 5, 128))
print('batch_first 时 output 形状', out_bf.shape)   # (32, 5, 256)

"""梯度消失为什么会发生：把 BPTT 的连乘用纯算术跑一遍。

反向传播沿时间轴往回走时，梯度要一段一段连乘：
    dL/dh_1 = dL/dh_T  ×  (dh_T/dh_{T-1}) × … × (dh_2/dh_1)
每一段的大小大致就是 |w_hh · tanh'|。tanh' 最大是 1，一般远小于 1，
所以这串乘积的数量级完全由「单段系数」的 T 次方决定。
本文件只用标准库，跑出来的数字就是 factor ** step，可以拿计算器核对。
"""


def chain(factor, steps):
    """连乘 steps 段，返回每一段之后的累计梯度倍率。"""
    value, trace = 1.0, []
    for _ in range(steps):
        value *= factor
        trace.append(value)
    return trace


def main():
    steps = 50

    print('时间步    系数 0.5        系数 0.9        系数 1.0        系数 1.1')
    print('-' * 68)
    t05, t09, t10, t11 = (chain(f, steps) for f in (0.5, 0.9, 1.0, 1.1))
    for t in (1, 5, 10, 20, 30, 50):
        print('%5d  %14.6e  %14.6e  %14.6e  %14.6e'
              % (t, t05[t - 1], t09[t - 1], t10[t - 1], t11[t - 1]))

    print()
    print('读法：')
    print('  系数 0.5 → 第 20 步梯度只剩 %.3e，20 步以前的词等于没参与训练'
          % t05[19])
    print('  系数 0.9 → 衰减慢一些，第 50 步也只剩 %.3e' % t09[49])
    print('  系数 1.0 → 唯一不衰减也不爆炸的临界点，现实中碰不到')
    print('  系数 1.1 → 第 50 步放大到 %.3e，这就是梯度爆炸' % t11[49])

    # ------------------------------------------------------------ 半衰期
    # 想知道「多少步之后梯度掉到千分之一」，解 factor**n = 0.001
    import math
    for f in (0.5, 0.8, 0.9, 0.95):
        n = math.log(0.001) / math.log(f)
        print('系数 %.2f：约 %.1f 步之后梯度掉到千分之一' % (f, n))

    # 三条可以直接带走的结论：
    #   1) 梯度消失不是 bug，是连乘的必然结果，网络越深/句子越长越明显
    #   2) 梯度爆炸能救（裁剪梯度），梯度消失救不了——信息本身没传过来
    #   3) LSTM / GRU 的门控就是给这条链加一条「系数接近 1」的旁路，
    #      让远处的梯度有机会不被连乘吃掉
    print()
    print('梯度爆炸可以用 torch.nn.utils.clip_grad_norm_(params, max_norm) 压住；')
    print('梯度消失压不住，只能换结构（LSTM / GRU / Transformer）。')


if __name__ == '__main__':
    main()

"""RNN / LSTM / GRU 三者的接口与参数量对照。

三个层的用法几乎一样，换名字就能换结构，区别只有两处：
    1) LSTM 的隐藏状态是 (h, c) 两个张量，RNN 和 GRU 只有 h
    2) 参数量：RNN 一份，GRU 三份，LSTM 四份 —— 门越多参数越多
"""
import torch
import torch.nn as nn


def param_count(layer):
    return sum(p.numel() for p in layer.parameters())


def expected(gates, input_size, hidden_size):
    """按公式推参数量：gates × (hidden×input + hidden×hidden + 2×hidden)。

    每个门都是一套完整的「输入权重 + 状态权重 + 两个偏置」。
    """
    return gates * (hidden_size * input_size
                    + hidden_size * hidden_size
                    + 2 * hidden_size)


def main():
    torch.manual_seed(0)

    input_size, hidden_size = 128, 256
    batch, seq_len = 8, 12
    x = torch.randn(batch, seq_len, input_size)      # batch_first 布局

    rnn = nn.RNN(input_size, hidden_size, batch_first=True)
    lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
    gru = nn.GRU(input_size, hidden_size, batch_first=True)

    # ---------------------------------------------------------- 返回值
    out_rnn, hn_rnn = rnn(x)
    print('RNN  output', out_rnn.shape, ' hn', hn_rnn.shape)

    out_gru, hn_gru = gru(x)
    print('GRU  output', out_gru.shape, ' hn', hn_gru.shape)

    # LSTM 多一个细胞状态 c：h 是「对外说的话」，c 是「自己记的账」
    out_lstm, (hn_lstm, cn_lstm) = lstm(x)
    print('LSTM output', out_lstm.shape, ' hn', hn_lstm.shape, ' cn', cn_lstm.shape)
    # 写成 out, hn = lstm(x) 也不报错，但 hn 会是一个元组，
    # 下一步拿它当张量用就会抛 AttributeError —— 换 LSTM 最常见的坑

    # ---------------------------------------------------------- 参数量
    print()
    print('%-6s %-10s %-12s %s' % ('层', '门数', '实际参数量', '公式推导'))
    for name, layer, gates in (('RNN', rnn, 1), ('GRU', gru, 3), ('LSTM', lstm, 4)):
        print('%-6s %-10d %-12d %d'
              % (name, gates, param_count(layer),
                 expected(gates, input_size, hidden_size)))
    # 三行数字完全对得上，说明「门控 = 复制几套同样的权重」这个理解是对的。
    # 代价也很直白：LSTM 的参数和计算量是同规模 RNN 的 4 倍。

    # ---------------------------------------------------------- 手动传状态
    print()
    h0 = torch.zeros(1, batch, hidden_size)
    c0 = torch.zeros(1, batch, hidden_size)
    out, (hn, cn) = lstm(x, (h0, c0))        # LSTM 要成对传，元组不能拆
    print('显式传 (h0, c0) 之后 output', out.shape)
    # 不传初始状态时 PyTorch 自动填一份全 0 的状态，效果一样；
    # 逐词生成时必须自己带着状态走，否则每一步都从零记忆开始。

    print()
    print('选型建议：短序列先用 GRU（参数比 LSTM 少 1/4，效果常常持平）；')
    print('需要更强的长期记忆再上 LSTM；纯 RNN 基本只用于教学和极短序列。')


if __name__ == '__main__':
    main()

"""构建词表：把一段文本变成「词 → 索引」的字典和一串索引。

这一步在任何 NLP 任务里都排第一位，做错了后面全错。
本文件不依赖 jieba，用一个极简的中文单字 + 英文单词切分器代替，
这样没装分词器也能跑；真实项目把 tokenize() 换成 jieba.lcut 即可。
"""


def tokenize(line):
    """极简切分：中文按单字切，连续的英文字母/数字算一个词。"""
    tokens, buf = [], ''
    for ch in line:
        if ch.isascii() and ch.isalnum():
            buf += ch                      # 英文数字先攒着，凑成一个完整单词
            continue
        if buf:
            tokens.append(buf)
            buf = ''
        if not ch.isspace():
            tokens.append(ch)              # 中文和标点一个字一个词
    if buf:
        tokens.append(buf)
    return tokens


def build_vocab(lines):
    """返回四件套：index_to_word / word_to_index / vocab_size / corpus_idx。"""
    index_to_word = []                     # 下标 i 对应第 i 个词，顺序即索引
    all_lines = []

    for line in lines:
        words = tokenize(line)
        all_lines.append(words)
        for word in words:
            # 用「不在就追加」保持首次出现的顺序，结果可复现；
            # 换成 set() 会丢顺序，每次跑索引都不一样，调试时对不上号
            if word not in index_to_word:
                index_to_word.append(word)

    # 换行符本身也要进词表，否则拼接语料时无法表示句子边界
    if '\n' not in index_to_word:
        index_to_word.append('\n')

    word_to_index = {w: i for i, w in enumerate(index_to_word)}

    corpus_idx = []                        # 整个语料摊平成一串索引
    for words in all_lines:
        corpus_idx.extend(word_to_index[w] for w in words)
        corpus_idx.append(word_to_index['\n'])

    return index_to_word, word_to_index, len(index_to_word), corpus_idx


if __name__ == '__main__':
    corpus = ['我爱你', '我爱北京']

    i2w, w2i, vocab_size, idx_seq = build_vocab(corpus)

    print('词表      :', i2w)
    print('词表大小  :', vocab_size)
    print('词到索引  :', w2i)
    print('语料索引串:', idx_seq)

    # 两行文本一共 3 + 4 = 7 个词 + 2 个换行 = 9 个索引，
    # 去重后是 我 爱 你 北 京 \n 共 6 个词，所以 vocab_size = 6。
    assert vocab_size == 6, '词表大小推导不符，先检查 tokenize'
    assert len(idx_seq) == 9, '语料索引串长度推导不符'

    # 反查：索引串还原回文本，验证映射是双向可逆的
    print('还原文本  :', ''.join(i2w[i] for i in idx_seq))

"""词嵌入最小例子：文字 -> 索引 -> 向量。"""
import torch
import torch.nn as nn
import jieba

text = '北京冬奥的进度条已经过半，不少外国运动员在完成自己的比赛后踏上归途。'

# 1. 分词：中文没有天然空格，先切成词
words = jieba.lcut(text)
print('分词结果:', words)

# 2. 去重得到词表：词表有几个词，嵌入矩阵就有几行
unique_words = list(set(words))
print('去重后词数:', len(unique_words))

# 3. 建词嵌入层：num_embeddings 是词表大小，embedding_dim 是每个词用几维表示
embed = nn.Embedding(num_embeddings=len(unique_words), embedding_dim=4)
print('嵌入矩阵形状:', embed.weight.shape)     # (词表大小, 4)

# 4. 查表：传进去的是索引（int64），拿回来的是那一行向量
for i, word in enumerate(unique_words[:5]):
    print('%4s' % word, embed(torch.tensor(i)).detach())

# 这些向量一开始是随机的，靠训练才会把语义学进去——
# 嵌入矩阵本身就是网络参数，会被反向传播更新

"""完整案例：用 RNN 生成歌词。

流程：建词表 -> 切成定长样本 -> 嵌入层 + RNN + 全连接 -> 训练 -> 逐词生成。
"""
import time
import jieba
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

DATA = 'data/jaychou_lyrics.txt'


# ------------------------------------------------------------------ 词表
def build_vocab(file_name=DATA):
    unique_words, all_words = [], []

    for line in open(file_name, 'r', encoding='utf-8'):
        words = jieba.lcut(line)
        all_words.append(words)
        for word in words:
            if word not in unique_words:
                unique_words.append(word)

    word_count = len(unique_words)
    word_to_index = {word: idx for idx, word in enumerate(unique_words)}

    # 把整个语料摊平成一串索引，行与行之间用空格这个词隔开
    corpus_idx = []
    for words in all_words:
        temp = [word_to_index[word] for word in words]
        temp.append(word_to_index[' '])
        corpus_idx.extend(temp)

    return unique_words, word_to_index, word_count, corpus_idx


# ------------------------------------------------------------------ 数据集
class LyricsDataset(torch.utils.data.Dataset):
    """从语料里切定长片段：输入是第 i 到 i+n 个词，目标是整体右移一位。"""

    def __init__(self, corpus_idx, num_chars):
        self.corpus_idx = corpus_idx
        self.num_chars = num_chars
        self.word_count = len(corpus_idx)
        self.number = self.word_count // self.num_chars

    def __len__(self):
        return self.number

    def __getitem__(self, idx):
        # 夹一下起点，避免尾部越界
        start = min(max(idx, 0), self.word_count - self.num_chars - 2)
        x = self.corpus_idx[start: start + self.num_chars]
        y = self.corpus_idx[start + 1: start + 1 + self.num_chars]   # 右移一位
        return torch.tensor(x), torch.tensor(y)


# ------------------------------------------------------------------ 模型
class TextGenerator(nn.Module):
    def __init__(self, word_count, embed_dim=128, hidden_dim=128):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.ebd = nn.Embedding(word_count, embed_dim)       # 词 -> 向量
        self.rnn = nn.RNN(embed_dim, hidden_dim, 1)          # 提取上下文
        self.out = nn.Linear(hidden_dim, word_count)         # 给词表里每个词打分

    def forward(self, inputs, hidden):
        # inputs: (batch, seq_len) -> embed: (batch, seq_len, embed_dim)
        embed = self.ebd(inputs)
        # nn.RNN 默认要 (seq_len, batch, embed_dim)，所以转置前两维
        output, hidden = self.rnn(embed.transpose(0, 1), hidden)
        # 把时间步和 batch 压到一起再过全连接：(seq_len*batch, hidden_dim)
        output = self.out(output.reshape(-1, output.shape[-1]))
        return output, hidden

    def init_hidden(self, bs=1):
        return torch.zeros(1, bs, self.hidden_dim)


# ------------------------------------------------------------------ 训练
def train(epochs=10, seq_len=32, lr=1e-3):
    _, _, word_count, corpus_idx = build_vocab()
    lyrics = LyricsDataset(corpus_idx, seq_len)

    model = TextGenerator(word_count)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)

    for epoch_idx in range(epochs):
        dataloader = DataLoader(lyrics, shuffle=True, batch_size=1)
        start, iter_num, total_loss = time.time(), 0, 0.0

        for x, y in dataloader:
            hidden = model.init_hidden(bs=1)
            output, hidden = model(x, hidden)

            # 目标值要和 output 的排列方式对齐：
            # y 是 (batch, seq_len) -> 转成 (seq_len, batch) -> 拉平成一维
            y = torch.transpose(y, 0, 1).contiguous().view(-1)
            loss = criterion(output, y)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            iter_num += 1
            total_loss += loss.item()

        print('epoch %3d  loss %.5f  time %.2fs'
              % (epoch_idx + 1, total_loss / iter_num, time.time() - start))
        torch.save(model.state_dict(), 'data/lyrics_model_%d.pth' % (epoch_idx + 1))


# ------------------------------------------------------------------ 生成
def predict(start_word, sentence_length, weight='data/lyrics_model_10.pth'):
    index_to_word, word_to_index, word_count, _ = build_vocab()

    model = TextGenerator(word_count)
    model.load_state_dict(torch.load(weight))
    model.eval()

    hidden = model.init_hidden()
    word_idx = word_to_index[start_word]      # 起始词必须在词表里，否则 KeyError
    generated = [word_idx]

    with torch.no_grad():
        for _ in range(sentence_length):
            # 每一步只送一个词，但 hidden 要一路带着走，记忆才不断
            output, hidden = model(torch.tensor([[word_idx]]), hidden)
            word_idx = torch.argmax(output).item()
            generated.append(word_idx)

    return ''.join(index_to_word[i] for i in generated)


if __name__ == '__main__':
    train()
    print(predict('分手', 50))

"""argmax：把一排分数变成「选哪个词」。"""
import torch

# 假设词表只有 5 个词，模型给出这一步的 5 个分数
logits = torch.tensor([1.2, -0.3, 3.4, 0.8, 2.9])
vocab = ['我', '爱', '你', '的', '歌']

prob = torch.softmax(logits, dim=0)
print('概率分布', prob)
print('概率之和', prob.sum().item())          # 1.0

idx = torch.argmax(logits)                    # 分数最大的下标
print('选中下标', idx.item(), '对应词', vocab[idx])

# 注意：argmax 每次都挑最高分，生成出来的句子会很单调、容易循环。
# 要多样性就改成按概率采样：
sampled = torch.multinomial(prob, num_samples=1)
print('按概率采样到', vocab[sampled.item()])

"""采样策略：argmax、温度采样、top-k 三种选词方式的差别。

生成任务里，模型每一步给出的是「词表上的一组分数」，
真正决定生成效果的，是你怎么把这组分数变成一个词。
这一步不参与训练，纯粹是推理时的选择，改它不用重训模型。
"""
import torch


VOCAB = ['我', '爱', '你', '的', '歌', '雨', '天']
# 假设某一步模型给出这 7 个分数（logits，没过 softmax 的原始分）
LOGITS = torch.tensor([1.2, -0.3, 3.4, 0.8, 2.9, 0.1, -1.5])


def softmax_with_temperature(logits, temperature):
    """温度缩放：先把 logits 除以 T，再做 softmax。

    T < 1：差距被放大，分布更尖，输出更保守
    T = 1：原样
    T > 1：差距被压缩，分布更平，输出更发散
    """
    return torch.softmax(logits / temperature, dim=-1)


def main():
    torch.manual_seed(0)

    print('原始分数 :', LOGITS.tolist())
    prob = torch.softmax(LOGITS, dim=-1)
    print('softmax  :', [round(p, 4) for p in prob.tolist()])
    print('概率之和 :', round(prob.sum().item(), 6))    # 恒为 1.0

    # ---------------------------------------------------------- argmax
    idx = torch.argmax(LOGITS).item()
    print('\nargmax 选中 :', VOCAB[idx], '（分数最高的 3.4）')
    # argmax 是确定性的：同样的输入永远给同样的词。
    # 后果是生成的句子容易原地打转 —— 「你你你你」这种循环就是这么来的。

    # ---------------------------------------------------------- 温度
    print('\n不同温度下的概率分布：')
    print('%-6s' % 'T', ''.join('%7s' % w for w in VOCAB))
    for t in (0.5, 1.0, 1.5, 3.0):
        p = softmax_with_temperature(LOGITS, t)
        print('%-6.1f' % t, ''.join('%7.3f' % v for v in p.tolist()))
    # 看最后一行：T=3 时最高分和最低分的概率已经拉不开，
    # 生成会开始胡言乱语；T 太小又退化成 argmax。常用区间是 0.7 ~ 1.0。

    # ---------------------------------------------------------- 按概率采样
    print('\n按 T=1.0 概率采样 10 次：')
    p = softmax_with_temperature(LOGITS, 1.0)
    picks = [VOCAB[torch.multinomial(p, 1).item()] for _ in range(10)]
    print(' '.join(picks))
    # multinomial 按概率抽签：高分词抽中的次数多，但低分词也有机会，
    # 这就是生成结果有变化的来源。注意它要的是概率不是 logits。

    # ---------------------------------------------------------- top-k
    print('\ntop-k 采样（k=3）：')
    k = 3
    topk_vals, topk_idx = torch.topk(LOGITS, k)
    print('候选词  :', [VOCAB[i] for i in topk_idx.tolist()])
    topk_prob = torch.softmax(topk_vals, dim=-1)        # 只在候选里重新归一化
    print('候选概率:', [round(v, 4) for v in topk_prob.tolist()])
    chosen = topk_idx[torch.multinomial(topk_prob, 1)].item()
    print('抽中    :', VOCAB[chosen])
    # top-k 的价值：先砍掉长尾里那些明显不通顺的词，再在靠谱的候选里随机，
    # 兼顾「不跑题」和「不重复」。工程上常和温度一起用。


if __name__ == '__main__':
    main()

"""序列模型骨架模板：词表、数据集、模型三件套。"""
import torch
import torch.nn as nn


class SeqModel(nn.Module):
    def __init__(self, vocab_size, embed_dim=128, hidden_dim=256,
                 num_layers=1, n_class=None):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers

        self.embedding = nn.Embedding(vocab_size, embed_dim)

        # TODO: 序列长、要记更久就把 nn.RNN 换成 nn.LSTM / nn.GRU，
        #       接口一致，LSTM 的隐藏状态是 (h, c) 两个张量
        self.rnn = nn.RNN(embed_dim, hidden_dim, num_layers, batch_first=True)

        # TODO: 逐词生成时输出维度 = 词表大小；整句分类时 = 类别数
        self.out = nn.Linear(hidden_dim, n_class or vocab_size)

    def forward(self, inputs, hidden=None):
        # inputs: (batch, seq_len) 的词索引
        x = self.embedding(inputs)               # (batch, seq_len, embed_dim)
        output, hidden = self.rnn(x, hidden)     # (batch, seq_len, hidden_dim)

        # TODO: 逐词预测保留所有时间步；整句分类只取最后一步 output[:, -1, :]
        logits = self.out(output)
        return logits, hidden

    def init_hidden(self, batch_size=1):
        return torch.zeros(self.num_layers, batch_size, self.hidden_dim)


if __name__ == '__main__':
    model = SeqModel(vocab_size=1000)
    x = torch.randint(0, 1000, (4, 10))          # 4 句话，每句 10 个词
    logits, hidden = model(x, model.init_hidden(batch_size=4))
    print(logits.shape, hidden.shape)            # (4, 10, 1000)  (1, 4, 256)

"""RNN 四类形状报错的原地复现：每一段都用 try/except 抓住，跑到底不中断。

序列模型的报错九成出在「谁在第几维」上，自己触发一遍比背十遍管用。
"""
import torch
import torch.nn as nn


def show(title):
    print('\n' + '=' * 12 + ' ' + title)


def main():
    torch.manual_seed(0)
    vocab_size, embed_dim, hidden_dim = 100, 16, 32
    batch, seq_len = 4, 7

    embedding = nn.Embedding(vocab_size, embed_dim)
    rnn_seq_first = nn.RNN(embed_dim, hidden_dim, num_layers=1)
    rnn_batch_first = nn.RNN(embed_dim, hidden_dim, num_layers=1, batch_first=True)

    tokens = torch.randint(0, vocab_size, (batch, seq_len))   # (batch, seq_len)
    x = embedding(tokens)                                     # (batch, seq_len, embed_dim)
    print('嵌入之后的形状', x.shape)

    # ---------------------------------------------------------------- 一
    show('一、batch 和 seq_len 放反')
    # nn.RNN 默认要 (seq_len, batch, input_size)，直接把 (batch, seq_len, dim) 丢进去
    # 不一定报错——它会把 batch 当成句长、把句长当成 batch，静默算错
    out_wrong, _ = rnn_seq_first(x)
    print('不报错，但 output 形状是', out_wrong.shape,
          '→ 它把', batch, '当成了句长')
    out_right, _ = rnn_seq_first(x.transpose(0, 1))
    print('转置之后才对   ', out_right.shape)
    out_bf, _ = rnn_batch_first(x)
    print('或者建层时写 batch_first=True', out_bf.shape)
    # 这是本讲最阴险的一类问题：没有任何异常，只是模型永远学不好

    # ---------------------------------------------------------------- 二
    show('二、h0 的第一维写成了 batch')
    try:
        bad_h0 = torch.zeros(batch, 1, hidden_dim)     # 写反了
        rnn_seq_first(x.transpose(0, 1), bad_h0)
    except RuntimeError as e:
        print('RuntimeError:', str(e).split('\n')[0])
    good_h0 = torch.zeros(1, batch, hidden_dim)        # (num_layers, batch, hidden)
    out, hn = rnn_seq_first(x.transpose(0, 1), good_h0)
    print('正确的 h0 形状', good_h0.shape, '→ hn', hn.shape)

    # ---------------------------------------------------------------- 三
    show('三、num_layers 改了但 h0 没跟着改')
    rnn2 = nn.RNN(embed_dim, hidden_dim, num_layers=2)
    try:
        rnn2(x.transpose(0, 1), torch.zeros(1, batch, hidden_dim))
    except RuntimeError as e:
        print('RuntimeError:', str(e).split('\n')[0])
    out2, hn2 = rnn2(x.transpose(0, 1), torch.zeros(2, batch, hidden_dim))
    print('两层时 output', out2.shape, ' hn', hn2.shape,
          '→ output 永远只有最后一层，hn 才是每一层都有')

    # ---------------------------------------------------------------- 四
    show('四、逐词预测时 CrossEntropyLoss 的形状没对齐')
    head = nn.Linear(hidden_dim, vocab_size)
    logits = head(out_bf)                       # (batch, seq_len, vocab_size)
    targets = torch.randint(0, vocab_size, (batch, seq_len))
    criterion = nn.CrossEntropyLoss()
    try:
        criterion(logits, targets)              # 三维预测 + 二维标签，直接不收
    except (RuntimeError, ValueError) as e:
        print('报错:', str(e).split('\n')[0])
    # 正确做法：把 batch 和时间步压平，变成「一共 batch*seq_len 条独立的多分类」
    loss = criterion(logits.reshape(-1, vocab_size), targets.reshape(-1))
    print('压平之后可以算:', logits.reshape(-1, vocab_size).shape,
          targets.reshape(-1).shape, '→ loss 是标量:', loss.dim() == 0)


if __name__ == '__main__':
    main()