Decoder Architecture › The Generation Loop

def generate(model, prompt_tokens, max_new_tokens=20, temperature=1.0):
    """
    Generate tokens autoregressively.

    Args:
        model: GPTDecoder instance
        prompt_tokens: Initial token indices
        max_new_tokens: Maximum tokens to generate
        temperature: Sampling temperature (higher = more random)

    Returns:
        Complete sequence including prompt and generated tokens
    """
    tokens = list(prompt_tokens)

    for _ in range(max_new_tokens):
        # Get current sequence length
        current_len = len(tokens)
        if current_len >= model.max_seq_len:
            break

        # Forward pass
        token_array = np.array(tokens)
        logits = model(token_array)

        # Get logits for next token prediction (last position)
        next_logits = logits[-1]  # (vocab_size,)

        # Apply temperature
        if temperature != 1.0:
            next_logits = next_logits / temperature

        # Convert to probabilities
        probs = softmax(next_logits)

        # Sample next token
        next_token = np.random.choice(len(probs), p=probs)
        tokens.append(next_token)

        # Simple stop condition: token 0 is end-of-sequence
        if next_token == 0:
            break

    return np.array(tokens)


# Generate from a prompt
prompt = np.array([42, 100, 256])
generated = generate(model, prompt, max_new_tokens=10, temperature=0.8)