Sample one token per row via multinomial distribution.
(probs: np.ndarray)
| 92 | |
| 93 | |
| 94 | def multinomial(probs: np.ndarray) -> np.ndarray: |
| 95 | """Sample one token per row via multinomial distribution.""" |
| 96 | N = probs.shape[0] |
| 97 | cum = np.cumsum(probs, axis=-1) |
| 98 | r = np.random.random(N)[:, None] |
| 99 | # Clamp to valid range: float32 cumsum may not reach 1.0 exactly, |
| 100 | # so r > cumsum[-1] would produce an out-of-bounds index. |
| 101 | return np.minimum((cum < r).sum(axis=-1), probs.shape[-1] - 1).astype(np.int64) |
| 102 | |
| 103 | |
| 104 | def sample_token( |