File size: 7,407 Bytes
9e30969 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | """
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)
|