318 lines
11 KiB
C
318 lines
11 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* test_rosa_db_qa.c — Phase 4: Model answering with ROSA context
|
|
*
|
|
* The KEY TEST: Can the 3.6B model answer questions about past conversations
|
|
* when given ROSA-retrieved context?
|
|
*
|
|
* For each question, runs:
|
|
* A) Baseline: just the question (no context)
|
|
* B) ROSA: retrieved context + question
|
|
*
|
|
* Usage:
|
|
* ./examples/test_rosa_db_qa \
|
|
* --model <path> --vocab <path> --db <path> \
|
|
* --context-window 256 --max-tokens 64 --quantize
|
|
*/
|
|
|
|
#include "rosa/rosa.h"
|
|
#include "rwkvx/rwkvx.h"
|
|
#include "data/data.h"
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <time.h>
|
|
|
|
#define MAX_TOKENS_PER_CONV (1024 * 64)
|
|
#define MAX_QUERY_TOKENS 512
|
|
#define MAX_PROMPT_TOKENS (1024 * 4)
|
|
|
|
/* ---- Index builder (reused from Phase 3) ---- */
|
|
|
|
typedef struct {
|
|
sa_token_automaton *sa;
|
|
sa_position_map *pmap;
|
|
int total_tokens;
|
|
int n_conv;
|
|
int *conv_start;
|
|
int *conv_len;
|
|
uint32_t *all_tokens;
|
|
int all_tokens_cap;
|
|
} rosa_index;
|
|
|
|
static rosa_index* build_index(fd_dataset *ds, rx_tokenizer *tok) {
|
|
rosa_index *idx = (rosa_index *)calloc(1, sizeof(rosa_index));
|
|
idx->n_conv = fd_len(ds);
|
|
idx->sa = sa_token_new(1024, 65536);
|
|
idx->pmap = sa_position_map_new(4096);
|
|
idx->conv_start = (int *)malloc(idx->n_conv * sizeof(int));
|
|
idx->conv_len = (int *)malloc(idx->n_conv * sizeof(int));
|
|
idx->all_tokens_cap = 1024 * 1024;
|
|
idx->all_tokens = (uint32_t *)malloc(idx->all_tokens_cap * sizeof(uint32_t));
|
|
|
|
uint32_t *buf = (uint32_t *)malloc(MAX_TOKENS_PER_CONV * sizeof(uint32_t));
|
|
|
|
for (int c = 0; c < idx->n_conv; c++) {
|
|
fd_sample *sample = fd_get(ds, c);
|
|
const char *text = fd_sample_get(sample, "text");
|
|
if (!text) { idx->conv_start[c] = idx->total_tokens; idx->conv_len[c] = 0; continue; }
|
|
|
|
int n = rx_tokenizer_encode(tok, text, buf, MAX_TOKENS_PER_CONV);
|
|
if (n <= 0) { idx->conv_start[c] = idx->total_tokens; idx->conv_len[c] = 0; continue; }
|
|
|
|
idx->conv_start[c] = idx->total_tokens;
|
|
idx->conv_len[c] = n;
|
|
|
|
while (idx->total_tokens + n > idx->all_tokens_cap) {
|
|
idx->all_tokens_cap *= 2;
|
|
idx->all_tokens = (uint32_t *)realloc(idx->all_tokens,
|
|
idx->all_tokens_cap * sizeof(uint32_t));
|
|
}
|
|
|
|
for (int t = 0; t < n; t++) {
|
|
int gpos = idx->total_tokens + t;
|
|
sa_token_extend(idx->sa, (int)buf[t], gpos);
|
|
sa_position_map_append(idx->pmap, c, 0, t);
|
|
idx->all_tokens[gpos] = buf[t];
|
|
}
|
|
idx->total_tokens += n;
|
|
}
|
|
|
|
free(buf);
|
|
return idx;
|
|
}
|
|
|
|
static void free_index(rosa_index *idx) {
|
|
sa_token_free(idx->sa);
|
|
sa_position_map_free(idx->pmap);
|
|
free(idx->conv_start);
|
|
free(idx->conv_len);
|
|
free(idx->all_tokens);
|
|
free(idx);
|
|
}
|
|
|
|
/* ---- Retrieve context for a query ---- */
|
|
|
|
static char* retrieve_context(rosa_index *idx, rx_tokenizer *tok,
|
|
const char *query, int context_window) {
|
|
uint32_t qtokens[MAX_QUERY_TOKENS];
|
|
int qlen = rx_tokenizer_encode(tok, query, qtokens, MAX_QUERY_TOKENS);
|
|
if (qlen <= 0) return NULL;
|
|
|
|
int state = 0, match_len = 0;
|
|
int best_endpos = -1, best_match = 0;
|
|
|
|
for (int i = 0; i < qlen; i++) {
|
|
int ep = sa_token_query_step(idx->sa, (int)qtokens[i], &state, &match_len);
|
|
if (match_len > best_match) {
|
|
best_match = match_len;
|
|
best_endpos = ep;
|
|
}
|
|
}
|
|
|
|
if (best_endpos < 0) return NULL;
|
|
|
|
const sa_position_info *info = sa_position_map_get(idx->pmap, best_endpos);
|
|
if (!info) return NULL;
|
|
|
|
int conv_start = idx->conv_start[info->conv_idx];
|
|
int conv_len = idx->conv_len[info->conv_idx];
|
|
|
|
int match_end = info->token_offset;
|
|
int ctx_start = match_end - best_match - context_window / 2;
|
|
if (ctx_start < 0) ctx_start = 0;
|
|
int ctx_end = match_end + context_window / 2;
|
|
if (ctx_end > conv_len) ctx_end = conv_len;
|
|
|
|
int ctx_len = ctx_end - ctx_start;
|
|
if (ctx_len <= 0) return NULL;
|
|
|
|
return rx_tokenizer_decode(tok, &idx->all_tokens[conv_start + ctx_start], ctx_len);
|
|
}
|
|
|
|
/* ---- Generate answer from model ---- */
|
|
|
|
static char* generate_answer(rx_model *model, rx_tokenizer *tok,
|
|
const char *prompt, int max_tokens) {
|
|
rx_state *state = rx_state_new(model);
|
|
|
|
/* Tokenize and prefill prompt */
|
|
uint32_t *prompt_tokens = (uint32_t *)malloc(MAX_PROMPT_TOKENS * sizeof(uint32_t));
|
|
int prompt_len = rx_tokenizer_encode(tok, prompt, prompt_tokens, MAX_PROMPT_TOKENS);
|
|
|
|
tm_tensor *logits = NULL;
|
|
for (int i = 0; i < prompt_len; i++) {
|
|
if (logits) tm_free(logits);
|
|
logits = rx_forward(model, prompt_tokens[i], state);
|
|
}
|
|
|
|
/* Generate tokens */
|
|
uint32_t *out_tokens = (uint32_t *)malloc(max_tokens * sizeof(uint32_t));
|
|
int n_out = 0;
|
|
|
|
rx_sampling_params params = rx_sampling_defaults();
|
|
params.temperature = 0.0f; /* greedy for reproducibility */
|
|
params.top_p = 1.0f;
|
|
|
|
rx_config cfg = rx_model_config(model);
|
|
|
|
for (int i = 0; i < max_tokens; i++) {
|
|
uint32_t tok_id = rx_sample(logits, ¶ms, NULL, cfg.vocab_size);
|
|
tm_free(logits);
|
|
logits = NULL;
|
|
|
|
if (tok_id == 0) break; /* EOS */
|
|
out_tokens[n_out++] = tok_id;
|
|
|
|
logits = rx_forward(model, tok_id, state);
|
|
}
|
|
|
|
if (logits) tm_free(logits);
|
|
|
|
char *answer = NULL;
|
|
if (n_out > 0)
|
|
answer = rx_tokenizer_decode(tok, out_tokens, n_out);
|
|
else
|
|
answer = strdup("(no output)");
|
|
|
|
free(prompt_tokens);
|
|
free(out_tokens);
|
|
rx_state_free(state);
|
|
return answer;
|
|
}
|
|
|
|
/* ---- Main ---- */
|
|
|
|
int main(int argc, char **argv) {
|
|
const char *model_path = NULL;
|
|
const char *vocab_path = NULL;
|
|
const char *db_path = NULL;
|
|
int context_window = 256;
|
|
int max_tokens = 64;
|
|
int quantize = 0;
|
|
|
|
for (int i = 1; i < argc; i++) {
|
|
if (strcmp(argv[i], "--model") == 0 && i + 1 < argc)
|
|
model_path = argv[++i];
|
|
else if (strcmp(argv[i], "--vocab") == 0 && i + 1 < argc)
|
|
vocab_path = argv[++i];
|
|
else if (strcmp(argv[i], "--db") == 0 && i + 1 < argc)
|
|
db_path = argv[++i];
|
|
else if (strcmp(argv[i], "--context-window") == 0 && i + 1 < argc)
|
|
context_window = atoi(argv[++i]);
|
|
else if (strcmp(argv[i], "--max-tokens") == 0 && i + 1 < argc)
|
|
max_tokens = atoi(argv[++i]);
|
|
else if (strcmp(argv[i], "--quantize") == 0)
|
|
quantize = 1;
|
|
}
|
|
|
|
if (!model_path || !vocab_path || !db_path) {
|
|
fprintf(stderr, "Usage: %s --model <path> --vocab <path> --db <path> "
|
|
"[--context-window N] [--max-tokens N] [--quantize]\n", argv[0]);
|
|
return 1;
|
|
}
|
|
|
|
/* Load tokenizer */
|
|
printf("Loading tokenizer...\n");
|
|
rx_tokenizer *tok = rx_tokenizer_new(vocab_path);
|
|
if (!tok) { fprintf(stderr, "ERROR: tokenizer\n"); return 1; }
|
|
|
|
/* Load conversations and build index */
|
|
printf("Loading conversations...\n");
|
|
fd_dataset *ds = fd_load_chat_sqlite(db_path, 0);
|
|
if (!ds) { fprintf(stderr, "ERROR: database\n"); return 1; }
|
|
|
|
printf("Building ROSA index (%d conversations)...\n", fd_len(ds));
|
|
rosa_index *idx = build_index(ds, tok);
|
|
printf("Indexed %d tokens\n", idx->total_tokens);
|
|
|
|
/* Load model */
|
|
printf("Loading model from %s%s...\n", model_path, quantize ? " (INT8)" : "");
|
|
rx_model *model = quantize ? rx_model_load_q8(model_path) : rx_model_load(model_path);
|
|
if (!model) { fprintf(stderr, "ERROR: model\n"); return 1; }
|
|
|
|
/* ---- Test questions ---- */
|
|
typedef struct { const char *question; const char *hint; } qa_pair;
|
|
qa_pair questions[] = {
|
|
{"What is the WKV-7 kernel used for in RWKV-X?",
|
|
"recurrent state evolution / generalized delta rule"},
|
|
{"What does HybridMem use HDC binary vectors for?",
|
|
"approximate similarity search / retrieval"},
|
|
{"How many layers does the 0.2B RWKV-X model have?",
|
|
"12 RWKV + 4 attention = 16 total"},
|
|
{"What build flag enables BLAS acceleration in Foundry?",
|
|
"USE_BLAS=1"},
|
|
{"What is the purpose of LoRA adapters in the training system?",
|
|
"fine-tune attention blocks / parameter-efficient"},
|
|
{"What is the name chosen for the RWKV-X product?",
|
|
"Engram"},
|
|
{"What framework does Metal GPU use on macOS?",
|
|
"Metal / MetalPerformanceShaders / Accelerate"},
|
|
{"What format are model weights stored in?",
|
|
"safetensors"},
|
|
{"What is BootAI?",
|
|
"bare-metal UEFI AI app / boots into LLM inference"},
|
|
{"What tokenizer does RWKV-X use?",
|
|
"RWKV World / BPE tokenizer"},
|
|
};
|
|
int n_questions = 10;
|
|
|
|
printf("\n=== Phase 4: Model Answering (KEY TEST) ===\n\n");
|
|
|
|
for (int q = 0; q < n_questions; q++) {
|
|
printf("--- Question %d/%d ---\n", q + 1, n_questions);
|
|
printf("Q: %s\n", questions[q].question);
|
|
printf("Expected: %s\n\n", questions[q].hint);
|
|
|
|
/* A) Baseline: no context */
|
|
char baseline_prompt[2048];
|
|
snprintf(baseline_prompt, sizeof(baseline_prompt),
|
|
"<|user|>\n%s\n<|assistant|>\n", questions[q].question);
|
|
|
|
printf(" [Baseline] generating...\n");
|
|
clock_t t0 = clock();
|
|
char *baseline_answer = generate_answer(model, tok, baseline_prompt, max_tokens);
|
|
clock_t t1 = clock();
|
|
printf(" Baseline (%.1fs): %.200s\n\n",
|
|
(double)(t1 - t0) / CLOCKS_PER_SEC, baseline_answer);
|
|
|
|
/* B) ROSA: retrieve context + question */
|
|
char *context = retrieve_context(idx, tok, questions[q].question, context_window);
|
|
|
|
char *rosa_prompt = (char *)malloc(context_window * 16 + 2048);
|
|
if (context) {
|
|
snprintf(rosa_prompt, context_window * 16 + 2048,
|
|
"<|user|>\nContext from previous conversations:\n---\n%.4000s\n---\n\n"
|
|
"Based on the context above, answer the following question:\n%s\n<|assistant|>\n",
|
|
context, questions[q].question);
|
|
} else {
|
|
snprintf(rosa_prompt, 2048,
|
|
"<|user|>\n%s\n<|assistant|>\n", questions[q].question);
|
|
}
|
|
|
|
printf(" [ROSA] generating (context: %s)...\n", context ? "yes" : "no");
|
|
clock_t t2 = clock();
|
|
char *rosa_answer = generate_answer(model, tok, rosa_prompt, max_tokens);
|
|
clock_t t3 = clock();
|
|
printf(" ROSA (%.1fs): %.200s\n\n",
|
|
(double)(t3 - t2) / CLOCKS_PER_SEC, rosa_answer);
|
|
|
|
free(baseline_answer);
|
|
free(rosa_answer);
|
|
free(context);
|
|
free(rosa_prompt);
|
|
}
|
|
|
|
printf("=== Phase 4 Complete ===\n");
|
|
printf("Review answers above to score manually.\n");
|
|
|
|
/* Cleanup */
|
|
rx_model_free(model);
|
|
free_index(idx);
|
|
fd_free(ds);
|
|
rx_tokenizer_free(tok);
|
|
|
|
return 0;
|
|
}
|