(
logits,
prev_tokens: Optional[torch.LongTensor] = None,
repetition_penalty: float = 1.0,
top_p=None,
top_k=None,
do_sample=True,
)
| 109 | |
| 110 | |
| 111 | def sample_token( |
| 112 | logits, |
| 113 | prev_tokens: Optional[torch.LongTensor] = None, |
| 114 | repetition_penalty: float = 1.0, |
| 115 | top_p=None, |
| 116 | top_k=None, |
| 117 | do_sample=True, |
| 118 | ): |
| 119 | vocab_size = logits.size(-1) |
| 120 | |
| 121 | # ===== Repetition Penalty (before reshaping!) ===== |
| 122 | if prev_tokens is not None and repetition_penalty != 1.0: |
| 123 | logits = apply_repetition_penalty_delay_pattern( |
| 124 | logits, |
| 125 | prev_tokens, |
| 126 | repetition_penalty, |
| 127 | ) |
| 128 | |
| 129 | if not do_sample: |
| 130 | return torch.argmax(logits, dim=-1) |
| 131 | |
| 132 | # ===== Only flatten after this, for top-k / top-p / multinomial ===== |
| 133 | original_shape = logits.shape |
| 134 | reshaped_logits = logits.view(-1, vocab_size) |
| 135 | |
| 136 | if top_k is not None and top_k > 0: |
| 137 | reshaped_logits = apply_top_k(reshaped_logits, top_k) |
| 138 | |
| 139 | if top_p is not None and top_p < 1.0: |
| 140 | reshaped_logits = apply_top_p_optimized(reshaped_logits, top_p) |
| 141 | |
| 142 | probs = F.softmax(reshaped_logits, dim=-1) |
| 143 | next_tokens = torch.multinomial(probs, num_samples=1) |
| 144 | |
| 145 | return next_tokens.view(original_shape[:-1]) |
| 146 | |
| 147 | |
| 148 | def find_last_equal_C(tensor, C): |
no test coverage detected