| 215 | } |
| 216 | |
| 217 | void Sampler::sampler_topk_apply(int k) { |
| 218 | if (k == 0) { |
| 219 | return; |
| 220 | } |
| 221 | |
| 222 | logits_list_t pairs; |
| 223 | pairs.resize(in_features); |
| 224 | for (int i = 0; i < in_features; i++) { |
| 225 | pairs[i].logits = this->logits[i]; |
| 226 | pairs[i].token_id = i; |
| 227 | pairs[i].prob = 0.0f; |
| 228 | } |
| 229 | |
| 230 | std::partial_sort( |
| 231 | pairs.begin(), |
| 232 | pairs.begin() + k, |
| 233 | pairs.end(), |
| 234 | [](const logits_t& a, const logits_t& b) { |
| 235 | return a.logits > b.logits; |
| 236 | } |
| 237 | ); |
| 238 | |
| 239 | this->top_k_logits.assign(pairs.begin(), pairs.begin() + k); |
| 240 | } |
| 241 | |
| 242 | void Sampler::sampler_topp_apply(float p) { |
| 243 | if (p >= 1.0f) { |