停止条件 - 从零开始大模型

上一章，我们学会了控制模型推理采样的质量和多样性。现在思考另一个问题：模型什么时候停止生成？

默认情况下，模型会一直生成直到最大长度。但很多时候，我们希望模型在合适的时机停下来：

生成一句话就停；
生成一个段落就停；
遇到特定的词元或序列就停。

停止条件¶

停止条件（Stopping Criteria）是控制终止生成的规则。常见的停止条件有：

最大长度：生成固定数量的词元后停止。这也是最基本的默认停止条件；
停止词元：遇到特定词元（比如：句号、换行符、EOF）后停止；
停止序列：遇到特定序列（比如：空行）后结束。

实现原理¶

每生成一个词元，检查是否满足某个停止条件。满足条件就提前终止，不再继续生成。

import numpy as np

from src.core import *
from src.gpt import *

np.random.seed(42)

GPT 停止条件训练器¶

GPT 训练器的 generate 函数支持最大长度停止条件。

GPT 停止条件训练器 GPTTrainerWithStopping 继承自 GPT 训练器，扩展了 generate 函数：

停止序列：新增参数 stop_sequences 接受一个字符串（停止序列）列表；
匹配词元：每生成一个词元，检查生成的后缀是否匹配任何一个停止序列；
提前终止：匹配成功就跳出循环，结束生成。

class GPTTrainerWithStopping(GPTTrainer):

    def generate(self, dataset, prompt, steps=256, stop_sequences=None):
        # 推理模式
        self.model.eval()
        tokens = dataset.encode(prompt)
        prompt_len = len(tokens)
        # 停止序列
        stops = [dataset.encode(s) for s in (stop_sequences or [])]

        with Tensor.no_grad():
            # 最大长度
            for _ in range(steps):
                # 推理
                feature = Tensor([tokens[-dataset.context_size:]])
                logits = self.model(feature)

                # 采样
                last_logits = logits.data[0, -1]
                exp = np.exp(last_logits - np.max(last_logits))
                probs = exp / np.sum(exp)
                token = np.random.choice(len(probs), p=probs)
                tokens.append(token)

                generated = tokens[prompt_len:]
                # 匹配词元
                hit = next((s for s in stops if self._ends_with(generated, s)), None)
                # 提前终止
                if hit is not None:
                    break

        return dataset.decode(tokens)

    @staticmethod
    def _ends_with(tokens, sequence):
        return len(tokens) >= len(sequence) and tokens[-len(sequence):] == sequence

超参数¶

模型参数¶

模型参数需要保持和预训练阶段一致。

BATCH_SIZE = 16
CONTEXT_SIZE = 256
EMBEDDING_SIZE = 384
HEADS = 6
BLOCKS = 6

建模¶

加载模型¶

DATA_FILE = "../../../../data/tinyshakespeare.txt"
MODEL_FILE = "../../../../data/tinyshakespeare-gpt.npz"

dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = GPT(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, HEADS, BLOCKS)
loss_fn = GPTLoss()
optimizer = AdamWOptimizer(model.parameters)
trainer = GPTTrainerWithStopping(model, loss_fn, optimizer)
trainer.load(MODEL_FILE)

最大长度¶

print(trainer.generate(dataset, prompt="ROMEO:", steps=64))

ROMEO:
Cut on Duke of York. Whence news
but thinkful should he divines

停止词元¶

print(trainer.generate(dataset, prompt="ROMEO:", stop_sequences=[".", "?", "!"]))

ROMEO:
His many choses!

停止序列¶

print(trainer.generate(dataset, prompt="ROMEO:", stop_sequences=["\n\n"]))

ROMEO:
Ay, see your children will be frails
Tybalt's fear; for himself, here this duck?

本章，我们实现了停止条件：按停止词元（如句号）或者停止序列（如空行）让模型在合适的时机停止生成。

课后练习¶

尝试添加一个最大重复次数的停止条件：如果同一个词元连续出现超过 3 次（比如：aaaa），就停止生成。