218 lines
8.9 KiB
C
218 lines
8.9 KiB
C
/* 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;
|
|
}
|