File size: 2,190 Bytes
9df3dc4 0cf8154 9df3dc4 0cf8154 | 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 | """Autoregressive sampling, with a KV cache while the text fits the context."""
from __future__ import annotations
from collections.abc import Callable
import torch
from zenith.model.model import ZenithLM
@torch.no_grad()
def generate(
model: ZenithLM,
idx: torch.Tensor,
max_new_tokens: int,
temperature: float = 1.0,
top_k: int | None = None,
eos_id: int | None = None,
generator: torch.Generator | None = None,
logits_hook: Callable[[list[int], torch.Tensor], torch.Tensor] | None = None,
stop_fn: Callable[[list[int]], bool] | None = None,
) -> torch.Tensor:
"""Extend `idx` (shape (1, T)) by up to `max_new_tokens`; temperature 0 is greedy.
`logits_hook(ids, logits)` may rewrite the next-token logits (e.g. to mask tokens);
`stop_fn(ids)` ends generation after the token for which it returns True.
"""
# The cache only works while nothing slides out of the context window.
use_cache = idx.shape[1] + max_new_tokens <= model.cfg.context_len
cache = model.new_cache() if use_cache else None
for step in range(max_new_tokens):
if not use_cache:
logits, _ = model(idx[:, -model.cfg.context_len :])
elif step == 0:
logits, _ = model(idx, cache=cache)
else:
logits, _ = model(idx[:, -1:], cache=cache, start_pos=idx.shape[1] - 1)
logits = logits[:, -1, :].float()
if logits_hook is not None:
logits = logits_hook(idx[0].tolist(), logits)
if temperature == 0:
next_id = logits.argmax(dim=-1, keepdim=True)
else:
logits = logits / temperature
if top_k is not None:
kth = torch.topk(logits, min(top_k, logits.size(-1))).values[:, -1:]
logits = logits.masked_fill(logits < kth, float("-inf"))
probs = logits.softmax(dim=-1)
next_id = torch.multinomial(probs.cpu(), 1, generator=generator).to(idx.device)
idx = torch.cat([idx, next_id], dim=1)
if eos_id is not None and next_id.item() == eos_id:
break
if stop_fn is not None and stop_fn(idx[0].tolist()):
break
return idx
|