From 4575db17c2029ec3f4483c9a4718c59c9c3712ad Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Mon, 22 Jun 2026 14:03:15 -0400 Subject: [PATCH 1/2] feat(embed): embed CLI binary with --download-only model fetch Add the embed binary (gated required-features = native), used by Lattice Studio. The --download-only flag fetches an embedding model and emits an @@lattice download_done event consumed by the Models "Get Models" surface. Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/embed/Cargo.toml | 5 + crates/embed/src/bin/embed.rs | 201 ++++++++++++++++++++++++++++++++++ 2 files changed, 206 insertions(+) create mode 100644 crates/embed/src/bin/embed.rs diff --git a/crates/embed/Cargo.toml b/crates/embed/Cargo.toml index 2de98fb3a3..6cff4ada48 100644 --- a/crates/embed/Cargo.toml +++ b/crates/embed/Cargo.toml @@ -46,6 +46,11 @@ metal-gpu = ["lattice-inference/metal-gpu"] avx512 = [] # Requires nightly Rust (stdarch_x86_avx512) local = [] +[[bin]] +name = "embed" +path = "src/bin/embed.rs" +required-features = ["native"] + [[example]] name = "basic_embed" required-features = ["native"] diff --git a/crates/embed/src/bin/embed.rs b/crates/embed/src/bin/embed.rs new file mode 100644 index 0000000000..a801a16d36 --- /dev/null +++ b/crates/embed/src/bin/embed.rs @@ -0,0 +1,201 @@ +//! CLI tool for generating text embeddings using lattice-embed. +//! +//! # Usage +//! +//! ```text +//! embed --model bge-small-en-v1.5 --text "hello" --text "world" [--json] +//! ``` +//! +//! When `--json` is set, emits a single `@@lattice {"ev":"embed_done",...}` line +//! to stdout in addition to the human-readable summary. + +use std::process::ExitCode; +use std::str::FromStr; +use std::time::Instant; + +use lattice_embed::{EmbeddingModel, EmbeddingService, NativeEmbeddingService}; + +fn usage(msg: &str) -> ExitCode { + eprintln!("ERROR: {msg}\n"); + eprintln!("{USAGE}"); + ExitCode::FAILURE +} + +const USAGE: &str = "\ +usage: embed [--model ] --text [--text ...] [--json] + +Generate embeddings for one or more text strings. + +options: + --model Embedding model to use. Default: bge-small-en-v1.5 + Accepted: bge-small-en-v1.5, bge-base-en-v1.5, bge-large-en-v1.5, + multilingual-e5-small, multilingual-e5-base, all-minilm-l6-v2, + paraphrase-multilingual-minilm-l12-v2 + Also accepts HuggingFace IDs like BAAI/bge-small-en-v1.5. + --text Text to embed. Repeat for multiple texts. + --json Emit a structured @@lattice {\"ev\":\"embed_done\",...} line to stdout. + --download-only Ensure the model is downloaded and loadable, then exit (no --text needed). + Emits @@lattice {\"ev\":\"download_done\",\"ok\":bool} with --json. + -h, --help Print this help and exit. +"; + +#[tokio::main] +async fn main() -> ExitCode { + let args: Vec = std::env::args().collect(); + + let mut model_name: Option = None; + let mut texts: Vec = Vec::new(); + let mut emit_json: bool = false; + let mut download_only: bool = false; + + let mut i = 1; + while i < args.len() { + match args[i].as_str() { + "--model" => { + i += 1; + let Some(v) = args.get(i) else { + return usage("--model requires an argument"); + }; + model_name = Some(v.clone()); + } + "--text" => { + i += 1; + let Some(v) = args.get(i) else { + return usage("--text requires an argument"); + }; + texts.push(v.clone()); + } + "--json" => { + emit_json = true; + } + "--download-only" => { + download_only = true; + } + "--help" | "-h" => { + eprintln!("{USAGE}"); + return ExitCode::SUCCESS; + } + other => return usage(&format!("unknown argument: {other}")), + } + i += 1; + } + + if !download_only && texts.is_empty() { + return usage("at least one --text argument is required"); + } + + let model = match model_name { + Some(ref name) => match EmbeddingModel::from_str(name) { + Ok(m) => m, + Err(_) => { + return usage(&format!( + "--model '{name}' is not a recognised embedding model" + )); + } + }, + None => EmbeddingModel::default(), + }; + + eprintln!("Model: {model}"); + eprintln!("Dimensions: {}", model.dimensions()); + eprintln!("Texts: {}", texts.len()); + eprintln!(); + eprintln!("Generating embeddings (model loads on first call — may download ~130 MB)..."); + + let service = NativeEmbeddingService::with_model(model); + + // --download-only: ensure the model is present (downloading + checksum-verifying if + // needed) and loadable, then exit. A single warmup embedding forces the lazy + // ensure_model_files + model load without printing the similarity report. + if download_only { + match service.embed(&["warmup".to_string()], model).await { + Ok(_) => { + eprintln!("Model {model} is downloaded and ready."); + if emit_json { + let obj = serde_json::json!({ + "ev": "download_done", + "model": model.to_string(), + "ok": true, + }); + println!("@@lattice {obj}"); + } + return ExitCode::SUCCESS; + } + Err(err) => { + eprintln!("ERROR: model download/load failed: {err}"); + if emit_json { + let obj = serde_json::json!({ + "ev": "download_done", + "model": model.to_string(), + "ok": false, + "error": err.to_string(), + }); + println!("@@lattice {obj}"); + } + return ExitCode::FAILURE; + } + } + } + + let t0 = Instant::now(); + let embeddings = match service.embed(&texts, model).await { + Ok(e) => e, + Err(err) => { + eprintln!("ERROR: embedding failed: {err}"); + return ExitCode::FAILURE; + } + }; + let elapsed_ms = t0.elapsed().as_millis(); + + if embeddings.is_empty() { + eprintln!("ERROR: service returned zero embeddings"); + return ExitCode::FAILURE; + } + + let dims = embeddings[0].len(); + let count = embeddings.len(); + + // Build NxN pairwise cosine matrix. + let mut cosine: Vec> = Vec::with_capacity(count); + for i in 0..count { + let mut row = Vec::with_capacity(count); + for j in 0..count { + let sim = lattice_embed::utils::cosine_similarity(&embeddings[i], &embeddings[j]); + row.push(sim); + } + cosine.push(row); + } + + // Build preview: first 8 dims of each vector. + let preview_len = dims.min(8); + let preview: Vec> = embeddings + .iter() + .map(|e| e[..preview_len].to_vec()) + .collect(); + + eprintln!("=== Embedding Results ==="); + eprintln!("Dims: {dims}"); + eprintln!("Count: {count}"); + eprintln!("Elapsed: {elapsed_ms}ms"); + eprintln!(); + eprintln!("Pairwise cosine similarity:"); + for (i, row) in cosine.iter().enumerate() { + let vals: Vec = row.iter().map(|v| format!("{v:.4}")).collect(); + eprintln!(" [{i}] {}", vals.join(" ")); + } + + if emit_json { + let obj = serde_json::json!({ + "ev": "embed_done", + "model": model.to_string(), + "dims": dims, + "count": count, + "cosine": cosine, + "preview": preview, + "ms": elapsed_ms, + }); + println!("@@lattice {obj}"); + } + + ExitCode::SUCCESS +} From 63c6b3a61308b24e0167c45bb968a11bfde49685 Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Mon, 22 Jun 2026 15:37:28 -0400 Subject: [PATCH 2/2] fix(embed): --download-only skips the warmup embedding pass MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --download-only exists to fetch + checksum-verify + load a model without running inference, but it called service.embed(&["warmup"], model), which runs a full encode_batch after the load. Add NativeEmbeddingService:: ensure_loaded — it drives the same download + model-load path embed() does and returns without encoding — and call it from the download-only branch. The flag now stops at load, saving a forward pass; error handling and the JSON success event are unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/embed/src/bin/embed.rs | 7 +++---- crates/embed/src/service/native.rs | 13 +++++++++++++ 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/crates/embed/src/bin/embed.rs b/crates/embed/src/bin/embed.rs index a801a16d36..58d1f0b58e 100644 --- a/crates/embed/src/bin/embed.rs +++ b/crates/embed/src/bin/embed.rs @@ -105,11 +105,10 @@ async fn main() -> ExitCode { let service = NativeEmbeddingService::with_model(model); // --download-only: ensure the model is present (downloading + checksum-verifying if - // needed) and loadable, then exit. A single warmup embedding forces the lazy - // ensure_model_files + model load without printing the similarity report. + // needed) and loadable, then exit without running any encode pass. if download_only { - match service.embed(&["warmup".to_string()], model).await { - Ok(_) => { + match service.ensure_loaded().await { + Ok(()) => { eprintln!("Model {model} is downloaded and ready."); if emit_json { let obj = serde_json::json!({ diff --git a/crates/embed/src/service/native.rs b/crates/embed/src/service/native.rs index 1f2ece04f7..f552b85ebe 100644 --- a/crates/embed/src/service/native.rs +++ b/crates/embed/src/service/native.rs @@ -152,6 +152,19 @@ impl NativeEmbeddingService { .unwrap_or(0) } + /// **Unstable**: download and load the model without producing any embeddings. + /// + /// Performs the same download + checksum-verify + model-load sequence as the + /// first call to `embed`, then returns `Ok(())` without running an encode pass. + /// Intended for use by the `--download-only` CLI flag so that the model is + /// warmed into the file-system cache without wasting a forward pass. + /// + /// Errors are the same as those from `embed`: network failures, checksum + /// mismatches, and unsupported model variants surface as `EmbedError`. + pub async fn ensure_loaded(&self) -> Result<()> { + self.ensure_model().await.map(|_| ()) + } + /// Ensure the model is loaded (cancellation-safe). /// /// Uses `std::sync::OnceLock` so the model loading runs to completion