inference/generate.py

"""
Text generation with top-k + nucleus (top-p) sampling.

SAMPLING STRATEGIES:
─────────────────────

GREEDY (temperature=0):
  Always picks the highest-probability token. Fast but repetitive — the model
  gets stuck in loops because the same high-probability token is chosen forever.

TEMPERATURE SCALING:
  Divides logits by `temperature` before softmax:
    - temperature < 1.0 → sharper distribution → more confident, focused
    - temperature > 1.0 → flatter distribution → more random, creative
    - temperature = 1.0 → unmodified model distribution

TOP-K FILTERING:
  Zeros out all but the K highest-probability tokens. Prevents sampling from
  the long tail of very unlikely tokens. Typical values: 40–100.

TOP-P (NUCLEUS) SAMPLING:
  Sort tokens by probability (descending). Keep only the smallest set whose
  cumulative probability sum reaches `top_p`. The cutoff adapts to the shape
  of the distribution — if the model is confident (one token has 80% prob),
  nucleus is tight. If the model is uncertain, more tokens are included.

COMBINING TOP-K AND TOP-P:
  Apply top-k first, then top-p on the survivors. This clips both the
  absolute count and the probability mass of the sampling pool.

KV-CACHE AT INFERENCE:
  During training, the full sequence is processed in one forward pass.
  During autoregressive generation, each step appends one token and re-runs
  the model over the entire grown sequence — O(T²) total compute.
  The KV-cache stores K, V tensors from all previous steps. Each new step
  only processes the single new token, then appends to the cache.
  Total compute drops to O(T) — a massive speedup for long generations.
  This implementation uses the cache from MiniTransformer.
"""

import torch
import torch.nn.functional as F

from tokenizer.bpe_tokenizer import BPETokenizer


def sample_logits(
    logits:      torch.Tensor,
    temperature: float = 1.0,
    top_k:       int   = 50,
    top_p:       float = 0.9,
) -> int:
    """
    Sample a single token ID from a logits tensor.

    Args:
        logits:      Raw logits, shape [1, vocab_size].
        temperature: Scaling factor. 0 → greedy argmax.
        top_k:       Keep only top-k tokens (None or 0 to disable).
        top_p:       Nucleus probability threshold (None or 1.0 to disable).

    Returns:
        A single integer token ID.
    """
    # Greedy shortcut
    if temperature <= 0:
        return int(logits.argmax(dim=-1).item())

    logits = logits / temperature

    # ── Top-k filtering ───────────────────────────────────────────────────────
    if top_k and top_k > 0:
        k = min(top_k, logits.size(-1))
        v, _ = torch.topk(logits, k)
        # v[:, -1:] is the k-th highest value per row → shape [B, 1] for broadcast
        logits = logits.masked_fill(logits < v[:, -1:], float("-inf"))

    # ── Top-p (nucleus) filtering ─────────────────────────────────────────────
    if top_p and top_p < 1.0:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

        # Remove tokens once the cumulative probability exceeds top_p.
        # Shift right by 1 so we always keep at least the top-1 token.
        remove_mask = (cumulative_probs - F.softmax(sorted_logits, dim=-1)) > top_p
        sorted_logits = sorted_logits.masked_fill(remove_mask, float("-inf"))

        # Scatter back to original token order
        logits = torch.zeros_like(logits).scatter_(1, sorted_indices, sorted_logits)

    probs = F.softmax(logits, dim=-1)
    return int(torch.multinomial(probs, num_samples=1).item())


def generate(
    model,
    prompt:         str,
    tokenizer:      "BPETokenizer | None" = None,
    max_new_tokens: int            = 50,
    temperature:    float          = 0.8,
    top_k:          int            = 50,
    top_p:          float          = 0.9,
    device:         str            = "cuda" if torch.cuda.is_available() else "cpu",
    use_kv_cache:   bool           = True,
    repetition_penalty: float      = 1.15,
    penalty_window: int            = 64,
) -> str:
    """
    Autoregressively generate text from a prompt using the trained model.

    Args:
        model:          A trained MiniTransformer.
        prompt:         Input text string to condition generation on.
        tokenizer:      BPETokenizer instance. Created if not provided.
                        Pass one in to avoid reloading from disk each call.
        max_new_tokens: Maximum number of tokens to generate.
        temperature:    Sampling temperature (0 = greedy).
        top_k:          Top-k filtering. Disable with None or 0.
        top_p:          Nucleus sampling. Disable with None or 1.0.
        device:         Device for inference ("cuda" or "cpu").
        use_kv_cache:   Use KV-cache for faster autoregressive inference.
        repetition_penalty: Multiplier to penalize repeated tokens (> 1.0 discourages loops).
        penalty_window: How many recent tokens to consider for the repetition penalty.

    Returns:
        The decoded generated string (prompt + continuation), cleaned.
    """
    if tokenizer is None:
        tokenizer = BPETokenizer("dataProcessed/tokenizer.json")

    context_len = model.config.context_len

    model.eval()
    model.to(device)

    # Encode prompt to token IDs
    tokens = torch.tensor(
        tokenizer.encode(prompt), dtype=torch.long, device=device
    ).unsqueeze(0)                               # [1, prompt_len]

    past_kvs: list | None = None
    temperature = max(temperature, 1e-5)

    with torch.no_grad():
        for _ in range(max_new_tokens):
            if use_kv_cache and past_kvs is not None:
                # KV-cache mode: only feed the single newest token
                context = tokens[:, -1:]
            else:
                # Full-context mode (first step or no cache): truncate to window
                context = tokens[:, -context_len:]

            logits, past_kvs_new = model(context, past_kvs if use_kv_cache else None)

            if use_kv_cache:
                past_kvs = past_kvs_new

            next_logits = logits[:, -1, :].clone()  # [1, vocab_size]

            # Apply repetition penalty to recent tokens within the window
            if repetition_penalty > 1.0 and tokens.shape[1] > 0:
                recent_tokens = tokens[0, -penalty_window:].tolist()
                for token_id in set(recent_tokens):
                    if next_logits[0, token_id] > 0:
                        next_logits[0, token_id] /= repetition_penalty
                    else:
                        next_logits[0, token_id] *= repetition_penalty

            next_id = sample_logits(next_logits, temperature, top_k, top_p)

            # Stop generation when the model produces the EOS token
            if next_id == tokenizer.eos_id:
                break

            next_tensor = torch.tensor([[next_id]], dtype=torch.long, device=device)
            tokens = torch.cat([tokens, next_tensor], dim=1)

    decoded = tokenizer.decode(tokens[0].tolist())
    return tokenizer.clean(decoded)