test_embedding.c

/*
 * test_embedding.c — Test EmbeddingGemma-300M embedding + intent recognition.
 *
 * Usage: ./test_embedding <model_dir>
 * Example: ./test_embedding models/embeddinggemma
 */

#include "embedding.h"

#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#define EMB_DIM 768

static void print_embedding(const char *label, const float *emb, int n) {
    printf("  %s first %d: [", label, n);
    for (int i = 0; i < n; i++) printf("%s%.6f", i ? ", " : "", emb[i]);
    printf("]\n");
}

static float vec_norm(const float *v, int dim) {
    float s = 0;
    for (int i = 0; i < dim; i++) s += v[i] * v[i];
    return sqrtf(s);
}

/* Reference values from Python (sentence-transformers) */
static const float ref_first8_lights[] = {
    -0.188852f, 0.019110f, 0.053933f, 0.006253f,
    -0.017149f, 0.037513f, -0.021670f, 0.015921f
};
static const float ref_sim_lights = 0.958773f;
static const float ref_sim_weather = 0.407931f;


int main(int argc, char **argv) {
    if (argc < 2) {
        fprintf(stderr, "Usage: %s <model_dir>\n", argv[0]);
        return 1;
    }

    printf("Loading embedding model from %s...\n", argv[1]);
    embedding_model *model = embedding_model_load(argv[1]);
    if (!model) { fprintf(stderr, "Failed to load model\n"); return 1; }

    embedding_state *state = embedding_state_create(model, 128);

    /* --- Test 1: Basic embedding --- */
    printf("\n=== Test 1: Embedding generation ===\n");

    float emb1[EMB_DIM], emb2[EMB_DIM], emb3[EMB_DIM];

    printf("Embedding \"turn on the lights\"...\n");
    if (embedding_model_embed(model, state, "turn on the lights", emb1, EMB_DIM) != 0) {
        fprintf(stderr, "Failed to embed\n"); return 1;
    }
    print_embedding("", emb1, 8);
    printf("  Norm: %.6f\n", vec_norm(emb1, EMB_DIM));

    printf("Embedding \"switch on the lights\"...\n");
    embedding_model_embed(model, state, "switch on the lights", emb2, EMB_DIM);
    print_embedding("", emb2, 8);

    printf("Embedding \"what is the weather\"...\n");
    embedding_model_embed(model, state, "what is the weather", emb3, EMB_DIM);
    print_embedding("", emb3, 8);

    /* --- Test 2: Similarity --- */
    printf("\n=== Test 2: Cosine similarity ===\n");
    float sim_12 = embedding_cosine_similarity(emb1, emb2, EMB_DIM);
    float sim_13 = embedding_cosine_similarity(emb1, emb3, EMB_DIM);
    printf("  sim(\"turn on the lights\", \"switch on the lights\") = %.6f (ref: %.6f)\n",
           sim_12, ref_sim_lights);
    printf("  sim(\"turn on the lights\", \"what is the weather\")  = %.6f (ref: %.6f)\n",
           sim_13, ref_sim_weather);

    /* Check against reference */
    printf("\n=== Test 3: Reference comparison ===\n");
    float max_diff = 0;
    for (int i = 0; i < 8; i++) {
        float diff = fabsf(emb1[i] - ref_first8_lights[i]);
        if (diff > max_diff) max_diff = diff;
    }
    printf("  Max diff in first 8 dims: %.6f\n", max_diff);
    printf("  Similarity diff (lights):  %.6f\n", fabsf(sim_12 - ref_sim_lights));
    printf("  Similarity diff (weather): %.6f\n", fabsf(sim_13 - ref_sim_weather));

    int pass = (max_diff < 0.05f) && (fabsf(sim_12 - ref_sim_lights) < 0.1f);
    printf("  Result: %s\n", pass ? "PASS" : "FAIL (embeddings differ from reference)");

    /* --- Test 4: Intent recognition --- */
    printf("\n=== Test 4: Intent recognition ===\n");
    intent_recognizer *ir = intent_recognizer_create(model, state);

    printf("Registering intents...\n");
    intent_recognizer_register(ir, "turn on the lights", NULL, 0, 0);
    intent_recognizer_register(ir, "what is the weather", NULL, 0, 0);
    printf("  Registered %zu intents\n", intent_recognizer_get_count(ir));

    intent_match *matches = NULL;
    size_t count = 0;

    printf("Ranking \"switch on the lights\":\n");
    intent_recognizer_rank(ir, model, state, "switch on the lights",
                           0.7f, 3, &matches, &count);
    if (count > 0) {
        for (size_t i = 0; i < count; i++)
            printf("  [%zu] trigger=\"%s\" sim=%.4f\n",
                   i, matches[i].trigger_phrase, matches[i].similarity);
    } else {
        printf("  No match\n");
    }
    intent_recognizer_free_matches(matches, count);

    printf("Ranking \"play some music\":\n");
    matches = NULL; count = 0;
    intent_recognizer_rank(ir, model, state, "play some music",
                           0.7f, 3, &matches, &count);
    printf("  Matched: %s\n", count > 0 ? "yes" : "no");
    intent_recognizer_free_matches(matches, count);

    intent_recognizer_clear(ir);
    printf("  After clear: %zu intents\n", intent_recognizer_get_count(ir));
    intent_recognizer_free(ir);
    embedding_state_free(state);
    embedding_model_free(model);

    printf("\nDone.\n");
    return pass ? 0 : 1;
}