File size: 1,146 Bytes
4d5ab1c | 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 | /*
Implements a simple Sampler, used during model inference to sample tokens.
*/
#ifndef SAMPLER_H
#define SAMPLER_H
#include <math.h>
// Simple xorshift RNG
unsigned int random_u32(unsigned long long *state) {
// xorshift rng: https://en.wikipedia.org/wiki/Xorshift#xorshift.2A
*state ^= *state >> 12;
*state ^= *state << 25;
*state ^= *state >> 27;
return (*state * 0x2545F4914F6CDD1Dull) >> 32;
}
float random_f32(unsigned long long *state) { // random float32 in [0,1)
return (random_u32(state) >> 8) / 16777216.0f;
}
int sample_softmax(const float* logits, int n, float coin) {
// sample index from logits (converted to probabilities using softmax)
// coin is a random number in [0, 1), usually from random_f32()
double norm = 0;
for (int i = 0; i < n; i++) {
norm += expf(logits[i]);
}
// instead of dividing all exp(logits), we can just multiply coin.
coin *= norm;
float cdf = 0.0f;
for (int i = 0; i < n; i++) {
cdf += expf(logits[i]);
if (coin < cdf) {
return i;
}
}
return n - 1; // in case of rounding errors
}
#endif |