Stopping Criteria - From Neuron to GPT

In the previous chapter, we learned to control the quality and diversity of sampling at inference. Now consider another question: when should the model stop generating?

By default, the model keeps generating until the maximum length. But often we want the model to stop at the right moment:

stop after one sentence;
stop after one paragraph;
stop when a specific token or sequence appears.

Stopping criteria¶

Stopping criteria are the rules that control when generation terminates. Common stopping criteria include:

Maximum length: stop after generating a fixed number of tokens. This is also the most basic default stopping criterion;
Stop tokens: stop after a specific token (for example a period, a newline, or EOF) appears;
Stop sequences: stop after a specific sequence (for example a blank line) appears.

How it works¶

After each token is generated, check whether any stopping criterion is met. If one is met, terminate early and do not generate any further.

import numpy as np

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

np.random.seed(42)

GPT stopping criteria trainer¶

class GPTTrainerWithStopping(GPTTrainer):

    def generate(self, dataset, prompt, steps=256, stop_sequences=None):
        # Inference mode
        self.model.eval()
        tokens = dataset.encode(prompt)
        prompt_len = len(tokens)
        # Stop sequences
        stops = [dataset.encode(s) for s in (stop_sequences or [])]

        with Tensor.no_grad():
            # Maximum length
            for _ in range(steps):
                # Inference
                feature = Tensor([tokens[-dataset.context_size:]])
                logits = self.model(feature)

                # Sampling
                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:]
                # Match tokens
                hit = next((s for s in stops if self._ends_with(generated, s)), None)
                # Early termination
                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

Hyperparameters¶

Model parameters¶

The model parameters must stay consistent with the pre-training stage.

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

Modeling¶

Loading the model¶

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)

Maximum length¶

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

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

Stop tokens¶

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

ROMEO:
His many choses!

Stop sequences¶

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?

In this chapter, we implemented stopping criteria: by a stop token (such as a period) or a stop sequence (such as a blank line), the model stops generating at the right moment.

Exercises¶

Try adding a maximum repetition stopping criterion: if the same token appears more than 3 times in a row (for example aaaa), stop generating.