377 lines
21 KiB
C++
377 lines
21 KiB
C++
// 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();
|
||
static ggml_backend_dev_t cpu_devices[] = {nullptr};
|
||
if (!gpu) mp.devices = cpu_devices;
|
||
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.offload_kqv = gpu; cp.op_offload = gpu;
|
||
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";
|
||
// Disable backend registration before llama initializes its registry.
|
||
// A GPU-capable distribution can therefore be tested on CPU without
|
||
// probing/enabling the Vulkan provider in the Windows VM.
|
||
if (force_cpu || device == "cpu") {
|
||
#ifdef _WIN32
|
||
require(_putenv_s("GGML_DISABLE_VULKAN", "1") == 0, "Cannot disable Vulkan for CPU startup");
|
||
#else
|
||
require(setenv("GGML_DISABLE_VULKAN", "1", 1) == 0, "Cannot disable Vulkan for CPU startup");
|
||
#endif
|
||
}
|
||
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;
|
||
}
|
||
}
|