karl/searchi
1
0 Fork 0
Code Issues Pull Requests Packages Projects Releases 10 Wiki Activity

feat: integrate Google semantic search and branded desktop launcher
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

Browse Source
This commit is contained in:
karl committed 2026-10-12 17:44:56 +02:00
1 parent a90df15503
commit caf6a2f82e
80 files changed
+14971 -837
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/searchi-semantic/Cargo.toml
+9 -3
View File
@@ -3,13 +3,19 @@ name = "searchi-semantic"
version.workspace = true
edition.workspace = true
license.workspace = true
description = "Bounded offline CPU ONNX embedding adapter for Searchi"
description = "Bounded local Google EmbeddingGemma 2 worker adapter for Searchi"
[features]
default = []
minilm-proof = ["dep:ort", "dep:tokenizers"]
[dependencies]
anyhow.workspace = true
ort = { version = "=2.0.0-rc.13", default-features = false, features = ["std", "load-dynamic", "api-27"] }
ort = { version = "=2.0.0-rc.13", optional = true, default-features = false, features = ["std", "load-dynamic", "api-27"] }
sha2 = "0.10"
tokenizers = { version = "=0.22.2", default-features = false, features = ["fancy-regex"] }
tokenizers = { version = "=0.22.2", optional = true, default-features = false, features = ["fancy-regex"] }
serde.workspace = true
serde_json.workspace = true
[dev-dependencies]
serde_json.workspace = true
crates/searchi-semantic/native/CMakeLists.txt
+51
View File
@@ -0,0 +1,51 @@
cmake_minimum_required(VERSION 3.20)
project(searchi_gemma_worker LANGUAGES C CXX)
if(NOT DEFINED LLAMA_CPP_DIR OR NOT EXISTS "${LLAMA_CPP_DIR}/include/llama.h")
message(FATAL_ERROR "Supply -DLLAMA_CPP_DIR=<pinned patched llama.cpp b11461 source>")
endif()
set(CMAKE_CXX_STANDARD 17)
if(MSVC)
set(CMAKE_MSVC_RUNTIME_LIBRARY "MultiThreaded$<$<CONFIG:Debug>:Debug>" CACHE STRING "Standalone worker static MSVC runtime" FORCE)
endif()
set(BUILD_SHARED_LIBS OFF CACHE BOOL "Static local worker runtime" FORCE)
set(GGML_OPENMP OFF CACHE BOOL "Use bounded native llama thread pools without external OpenMP runtime" FORCE)
set(GGML_NATIVE OFF CACHE BOOL "Portable CPU kernels")
# GPU backends are explicit production build options. The Windows VM builds
# with all of these OFF; runtime device selection always retains CPU fallback.
set(GGML_CUDA OFF CACHE BOOL "Production CUDA backend")
set(GGML_VULKAN OFF CACHE BOOL "Production Vulkan backend")
set(GGML_METAL OFF CACHE BOOL "Production Metal backend")
set(SEARCHI_ACCELERATION "cpu" CACHE STRING "cpu, auto, metal, vulkan, or cuda")
if(SEARCHI_ACCELERATION STREQUAL "cpu")
set(GGML_CUDA OFF CACHE BOOL "CPU test build" FORCE)
set(GGML_VULKAN OFF CACHE BOOL "CPU test build" FORCE)
set(GGML_METAL OFF CACHE BOOL "CPU test build" FORCE)
elseif(SEARCHI_ACCELERATION STREQUAL "metal" OR (SEARCHI_ACCELERATION STREQUAL "auto" AND APPLE))
set(GGML_METAL ON CACHE BOOL "Production Metal backend" FORCE)
elseif(SEARCHI_ACCELERATION STREQUAL "cuda")
set(GGML_CUDA ON CACHE BOOL "Production CUDA backend" FORCE)
elseif(SEARCHI_ACCELERATION STREQUAL "vulkan")
set(GGML_VULKAN ON CACHE BOOL "Production Vulkan backend" FORCE)
elseif(SEARCHI_ACCELERATION STREQUAL "auto")
find_package(Vulkan QUIET)
if(Vulkan_FOUND)
set(GGML_VULKAN ON CACHE BOOL "Available production Vulkan backend" FORCE)
else()
message(STATUS "No production acceleration SDK detected; retaining CPU backend")
endif()
else()
message(FATAL_ERROR "Unknown SEARCHI_ACCELERATION=${SEARCHI_ACCELERATION}")
endif()
set(LLAMA_BUILD_COMMON OFF CACHE BOOL "No extra common runtime" FORCE)
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_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_compile_definitions(searchi-gemma-worker PRIVATE SEARCHI_LLAMA_TAG="b11461-no-logits-v1")
install(TARGETS searchi-gemma-worker RUNTIME DESTINATION bin)
crates/searchi-semantic/native/worker.cpp
+221
View File
@@ -0,0 +1,221 @@
// 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;
}
}
crates/searchi-semantic/src/gemma.rs
+398
View File
@@ -0,0 +1,398 @@
use crate::TextWindow;
use anyhow::{ensure, Context, Result};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::{
fs::File,
io::{BufRead, BufReader, Read, Write},
path::{Path, PathBuf},
process::{Child, ChildStdin, Command, Stdio},
sync::mpsc::{self, Receiver},
thread,
time::Duration,
};
pub const GEMMA_DIMENSIONS: usize = 768;
pub const GEMMA_FINGERPRINT: &str = "google/embeddinggemma-2:bfcd298762cc34d0357ece5ebdd31791a3a374d8:q8_0:llama-b11461-no-logits-v1:mean:l2:query512:doc2048:windows-overlap128:prompts-v1";
const MODEL_FILE: &str = "embeddinggemma-2-Q8_0.gguf";
const MODEL_BYTES: u64 = 309855456;
const MODEL_SHA256: &str = "2188ac1deca4b77dffefd603c2776a9d76d9d74ec01841392982ebb840b09135";
const MAX_REQUEST: usize = 2 * 1024 * 1024;
const MAX_RESPONSE: usize = 4 * 1024 * 1024;
#[derive(Debug, Clone, Copy)]
pub enum Device {
Auto,
Cpu,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct WorkerInfo {
pub ready: bool,
pub model: String,
pub dimensions: usize,
pub device: String,
pub threads: usize,
pub runtime: String,
}
/// Persistent isolated native worker. No HTTP or remote model fallback.
pub struct GemmaEmbedder {
child: Child,
stdin: ChildStdin,
responses: Receiver<std::result::Result<Value, String>>,
pub info: WorkerInfo,
healthy: bool,
force_cpu: bool,
model_dir: PathBuf,
worker_path: PathBuf,
}
impl Drop for GemmaEmbedder {
fn drop(&mut self) {
let _ = self.child.kill();
let _ = self.child.wait();
}
}
impl GemmaEmbedder {
/// Load the verified Google model and explicitly supplied local worker.
/// Production selects available acceleration with CPU fallback. Tests set
/// `SEARCHI_FORCE_CPU=1` or call `load_with_device(..., Device::Cpu)`.
pub fn load(
model_dir: impl AsRef<Path>,
worker: impl AsRef<Path>,
threads: usize,
) -> Result<Self> {
Self::load_with_device(model_dir, worker, threads, Device::Auto)
}
pub fn load_with_device(
model_dir: impl AsRef<Path>,
worker: impl AsRef<Path>,
threads: usize,
device: Device,
) -> Result<Self> {
ensure!(
(1..=8).contains(&threads),
"CPU threads must be between 1 and 8"
);
let model_dir = model_dir.as_ref().to_path_buf();
let model = model_dir.join(MODEL_FILE);
verify_model(&model)?;
let worker = worker
.as_ref()
.canonicalize()
.context("Cannot resolve local Gemma worker executable")?;
ensure!(
worker.is_file(),
"Gemma worker must be a local executable file"
);
let force_cpu = std::env::var("SEARCHI_FORCE_CPU").is_ok_and(|value| value == "1");
let cpu = force_cpu || matches!(device, Device::Cpu);
let mut command = Command::new(&worker);
command
.args(["--model"])
.arg(model)
.args([
"--threads",
&threads.to_string(),
"--device",
if cpu { "cpu" } else { "auto" },
])
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::inherit());
// Background helpers must not create a visible console on Windows.
#[cfg(windows)]
{
use std::os::windows::process::CommandExt;
command.creation_flags(0x08000000); // CREATE_NO_WINDOW
}
let mut child = command.spawn().context("Cannot start local Gemma worker")?;
let stdin = child.stdin.take().context("Worker stdin missing")?;
let stdout = child.stdout.take().context("Worker stdout missing")?;
let (sender, responses) = mpsc::sync_channel(1);
thread::spawn(move || {
let mut reader = BufReader::new(stdout);
loop {
let response = read_response(&mut reader);
let failed = response.is_err();
if sender.send(response).is_err() || failed {
break;
}
}
});
let ready = match responses.recv_timeout(Duration::from_secs(120)) {
Ok(Ok(value)) => value,
other => {
let _ = child.kill();
let _ = child.wait();
anyhow::bail!("Gemma worker did not initialize: {other:?}");
}
};
let info: WorkerInfo = match serde_json::from_value(ready) {
Ok(info) => info,
Err(error) => {
let _ = child.kill();
let _ = child.wait();
return Err(error).context("Invalid Gemma worker handshake");
}
};
if !info.ready
|| info.model != "google/embeddinggemma-2"
|| info.dimensions != GEMMA_DIMENSIONS
|| info.runtime != "b11461-no-logits-v1"
|| info.threads != threads
|| !["cpu", "gpu"].contains(&info.device.as_str())
|| (cpu && info.device != "cpu")
{
let _ = child.kill();
let _ = child.wait();
anyhow::bail!("Gemma worker handshake violates configured model/device contract");
}
Ok(Self {
child,
stdin,
responses,
info,
healthy: true,
force_cpu: cpu,
model_dir,
worker_path: worker,
})
}
fn request(&mut self, request: Value) -> Result<Value> {
ensure!(
self.healthy,
"Gemma worker session failed; reopen it explicitly"
);
let encoded = serde_json::to_vec(&request)?;
ensure!(
encoded.len() < MAX_REQUEST,
"Gemma request exceeds 2 MiB protocol budget"
);
let response = match self.exchange(&encoded) {
Ok(response) => response,
Err(error) if can_retry_cpu(&self.info.device, self.force_cpu) => {
// GPU runtimes can abort rather than return an encoder error.
// Restart once as CPU and retry this same bounded request.
let _ = self.child.wait();
let replacement = Self::load_with_device(
&self.model_dir,
&self.worker_path,
self.info.threads,
Device::Cpu,
)
.with_context(|| {
format!("CPU worker restart failed after GPU transport error: {error}")
})?;
*self = replacement;
self.exchange(&encoded)?
}
Err(error) => return Err(error),
};
if let Some(device) = response["device"].as_str() {
ensure!(
["cpu", "gpu"].contains(&device) && (!self.force_cpu || device == "cpu"),
"Worker response violates configured device policy"
);
self.info.device = device.into();
}
if response["fatal"] == true {
self.healthy = false;
self.info.ready = false;
let _ = self.child.kill();
}
ensure!(
response["ok"] == true,
"Gemma worker error: {}",
response["error"].as_str().unwrap_or("unknown error")
);
Ok(response)
}
fn exchange(&mut self, encoded: &[u8]) -> Result<Value> {
let sent = (|| {
self.stdin.write_all(encoded)?;
self.stdin.write_all(b"\n")?;
self.stdin.flush()
})()
.context("Cannot write local Gemma worker request");
if let Err(error) = sent {
self.healthy = false;
self.info.ready = false;
let _ = self.child.kill();
return Err(error);
}
let response = match self.responses.recv_timeout(Duration::from_secs(180)) {
Ok(Ok(value)) => value,
other => {
self.healthy = false;
self.info.ready = false;
let _ = self.child.kill();
anyhow::bail!("Gemma worker request failed: {other:?}");
}
};
Ok(response)
}
pub fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>> {
ensure!(
!text.trim().is_empty() && text.len() <= 1024 * 1024,
"Passage must be nonempty and at most 1 MiB"
);
let response = self.request(json!({"op":"windows","text":text,"title":title}))?;
let windows: Vec<TextWindow> = serde_json::from_value(response["windows"].clone())?;
ensure!(
!windows.is_empty() && windows.len() <= 256,
"Invalid Gemma window count"
);
for window in &windows {
ensure!(
text.get(window.start_byte..window.end_byte) == Some(window.text.as_str()),
"Gemma window does not match original source range"
);
}
ensure!(
windows
.last()
.is_some_and(|window| window.end_byte == text.len()),
"Gemma windows did not cover the passage tail"
);
Ok(windows)
}
pub fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>> {
ensure!(
!texts.is_empty() && texts.len() <= 8,
"Document batch size must be between 1 and 8"
);
ensure!(
texts
.iter()
.all(|text| !text.trim().is_empty() && text.len() <= 65536),
"Document window must be nonempty and at most 64 KiB"
);
let response = self.request(json!({"op":"documents","texts":texts,"title":title}))?;
extract_vectors(response, texts.len())
}
pub fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
ensure!(
!query.trim().is_empty() && query.len() <= 4096,
"Query must be nonempty and at most 4096 UTF-8 bytes"
);
let response = self.request(json!({"op":"query","text":query}))?;
Ok(extract_vectors(response, 1)?.remove(0))
}
}
fn can_retry_cpu(actual_device: &str, force_cpu: bool) -> bool {
actual_device == "gpu" && !force_cpu
}
fn read_response(reader: &mut impl BufRead) -> std::result::Result<Value, String> {
let mut line = Vec::new();
match reader
.take((MAX_RESPONSE + 1) as u64)
.read_until(b'\n', &mut line)
{
Ok(0) => Err("Gemma worker closed its output".into()),
Ok(_) if line.len() > MAX_RESPONSE || line.last() != Some(&b'\n') => {
Err("Gemma worker response exceeds protocol byte budget".into())
}
Ok(_) => serde_json::from_slice(&line)
.map_err(|error| format!("Invalid Gemma worker JSON: {error}")),
Err(error) => Err(format!("Gemma worker read failed: {error}")),
}
}
fn extract_vectors(response: Value, count: usize) -> Result<Vec<Vec<f32>>> {
let vectors: Vec<Vec<f32>> = serde_json::from_value(response["vectors"].clone())?;
ensure!(
vectors.len() == count,
"Gemma returned incorrect embedding batch length"
);
for vector in &vectors {
ensure!(
vector.len() == GEMMA_DIMENSIONS && vector.iter().all(|value| value.is_finite()),
"Invalid Gemma embedding shape or values"
);
let norm = vector.iter().map(|v| v * v).sum::<f32>().sqrt();
ensure!(
(norm - 1.0).abs() < 1e-4,
"Gemma embedding is not L2 normalized"
);
}
Ok(vectors)
}
fn verify_model(path: &Path) -> Result<()> {
let file = File::open(path).with_context(|| {
format!(
"Missing Google EmbeddingGemma 2 model {}; run explicit model setup first",
path.display()
)
})?;
ensure!(
file.metadata()?.len() == MODEL_BYTES,
"EmbeddingGemma 2 model size mismatch"
);
let mut reader = BufReader::new(file);
let mut hash = Sha256::new();
let mut buffer = [0; 65536];
loop {
let count = reader.read(&mut buffer)?;
if count == 0 {
break;
}
hash.update(&buffer[..count]);
}
ensure!(
format!("{:x}", hash.finalize()) == MODEL_SHA256,
"EmbeddingGemma 2 model SHA256 mismatch"
);
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn missing_google_model_fails_without_starting_worker() {
assert!(GemmaEmbedder::load("missing-google-model", "missing-worker", 1).is_err());
}
#[test]
fn retry_policy_never_enables_gpu_or_retries_cpu_as_gpu() {
assert!(can_retry_cpu("gpu", false));
assert!(!can_retry_cpu("gpu", true));
assert!(!can_retry_cpu("cpu", false));
assert!(!can_retry_cpu("cpu", true));
}
#[test]
fn protocol_rejects_oversized_truncated_and_malformed_frames() {
assert!(read_response(&mut std::io::Cursor::new(vec![b'x'; MAX_RESPONSE + 1])).is_err());
assert!(read_response(&mut std::io::Cursor::new(b"{\"ok\":true}")).is_err());
assert!(read_response(&mut std::io::Cursor::new(b"not-json\n")).is_err());
assert_eq!(
read_response(&mut std::io::Cursor::new(b"{\"ok\":true}\n")).unwrap()["ok"],
true
);
}
#[test]
fn protocol_rejects_wrong_embedding_shape_count_and_normalization() {
assert!(extract_vectors(json!({"vectors":[]}), 1).is_err());
assert!(extract_vectors(json!({"vectors":[[1.0,0.0]]}), 1).is_err());
assert!(extract_vectors(json!({"vectors":[vec![0.0;GEMMA_DIMENSIONS]]}), 1).is_err());
let mut invalid = vec![0.0; GEMMA_DIMENSIONS];
invalid[0] = 1.0;
invalid[1] = 1.0;
assert!(extract_vectors(json!({"vectors":[invalid]}), 1).is_err());
}
}
crates/searchi-semantic/src/lib.rs
+17 -288
View File
@@ -1,290 +1,19 @@
//! Offline text embeddings using a pinned quantized multilingual model.
//!
//! Model provisioning is explicit and separate. This crate never downloads a
//! model or runtime and exposes no GPU execution-provider configuration.
//! Local Google EmbeddingGemma 2 inference through an isolated native worker.
//! The rejected MiniLM experiment is available only under `minilm-proof`.
mod gemma;
pub use gemma::{Device, GemmaEmbedder, WorkerInfo, GEMMA_DIMENSIONS, GEMMA_FINGERPRINT};
use std::{
fs::File,
io::{BufReader, Read},
path::{Path, PathBuf},
sync::Mutex,
/// A verified source byte range covered by a model token window.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TextWindow {
pub text: String,
pub start_byte: usize,
pub end_byte: usize,
}
#[cfg(feature = "minilm-proof")]
mod minilm;
#[cfg(feature = "minilm-proof")]
pub use minilm::{
CpuEmbedder, DIMENSIONS, MAX_BATCH, MAX_TEXT_BYTES, MAX_TOKENS, MODEL_FINGERPRINT,
};
use anyhow::{bail, ensure, Context, Result};
use ort::{ep, session::Session, value::Tensor};
use sha2::{Digest, Sha256};
use tokenizers::{PaddingDirection, PaddingParams, Tokenizer, TruncationParams};
/// Embedding dimensionality; stored vectors must carry the model fingerprint.
pub const DIMENSIONS: usize = 384;
pub const MAX_TOKENS: usize = 128;
pub const MAX_BATCH: usize = 8;
pub const MAX_TEXT_BYTES: usize = 65536;
/// Includes model revision, quantization, pooling, normalization and token cap.
pub const MODEL_FINGERPRINT: &str =
"minilm-multilingual:2c4055b12046f11709e9df2c122e59ffbdc2f900:int8:mean:l2:128";
const ARTIFACTS: [(&str, u64, &str); 2] = [
(
"onnx/model_quantized.onnx",
118308126,
"66fc00f5f29afcaff34092e1bdd20008ca3918265a82fb9695a551e510cc4ebc",
),
(
"tokenizer.json",
17082913,
"b60b6b43406a48bf3638526314f3d232d97058bc93472ff2de930d43686fa441",
),
];
// ONNX Runtime's library and environment are process-global. Prevent concurrent
// initializers from selecting different libraries or silently overriding state.
static RUNTIME_PATH: Mutex<Option<PathBuf>> = Mutex::new(None);
/// Local, mutable inference session. Call from a bounded background worker.
pub struct CpuEmbedder {
session: Session,
tokenizer: Tokenizer,
}
impl CpuEmbedder {
/// Verify local artifacts, then load the explicitly supplied runtime.
///
/// The runtime must be initialized through this adapter before other `ort`
/// code in this process. Threads are limited to 1–8; default caller policy
/// should use one thread in a VM. Missing or damaged assets are errors.
pub fn load(
model_dir: impl AsRef<Path>,
runtime_library: impl AsRef<Path>,
threads: usize,
) -> Result<Self> {
ensure!(
(1..=8).contains(&threads),
"CPU threads must be between 1 and 8"
);
let model_dir = model_dir.as_ref();
verify_artifacts(model_dir)?;
let runtime_path = runtime_library
.as_ref()
.canonicalize()
.context("Cannot resolve explicitly supplied ONNX Runtime library")?;
ensure!(
runtime_path.is_file(),
"ONNX Runtime library must be a file"
);
{
let mut initialized = RUNTIME_PATH
.lock()
.map_err(|_| anyhow::anyhow!("Runtime initialization lock poisoned"))?;
if let Some(existing) = initialized.as_ref() {
ensure!(
*existing == runtime_path,
"ONNX Runtime already initialized with a different library"
);
} else {
let committed = ort::init_from(&runtime_path)?
.with_telemetry(false)
.with_execution_providers([ep::CPU::default().build()])
.commit();
ensure!(
committed,
"ONNX Runtime environment was already initialized outside the CPU adapter"
);
*initialized = Some(runtime_path);
}
}
let session = Session::builder()?
.with_execution_providers([ep::CPU::default().build()])
.map_err(|e| anyhow::anyhow!("CPU session configuration: {e}"))?
.with_intra_threads(threads)
.map_err(|e| anyhow::anyhow!("CPU thread configuration: {e}"))?
.with_inter_threads(1)
.map_err(|e| anyhow::anyhow!("Inter-thread configuration: {e}"))?
.with_parallel_execution(false)
.map_err(|e| anyhow::anyhow!("Sequential session configuration: {e}"))?
.commit_from_file(model_dir.join("onnx/model_quantized.onnx"))?;
let mut tokenizer = Tokenizer::from_file(model_dir.join("tokenizer.json"))
.map_err(|error| anyhow::anyhow!("Cannot load tokenizer: {error}"))?;
tokenizer
.with_truncation(Some(TruncationParams {
max_length: MAX_TOKENS,
..Default::default()
}))
.map_err(|error| anyhow::anyhow!("Cannot configure tokenizer: {error}"))?;
tokenizer.with_padding(Some(PaddingParams {
pad_id: 1,
pad_token: "<pad>".into(),
..Default::default()
}));
Ok(Self { session, tokenizer })
}
/// Produce finite, L2-normalized vectors from genuine local model inference.
///
/// Inputs above 128 tokens are truncated. Production indexing should chunk
/// documents first and keep chunk locations; never silently embed entire PDFs.
pub fn encode(&mut self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
validate_texts(texts)?;
// Encode sequentially so tokenizers cannot create an unbounded global
// Rayon worker pool independently of the caller's inference budget.
let mut encoded = texts
.iter()
.map(|text| self.tokenizer.encode(*text, true))
.collect::<std::result::Result<Vec<_>, _>>()
.map_err(|error| anyhow::anyhow!("Tokenization failed: {error}"))?;
let sequence_len = encoded.iter().map(|item| item.len()).max().unwrap_or(0);
for item in &mut encoded {
item.pad(sequence_len, 1, 0, "<pad>", PaddingDirection::Right);
}
ensure!(
sequence_len > 0 && sequence_len <= MAX_TOKENS,
"Invalid token sequence length"
);
let ids: Vec<i64> = encoded
.iter()
.flat_map(|item| item.get_ids().iter().map(|&v| i64::from(v)))
.collect();
let masks: Vec<i64> = encoded
.iter()
.flat_map(|item| item.get_attention_mask().iter().map(|&v| i64::from(v)))
.collect();
let types: Vec<i64> = encoded
.iter()
.flat_map(|item| item.get_type_ids().iter().map(|&v| i64::from(v)))
.collect();
let shape = [texts.len(), sequence_len];
let outputs = self.session.run(ort::inputs![
"input_ids" => Tensor::from_array((shape, ids))?,
"attention_mask" => Tensor::from_array((shape, masks.clone()))?,
"token_type_ids" => Tensor::from_array((shape, types))?
])?;
let (output_shape, hidden) = outputs[0].try_extract_tensor::<f32>()?;
ensure!(
**output_shape == [texts.len() as i64, sequence_len as i64, DIMENSIONS as i64],
"Unexpected embedding output shape"
);
pool_and_normalize(hidden, &masks, texts.len(), sequence_len)
}
}
fn validate_texts(texts: &[&str]) -> Result<()> {
ensure!(
!texts.is_empty() && texts.len() <= MAX_BATCH,
"Batch size must be between 1 and {MAX_BATCH}"
);
ensure!(
texts
.iter()
.all(|text| !text.trim().is_empty() && text.len() <= MAX_TEXT_BYTES),
"Text must be nonempty and at most {MAX_TEXT_BYTES} UTF-8 bytes"
);
Ok(())
}
fn verify_artifacts(model_dir: &Path) -> Result<()> {
for (relative, bytes, expected) in ARTIFACTS {
let path = model_dir.join(relative);
let file = File::open(&path).with_context(|| {
format!(
"Missing local model artifact {}; run the explicit model downloader first",
path.display()
)
})?;
ensure!(
file.metadata()?.len() == bytes,
"Model artifact size mismatch: {}",
path.display()
);
let mut reader = BufReader::new(file);
let mut hash = Sha256::new();
let mut buffer = [0_u8; 65536];
loop {
let read = reader.read(&mut buffer)?;
if read == 0 {
break;
}
hash.update(&buffer[..read]);
}
let actual = format!("{:x}", hash.finalize());
ensure!(
actual == expected,
"Model artifact SHA256 mismatch: {}",
path.display()
);
}
Ok(())
}
fn pool_and_normalize(
hidden: &[f32],
masks: &[i64],
batch: usize,
tokens: usize,
) -> Result<Vec<Vec<f32>>> {
ensure!(
hidden.len() == batch * tokens * DIMENSIONS && masks.len() == batch * tokens,
"Invalid pooling input dimensions"
);
let mut vectors = Vec::with_capacity(batch);
for item in 0..batch {
let mut vector = vec![0_f32; DIMENSIONS];
let mut count = 0;
for token in 0..tokens {
if masks[item * tokens + token] == 0 {
continue;
}
count += 1;
let offset = (item * tokens + token) * DIMENSIONS;
for dimension in 0..DIMENSIONS {
vector[dimension] += hidden[offset + dimension];
}
}
ensure!(count > 0, "Cannot pool a fully masked sequence");
for value in &mut vector {
*value /= count as f32;
}
let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
if !norm.is_finite() || norm <= f32::EPSILON {
bail!("Model produced nonfinite or zero embeddings");
}
for value in &mut vector {
*value /= norm;
}
ensure!(
vector.iter().all(|value| value.is_finite()),
"Model produced nonfinite embeddings"
);
vectors.push(vector);
}
Ok(vectors)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn padding_is_excluded_and_output_is_normalized() {
let mut hidden = vec![0.0; 2 * DIMENSIONS];
hidden[0] = 3.0;
hidden[1] = 4.0;
hidden[DIMENSIONS] = 999.0;
let pooled = pool_and_normalize(&hidden, &[1, 0], 1, 2).unwrap();
assert!((pooled[0][0] - 0.6).abs() < 1e-6);
assert!((pooled[0][1] - 0.8).abs() < 1e-6);
}
#[test]
fn rejects_unbounded_inputs_and_invalid_outputs() {
assert!(validate_texts(&[]).is_err());
assert!(validate_texts(&["x"; MAX_BATCH + 1]).is_err());
assert!(validate_texts(&[" "]).is_err());
assert!(validate_texts(&[&"x".repeat(MAX_TEXT_BYTES + 1)]).is_err());
assert!(pool_and_normalize(&vec![0.0; DIMENSIONS], &[0], 1, 1).is_err());
assert!(pool_and_normalize(&vec![f32::NAN; DIMENSIONS], &[1], 1, 1).is_err());
}
#[test]
fn missing_model_does_not_load_or_download_runtime() {
assert!(CpuEmbedder::load("this-model-does-not-exist", "missing-runtime", 1).is_err());
}
}
crates/searchi-semantic/src/minilm.rs
+347
View File
@@ -0,0 +1,347 @@
//! Offline text embeddings using a pinned quantized multilingual model.
//!
//! Model provisioning is explicit and separate. This crate never downloads a
//! model or runtime and exposes no GPU execution-provider configuration.
use std::{
fs::File,
io::{BufReader, Read},
path::{Path, PathBuf},
sync::Mutex,
};
use anyhow::{bail, ensure, Context, Result};
use ort::{ep, session::Session, value::Tensor};
use sha2::{Digest, Sha256};
use tokenizers::{PaddingDirection, PaddingParams, Tokenizer, TruncationParams};
/// Embedding dimensionality; stored vectors must carry the model fingerprint.
pub const DIMENSIONS: usize = 384;
pub const MAX_TOKENS: usize = 128;
pub const MAX_BATCH: usize = 8;
pub const MAX_TEXT_BYTES: usize = 65536;
pub const MAX_PASSAGE_BYTES: usize = 1024 * 1024;
pub const MAX_WINDOWS: usize = 256;
pub const WINDOW_OVERLAP_TOKENS: usize = 24;
/// A source byte range covered by one overlapping model token window.
use crate::TextWindow;
/// Includes model revision, quantization, pooling, normalization and token cap.
pub const MODEL_FINGERPRINT: &str =
"minilm-multilingual:2c4055b12046f11709e9df2c122e59ffbdc2f900:int8:mean:l2:128";
const ARTIFACTS: [(&str, u64, &str); 2] = [
(
"onnx/model_quantized.onnx",
118308126,
"66fc00f5f29afcaff34092e1bdd20008ca3918265a82fb9695a551e510cc4ebc",
),
(
"tokenizer.json",
17082913,
"b60b6b43406a48bf3638526314f3d232d97058bc93472ff2de930d43686fa441",
),
];
// ONNX Runtime's library and environment are process-global. Prevent concurrent
// initializers from selecting different libraries or silently overriding state.
static RUNTIME_PATH: Mutex<Option<PathBuf>> = Mutex::new(None);
/// Local, mutable inference session. Call from a bounded background worker.
pub struct CpuEmbedder {
session: Session,
tokenizer: Tokenizer,
}
impl CpuEmbedder {
/// Split a bounded passage into overlapping token windows, retaining its tail.
/// Special tokens have no source range; source offsets remain UTF-8 byte
/// offsets into the original passage for accurate preview highlighting.
pub fn split_windows(&self, text: &str) -> Result<Vec<TextWindow>> {
ensure!(
!text.trim().is_empty() && text.len() <= MAX_PASSAGE_BYTES,
"Passage must be nonempty and at most {MAX_PASSAGE_BYTES} UTF-8 bytes"
);
let mut tokenizer = self.tokenizer.clone();
tokenizer
.with_truncation(Some(TruncationParams {
max_length: MAX_TOKENS,
stride: WINDOW_OVERLAP_TOKENS,
..Default::default()
}))
.map_err(|error| anyhow::anyhow!("Window tokenizer configuration: {error}"))?;
tokenizer.with_padding(None);
let encoded = tokenizer
.encode(text, true)
.map_err(|error| anyhow::anyhow!("Window tokenization failed: {error}"))?;
ensure!(
encoded.get_overflowing().len() < MAX_WINDOWS,
"Passage exceeds {MAX_WINDOWS} model windows; split it before embedding"
);
let mut windows = Vec::new();
for item in std::iter::once(&encoded).chain(encoded.get_overflowing()) {
let offsets: Vec<_> = item
.get_offsets()
.iter()
.zip(item.get_special_tokens_mask())
.filter(|((start, end), special)| **special == 0 && end > start)
.map(|(offset, _)| *offset)
.collect();
let start_byte = offsets.iter().map(|(start, _)| *start).min().unwrap_or(0);
let end_byte = offsets.iter().map(|(_, end)| *end).max().unwrap_or(0);
ensure!(start_byte < end_byte, "Token window has no source text");
let chunk = text
.get(start_byte..end_byte)
.context("Invalid UTF-8 tokenizer offsets")?;
ensure!(
chunk.len() <= MAX_TEXT_BYTES,
"Token window exceeds embedding text byte budget"
);
windows.push(TextWindow {
text: chunk.to_owned(),
start_byte,
end_byte,
});
}
Ok(windows)
}
/// Verify local artifacts, then load the explicitly supplied runtime.
///
/// The runtime must be initialized through this adapter before other `ort`
/// code in this process. Threads are limited to 1–8; default caller policy
/// should use one thread in a VM. Missing or damaged assets are errors.
pub fn load(
model_dir: impl AsRef<Path>,
runtime_library: impl AsRef<Path>,
threads: usize,
) -> Result<Self> {
ensure!(
(1..=8).contains(&threads),
"CPU threads must be between 1 and 8"
);
let model_dir = model_dir.as_ref();
verify_artifacts(model_dir)?;
let runtime_path = runtime_library
.as_ref()
.canonicalize()
.context("Cannot resolve explicitly supplied ONNX Runtime library")?;
ensure!(
runtime_path.is_file(),
"ONNX Runtime library must be a file"
);
{
let mut initialized = RUNTIME_PATH
.lock()
.map_err(|_| anyhow::anyhow!("Runtime initialization lock poisoned"))?;
if let Some(existing) = initialized.as_ref() {
ensure!(
*existing == runtime_path,
"ONNX Runtime already initialized with a different library"
);
} else {
let committed = ort::init_from(&runtime_path)?
.with_telemetry(false)
.with_execution_providers([ep::CPU::default().build()])
.commit();
ensure!(
committed,
"ONNX Runtime environment was already initialized outside the CPU adapter"
);
*initialized = Some(runtime_path);
}
}
let session = Session::builder()?
.with_execution_providers([ep::CPU::default().build()])
.map_err(|e| anyhow::anyhow!("CPU session configuration: {e}"))?
.with_intra_threads(threads)
.map_err(|e| anyhow::anyhow!("CPU thread configuration: {e}"))?
.with_inter_threads(1)
.map_err(|e| anyhow::anyhow!("Inter-thread configuration: {e}"))?
.with_parallel_execution(false)
.map_err(|e| anyhow::anyhow!("Sequential session configuration: {e}"))?
.commit_from_file(model_dir.join("onnx/model_quantized.onnx"))?;
let mut tokenizer = Tokenizer::from_file(model_dir.join("tokenizer.json"))
.map_err(|error| anyhow::anyhow!("Cannot load tokenizer: {error}"))?;
tokenizer
.with_truncation(Some(TruncationParams {
max_length: MAX_TOKENS,
..Default::default()
}))
.map_err(|error| anyhow::anyhow!("Cannot configure tokenizer: {error}"))?;
tokenizer.with_padding(Some(PaddingParams {
pad_id: 1,
pad_token: "<pad>".into(),
..Default::default()
}));
Ok(Self { session, tokenizer })
}
/// Produce finite, L2-normalized vectors from genuine local model inference.
///
/// Inputs above 128 tokens are truncated. Production indexing should chunk
/// documents first and keep chunk locations; never silently embed entire PDFs.
pub fn encode(&mut self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
validate_texts(texts)?;
// Encode sequentially so tokenizers cannot create an unbounded global
// Rayon worker pool independently of the caller's inference budget.
let mut encoded = texts
.iter()
.map(|text| self.tokenizer.encode(*text, true))
.collect::<std::result::Result<Vec<_>, _>>()
.map_err(|error| anyhow::anyhow!("Tokenization failed: {error}"))?;
let sequence_len = encoded.iter().map(|item| item.len()).max().unwrap_or(0);
for item in &mut encoded {
item.pad(sequence_len, 1, 0, "<pad>", PaddingDirection::Right);
}
ensure!(
sequence_len > 0 && sequence_len <= MAX_TOKENS,
"Invalid token sequence length"
);
let ids: Vec<i64> = encoded
.iter()
.flat_map(|item| item.get_ids().iter().map(|&v| i64::from(v)))
.collect();
let masks: Vec<i64> = encoded
.iter()
.flat_map(|item| item.get_attention_mask().iter().map(|&v| i64::from(v)))
.collect();
let types: Vec<i64> = encoded
.iter()
.flat_map(|item| item.get_type_ids().iter().map(|&v| i64::from(v)))
.collect();
let shape = [texts.len(), sequence_len];
let outputs = self.session.run(ort::inputs![
"input_ids" => Tensor::from_array((shape, ids))?,
"attention_mask" => Tensor::from_array((shape, masks.clone()))?,
"token_type_ids" => Tensor::from_array((shape, types))?
])?;
let (output_shape, hidden) = outputs[0].try_extract_tensor::<f32>()?;
ensure!(
**output_shape == [texts.len() as i64, sequence_len as i64, DIMENSIONS as i64],
"Unexpected embedding output shape"
);
pool_and_normalize(hidden, &masks, texts.len(), sequence_len)
}
}
fn validate_texts(texts: &[&str]) -> Result<()> {
ensure!(
!texts.is_empty() && texts.len() <= MAX_BATCH,
"Batch size must be between 1 and {MAX_BATCH}"
);
ensure!(
texts
.iter()
.all(|text| !text.trim().is_empty() && text.len() <= MAX_TEXT_BYTES),
"Text must be nonempty and at most {MAX_TEXT_BYTES} UTF-8 bytes"
);
Ok(())
}
fn verify_artifacts(model_dir: &Path) -> Result<()> {
for (relative, bytes, expected) in ARTIFACTS {
let path = model_dir.join(relative);
let file = File::open(&path).with_context(|| {
format!(
"Missing local model artifact {}; run the explicit model downloader first",
path.display()
)
})?;
ensure!(
file.metadata()?.len() == bytes,
"Model artifact size mismatch: {}",
path.display()
);
let mut reader = BufReader::new(file);
let mut hash = Sha256::new();
let mut buffer = [0_u8; 65536];
loop {
let read = reader.read(&mut buffer)?;
if read == 0 {
break;
}
hash.update(&buffer[..read]);
}
let actual = format!("{:x}", hash.finalize());
ensure!(
actual == expected,
"Model artifact SHA256 mismatch: {}",
path.display()
);
}
Ok(())
}
fn pool_and_normalize(
hidden: &[f32],
masks: &[i64],
batch: usize,
tokens: usize,
) -> Result<Vec<Vec<f32>>> {
ensure!(
hidden.len() == batch * tokens * DIMENSIONS && masks.len() == batch * tokens,
"Invalid pooling input dimensions"
);
let mut vectors = Vec::with_capacity(batch);
for item in 0..batch {
let mut vector = vec![0_f32; DIMENSIONS];
let mut count = 0;
for token in 0..tokens {
if masks[item * tokens + token] == 0 {
continue;
}
count += 1;
let offset = (item * tokens + token) * DIMENSIONS;
for dimension in 0..DIMENSIONS {
vector[dimension] += hidden[offset + dimension];
}
}
ensure!(count > 0, "Cannot pool a fully masked sequence");
for value in &mut vector {
*value /= count as f32;
}
let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
if !norm.is_finite() || norm <= f32::EPSILON {
bail!("Model produced nonfinite or zero embeddings");
}
for value in &mut vector {
*value /= norm;
}
ensure!(
vector.iter().all(|value| value.is_finite()),
"Model produced nonfinite embeddings"
);
vectors.push(vector);
}
Ok(vectors)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn padding_is_excluded_and_output_is_normalized() {
let mut hidden = vec![0.0; 2 * DIMENSIONS];
hidden[0] = 3.0;
hidden[1] = 4.0;
hidden[DIMENSIONS] = 999.0;
let pooled = pool_and_normalize(&hidden, &[1, 0], 1, 2).unwrap();
assert!((pooled[0][0] - 0.6).abs() < 1e-6);
assert!((pooled[0][1] - 0.8).abs() < 1e-6);
}
#[test]
fn rejects_unbounded_inputs_and_invalid_outputs() {
assert!(validate_texts(&[]).is_err());
assert!(validate_texts(&["x"; MAX_BATCH + 1]).is_err());
assert!(validate_texts(&[" "]).is_err());
assert!(validate_texts(&[&"x".repeat(MAX_TEXT_BYTES + 1)]).is_err());
assert!(pool_and_normalize(&vec![0.0; DIMENSIONS], &[0], 1, 1).is_err());
assert!(pool_and_normalize(&vec![f32::NAN; DIMENSIONS], &[1], 1, 1).is_err());
}
#[test]
fn missing_model_does_not_load_or_download_runtime() {
assert!(CpuEmbedder::load("this-model-does-not-exist", "missing-runtime", 1).is_err());
}
}
crates/searchi-semantic/tests/cpu_model.rs
+1
View File
@@ -1,4 +1,5 @@
//! Opt-in real-artifact test. Default tests never acquire model/runtime assets.
#![cfg(feature = "minilm-proof")]
use searchi_semantic::{CpuEmbedder, DIMENSIONS};
use std::path::PathBuf;
crates/searchi-semantic/tests/fixtures/LICENSE-DigUp.txt
Vendored
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2026 Abdur Rahim
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
crates/searchi-semantic/tests/fixtures/digup-text-reference.json
Vendored
+3858
View File
File diff suppressed because it is too large. Load diff
crates/searchi-semantic/tests/gemma_cpu.rs
+102
View File
@@ -0,0 +1,102 @@
use anyhow::Result;
use searchi_semantic::{Device, GemmaEmbedder, GEMMA_DIMENSIONS};
#[test]
#[ignore = "Requires explicitly provisioned Google model and CPU worker"]
fn google_cpu_real_embeddings_and_verified_tail_windows() -> Result<()> {
let model = std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Set SEARCHI_GEMMA_MODEL_DIR");
let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Set SEARCHI_GEMMA_WORKER");
let mut embedder = GemmaEmbedder::load_with_device(model, worker, 1, Device::Cpu)?;
assert_eq!(embedder.info.device, "cpu");
assert_eq!(embedder.info.dimensions, GEMMA_DIMENSIONS);
let reference: serde_json::Value =
serde_json::from_str(include_str!("fixtures/digup-text-reference.json"))?;
let long_reference=(1..=60).map(|i|format!("Note {i}: the review moved to room {} because the projector in the main hall failed again, and the team agreed to send the slides the night before.",i%7+2)).collect::<Vec<_>>().join(" ");
let mut made = Vec::new();
for (name, query) in [
("query", "payment declined error"),
("query-bn", "ইলিশ মাছ রান্না"),
("query-ar", "متى تغلق المكتبة"),
] {
made.push((name, embedder.embed_query(query)?));
}
made.push(("document",embedder.embed_documents(&["Flight to Tokyo Haneda on November 3. Then the train to Kamakura to see the Great Buddha, and back for dinner in Shibuya."],"Tokyo itinerary")?.remove(0)));
made.push((
"document-long",
embedder.embed_documents(&[&long_reference], "")?.remove(0),
));
for (name, vector) in made {
let expected = reference["vectors"][name].as_array().unwrap();
let cosine: f64 = vector
.iter()
.zip(expected)
.map(|(a, b)| f64::from(*a) * b.as_f64().unwrap())
.sum();
println!("DigUp text reference {name}: cosine {cosine:.8}");
assert!(
cosine >= 0.999,
"Google CPU reference parity failed for {name}: {cosine}"
);
}
let documents = [
"The supplier agreement may be ended with sixty days written notice.",
"Business train tickets and hotel stays can be reimbursed by submitting receipts.",
"Die Gartenplanung sieht ein Hochbeet für Gemüse und eine Regentonne vor.",
];
let vectors = embedder.embed_documents(&documents, "none")?;
for (query, expected) in [
("Wie beenden wir die Zusammenarbeit mit dem Lieferanten?", 0),
("How do I claim my accommodation and railway expenses?", 1),
(
"Where is the plan for the vegetable bed and rainwater barrel?",
2,
),
] {
let query = embedder.embed_query(query)?;
let scores: Vec<f32> = vectors
.iter()
.map(|vector| vector.iter().zip(&query).map(|(a, b)| a * b).sum())
.collect();
let best = scores
.iter()
.enumerate()
.max_by(|a, b| a.1.total_cmp(b.1))
.unwrap()
.0;
assert_eq!(
best, expected,
"Semantic cross-language ranking failed: {scores:?}"
);
}
let long = format!("{}\nHotel accommodation on business trips is reimbursed after submitting the overnight stay receipt.",
"The routine maintenance log records ordinary uptime checks and cable inspections. ".repeat(350));
let windows = embedder.split_windows(&long, "maintenance-log.txt")?;
assert!(windows.len() > 1);
assert_eq!(windows.last().unwrap().end_byte, long.len());
assert!(windows
.last()
.unwrap()
.text
.contains("overnight stay receipt"));
let texts: Vec<_> = windows.iter().map(|window| window.text.as_str()).collect();
let vectors = embedder.embed_documents(&texts, "maintenance-log.txt")?;
let query = embedder.embed_query(
"Where does the policy say that lodging costs during work travel can be claimed?",
)?;
let scores: Vec<f32> = vectors
.iter()
.map(|vector| vector.iter().zip(&query).map(|(a, b)| a * b).sum())
.collect();
let best = scores
.iter()
.enumerate()
.max_by(|a, b| a.1.total_cmp(b.1))
.unwrap()
.0;
assert_eq!(
best,
windows.len() - 1,
"Tail retrieval did not find final window: {scores:?}"
);
Ok(())
}