Migrate embedding pipeline and model download to bread-onnx
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.
This commit is contained in:
Breadway
parent
d35d9a1703
commit
8a2936b8fd
4 changed files with 342 additions and 145 deletions
|
|
@ -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<Self> {
|
||||
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<Vec<f32>> {
|
||||
let encoding = self
|
||||
.tokenizer
|
||||
.encode(text, true)
|
||||
.map_err(|e| anyhow::anyhow!("tokenization failed: {e}"))?;
|
||||
|
||||
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();
|
||||
|
||||
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))
|
||||
.map_err(|e| anyhow::anyhow!("failed to build input_ids tensor: {e}"))?;
|
||||
let mask_tensor = Tensor::<i64>::from_array((vec![1i64, seq_len], mask.clone()))
|
||||
.map_err(|e| anyhow::anyhow!("failed to build attention_mask tensor: {e}"))?;
|
||||
let type_tensor = Tensor::<i64>::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::<f32>()
|
||||
.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::<f32>().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::<f32>().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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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(())
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue