205 lines
7.1 KiB
C
205 lines
7.1 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* test_rosa_db_index.c — Phase 2: Index chat history with suffix automaton
|
|
*
|
|
* Loads conversations from SQLite, tokenizes them, builds a suffix automaton
|
|
* over the full token stream, and creates a position map for retrieval.
|
|
*
|
|
* Usage:
|
|
* ./examples/test_rosa_db_index --vocab <vocab_path> --db <db_path>
|
|
*/
|
|
|
|
#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) /* 64K tokens max per conversation */
|
|
|
|
static void print_usage(const char *prog) {
|
|
fprintf(stderr, "Usage: %s --vocab <path> --db <path> [--max-conv N]\n", prog);
|
|
}
|
|
|
|
int main(int argc, char **argv) {
|
|
const char *vocab_path = NULL;
|
|
const char *db_path = NULL;
|
|
int max_conv = 0;
|
|
|
|
for (int i = 1; i < argc; i++) {
|
|
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], "--max-conv") == 0 && i + 1 < argc)
|
|
max_conv = atoi(argv[++i]);
|
|
else {
|
|
print_usage(argv[0]);
|
|
return 1;
|
|
}
|
|
}
|
|
|
|
if (!vocab_path || !db_path) {
|
|
print_usage(argv[0]);
|
|
return 1;
|
|
}
|
|
|
|
/* ---- Load tokenizer ---- */
|
|
printf("Loading tokenizer from %s...\n", vocab_path);
|
|
rx_tokenizer *tok = rx_tokenizer_new(vocab_path);
|
|
if (!tok) {
|
|
fprintf(stderr, "ERROR: Failed to load tokenizer\n");
|
|
return 1;
|
|
}
|
|
|
|
/* ---- Load conversations from SQLite ---- */
|
|
printf("Loading conversations from %s...\n", db_path);
|
|
fd_dataset *ds = fd_load_chat_sqlite(db_path, max_conv);
|
|
if (!ds) {
|
|
fprintf(stderr, "ERROR: Failed to load database\n");
|
|
rx_tokenizer_free(tok);
|
|
return 1;
|
|
}
|
|
|
|
int n_conv = fd_len(ds);
|
|
printf("Loaded %d conversations\n\n", n_conv);
|
|
|
|
/* ---- Tokenize all conversations and build suffix automaton ---- */
|
|
printf("Tokenizing and indexing...\n");
|
|
|
|
uint32_t *token_buf = (uint32_t *)malloc(MAX_TOKENS_PER_CONV * sizeof(uint32_t));
|
|
sa_token_automaton *sa = sa_token_new(1024, 65536);
|
|
sa_position_map *pmap = sa_position_map_new(4096);
|
|
|
|
int total_tokens = 0;
|
|
int total_messages = 0;
|
|
clock_t t_start = clock();
|
|
|
|
for (int c = 0; c < n_conv; c++) {
|
|
fd_sample *sample = fd_get(ds, c);
|
|
const char *text = fd_sample_get(sample, "text");
|
|
if (!text) continue;
|
|
|
|
/* Count messages (approximate: count <|user|> and <|assistant|> tags) */
|
|
int msg_count = 0;
|
|
const char *p = text;
|
|
while ((p = strstr(p, "<|")) != NULL) {
|
|
if (strncmp(p, "<|user|>", 8) == 0 || strncmp(p, "<|assistant|>", 13) == 0)
|
|
msg_count++;
|
|
p += 2;
|
|
}
|
|
total_messages += msg_count;
|
|
|
|
/* Tokenize the full conversation text */
|
|
int n_tokens = rx_tokenizer_encode(tok, text, token_buf, MAX_TOKENS_PER_CONV);
|
|
if (n_tokens <= 0) continue;
|
|
|
|
/* Extend suffix automaton and position map for each token */
|
|
for (int t = 0; t < n_tokens; t++) {
|
|
int global_pos = total_tokens + t;
|
|
sa_token_extend(sa, (int)token_buf[t], global_pos);
|
|
sa_position_map_append(pmap, c, 0, t);
|
|
}
|
|
|
|
total_tokens += n_tokens;
|
|
|
|
/* Progress every 5 conversations */
|
|
if ((c + 1) % 5 == 0 || c == n_conv - 1) {
|
|
printf(" [%d/%d] %d tokens, %d SA states\r",
|
|
c + 1, n_conv, total_tokens, sa->n_states);
|
|
fflush(stdout);
|
|
}
|
|
}
|
|
|
|
clock_t t_end = clock();
|
|
double index_time = (double)(t_end - t_start) / CLOCKS_PER_SEC;
|
|
|
|
printf("\n\n");
|
|
|
|
/* ---- Compute memory usage ---- */
|
|
/* SA memory: n_states * sizeof(sa_token_state) + transitions */
|
|
size_t sa_mem = (size_t)sa->n_states * sizeof(sa_token_state);
|
|
for (int i = 0; i < sa->n_states; i++)
|
|
sa_mem += (size_t)sa->states[i].trans_cap * sizeof(sa_trans_entry);
|
|
double sa_mem_mb = (double)sa_mem / (1024.0 * 1024.0);
|
|
|
|
/* Position map memory */
|
|
size_t pmap_mem = (size_t)pmap->count * sizeof(sa_position_info);
|
|
double pmap_mem_mb = (double)pmap_mem / (1024.0 * 1024.0);
|
|
|
|
/* ---- Verify SA property: states <= 2*N + 1 ---- */
|
|
int sa_limit = 2 * total_tokens + 1;
|
|
int sa_ok = sa->n_states <= sa_limit;
|
|
|
|
/* ---- Quick query test ---- */
|
|
printf("Running quick query test...\n");
|
|
int query_state = 0, query_match = 0;
|
|
int query_endpos = -1;
|
|
|
|
/* Query with the first 10 tokens of the first conversation */
|
|
fd_sample *first = fd_get(ds, 0);
|
|
const char *first_text = fd_sample_get(first, "text");
|
|
int first_n = rx_tokenizer_encode(tok, first_text, token_buf, MAX_TOKENS_PER_CONV);
|
|
int query_len = first_n < 10 ? first_n : 10;
|
|
|
|
for (int t = 0; t < query_len; t++)
|
|
query_endpos = sa_token_query_step(sa, (int)token_buf[t], &query_state, &query_match);
|
|
|
|
printf(" Query: first %d tokens of conv 0\n", query_len);
|
|
printf(" Match length: %d tokens\n", query_match);
|
|
printf(" End position: %d\n", query_endpos);
|
|
if (query_endpos >= 0 && query_endpos < pmap->count) {
|
|
const sa_position_info *info = sa_position_map_get(pmap, query_endpos);
|
|
printf(" Mapped to: conv %d, token offset %d\n", info->conv_idx, info->token_offset);
|
|
}
|
|
printf("\n");
|
|
|
|
/* ---- Print results ---- */
|
|
printf("=== Phase 2 Results ===\n\n");
|
|
printf("| Metric | Value |\n");
|
|
printf("|-------------------------|----------------|\n");
|
|
printf("| Conversations loaded | %d |\n", n_conv);
|
|
printf("| Messages loaded | %d |\n", total_messages);
|
|
printf("| Total tokens | %d |\n", total_tokens);
|
|
printf("| SA states | %d |\n", sa->n_states);
|
|
printf("| SA states <= 2*N+1 | %s |\n", sa_ok ? "YES" : "NO");
|
|
printf("| SA memory (MB) | %.2f |\n", sa_mem_mb);
|
|
printf("| Position map memory (MB)| %.2f |\n", pmap_mem_mb);
|
|
printf("| Total memory (MB) | %.2f |\n", sa_mem_mb + pmap_mem_mb);
|
|
printf("| Index time (s) | %.2f |\n", index_time);
|
|
printf("\n");
|
|
|
|
/* ---- Pass/fail ---- */
|
|
int pass = 1;
|
|
if (!sa_ok) {
|
|
printf("FAIL: SA states (%d) > 2*N+1 (%d)\n", sa->n_states, sa_limit);
|
|
pass = 0;
|
|
}
|
|
if (sa_mem_mb + pmap_mem_mb > 500.0) {
|
|
printf("FAIL: Memory %.2f MB > 500 MB limit\n", sa_mem_mb + pmap_mem_mb);
|
|
pass = 0;
|
|
}
|
|
if (index_time > 30.0) {
|
|
printf("FAIL: Index time %.2f s > 30 s limit\n", index_time);
|
|
pass = 0;
|
|
}
|
|
if (query_match < query_len) {
|
|
printf("FAIL: Query self-match %d < %d (should match own tokens)\n", query_match, query_len);
|
|
pass = 0;
|
|
}
|
|
|
|
printf("**Verdict: %s**\n", pass ? "PASS" : "FAIL");
|
|
|
|
/* ---- Cleanup ---- */
|
|
sa_token_free(sa);
|
|
sa_position_map_free(pmap);
|
|
free(token_buf);
|
|
fd_free(ds);
|
|
rx_tokenizer_free(tok);
|
|
|
|
return pass ? 0 : 1;
|
|
}
|