Migrate embedding pipeline, EP session building, and model download to bread-onnx
embed.rs's OrtEmbedder now delegates its tokenize -> tensor build -> mean-pool -> L2-normalize pipeline to bread_onnx::embedding::EmbeddingSession (near-byte-identical to breadarrd's own OrtEmbedder — same duplication, now shared, path dependency for now, see the TODO in breadmill/Cargo.toml), and its per-EP session builders (npu_session/rocm_session/cuda_session/ openvino_session) collapse into a single to_provider() mapping onto bread_onnx::Provider, which this crate's own breadmill/src/session.rs counterpart now handles generically. This crate's Backend enum, cargo feature gates (npu/rocm/cuda/openvino/full - unchanged, still control which onnxruntime EPs actually link/load), and NPU vaip_config.json discovery all stay local since they're genuinely breadsearch-specific. main.rs's Backend construction and CLI flag handling are untouched. This is also the reference implementation the MIGraphX-not-ROCm default in bread-onnx's provider module was promoted from (see this machine's own breadsearch-gpu-backends operator notes) — breadpad's ONNX migration, which had the actual silent-fallback bug, follows in a later commit. download_if_missing is replaced with bread_onnx::download::ensure_file (same sync/ureq approach, now shared with breadarrd's downloader). Builds clean with default features AND --features full (npu+rocm+cuda+ openvino all compiling together, matching how this crate already combined them). All existing tests pass across the whole workspace.
This commit is contained in:
Breadway
parent
d01a3841d9
commit
0990969722
4 changed files with 88 additions and 231 deletions
44
Cargo.lock
generated
44
Cargo.lock
generated
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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<Self, String> {
|
||||
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<Vec<f32>, String> {
|
||||
|
|
@ -64,229 +53,71 @@ impl OrtEmbedder {
|
|||
|
||||
fn embed_with_prefix(&mut self, text: &str, prefix: &str) -> Result<Vec<f32>, String> {
|
||||
let input = format!("{}{}", prefix, text);
|
||||
|
||||
let encoding = self
|
||||
.tokenizer
|
||||
.encode(input, true)
|
||||
.map_err(|e| e.to_string())?;
|
||||
|
||||
let mut ids: Vec<i64> = encoding.get_ids().iter().map(|&x| x as i64).collect();
|
||||
let mut mask: Vec<i64> = encoding
|
||||
.get_attention_mask()
|
||||
.iter()
|
||||
.map(|&x| x as i64)
|
||||
.collect();
|
||||
let mut type_ids: Vec<i64> = 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::<i64>::from_array((vec![1i64, seq_len], ids.clone())).map_err(|e| e.to_string())?;
|
||||
let mask_tensor =
|
||||
Tensor::<i64>::from_array((vec![1i64, seq_len], mask.clone())).map_err(|e| e.to_string())?;
|
||||
let type_tensor =
|
||||
Tensor::<i64>::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::<f32>()
|
||||
.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;
|
||||
self.inner.embed(&input).map_err(|e| e.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
if count > 0 {
|
||||
for x in &mut result {
|
||||
*x /= count as f32;
|
||||
}
|
||||
}
|
||||
// ---- 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.
|
||||
|
||||
l2_normalize(&mut result);
|
||||
|
||||
// Clamp/pad to configured dim
|
||||
result.truncate(self.dim);
|
||||
while result.len() < self.dim {
|
||||
result.push(0.0);
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
}
|
||||
|
||||
fn l2_normalize(v: &mut Vec<f32>) {
|
||||
let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
if norm > 1e-10 {
|
||||
for x in v.iter_mut() {
|
||||
*x /= norm;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---- Execution provider selection -------------------------------------------
|
||||
|
||||
fn configure_eps(builder: SessionBuilder, backend: &Backend) -> Result<SessionBuilder, String> {
|
||||
fn to_provider(backend: Backend) -> Result<Provider, String> {
|
||||
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<SessionBuilder, String> {
|
||||
let vitis_ep = build_vitis_ep(cache_dir)?;
|
||||
eprintln!("breadmill: using NPU (VitisAI) execution provider");
|
||||
fn npu_provider(cache_dir: PathBuf) -> Result<Provider, String> {
|
||||
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<SessionBuilder, String> {
|
||||
fn npu_provider(_cache_dir: PathBuf) -> Result<Provider, String> {
|
||||
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<ort::ep::ExecutionProviderDispatch, String> {
|
||||
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<SessionBuilder, String> {
|
||||
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<Provider, String> {
|
||||
Ok(Provider::MiGraphX { device_id: 0 })
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "rocm"))]
|
||||
fn rocm_session(builder: SessionBuilder) -> Result<SessionBuilder, String> {
|
||||
fn rocm_provider() -> Result<Provider, String> {
|
||||
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<SessionBuilder, String> {
|
||||
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<Provider, String> {
|
||||
Ok(Provider::Cuda { device_id: 0 })
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
fn cuda_session(builder: SessionBuilder) -> Result<SessionBuilder, String> {
|
||||
fn cuda_provider() -> Result<Provider, String> {
|
||||
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<SessionBuilder, String> {
|
||||
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<SessionBuilder, String> {
|
||||
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.
|
||||
|
|
|
|||
|
|
@ -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(())
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue