| 19 | } |
| 20 | |
| 21 | float stable_softmax_pick(const float * row, int64_t vocab_size, float uniform) { |
| 22 | float max_value = row[0]; |
| 23 | for (int64_t i = 1; i < vocab_size; ++i) { |
| 24 | max_value = std::max(max_value, row[i]); |
| 25 | } |
| 26 | |
| 27 | std::vector<float> probs(static_cast<size_t>(vocab_size)); |
| 28 | float sum = 0.0f; |
| 29 | for (int64_t i = 0; i < vocab_size; ++i) { |
| 30 | probs[static_cast<size_t>(i)] = std::exp(row[i] - max_value); |
| 31 | sum += probs[static_cast<size_t>(i)]; |
| 32 | } |
| 33 | |
| 34 | float cumulative = 0.0f; |
| 35 | for (int64_t i = 0; i < vocab_size; ++i) { |
| 36 | cumulative += probs[static_cast<size_t>(i)] / sum; |
| 37 | if (uniform <= cumulative || i == vocab_size - 1) { |
| 38 | return static_cast<float>(i); |
| 39 | } |
| 40 | } |
| 41 | return 0.0f; |
| 42 | } |
| 43 | |
| 44 | } // namespace |
| 45 | |