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

@@ -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);
 

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<mtmd_input_part> parts;
+        std::vector<mtmd_internal_part> parts;
         for (char c : spec) {
             if (c == 't') {
                 parts.push_back({ "hello", nullptr });

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

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<std::string> 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<mtmd_input_text> texts(segments.size());
+    std::vector<mtmd_input_part> 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<const mtmd_input_part *> 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;

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<std::vector<const mtmd_bitmap *>> mtmd_group_mergeable_bitmaps(std::vector<mtmd_input_part> & parts, int n_merge);
+std::vector<std::vector<const mtmd_bitmap *>> mtmd_group_mergeable_bitmaps(std::vector<mtmd_internal_part> & parts, int n_merge);

tools/mtmd/mtmd.cpp

@@ -1097,7 +1097,7 @@ void mtmd_free(mtmd_context * ctx) {
     delete ctx;
 }
 
-std::vector<std::vector<const mtmd_bitmap *>> mtmd_group_mergeable_bitmaps(std::vector<mtmd_input_part> & parts, int n_merge) {
+std::vector<std::vector<const mtmd_bitmap *>> mtmd_group_mergeable_bitmaps(std::vector<mtmd_internal_part> & parts, int n_merge) {
     std::vector<std::vector<const mtmd_bitmap *>> 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<part> parts;
     // these will be freed when mtmd_tokenizer finishes
     std::vector<mtmd::bitmap> 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<part> 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<float> & out_embd) {
     clip_ctx * ctx_clip = ctx->ctx_v;
     if (!ctx_clip) {

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");