gte 0.0.6 → 0.0.8

data/ext/gte/src/session.rs

@@ -3,12 +3,14 @@ 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, Ix2};
+use ndarray::{Array2, ArrayView2, ArrayViewD, Ix2};
 use ort::execution_providers::{
     CoreMLExecutionProvider, ExecutionProviderDispatch, XNNPACKExecutionProvider,
 };
 use ort::session::Session;
-use std::path::Path;
+use std::path::{Path, PathBuf};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Condvar, Mutex};
 
 pub fn build_session<P: AsRef<Path>>(model_path: P, config: &ModelConfig) -> Result<Session> {
     let opt_level = match config.optimization_level {
@@ -18,22 +20,176 @@ pub fn build_session<P: AsRef<Path>>(model_path: P, config: &ModelConfig) -> Res
         _ => ort::session::builder::GraphOptimizationLevel::Level3,
     };
 
-    let mut builder = Session::builder()?
-        .with_optimization_level(opt_level)?
-        .with_memory_pattern(true)?;
+    fn ort_err(e: impl std::fmt::Display) -> GteError {
+        GteError::Ort(e.to_string())
+    }
+
+    let mut builder = Session::builder()
+        .map_err(ort_err)?
+        .with_optimization_level(opt_level)
+        .map_err(ort_err)?
+        .with_memory_pattern(true)
+        .map_err(ort_err)?;
 
     let providers = preferred_execution_providers(config.execution_providers.as_deref());
     if !providers.is_empty() {
-        builder = builder.with_execution_providers(providers)?;
+        builder = builder
+            .with_execution_providers(providers)
+            .map_err(ort_err)?;
     }
 
     if config.num_threads > 0 {
-        builder = builder.with_intra_threads(config.num_threads)?;
+        builder = builder
+            .with_intra_threads(config.num_threads)
+            .map_err(ort_err)?;
+        builder = builder
+            .with_inter_threads(config.num_threads)
+            .map_err(ort_err)?;
+    }
+
+    builder.commit_from_file(model_path).map_err(ort_err)
+}
+
+// ---------------------------------------------------------------------------
+// Session pool
+// ---------------------------------------------------------------------------
+
+const AUTO_THREAD_POOL_CAP: usize = 6;
+
+/// Keep enough sessions to cover the configured thread budget without
+/// oversubscribing CPU parallelism. In ORT auto-thread mode (`num_threads == 0`)
+/// we still keep a modest pool because request-level concurrency benefits from
+/// more than one session even when ORT manages thread counts internally.
+fn pool_capacity(num_threads: usize) -> usize {
+    let available_parallelism = std::thread::available_parallelism()
+        .map(|n| n.get())
+        .unwrap_or(1);
+    pool_capacity_with_parallelism(num_threads, available_parallelism)
+}
+
+fn pool_capacity_with_parallelism(num_threads: usize, available_parallelism: usize) -> usize {
+    if available_parallelism == 0 {
+        return 1;
+    }
+
+    if num_threads == 0 {
+        return available_parallelism.clamp(1, AUTO_THREAD_POOL_CAP);
+    }
+
+    available_parallelism.div_ceil(num_threads).max(1)
+}
+
+pub struct SessionPool {
+    sessions: Mutex<Vec<Session>>,
+    available: Condvar,
+    created: AtomicUsize,
+    capacity: usize,
+    model_path: PathBuf,
+    build_config: ModelConfig,
+}
+
+impl SessionPool {
+    pub fn new(initial: Session, model_path: PathBuf, build_config: ModelConfig) -> Self {
+        let capacity = pool_capacity(build_config.num_threads);
+        Self {
+            sessions: Mutex::new(vec![initial]),
+            available: Condvar::new(),
+            created: AtomicUsize::new(1),
+            capacity,
+            model_path,
+            build_config,
+        }
+    }
+
+    pub fn acquire(&self) -> Result<PooledSession<'_>> {
+        if let Some(session) = self.take_available_session() {
+            return Ok(PooledSession {
+                pool: self,
+                session: Some(session),
+            });
+        }
+
+        if let Some(session) = self.try_grow()? {
+            return Ok(PooledSession {
+                pool: self,
+                session: Some(session),
+            });
+        }
+
+        let session = self.wait_for_session();
+        Ok(PooledSession {
+            pool: self,
+            session: Some(session),
+        })
+    }
+
+    fn release(&self, session: Session) {
+        self.sessions.lock().unwrap().push(session);
+        self.available.notify_one();
+    }
+
+    fn take_available_session(&self) -> Option<Session> {
+        self.sessions.lock().unwrap().pop()
+    }
+
+    fn try_grow(&self) -> Result<Option<Session>> {
+        let grew = self
+            .created
+            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |count| {
+                (count < self.capacity).then_some(count + 1)
+            });
+        if grew.is_err() {
+            return Ok(None);
+        }
+
+        match build_session(&self.model_path, &self.build_config) {
+            Ok(session) => Ok(Some(session)),
+            Err(error) => {
+                self.created.fetch_sub(1, Ordering::AcqRel);
+                Err(error)
+            }
+        }
+    }
+
+    fn wait_for_session(&self) -> Session {
+        let mut lock = self.sessions.lock().unwrap();
+        loop {
+            if let Some(session) = lock.pop() {
+                return session;
+            }
+            lock = self.available.wait(lock).unwrap();
+        }
+    }
+}
+
+pub struct PooledSession<'a> {
+    pool: &'a SessionPool,
+    session: Option<Session>,
+}
+
+impl std::ops::Deref for PooledSession<'_> {
+    type Target = Session;
+    fn deref(&self) -> &Session {
+        self.session.as_ref().unwrap()
     }
+}
 
-    Ok(builder.commit_from_file(model_path)?)
+impl std::ops::DerefMut for PooledSession<'_> {
+    fn deref_mut(&mut self) -> &mut Session {
+        self.session.as_mut().unwrap()
+    }
 }
 
+impl Drop for PooledSession<'_> {
+    fn drop(&mut self) {
+        if let Some(s) = self.session.take() {
+            self.pool.release(s);
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+
 fn preferred_execution_providers(order_override: Option<&str>) -> Vec<ExecutionProviderDispatch> {
     let order = resolve_provider_order(order_override);
 
@@ -55,7 +211,10 @@ fn resolve_provider_order(order_override: Option<&str>) -> String {
     resolve_provider_order_with_env(order_override, env_order.as_deref())
 }
 
-fn resolve_provider_order_with_env(order_override: Option<&str>, env_order: Option<&str>) -> String {
+fn resolve_provider_order_with_env(
+    order_override: Option<&str>,
+    env_order: Option<&str>,
+) -> String {
     order_override
         .or(env_order)
         .unwrap_or("cpu")
@@ -75,14 +234,24 @@ fn parse_provider_registrations(order: &str) -> Vec<&str> {
 }
 
 pub fn run_session(
-    session: &Session,
+    session: &mut Session,
     tokenized: &Tokenized,
     config: &ModelConfig,
 ) -> Result<Array2<f32>> {
     let input_tensors = InputTensors::from_tokenized(tokenized, config.with_attention_mask)?;
-    let outputs = session.run(input_tensors.inputs)?;
+    let outputs = session
+        .run(input_tensors.inputs)
+        .map_err(|e| GteError::Ort(e.to_string()))?;
     let array = extract_output_tensor(&outputs, config.output_tensor.as_str())?;
 
+    extract_embeddings(array, input_tensors.attention_mask, config)
+}
+
+fn extract_embeddings(
+    array: ArrayViewD<'_, f32>,
+    attention_mask: ArrayView2<'_, i64>,
+    config: &ModelConfig,
+) -> Result<Array2<f32>> {
     match config.mode {
         ExtractorMode::Token(idx) => {
             let shape = array.shape();
@@ -102,15 +271,43 @@ pub fn run_session(
                     ndim
                 ))
             })?;
-            mean_pool(hidden_states.view(), input_tensors.attention_mask)
+            mean_pool(hidden_states, attention_mask)
         }
-        ExtractorMode::Raw => Ok(array.into_dimensionality::<Ix2>()?.into_owned()),
+        ExtractorMode::Raw => array
+            .into_dimensionality::<Ix2>()
+            .map(|view| view.to_owned())
+            .map_err(|e| GteError::Shape(e.to_string())),
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use super::{parse_provider_registrations, resolve_provider_order_with_env};
+    use crate::model_config::{ExtractorMode, ModelConfig, PaddingMode};
+    use ndarray::{array, ArrayView2};
+
+    use super::{
+        extract_embeddings, parse_provider_registrations, pool_capacity_with_parallelism,
+        resolve_provider_order_with_env,
+    };
+
+    fn test_config(mode: ExtractorMode) -> ModelConfig {
+        ModelConfig {
+            max_length: 8,
+            padding_mode: PaddingMode::BatchLongest,
+            output_tensor: "output".to_string(),
+            mode,
+            with_type_ids: false,
+            with_attention_mask: true,
+            num_threads: 1,
+            optimization_level: 3,
+            execution_providers: None,
+        }
+    }
+
+    fn empty_attention_mask() -> ArrayView2<'static, i64> {
+        static EMPTY: [i64; 0] = [];
+        ArrayView2::from_shape((0, 0), &EMPTY).unwrap()
+    }
 
     #[test]
     fn parse_provider_registrations_keeps_supported_order() {
@@ -142,7 +339,74 @@ mod tests {
 
     #[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, Some("coreml")),
+            "coreml"
+        );
         assert_eq!(resolve_provider_order_with_env(None, None), "cpu");
     }
+
+    #[test]
+    fn pool_capacity_uses_bounded_parallel_pool_for_auto_thread_mode() {
+        assert_eq!(pool_capacity_with_parallelism(0, 1), 1);
+        assert_eq!(pool_capacity_with_parallelism(0, 4), 4);
+        assert_eq!(pool_capacity_with_parallelism(0, 8), 6);
+    }
+
+    #[test]
+    fn pool_capacity_scales_with_available_parallelism() {
+        assert_eq!(pool_capacity_with_parallelism(1, 1), 1);
+        assert_eq!(pool_capacity_with_parallelism(1, 8), 8);
+        assert_eq!(pool_capacity_with_parallelism(2, 8), 4);
+        assert_eq!(pool_capacity_with_parallelism(3, 8), 3);
+        assert_eq!(pool_capacity_with_parallelism(8, 4), 1);
+    }
+
+    #[test]
+    fn extract_embeddings_raw_copies_only_final_matrix() {
+        let output = array![[1.0f32, 2.0], [3.0, 4.0]];
+        let extracted = extract_embeddings(
+            output.view().into_dyn(),
+            empty_attention_mask(),
+            &test_config(ExtractorMode::Raw),
+        )
+        .unwrap();
+
+        assert_eq!(extracted, output);
+    }
+
+    #[test]
+    fn extract_embeddings_token_selects_without_copying_full_sequence() {
+        let output = array![
+            [[1.0f32, 2.0], [3.0, 4.0], [5.0, 6.0]],
+            [[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]]
+        ];
+        let expected = array![[3.0f32, 4.0], [9.0, 10.0]];
+        let extracted = extract_embeddings(
+            output.view().into_dyn(),
+            empty_attention_mask(),
+            &test_config(ExtractorMode::Token(1)),
+        )
+        .unwrap();
+
+        assert_eq!(extracted, expected);
+    }
+
+    #[test]
+    fn extract_embeddings_mean_pool_uses_output_view_and_attention_mask() {
+        let output = array![
+            [[1.0f32, 3.0], [5.0, 7.0], [100.0, 100.0]],
+            [[2.0, 4.0], [6.0, 8.0], [10.0, 12.0]]
+        ];
+        let attention_mask = array![[1_i64, 1, 0], [0, 1, 1]];
+        let expected = array![[3.0f32, 5.0], [8.0, 10.0]];
+        let extracted = extract_embeddings(
+            output.view().into_dyn(),
+            attention_mask.view(),
+            &test_config(ExtractorMode::MeanPool),
+        )
+        .unwrap();
+
+        assert_eq!(extracted, expected);
+    }
 }

data/ext/gte/src/tokenizer.rs

@@ -1,4 +1,5 @@
 use crate::error::{GteError, Result};
+use crate::model_config::PaddingMode;
 use std::path::Path;
 use tokenizers::{PaddingParams, PaddingStrategy, TruncationParams};
 
@@ -20,6 +21,8 @@ impl Tokenizer {
         tokenizer_path: P,
         max_length: usize,
         with_type_ids: bool,
+        padding_mode: PaddingMode,
+        fixed_padding_length: Option<usize>,
     ) -> Result<Self> {
         let mut tokenizer = tokenizers::Tokenizer::from_file(tokenizer_path)
             .map_err(|e| GteError::Tokenizer(e.to_string()))?;
@@ -33,7 +36,7 @@ impl Tokenizer {
             .map_err(|e| GteError::Tokenizer(e.to_string()))?;
 
         let padding = PaddingParams {
-            strategy: PaddingStrategy::BatchLongest,
+            strategy: resolve_padding_strategy(padding_mode, max_length, fixed_padding_length),
             ..Default::default()
         };
         tokenizer.with_padding(Some(padding));
@@ -73,6 +76,56 @@ impl Tokenizer {
             .map_err(|e| GteError::Tokenizer(e.to_string()))?;
         build_tokenized(&encodings, self.with_type_ids)
     }
+
+    pub fn tokenize_query_candidates(&self, query: &str, candidates: &[String]) -> Result<Tokenized> {
+        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()))?;
+        build_tokenized(&encodings, self.with_type_ids)
+    }
+}
+
+pub fn parse_padding_mode_override(value: Option<&str>) -> Result<Option<PaddingMode>> {
+    let Some(raw) = value.map(str::trim).filter(|v| !v.is_empty()) else {
+        return Ok(None);
+    };
+
+    let normalized = raw.to_ascii_lowercase().replace('-', "_");
+    let parsed = match normalized.as_str() {
+        "auto" => PaddingMode::Auto,
+        "batch_longest" | "batchlongest" => PaddingMode::BatchLongest,
+        "fixed" => PaddingMode::Fixed,
+        _ => {
+            return Err(GteError::Inference(format!(
+                "invalid padding mode '{}'; expected one of: auto, batch_longest, fixed",
+                raw
+            )))
+        }
+    };
+    Ok(Some(parsed))
+}
+
+fn resolve_padding_strategy(
+    padding_mode: PaddingMode,
+    max_length: usize,
+    fixed_padding_length: Option<usize>,
+) -> PaddingStrategy {
+    match padding_mode {
+        PaddingMode::BatchLongest => PaddingStrategy::BatchLongest,
+        PaddingMode::Fixed => PaddingStrategy::Fixed(max_length),
+        PaddingMode::Auto => {
+            if fixed_padding_length.is_some() {
+                PaddingStrategy::Fixed(max_length)
+            } else {
+                PaddingStrategy::BatchLongest
+            }
+        }
+    }
 }
 
 fn build_tokenized_single(
@@ -121,21 +174,17 @@ fn build_tokenized(encodings: &[tokenizers::Encoding], with_type_ids: bool) -> R
     let mut type_ids = with_type_ids.then(|| Vec::with_capacity(len));
 
     for encoding in encodings {
-        input_ids.extend(encoding.get_ids().iter().map(|&value| i64::from(value)));
-        attn_masks.extend(
-            encoding
-                .get_attention_mask()
-                .iter()
-                .map(|&value| i64::from(value)),
-        );
+        for &value in encoding.get_ids() {
+            input_ids.push(i64::from(value));
+        }
+        for &value in encoding.get_attention_mask() {
+            attn_masks.push(i64::from(value));
+        }
 
         if let Some(type_ids) = type_ids.as_mut() {
-            type_ids.extend(
-                encoding
-                    .get_type_ids()
-                    .iter()
-                    .map(|&value| i64::from(value)),
-            );
+            for &value in encoding.get_type_ids() {
+                type_ids.push(i64::from(value));
+            }
         }
     }
 
@@ -147,3 +196,39 @@ fn build_tokenized(encodings: &[tokenizers::Encoding], with_type_ids: bool) -> R
         type_ids,
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{parse_padding_mode_override, resolve_padding_strategy};
+    use crate::model_config::PaddingMode;
+    use tokenizers::PaddingStrategy;
+
+    #[test]
+    fn parse_padding_mode_override_accepts_expected_values() {
+        assert_eq!(
+            parse_padding_mode_override(Some("auto")).unwrap(),
+            Some(PaddingMode::Auto)
+        );
+        assert_eq!(
+            parse_padding_mode_override(Some("batch-longest")).unwrap(),
+            Some(PaddingMode::BatchLongest)
+        );
+        assert_eq!(
+            parse_padding_mode_override(Some("fixed")).unwrap(),
+            Some(PaddingMode::Fixed)
+        );
+    }
+
+    #[test]
+    fn parse_padding_mode_override_rejects_invalid_values() {
+        assert!(parse_padding_mode_override(Some("unknown")).is_err());
+    }
+
+    #[test]
+    fn resolve_padding_strategy_uses_fixed_for_auto_when_model_has_fixed_padding() {
+        match resolve_padding_strategy(PaddingMode::Auto, 64, Some(64)) {
+            PaddingStrategy::Fixed(64) => {}
+            other => panic!("expected Fixed(64), got {:?}", other),
+        }
+    }
+}

data/ext/gte/tests/inference_integration_test.rs

@@ -1,11 +1,12 @@
 use gte::embedder::Embedder;
+use gte::model_config::ModelLoadOverrides;
 
 #[test]
 #[ignore = "requires ext/gte/tests/fixtures/e5/tokenizer.json and model.onnx"]
 fn test_e5_single_embedding_shape() {
     const DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/e5");
 
-    let embedder = Embedder::from_dir(DIR, 0, 3, None, None, None, None)
+    let embedder = Embedder::from_dir(DIR, 0, 3, ModelLoadOverrides::default())
         .expect("embedder should initialize");
     let result = embedder
         .embed(vec!["query: Hello world".to_string()])
@@ -20,7 +21,7 @@ fn test_e5_single_embedding_shape() {
 fn test_clip_single_embedding_shape() {
     const DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/clip");
 
-    let embedder = Embedder::from_dir(DIR, 0, 3, None, None, None, None)
+    let embedder = Embedder::from_dir(DIR, 0, 3, ModelLoadOverrides::default())
         .expect("embedder should initialize");
     let result = embedder
         .embed(vec!["a photo of a cat".to_string()])
@@ -35,7 +36,7 @@ fn test_clip_single_embedding_shape() {
 fn test_e5_batch_embedding_shape() {
     const DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/e5");
 
-    let embedder = Embedder::from_dir(DIR, 0, 3, None, None, None, None)
+    let embedder = Embedder::from_dir(DIR, 0, 3, ModelLoadOverrides::default())
         .expect("embedder should initialize");
     let texts = vec![
         "query: first sentence".to_string(),
@@ -54,7 +55,7 @@ fn test_e5_batch_embedding_shape() {
 fn test_e5_long_input_truncation_no_error() {
     const DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/e5");
 
-    let embedder = Embedder::from_dir(DIR, 0, 3, None, None, None, None)
+    let embedder = Embedder::from_dir(DIR, 0, 3, ModelLoadOverrides::default())
         .expect("embedder should initialize");
     let very_long_text = "word ".repeat(1000);
     let result = embedder

data/ext/gte/tests/tokenizer_unit_test.rs

@@ -1,3 +1,4 @@
+use gte::model_config::PaddingMode;
 use gte::tokenizer::Tokenizer;
 
 #[test]
@@ -8,7 +9,8 @@ fn test_e5_tokenizer_output_shape() {
         "/tests/fixtures/e5/tokenizer.json"
     );
 
-    let tokenizer = Tokenizer::new(TOKENIZER, 512, true).expect("tokenizer should load");
+    let tokenizer = Tokenizer::new(TOKENIZER, 512, true, PaddingMode::BatchLongest, None)
+        .expect("tokenizer should load");
     let texts = vec![
         "Hello, world!".to_string(),
         "A second, longer sentence to test padding behavior.".to_string(),
@@ -33,7 +35,8 @@ fn test_e5_truncation_at_max_length() {
         "/tests/fixtures/e5/tokenizer.json"
     );
 
-    let tokenizer = Tokenizer::new(TOKENIZER, 16, false).expect("tokenizer should load");
+    let tokenizer = Tokenizer::new(TOKENIZER, 16, false, PaddingMode::BatchLongest, None)
+        .expect("tokenizer should load");
     let long_text = "word ".repeat(200);
     let tokenized = tokenizer
         .tokenize(&[long_text])

data/lib/gte/config.rb

@@ -4,12 +4,12 @@ module GTE
   module Config
     Text = Data.define(
       :model_dir, :threads, :optimization_level,
-      :model_name, :normalize, :output_tensor, :max_length, :execution_providers
+      :model_name, :normalize, :output_tensor, :max_length, :padding, :execution_providers
     )
 
     Reranker = Data.define(
       :model_dir, :threads, :optimization_level,
-      :model_name, :sigmoid, :output_tensor, :max_length, :execution_providers
+      :model_name, :sigmoid, :output_tensor, :max_length, :padding, :execution_providers
     )
   end
 end

data/lib/gte/embedder.rb

@@ -2,6 +2,9 @@
 
 module GTE
   class Embedder
+    DEFAULT_THREADS = 1
+    DEFAULT_OPTIMIZATION_LEVEL = 3
+
     class << self
       def config(model_dir)
         cfg = default_config(model_dir)
@@ -18,21 +21,21 @@ module GTE
           config.normalize,
           config.output_tensor.to_s,
           config.max_length || 0,
+          config.padding.to_s,
           config.execution_providers.to_s
         )
       end
 
-      private
-
       def default_config(model_dir)
         Config::Text.new(
           model_dir: File.expand_path(model_dir),
-          threads: 3,
-          optimization_level: 3,
+          threads: DEFAULT_THREADS,
+          optimization_level: DEFAULT_OPTIMIZATION_LEVEL,
           model_name: nil,
           normalize: true,
           output_tensor: nil,
           max_length: nil,
+          padding: nil,
           execution_providers: nil
         )
       end

data/lib/gte/reranker.rb

@@ -19,12 +19,13 @@ module GTE
       def default_config(model_dir)
         Config::Reranker.new(
           model_dir: File.expand_path(model_dir),
-          threads: 3,
+          threads: 1,
           optimization_level: 3,
           model_name: nil,
           sigmoid: false,
           output_tensor: nil,
           max_length: nil,
+          padding: nil,
           execution_providers: nil
         )
       end
@@ -38,6 +39,7 @@ module GTE
           cfg.sigmoid,
           cfg.output_tensor.to_s,
           cfg.max_length || 0,
+          cfg.padding.to_s,
           cfg.execution_providers.to_s
         )
       end

data/lib/gte.rb

@@ -19,16 +19,7 @@ module GTE
 
   class << self
     def config(model_dir)
-      cfg = Config::Text.new(
-        model_dir: File.expand_path(model_dir),
-        threads: 3,
-        optimization_level: 3,
-        model_name: nil,
-        normalize: true,
-        output_tensor: nil,
-        max_length: nil,
-        execution_providers: nil
-      )
+      cfg = Embedder.default_config(model_dir)
 
       cfg = yield(cfg) if block_given?