Autoregressive Generation › Complete Generation Implementation

from dataclasses import dataclass


@dataclass
class GenerationConfig:
    """Configuration for text generation."""

    max_new_tokens: int = 50
    temperature: float = 1.0
    top_k: int | None = None
    top_p: float | None = None
    repetition_penalty: float = 1.0
    stop_sequences: list[str] | None = None


def generate_with_config(
    model,
    tokenizer,
    prompt: str,
    config: GenerationConfig,
) -> str:
    """
    Complete generation function with configurable sampling strategies.

    Supports temperature scaling, top-k sampling, nucleus sampling,
    repetition penalty, and stop sequences.
    """
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    generated_ids = input_ids.clone()

    for _ in range(config.max_new_tokens):
        with torch.no_grad():
            outputs = model(generated_ids)
            logits = outputs.logits[:, -1, :]

            # Apply repetition penalty
            if config.repetition_penalty != 1.0:
                for token_id in generated_ids[0].unique():
                    logits[0, token_id] /= config.repetition_penalty

            # Apply temperature
            if config.temperature != 1.0:
                logits = logits / config.temperature

            # Apply top-k filtering
            if config.top_k is not None:
                top_k_logits, top_k_indices = torch.topk(logits, config.top_k)
                logits = torch.full_like(logits, float("-inf"))
                logits.scatter_(1, top_k_indices, top_k_logits)

            # Apply nucleus (top-p) filtering
            if config.top_p is not None:
                sorted_logits, sorted_indices = torch.sort(
                    logits, descending=True
                )
                cumulative_probs = torch.cumsum(
                    F.softmax(sorted_logits, dim=-1), dim=-1
                )
                sorted_indices_to_remove = cumulative_probs > config.top_p
                sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[
                    :, :-1
                ].clone()
                sorted_indices_to_remove[:, 0] = False
                indices_to_remove = sorted_indices_to_remove.scatter(
                    1, sorted_indices, sorted_indices_to_remove
                )
                logits[indices_to_remove] = float("-inf")

            # Greedy selection of highest-probability token
            next_token = torch.argmax(logits, dim=-1, keepdim=True)

        generated_ids = torch.cat([generated_ids, next_token], dim=1)

        # Check EOS
        if next_token.item() == tokenizer.eos_token_id:
            break

        # Check stop sequences
        if config.stop_sequences:
            current_text = tokenizer.decode(
                generated_ids[0], skip_special_tokens=True
            )
            for stop_seq in config.stop_sequences:
                if stop_seq in current_text[len(prompt) :]:
                    idx = current_text.find(stop_seq, len(prompt))
                    return current_text[:idx]

    return tokenizer.decode(generated_ids[0], skip_special_tokens=True)