karl
caf6a2f82e
Core CI / core (ubuntu-latest) (push) Failing after 3m27s
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 5m45s
Core CI / desktop (macos-latest) (push) Canceled after 0s
Core CI / desktop (windows-latest) (push) Canceled after 0s
222 lines
12 KiB
C++
222 lines
12 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 "nlohmann/json.hpp"
|
||
#include <algorithm>
|
||
#include <cmath>
|
||
#include <cstdlib>
|
||
#include <cstdio>
|
||
#include <iostream>
|
||
#include <memory>
|
||
#include <stdexcept>
|
||
#include <string>
|
||
#include <vector>
|
||
|
||
using json = nlohmann::json;
|
||
constexpr size_t max_request_bytes = 2 * 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;
|
||
std::string model_path;
|
||
void release() {
|
||
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;
|
||
}
|
||
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);
|
||
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 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) {
|
||
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");
|
||
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");
|
||
}
|
||
require_inference(llama_process(context,LLAMA_PROCESS_TYPE_ENCODE,batch) == 0, "llama_process encoder failed");
|
||
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));
|
||
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");
|
||
// Preserve actual source offsets, rather than inventing locations from
|
||
// a lossy decoded representation. Unexpected normalization is an error.
|
||
require(decode(tokens,tokens.size()) == text, "Tokenizer roundtrip changed source text; cannot verify source offsets");
|
||
size_t prompt = tokenize(document_prefix(title),true).size();
|
||
require(prompt + 132 < size_t(context_tokens), "Document prefix exceeds token budget");
|
||
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 == "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 == "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}}).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);
|
||
}
|
||
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;
|
||
}
|
||
}
|