Sample tokens from logits with optional filtering and repetition penalty.
(
logits: np.ndarray,
prev_tokens: np.ndarray | None = None,
repetition_penalty: float = 1.0,
top_p: float | None = None,
top_k: int | None = None,
do_sample: bool = True,
)
| 102 | |
| 103 | |
| 104 | def sample_token( |
| 105 | logits: np.ndarray, |
| 106 | prev_tokens: np.ndarray | None = None, |
| 107 | repetition_penalty: float = 1.0, |
| 108 | top_p: float | None = None, |
| 109 | top_k: int | None = None, |
| 110 | do_sample: bool = True, |
| 111 | ) -> np.ndarray: |
| 112 | """Sample tokens from logits with optional filtering and repetition penalty.""" |
| 113 | vocab_size = logits.shape[-1] |
| 114 | |
| 115 | if prev_tokens is not None and repetition_penalty != 1.0: |
| 116 | logits = apply_repetition_penalty(logits, prev_tokens, repetition_penalty) |
| 117 | |
| 118 | if not do_sample: |
| 119 | return np.argmax(logits, axis=-1).astype(np.int64) |
| 120 | |
| 121 | original_shape = logits.shape |
| 122 | flat = logits.reshape(-1, vocab_size) |
| 123 | N = flat.shape[0] |
| 124 | |
| 125 | if top_k is not None and top_k > 0: |
| 126 | k = min(top_k, vocab_size) |
| 127 | top_idx = np.argpartition(flat, -k, axis=-1)[:, -k:] |
| 128 | top_vals = np.take_along_axis(flat, top_idx, axis=-1) |
| 129 | |
| 130 | if top_p is not None and top_p < 1.0: |
| 131 | top_vals = apply_top_p(top_vals, top_p) |
| 132 | |
| 133 | probs = softmax(top_vals) |
| 134 | local = multinomial(probs) |
| 135 | tokens = top_idx[np.arange(N), local] |
| 136 | else: |
| 137 | if top_p is not None and top_p < 1.0: |
| 138 | flat = apply_top_p(flat, top_p) |
| 139 | probs = softmax(flat) |
| 140 | tokens = multinomial(probs) |
| 141 | |
| 142 | return tokens.reshape(original_shape[:-1]).astype(np.int64) |
no test coverage detected