From 8a2936b8fd0509917a43a3e02d35a5ebc86d237f Mon Sep 17 00:00:00 2001 From: Breadway Date: Fri, 17 Jul 2026 09:37:55 +0800 Subject: [PATCH] Migrate embedding pipeline and model download to bread-onnx MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit OrtEmbedder's tokenize -> tensor build -> mean-pool -> L2-normalize pipeline was near-byte-identical to breadmill's own OrtEmbedder (same truncation, same actual_seq.min(mask.len()) padding guard, same 1e-10 epsilon) — now both share bread_onnx::embedding::EmbeddingSession (path dependency for now, see the TODO in breadarrd/Cargo.toml). This crate stays CPU-only (Provider::Cpu), matching its existing documented rationale. ensure_model's reqwest-based download function is replaced with bread_onnx::download::ensure_file (sync/ureq, matching breadmill's own downloader and this workspace's bakery convention) dispatched via spawn_blocking from this async context. Builds and tests clean across the whole breadarr workspace: 205 passed, 1 pre-existing network-dependent test ignored, 0 failed. --- Cargo.lock | 333 +++++++++++++++++++++++++++++++-- breadarrd/Cargo.toml | 2 + breadarrd/src/matcher/embed.rs | 126 +------------ breadarrd/src/matcher/mod.rs | 26 ++- 4 files changed, 342 insertions(+), 145 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ce77140..d8cee0c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,12 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + [[package]] name = "ahash" version = "0.8.12" @@ -160,6 +166,38 @@ version = "2.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +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-utils" +version = "0.2.3" +dependencies = [ + "dirs", + "serde", + "serde_json", +] + [[package]] name = "breadarr-shared" version = "0.1.0" @@ -192,6 +230,7 @@ dependencies = [ "anyhow", "async-trait", "axum", + "bread-onnx", "breadarr-shared", "chrono", "fastrand", @@ -375,6 +414,24 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + [[package]] name = "crossbeam-deque" version = "0.8.7" @@ -425,6 +482,16 @@ dependencies = [ "winapi", ] +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + [[package]] name = "cssparser" version = "0.31.2" @@ -590,6 +657,37 @@ dependencies = [ "syn", ] +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "dirs" +version = "5.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44c45a9d03d6676652bcb5e724c7e988de1acad23a711b5217ab9cbecbec2225" +dependencies = [ + "dirs-sys", +] + +[[package]] +name = "dirs-sys" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520f05a5cbd335fae5a99ff7a6ab8627577660ee5cfd6a94a6a929b52ff0321c" +dependencies = [ + "libc", + "option-ext", + "redox_users", + "windows-sys 0.48.0", +] + [[package]] name = "displaydoc" version = "0.2.6" @@ -701,6 +799,16 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "crc32fast", + "miniz_oxide", +] + [[package]] name = "fnv" version = "1.0.7" @@ -795,6 +903,16 @@ dependencies = [ "byteorder", ] +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + [[package]] name = "getopts" version = "0.2.24" @@ -898,6 +1016,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + [[package]] name = "hmac-sha256" version = "1.1.14" @@ -1272,6 +1396,15 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libredox" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c943259e342f1e06ff2da7a83eabdfe7f92ce10262688dbf1895ff0b3e6e4652" +dependencies = [ + "libc", +] + [[package]] name = "libsqlite3-sys" version = "0.28.0" @@ -1426,6 +1559,16 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "mio" version = "1.2.1" @@ -1633,6 +1776,12 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "option-ext" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" + [[package]] name = "ort" version = "2.0.0-rc.12" @@ -1643,7 +1792,7 @@ dependencies = [ "ort-sys", "smallvec", "tracing", - "ureq", + "ureq 3.3.0", ] [[package]] @@ -1654,7 +1803,7 @@ checksum = "d7b497d21a8b6fbb4b5a544f8fadb77e801a09ae0add9e411d31c6f89e3c1e90" dependencies = [ "hmac-sha256", "lzma-rust2", - "ureq", + "ureq 3.3.0", ] [[package]] @@ -2029,6 +2178,17 @@ dependencies = [ "bitflags", ] +[[package]] +name = "redox_users" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba009ff324d1fc1b900bd1fdb31564febe58a8ccc8a6fdbb93b543d33b13ca43" +dependencies = [ + "getrandom 0.2.17", + "libredox", + "thiserror 1.0.69", +] + [[package]] name = "regex" version = "1.13.0" @@ -2162,7 +2322,9 @@ version = "0.23.41" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6b92b125634d9b795e7beca796cc790df15a7fb38323bf3196fda83292d06b1f" dependencies = [ + "log", "once_cell", + "ring", "rustls-pki-types", "rustls-webpki", "subtle", @@ -2358,6 +2520,17 @@ dependencies = [ "stable_deref_trait", ] +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "sharded-slab" version = "0.1.7" @@ -2404,6 +2577,12 @@ dependencies = [ "libc", ] +[[package]] +name = "simd-adler32" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea" + [[package]] name = "siphasher" version = "0.3.11" @@ -2608,13 +2787,33 @@ dependencies = [ "utf-8", ] +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + [[package]] name = "thiserror" version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" dependencies = [ - "thiserror-impl", + "thiserror-impl 2.0.18", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn", ] [[package]] @@ -2705,7 +2904,7 @@ dependencies = [ "serde", "serde_json", "spm_precompiled", - "thiserror", + "thiserror 2.0.18", "unicode-normalization-alignments", "unicode-segmentation", "unicode_categories", @@ -2927,6 +3126,12 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + [[package]] name = "unicase" version = "2.9.0" @@ -2995,6 +3200,24 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" +[[package]] +name = "ureq" +version = "2.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02d1a66277ed75f640d608235660df48c8e3c19f3b4edb6a263315626cc3c01d" +dependencies = [ + "base64 0.22.1", + "flate2", + "log", + "once_cell", + "rustls", + "rustls-pki-types", + "serde", + "serde_json", + "url", + "webpki-roots 0.26.11", +] + [[package]] name = "ureq" version = "3.3.0" @@ -3181,6 +3404,24 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "webpki-roots" +version = "0.26.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" +dependencies = [ + "webpki-roots 1.0.8", +] + +[[package]] +name = "webpki-roots" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf85cb06032201fa7c6f829d7db5a7e5aa45bcc0655327713065f6f0576731bf" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "winapi" version = "0.3.9" @@ -3273,13 +3514,22 @@ dependencies = [ "windows-link", ] +[[package]] +name = "windows-sys" +version = "0.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" +dependencies = [ + "windows-targets 0.48.5", +] + [[package]] name = "windows-sys" version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets", + "windows-targets 0.52.6", ] [[package]] @@ -3288,7 +3538,7 @@ version = "0.59.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" dependencies = [ - "windows-targets", + "windows-targets 0.52.6", ] [[package]] @@ -3300,34 +3550,67 @@ dependencies = [ "windows-link", ] +[[package]] +name = "windows-targets" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" +dependencies = [ + "windows_aarch64_gnullvm 0.48.5", + "windows_aarch64_msvc 0.48.5", + "windows_i686_gnu 0.48.5", + "windows_i686_msvc 0.48.5", + "windows_x86_64_gnu 0.48.5", + "windows_x86_64_gnullvm 0.48.5", + "windows_x86_64_msvc 0.48.5", +] + [[package]] name = "windows-targets" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm", - "windows_aarch64_msvc", - "windows_i686_gnu", + "windows_aarch64_gnullvm 0.52.6", + "windows_aarch64_msvc 0.52.6", + "windows_i686_gnu 0.52.6", "windows_i686_gnullvm", - "windows_i686_msvc", - "windows_x86_64_gnu", - "windows_x86_64_gnullvm", - "windows_x86_64_msvc", + "windows_i686_msvc 0.52.6", + "windows_x86_64_gnu 0.52.6", + "windows_x86_64_gnullvm 0.52.6", + "windows_x86_64_msvc 0.52.6", ] +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" + [[package]] name = "windows_aarch64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" +[[package]] +name = "windows_aarch64_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" + [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" +[[package]] +name = "windows_i686_gnu" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" + [[package]] name = "windows_i686_gnu" version = "0.52.6" @@ -3340,24 +3623,48 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" +[[package]] +name = "windows_i686_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" + [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" +[[package]] +name = "windows_x86_64_gnu" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" + [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" + [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" +[[package]] +name = "windows_x86_64_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" + [[package]] name = "windows_x86_64_msvc" version = "0.52.6" diff --git a/breadarrd/Cargo.toml b/breadarrd/Cargo.toml index 77824c6..c6754c4 100644 --- a/breadarrd/Cargo.toml +++ b/breadarrd/Cargo.toml @@ -19,6 +19,8 @@ regex.workspace = true serde_json.workspace = true ort.workspace = true tokenizers.workspace = true +# 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" } scraper.workspace = true chrono.workspace = true fastrand.workspace = true diff --git a/breadarrd/src/matcher/embed.rs b/breadarrd/src/matcher/embed.rs index 4f9bdf9..13902a9 100644 --- a/breadarrd/src/matcher/embed.rs +++ b/breadarrd/src/matcher/embed.rs @@ -1,10 +1,8 @@ use std::path::Path; use anyhow::Result; -use ort::session::builder::GraphOptimizationLevel; -use ort::session::Session; -use ort::value::Tensor; -use tokenizers::Tokenizer; +use bread_onnx::embedding::EmbeddingSession; +use bread_onnx::Provider; /// all-MiniLM-L6-v2's trained max sequence length. Release/show titles are /// always far shorter than this, but truncate defensively rather than let a @@ -12,9 +10,7 @@ use tokenizers::Tokenizer; const MAX_SEQ_LEN: usize = 256; pub struct OrtEmbedder { - session: Session, - tokenizer: Tokenizer, - dim: usize, + inner: EmbeddingSession, } impl OrtEmbedder { @@ -23,130 +19,24 @@ impl OrtEmbedder { /// MiniLM-class model is cheap enough on CPU that a multi-backend GPU /// setup isn't worth the added complexity for a model this small. pub fn load(model_path: &Path, tokenizer_path: &Path, dim: usize) -> Result { - let session = Session::builder() - .map_err(|e| anyhow::anyhow!("failed to create ort session builder: {e}"))? - .with_optimization_level(GraphOptimizationLevel::Level3) - .map_err(|e| anyhow::anyhow!("failed to set optimization level: {e}"))? - .commit_from_file(model_path) - .map_err(|e| { - anyhow::anyhow!("failed to load model from {}: {e}", model_path.display()) - })?; - - let tokenizer = Tokenizer::from_file(tokenizer_path) - .map_err(|e| anyhow::anyhow!("failed to load tokenizer: {e}"))?; - - Ok(Self { - session, - tokenizer, - dim, - }) + let inner = EmbeddingSession::load(model_path, tokenizer_path, dim, MAX_SEQ_LEN, &[Provider::Cpu])?; + Ok(Self { inner }) } pub fn embed(&mut self, text: &str) -> Result> { - let encoding = self - .tokenizer - .encode(text, true) - .map_err(|e| anyhow::anyhow!("tokenization failed: {e}"))?; - - 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(); - - 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)) - .map_err(|e| anyhow::anyhow!("failed to build input_ids tensor: {e}"))?; - let mask_tensor = Tensor::::from_array((vec![1i64, seq_len], mask.clone())) - .map_err(|e| anyhow::anyhow!("failed to build attention_mask tensor: {e}"))?; - let type_tensor = Tensor::::from_array((vec![1i64, seq_len], type_ids)) - .map_err(|e| anyhow::anyhow!("failed to build token_type_ids tensor: {e}"))?; - - let outputs = self - .session - .run(ort::inputs! { - "input_ids" => id_tensor, - "attention_mask" => mask_tensor, - "token_type_ids" => type_tensor, - }) - .map_err(|e| anyhow::anyhow!("ort inference failed: {e}"))?; - - let (shape, data) = outputs["last_hidden_state"] - .try_extract_tensor::() - .map_err(|e| anyhow::anyhow!("failed to extract last_hidden_state: {e}"))?; - - let actual_seq = shape[1] as usize; - let actual_dim = shape[2] as usize; - - // Mean-pool over non-padded positions only. - 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); - result.truncate(self.dim); - while result.len() < self.dim { - result.push(0.0); - } - - Ok(result) + self.inner.embed(text) } } -fn l2_normalize(v: &mut [f32]) { - let norm: f32 = v.iter().map(|x| x * x).sum::().sqrt(); - if norm > 1e-10 { - for x in v.iter_mut() { - *x /= norm; - } - } -} - -pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 { - a.iter().zip(b).map(|(x, y)| x * y).sum() -} +pub use bread_onnx::embedding::cosine_similarity; #[cfg(test)] mod tests { use super::*; - #[test] - fn l2_normalize_produces_unit_vector() { - let mut v = vec![3.0, 4.0]; - l2_normalize(&mut v); - let norm: f32 = v.iter().map(|x| x * x).sum::().sqrt(); - assert!((norm - 1.0).abs() < 1e-6); - } - - #[test] - fn l2_normalize_leaves_zero_vector_untouched() { - let mut v = vec![0.0, 0.0, 0.0]; - l2_normalize(&mut v); - assert_eq!(v, vec![0.0, 0.0, 0.0]); - } - #[test] fn cosine_similarity_of_identical_unit_vectors_is_one() { - let mut v = vec![1.0, 2.0, 3.0]; - l2_normalize(&mut v); + let v = vec![0.6, 0.8]; // already unit length let sim = cosine_similarity(&v, &v); assert!((sim - 1.0).abs() < 1e-6); } diff --git a/breadarrd/src/matcher/mod.rs b/breadarrd/src/matcher/mod.rs index ef19a56..f503d05 100644 --- a/breadarrd/src/matcher/mod.rs +++ b/breadarrd/src/matcher/mod.rs @@ -16,6 +16,13 @@ const TOKENIZER_URL: &str = /// Downloads the embedding model into `model_dir` if it isn't already /// there — keeps setup to "run the daemon," no separate fetch step, in /// keeping with the project's minimal-setup goal. +/// +/// `bread_onnx::download::ensure_file` is sync/blocking (`ureq`, matching +/// this workspace's `bakery` download convention) — this used to be a +/// `reqwest`-based async implementation local to this crate, genuinely +/// duplicating breadmill's own sync/`ureq` downloader. Since this fn is +/// called from an async context, each call is dispatched via +/// `spawn_blocking` rather than blocking the async runtime directly. pub async fn ensure_model(model_dir: &Path) -> Result<(PathBuf, PathBuf)> { std::fs::create_dir_all(model_dir) .with_context(|| format!("failed to create {}", model_dir.display()))?; @@ -25,28 +32,19 @@ pub async fn ensure_model(model_dir: &Path) -> Result<(PathBuf, PathBuf)> { if !model_path.exists() { tracing::info!("downloading title-matching model (~90MB, one-time)"); - download(MODEL_URL, &model_path).await?; + download(MODEL_URL, model_path.clone()).await?; } if !tokenizer_path.exists() { - download(TOKENIZER_URL, &tokenizer_path).await?; + download(TOKENIZER_URL, tokenizer_path.clone()).await?; } Ok((model_path, tokenizer_path)) } -async fn download(url: &str, dest: &Path) -> Result<()> { - let bytes = reqwest::get(url) +async fn download(url: &'static str, dest: PathBuf) -> Result<()> { + tokio::task::spawn_blocking(move || bread_onnx::download::ensure_file(url, &dest, None)) .await - .with_context(|| format!("failed to download {url}"))? - .error_for_status() - .with_context(|| format!("{url} returned an error status"))? - .bytes() - .await - .with_context(|| format!("failed to read response body from {url}"))?; - let tmp = dest.with_extension("part"); - std::fs::write(&tmp, &bytes).with_context(|| format!("failed to write {}", tmp.display()))?; - std::fs::rename(&tmp, dest) - .with_context(|| format!("failed to finalize {}", dest.display()))?; + .context("download task panicked")??; Ok(()) }