258 lines
8.8 KiB
C
258 lines
8.8 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* test_precision_at_scale.c — HDC Memory Viability Test 4
|
|
*
|
|
* Retrieval Precision at Scale
|
|
*
|
|
* Method:
|
|
* 1. Generate N random embeddings (dim=768), store all via fm_put()
|
|
* 2. For 100 queries: brute-force cosine top-5 vs HDC top-5
|
|
* 3. N = 100, 1K, 10K, 100K
|
|
* 4. Report precision@5 at each N
|
|
*
|
|
* Usage:
|
|
* ./examples/test_precision_at_scale [--max-n 100000]
|
|
*/
|
|
|
|
#include "hybridmem/hybridmem.h"
|
|
#include "hdc/hdc.h"
|
|
#include "tensor/tensor.h"
|
|
|
|
#include <math.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <time.h>
|
|
|
|
/* Simple xorshift PRNG */
|
|
static uint32_t rng_state = 42;
|
|
static float randf(void) {
|
|
rng_state ^= rng_state << 13;
|
|
rng_state ^= rng_state >> 17;
|
|
rng_state ^= rng_state << 5;
|
|
return (float)(rng_state & 0x7FFFFFFF) / (float)0x7FFFFFFF * 2.0f - 1.0f;
|
|
}
|
|
|
|
/* Cosine similarity */
|
|
static float cosine_sim(const float *a, const float *b, int dim) {
|
|
float dot = 0, na = 0, nb = 0;
|
|
for (int i = 0; i < dim; i++) {
|
|
dot += a[i] * b[i];
|
|
na += a[i] * a[i];
|
|
nb += b[i] * b[i];
|
|
}
|
|
float d = sqrtf(na) * sqrtf(nb);
|
|
return d > 1e-12f ? dot / d : 0;
|
|
}
|
|
|
|
typedef struct { int idx; float score; } scored_item;
|
|
|
|
static int cmp_scored_desc(const void *a, const void *b) {
|
|
float diff = ((const scored_item*)b)->score - ((const scored_item*)a)->score;
|
|
return diff > 0 ? 1 : (diff < 0 ? -1 : 0);
|
|
}
|
|
|
|
/* Brute-force cosine top-k. Returns indices. */
|
|
static void brute_force_topk(const float *query, const float *corpus,
|
|
int n, int dim, int k, int *out_indices) {
|
|
scored_item *items = (scored_item*)malloc(n * sizeof(scored_item));
|
|
for (int i = 0; i < n; i++) {
|
|
items[i].idx = i;
|
|
items[i].score = cosine_sim(query, &corpus[i * dim], dim);
|
|
}
|
|
qsort(items, n, sizeof(scored_item), cmp_scored_desc);
|
|
for (int i = 0; i < k && i < n; i++)
|
|
out_indices[i] = items[i].idx;
|
|
free(items);
|
|
}
|
|
|
|
/* Run precision experiment at given N */
|
|
static float run_precision_experiment(int N, int dim, int n_queries, int top_k) {
|
|
fprintf(stderr, " N=%d: generating...", N);
|
|
fflush(stderr);
|
|
|
|
/* Generate embeddings */
|
|
float *corpus = (float*)malloc((size_t)N * dim * sizeof(float));
|
|
rng_state = 42;
|
|
for (int i = 0; i < N * dim; i++)
|
|
corpus[i] = randf();
|
|
|
|
/* Create HybridMem store and insert all entries */
|
|
char store_dir[256];
|
|
snprintf(store_dir, sizeof(store_dir), "/tmp/hm_precision_%d", N);
|
|
char cmd[512];
|
|
snprintf(cmd, sizeof(cmd), "rm -rf %s", store_dir);
|
|
system(cmd);
|
|
|
|
fm_config cfg = fm_config_default(dim);
|
|
cfg.store_dir = store_dir;
|
|
fm_store *store = fm_store_new(&cfg);
|
|
if (!store) {
|
|
fprintf(stderr, " ERROR: failed to create store\n");
|
|
free(corpus);
|
|
return -1;
|
|
}
|
|
|
|
fprintf(stderr, " storing...");
|
|
fflush(stderr);
|
|
|
|
/* Store entries individually (each is a 1-token entry) */
|
|
for (int i = 0; i < N; i++) {
|
|
int64_t shape_kv[] = {1, dim};
|
|
tm_tensor *k = tm_from_data(&corpus[i * dim], shape_kv, 2, TM_F32);
|
|
tm_tensor *k_owned = tm_clone(k);
|
|
tm_tensor *v_owned = tm_clone(k);
|
|
tm_free(k);
|
|
fm_put(store, k_owned, v_owned, 0, 1.0f);
|
|
tm_free(k_owned);
|
|
tm_free(v_owned);
|
|
}
|
|
fm_store_flush(store);
|
|
|
|
fprintf(stderr, " querying...");
|
|
fflush(stderr);
|
|
|
|
/* Run queries: pick random query embeddings and compare */
|
|
int total_hits = 0;
|
|
int total_possible = 0;
|
|
|
|
uint32_t query_rng = 12345;
|
|
for (int q = 0; q < n_queries; q++) {
|
|
/* Pick a random embedding as query */
|
|
query_rng ^= query_rng << 13;
|
|
query_rng ^= query_rng >> 17;
|
|
query_rng ^= query_rng << 5;
|
|
int query_idx = query_rng % N;
|
|
float *query = &corpus[query_idx * dim];
|
|
|
|
/* Brute-force cosine top-k */
|
|
int bf_indices[top_k];
|
|
brute_force_topk(query, corpus, N, dim, top_k, bf_indices);
|
|
|
|
/* HDC retrieval via fm_query */
|
|
int64_t shape_q[] = {dim};
|
|
tm_tensor *query_t = tm_from_data(query, shape_q, 1, TM_F32);
|
|
fm_results *results = fm_query(store, query_t, 0, top_k);
|
|
tm_free(query_t);
|
|
|
|
if (results) {
|
|
/* Collect retrieved indices by matching embedding content */
|
|
for (int r = 0; r < results->count; r++) {
|
|
float *rk = tm_data_f32(results->items[r].keys);
|
|
int n_ret = (int)tm_size(results->items[r].keys, 0);
|
|
for (int t = 0; t < n_ret; t++) {
|
|
/* Check against brute-force top-k */
|
|
for (int bi = 0; bi < top_k; bi++) {
|
|
float *bf_emb = &corpus[bf_indices[bi] * dim];
|
|
float diff_sq = 0;
|
|
for (int j = 0; j < dim; j++) {
|
|
float d = rk[t * dim + j] - bf_emb[j];
|
|
diff_sq += d * d;
|
|
}
|
|
if (diff_sq < 1e-4f) {
|
|
total_hits++;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
fm_results_free(results);
|
|
}
|
|
|
|
total_possible += top_k;
|
|
}
|
|
|
|
fm_store_free(store);
|
|
free(corpus);
|
|
snprintf(cmd, sizeof(cmd), "rm -rf %s", store_dir);
|
|
system(cmd);
|
|
|
|
float precision = total_possible > 0 ? (float)total_hits / total_possible : 0;
|
|
fprintf(stderr, " precision@%d = %.1f%%\n", top_k, precision * 100.0f);
|
|
return precision;
|
|
}
|
|
|
|
int main(int argc, char **argv) {
|
|
int max_n = 100000;
|
|
int dim = 768;
|
|
int n_queries = 100;
|
|
int top_k = 5;
|
|
|
|
for (int i = 1; i < argc; i++) {
|
|
if (strcmp(argv[i], "--max-n") == 0 && i + 1 < argc)
|
|
max_n = atoi(argv[++i]);
|
|
else if (strcmp(argv[i], "--dim") == 0 && i + 1 < argc)
|
|
dim = atoi(argv[++i]);
|
|
else if (strcmp(argv[i], "--queries") == 0 && i + 1 < argc)
|
|
n_queries = atoi(argv[++i]);
|
|
else if (strcmp(argv[i], "--top-k") == 0 && i + 1 < argc)
|
|
top_k = atoi(argv[++i]);
|
|
else if (strcmp(argv[i], "--help") == 0) {
|
|
fprintf(stderr,
|
|
"Usage: %s [options]\n"
|
|
" --max-n N Maximum store size (default: 100000)\n"
|
|
" --dim N Embedding dimension (default: 768)\n"
|
|
" --queries N Number of queries (default: 100)\n"
|
|
" --top-k N Top-k for precision (default: 5)\n",
|
|
argv[0]);
|
|
return 0;
|
|
}
|
|
}
|
|
|
|
fprintf(stderr, "==========================================\n");
|
|
fprintf(stderr, "HDC Memory Viability — Test 4\n");
|
|
fprintf(stderr, "Retrieval Precision at Scale\n");
|
|
fprintf(stderr, "==========================================\n");
|
|
fprintf(stderr, "Embedding dim: %d\n", dim);
|
|
fprintf(stderr, "Queries: %d\n", n_queries);
|
|
fprintf(stderr, "Top-k: %d\n", top_k);
|
|
fprintf(stderr, "Max N: %d\n\n", max_n);
|
|
|
|
int test_sizes[] = {100, 1000, 10000, 100000};
|
|
int n_sizes = sizeof(test_sizes) / sizeof(test_sizes[0]);
|
|
|
|
float precisions[4];
|
|
int n_run = 0;
|
|
|
|
for (int s = 0; s < n_sizes; s++) {
|
|
if (test_sizes[s] > max_n) break;
|
|
|
|
struct timespec t0, t1;
|
|
clock_gettime(CLOCK_MONOTONIC, &t0);
|
|
precisions[s] = run_precision_experiment(test_sizes[s], dim, n_queries, top_k);
|
|
clock_gettime(CLOCK_MONOTONIC, &t1);
|
|
double secs = (t1.tv_sec - t0.tv_sec) + (t1.tv_nsec - t0.tv_nsec) / 1e9;
|
|
fprintf(stderr, " Time: %.1fs\n", secs);
|
|
n_run++;
|
|
}
|
|
|
|
/* Summary */
|
|
fprintf(stderr, "\n--- Summary ---\n");
|
|
fprintf(stderr, "%-10s %-15s\n", "N", "precision@5");
|
|
for (int i = 0; i < 30; i++) fprintf(stderr, "-");
|
|
fprintf(stderr, "\n");
|
|
for (int s = 0; s < n_run; s++) {
|
|
fprintf(stderr, "%-10d %.1f%%\n", test_sizes[s], precisions[s] * 100.0f);
|
|
}
|
|
|
|
fprintf(stderr, "\nVERDICT:\n");
|
|
/* Check precision at largest tested scale */
|
|
float last_precision = precisions[n_run - 1];
|
|
if (last_precision >= 0.80f) {
|
|
fprintf(stderr, " PASS — precision@5 = %.1f%% at N=%d\n",
|
|
last_precision * 100.0f, test_sizes[n_run - 1]);
|
|
} else if (last_precision >= 0.50f) {
|
|
fprintf(stderr, " MARGINAL — precision@5 = %.1f%% at N=%d\n",
|
|
last_precision * 100.0f, test_sizes[n_run - 1]);
|
|
fprintf(stderr, " Consider: increase HDC dim, add re-ranking, or use multi-probe\n");
|
|
} else {
|
|
fprintf(stderr, " FAIL — precision@5 = %.1f%% at N=%d\n",
|
|
last_precision * 100.0f, test_sizes[n_run - 1]);
|
|
fprintf(stderr, " HDC index quality degrades too much at this scale\n");
|
|
}
|
|
|
|
fprintf(stderr, "\n==========================================\n");
|
|
return 0;
|
|
}
|