| """Host-side next-token selection for the chain-of-causation decode. |
| |
| Mirrors the reference: discrete trajectory-token logits and the text EOS ids are masked to -inf, |
| then either greedy argmax (parity tests) or temperature + top-p sampling (default top_p 0.98, T 0.6). |
| Logits arrive as float32 [V] (or [1, V]) from the device. |
| """ |
| from __future__ import annotations |
|
|
| import torch |
|
|
| from .config import ModelConfig |
|
|
|
|
| class TokenSampler: |
| def __init__(self, cfg: ModelConfig, top_p: float = 0.98, temperature: float = 0.6, greedy: bool = False, |
| seed: int | None = None, top_k: int | None = None): |
| self.cfg = cfg |
| self.top_p = top_p |
| self.temperature = temperature |
| self.greedy = greedy |
| self.top_k = top_k |
| self.gen = torch.Generator() |
| if seed is not None: |
| self.gen.manual_seed(seed) |
| self.mask = torch.zeros(cfg.text.vocab, dtype=torch.bool) |
| for lo, hi in cfg.masked_logit_ranges: |
| self.mask[lo:hi] = True |
| self.stop_id = cfg.traj.future_start |
|
|
| def __call__(self, logits: torch.Tensor) -> int: |
| logits = logits.reshape(-1).float().clone() |
| logits[self.mask] = float("-inf") |
| if self.greedy or self.temperature <= 0: |
| return int(torch.argmax(logits)) |
| logits = logits / self.temperature |
| if self.top_k: |
| kth = torch.topk(logits, self.top_k).values[-1] |
| logits[logits < kth] = float("-inf") |
| probs = torch.softmax(logits, dim=-1) |
| if self.top_p < 1.0: |
| sp, si = torch.sort(probs, descending=True) |
| cum = torch.cumsum(sp, dim=-1) |
| |
| remove = cum - sp > self.top_p |
| sp[remove] = 0.0 |
| sp = sp / sp.sum() |
| idx = torch.multinomial(sp, 1, generator=self.gen) |
| return int(si[idx]) |
| return int(torch.multinomial(probs, 1, generator=self.gen)) |
|
|
| def is_stop(self, token: int) -> bool: |
| return token == self.stop_id |
|
|
