112 lines
4.3 KiB
C++
112 lines
4.3 KiB
C++
#include "sampler.h"
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <vector>
|
|
|
|
void sampler_seed(Sampler& s, uint32_t seed) { s.rng_state = seed ? seed : 1u; s.recent.clear(); }
|
|
void sampler_reset_history(Sampler& s) { s.recent.clear(); }
|
|
|
|
static inline float xs_rand(uint32_t& st) {
|
|
st ^= st << 13; st ^= st >> 17; st ^= st << 5;
|
|
return (st & 0x00FFFFFF) / (float)0x01000000;
|
|
}
|
|
|
|
// HF transformers RepetitionPenaltyLogitsProcessor :
|
|
// if logit > 0 : logit /= penalty
|
|
// else : logit *= penalty
|
|
static void apply_rep_penalty_(float* logits, int vocab, const int* tokens, int n_tokens, float penalty) {
|
|
if (penalty == 1.0f || n_tokens == 0) return;
|
|
for (int i = 0; i < n_tokens; i++) {
|
|
int t = tokens[i];
|
|
if (t < 0 || t >= vocab) continue;
|
|
float v = logits[t];
|
|
logits[t] = (v > 0.0f) ? (v / penalty) : (v * penalty);
|
|
}
|
|
}
|
|
|
|
// Garde les top_k logits, met les autres à -INFINITY.
|
|
// Si top_k <= 0 ou >= vocab, no-op.
|
|
static void apply_top_k_(float* logits, int vocab, int top_k) {
|
|
if (top_k <= 0 || top_k >= vocab) return;
|
|
std::vector<float> sorted(logits, logits + vocab);
|
|
std::nth_element(sorted.begin(), sorted.begin() + (vocab - top_k), sorted.end());
|
|
float threshold = sorted[vocab - top_k];
|
|
for (int i = 0; i < vocab; ++i)
|
|
if (logits[i] < threshold) logits[i] = -INFINITY;
|
|
}
|
|
|
|
// Nucleus sampling : sort logits desc, garde la masse cumsum jusqu'à top_p, met le reste à -inf.
|
|
// Garde au moins 1 token.
|
|
static void apply_top_p_(float* logits, int vocab, float top_p) {
|
|
if (top_p >= 1.0f || top_p <= 0.0f) return;
|
|
std::vector<int> idx(vocab);
|
|
for (int i = 0; i < vocab; ++i) idx[i] = i;
|
|
std::sort(idx.begin(), idx.end(), [&](int a, int b){ return logits[a] > logits[b]; });
|
|
// softmax pour cumsum (stable : soustrait max)
|
|
float mx = logits[idx[0]];
|
|
std::vector<float> probs(vocab);
|
|
double sum = 0.0;
|
|
for (int i = 0; i < vocab; ++i) {
|
|
float v = std::exp(logits[idx[i]] - mx);
|
|
probs[i] = v;
|
|
sum += v;
|
|
}
|
|
if (sum <= 0) return;
|
|
double inv = 1.0 / sum;
|
|
double cum = 0.0;
|
|
int cutoff = 0;
|
|
for (int i = 0; i < vocab; ++i) {
|
|
cum += probs[i] * inv;
|
|
if (cum >= top_p) { cutoff = i; break; }
|
|
}
|
|
// garde [0..cutoff] (inclus), met le reste à -inf
|
|
for (int i = cutoff + 1; i < vocab; ++i) logits[idx[i]] = -INFINITY;
|
|
}
|
|
|
|
static int multinomial_(float* logits, int vocab, uint32_t& rng_state) {
|
|
// softmax + sample
|
|
float mx = -INFINITY;
|
|
for (int i = 0; i < vocab; ++i) if (logits[i] > mx) mx = logits[i];
|
|
if (!std::isfinite(mx)) return 0;
|
|
double sum = 0.0;
|
|
std::vector<double> w(vocab);
|
|
for (int i = 0; i < vocab; ++i) {
|
|
if (std::isfinite(logits[i])) { w[i] = std::exp(logits[i] - mx); sum += w[i]; }
|
|
else { w[i] = 0.0; }
|
|
}
|
|
double r = xs_rand(rng_state) * sum;
|
|
double acc = 0.0;
|
|
for (int i = 0; i < vocab; ++i) { acc += w[i]; if (r <= acc) return i; }
|
|
// fallback : argmax (au cas où numérique)
|
|
int b = 0; for (int i = 1; i < vocab; ++i) if (logits[i] > logits[b]) b = i;
|
|
return b;
|
|
}
|
|
|
|
int sampler_sample(Sampler& s, float* logits, int vocab) {
|
|
// rep_penalty depuis l'historique (Talker = phrase entière)
|
|
if (s.rep_penalty != 1.0f && !s.recent.empty()) {
|
|
std::vector<int> recent_v(s.recent.begin(), s.recent.end());
|
|
apply_rep_penalty_(logits, vocab, recent_v.data(), (int)recent_v.size(), s.rep_penalty);
|
|
}
|
|
// temperature
|
|
if (s.temp > 0.0f && s.temp != 1.0f)
|
|
for (int i = 0; i < vocab; ++i) logits[i] /= s.temp;
|
|
apply_top_k_(logits, vocab, s.top_k);
|
|
apply_top_p_(logits, vocab, s.top_p);
|
|
int tok = multinomial_(logits, vocab, s.rng_state);
|
|
s.recent.push_back(tok);
|
|
while ((int)s.recent.size() > s.rep_window) s.recent.pop_front();
|
|
return tok;
|
|
}
|
|
|
|
int sampler_sample_local(Sampler& s, float* logits, int vocab,
|
|
const int* recent_step, int n_recent) {
|
|
if (s.rep_penalty != 1.0f && n_recent > 0)
|
|
apply_rep_penalty_(logits, vocab, recent_step, n_recent, s.rep_penalty);
|
|
if (s.temp > 0.0f && s.temp != 1.0f)
|
|
for (int i = 0; i < vocab; ++i) logits[i] /= s.temp;
|
|
apply_top_k_(logits, vocab, s.top_k);
|
|
apply_top_p_(logits, vocab, s.top_p);
|
|
return multinomial_(logits, vocab, s.rng_state);
|
|
}
|