diff --git a/Cargo.lock b/Cargo.lock index a5395bb..9b54221 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -57,6 +57,12 @@ version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" +[[package]] +name = "anyhow" +version = "1.0.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" + [[package]] name = "arbitrary" version = "1.4.2" @@ -120,6 +126,20 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bread-onnx" +version = "0.2.3" +dependencies = [ + "anyhow", + "bread-utils", + "hex", + "ort", + "sha2", + "tokenizers", + "tracing", + "ureq 2.12.1", +] + [[package]] name = "bread-theme" version = "0.2.3" @@ -131,10 +151,20 @@ dependencies = [ "serde_json", ] +[[package]] +name = "bread-utils" +version = "0.2.3" +dependencies = [ + "dirs", + "serde", + "serde_json", +] + [[package]] name = "breadmill" version = "0.2.4" dependencies = [ + "bread-onnx", "breadsearch-shared", "hex", "ignore", @@ -2966,9 +2996,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ "pin-project-lite", + "tracing-attributes", "tracing-core", ] +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tracing-core" version = "0.1.36" @@ -3107,6 +3149,8 @@ dependencies = [ "once_cell", "rustls", "rustls-pki-types", + "serde", + "serde_json", "url", "webpki-roots 0.26.11", ] diff --git a/breadmill/Cargo.toml b/breadmill/Cargo.toml index 768735d..929fd8d 100644 --- a/breadmill/Cargo.toml +++ b/breadmill/Cargo.toml @@ -42,6 +42,8 @@ breadsearch-shared = { path = "../breadsearch-shared" } # needed for the plain CPU path even in the npu build. ort = { version = "2.0.0-rc.12", default-features = false, features = ["std", "tracing", "download-binaries", "tls-native", "copy-dylibs", "api-23"] } tokenizers = "0" +# TODO(owner): switch to tag-pinned git dependency once bread-onnx is merged and tagged, matching the bread-theme pattern +bread-onnx = { path = "../../bread-ecosystem-fix-worktree/bread-onnx" } # Surfaces ort's own EP-registration tracing (e.g. a GPU EP silently failing to # register and falling back to CPU) as visible log output instead of nowhere. diff --git a/breadmill/src/embed.rs b/breadmill/src/embed.rs index 3020e58..d507951 100644 --- a/breadmill/src/embed.rs +++ b/breadmill/src/embed.rs @@ -1,10 +1,7 @@ use std::path::{Path, PathBuf}; -use ort::{ - session::{Session, builder::{GraphOptimizationLevel, SessionBuilder}}, - value::Tensor, -}; -use tokenizers::Tokenizer; +use bread_onnx::embedding::EmbeddingSession; +use bread_onnx::Provider; const DOCUMENT_PREFIX: &str = "search_document: "; const QUERY_PREFIX: &str = "search_query: "; @@ -35,23 +32,15 @@ pub enum Backend { } pub struct OrtEmbedder { - session: Session, - tokenizer: Tokenizer, - dim: usize, + inner: EmbeddingSession, } impl OrtEmbedder { pub fn load(model_path: &Path, tokenizer_path: &Path, dim: usize, backend: Backend) -> Result { - let builder = Session::builder() - .map_err(|e| e.to_string())? - .with_optimization_level(GraphOptimizationLevel::All) + let provider = to_provider(backend)?; + let inner = EmbeddingSession::load(model_path, tokenizer_path, dim, MAX_SEQ_LEN, &[provider]) .map_err(|e| e.to_string())?; - - let mut builder = configure_eps(builder, &backend)?; - let session = builder.commit_from_file(model_path).map_err(|e| e.to_string())?; - let tokenizer = Tokenizer::from_file(tokenizer_path).map_err(|e| e.to_string())?; - - Ok(Self { session, tokenizer, dim }) + Ok(Self { inner }) } pub fn embed_document(&mut self, text: &str) -> Result, String> { @@ -64,229 +53,71 @@ impl OrtEmbedder { fn embed_with_prefix(&mut self, text: &str, prefix: &str) -> Result, String> { let input = format!("{}{}", prefix, text); - - let encoding = self - .tokenizer - .encode(input, true) - .map_err(|e| e.to_string())?; - - let mut ids: Vec = encoding.get_ids().iter().map(|&x| x as i64).collect(); - let mut mask: Vec = encoding - .get_attention_mask() - .iter() - .map(|&x| x as i64) - .collect(); - let mut type_ids: Vec = encoding - .get_type_ids() - .iter() - .map(|&x| x as i64) - .collect(); - - if ids.len() > MAX_SEQ_LEN { - eprintln!( - "breadmill: truncating {} tokens to {} (chunk too large)", - ids.len(), - MAX_SEQ_LEN - ); - ids.truncate(MAX_SEQ_LEN); - mask.truncate(MAX_SEQ_LEN); - type_ids.truncate(MAX_SEQ_LEN); - } - - let seq_len = ids.len() as i64; - - let id_tensor = - Tensor::::from_array((vec![1i64, seq_len], ids.clone())).map_err(|e| e.to_string())?; - let mask_tensor = - Tensor::::from_array((vec![1i64, seq_len], mask.clone())).map_err(|e| e.to_string())?; - let type_tensor = - Tensor::::from_array((vec![1i64, seq_len], type_ids)).map_err(|e| e.to_string())?; - - let outputs = self - .session - .run(ort::inputs! { - "input_ids" => id_tensor, - "attention_mask" => mask_tensor, - "token_type_ids" => type_tensor, - }) - .map_err(|e| e.to_string())?; - - // last_hidden_state: shape [1, seq_len, dim] - let (shape, data) = outputs["last_hidden_state"] - .try_extract_tensor::() - .map_err(|e| e.to_string())?; - - let actual_seq = shape[1] as usize; - let actual_dim = shape[2] as usize; - - // Mean-pool over non-padding positions. Some execution providers (e.g. - // MIGraphX) pad the output sequence dimension for kernel efficiency, so - // actual_seq can exceed mask.len() — only positions covered by our own - // attention mask are meaningful, so cap the loop at whichever is shorter. - let mut result = vec![0.0f32; actual_dim]; - let mut count = 0usize; - - for t in 0..actual_seq.min(mask.len()) { - if mask[t] > 0 { - for d in 0..actual_dim { - result[d] += data[t * actual_dim + d]; - } - count += 1; - } - } - - if count > 0 { - for x in &mut result { - *x /= count as f32; - } - } - - l2_normalize(&mut result); - - // Clamp/pad to configured dim - result.truncate(self.dim); - while result.len() < self.dim { - result.push(0.0); - } - - Ok(result) + self.inner.embed(&input).map_err(|e| e.to_string()) } } -fn l2_normalize(v: &mut Vec) { - let norm: f32 = v.iter().map(|x| x * x).sum::().sqrt(); - if norm > 1e-10 { - for x in v.iter_mut() { - *x /= norm; - } - } -} +// ---- Backend -> bread_onnx::Provider ---------------------------------------- +// +// `bread_onnx::session::build_session` (via `EmbeddingSession::load`) is the +// shared session-builder + EP-fallback + loud-logging code every EP branch +// below used to hand-roll separately (`configure_eps`/`npu_session`/ +// `rocm_session`/`cuda_session`/`openvino_session`). What's genuinely +// specific to this crate — its cargo feature gates (npu/rocm/cuda/openvino), +// and NPU's `vaip_config.json` discovery — stays here. -// ---- Execution provider selection ------------------------------------------- - -fn configure_eps(builder: SessionBuilder, backend: &Backend) -> Result { +fn to_provider(backend: Backend) -> Result { match backend { - Backend::Cpu => Ok(builder), - Backend::Npu { cache_dir } => npu_session(builder, cache_dir), - Backend::Rocm => rocm_session(builder), - Backend::Cuda => cuda_session(builder), - Backend::OpenVino { cache_dir } => openvino_session(builder, cache_dir), + Backend::Cpu => Ok(Provider::Cpu), + Backend::Npu { cache_dir } => npu_provider(cache_dir), + Backend::Rocm => rocm_provider(), + Backend::Cuda => cuda_provider(), + Backend::OpenVino { cache_dir } => Ok(Provider::OpenVino { device_type: "GPU".to_string(), cache_dir }), } } #[cfg(feature = "npu")] -fn npu_session(builder: SessionBuilder, cache_dir: &Path) -> Result { - let vitis_ep = build_vitis_ep(cache_dir)?; - eprintln!("breadmill: using NPU (VitisAI) execution provider"); +fn npu_provider(cache_dir: PathBuf) -> Result { + let vaip_config = find_vaip_config()?; if std::env::var("ORT_DYLIB_PATH").is_err() { eprintln!( "breadmill: hint — set ORT_DYLIB_PATH to the Ryzen AI SDK ORT, e.g.:\n \ ORT_DYLIB_PATH=~/.local/share/ryzen-ai-1.7.1/lib/libonnxruntime.so" ); } - builder - .with_execution_providers([vitis_ep, ort::ep::CPU::default().build()]) - .map_err(|e| e.to_string()) + Ok(Provider::Vitis { + config_file: vaip_config, + cache_dir: cache_dir.join("npu"), + cache_key: "nomic-embed-text-v1.5".to_string(), + }) } #[cfg(not(feature = "npu"))] -fn npu_session(builder: SessionBuilder, _cache_dir: &Path) -> Result { +fn npu_provider(_cache_dir: PathBuf) -> Result { eprintln!("breadmill: NPU backend requested but not compiled in (rebuild with --features npu); using CPU"); - Ok(builder) + Ok(Provider::Cpu) } -// ---- VitisAI EP (NPU) ------------------------------------------------------- - -#[cfg(feature = "npu")] -fn build_vitis_ep(cache_dir: &Path) -> Result { - let vaip_config = find_vaip_config()?; - let npu_cache = cache_dir.join("npu"); - std::fs::create_dir_all(&npu_cache).map_err(|e| e.to_string())?; - eprintln!("breadmill: vaip_config: {}", vaip_config.display()); - eprintln!("breadmill: NPU model cache: {}", npu_cache.display()); - Ok(ort::ep::Vitis::default() - .with_config_file(vaip_config.to_string_lossy()) - .with_cache_dir(npu_cache.to_string_lossy()) - .with_cache_key("nomic-embed-text-v1.5") - .build()) -} - -// ---- MIGraphX EP (AMD iGPU, ROCm-backed) ------------------------------------- - #[cfg(feature = "rocm")] -fn rocm_session(builder: SessionBuilder) -> Result { - eprintln!("breadmill: using MIGraphX execution provider (device 0)"); - eprintln!( - "breadmill: note — check the log line above/below for \"Successfully registered \ - `MIGraphXExecutionProvider`\"; if it's missing, the ONNX Runtime in use wasn't built \ - with MIGraphX support and inference silently fell back to CPU" - ); - builder - .with_execution_providers([ - ort::ep::MIGraphX::default().with_device_id(0).build(), - ort::ep::CPU::default().build(), - ]) - .map_err(|e| e.to_string()) +fn rocm_provider() -> Result { + Ok(Provider::MiGraphX { device_id: 0 }) } #[cfg(not(feature = "rocm"))] -fn rocm_session(builder: SessionBuilder) -> Result { +fn rocm_provider() -> Result { eprintln!("breadmill: ROCm backend requested but not compiled in (rebuild with --features rocm); using CPU"); - Ok(builder) + Ok(Provider::Cpu) } -// ---- CUDA EP (NVIDIA GPU) ---------------------------------------------------- - #[cfg(feature = "cuda")] -fn cuda_session(builder: SessionBuilder) -> Result { - eprintln!("breadmill: using CUDA execution provider (device 0)"); - eprintln!( - "breadmill: note — check the log line above/below for \"Successfully registered \ - `CUDAExecutionProvider`\"; if it's missing, the ONNX Runtime in use wasn't built \ - with CUDA support and inference silently fell back to CPU" - ); - builder - .with_execution_providers([ - ort::ep::CUDA::default().with_device_id(0).build(), - ort::ep::CPU::default().build(), - ]) - .map_err(|e| e.to_string()) +fn cuda_provider() -> Result { + Ok(Provider::Cuda { device_id: 0 }) } #[cfg(not(feature = "cuda"))] -fn cuda_session(builder: SessionBuilder) -> Result { +fn cuda_provider() -> Result { eprintln!("breadmill: CUDA backend requested but not compiled in (rebuild with --features cuda); using CPU"); - Ok(builder) -} - -// ---- OpenVINO EP (Intel iGPU/dGPU) ------------------------------------------- - -#[cfg(feature = "openvino")] -fn openvino_session(builder: SessionBuilder, cache_dir: &Path) -> Result { - let ov_cache = cache_dir.join("openvino"); - std::fs::create_dir_all(&ov_cache).map_err(|e| e.to_string())?; - eprintln!("breadmill: using OpenVINO execution provider (device_type GPU)"); - eprintln!("breadmill: OpenVINO model cache: {}", ov_cache.display()); - eprintln!( - "breadmill: note — check the log line above/below for \"Successfully registered \ - `OpenVINOExecutionProvider`\"; if it's missing, the ONNX Runtime in use wasn't built \ - with OpenVINO support and inference silently fell back to CPU" - ); - builder - .with_execution_providers([ - ort::ep::OpenVINO::default() - .with_device_type("GPU") - .with_cache_dir(ov_cache.to_string_lossy()) - .build(), - ort::ep::CPU::default().build(), - ]) - .map_err(|e| e.to_string()) -} - -#[cfg(not(feature = "openvino"))] -fn openvino_session(builder: SessionBuilder, _cache_dir: &Path) -> Result { - eprintln!("breadmill: OpenVINO backend requested but not compiled in (rebuild with --features openvino); using CPU"); - Ok(builder) + Ok(Provider::Cpu) } /// Locate the VitisAI EP config file required by the AMD Ryzen AI SDK. diff --git a/breadmill/src/main.rs b/breadmill/src/main.rs index 08452e8..0b79eef 100644 --- a/breadmill/src/main.rs +++ b/breadmill/src/main.rs @@ -1,5 +1,4 @@ use std::{ - io::Read, path::{Path, PathBuf}, sync::{Arc, atomic::Ordering}, }; @@ -226,29 +225,10 @@ fn download_if_missing(url: &str, dest: &Path) -> Result<(), String> { eprintln!(" already present: {}", dest.display()); return Ok(()); } - - eprintln!(" downloading {} ...", url); - let agent = ureq::AgentBuilder::new() - .timeout(std::time::Duration::from_secs(300)) - .build(); - - let response = agent.get(url).call().map_err(|e| e.to_string())?; - let mut bytes = Vec::new(); - response - .into_reader() - .read_to_end(&mut bytes) - .map_err(|e| e.to_string())?; - - if bytes.is_empty() { - return Err(format!("empty download from {}", url)); - } - - // Write atomically via temp file - let tmp = dest.with_extension("tmp"); - std::fs::write(&tmp, &bytes).map_err(|e| e.to_string())?; - std::fs::rename(&tmp, dest).map_err(|e| e.to_string())?; - - eprintln!(" saved {} ({:.1} MB)", dest.display(), bytes.len() as f64 / 1_048_576.0); + // Shared with breadarrd's own (previously reqwest/async, now also this + // same sync/ureq implementation) model downloader — see + // bread_onnx::download's doc comment. + bread_onnx::download::ensure_file(url, dest, None).map_err(|e| e.to_string())?; Ok(()) }