feat: up batch size

README.md

@@ -74,7 +74,7 @@ fn main() -> Result<()> {
     ];
 
     // Generate embeddings using default parameters
-    // (Default max_length: Some(512), Default batch_size: 1024)
+    // (Default max_length: Some(512), all sentences in one batch)
     let embeddings = model.encode(&sentences);
     // `embeddings` is a Vec<Vec<f32>>
     println!("Generated {} embeddings.", embeddings.len());
@@ -87,6 +87,15 @@ fn main() -> Result<()> {
     );
     println!("Generated {} custom embeddings.", custom_embeddings.len());
 
+    // `encode_flat` returns one row-major `Vec<f32>` of
+    // `sentences.len() * model.dim()` values. Use it when you want contiguous
+    // output or only read the rows; prefer `encode` if you need an owned
+    // `Vec<f32>` per sentence.
+    let flat = model.encode_flat(&sentences);
+    for row in flat.chunks_exact(model.dim()) {
+        println!("First value: {}", row[0]);
+    }
+
     Ok(())
 }
 ```

src/model.rs

@@ -11,7 +11,11 @@ use std::{
     path::{Path, PathBuf},
     sync::Arc,
 };
-use tokenizers::{Tokenizer, models::ModelWrapper, parallelism::MaybeParallelRefIterator};
+use tokenizers::{
+    Tokenizer,
+    models::ModelWrapper,
+    parallelism::{MaybeParallelRefIterator, MaybeParallelRefMutIterator},
+};
 
 use crate::fast_unigram::FastUnigram;
 use crate::fast_wordpiece::FastWordPiece;
@@ -410,6 +414,9 @@ impl StaticModel {
         weights: Option<Cow<'static, [f32]>>,
         token_mapping: Option<Cow<'static, [usize]>>,
     ) -> Result<Self> {
+        if cols == 0 {
+            return Err(anyhow!("embedding dimension must be greater than 0"));
+        }
         if embeddings.len() != rows * cols {
             return Err(anyhow!(
                 "embeddings length {} != rows {} * cols {}",
@@ -535,9 +542,33 @@ impl StaticModel {
         embeddings
     }
 
-    /// Default encode: `max_length=512`, `batch_size=1024`
+    /// Default encode: `max_length=512`, all sentences in one batch
     pub fn encode(&self, sentences: &[String]) -> Vec<Vec<f32>> {
-        self.encode_with_args(sentences, Some(512), 1024)
+        self.encode_with_args(sentences, Some(512), sentences.len().max(1))
+    }
+
+    /// Embedding dimension of the vectors this model produces.
+    pub fn dim(&self) -> usize {
+        self.inner.dim
+    }
+
+    /// Encode texts into one row-major buffer of `sentences.len() * dim` floats.
+    ///
+    /// # Arguments
+    /// * `sentences` - the list of sentences to encode.
+    /// * `max_length` - max tokens per text.
+    pub fn encode_flat_with_args(&self, sentences: &[String], max_length: Option<usize>) -> Vec<f32> {
+        let dim = self.inner.dim;
+        let mut out = vec![0.0_f32; sentences.len() * dim];
+        let mut jobs: Vec<(&str, &mut [f32])> = sentences.iter().map(String::as_str).zip(out.chunks_mut(dim)).collect();
+        jobs.maybe_par_iter_mut()
+            .for_each(|(text, row)| self.encode_into(text, max_length, row));
+        out
+    }
+
+    /// Default flat encode: `max_length=512`
+    pub fn encode_flat(&self, sentences: &[String]) -> Vec<f32> {
+        self.encode_flat_with_args(sentences, Some(512))
     }
 
     /// Encode a single sentence into a vector.
@@ -549,6 +580,13 @@ impl StaticModel {
     /// Uses the custom WordPiece tokenizer or Unigram model when the model
     /// matches that pipeline; other pipelines use the `tokenizers` crate.
     fn encode_one(&self, text: &str, max_length: Option<usize>) -> Vec<f32> {
+        let mut out = vec![0.0_f32; self.inner.dim];
+        self.encode_into(text, max_length, &mut out);
+        out
+    }
+
+    /// Truncate, tokenize and pool one text into `out`, which must hold `dim` zeros.
+    fn encode_into(&self, text: &str, max_length: Option<usize>, out: &mut [f32]) {
         let model = &self.inner;
         let text = max_length.map_or(text, |max_tok| {
             Self::truncate_str(text, max_tok, model.median_token_length)
@@ -557,59 +595,58 @@ impl StaticModel {
             (Some(ft), _) => TOKEN_IDS.with_borrow_mut(|ids| {
                 ids.clear();
                 ft.encode_ids_into(text, ids);
-                self.pool_ids(ids, max_length)
+                self.pool_ids(ids, max_length, out)
             }),
             (None, Some(fu)) => {
                 let ids = fu.encode_ids(&model.tokenizer, text).expect("tokenization failed");
-                self.pool_ids(&ids, max_length)
+                self.pool_ids(&ids, max_length, out)
             }
             (None, None) => {
                 let encoding = model.tokenizer.encode_fast(text, false).expect("tokenization failed");
-                self.pool_ids(encoding.get_ids(), max_length)
+                self.pool_ids(encoding.get_ids(), max_length, out)
             }
         }
     }
 
     /// Mean-pool the first `max_length` token IDs, skipping the unk token.
-    fn pool_ids(&self, ids: &[u32], max_length: Option<usize>) -> Vec<f32> {
+    fn pool_ids(&self, ids: &[u32], max_length: Option<usize>, sum: &mut [f32]) {
         let model = &self.inner;
         let ids = &ids[..max_length.map_or(ids.len(), |max_tok| max_tok.min(ids.len()))];
-        let (mut sum, cnt) = match &model.embeddings {
-            Embeddings::F32(table) => self.sum_rows(table, ids),
-            Embeddings::I8(table) => self.sum_rows(table, ids),
+        let cnt = match &model.embeddings {
+            Embeddings::F32(table) => self.sum_rows(table, ids, sum),
+            Embeddings::I8(table) => self.sum_rows(table, ids, sum),
         };
         let denom = cnt.max(1) as f32;
-        for x in &mut sum {
+        for x in sum.iter_mut() {
             *x /= denom;
         }
         if model.normalize {
             let norm = sum.iter().map(|&v| v * v).sum::<f32>().sqrt().max(1e-12);
-            for x in &mut sum {
+            for x in sum.iter_mut() {
                 *x /= norm;
             }
         }
-        sum
     }
 
     /// Weighted row sum, using a compile-time dimension when one matches.
-    fn sum_rows<T: Copy>(&self, table: &[T], ids: &[u32]) -> (Vec<f32>, usize)
+    fn sum_rows<T: Copy>(&self, table: &[T], ids: &[u32], out: &mut [f32]) -> usize
     where
         f32: From<T>,
     {
         match self.inner.dim {
-            64 => self.sum_rows_fixed::<T, 64>(table, ids),
-            128 => self.sum_rows_fixed::<T, 128>(table, ids),
-            256 => self.sum_rows_fixed::<T, 256>(table, ids),
-            384 => self.sum_rows_fixed::<T, 384>(table, ids),
-            512 => self.sum_rows_fixed::<T, 512>(table, ids),
-            768 => self.sum_rows_fixed::<T, 768>(table, ids),
-            1024 => self.sum_rows_fixed::<T, 1024>(table, ids),
-            _ => self.sum_rows_dyn(table, ids),
+            64 => self.sum_rows_fixed::<T, 64>(table, ids, out),
+            128 => self.sum_rows_fixed::<T, 128>(table, ids, out),
+            256 => self.sum_rows_fixed::<T, 256>(table, ids, out),
+            384 => self.sum_rows_fixed::<T, 384>(table, ids, out),
+            512 => self.sum_rows_fixed::<T, 512>(table, ids, out),
+            768 => self.sum_rows_fixed::<T, 768>(table, ids, out),
+            1024 => self.sum_rows_fixed::<T, 1024>(table, ids, out),
+            _ => self.sum_rows_dyn(table, ids, out),
         }
     }
 
     /// Weighted row sum for a compile-time embedding dimension.
-    fn sum_rows_fixed<T: Copy, const D: usize>(&self, table: &[T], ids: &[u32]) -> (Vec<f32>, usize)
+    fn sum_rows_fixed<T: Copy, const D: usize>(&self, table: &[T], ids: &[u32], out: &mut [f32]) -> usize
     where
         f32: From<T>,
     {
@@ -620,21 +657,20 @@ impl StaticModel {
                 *s += f32::from(v) * scale;
             }
         });
-        (sum.to_vec(), cnt)
+        out.copy_from_slice(&sum);
+        cnt
     }
 
     /// Weighted row sum for any embedding dimension.
-    fn sum_rows_dyn<T: Copy>(&self, table: &[T], ids: &[u32]) -> (Vec<f32>, usize)
+    fn sum_rows_dyn<T: Copy>(&self, table: &[T], ids: &[u32], sum: &mut [f32]) -> usize
     where
         f32: From<T>,
     {
-        let mut sum = vec![0.0_f32; self.inner.dim];
-        let cnt = self.for_each_row(table, ids, |row, scale| {
+        self.for_each_row(table, ids, |row, scale| {
             for (s, &v) in sum.iter_mut().zip(row) {
                 *s += f32::from(v) * scale;
             }
-        });
-        (sum, cnt)
+        })
     }
 
     /// Call `f` with the embedding row and scale (token weight times row

tests/test_model.rs

@@ -241,3 +241,14 @@ fn test_unigram_unknown_token_is_not_pooled() {
 
     assert_eq!(model.encode_single("axb"), model.encode_single("ab"));
 }
+
+#[test]
+fn test_zero_dim_model_is_rejected() {
+    use tokenizers::Tokenizer;
+
+    let tokenizer = Tokenizer::from_file("tests/fixtures/test-model-float32/tokenizer.json").unwrap();
+    let rows = tokenizer.get_vocab_size(true);
+    let result = StaticModel::from_owned(tokenizer, Vec::new(), rows, 0, true, None, None);
+    let msg = result.err().unwrap().to_string();
+    assert!(msg.contains("dimension"), "error should mention dimension: {msg}");
+}

tests/test_pooling.rs

@@ -242,15 +242,27 @@ fn check_all(tokenizer_path: &str, dims: &[usize]) {
                 let case = Case::new(&tokenizer, dim, extras, normalize);
                 let model = case.model();
 
+                assert_eq!(model.dim(), dim, "{} dim()", case.label());
                 for (max_length, ids) in MAX_LENGTHS.iter().zip(&ids) {
                     let want = case.reference(ids);
+                    let flat = model.encode_flat_with_args(&texts, *max_length);
+                    assert_eq!(
+                        flat.len(),
+                        texts.len() * dim,
+                        "{} encode_flat_with_args length",
+                        case.label()
+                    );
+                    let rows: Vec<Vec<f32>> = flat.chunks_exact(dim).map(<[f32]>::to_vec).collect();
+                    let what = format!("encode_flat_with_args(max_length={max_length:?})");
+                    assert_bit_identical(&case, &what, &rows, &want, &texts);
                     for batch_size in [1, 1024] {
                         let got = model.encode_with_args(&texts, *max_length, batch_size);
                         let what = format!("encode_with_args(max_length={max_length:?}, batch_size={batch_size})");
                         assert_bit_identical(&case, &what, &got, &want, &texts);
                     }
                     if *max_length == Some(512) {
                         assert_bit_identical(&case, "encode", &model.encode(&texts), &want, &texts);
+                        assert_eq!(model.encode_flat(&texts), want.concat(), "{} encode_flat", case.label());
                         let singles: Vec<Vec<f32>> = texts.iter().map(|t| model.encode_single(t)).collect();
                         assert_bit_identical(&case, "encode_single", &singles, &want, &texts);
                     }
@@ -270,6 +282,14 @@ fn unigram_pooling_matches_reference() {
     check_all(UNIGRAM_TOKENIZER, &[2, 64, 256, 1000]);
 }
 
+#[test]
+fn encode_flat_of_no_texts_is_empty() {
+    let tokenizer = Tokenizer::from_file(WORDPIECE_TOKENIZER).expect("load fixture tokenizer");
+    let model = Case::new(&tokenizer, 64, Extras::None, true).model();
+    assert!(model.encode_flat(&[]).is_empty());
+    assert!(model.encode(&[]).is_empty());
+}
+
 #[test]
 fn interleaved_lengths_reuse_buffers_correctly() {
     let tokenizer = Tokenizer::from_file(WORDPIECE_TOKENIZER).expect("load fixture tokenizer");
@@ -289,6 +309,7 @@ fn interleaved_lengths_reuse_buffers_correctly() {
     let singles: Vec<Vec<f32>> = texts.iter().map(|t| model.encode_single(t)).collect();
     assert_bit_identical(&case, "encode_single", &singles, &want, &texts);
     assert_bit_identical(&case, "encode", &model.encode(&texts), &want, &texts);
+    assert_eq!(model.encode_flat(&texts), want.concat(), "{} encode_flat", case.label());
 }
 
 #[test]

tests/test_quantize_int8.rs

@@ -186,6 +186,12 @@ fn check(tokenizer_path: &str, dim: usize, combo: (bool, bool, bool)) {
     }
     let singles = |m: &StaticModel| texts.iter().map(|t| m.encode_single(t)).collect::<Vec<_>>();
     assert_bit_identical(&case.label(), &singles(&quantized), &singles(&oracle));
+    let flat: Vec<Vec<f32>> = quantized
+        .encode_flat(&texts)
+        .chunks_exact(dim)
+        .map(<[f32]>::to_vec)
+        .collect();
+    assert_bit_identical(&format!("{} encode_flat", case.label()), &flat, &oracle.encode(&texts));
 
     let cos: Vec<f64> = quantized
         .encode(&texts)