-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathsampling.cpp
More file actions
135 lines (121 loc) · 4.21 KB
/
Copy pathsampling.cpp
File metadata and controls
135 lines (121 loc) · 4.21 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
#include "sampling.h"
#include <algorithm>
#include <cmath>
#include <numeric>
#include <random>
#include <stdexcept>
namespace muscriptor {
namespace {
std::mt19937& rng() {
static thread_local std::mt19937 gen{std::random_device{}()};
return gen;
}
int multinomial_1(const std::vector<float>& probs) {
float sum = std::accumulate(probs.begin(), probs.end(), 0.f);
if (sum <= 0.f) return 0;
std::uniform_real_distribution<float> dist(0.f, sum);
float r = dist(rng());
float cum = 0.f;
for (int i = 0; i < static_cast<int>(probs.size()); ++i) {
cum += probs[i];
if (r <= cum) return i;
}
return static_cast<int>(probs.size()) - 1;
}
} // namespace
int sample_top_k(std::vector<float> probs, int k) {
if (k <= 0 || k >= static_cast<int>(probs.size())) {
return multinomial_1(probs);
}
std::vector<int> idx(probs.size());
std::iota(idx.begin(), idx.end(), 0);
std::partial_sort(idx.begin(), idx.begin() + k, idx.end(),
[&](int a, int b) { return probs[a] > probs[b]; });
float min_val = probs[idx[k - 1]];
for (float& p : probs) {
if (p < min_val) p = 0.f;
}
float sum = std::accumulate(probs.begin(), probs.end(), 0.f);
if (sum > 0.f) {
for (float& p : probs) p /= sum;
}
return multinomial_1(probs);
}
int sample_top_p(std::vector<float> probs, float p) {
std::vector<int> idx(probs.size());
std::iota(idx.begin(), idx.end(), 0);
std::sort(idx.begin(), idx.end(), [&](int a, int b) { return probs[a] > probs[b]; });
float cum = 0.f;
std::vector<float> sorted_probs(probs.size(), 0.f);
for (size_t i = 0; i < idx.size(); ++i) {
float pi = probs[idx[i]];
if (cum - pi > p && i > 0) break; // mask when cumsum - current > p
// Match Python: mask = probs_sum - probs_sort > p
sorted_probs[i] = pi;
cum += pi;
if (cum - pi > p) {
sorted_probs[i] = 0.f;
}
}
// Rebuild properly matching Python
std::vector<float> probs_sort(probs.size());
for (size_t i = 0; i < idx.size(); ++i) probs_sort[i] = probs[idx[i]];
float probs_sum = 0.f;
for (size_t i = 0; i < probs_sort.size(); ++i) {
float prev = probs_sum;
probs_sum += probs_sort[i];
if (probs_sum - probs_sort[i] > p) {
probs_sort[i] = 0.f;
}
(void)prev;
}
float sum = std::accumulate(probs_sort.begin(), probs_sort.end(), 0.f);
if (sum > 0.f) {
for (float& v : probs_sort) v /= sum;
}
int local = multinomial_1(probs_sort);
return idx[local];
}
int sample_from_probs(const std::vector<float>& probs, float top_p, int top_k) {
if (top_p > 0.f) return sample_top_p(probs, top_p);
if (top_k > 0) return sample_top_k(probs, top_k);
return multinomial_1(probs);
}
int sample_stratified(
std::vector<float> probs,
int special_token,
float first_temp,
float second_temp,
float top_p,
int top_k) {
constexpr float eps = 1e-12f;
float p_special = std::clamp(probs[special_token], eps, 1.f - eps);
float logits0 = std::log(p_special) / std::max(first_temp, eps);
float logits1 = std::log(1.f - p_special) / std::max(first_temp, eps);
float m = std::max(logits0, logits1);
float e0 = std::exp(logits0 - m);
float e1 = std::exp(logits1 - m);
float p_special_temp = e0 / (e0 + e1);
std::uniform_real_distribution<float> dist(0.f, 1.f);
bool is_special = dist(rng()) < p_special_temp;
if (is_special) return special_token;
float denom = std::max(1.f - p_special, eps);
for (float& p : probs) p /= denom;
probs[special_token] = 0.f;
if (second_temp > 0.f) {
float max_log = -1e30f;
std::vector<float> logs(probs.size());
for (size_t i = 0; i < probs.size(); ++i) {
logs[i] = std::log(std::max(probs[i], eps)) / second_temp;
max_log = std::max(max_log, logs[i]);
}
float sum = 0.f;
for (size_t i = 0; i < logs.size(); ++i) {
probs[i] = std::exp(logs[i] - max_log);
sum += probs[i];
}
for (float& p : probs) p /= sum;
}
return sample_from_probs(probs, top_p, top_k);
}
} // namespace muscriptor