searchi/crates/searchi-semantic/native/worker.cpp

// 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>
#include <string>
#include <vector>

using json = nlohmann::json;
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;
constexpr int context_tokens = 2048;
constexpr int dimensions = 768;

static void require(bool condition, const char * message) {
    if (!condition) throw std::runtime_error(message);
}

struct InferenceError : std::runtime_error { using std::runtime_error::runtime_error; };
static void require_inference(bool condition, const char * message) {
    if (!condition) throw InferenceError(message);
}

class Encoder {
    llama_model * model = nullptr;
    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; media = nullptr;
    }
    bool open(bool gpu) {
        auto mp = llama_model_default_params();
        mp.n_gpu_layers = gpu ? 999 : 0;
        mp.progress_callback = [](float,void *) {return true;};
        model = llama_model_load_from_file(model_path.c_str(),mp);
        if (!model) return false;
        auto cp = llama_context_default_params();
        cp.n_ctx = context_tokens; cp.n_batch = context_tokens; cp.n_ubatch = context_tokens;
        cp.n_seq_max = 1; cp.kv_unified = true; cp.embeddings = true;
        cp.pooling_type = LLAMA_POOLING_TYPE_MEAN;
        cp.n_threads = threads; cp.n_threads_batch = threads;
        context = llama_init_from_model(model,cp);
        if (!context) {release(); return false;}
        require(llama_model_n_embd_out(model) == dimensions,"EmbeddingGemma 2 output must have 768 dimensions");
        vocab = llama_model_get_vocab(model);
        batch = llama_batch_ext_init(context);
        if (!batch) {release(); return false;}
        device = gpu ? "gpu" : "cpu";
        return true;
    }
public:
    std::string device;
    int threads;
    Encoder(const std::string & path, int cpu_threads, bool want_gpu) : model_path(path),threads(cpu_threads) {
        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");
    }
    ~Encoder() {
        release();
        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();
        require(open(false),"CPU fallback initialization failed after GPU inference failure");
    }
    std::vector<llama_token> tokenize(const std::string & text, bool specials) const {
        std::vector<llama_token> tokens(text.size() + 8);
        int count = llama_tokenize(vocab, text.data(), int(text.size()), tokens.data(), int(tokens.size()), specials, false);
        if (count < 0) {
            tokens.resize(size_t(-count));
            count = llama_tokenize(vocab, text.data(), int(text.size()), tokens.data(), int(tokens.size()), specials, false);
        }
        require(count >= 0, "Tokenizer failed");
        tokens.resize(size_t(count));
        return tokens;
    }
    std::string decode(const std::vector<llama_token> & tokens, size_t count) const {
        std::vector<char> text(max_passage_bytes + 64);
        int length = llama_detokenize(vocab, tokens.data(), int(count), text.data(), int(text.size()), false, false);
        require(length >= 0, "Detokenized passage exceeds byte budget");
        return std::string(text.data(),size_t(length));
    }
    static std::string document_prefix(const std::string & title) {
        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, 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]);
            require(index >= 0, "Embedding batch is full");
            llama_pos position[4] = {llama_pos(i),0,0,0};
            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;
        for (int i = 0; i < dimensions; ++i) { require_inference(std::isfinite(pooled[i]), "Nonfinite encoder output"); norm += double(pooled[i])*pooled[i]; }
        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 {
        require(!text.empty() && text.size() <= max_passage_bytes, "Passage exceeds 1 MiB or is empty");
        auto tokens = tokenize(text,false);
        require(!tokens.empty(), "Passage has no tokens");
        size_t prompt = tokenize(document_prefix(title),true).size();
        require(prompt + 132 < size_t(context_tokens), "Document prefix exceeds token budget");
        // SentencePiece may normalize source characters. In that case decoded
        // token lengths are not source offsets: split the original UTF-8 bytes
        // and verify each complete prefixed fragment against the real tokenizer.
        if (decode(tokens,tokens.size()) != text) {
            const auto prefix = document_prefix(title);
            json result = json::array();
            if (prefix.size()+text.size() <= max_text_bytes && tokenize(prefix+text,true).size() <= size_t(context_tokens)) {
                return json::array({{{"text",text},{"start_byte",0},{"end_byte",text.size()}}});
            }
            std::vector<size_t> boundaries;
            for (size_t i=0;i<text.size();++i) if ((uint8_t(text[i]) & 0xc0) != 0x80) boundaries.push_back(i);
            boundaries.push_back(text.size());
            size_t start_index = 0;
            while (start_index+1 < boundaries.size()) {
                const size_t start = boundaries[start_index];
                require(result.size() < max_windows, "Passage exceeds 256 windows");
                size_t low = start_index;
                size_t high = size_t(std::upper_bound(boundaries.begin(),boundaries.end(),std::min(text.size(),start+max_text_bytes-prefix.size()))-boundaries.begin())-1;
                size_t best_index = start_index;
                while (low < high) {
                    size_t candidate_index = low + (high-low+1)/2;
                    size_t candidate = boundaries[candidate_index];
                    if (tokenize(prefix+text.substr(start,candidate-start),true).size() <= size_t(context_tokens)-4) {
                        best_index = candidate_index;
                        low = candidate_index;
                    } else {
                        high = candidate_index-1;
                    }
                }
                size_t best = boundaries[best_index];
                require(best > start, "Cannot fit source fragment in model context");
                auto fragment = text.substr(start,best-start);
                require(tokenize(prefix+fragment,true).size() <= size_t(context_tokens), "Source fragment exceeds context budget");
                result.push_back({{"text",fragment},{"start_byte",start},{"end_byte",best}});
                if (best == text.size()) break;
                // Bounded original-byte overlap; always advance, even for
                // unusual Unicode/token expansion. No decoded text is indexed.
                size_t next = best - std::min(size_t(128),(best-start)/2);
                start_index = size_t(std::lower_bound(boundaries.begin(),boundaries.end(),next)-boundaries.begin());
            }
            return result;
        }
        size_t budget = size_t(context_tokens) - prompt - 4;
        size_t stride = budget - 128;
        json result = json::array();
        for (size_t start = 0; start < tokens.size(); start += stride) {
            require(result.size() < max_windows, "Passage exceeds 256 windows");
            size_t end = std::min(start + budget,tokens.size());
            size_t begin_byte = start ? decode(tokens,start).size() : 0;
            size_t end_byte = decode(tokens,end).size();
            // Token pieces can split a UTF-8 character. Enlarge to valid source
            // boundaries; the four-token reserve covers boundary re-tokenization.
            while (begin_byte > 0 && begin_byte < text.size() && (uint8_t(text[begin_byte]) & 0xc0) == 0x80) --begin_byte;
            while (end_byte < text.size() && (uint8_t(text[end_byte]) & 0xc0) == 0x80) ++end_byte;
            require(begin_byte < end_byte && end_byte <= text.size(), "Invalid source window byte range");
            auto fragment = text.substr(begin_byte,end_byte-begin_byte);
            require(fragment.size() + document_prefix(title).size() <= max_text_bytes, "Window exceeds text byte budget");
            require(tokenize(document_prefix(title)+fragment,true).size() <= size_t(context_tokens), "Window re-tokenization exceeds context budget");
            result.push_back({{"text",fragment},{"start_byte",begin_byte},{"end_byte",end_byte}});
            if (end == tokens.size()) break;
        }
        return result;
    }
    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" || op == "code_query") {
            bool debug = request.value("debug_performance",false);
            json timings;
            auto prefix = op == "code_query" ? "task: code retrieval | query: " : "task: search result | query: ";
            auto vector = embed(prefix+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");
            auto prefix = document_prefix(request.value("title",std::string()));
            json vectors = json::array();
            for (const auto & text : texts) vectors.push_back(embed(prefix+text.get<std::string>(),context_tokens));
            return {{"vectors",vectors}};
        }
        throw std::runtime_error("Unknown embedding operation");
    }
};

int main(int argc, char ** argv) {
    try {
        std::string path, device = "auto"; int threads = 1;
        for (int i = 1; i < argc; ++i) {
            require(i+1 < argc,"Worker options require values");
            std::string option = argv[i++], value = argv[i];
            if (option == "--model") path = value;
            else if (option == "--threads") threads = std::stoi(value);
            else if (option == "--device") device = value;
            else throw std::runtime_error("Unknown worker option");
        }
        require(!path.empty() && threads >= 1 && threads <= 8,"Model path required; threads must be 1–8");
        require(device == "cpu" || device == "auto","Device must be cpu or auto");
        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},{"media_supported",true},{"code_query_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 16 MiB"); line.push_back(c);
            }
            if (!std::cin && line.empty()) break;
            json response;
            try {
                auto request = json::parse(line);
                try {response = encoder.handle(request);}
                catch (const InferenceError &) {
                    if (encoder.device != "gpu") throw;
                    encoder.retry_cpu();
                    response = encoder.handle(request); // Exactly one CPU retry.
                    response["fallback"] = "gpu_inference_failure";
                }
                response["ok"] = true;
            }
            catch (const std::exception & error) { response = {{"ok",false},{"error",error.what()},{"fatal",!encoder.healthy()}}; }
            response["device"] = encoder.device;
            auto encoded = response.dump();
            require(encoded.size() <= 4*1024*1024,"Response exceeds 4 MiB");
            std::cout << encoded << '\n' << std::flush;
            if (!encoder.healthy()) return 1;
        }
        return 0;
    } catch (const std::exception & error) {
        std::cerr << "Searchi Gemma worker: " << error.what() << '\n';
        return 1;
    }
}