From 7339054744f109c4cd89b75689dbb8a2c154d60e Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Wed, 2 Sep 2026 21:20:10 +0200 Subject: [PATCH] mtmd: add mtmd_tokenize_from_parts() (#28250) * add mtmd_tokenize_from_parts * use it in mtmd-cli * move add_special to call level --- tests/test-mtmd-c-api.c | 33 ++++++++++++++++++++++ tests/test-mtmd-impl.cpp | 2 +- tools/mtmd/README-dev.md | 1 + tools/mtmd/mtmd-cli.cpp | 58 ++++++++++++++++++++++++++++---------- tools/mtmd/mtmd-internal.h | 6 ++-- tools/mtmd/mtmd.cpp | 54 +++++++++++++++++++++++++++++++---- tools/mtmd/mtmd.h | 27 +++++++++++++++--- 7 files changed, 154 insertions(+), 27 deletions(-) diff --git a/tests/test-mtmd-c-api.c b/tests/test-mtmd-c-api.c index 970d8a6000..664c56cf80 100644 --- a/tests/test-mtmd-c-api.c +++ b/tests/test-mtmd-c-api.c @@ -130,6 +130,39 @@ int main(void) { } printf("Chunk save/load round-trip OK\n"); + // test input validation of mtmd_tokenize_from_parts() + // invalid parts are rejected before the ctx is used, so NULL ctx is OK here + { + mtmd_input_chunks * out = mtmd_input_chunks_init(); + mtmd_bitmap * bmp = mtmd_bitmap_init(4, 4, NULL); // placeholder bitmap + struct mtmd_input_text txt = { "hello", 5, false, false }; + struct mtmd_input_text txt_null = { NULL, 0, false, false }; + + struct mtmd_input_part part_both = { &txt, bmp }; + struct mtmd_input_part part_neither = { NULL, NULL }; + struct mtmd_input_part part_null_text = { &txt_null, NULL }; + const mtmd_input_part * parts[1]; + int32_t rc; + + parts[0] = &part_both; + rc = mtmd_tokenize_from_parts(NULL, out, parts, 1, false); + printf("tokenize part with both text and bitmap rc = %d (expect 1)\n", rc); + assert(rc == 1); + + parts[0] = &part_neither; + rc = mtmd_tokenize_from_parts(NULL, out, parts, 1, false); + printf("tokenize part with neither text nor bitmap rc = %d (expect 1)\n", rc); + assert(rc == 1); + + parts[0] = &part_null_text; + rc = mtmd_tokenize_from_parts(NULL, out, parts, 1, false); + printf("tokenize part with null text pointer rc = %d (expect 1)\n", rc); + assert(rc == 1); + + mtmd_bitmap_free(bmp); + mtmd_input_chunks_free(out); + } + // Free the chunks mtmd_input_chunks_free(chunks); diff --git a/tests/test-mtmd-impl.cpp b/tests/test-mtmd-impl.cpp index df18b0a42e..2ec6b23915 100644 --- a/tests/test-mtmd-impl.cpp +++ b/tests/test-mtmd-impl.cpp @@ -80,7 +80,7 @@ MAKE_TEST(test_temporal_merge_grouping) { // spec chars: // v = video frame, w = video frame of another size, a = audio, i = plain image, t = text auto make_parts = [&pool](const std::string & spec) { - std::vector parts; + std::vector parts; for (char c : spec) { if (c == 't') { parts.push_back({ "hello", nullptr }); diff --git a/tools/mtmd/README-dev.md b/tools/mtmd/README-dev.md index e14906823a..b85627d2a4 100644 --- a/tools/mtmd/README-dev.md +++ b/tools/mtmd/README-dev.md @@ -20,6 +20,7 @@ In short: A typical pipeline of the core libmtmd is as follows: - A bitmap (RGB image or PCM audio) is created - Bitmap and the text prompt is provided to `mtmd_tokenize()` that breaks the input into chunks + - Alternatively, `mtmd_tokenize_from_parts()` takes a list of pre-split text/media parts instead of a marker-based prompt - The tokenizer function first expands a "lazy" bitmap if it finds one. Typically, this is used by video, so that one media token corresponds to one input bitmap - For models that support "fused" temporal frames like Qwen-VL, the tokenizer tries to merge pair of consecutive frames into one batch. Only bitmaps marked by `mtmd_bitmap_set_mergeable()` are merged - The preprocessor will then be called, which produces a list of chunks diff --git a/tools/mtmd/mtmd-cli.cpp b/tools/mtmd/mtmd-cli.cpp index 97678c6b2e..ba18b3e32b 100644 --- a/tools/mtmd/mtmd-cli.cpp +++ b/tools/mtmd/mtmd-cli.cpp @@ -109,16 +109,15 @@ struct mtmd_cli_context { mtmd_cli_context(common_params & params) : llama_init(common_init_from_params(params)) { model = llama_init->model(); lctx = llama_init->context(); + if (!model || !lctx) { + exit(1); + } vocab = llama_model_get_vocab(model); smpl = common_sampler_init(model, params.sampling); n_threads = params.cpuparams.n_threads; batch = llama_batch_init(1, 0, 1); // batch for next token generation n_batch = params.n_batch; - if (!model || !lctx) { - exit(1); - } - init_vision_context(params); if (!mtmd_helper_model_can_chat(lctx, ctx_vision.get())) { @@ -265,21 +264,50 @@ static int eval_message(mtmd_cli_context & ctx, common_chat_msg & msg) { auto formatted_chat = chat_add_and_format(ctx, msg); LOG_DBG("formatted_chat.prompt: %s\n", formatted_chat.c_str()); - mtmd_input_text text; - text.text = formatted_chat.data(); - text.text_len = formatted_chat.size(); - text.add_special = add_bos; - text.parse_special = true; - if (g_is_interrupted) return 0; - mtmd::input_chunks chunks(mtmd_input_chunks_init()); + // note: we replace the marker here instead of letting mtmd_tokenize() to do that + // because we want to demonstrate how to use mtmd_tokenize_from_parts() + + // split the formatted chat on the media marker to get text segments + const std::string marker = mtmd_default_marker(); + std::vector segments; + size_t start = 0; + size_t pos; + while ((pos = formatted_chat.find(marker, start)) != std::string::npos) { + segments.push_back(formatted_chat.substr(start, pos - start)); + start = pos + marker.size(); + } + segments.push_back(formatted_chat.substr(start)); + auto bitmaps_c_ptr = ctx.bitmaps.c_ptr(); - int32_t res = mtmd_tokenize(ctx.ctx_vision.get(), + if (segments.size() - 1 != bitmaps_c_ptr.size()) { + LOG_ERR("Number of media markers (%zu) does not match number of loaded media (%zu)\n", + segments.size() - 1, bitmaps_c_ptr.size()); + return 1; + } + + // interleave text and media parts + std::vector texts(segments.size()); + std::vector parts; + for (size_t i = 0; i < segments.size(); i++) { + texts[i] = {segments[i].data(), segments[i].size(), /* add_special */ false, /* parse_special */ true}; + parts.push_back({&texts[i], nullptr}); + if (i < bitmaps_c_ptr.size()) { + parts.push_back({nullptr, bitmaps_c_ptr[i]}); + } + } + std::vector parts_ptr; + for (const auto & p : parts) { + parts_ptr.push_back(&p); + } + + mtmd::input_chunks chunks(mtmd_input_chunks_init()); + int32_t res = mtmd_tokenize_from_parts(ctx.ctx_vision.get(), chunks.ptr.get(), // output - &text, // text - bitmaps_c_ptr.data(), - bitmaps_c_ptr.size()); + parts_ptr.data(), + parts_ptr.size(), + add_bos); if (res != 0) { LOG_ERR("Unable to tokenize prompt, res = %d\n", res); return 1; diff --git a/tools/mtmd/mtmd-internal.h b/tools/mtmd/mtmd-internal.h index 067fa88b99..e7c62773e3 100644 --- a/tools/mtmd/mtmd-internal.h +++ b/tools/mtmd/mtmd-internal.h @@ -10,10 +10,12 @@ #define MTMD_INTERNAL_HEADER // bitmap is null for text parts -struct mtmd_input_part { +struct mtmd_internal_part { std::string text; const mtmd_bitmap * bitmap; + // only used for text parts + bool parse_special = false; }; // [QWEN_VIDEO] merged parts are erased from `parts`, so one group always maps to one part -std::vector> mtmd_group_mergeable_bitmaps(std::vector & parts, int n_merge); +std::vector> mtmd_group_mergeable_bitmaps(std::vector & parts, int n_merge); diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index f39f0974e3..a13b0556b0 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -1097,7 +1097,7 @@ void mtmd_free(mtmd_context * ctx) { delete ctx; } -std::vector> mtmd_group_mergeable_bitmaps(std::vector & parts, int n_merge) { +std::vector> mtmd_group_mergeable_bitmaps(std::vector & parts, int n_merge) { std::vector> output; for (size_t i = 0; i < parts.size(); i++) { if (parts[i].bitmap == nullptr) { @@ -1124,7 +1124,7 @@ struct mtmd_tokenizer { bool parse_special; const llama_vocab * vocab; - using part = mtmd_input_part; + using part = mtmd_internal_part; std::vector parts; // these will be freed when mtmd_tokenizer finishes std::vector bm_from_lazy; // TODO @ngxson : refactor, free bm_from_lazy progressively @@ -1160,7 +1160,7 @@ struct mtmd_tokenizer { } parts.push_back({"", bitmaps[i_bm++]}); } else { - parts.push_back({std::move(part), nullptr}); + parts.push_back({std::move(part), nullptr, parse_special}); } } @@ -1177,6 +1177,26 @@ struct mtmd_tokenizer { expand_lazy_bitmaps(); } + mtmd_tokenizer(mtmd_context * ctx, + const mtmd_input_part ** input_parts, + size_t n_parts, + bool add_special) : ctx(ctx) { + this->add_special = add_special; + parse_special = true; // only used for text returned by lazy bitmaps + vocab = ctx->vocab; + + for (size_t i = 0; i < n_parts; i++) { + const mtmd_input_part * p = input_parts[i]; + if (p->text != nullptr) { + parts.push_back({std::string(p->text->text, p->text->text_len), nullptr, p->text->parse_special}); + } else { + parts.push_back({"", p->bitmap}); + } + } + + expand_lazy_bitmaps(); + } + void expand_lazy_bitmaps() { std::vector expanded; expanded.reserve(parts.size()); @@ -1201,7 +1221,7 @@ struct mtmd_tokenizer { LOG_DBG("%s: lazy callback returned bitmap with dimensions %d x %d\n", __func__, out_bm->nx, out_bm->ny); } else if (out_str) { auto & ptr = text_from_lazy.emplace_back(out_str); // remember to free it later - expanded.push_back({ptr, nullptr}); + expanded.push_back({ptr, nullptr, parse_special}); LOG_DBG("%s: lazy callback returned text: %s\n", __func__, out_str); } } else if (res == -1) { @@ -1245,7 +1265,7 @@ struct mtmd_tokenizer { return res; } } else { - add_text(p.text, parse_special); + add_text(p.text, p.parse_special); } } @@ -1727,6 +1747,30 @@ int32_t mtmd_tokenize(mtmd_context * ctx, } } +int32_t mtmd_tokenize_from_parts(mtmd_context * ctx, + mtmd_input_chunks * output, + const mtmd_input_part ** parts, + size_t n_parts, + bool add_special) { + for (size_t i = 0; i < n_parts; i++) { + if ((parts[i]->text == nullptr) == (parts[i]->bitmap == nullptr)) { + LOG_ERR("%s: part %zu must have either text or bitmap set, not both\n", __func__, i); + return 1; + } + if (parts[i]->text != nullptr && parts[i]->text->text == nullptr) { + LOG_ERR("%s: part %zu has null text pointer\n", __func__, i); + return 1; + } + } + try { + mtmd_tokenizer tokenizer(ctx, parts, n_parts, add_special); + return tokenizer.tokenize(output); + } catch (const std::exception & e) { + LOG_ERR("%s: error: %s\n", __func__, e.what()); + return 2; + } +} + static int32_t mtmd_encode_impl(mtmd_context * ctx, const mtmd_image_tokens * image_tokens, std::vector & out_embd) { clip_ctx * ctx_clip = ctx->ctx_v; if (!ctx_clip) { diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index ef88efd316..0c2f9886e1 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -73,6 +73,12 @@ struct mtmd_input_text { bool parse_special; }; +struct mtmd_input_part { + // only text or bitmap can be set, not both + const struct mtmd_input_text * text; + const struct mtmd_bitmap * bitmap; +}; + // // C API // @@ -83,6 +89,7 @@ typedef struct mtmd_image_tokens mtmd_image_tokens; typedef struct mtmd_input_chunk mtmd_input_chunk; typedef struct mtmd_input_chunks mtmd_input_chunks; typedef struct mtmd_input_text mtmd_input_text; +typedef struct mtmd_input_part mtmd_input_part; typedef struct mtmd_batch mtmd_batch; typedef bool (*mtmd_progress_callback)(float progress, void * user_data); @@ -276,10 +283,10 @@ struct mtmd_decoder_pos { // return relative position (for example, embedding 0 will have position (0, 0, 0); remember to adjust it to the current absolute position) MTMD_API struct mtmd_decoder_pos mtmd_image_tokens_get_decoder_pos(const mtmd_image_tokens * image_tokens, llama_pos pos_0, size_t i); -// tokenize an input text prompt and a list of bitmaps (images/audio) -// the prompt must have the input image marker (default: "<__media__>") in it +// tokenize an input text prompt and a list of bitmaps (image/audio) +// the prompt must have the input media marker (default: "<__media__>") in it // the default marker is defined by mtmd_default_marker() -// the marker will be replaced with the image/audio chunk +// the marker will be replaced with the media chunk // for example: // "here is an image: <__media__>\ndescribe it in detail." // this will gives 3 chunks: @@ -291,13 +298,25 @@ MTMD_API struct mtmd_decoder_pos mtmd_image_tokens_get_decoder_pos(const mtmd_im // return values: // 0 on success // 1 on number of bitmaps not matching the number of markers -// 2 on image preprocessing error +// 2 on media preprocessing error MTMD_API int32_t mtmd_tokenize(mtmd_context * ctx, mtmd_input_chunks * output, const mtmd_input_text * text, const mtmd_bitmap ** bitmaps, size_t n_bitmaps); +// same as mtmd_tokenize(), but takes an array of mtmd_input_part +// use cases: +// - when you don't want to use media markers (they will be tokenized as normal text) +// - when you want to control parse_special for each text part +// note: per-part add_special will be ignored +// return 1 if a part has both text and bitmap set (or neither) +MTMD_API int32_t mtmd_tokenize_from_parts(mtmd_context * ctx, + mtmd_input_chunks * output, + const mtmd_input_part ** parts, + size_t n_parts, + bool add_special); + DEPRECATED(MTMD_API int32_t mtmd_encode(mtmd_context * ctx, const mtmd_image_tokens * image_tokens), "use mtmd_encode_chunk() instead");