""" 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)