gte 0.0.14-aarch64-linux → 0.0.16-aarch64-linux

data/ext/gte/src/reranker.rs

@@ -6,7 +6,7 @@ use crate::model_profile::{
 };
 use crate::pipeline::{extract_output_tensor, InputTensors};
 use crate::postprocess::sigmoid_scores;
-use crate::session::{build_session, SessionPool};
+use crate::session::{build_session, resolve_pool_size, SessionPool};
 use crate::tokenizer::{parse_padding_mode_override, Tokenizer};
 use std::path::{Path, PathBuf};
 
@@ -54,8 +54,6 @@ impl Reranker {
             with_attention_mask: true,
             optimization_level,
             execution_providers: overrides.execution_providers.map(str::to_string),
-            lowercase_input: false,
-            max_input_chars: None,
         };
         let session = build_session(&model_path, &probe_config)?;
 
@@ -83,10 +81,8 @@ impl Reranker {
             with_attention_mask: config.with_attention_mask,
             optimization_level,
             execution_providers: None,
-            lowercase_input: false,
-            max_input_chars: None,
         };
-        let pool = SessionPool::new(session, &model_path, &model_config)?;
+        let pool = SessionPool::new(&model_path, &model_config, resolve_pool_size())?;
         Ok(Self { tokenizer, pool, config })
     }
 
@@ -102,13 +98,12 @@ impl Reranker {
 
     fn score_tokenized(&self, tokenized: &crate::tokenizer::Tokenized, apply_sigmoid: bool) -> Result<Vec<f32>> {
         let input_tensors = InputTensors::from_tokenized(tokenized, self.config.with_attention_mask)?;
-        let output_name = self.config.output_tensor.clone();
         let inputs = input_tensors.inputs;
 
         self.pool.with_session(|session| {
             let outputs = session.run(inputs).map_err(|e| GteError::Ort(e.to_string()))?;
 
-            let array = extract_output_tensor(&outputs, output_name.as_str())?;
+            let array = extract_output_tensor(&outputs, self.config.output_tensor.as_str())?;
 
             let mut scores = match array.ndim() {
                 1 => array.into_dimensionality::<ndarray::Ix1>()?.to_vec(),
@@ -116,14 +111,16 @@ impl Reranker {
                     let shape = array.shape();
                     if shape[1] == 0 {
                         return Err(GteError::Inference(format!(
-                            "reranker output '{output_name}' has invalid shape {shape:?}"
+                            "reranker output '{}' has invalid shape {shape:?}",
+                            self.config.output_tensor
                         )));
                     }
                     array.slice(ndarray::s![.., 0]).to_vec()
                 }
                 n => {
                     return Err(GteError::Inference(format!(
-                        "reranker output '{output_name}' rank {n} is unsupported; expected rank 1 or 2"
+                        "reranker output '{}' rank {n} is unsupported; expected rank 1 or 2",
+                        self.config.output_tensor
                     )))
                 }
             };

data/ext/gte/src/ruby_embedder.rs

@@ -3,7 +3,7 @@
 #![allow(unused_results)]
 #![allow(unused_qualifications)]
 
-use crate::embedder::{normalize_l2, output_name_suggests_normalized, Embedder};
+use crate::embedder::Embedder;
 use crate::error::GteError;
 use crate::model_config::ModelLoadOverrides;
 use crate::reranker::Reranker;
@@ -15,7 +15,6 @@ use std::sync::Arc;
 #[wrap(class = "GTE::Embedder", free_immediately, size)]
 pub struct RbEmbedder {
     inner: Arc<Embedder>,
-    normalize: bool,
 }
 
 #[wrap(class = "GTE::Reranker", free_immediately, size)]
@@ -38,7 +37,6 @@ pub struct RbTensor {
 struct InferArgs {
     embedder: *const Embedder,
     texts: *const Vec<String>,
-    normalize: bool,
     result: Option<crate::error::Result<ndarray::Array2<f32>>>,
 }
 
@@ -66,15 +64,7 @@ fn panic_payload_to_string(payload: Box<dyn std::any::Any + Send>) -> String {
 
 unsafe extern "C" fn run_embed_without_gvl(ptr: *mut c_void) -> *mut c_void {
     let args = &mut *(ptr as *mut InferArgs);
-    let run_result = catch_unwind(AssertUnwindSafe(|| {
-        // Full embedding path (tokenization + inference) runs without the GVL.
-        let embeddings = (*args.embedder).embed_ref(&*args.texts)?;
-        if args.normalize {
-            Ok(normalize_l2(embeddings))
-        } else {
-            Ok(embeddings)
-        }
-    }));
+    let run_result = catch_unwind(AssertUnwindSafe(|| (*args.embedder).embed(&*args.texts)));
     args.result = Some(match run_result {
         Ok(result) => result,
         Err(payload) => {
@@ -97,14 +87,9 @@ unsafe extern "C" fn run_score_without_gvl(ptr: *mut c_void) -> *mut c_void {
     std::ptr::null_mut()
 }
 
-fn infer_without_gvl(
-    embedder: &Arc<Embedder>,
-    normalize: bool,
-    texts: Vec<String>,
-) -> Result<ndarray::Array2<f32>, Error> {
+fn infer_without_gvl(embedder: &Arc<Embedder>, texts: Vec<String>) -> Result<ndarray::Array2<f32>, Error> {
     let embeddings = unsafe {
-        let mut args =
-            InferArgs { embedder: Arc::as_ptr(embedder), texts: &texts as *const Vec<String>, normalize, result: None };
+        let mut args = InferArgs { embedder: Arc::as_ptr(embedder), texts: &texts as *const Vec<String>, result: None };
         rb_sys::rb_thread_call_without_gvl(
             Some(run_embed_without_gvl),
             &mut args as *mut InferArgs as *mut c_void,
@@ -167,13 +152,10 @@ impl RbEmbedder {
         dir_path: String,
         optimization_level: u8,
         model_name: String,
-        normalize: bool,
         output_tensor: String,
         max_length: usize,
         padding: String,
         execution_providers: String,
-        lowercase_input: bool,
-        max_input_chars: usize,
     ) -> Result<Self, Error> {
         let name = if model_name.is_empty() { None } else { Some(model_name.as_str()) };
         let output_override = if output_tensor.is_empty() { None } else { Some(output_tensor.as_str()) };
@@ -181,29 +163,26 @@ impl RbEmbedder {
         let execution_providers_override =
             if execution_providers.is_empty() { None } else { Some(execution_providers.as_str()) };
         let padding_override = if padding.is_empty() { None } else { Some(padding.as_str()) };
-        let max_input_chars_override = if max_input_chars == 0 { None } else { Some(max_input_chars) };
         let overrides = ModelLoadOverrides {
             model_name: name,
             output_tensor: output_override,
             max_length: max_length_override,
             padding: padding_override,
             execution_providers: execution_providers_override,
-            lowercase_input: Some(lowercase_input),
-            max_input_chars: max_input_chars_override,
+            ..ModelLoadOverrides::default()
         };
         let embedder = Embedder::from_dir(&dir_path, optimization_level, overrides).map_err(magnus::Error::from)?;
-        let skip_normalize = normalize && output_name_suggests_normalized(&embedder.config.output_tensor);
-        Ok(RbEmbedder { inner: Arc::new(embedder), normalize: normalize && !skip_normalize })
+        Ok(RbEmbedder { inner: Arc::new(embedder) })
     }
 
     pub fn rb_embed(_ruby: &Ruby, rb_self: &Self, texts: RArray) -> Result<RbTensor, Error> {
         let texts: Vec<String> = texts.to_vec()?;
-        let embeddings = infer_without_gvl(&rb_self.inner, rb_self.normalize, texts)?;
+        let embeddings = infer_without_gvl(&rb_self.inner, texts)?;
         tensor_from_array(embeddings)
     }
 
     pub fn rb_embed_one(_ruby: &Ruby, rb_self: &Self, text: String) -> Result<RbTensor, Error> {
-        let embeddings = infer_without_gvl(&rb_self.inner, rb_self.normalize, vec![text])?;
+        let embeddings = infer_without_gvl(&rb_self.inner, vec![text])?;
         tensor_from_array(embeddings)
     }
 }
@@ -219,8 +198,6 @@ impl RbReranker {
         max_length: usize,
         padding: String,
         execution_providers: String,
-        _lowercase_input: bool,
-        _max_input_chars: usize,
     ) -> Result<Self, Error> {
         let name = if model_name.is_empty() { None } else { Some(model_name.as_str()) };
         let output_override = if output_tensor.is_empty() { None } else { Some(output_tensor.as_str()) };
@@ -335,12 +312,12 @@ impl RbTensor {
 pub fn register(ruby: &Ruby) -> Result<(), Error> {
     let module = ruby.define_module("GTE")?;
     let embedder_class = module.define_class("Embedder", ruby.class_object())?;
-    embedder_class.define_singleton_method("new", function!(RbEmbedder::rb_new, 10))?;
+    embedder_class.define_singleton_method("new", function!(RbEmbedder::rb_new, 7))?;
     embedder_class.define_method("embed", method!(RbEmbedder::rb_embed, 1))?;
     embedder_class.define_method("embed_one", method!(RbEmbedder::rb_embed_one, 1))?;
 
     let reranker_class = module.define_class("Reranker", ruby.class_object())?;
-    reranker_class.define_singleton_method("new", function!(RbReranker::rb_new, 10))?;
+    reranker_class.define_singleton_method("new", function!(RbReranker::rb_new, 8))?;
     reranker_class.define_method("score", method!(RbReranker::rb_score, 2))?;
 
     let tensor_class = module.define_class("Tensor", ruby.class_object())?;

data/ext/gte/src/session.rs

@@ -3,77 +3,58 @@ use crate::model_config::{ExtractorMode, ModelConfig};
 use crate::pipeline::{extract_output_tensor, InputTensors};
 use crate::postprocess::mean_pool;
 use crate::tokenizer::Tokenized;
-use ndarray::{Array2, ArrayView2, ArrayViewD, Ix2};
+use ndarray::{Array2, ArrayViewD, Ix2};
 use ort::execution_providers::{CoreMLExecutionProvider, ExecutionProviderDispatch, XNNPACKExecutionProvider};
 use ort::session::{OutputSelector, RunOptions, Session};
-use std::cell::RefCell;
-use std::collections::hash_map::Entry;
-use std::collections::HashMap;
-use std::path::{Path, PathBuf};
+use parking_lot::Mutex;
+use std::path::Path;
 use std::sync::atomic::{AtomicUsize, Ordering};
 
-// ---------------------------------------------------------------------------
-// Thread-local session storage — each OS thread lazily creates its own ONNX
-// session the first time it calls into a given pool.  No Mutex, no contention.
-// ---------------------------------------------------------------------------
-
-static NEXT_POOL_ID: AtomicUsize = AtomicUsize::new(1);
-
-struct SessionRecipe {
-    model_path: PathBuf,
-    build_config: ModelConfig,
-}
-
-thread_local! {
-    static SESSIONS: RefCell<HashMap<usize, Session>> = RefCell::new(HashMap::new());
+pub(crate) fn resolve_pool_size() -> usize {
+    if let Some(n) =
+        std::env::var("GTE_SESSION_POOL_SIZE").ok().and_then(|v| v.trim().parse::<usize>().ok()).filter(|&n| n > 0)
+    {
+        return n;
+    }
+    let cpus = std::thread::available_parallelism().map(std::num::NonZero::get).unwrap_or(2);
+    cpus.min(4).max(1)
 }
 
 pub struct SessionPool {
-    pool_id: usize,
-    recipe: SessionRecipe,
+    sessions: Vec<Mutex<Session>>,
+    next_idx: AtomicUsize,
 }
 
 impl SessionPool {
-    pub fn new(initial: Session, model_path: &Path, build_config: &ModelConfig) -> Result<Self> {
-        let pool_id = NEXT_POOL_ID.fetch_add(1, Ordering::Relaxed);
-
-        SESSIONS.with(|map| {
-            _ = map.borrow_mut().insert(pool_id, initial);
-        });
-
-        Ok(Self {
-            pool_id,
-            recipe: SessionRecipe { model_path: model_path.to_path_buf(), build_config: build_config.clone() },
-        })
-    }
-
-    pub fn run(&self, tokenized: &Tokenized, config: &ModelConfig) -> Result<Array2<f32>> {
-        self.with_session(|session| run_session(session, tokenized, config))
+    pub fn new(model_path: &Path, config: &ModelConfig, pool_size: usize) -> Result<Self> {
+        let sessions = (0..pool_size)
+            .map(|_| build_session(model_path, config))
+            .collect::<Result<Vec<_>>>()?
+            .into_iter()
+            .map(Mutex::new)
+            .collect();
+        Ok(Self { sessions, next_idx: AtomicUsize::new(0) })
     }
 
     pub fn with_session<F, R>(&self, f: F) -> Result<R>
     where
         F: FnOnce(&mut Session) -> Result<R>,
     {
-        SESSIONS.with(|map| {
-            let mut map = map.borrow_mut();
-            let session = match map.entry(self.pool_id) {
-                Entry::Occupied(e) => e.into_mut(),
-                Entry::Vacant(e) => {
-                    let session = build_session(&self.recipe.model_path, &self.recipe.build_config)?;
-                    e.insert(session)
-                }
-            };
-            f(session)
-        })
+        let idx = if self.sessions.len() == 1 {
+            0
+        } else {
+            self.next_idx.fetch_add(1, Ordering::Relaxed) % self.sessions.len()
+        };
+        let mut session = self.sessions[idx].lock();
+        f(&mut session)
     }
-}
 
-// ---------------------------------------------------------------------------
-// Session construction
-// ---------------------------------------------------------------------------
+    pub fn len(&self) -> usize {
+        self.sessions.len()
+    }
+}
 
-pub fn build_session<P: AsRef<Path>>(model_path: P, config: &ModelConfig) -> Result<Session> {
+pub(crate) fn build_session<P: AsRef<Path>>(model_path: P, config: &ModelConfig) -> Result<Session> {
     fn ort_err(e: impl std::fmt::Display) -> GteError {
         GteError::Ort(e.to_string())
     }
@@ -109,10 +90,14 @@ pub fn build_session<P: AsRef<Path>>(model_path: P, config: &ModelConfig) -> Res
 }
 
 fn auto_detect_providers() -> Vec<ExecutionProviderDispatch> {
-    let mut providers = Vec::new();
     #[cfg(target_arch = "aarch64")]
-    providers.push(XNNPACKExecutionProvider::default().build().fail_silently());
-    providers
+    {
+        vec![XNNPACKExecutionProvider::default().build().fail_silently()]
+    }
+    #[cfg(not(target_arch = "aarch64"))]
+    {
+        Vec::new()
+    }
 }
 
 fn preferred_execution_providers(order_override: Option<&str>) -> Vec<ExecutionProviderDispatch> {
@@ -138,11 +123,7 @@ fn preferred_execution_providers(order_override: Option<&str>) -> Vec<ExecutionP
     providers
 }
 
-// ---------------------------------------------------------------------------
-// Run a single inference
-// ---------------------------------------------------------------------------
-
-pub fn run_session(session: &mut Session, tokenized: &Tokenized, config: &ModelConfig) -> Result<Array2<f32>> {
+pub(crate) fn run_session(session: &mut Session, tokenized: &Tokenized, config: &ModelConfig) -> Result<Array2<f32>> {
     let input_tensors = InputTensors::from_tokenized(tokenized, config.with_attention_mask)?;
     let run_opts = RunOptions::new()
         .map_err(|e| GteError::Ort(e.to_string()))?
@@ -156,7 +137,7 @@ pub fn run_session(session: &mut Session, tokenized: &Tokenized, config: &ModelC
 
 fn extract_embeddings(
     array: ArrayViewD<'_, f32>,
-    attention_mask: ArrayView2<'_, i64>,
+    attention_mask: ndarray::ArrayView2<'_, i64>,
     config: &ModelConfig,
 ) -> Result<Array2<f32>> {
     match config.mode {
@@ -189,21 +170,6 @@ mod tests {
 
     use super::extract_embeddings;
 
-    fn resolve_provider_order_with_env(order_override: Option<&str>, env_order: Option<&str>) -> String {
-        order_override.or(env_order).unwrap_or("cpu").to_ascii_lowercase()
-    }
-
-    fn parse_provider_registrations(order: &str) -> Vec<&str> {
-        let mut providers = Vec::new();
-        for provider in order.split(',').map(str::trim).filter(|p| !p.is_empty()) {
-            match provider {
-                "xnnpack" | "coreml" => providers.push(provider),
-                _ => {}
-            }
-        }
-        providers
-    }
-
     fn test_config(mode: ExtractorMode) -> ModelConfig {
         ModelConfig {
             max_length: 8,
@@ -214,8 +180,6 @@ mod tests {
             with_attention_mask: true,
             optimization_level: 3,
             execution_providers: None,
-            lowercase_input: false,
-            max_input_chars: None,
         }
     }
 
@@ -224,37 +188,6 @@ mod tests {
         ArrayView2::from_shape((0, 0), &EMPTY).unwrap()
     }
 
-    #[test]
-    fn parse_provider_registrations_keeps_supported_order() {
-        let parsed = parse_provider_registrations("xnnpack,coreml");
-        assert_eq!(parsed, vec!["xnnpack", "coreml"]);
-    }
-
-    #[test]
-    fn parse_provider_registrations_treats_cpu_and_none_as_fallback() {
-        assert!(parse_provider_registrations("cpu").is_empty());
-        assert!(parse_provider_registrations("none").is_empty());
-        assert!(parse_provider_registrations("none,cpu").is_empty());
-    }
-
-    #[test]
-    fn parse_provider_registrations_ignores_unknowns_and_empties() {
-        let parsed = parse_provider_registrations(" ,xnnpak,,xnnpack,unknown,coreml,");
-        assert_eq!(parsed, vec!["xnnpack", "coreml"]);
-    }
-
-    #[test]
-    fn resolve_provider_order_prefers_override() {
-        assert_eq!(resolve_provider_order_with_env(Some("xnnpack"), Some("coreml")), "xnnpack");
-        assert_eq!(resolve_provider_order_with_env(Some("CPU"), None), "cpu");
-    }
-
-    #[test]
-    fn resolve_provider_order_falls_back_to_env_then_cpu_default() {
-        assert_eq!(resolve_provider_order_with_env(None, Some("coreml")), "coreml");
-        assert_eq!(resolve_provider_order_with_env(None, None), "cpu");
-    }
-
     #[test]
     fn extract_embeddings_raw_copies_only_final_matrix() {
         let output = array![[1.0f32, 2.0], [3.0, 4.0]];
@@ -287,4 +220,20 @@ mod tests {
 
         assert_eq!(extracted, expected);
     }
+
+    #[test]
+    fn resolve_pool_size_uses_env_var() {
+        std::env::set_var("GTE_SESSION_POOL_SIZE", "16");
+        let size = super::resolve_pool_size();
+        assert_eq!(size, 16);
+        std::env::remove_var("GTE_SESSION_POOL_SIZE");
+    }
+
+    #[test]
+    fn resolve_pool_size_defaults_to_cpu_count_capped_at_4() {
+        // Without GTE_SESSION_POOL_SIZE, the default is min(available_parallelism, 4).max(1).
+        // On any machine with >= 1 CPU, this should return between 1 and 4.
+        let size = super::resolve_pool_size();
+        assert!((1..=4).contains(&size), "expected 1-4, got {size}");
+    }
 }

data/ext/gte/src/tokenizer.rs

@@ -1,14 +1,13 @@
 use crate::error::{GteError, Result};
 use crate::model_config::PaddingMode;
+use ndarray::Array2;
 use std::path::Path;
 use tokenizers::{PaddingParams, PaddingStrategy, TruncationParams};
 
 pub struct Tokenized {
-    pub rows: usize,
-    pub cols: usize,
-    pub input_ids: Vec<i64>,
-    pub attn_masks: Vec<i64>,
-    pub type_ids: Option<Vec<i64>>,
+    pub input_ids: Array2<i64>,
+    pub attn_masks: Array2<i64>,
+    pub type_ids: Option<Array2<i64>>,
 }
 
 pub struct Tokenizer {
@@ -24,7 +23,6 @@ impl Tokenizer {
         padding_mode: PaddingMode,
         fixed_padding_length: Option<usize>,
     ) -> Result<Self> {
-        #[allow(unused_results)]
         {
             let mut tokenizer =
                 tokenizers::Tokenizer::from_file(tokenizer_path).map_err(|e| GteError::Tokenizer(e.to_string()))?;
@@ -34,41 +32,59 @@ impl Tokenizer {
                 strategy: resolve_padding_strategy(padding_mode, max_length, fixed_padding_length),
                 ..Default::default()
             };
-            tokenizer.with_truncation(Some(truncation)).map_err(|e| GteError::Tokenizer(e.to_string()))?;
-            tokenizer.with_padding(Some(padding));
+            let _ = tokenizer.with_truncation(Some(truncation)).map_err(|e| GteError::Tokenizer(e.to_string()))?;
+            let _ = tokenizer.with_padding(Some(padding));
 
             Ok(Self { tokenizer, with_type_ids })
         }
     }
 
     pub fn tokenize(&self, texts: &[String]) -> Result<Tokenized> {
-        if texts.len() == 1 {
-            let encoding =
-                self.tokenizer.encode_fast(texts[0].as_str(), true).map_err(|e| GteError::Tokenizer(e.to_string()))?;
-            return Ok(build_tokenized_single(&encoding, self.with_type_ids));
+        if texts.is_empty() {
+            return Ok(Tokenized {
+                input_ids: Array2::zeros((0, 0)),
+                attn_masks: Array2::zeros((0, 0)),
+                type_ids: None,
+            });
         }
 
         let encode_inputs: Vec<&str> = texts.iter().map(String::as_str).collect();
         let encodings =
             self.tokenizer.encode_batch_fast(encode_inputs, true).map_err(|e| GteError::Tokenizer(e.to_string()))?;
 
-        Ok(build_tokenized(&encodings, self.with_type_ids))
+        build_tokenized(&encodings, self.with_type_ids)
     }
 
     pub fn tokenize_pairs(&self, pairs: &[(String, String)]) -> Result<Tokenized> {
+        if pairs.is_empty() {
+            return Ok(Tokenized {
+                input_ids: Array2::zeros((0, 0)),
+                attn_masks: Array2::zeros((0, 0)),
+                type_ids: None,
+            });
+        }
+
         let encode_inputs: Vec<tokenizers::EncodeInput<'_>> =
             pairs.iter().map(|(left, right)| (left.as_str(), right.as_str()).into()).collect();
         let encodings =
             self.tokenizer.encode_batch_fast(encode_inputs, true).map_err(|e| GteError::Tokenizer(e.to_string()))?;
-        Ok(build_tokenized(&encodings, self.with_type_ids))
+        build_tokenized(&encodings, self.with_type_ids)
     }
 
     pub fn tokenize_query_candidates(&self, query: &str, candidates: &[String]) -> Result<Tokenized> {
+        if candidates.is_empty() {
+            return Ok(Tokenized {
+                input_ids: Array2::zeros((0, 0)),
+                attn_masks: Array2::zeros((0, 0)),
+                type_ids: None,
+            });
+        }
+
         let encode_inputs: Vec<tokenizers::EncodeInput<'_>> =
             candidates.iter().map(|candidate| (query, candidate.as_str()).into()).collect();
         let encodings =
             self.tokenizer.encode_batch_fast(encode_inputs, true).map_err(|e| GteError::Tokenizer(e.to_string()))?;
-        Ok(build_tokenized(&encodings, self.with_type_ids))
+        build_tokenized(&encodings, self.with_type_ids)
     }
 }
 
@@ -102,36 +118,30 @@ fn resolve_padding_strategy(
     }
 }
 
-fn build_tokenized_single(encoding: &tokenizers::Encoding, with_type_ids: bool) -> Tokenized {
-    let cols = encoding.len();
-
-    let input_ids: Vec<i64> = encoding.get_ids().iter().map(|&v| i64::from(v)).collect();
-    let attn_masks: Vec<i64> = encoding.get_attention_mask().iter().map(|&v| i64::from(v)).collect();
-    let type_ids: Option<Vec<i64>> =
-        with_type_ids.then(|| encoding.get_type_ids().iter().map(|&v| i64::from(v)).collect());
-
-    Tokenized { rows: 1, cols, input_ids, attn_masks, type_ids }
+fn to_i64(array: &[u32]) -> Vec<i64> {
+    array.iter().map(|&v| v as i64).collect()
 }
 
-fn build_tokenized(encodings: &[tokenizers::Encoding], with_type_ids: bool) -> Tokenized {
+fn build_tokenized(encodings: &[tokenizers::Encoding], with_type_ids: bool) -> Result<Tokenized> {
     let rows = encodings.len();
     let cols = encodings.first().map_or(0, tokenizers::Encoding::len);
-    let len = rows * cols;
+    if rows == 0 || cols == 0 {
+        return Ok(Tokenized { input_ids: Array2::zeros((0, 0)), attn_masks: Array2::zeros((0, 0)), type_ids: None });
+    }
 
-    let mut input_ids = Vec::with_capacity(len);
-    let mut attn_masks = Vec::with_capacity(len);
-    let mut type_ids = with_type_ids.then(|| Vec::with_capacity(len));
+    let mut input_ids = Array2::zeros((0, cols));
+    let mut attn_masks = Array2::zeros((0, cols));
+    let mut type_ids = with_type_ids.then(|| Array2::zeros((0, cols)));
 
     for encoding in encodings {
-        input_ids.extend(encoding.get_ids().iter().map(|&v| i64::from(v)));
-        attn_masks.extend(encoding.get_attention_mask().iter().map(|&v| i64::from(v)));
-
-        if let Some(type_ids) = type_ids.as_mut() {
-            type_ids.extend(encoding.get_type_ids().iter().map(|&v| i64::from(v)));
+        input_ids.push_row(ndarray::ArrayView::from(&to_i64(encoding.get_ids())))?;
+        attn_masks.push_row(ndarray::ArrayView::from(&to_i64(encoding.get_attention_mask())))?;
+        if let Some(ref mut type_ids) = type_ids {
+            type_ids.push_row(ndarray::ArrayView::from(&to_i64(encoding.get_type_ids())))?;
         }
     }
 
-    Tokenized { rows, cols, input_ids, attn_masks, type_ids }
+    Ok(Tokenized { input_ids, attn_masks, type_ids })
 }
 
 #[cfg(test)]
@@ -154,9 +164,6 @@ mod tests {
 
     #[test]
     fn resolve_padding_strategy_auto_always_uses_batch_longest() {
-        // Auto ignores fixed_padding_length from tokenizer.json — BatchLongest is
-        // always faster for inference and correct for variable-length inputs.
-        // Use PaddingMode::Fixed explicitly when fixed-length padding is required.
         assert!(matches!(resolve_padding_strategy(PaddingMode::Auto, 64, Some(64)), PaddingStrategy::BatchLongest));
         assert!(matches!(resolve_padding_strategy(PaddingMode::Auto, 512, None), PaddingStrategy::BatchLongest));
     }

data/ext/gte/tests/embedder_unit_test.rs

@@ -1,4 +1,4 @@
-use gte::embedder::normalize_l2;
+use gte::postprocess::normalize_l2;
 use ndarray::array;
 
 #[test]