Skip to content

Support ONNX models without token types (e.g. BAAI/bge-m3) #208

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sanikolaev merged 3 commits into from
Sep 19, 2026
Merged

Support ONNX models without token types (e.g. BAAI/bge-m3) #208

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4a443b9
fix(embeddings): support ONNX models without token types (e.g. BAAI/b…
sanikolaev Sep 18, 2026
9f83275
fix(embeddings): download ONNX external weights
sanikolaev Sep 18, 2026
d8dd5ef
fix(embeddings): honor ONNX model variants
donhardman Sep 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
214 changes: 157 additions & 57 deletions embeddings/src/model/local.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -245,6 +245,8 @@ pub struct LocalModelInfo {
pub gguf_path: Option<PathBuf>,
/// Path to ONNX file if using ONNX model
pub onnx_path: Option<PathBuf>,
/// sentence-transformers `1_Pooling/config.json`, fetched for ONNX models
pub pooling_config_path: Option<PathBuf>,
}

fn model_download_lock(model_id: &str) -> Arc<Mutex<()>> {
Expand Down Expand Up @@ -399,12 +401,21 @@ pub fn build_model_info(
}
};

// Only the ONNX path pools token states itself, so only it needs the
// model's published pooling mode.
let pooling_config_path = if onnx_path.is_some() {
api.get(POOLING_CONFIG_FILE).ok()
} else {
None
};

Ok(LocalModelInfo {
config_path,
tokenizer_path,
weights_paths,
gguf_path,
onnx_path,
pooling_config_path,
})
}

Expand Down Expand Up @@ -455,12 +466,39 @@ fn try_find_gguf_file(api: &hf_hub::api::sync::ApiRepo) -> Option<PathBuf> {
api.get(&gguf_files[0]).ok()
}

/// Try to find an ONNX model file in the repository
/// Looks for model.onnx at root or in onnx/ subdirectory
/// Try to find an ONNX model file in the repository.
/// Looks for model.onnx at root or in onnx/ subdirectory. External-data
/// exports keep tensors in sibling files ORT resolves relative to model.onnx,
/// so those are fetched too; a failed sidecar download fails the lookup.
fn try_find_onnx_file(api: &hf_hub::api::sync::ApiRepo) -> Option<PathBuf> {
api.get("model.onnx")
.ok()
.or_else(|| api.get("onnx/model.onnx").ok())
for model_path in ["model.onnx", "onnx/model.onnx"] {
let Ok(onnx_path) = api.get(model_path) else {
continue;
};
// Listing needs the network; a warm cache still loads offline.
if let Ok(repo_info) = api.info() {
let siblings: Vec<String> = repo_info
.siblings
.iter()
.map(|s| s.rfilename.clone())
.collect();
for sidecar in onnx_sidecar_files(model_path, &siblings) {
api.get(&sidecar).ok()?;
}
}
return Some(onnx_path);
}
None
}

/// External-data files of `model_path`: `model.onnx_data`, `model.onnx.data`,
/// `model.onnx_data_1`, ... share the model file name as prefix.
pub(crate) fn onnx_sidecar_files(model_path: &str, siblings: &[String]) -> Vec<String> {
siblings
.iter()
.filter(|file| file.starts_with(model_path) && file.as_str() != model_path)
.cloned()
.collect()
}

/// Load tokenizer with fallback for BPE format
Expand Down Expand Up @@ -923,13 +961,48 @@ impl QuantizedEmbeddingModel {
}
}

const POOLING_CONFIG_FILE: &str = "1_Pooling/config.json";
const ONNX_INPUT_IDS: &str = "input_ids";
const ONNX_ATTENTION_MASK: &str = "attention_mask";
const ONNX_TOKEN_TYPE_IDS: &str = "token_type_ids";
/// Output of sentence-transformers ONNX exports: pooled and normalized in-graph.
const ONNX_POOLED_OUTPUT: &str = "sentence_embedding";

/// How `[batch, seq, hidden]` token states are reduced when the graph does not
/// pool them itself. Comes from sentence-transformers `1_Pooling/config.json`.
#[derive(Debug, Clone, Copy, PartialEq)]
pub(crate) enum OnnxPooling {
Cls,
Mean,
}

pub(crate) fn parse_pooling_config(contents: &str) -> Result<OnnxPooling, Box<dyn Error>> {
let config: Value =
serde_json::from_str(contents).map_err(|_| LibError::ModelConfigParseFailed)?;
let flag = |key: &str| config.get(key).and_then(Value::as_bool).unwrap_or(false);
match (
flag("pooling_mode_cls_token"),
flag("pooling_mode_mean_tokens"),
) {
(true, false) => Ok(OnnxPooling::Cls),
(false, true) => Ok(OnnxPooling::Mean),
// max / last-token / concatenated modes would silently yield a vector
// space different from the published model's.
_ => Err(Box::new(LibError::ModelLoadFailed)),
}
}

/// ONNX embedding model using ORT (onnxruntime) for optimized inference.
/// Uses a single session with concurrent inference via SessionWrapper wrapper.
pub struct OnnxEmbeddingModel {
session: SessionWrapper,
tokenizer: Tokenizer,
max_input_len: usize,
hidden_size: usize,
/// Graph inputs in declaration order; exports differ (XLM-R graphs have no token_type_ids).
input_names: Vec<String>,
output_name: String,
pooling: OnnxPooling,
}

impl OnnxEmbeddingModel {
Expand Down Expand Up @@ -964,26 +1037,56 @@ impl OnnxEmbeddingModel {
.commit_from_file(&onnx_path)
.map_err(|_| LibError::ModelWeightsLoadFailed)?;

let input_names: Vec<String> = session
.inputs()
.iter()
.map(|input| input.name().to_string())
.collect();
if input_names.iter().any(|name| {
!matches!(
name.as_str(),
ONNX_INPUT_IDS | ONNX_ATTENTION_MASK | ONNX_TOKEN_TYPE_IDS
)
}) {
return Err(Box::new(LibError::ModelLoadFailed));
}
let output_names: Vec<&str> = session.outputs().iter().map(|o| o.name()).collect();
let output_name = output_names
.iter()
.copied()
.find(|name| *name == ONNX_POOLED_OUTPUT)
.or(output_names.first().copied())
.ok_or(LibError::ModelLoadFailed)?
.to_string();
// Token-state graphs are pooled here the way the model publishes it;
// mean pooling when the repo has no pooling config.
let pooling = match &model_info.pooling_config_path {
Some(path) if output_name != ONNX_POOLED_OUTPUT => parse_pooling_config(
&std::fs::read_to_string(path).map_err(|_| LibError::ModelConfigReadFailed)?,
)?,
_ => OnnxPooling::Mean,
};

Ok(Self {
session: SessionWrapper::new(session),
tokenizer,
max_input_len,
hidden_size,
input_names,
output_name,
pooling,
})
}

/// Run a single ONNX forward pass on one batch.
fn run_batch(
session: &SessionWrapper,
batch: &[Vec<u32>],
) -> Result<Vec<Vec<f32>>, Box<dyn Error>> {
fn run_batch(&self, batch: &[Vec<u32>]) -> Result<Vec<Vec<f32>>, Box<dyn Error>> {
let session = &self.session;
session.with_session(|sess| -> Result<Vec<Vec<f32>>, Box<dyn Error>> {
let batch_size = batch.len();
let max_len = batch.iter().map(|c| c.len()).max().unwrap_or(0);

let mut flat_ids: Vec<i64> = Vec::with_capacity(batch_size * max_len);
let mut flat_mask: Vec<i64> = Vec::with_capacity(batch_size * max_len);
let mut flat_type_ids: Vec<i64> = Vec::with_capacity(batch_size * max_len);

for chunk in batch {
let real_len = chunk.len();
Expand All @@ -993,27 +1096,25 @@ impl OnnxEmbeddingModel {
flat_ids.extend(std::iter::repeat_n(0i64, max_len - real_len));
flat_mask.extend(std::iter::repeat_n(1i64, real_len));
flat_mask.extend(std::iter::repeat_n(0i64, max_len - real_len));
flat_type_ids.extend(std::iter::repeat_n(0i64, max_len));
}

let input_ids = ort::value::Tensor::from_array((vec![batch_size, max_len], flat_ids))
.map_err(|_| LibError::OnnxModelEvalFailed)?;
let attention_mask =
ort::value::Tensor::from_array((vec![batch_size, max_len], flat_mask.clone()))
.map_err(|_| LibError::OnnxModelEvalFailed)?;
let token_type_ids =
ort::value::Tensor::from_array((vec![batch_size, max_len], flat_type_ids))
let mut inputs = Vec::with_capacity(self.input_names.len());
for name in &self.input_names {
let data = match name.as_str() {
ONNX_INPUT_IDS => std::mem::take(&mut flat_ids),
ONNX_ATTENTION_MASK => flat_mask.clone(),
ONNX_TOKEN_TYPE_IDS => vec![0i64; batch_size * max_len],
_ => unreachable!("input names are validated in new()"),
};
let tensor = ort::value::Tensor::from_array((vec![batch_size, max_len], data))
.map_err(|_| LibError::OnnxModelEvalFailed)?;

inputs.push((name.as_str(), tensor));
}
let outputs = sess
.run(ort::inputs![
"input_ids" => input_ids,
"attention_mask" => attention_mask,
"token_type_ids" => token_type_ids,
])
.run(inputs)
.map_err(|_| LibError::OnnxModelEvalFailed)?;

let (shape, data) = outputs[0]
let (shape, data) = outputs[self.output_name.as_str()]
.try_extract_tensor::<f32>()
.map_err(|_| LibError::OnnxModelEvalFailed)?;

Expand All @@ -1032,21 +1133,30 @@ impl OnnxEmbeddingModel {
let seq_len = shape[1] as usize;
let hidden_dim = shape[2] as usize;
for i in 0..batch_size {
let mut emb = vec![0.0f32; hidden_dim];
let mut count = 0.0f32;
for j in 0..seq_len {
let mask_val = flat_mask[i * max_len + j] as f32;
if mask_val > 0.0 {
let offset = (i * seq_len + j) * hidden_dim;
for k in 0..hidden_dim {
emb[k] += data[offset + k];
let mut emb = match self.pooling {
OnnxPooling::Cls => {
let offset = i * seq_len * hidden_dim;
data[offset..offset + hidden_dim].to_vec()
}
OnnxPooling::Mean => {
let mut emb = vec![0.0f32; hidden_dim];
let mut count = 0.0f32;
for j in 0..seq_len {
let mask_val = flat_mask[i * max_len + j] as f32;
if mask_val > 0.0 {
let offset = (i * seq_len + j) * hidden_dim;
for k in 0..hidden_dim {
emb[k] += data[offset + k];
}
count += 1.0;
}
}
if count > 0.0 {
emb.iter_mut().for_each(|v| *v /= count);
}
count += 1.0;
emb
}
}
if count > 0.0 {
emb.iter_mut().for_each(|v| *v /= count);
}
};
normalize(&mut emb);
embeddings.push(emb);
}
Expand All @@ -1059,17 +1169,14 @@ impl OnnxEmbeddingModel {
}

/// Tokenize one batch and run inference.
fn tokenize_and_infer(
session: &SessionWrapper,
tokenizer: &Tokenizer,
texts: &[&str],
max_input: usize,
) -> Result<Vec<Vec<f32>>, Box<dyn Error>> {
fn tokenize_and_infer(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, Box<dyn Error>> {
let max_input = self.max_input_len;
let truncated: Vec<&str> = texts
.iter()
.map(|t| pre_truncate_text(t, max_input))
.collect();
let encoded = tokenizer
let encoded = self
.tokenizer
.encode_batch(truncated, true)
.map_err(|_| LibError::ModelTokenizerEncodeFailed)?;
let chunks: Vec<Vec<u32>> = encoded
Expand All @@ -1079,7 +1186,7 @@ impl OnnxEmbeddingModel {
ids[..ids.len().min(max_input)].to_vec()
})
.collect();
Self::run_batch(session, &chunks)
self.run_batch(&chunks)
}

/// Adaptive predict: automatically chooses the best strategy based on input size.
Expand All @@ -1095,17 +1202,14 @@ impl OnnxEmbeddingModel {
threads: usize,
) -> Result<Vec<Vec<f32>>, Box<dyn Error>> {
let bs = batch_size();
let max_input = self.max_input_len;
let session = &self.session;
let tokenizer = &self.tokenizer;

if texts.is_empty() {
return Ok(Vec::new());
}

// Small input — single tokenize + infer, no threading overhead
if texts.len() <= bs {
return Self::tokenize_and_infer(session, tokenizer, texts, max_input);
return self.tokenize_and_infer(texts);
}

// Adaptive parallelism: scale workers with input size.
Expand All @@ -1130,13 +1234,9 @@ impl OnnxEmbeddingModel {
s.spawn(move || -> Result<Vec<Vec<f32>>, LibError> {
let mut embeddings = Vec::with_capacity(worker_texts.len());
for text in worker_texts {
let embs = Self::tokenize_and_infer(
session,
tokenizer,
std::slice::from_ref(text),
max_input,
)
.map_err(|_| LibError::OnnxModelEvalFailed)?;
let embs = self
.tokenize_and_infer(std::slice::from_ref(text))
.map_err(|_| LibError::OnnxModelEvalFailed)?;
embeddings.extend(embs);
}
Ok(embeddings)
Expand Down
Loading