| 26 | } |
| 27 | |
| 28 | int32_t sample_topk_row( |
| 29 | const float * logits, |
| 30 | int64_t vocab_size, |
| 31 | int64_t topk, |
| 32 | float temperature, |
| 33 | uint64_t seed, |
| 34 | uint64_t call_index, |
| 35 | const TorchCudaSamplingPolicy & policy, |
| 36 | TopKSamplerScratch & scratch) { |
| 37 | if (logits == nullptr || vocab_size <= 0) { |
| 38 | throw std::runtime_error("HeartMuLa sampler requires logits"); |
| 39 | } |
| 40 | if (!(temperature > 0.0F)) { |
| 41 | throw std::runtime_error("HeartMuLa sampler temperature must be positive"); |
| 42 | } |
| 43 | if (topk <= 0 || topk > vocab_size) { |
| 44 | throw std::runtime_error("HeartMuLa sampler topk is out of range"); |
| 45 | } |
| 46 | scratch.scores.resize(static_cast<size_t>(vocab_size)); |
| 47 | for (int64_t i = 0; i < vocab_size; ++i) { |
| 48 | scratch.scores[static_cast<size_t>(i)] = logits[i] / temperature; |
| 49 | } |
| 50 | if (topk < vocab_size) { |
| 51 | scratch.threshold_values = scratch.scores; |
| 52 | auto nth = scratch.threshold_values.begin() + static_cast<std::ptrdiff_t>(topk - 1); |
| 53 | std::nth_element(scratch.threshold_values.begin(), nth, scratch.threshold_values.end(), std::greater<float>()); |
| 54 | const float threshold = *nth; |
| 55 | for (float & score : scratch.scores) { |
| 56 | if (score < threshold) { |
| 57 | score = -std::numeric_limits<float>::infinity(); |
| 58 | } |
| 59 | } |
| 60 | } |
| 61 | float max_score = -std::numeric_limits<float>::infinity(); |
| 62 | for (const float score : scratch.scores) { |
| 63 | if (std::isfinite(score)) { |
| 64 | max_score = std::max(max_score, score); |
| 65 | } |
| 66 | } |
| 67 | if (!std::isfinite(max_score)) { |
| 68 | throw std::runtime_error("HeartMuLa sampler kept no finite logits"); |
| 69 | } |
| 70 | double total = 0.0; |
| 71 | for (const float score : scratch.scores) { |
| 72 | if (std::isfinite(score)) { |
| 73 | total += std::exp(static_cast<double>(score - max_score)); |
| 74 | } |
| 75 | } |
| 76 | if (!(total > 0.0) || !std::isfinite(total)) { |
| 77 | throw std::runtime_error("HeartMuLa sampler probability mass is invalid"); |
| 78 | } |
| 79 | double best_rank = -std::numeric_limits<double>::infinity(); |
| 80 | int32_t best = -1; |
| 81 | for (int64_t i = 0; i < vocab_size; ++i) { |
| 82 | const float score = scratch.scores[static_cast<size_t>(i)]; |
| 83 | if (!std::isfinite(score)) { |
| 84 | continue; |
| 85 | } |
no test coverage detected