feat: add multimodal search and granular performance diagnostics
Core CI / core (ubuntu-latest) (push) Failing after 3m31s
Core CI / core (macos-latest) (push) Canceled after 0s
Core CI / core (windows-latest) (push) Canceled after 0s
Core CI / desktop (ubuntu-latest) (push) Successful in 5m39s
Core CI / desktop (macos-latest) (push) Canceled after 0s
Core CI / desktop (windows-latest) (push) Canceled after 0s

This commit is contained in:
karl committed 2026-10-12 19:16:23 +02:00
1 parent d3bb8c7618
commit b8400c4fbf
110 files changed
+79395 -250

No files matched your search

@@ -41,11 +41,12 @@ set(LLAMA_BUILD_TOOLS OFF CACHE BOOL "No runtime tools" FORCE)
set(LLAMA_BUILD_SERVER OFF CACHE BOOL "No HTTP server" FORCE)
set(LLAMA_BUILD_EXAMPLES OFF CACHE BOOL "No examples" FORCE)
set(LLAMA_BUILD_TESTS OFF CACHE BOOL "No upstream tests in worker build" FORCE)
set(LLAMA_BUILD_MTMD OFF CACHE BOOL "Text worker; media encoder follows separately" FORCE)
set(LLAMA_BUILD_MTMD ON CACHE BOOL "Google image/audio projector" FORCE)
set(MTMD_VIDEO OFF CACHE BOOL "Frames decoded by bounded media adapter" FORCE)
set(LLAMA_OPENSSL OFF CACHE BOOL "No network runtime dependencies" FORCE)
add_subdirectory("${LLAMA_CPP_DIR}" llama-build)
add_executable(searchi-gemma-worker worker.cpp)
target_include_directories(searchi-gemma-worker PRIVATE "${LLAMA_CPP_DIR}/vendor")
target_link_libraries(searchi-gemma-worker PRIVATE llama)
target_link_libraries(searchi-gemma-worker PRIVATE llama mtmd)
target_compile_definitions(searchi-gemma-worker PRIVATE SEARCHI_LLAMA_TAG="b11461-no-logits-v1")
install(TARGETS searchi-gemma-worker RUNTIME DESTINATION bin)
+108 -6
View File
@@ -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;