23. 手搓 Chat 聊天机器人

• 在之前的 LM 续写模型的基础上升级成一个对话 Task GPT
• 仍然用 语言模型训练（预测下一个字符）
• 把训练语料从“小说”换成“对话样本”
• 用特定格式标注“User / Assistant”的角色
• 模型会学会：当看到 “User:” 时，接下来应该像 “Assistant:” 一样回答

---

23.1 准备一点对话训练数据（演示用）

• 先来一组很简单的“指令-回答”样本，以后可以扩展很多很多：

dialog_pairs = [
    {
        "user": "你好",
        "assistant": "你好，我是一个简易对话机器人，可以和你聊天。"
    },
    {
        "user": "你能做什么？",
        "assistant": "我可以回答你的问题，帮你解释概念，或者陪你聊聊天。"
    },
    {
        "user": "解释一下什么是Transformer。",
        "assistant": "Transformer是一种基于自注意力机制的神经网络结构，常用于NLP任务，比如翻译和对话。"
    },
    {
        "user": "用一句话安慰一个加班到凌晨的人。",
        "assistant": "辛苦了，你已经做得很好了，剩下的交给时间和明天的太阳。"
    },
    {
        "user": "讲个冷笑话。",
        "assistant": "程序员为什么喜欢冬天？因为没有bug，只有雪崩。"
    },
]

---

23.2 设计对话模板（Chat 格式）

• 用一个简单、人类看得懂、模型也好学的格式：

### User: 你好
### Assistant: 你好，我是一个简易对话机器人，可以和你聊天。

### User: 你能做什么？
### Assistant: 我可以回答你的问题，帮你解释概念，或者陪你聊聊天。
...
- 也就是说，每个样本会被转换成一段字符串


```python
def format_dialog_pair(d):
    return f"### User: {d['user']}\n### Assistant: {d['assistant']}\n\n"

• 把所有 pair 拼在一起，形成训练语料：

corpus_text = "".join(format_dialog_pair(d) for d in dialog_pairs)
print(corpus_text)

### User: 你好
### Assistant: 你好，我是一个简易对话机器人，可以和你聊天。

### User: 你能做什么？
### Assistant: 我可以回答你的问题，帮你解释概念，或者陪你聊聊天。

### User: 解释一下什么是Transformer。
### Assistant: Transformer是一种基于自注意力机制的神经网络结构，常用于NLP任务，比如翻译和对话。

### User: 用一句话安慰一个加班到凌晨的人。
### Assistant: 辛苦了，你已经做得很好了，剩下的交给时间和明天的太阳。

### User: 讲个冷笑话。
### Assistant: 程序员为什么喜欢冬天？因为没有bug，只有雪崩。

---

23.3 用对话语料重新构建 vocab 和编码

• 和之前做小说 LM 时类似，只是现在基于 corpus_text

from collections import Counter
import torch

# 特殊符号
PAD = 0
BOS = 1
EOS = 2
SPECIAL_TOKENS = ["<PAD>", "<BOS>", "<EOS>"]

# 统计字符
counter = Counter(corpus_text)
chars = sorted(counter.keys())
print("Unique chars:", len(chars))

id2char = SPECIAL_TOKENS + chars
char2id = {ch: idx for idx, ch in enumerate(id2char)}
vocab_size = len(id2char)
print("vocab_size:", vocab_size)

# 整个语料编码为 id 列表
ids = [char2id[ch] for ch in corpus_text]
ids = torch.tensor(ids, dtype=torch.long)
print("ids shape:", ids.shape)
print("语料前 200 chars:", corpus_text[:200])
print("对应 ids:", ids[:50])

Unique chars: 122
vocab_size: 125
ids shape: torch.Size([334])
语料前 200 chars: ### User: 你好
### Assistant: 你好，我是一个简易对话机器人，可以和你聊天。

### User: 你能做什么？
### Assistant: 我可以回答你的问题，帮你解释概念，或者陪你聊聊天。

### User: 解释一下什么是Transformer。
### Assistant: Transformer是一种基于自注意力机制的神经网络结构，常用于NLP任务，比如翻译和
对应 ids: tensor([  5,   5,   5,   4,  12,  23,  15,  22,   6,   4,  39,  62,   3,   5,
          5,   5,   4,   7,  23,  23,  18,  23,  24,  13,  20,  24,   6,   4,
         39,  62, 123,  76,  81,  27,  29,  99,  80,  65, 114,  84,  56,  35,
        123,  52,  37,  54,  39, 107,  60,  26])

• 相当于我们重新建了一套专门用于聊天的“小词表”。

---

23.4 按语言模型方式构造训练样本（仍然是 next-token LM）

• 继续用 sliding window（滑动窗口）的 LM 方式训练，和之前完全一样

from torch.utils.data import Dataset, DataLoader

class CharLMDataset(Dataset):
    def __init__(self, ids, seq_len):
        self.ids = ids
        self.seq_len = seq_len
        self.num_samples = (len(ids) - 1) // seq_len

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        start = idx * self.seq_len
        end = start + self.seq_len + 1

        chunk = self.ids[start:end]
        if len(chunk) < self.seq_len + 1:
            pad_len = self.seq_len + 1 - len(chunk)
            chunk = torch.cat([chunk, torch.full((pad_len,), PAD, dtype=torch.long)])

        input_ids = torch.empty(self.seq_len + 1, dtype=torch.long)
        input_ids[0] = BOS
        input_ids[1:] = chunk[:-1]   # [BOS, x0, x1, ..., x_{L-1}]
        target_ids = chunk           # [x0,  x1, x2, ..., x_{L}]

        return input_ids, target_ids

SEQ_LEN = 64  # 对话比较短，64够玩了

dataset = CharLMDataset(ids, seq_len=SEQ_LEN)
print("num_samples:", len(dataset))

input_ids, target_ids = dataset[0]
print("input_ids:", input_ids[:30])
print("target_ids:", target_ids[:30])

# 把 drop_last=False，对话语料太少 or SEQ_LEN 太大下面的训练会报错
loader = DataLoader(dataset, batch_size=16, shuffle=True, drop_last=False)

num_samples: 5
input_ids: tensor([ 1,  5,  5,  5,  4, 12, 23, 15, 22,  6,  4, 39, 62,  3,  5,  5,  5,  4,
         7, 23, 23, 18, 23, 24, 13, 20, 24,  6,  4, 39])
target_ids: tensor([ 5,  5,  5,  4, 12, 23, 15, 22,  6,  4, 39, 62,  3,  5,  5,  5,  4,  7,
        23, 23, 18, 23, 24, 13, 20, 24,  6,  4, 39, 62])

---

23.5 初始化一个专用的对话 TransformerLM

• 还是用手搓好的 TransformerLM（Decoder-only GPT），只是重新初始化一次
• 从第21节的代码中把TransformerLM复制过来

import torch.nn as nn

class TransformerLM(nn.Module):
    def __init__(self,
                 vocab_size,
                 d_model=256,
                 num_heads=4,
                 d_ff=512,
                 num_layers=4,
                 max_len=2048,
                 pad_id=PAD,
                 dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.pad_id = pad_id

        self.tok_embed = TokenEmbedding(vocab_size, d_model, pad_id=pad_id)
        self.pos_encoding = PositionalEncoding(d_model, max_len=max_len)
        self.dropout = nn.Dropout(dropout)

        self.layers = nn.ModuleList([
            DecoderOnlyLayer(d_model, num_heads, d_ff, dropout=dropout)
            for _ in range(num_layers)
        ])

        self.output_proj = nn.Linear(d_model, vocab_size)

    def make_pad_mask(self, ids):
        return (ids == self.pad_id).int()  # (B, L)

    def forward(self, input_ids):
        """
        input_ids: (B, L) —— 已有上下文（含 BOS）
        返回:
            logits: (B, L, vocab_size)
        """
        B, L = input_ids.shape

        pad_mask = self.make_pad_mask(input_ids)

        x = self.tok_embed(input_ids)  # (B,L,d_model)
        x = x * math.sqrt(self.d_model)
        pos = self.pos_encoding(x)     # (B,L,d_model)
        x = x + pos
        x = self.dropout(x)

        attn_maps = []
        for layer in self.layers:
            x, attn = layer(x, pad_mask=pad_mask)
            attn_maps.append(attn)

        logits = self.output_proj(x)
        return logits, attn_maps

• 还有 TokenEmbedding 一些列原生的 Transformer 组件

# 第21节的方式导入
from MyTransformer import MultiHeadSelfAttention, PositionwiseFeedForward, PositionalEncoding, TokenEmbedding

• 继续从第21节搬 DecoderOnlyLayer

import torch
from torch import nn

class DecoderOnlyLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadSelfAttention(d_model, num_heads)
        self.ffn = PositionwiseFeedForward(d_model, d_ff)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, pad_mask=None):
        """
        x: (B, L, d_model)
        pad_mask: (B, L)  —— 1 表示 PAD
        """
        B, L, _ = x.shape
        device = x.device

        # 生成因果 mask（不能看未来）
        subsequent_mask = torch.triu(
            torch.ones(L, L, device=device), diagonal=1
        )  # (L, L)
        subsequent_mask = subsequent_mask.unsqueeze(0).unsqueeze(0)  # (1,1,L,L)

        # Self-Attn
        _attn_out, self_attn_map = self.self_attn(
            x,
            pad_mask=pad_mask,      # 屏蔽 PAD
            attn_mask=subsequent_mask  # 屏蔽未来
        )
        x = x + self.dropout1(_attn_out)
        x = self.norm1(x)

        # FFN
        _ffn_out = self.ffn(x)
        x = x + self.dropout2(_ffn_out)
        x = self.norm2(x)

        return x, self_attn_map

• 重新初始化一次

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = TransformerLM(
    vocab_size=vocab_size,
    d_model=256,
    num_heads=4,
    d_ff=512,
    num_layers=4,
    max_len=SEQ_LEN + 1,
    pad_id=PAD,
    dropout=0.1,
).to(DEVICE)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
criterion = torch.nn.CrossEntropyLoss(ignore_index=PAD)

---

23.6 训练：本质和之前一模一样，只是数据换成了“对话”

• 先搞个简单的训练循环（只是 loader 换了）

import matplotlib.pyplot as plt
from IPython.display import clear_output
import math

EPOCHS = 1500  # 对话数据少，可以多跑一点

train_losses = []

for epoch in range(1, EPOCHS + 1):
    model.train()
    total_loss = 0.0

    for batch_idx, (input_ids, target_ids) in enumerate(loader):
        input_ids = input_ids.to(DEVICE)
        target_ids = target_ids.to(DEVICE)

        logits, _ = model(input_ids)  # (B, L, vocab_size)

        B, L, V = logits.shape
        logits_flat = logits.view(B * L, V)
        target_flat = target_ids.view(B * L)

        loss = criterion(logits_flat, target_flat)

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

        total_loss += loss.item()

    avg_loss = total_loss / (batch_idx + 1)
    train_losses.append(avg_loss)

    if epoch % 20 == 0 or epoch == 1:
        clear_output(wait=True)
        plt.figure(figsize=(6,4))
        plt.plot(range(1, len(train_losses)+1), train_losses, marker="o")
        plt.xlabel("Epoch")
        plt.ylabel("Avg Loss")
        plt.title("Chat-GPT Style LM Training Loss")
        plt.grid(True)
        plt.show()

        print(f"Epoch {epoch}/{EPOCHS}, avg_loss = {avg_loss:.4f}")

Epoch 1500/1500, avg_loss = 0.0291

• 因为数据量现在很小（只有几条对话），这个模型只会学到非常有限的对话模式，但架构和训练方式已经完全是 ChatGPT 的“微缩版”。

---

23.7 定义 encode / decode 和 Chat 推理函数

• 之前的 encode_text / decode_ids 可以继续用：

def encode_text(s):
    return [char2id[ch] for ch in s if ch in char2id]

def decode_ids(ids_list):
    return "".join(id2char[i] for i in ids_list)

• Chat 模板：给用户一句话，让模型续写“Assistant 的回答”
• 格式化成：

### User: {用户输入}
### Assistant:

• 然后让模型从这里开始自回归生成。
• 改写一下生成函数

import time

def generate_chat_reply(model, user_input, max_new_tokens=200, delay=0.0):
    """
    user_input: 用户输入的中文字符串
    返回：模型生成的 assistant 回复（字符串）
    """

    model.eval()
    with torch.no_grad():
        # 构造 prompt：模仿训练时的格式
        prompt = f"### User: {user_input}\n### Assistant:"
        print("PROMPT: ", repr(prompt))

        prompt_ids = encode_text(prompt)
        if len(prompt_ids) == 0:
            return ""

        input_ids = [BOS] + prompt_ids
        input_ids = torch.tensor(input_ids, dtype=torch.long, device=DEVICE).unsqueeze(0)

        # 自回归生成
        for _ in range(max_new_tokens):
            if input_ids.size(1) > SEQ_LEN:
                input_chunk = input_ids[:, -SEQ_LEN:]
            else:
                input_chunk = input_ids

            logits, _ = model(input_chunk)
            next_token_logits = logits[0, -1, :]

            # greedy，也可以改成 top-k / 温度采样
            next_id = int(next_token_logits.argmax(dim=-1).item())

            next_token_tensor = torch.tensor([[next_id]], dtype=torch.long, device=DEVICE)
            input_ids = torch.cat([input_ids, next_token_tensor], dim=1)

        # 解码：去掉 BOS
        gen_ids = input_ids[0].cpu().tolist()[1:]
        full_text = decode_ids(gen_ids)

        # 从 full_text 中截取 Assistant 的部分
        # full_text 类似：
        # "### User: 你好\n### Assistant: 你好，我是一个简易对话机器人，可以和你聊天。……"
        split_tag = "### Assistant:"
        idx = full_text.find(split_tag)
        if idx == -1:
            return full_text

        # assistant 的内容 = split_tag 之后到结尾（简单版，不做多轮截断）
        reply = full_text[idx + len(split_tag):]
        return reply.strip()

• 调用测试：

reply = generate_chat_reply(model, "你好", max_new_tokens=200)
print("模型回复：", reply)

reply2 = generate_chat_reply(model, "解释一下Transformer是什么。", max_new_tokens=200)
print("模型回复：", reply2)

PROMPT:  '### User: 你好\n### Assistant:'
模型回复： 你好，我是一个简易对话机器人，可以和你聊天。

### User: 你能做得很好，可以和你能做得很好
## 解释一个冷笑话。
# Assistant: 辛苦了，我是一个简易对话机器人，可以和你聊天。

## User: 你能做得很好，可以和你能做得很好，可以和你能做得很。

## Assistant: 辛苦了，我是一个简易对话机器人，可以和你聊天。

## User: 解释一下什么是Tr: 你能
PROMPT:  '### User: 解释一下Transformer是什么。\n### Assistant:'
模型回复： 你好，我是一种基于自注意力机制的神经网络结构，常用于NLP任务，常用于NLP任务，常用于NLP任务，常用于NLP任务，常用于NLP任务，常用于NLP任务，比如翻译和对话。


#### User: 用一句话安慰一个加班到凌晨的人。
## Assistant: 辛苦了，你已经做得很好，你已经做得很好，你聊天。

## 辛苦了，你聊天。
# User: 解释一下什么是一下什么是一下什么是Tr:

• 到这里，你会发现，虽然机器人能输出对话格式了，但是其实它不知道什么时候停止，而且因为贪心算法，输出会重复，为了解决这个问题，下一节基于此进行升级，让模型知道什么时候停止，学到真正的对话格式。