Brunobkr/llama.cpp_AlgMor24_github

1#include <stdio.h>2#include <stdlib.h>3#include <string.h>4#include <assert.h>5 6#include "mtmd.h"7#include "mtmd-helper.h"8 9int main(void) {10    printf("\n\nTesting libmtmd C API...\n");11    printf("--------\n\n");12 13    struct mtmd_context_params params = mtmd_context_params_default();14    printf("Default image marker: %s\n", params.image_marker);15 16    mtmd_input_chunks * chunks = mtmd_test_create_input_chunks();17 18    if (!chunks) {19        fprintf(stderr, "Failed to create input chunks\n");20        return 1;21    }22 23    // simple test for the helper24    size_t n_tokens_total = mtmd_helper_get_n_tokens(chunks);25    printf("Total tokens in chunks: %zu\n", n_tokens_total);26    assert(n_tokens_total > 0);27 28    size_t n_chunks = mtmd_input_chunks_size(chunks);29    printf("Number of chunks: %zu\n", n_chunks);30    assert(n_chunks > 0);31 32    for (size_t i = 0; i < n_chunks; i++) {33        const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i);34        assert(chunk != NULL);35        enum mtmd_input_chunk_type type = mtmd_input_chunk_get_type(chunk);36        printf("Chunk %zu type: %d\n", i, type);37 38        if (type == MTMD_INPUT_CHUNK_TYPE_TEXT) {39            size_t n_tokens;40            const llama_token * tokens = mtmd_input_chunk_get_tokens_text(chunk, &n_tokens);41            printf("    Text chunk with %zu tokens\n", n_tokens);42            assert(tokens != NULL);43            assert(n_tokens > 0);44            for (size_t j = 0; j < n_tokens; j++) {45                assert(tokens[j] >= 0);46                printf("    > Token %zu: %d\n", j, tokens[j]);47            }48 49        } else if (type == MTMD_INPUT_CHUNK_TYPE_IMAGE) {50            const mtmd_image_tokens * image_tokens = mtmd_input_chunk_get_tokens_image(chunk);51            size_t n_tokens = mtmd_image_tokens_get_n_tokens(image_tokens);52            // get position of the last token, which should be (nx - 1, ny - 1)53            struct mtmd_decoder_pos pos = mtmd_image_tokens_get_decoder_pos(image_tokens, 0, n_tokens - 1);54            size_t nx = pos.x + 1;55            size_t ny = pos.y + 1;56            const char * id = mtmd_image_tokens_get_id(image_tokens);57            assert(n_tokens > 0);58            assert(nx > 0);59            assert(ny > 0);60            assert(id != NULL);61            printf("    Image chunk with %zu tokens\n", n_tokens);62            printf("    Image size: %zu x %zu\n", nx, ny);63            printf("    Image ID: %s\n", id);64        }65    }66 67    // test chunk save/load round-trip68    for (size_t i = 0; i < n_chunks; i++) {69        const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i);70        assert(chunk != NULL);71        enum mtmd_input_chunk_type type = mtmd_input_chunk_get_type(chunk);72 73        // query the required buffer size (out_buf == NULL)74        size_t expected_len = 0;75        int32_t rc = mtmd_input_chunk_save(chunk, NULL, 0, &expected_len);76        printf("    Chunk %zu: save query rc = %d, expected_len = %zu\n", i, rc, expected_len);77        assert(rc == 0);78        assert(expected_len > 0);79 80        // saving into a too-small buffer must fail, not crash81        char tiny_buf[1];82        rc = mtmd_input_chunk_save(chunk, tiny_buf, sizeof(tiny_buf), NULL);83        printf("    Chunk %zu: save into too-small buffer rc = %d (expect non-zero)\n", i, rc);84        assert(rc != 0);85 86        // save into a properly-sized buffer87        char * buf = (char *) malloc(expected_len);88        assert(buf != NULL);89        rc = mtmd_input_chunk_save(chunk, buf, expected_len, NULL);90        assert(rc == 0);91 92        // loading from a truncated buffer must fail gracefully, not crash93        if (expected_len > 1) {94            mtmd_input_chunk * bad = mtmd_input_chunk_load(buf, expected_len - 1);95            printf("    Chunk %zu: load from truncated buffer = %p (expect NULL)\n", i, (void *) bad);96            assert(bad == NULL);97        }98 99        // load it back100        mtmd_input_chunk * loaded = mtmd_input_chunk_load(buf, expected_len);101        assert(loaded != NULL);102 103        // metadata must match the original chunk104        assert(mtmd_input_chunk_get_type(loaded) == type);105        assert(mtmd_input_chunk_get_n_tokens(loaded) == mtmd_input_chunk_get_n_tokens(chunk));106        assert(mtmd_input_chunk_get_n_pos(loaded) == mtmd_input_chunk_get_n_pos(chunk));107 108        if (type == MTMD_INPUT_CHUNK_TYPE_TEXT) {109            size_t n_tok_orig, n_tok_loaded;110            const llama_token * tok_orig   = mtmd_input_chunk_get_tokens_text(chunk, &n_tok_orig);111            const llama_token * tok_loaded = mtmd_input_chunk_get_tokens_text(loaded, &n_tok_loaded);112            printf("    Chunk %zu: loaded %zu text tokens (orig %zu), first token %d (orig %d)\n",113                i, n_tok_loaded, n_tok_orig,114                n_tok_loaded > 0 ? tok_loaded[0] : -1,115                n_tok_orig   > 0 ? tok_orig[0]   : -1);116            assert(n_tok_orig == n_tok_loaded);117            for (size_t j = 0; j < n_tok_orig; j++) {118                assert(tok_orig[j] == tok_loaded[j]);119            }120        } else if (type == MTMD_INPUT_CHUNK_TYPE_IMAGE || type == MTMD_INPUT_CHUNK_TYPE_AUDIO) {121            const char * id_orig   = mtmd_input_chunk_get_id(chunk);122            const char * id_loaded = mtmd_input_chunk_get_id(loaded);123            printf("    Chunk %zu: loaded id '%s' (orig '%s')\n", i, id_loaded, id_orig);124            assert(id_orig != NULL && id_loaded != NULL);125            assert(strcmp(id_orig, id_loaded) == 0);126        }127 128        mtmd_input_chunk_free(loaded);129        free(buf);130    }131    printf("Chunk save/load round-trip OK\n");132 133    // Free the chunks134    mtmd_input_chunks_free(chunks);135 136    printf("\n\nDONE: test libmtmd C API...\n");137 138    return 0;139}140