| """Autoregressive sampling (no KV cache).""" |
|
|
| 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. |
| """ |
| for _ in range(max_new_tokens): |
| logits, _ = model(idx[:, -model.cfg.context_len :]) |
| 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 |
|
|
