117 lines
5 KiB
C
117 lines
5 KiB
C
#include "gem/sampler.h"
|
|
|
|
#include <math.h>
|
|
#include <stdint.h>
|
|
#include <stdio.h>
|
|
|
|
int main(void)
|
|
{
|
|
char error[256];
|
|
struct gem_sampler_config greedy_config = {
|
|
.temperature = 0.0f, .top_k = 0, .top_p = 1.0f,
|
|
.repeat_penalty = 2.0f, .repeat_last_n = 8, .seed = 1,
|
|
};
|
|
struct gem_sampler *sampler = NULL;
|
|
if (gem_sampler_open(&sampler, 5, &greedy_config, error, sizeof(error)) != 0) {
|
|
fprintf(stderr, "%s\n", error);
|
|
return 1;
|
|
}
|
|
const float logits[] = {0.0f, 4.0f, 3.0f, -1.0f, 1.0f};
|
|
const uint32_t history[] = {1, 1};
|
|
uint32_t token;
|
|
if (gem_sampler_sample(sampler, logits, history, 2, &token, error, sizeof(error)) != 0 ||
|
|
token != 2) {
|
|
fprintf(stderr, "repeat-penalty greedy mismatch: token %u (%s)\n", token, error);
|
|
return 1;
|
|
}
|
|
gem_sampler_close(sampler);
|
|
|
|
struct gem_sampler_config random_config = {
|
|
.temperature = 0.8f, .top_k = 2, .top_p = 0.95f,
|
|
.repeat_penalty = 1.0f, .repeat_last_n = 0, .seed = 42,
|
|
};
|
|
if (gem_sampler_open(&sampler, 5, &random_config, error, sizeof(error)) != 0)
|
|
return 1;
|
|
for (size_t i = 0; i < 32; ++i) {
|
|
if (gem_sampler_sample(sampler, logits, NULL, 0, &token, error, sizeof(error)) != 0 ||
|
|
(token != 1 && token != 2)) {
|
|
fprintf(stderr, "top-k sample escaped candidate set\n");
|
|
return 1;
|
|
}
|
|
}
|
|
gem_sampler_close(sampler);
|
|
|
|
const float parity_logits[] = {4.0f, 3.0f, 2.0f, -INFINITY, -INFINITY,
|
|
-INFINITY, -INFINITY, -INFINITY};
|
|
struct gem_sampler_config heap_config = random_config;
|
|
heap_config.top_k = 3;
|
|
heap_config.seed = 1234;
|
|
struct gem_sampler_config sort_config = heap_config;
|
|
sort_config.top_k = 0;
|
|
struct gem_sampler *heap_sampler = NULL;
|
|
struct gem_sampler *sort_sampler = NULL;
|
|
if (gem_sampler_open(&heap_sampler, 8, &heap_config, error, sizeof(error)) != 0 ||
|
|
gem_sampler_open(&sort_sampler, 8, &sort_config, error, sizeof(error)) != 0)
|
|
return 1;
|
|
for (size_t i = 0; i < 128; ++i) {
|
|
uint32_t heap_token, sort_token;
|
|
if (gem_sampler_sample(heap_sampler, parity_logits, NULL, 0, &heap_token,
|
|
error, sizeof(error)) != 0 ||
|
|
gem_sampler_sample(sort_sampler, parity_logits, NULL, 0, &sort_token,
|
|
error, sizeof(error)) != 0 ||
|
|
heap_token != sort_token) {
|
|
fprintf(stderr, "heap/full-sort parity mismatch at sample %zu\n", i);
|
|
return 1;
|
|
}
|
|
}
|
|
gem_sampler_close(sort_sampler);
|
|
gem_sampler_close(heap_sampler);
|
|
|
|
const float tied_logits[] = {0.0f, 4.0f, 4.0f, 4.0f};
|
|
heap_config.top_k = 1;
|
|
if (gem_sampler_open(&sampler, 4, &heap_config, error, sizeof(error)) != 0 ||
|
|
gem_sampler_sample(sampler, tied_logits, NULL, 0, &token,
|
|
error, sizeof(error)) != 0 || token != 1) {
|
|
fprintf(stderr, "deterministic top-k tie mismatch: token %u\n", token);
|
|
return 1;
|
|
}
|
|
gem_sampler_close(sampler);
|
|
|
|
struct gem_sampler_config distribution_config = {
|
|
.temperature = 1.0f, .top_k = 3, .top_p = 1.0f,
|
|
.repeat_penalty = 1.0f, .repeat_last_n = 0, .seed = 7,
|
|
};
|
|
const float distribution_logits[] = {0.0f, logf(2.0f), -INFINITY, logf(3.0f)};
|
|
double probabilities[4];
|
|
if (gem_sampler_open(&sampler, 4, &distribution_config, error, sizeof(error)) != 0 ||
|
|
gem_sampler_distribution(sampler, distribution_logits, NULL, 0,
|
|
probabilities, error, sizeof(error)) != 0 ||
|
|
fabs(probabilities[0] - 1.0 / 6.0) > 1e-7 ||
|
|
fabs(probabilities[1] - 2.0 / 6.0) > 1e-7 || probabilities[2] != 0.0 ||
|
|
fabs(probabilities[3] - 3.0 / 6.0) > 1e-7) {
|
|
fprintf(stderr, "filtered sampler distribution mismatch: %s\n", error);
|
|
return 1;
|
|
}
|
|
uint32_t drawn;
|
|
bool accepted;
|
|
if (gem_sampler_draw_distribution(sampler, probabilities, &drawn,
|
|
error, sizeof(error)) != 0 || drawn >= 4 ||
|
|
gem_sampler_accept(sampler, 0.0, &accepted, error, sizeof(error)) != 0 || accepted ||
|
|
gem_sampler_accept(sampler, 1.0, &accepted, error, sizeof(error)) != 0 || !accepted) {
|
|
fprintf(stderr, "probability draw/accept mismatch: %s\n", error);
|
|
return 1;
|
|
}
|
|
const double target_probabilities[] = {0.25, 0.75};
|
|
const double draft_probabilities[] = {0.5, 0.5};
|
|
double residual[2];
|
|
if (fabs(gem_sampler_distribution_overlap(target_probabilities,
|
|
draft_probabilities, 2) - 0.75) > 1e-12 ||
|
|
gem_sampler_positive_difference(target_probabilities, draft_probabilities, 2,
|
|
residual, error, sizeof(error)) != 0 ||
|
|
residual[0] != 0.0 || residual[1] != 1.0) {
|
|
fprintf(stderr, "speculative residual distribution mismatch: %s\n", error);
|
|
return 1;
|
|
}
|
|
gem_sampler_close(sampler);
|
|
return 0;
|
|
}
|