feat: add multimodal search and granular performance diagnostics

crates/searchi-semantic/native/worker.cpp

@@ -1,11 +1,14 @@
 // Local NDJSON helper around the pinned llama.cpp encoder. stdout is exclusively
 // protocol output; llama diagnostics go to stderr. No model download or HTTP.
 #include "llama.h"
+#include "mtmd.h"
 #include "nlohmann/json.hpp"
 #include <algorithm>
 #include <cmath>
+#include <chrono>
 #include <cstdlib>
 #include <cstdio>
+#include <cstring>
 #include <iostream>
 #include <memory>
 #include <stdexcept>
@@ -13,7 +16,7 @@
 #include <vector>
 
 using json = nlohmann::json;
-constexpr size_t max_request_bytes = 2 * 1024 * 1024;
+constexpr size_t max_request_bytes = 16 * 1024 * 1024;
 constexpr size_t max_passage_bytes = 1024 * 1024;
 constexpr size_t max_text_bytes = 65536;
 constexpr size_t max_windows = 256;
@@ -34,12 +37,14 @@ class Encoder {
     llama_context * context = nullptr;
     llama_batch_ext * batch = nullptr;
     const llama_vocab * vocab = nullptr;
+    mtmd_context * media = nullptr;
     std::string model_path;
     void release() {
+        if (media) mtmd_free(media);
         if (batch) llama_batch_ext_free(batch);
         if (context) llama_free(context);
         if (model) llama_model_free(model);
-        batch = nullptr; context = nullptr; model = nullptr; vocab = nullptr;
+        batch = nullptr; context = nullptr; model = nullptr; vocab = nullptr; media = nullptr;
     }
     bool open(bool gpu) {
         auto mp = llama_model_default_params();
@@ -68,6 +73,9 @@ public:
         llama_log_set([](ggml_log_level level, const char * text, void *) {
             if (level >= GGML_LOG_LEVEL_WARN) std::fputs(text,stderr);
         },nullptr);
+        mtmd_log_set([](ggml_log_level level, const char * text, void *) {
+            if (level >= GGML_LOG_LEVEL_WARN) std::fputs(text,stderr);
+        },nullptr);
         llama_backend_init();
         bool use_gpu = want_gpu && llama_supports_gpu_offload();
         require(open(use_gpu) || (use_gpu && open(false)),"Cannot initialize local EmbeddingGemma 2 model/context");
@@ -77,6 +85,52 @@ public:
         llama_backend_free();
     }
     bool healthy() const {return context != nullptr && batch != nullptr;}
+    void ensure_media(const std::string & projector) {
+        if (media) return;
+        auto params = mtmd_context_params_default();
+        params.use_gpu = device == "gpu"; params.n_threads = threads;
+        params.print_timings = false; params.warmup = false;
+        params.image_min_tokens = 1; params.image_max_tokens = 1120;
+        media = mtmd_init_from_file(projector.c_str(),model,params);
+        require_inference(media != nullptr,"Cannot initialize Google media projector");
+    }
+    std::vector<float> embed_bitmap(mtmd_bitmap * bitmap) {
+        std::unique_ptr<mtmd_input_chunks,decltype(&mtmd_input_chunks_free)> chunks(mtmd_input_chunks_init(),mtmd_input_chunks_free);
+        require_inference(chunks != nullptr,"Cannot allocate media chunks");
+        const char * marker = mtmd_default_marker();
+        mtmd_input_text text{marker,std::strlen(marker),true,true};
+        const mtmd_bitmap * bitmaps[] = {bitmap};
+        require_inference(mtmd_tokenize(media,chunks.get(),&text,bitmaps,1) == 0,"Media tokenization failed");
+        llama_batch_ext_clear(batch); size_t position = 0;
+        auto place = [&](int index) {
+            require_inference(index >= 0 && position < context_tokens,"Media batch exceeds context");
+            llama_pos pos[4] = {llama_pos(position++),0,0,0};
+            require_inference(llama_batch_ext_set_pos(batch,index,pos) && llama_batch_ext_set_output_embd(batch,index,true),"Cannot place media row");
+        };
+        for (size_t i=0;i<mtmd_input_chunks_size(chunks.get());++i) {
+            const auto * chunk = mtmd_input_chunks_get(chunks.get(),i);
+            if (mtmd_input_chunk_get_type(chunk) == MTMD_INPUT_CHUNK_TYPE_TEXT) {
+                size_t count=0; const auto * tokens=mtmd_input_chunk_get_tokens_text(chunk,&count);
+                for(size_t j=0;j<count;++j) place(llama_batch_ext_add_token(batch,0,tokens[j]));
+            } else {
+                require_inference(mtmd_encode_chunk(media,chunk)==0,"Google media encoder failed");
+                const float * rows=mtmd_get_output_embd(media);
+                require_inference(rows != nullptr,"No projected media rows");
+                size_t width=llama_model_n_embd_inp(model);
+                for(size_t j=0;j<mtmd_input_chunk_get_n_tokens(chunk);++j) {
+                    llama_embd embd{rows+j*width,1,width};
+                    place(llama_batch_ext_add_embd(batch,0,embd));
+                }
+            }
+        }
+        require_inference(llama_process(context,LLAMA_PROCESS_TYPE_ENCODE,batch)==0,"Media text encoder failed");
+        const float * pooled=llama_get_embeddings_seq(context,0);
+        require_inference(pooled != nullptr,"No pooled media vector");
+        double norm=0; for(int i=0;i<dimensions;++i){require_inference(std::isfinite(pooled[i]),"Nonfinite media vector");norm+=double(pooled[i])*pooled[i];}
+        require_inference(norm>0 && std::isfinite(norm),"Invalid media vector norm");
+        std::vector<float> vector(dimensions);for(int i=0;i<dimensions;++i)vector[i]=float(pooled[i]/std::sqrt(norm));
+        return vector;
+    }
     void retry_cpu() {
         require(device == "gpu","Only a GPU inference failure can request a CPU retry");
         release();
@@ -103,10 +157,21 @@ public:
         require(title.size() <= 512, "Document title exceeds 512 UTF-8 bytes");
         return "title: " + (title.empty() ? "none" : title) + " | text: ";
     }
-    std::vector<float> embed(const std::string & text, int cap) {
+    std::vector<float> embed(const std::string & text, int cap, json * timings = nullptr) {
+        using clock = std::chrono::steady_clock;
+        auto started = timings ? clock::now() : clock::time_point{};
+        auto mark = [&](const char * name) {
+            if (timings) {
+                auto now = clock::now();
+                (*timings)[name] = std::chrono::duration_cast<std::chrono::microseconds>(now-started).count();
+                started = now;
+            }
+        };
         require(!text.empty() && text.size() <= max_text_bytes, "Embedding text must be nonempty and at most 64 KiB");
         auto tokens = tokenize(text,true);
         require(!tokens.empty() && tokens.size() <= size_t(cap), "Prompt exceeds model token budget; split document windows first");
+        if (timings) (*timings)["token_count"] = tokens.size();
+        mark("tokenization_us");
         llama_batch_ext_clear(batch);
         for (size_t i = 0; i < tokens.size(); ++i) {
             int index = llama_batch_ext_add_token(batch,0,tokens[i]);
@@ -115,7 +180,9 @@ public:
             require(llama_batch_ext_set_pos(batch,index,position), "Cannot set encoder position");
             require(llama_batch_ext_set_output_embd(batch,index,true), "Cannot select embedding output");
         }
+        mark("batch_prepare_us");
         require_inference(llama_process(context,LLAMA_PROCESS_TYPE_ENCODE,batch) == 0, "llama_process encoder failed");
+        mark("inference_us");
         const float * pooled = llama_get_embeddings_seq(context,0);
         require_inference(pooled != nullptr, "Encoder returned no pooled vector");
         double norm = 0;
@@ -123,6 +190,7 @@ public:
         require_inference(norm > 0 && std::isfinite(norm), "Zero or nonfinite embedding norm");
         std::vector<float> vector(dimensions);
         for (int i = 0; i < dimensions; ++i) vector[i] = float(pooled[i]/std::sqrt(norm));
+        mark("pool_normalize_us");
         return vector;
     }
     json windows(const std::string & text, const std::string & title) const {
@@ -157,8 +225,42 @@ public:
     }
     json handle(const json & request) {
         auto op = request.at("op").get<std::string>();
+        if(op=="image" || op=="audio") {
+            auto encoded=request.at("data").get<std::string>();
+            require(encoded.size()<=12000000 && encoded.size()%4==0,"Invalid media payload length");
+            static const std::string alphabet="ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
+            std::vector<unsigned char> bytes; bytes.reserve(encoded.size()/4*3);
+            unsigned accumulator=0;int bits=0;bool padding=false;
+            for(char c:encoded){if(c=='='){padding=true;continue;} require(!padding,"Invalid base64 padding");auto n=alphabet.find(c);require(n!=std::string::npos,"Invalid base64");accumulator=(accumulator<<6)|unsigned(n);bits+=6;if(bits>=8){bits-=8;bytes.push_back(static_cast<unsigned char>(accumulator>>bits));}}
+            auto projector=request.at("projector").get<std::string>();
+            require(!projector.empty() && projector.size()<32768,"Invalid projector path");
+            ensure_media(projector);
+            mtmd_bitmap * raw=nullptr;
+            if(op=="image") {
+                int w=request.at("width"),h=request.at("height");
+                require(w>=48 && h>=48 && w%48==0 && h%48==0 && uint64_t(w)*h<=1120*48*48,"Invalid bounded image dimensions");
+                require(bytes.size()==size_t(w)*h*3,"RGB image length mismatch");
+                raw=mtmd_bitmap_init(w,h,bytes.data());
+            } else {
+                require(bytes.size()%4==0 && !bytes.empty() && bytes.size()/4<=30*16000,"Audio must be at most30seconds mono16kHz PCM");
+                std::vector<float> pcm(bytes.size()/4);
+                std::memcpy(pcm.data(),bytes.data(),bytes.size());
+                for(float f:pcm)require(std::isfinite(f) && std::abs(f)<=1,"Audio PCM outside normalized range");
+                raw=mtmd_bitmap_init_from_audio(pcm.size(),pcm.data());
+            }
+            std::unique_ptr<mtmd_bitmap,decltype(&mtmd_bitmap_free)> bitmap(raw,mtmd_bitmap_free);
+            require_inference(bitmap != nullptr,"Cannot allocate media bitmap");
+            return {{"vectors",json::array({embed_bitmap(bitmap.get())})}};
+        }
         if (op == "windows") return {{"windows",windows(request.at("text").get<std::string>(),request.value("title",std::string()))}};
-        if (op == "query") return {{"vectors",json::array({embed("task: search result | query: "+request.at("text").get<std::string>(),512)})}};
+        if (op == "query") {
+            bool debug = request.value("debug_performance",false);
+            json timings;
+            auto vector = embed("task: search result | query: "+request.at("text").get<std::string>(),512,debug ? &timings : nullptr);
+            json response = {{"vectors",json::array({vector})}};
+            if (debug) response["performance"] = timings;
+            return response;
+        }
         if (op == "documents") {
             const auto & texts = request.at("texts");
             require(texts.is_array() && !texts.empty() && texts.size() <= 8,"Batch size must be between 1 and 8");
@@ -187,11 +289,11 @@ int main(int argc, char ** argv) {
         bool force_cpu = std::getenv("SEARCHI_FORCE_CPU") && std::string(std::getenv("SEARCHI_FORCE_CPU")) == "1";
         Encoder encoder(path,threads,device == "auto" && !force_cpu);
         std::cout << json({{"ready",true},{"model","google/embeddinggemma-2"},{"dimensions",dimensions},
-            {"device",encoder.device},{"threads",threads},{"runtime",SEARCHI_LLAMA_TAG}}).dump() << '\n' << std::flush;
+            {"device",encoder.device},{"threads",threads},{"runtime",SEARCHI_LLAMA_TAG},{"media_supported",true}}).dump() << '\n' << std::flush;
         for (;;) {
             std::string line; char c;
             while (std::cin.get(c) && c != '\n') {
-                require(line.size() < max_request_bytes,"Request exceeds 2 MiB"); line.push_back(c);
+                require(line.size() < max_request_bytes,"Request exceeds 16 MiB"); line.push_back(c);
             }
             if (!std::cin && line.empty()) break;
             json response;