David
/
Pantheon
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Pantheon/tests/test_grammar_mask_shortlist.c

218 lines
8.9 KiB
C
Raw Permalink Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
* Unit tests for the grammar mask's shortlist decision
* (src/core/grammar_mask_shortlist.c).
*
* Property, over thousands of random rows, allowed-sets and penalty settings:
* whenever the shortlist DECIDES, the draw it leaves is the one the full
* vocabulary mask leaves — greedy: the same argmax; sampling: the same top-K
* set after the penalties (the only tokens the sampler's draw can reach). Rows
* come both dense and as a backend's exact top-64 row (reference = the dense
* row it was cut from). The test also requires the shortlist to decide most
* rows (a decider that never decides would pass vacuously) and checks the
* undecidable cases (logit_bias, raising penalties, top_k off) leave the row
* untouched.
*
* Pure C, libc-only.
*/
#include "lib/test.h"
#include "../src/core/grammar_mask_shortlist.h"
#include "../src/core/sampling_modifiers.h"
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#define NV 3000
#define TOPN 64
static uint64_t g_rng = 88172645463325252ull;
static uint32_t rnd(void) {
g_rng ^= g_rng << 13; g_rng ^= g_rng >> 7; g_rng ^= g_rng << 17;
return (uint32_t)(g_rng >> 11);
}
static float rnd_unit(void) { return (float)(rnd() % 1000000) / 1000000.0f; }
typedef struct { uint32_t salt; uint32_t keep_per_100; } allow_rule;
static bool allowed_by_rule(void *ctx, uint32_t id) {
const allow_rule *r = (const allow_rule *)ctx;
return ((id * 2654435761u) ^ r->salt) % 100u < r->keep_per_100;
}
/* The sampler's view after a FULL mask: disallowed → -1e30, then the
* penalties in the sampler's order. */
static void full_mask_view(const float *row, float *out, const gms_sampling *s,
const allow_rule *rule) {
for (int i = 0; i < NV; i++)
out[i] = allowed_by_rule((void *)rule, (uint32_t)i) ? row[i] : -1e30f;
pn_apply_presence_frequency_penalties(out, NV, s->presence_penalty, s->frequency_penalty,
s->recent, s->n_recent);
if (s->temperature > 0.0f)
for (int i = 0; i < s->n_recent; i++) {
uint32_t t = s->recent[i];
out[t] = out[t] > 0.0f ? out[t] / s->repeat_penalty : out[t] * s->repeat_penalty;
}
}
/* The same view of the shortlist's rewritten row (-inf stays -inf). */
static void short_view(const float *rewritten, float *out, const gms_sampling *s) {
memcpy(out, rewritten, sizeof(float) * NV);
pn_apply_presence_frequency_penalties(out, NV, s->presence_penalty, s->frequency_penalty,
s->recent, s->n_recent);
if (s->temperature > 0.0f)
for (int i = 0; i < s->n_recent; i++) {
uint32_t t = s->recent[i];
if (out[t] < -1e37f) continue;
out[t] = out[t] > 0.0f ? out[t] / s->repeat_penalty : out[t] * s->repeat_penalty;
}
}
static int argmax(const float *v) {
int best = 0;
for (int i = 1; i < NV; i++) if (v[i] > v[best]) best = i;
return best;
}
/* Top-K membership: value >= the K-th largest (ties included, as the sampler). */
static void topk_set(const float *v, int k, bool *in) {
float *tmp = malloc(sizeof(float) * NV);
memcpy(tmp, v, sizeof(float) * NV);
for (int a = 0; a < k; a++) { /* partial selection, small k */
int best = a;
for (int i = a + 1; i < NV; i++) if (tmp[i] > tmp[best]) best = i;
float t = tmp[a]; tmp[a] = tmp[best]; tmp[best] = t;
}
float kth = tmp[k - 1];
for (int i = 0; i < NV; i++) in[i] = v[i] >= kth && v[i] > -1e29f;
free(tmp);
}
/* An LM-like row: a long tail of noise and a few hundred plausible tokens. */
static void make_row(float *row) {
for (int i = 0; i < NV; i++) row[i] = rnd_unit() * 8.0f - 4.0f;
for (int j = 0; j < 300; j++) row[rnd() % NV] = 4.0f + rnd_unit() * rnd_unit() * 16.0f;
}
/* The backend's exact top-N presentation of a dense row. */
static float neg_inf(void) {
uint32_t bits = 0xff800000u;
float v;
memcpy(&v, &bits, sizeof(v));
return v;
}
static void make_topn(const float *dense, float *topn) {
bool taken[NV] = { false };
for (int i = 0; i < NV; i++) topn[i] = neg_inf();
for (int k = 0; k < TOPN; k++) {
int best = -1;
for (int i = 0; i < NV; i++)
if (!taken[i] && (best < 0 || dense[i] > dense[best])) best = i;
taken[best] = true;
topn[best] = dense[best];
}
}
static void test_shortlist_matches_full_mask(void) {
static float dense[NV], row[NV], ref[NV], got[NV];
uint32_t recent[64], ids[GMS_MAX];
int decided = 0, trials = 0, mismatches = 0, decided_greedy = 0, decided_sampled = 0;
for (int t = 0; t < 3000; t++) {
make_row(dense);
allow_rule rule = { rnd(), 5u + rnd() % 96u };
/* The repeat window hits the row's TOP tokens (an agent re-emitting its
* own recent text), with repeats, so penalties push shortlisted tokens
* below the bound — the case the bound and the penalty replica exist
* for. */
uint32_t top[64], bound_ids[GMS_MAX];
float top_bound;
int n_top = gms_select_shortlist(dense, NV, bound_ids, &top_bound);
for (int i = 0; i < 64; i++) top[i] = bound_ids[rnd() % (uint32_t)n_top];
int n_recent = (int)(rnd() % 64);
for (int i = 0; i < n_recent; i++)
recent[i] = (i % 4 == 0) ? rnd() % NV : top[rnd() % 16];
gms_sampling s = {
.temperature = (t % 2) ? 0.7f : 0.0f, .top_k = 1 + (int)(rnd() % 40),
.repeat_penalty = 1.0f + (float)(rnd() % 6) * 0.1f,
.presence_penalty = (float)(rnd() % 3) * 1.0f,
.frequency_penalty = (float)(rnd() % 3) * 0.5f,
.recent = recent, .n_recent = n_recent,
};
bool use_topn = (t / 2) % 2;
if (use_topn) make_topn(dense, row); else memcpy(row, dense, sizeof(row));
float bound;
int n = gms_select_shortlist(row, NV, ids, &bound);
trials++;
if (!gms_mask_from_shortlist(row, NV, ids, n, bound, &s, allowed_by_rule, &rule))
continue;
decided++;
if (s.temperature <= 0.0f) decided_greedy++; else decided_sampled++;
full_mask_view(dense, ref, &s, &rule);
short_view(row, got, &s);
if (s.temperature <= 0.0f) {
if (argmax(ref) != argmax(got)) mismatches++;
} else {
bool in_ref[NV], in_got[NV];
topk_set(ref, s.top_k, in_ref);
topk_set(got, s.top_k, in_got);
if (memcmp(in_ref, in_got, sizeof(in_ref)) != 0) mismatches++;
}
}
printf(" trials=%d decided=%d (greedy %d, sampled %d) mismatches=%d\n", trials, decided,
decided_greedy, decided_sampled, mismatches);
ASSERT_INT_EQ(mismatches, 0);
ASSERT_TRUE(decided_greedy > trials / 4);
ASSERT_TRUE(decided_sampled > trials / 20);
}
static void test_undecidable_leaves_row(void) {
static float row[NV], before[NV];
uint32_t ids[GMS_MAX];
float bound;
allow_rule all = { 0, 100 };
make_row(row);
memcpy(before, row, sizeof(row));
int n = gms_select_shortlist(row, NV, ids, &bound);
gms_sampling bias = { .temperature = 0.0f, .top_k = 40, .repeat_penalty = 1.0f,
.has_logit_bias = true };
ASSERT_FALSE(gms_mask_from_shortlist(row, NV, ids, n, bound, &bias, allowed_by_rule, &all));
gms_sampling raising = { .temperature = 0.7f, .top_k = 40, .repeat_penalty = 0.9f };
ASSERT_FALSE(gms_mask_from_shortlist(row, NV, ids, n, bound, &raising, allowed_by_rule, &all));
gms_sampling no_topk = { .temperature = 0.7f, .top_k = 0, .repeat_penalty = 1.0f };
ASSERT_FALSE(gms_mask_from_shortlist(row, NV, ids, n, bound, &no_topk, allowed_by_rule, &all));
allow_rule none = { 0, 0 };
gms_sampling greedy = { .temperature = 0.0f, .top_k = 40, .repeat_penalty = 1.0f };
ASSERT_FALSE(gms_mask_from_shortlist(row, NV, ids, n, bound, &greedy, allowed_by_rule, &none));
ASSERT_TRUE(memcmp(row, before, sizeof(row)) == 0);
}
static void test_select_shortlist(void) {
static float row[NV], topn[NV];
uint32_t ids[GMS_MAX];
float bound;
make_row(row);
int n = gms_select_shortlist(row, NV, ids, &bound);
ASSERT_INT_EQ(n, GMS_MAX);
int above = 0;
for (int i = 0; i < NV; i++) if (row[i] > bound) above++;
ASSERT_TRUE(above <= GMS_MAX);
for (int i = 0; i < n; i++) ASSERT_TRUE(row[ids[i]] >= bound);
make_topn(row, topn);
n = gms_select_shortlist(topn, NV, ids, &bound);
ASSERT_INT_EQ(n, TOPN);
float min = 1e30f;
for (int i = 0; i < n; i++) if (topn[ids[i]] < min) min = topn[ids[i]];
ASSERT_TRUE(bound == min);
}
int main(void) {
printf("\n=== grammar_mask_shortlist ===\n\n");
RUN_TEST(select_shortlist);
RUN_TEST(undecidable_leaves_row);
RUN_TEST(shortlist_matches_full_mask);
TEST_REPORT();
return __ft_failed == 0 ? 0 : 1;
}
Reference in New Issue View Git Blame Copy Permalink