Skip to content

Commit 70627da

Browse files
committed
perf(local): increase batch size; fix empty input
- Increase default batch size from 8 to 32 - Handle empty text input arrays efficiently
1 parent 1540065 commit 70627da

1 file changed

Lines changed: 6 additions & 1 deletion

File tree

embeddings/src/model/local.rs

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ use std::sync::{Arc, Mutex, OnceLock};
3030
use tokenizers::Tokenizer;
3131

3232
/// Default batch size per ONNX forward pass.
33-
const DEFAULT_BATCH_SIZE: usize = 8;
33+
const DEFAULT_BATCH_SIZE: usize = 32;
3434

3535
/// Default intra-op threads: 0 = ORT default (all cores).
3636
const DEFAULT_INTRA_THREADS: usize = 0;
@@ -916,6 +916,10 @@ impl OnnxEmbeddingModel {
916916
let tokenizer = &self.tokenizer;
917917
let session = &self.session;
918918

919+
if texts.is_empty() {
920+
return Ok(Vec::new());
921+
}
922+
919923
let text_batches: Vec<&[&str]> = texts.chunks(bs).collect();
920924
if text_batches.is_empty() {
921925
return Ok(Vec::new());
@@ -987,6 +991,7 @@ impl OnnxEmbeddingModel {
987991
if all_embeddings.len() != texts.len() {
988992
return Err(Box::new(LibError::OnnxModelEvalFailed));
989993
}
994+
990995
Ok(all_embeddings)
991996
}
992997
}

0 commit comments

Comments
 (0)