gte 0.0.7 → 0.0.8

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 29659e3ab6072d858b1710a779c3d2e5981f7749782182d141ccd5e9790a1fbb
-  data.tar.gz: c42d51cfa1a2ba6a2e83249e8a725c978b11c7ef80c6d69f09a64e884be42031
+  metadata.gz: 2c754b4675ee105e9a280cd9deafa00a81b9e02ee629131f3e908400006b6ae4
+  data.tar.gz: 40a0d3e04c3d2943ae50910164d644ecb763eac99a02044dc962cc141a0e13c5
 SHA512:
-  metadata.gz: ff2c2b1450a6e82c07aacd2ec98437f03678d56eef9c5516f904021a54f59b2ba5c42b8f6af22b5c4b2dacea98615b99bc54d2c7cdc4e8fbccc1abc195fe9975
-  data.tar.gz: 04ca056458d40e2ba7fabcdbcab415a087d54802fb3bd86748dc901c2cf0ecb44072fd1820a73e3dcaca097f165df3e70bab747b38340cd738876af5f0ea7645
+  metadata.gz: 16614e01e7a33a53339ba9fe7cf32fe7606041518a24177258d7a6e5550516e8cff741d0f0df02b7e5863fc763c02ae81b943dc4b18295701a4cafdec6627cb0
+  data.tar.gz: 348e1fd1d9f4c44214b5101ba339109b5ececfbef18b48b7c11324a64481f476d8da831cc5148d17a85c41b525ee753c296d4421a4fb2adda269a3f5fe38cda6

data/README.md

@@ -33,8 +33,8 @@ raw_model = GTE.config(ENV.fetch("GTE_MODEL_DIR")) do |config|
   config.with(normalize: false)
 end
 
-full_throttle = GTE.config(ENV.fetch("GTE_MODEL_DIR")) do |config|
-  config.with(threads: 0)
+single_thread = GTE.config(ENV.fetch("GTE_MODEL_DIR")) do |config|
+  config.with(threads: 1)
 end
 
 custom = GTE.config(ENV.fetch("GTE_MODEL_DIR")) do |config|
@@ -50,7 +50,7 @@ end
 Config fields and defaults:
 
 - `model_dir`: absolute path to model directory
-- `threads`: `3` (set `0` for ONNX Runtime full-throttle threadpool)
+- `threads`: `1` (default tuned for p95 latency; use `0` for ONNX Runtime auto-thread mode)
 - `optimization_level`: `3`
 - `model_name`: `nil`
 - `normalize`: `true` (L2 normalization at Ruby-facing API)
@@ -68,7 +68,7 @@ Low-level embedder setup (without model cache):
 
 ```ruby
 embedder = GTE::Embedder.config(ENV.fetch("GTE_MODEL_DIR")) do |config|
-  config.with(threads: 0, execution_providers: "cpu")
+  config.with(threads: 1, execution_providers: "cpu")
 end
 ```
 
@@ -78,7 +78,7 @@ Use `GTE::Reranker.config(model_dir)` for cross-encoder reranking.
 
 ```ruby
 reranker = GTE::Reranker.config(ENV.fetch("GTE_RERANK_DIR")) do |config|
-  config.with(sigmoid: true, threads: 0)
+  config.with(sigmoid: true, threads: 1)
 end
 
 query = "how to train a neural network?"
@@ -102,7 +102,7 @@ ranked = reranker.rerank(query: query, candidates: candidates)
 Reranker config fields and defaults:
 
 - `model_dir`: absolute path to model directory
-- `threads`: `3`
+- `threads`: `1`
 - `optimization_level`: `3`
 - `model_name`: `nil`
 - `sigmoid`: `false` (set `true` if you want bounded [0,1] style scores)
@@ -174,7 +174,7 @@ make ci
 
 ## Benchmark
 
-The repo includes two benchmark paths:
+The repo includes a shared multi-runtime benchmark harness:
 
 ```bash
 make bench
@@ -183,6 +183,11 @@ nix develop -c bundle exec rake bench:matrix_sweep
 nix develop -c bundle exec ruby bench/memory_probe.rb --compare-pure
 ```
 
+- `make bench`: Puma-like single-request comparison at concurrency `16`
+- `rake bench:pure_compare`: batch amortization comparison
+- `rake bench:matrix_sweep`: GTE provider/thread sweep using the shared result schema
+- Optional Python comparisons use `bench/python_onnxruntime.py` and are skipped automatically if local dependencies are unavailable.
+
 To run benchmark + append a `RUNS.md` entry + enforce goal checks:
 
 ```bash
@@ -191,5 +196,5 @@ make bench-record
 
 `bench/runs_ledger.rb check` is goal-focused by default:
 
-- Enforces goal metric (`response_time_p95` ratio threshold).
+- Enforces the goal metric (`response_time_p95`) across every enabled competitor.
 - Does not require current-version coverage in `RUNS.md` unless explicitly enabled.

data/Rakefile

@@ -10,17 +10,52 @@ rescue LoadError
 end
 
 spec = Gem::Specification.load('gte.gemspec')
+cross_target = ENV.fetch('RUBY_TARGET', nil)
 
-Rake::ExtensionTask.new('gte', spec) do |ext|
+if cross_target == 'arm64-darwin'
+  # rb-sys-dock's darwin image can expose an unusable default LIBRARY_PATH.
+  # Force the compiler-rt darwin runtime directory so -lclang_rt.osx resolves.
+  ENV['LIBRARY_PATH'] = '/usr/lib/llvm-10/lib/clang/10.0.0/lib/darwin'
+end
+
+extension_task = Rake::ExtensionTask.new('gte', spec) do |ext|
   ext.lib_dir = 'lib/gte'
   ext.cross_compile = true
-  ext.cross_platform = %w[x86_64-linux arm64-darwin]
+  # rb-sys-dock invokes `rake native:$RUBY_TARGET gem` without the `cross` task,
+  # so scope platforms during dock builds to avoid host-Ruby fallback copy tasks.
+  cross_platforms = if cross_target && !cross_target.empty?
+                      [cross_target]
+                    else
+                      %w[x86_64-linux aarch64-linux arm64-darwin]
+                    end
+  ext.cross_platform = cross_platforms
+end
+
+if cross_target && !cross_target.empty? && ENV['RUBY_CC_VERSION']
+  ruby_version = ENV['RUBY_CC_VERSION'].split(':').first
+  lib_binary_path = File.join(extension_task.lib_dir, File.basename(extension_task.binary(cross_target)))
+  copy_task = "copy:gte:#{cross_target}:#{ruby_version}"
+
+  if Rake::Task.task_defined?(lib_binary_path) && Rake::Task.task_defined?(copy_task)
+    Rake::Task[lib_binary_path].prerequisites.clear
+    Rake::Task[lib_binary_path].enhance([copy_task])
+  end
 end
 
 task default: %i[compile spec]
 
+def bundler_env
+  root = File.expand_path(__dir__)
+  {
+    'BUNDLE_DISABLE_SHARED_GEMS' => '1',
+    'GEM_HOME' => File.join(root, '.bundle-gems'),
+    'GEM_PATH' => File.join(root, '.bundle-gems'),
+    'BUNDLE_PATH' => File.join(root, 'vendor/bundle')
+  }
+end
+
 def run_in_nix(*command)
-  sh('nix', 'develop', '-c', *command)
+  sh(bundler_env, 'nix', 'develop', '-c', *command)
 end
 
 namespace :bench do

data/VERSION

@@ -1 +1 @@
-0.0.7
+0.0.8

data/ext/gte/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "gte"
-version = "0.0.7"
+version = "0.0.8"
 edition = "2021"
 authors = ["elcuervo <elcuervo@elcuervo.net>"]
 license = "MIT"
@@ -21,10 +21,10 @@ ruby-ffi = ["dep:magnus", "dep:rb-sys"]
 [dependencies]
 rb-sys = { version = "0.9", features = ["stable-api-compiled-fallback"], optional = true }
 magnus = { version = "0.8", optional = true }
-ort = { version = "=2.0.0-rc.9", features = ["ndarray"] }
-ort-sys = "=2.0.0-rc.9"
+ort = { version = "=2.0.0-rc.12", features = ["ndarray", "xnnpack"] }
+ort-sys = "=2.0.0-rc.12"
 tokenizers = "0.21.0"
-ndarray = "0.16.0"
+ndarray = "0.17"
 half = "2"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"

data/ext/gte/src/embedder.rs

@@ -5,15 +5,14 @@ use crate::model_profile::{
     resolve_named_model, resolve_tokenizer_path, select_output_tensor, validate_supported_text_inputs,
 };
 use crate::postprocess::normalize_l2 as normalize_l2_rows;
-use crate::session::{build_session, run_session};
+use crate::session::{build_session, run_session, SessionPool};
 use crate::tokenizer::{parse_padding_mode_override, Tokenized, Tokenizer};
 use ndarray::Array2;
-use ort::session::Session;
-use std::path::Path;
+use std::path::{Path, PathBuf};
 
 pub struct Embedder {
     tokenizer: Tokenizer,
-    session: Session,
+    pool: SessionPool,
     config: ModelConfig,
 }
 
@@ -30,12 +29,10 @@ impl Embedder {
             config.padding_mode,
             None,
         )?;
-        let session = build_session(model_path, &config)?;
-        Ok(Self {
-            tokenizer,
-            session,
-            config,
-        })
+        let model_path = model_path.as_ref().to_path_buf();
+        let session = build_session(&model_path, &config)?;
+        let pool = SessionPool::new(session, model_path, config.clone());
+        Ok(Self { tokenizer, pool, config })
     }
 
     pub fn from_dir<P: AsRef<Path>>(
@@ -53,7 +50,7 @@ impl Embedder {
 
         let dir = dir.as_ref();
         let tokenizer_path = resolve_tokenizer_path(dir)?;
-        let model_path = match overrides.model_name.filter(|s| !s.is_empty()) {
+        let model_path: PathBuf = match overrides.model_name.filter(|s| !s.is_empty()) {
             Some(name) => resolve_named_model(dir, name)?,
             None => resolve_default_text_model(dir)?,
         };
@@ -117,11 +114,8 @@ impl Embedder {
             tokenizer_profile.fixed_padding_length,
         )?;
 
-        Ok(Self {
-            tokenizer,
-            session,
-            config,
-        })
+        let pool = SessionPool::new(session, model_path, session_config);
+        Ok(Self { tokenizer, pool, config })
     }
 
     pub fn embed(&self, texts: Vec<String>) -> Result<Array2<f32>> {
@@ -134,7 +128,8 @@ impl Embedder {
     }
 
     pub fn run(&self, tokenized: &Tokenized) -> crate::error::Result<Array2<f32>> {
-        run_session(&self.session, tokenized, &self.config)
+        let mut session = self.pool.acquire()?;
+        run_session(&mut session, tokenized, &self.config)
     }
 }
 

data/ext/gte/src/model_profile.rs

@@ -133,11 +133,9 @@ fn parse_fixed_padding_length_from_tokenizer_json(tokenizer_json: &Value) -> Opt
 }
 
 pub fn validate_supported_text_inputs(session: &Session, api_label: &str) -> Result<()> {
-    let unsupported: Vec<String> = session
-        .inputs
-        .iter()
-        .filter(|i| !SUPPORTED_INPUTS.contains(&i.name.as_str()))
-        .map(|i| i.name.clone())
+    let unsupported: Vec<String> = session.inputs().iter()
+        .filter(|i| !SUPPORTED_INPUTS.contains(&i.name()))
+        .map(|i| i.name().to_owned())
         .collect();
 
     if unsupported.is_empty() {
@@ -160,7 +158,7 @@ pub fn validate_supported_text_inputs(session: &Session, api_label: &str) -> Res
 }
 
 pub fn has_input(session: &Session, name: &str) -> bool {
-    session.inputs.iter().any(|input| input.name == name)
+    session.inputs().iter().any(|input| input.name() == name)
 }
 
 fn output_name_matches(name: &str, preferred: &str) -> bool {
@@ -175,16 +173,16 @@ pub fn select_output_tensor(
 ) -> Result<String> {
     if let Some(requested_name) = requested.map(str::trim).filter(|name| !name.is_empty()) {
         if let Some(output) = session
-            .outputs
+            .outputs()
             .iter()
-            .find(|o| output_name_matches(o.name.as_str(), requested_name))
+            .find(|o| output_name_matches(o.name(), requested_name))
         {
-            return Ok(output.name.clone());
+            return Ok(output.name().to_owned());
         }
         let available = session
-            .outputs
+            .outputs()
             .iter()
-            .map(|o| o.name.as_str())
+            .map(|o| o.name())
             .collect::<Vec<_>>()
             .join(", ");
         return Err(GteError::Inference(format!(
@@ -195,18 +193,18 @@ pub fn select_output_tensor(
 
     for preferred in preferred_outputs {
         if let Some(output) = session
-            .outputs
+            .outputs()
             .iter()
-            .find(|o| output_name_matches(o.name.as_str(), preferred))
+            .find(|o| output_name_matches(o.name(), preferred))
         {
-            return Ok(output.name.clone());
+            return Ok(output.name().to_owned());
         }
     }
 
     session
-        .outputs
+        .outputs()
         .first()
-        .map(|o| o.name.clone())
+        .map(|o| o.name().to_owned())
         .ok_or_else(|| GteError::Inference("model has no outputs".into()))
 }
 
@@ -216,9 +214,9 @@ fn output_basename(name: &str) -> &str {
 
 pub fn infer_extraction_mode(session: &Session, output_tensor: &str) -> Result<ExtractorMode> {
     let output = session
-        .outputs
+        .outputs()
         .iter()
-        .find(|o| o.name == output_tensor)
+        .find(|o| o.name() == output_tensor)
         .ok_or_else(|| {
             GteError::Inference(format!(
                 "output tensor '{}' not found in model outputs",
@@ -226,8 +224,8 @@ pub fn infer_extraction_mode(session: &Session, output_tensor: &str) -> Result<E
             ))
         })?;
 
-    let ndims = match &output.output_type {
-        ort::value::ValueType::Tensor { dimensions, .. } => dimensions.len(),
+    let ndims = match output.dtype() {
+        ort::value::ValueType::Tensor { shape, .. } => shape.len(),
         other => {
             return Err(GteError::Inference(format!(
                 "output is not a tensor: {:?}",

data/ext/gte/src/pipeline.rs

@@ -1,8 +1,8 @@
 use crate::error::{GteError, Result};
 use crate::tokenizer::Tokenized;
-use ndarray::ArrayView2;
+use ndarray::{ArrayView2, ArrayViewD};
 use ort::session::SessionInputValue;
-use ort::value::Value;
+use ort::value::TensorRef;
 
 pub struct InputTensors<'a> {
     pub inputs: Vec<(&'static str, SessionInputValue<'a>)>,
@@ -23,13 +23,13 @@ impl<'a> InputTensors<'a> {
         let mut inputs = Vec::with_capacity(2 + usize::from(tokenized.type_ids.is_some()));
         inputs.push((
             "input_ids",
-            SessionInputValue::from(Value::from_array(input_ids_view)?),
+            SessionInputValue::from(TensorRef::from_array_view(input_ids_view)?),
         ));
 
         if with_attention_mask {
             inputs.push((
                 "attention_mask",
-                SessionInputValue::from(Value::from_array(attention_mask)?),
+                SessionInputValue::from(TensorRef::from_array_view(attention_mask)?),
             ));
         }
 
@@ -38,7 +38,7 @@ impl<'a> InputTensors<'a> {
                 ArrayView2::from_shape((tokenized.rows, tokenized.cols), type_ids)?;
             inputs.push((
                 "token_type_ids",
-                SessionInputValue::from(Value::from_array(type_ids_view)?),
+                SessionInputValue::from(TensorRef::from_array_view(type_ids_view)?),
             ));
         }
 
@@ -50,11 +50,14 @@ impl<'a> InputTensors<'a> {
 }
 
 pub fn extract_output_tensor<'a>(
-    outputs: &'a ort::session::SessionOutputs<'a, 'a>,
+    outputs: &'a ort::session::SessionOutputs<'_>,
     output_name: &str,
-) -> Result<ndarray::CowArray<'a, f32, ndarray::IxDyn>> {
+) -> Result<ArrayViewD<'a, f32>> {
     let tensor_value = outputs.get(output_name).ok_or_else(|| {
-        GteError::Inference(format!("output tensor '{}' not found in model outputs", output_name))
+        GteError::Inference(format!(
+            "output tensor '{}' not found in model outputs",
+            output_name
+        ))
     })?;
-    Ok(tensor_value.try_extract_tensor::<f32>()?.into())
+    Ok(tensor_value.try_extract_array::<f32>()?)
 }

data/ext/gte/src/reranker.rs

@@ -6,10 +6,9 @@ use crate::model_profile::{
 };
 use crate::pipeline::{extract_output_tensor, InputTensors};
 use crate::postprocess::sigmoid_scores;
-use crate::session::build_session;
+use crate::session::{build_session, SessionPool};
 use crate::tokenizer::{parse_padding_mode_override, Tokenizer};
-use ort::session::Session;
-use std::path::Path;
+use std::path::{Path, PathBuf};
 
 #[derive(Debug, Clone)]
 struct RerankerConfig {
@@ -22,7 +21,7 @@ struct RerankerConfig {
 
 pub struct Reranker {
     tokenizer: Tokenizer,
-    session: Session,
+    pool: SessionPool,
     config: RerankerConfig,
 }
 
@@ -35,7 +34,7 @@ impl Reranker {
     ) -> Result<Self> {
         let dir = dir.as_ref();
         let tokenizer_path = resolve_tokenizer_path(dir)?;
-        let model_path = match overrides.model_name.filter(|s| !s.is_empty()) {
+        let model_path: PathBuf = match overrides.model_name.filter(|s| !s.is_empty()) {
             Some(name) => resolve_named_model(dir, name)?,
             None => resolve_default_text_model(dir)?,
         };
@@ -88,11 +87,8 @@ impl Reranker {
             tokenizer_profile.fixed_padding_length,
         )?;
 
-        Ok(Self {
-            tokenizer,
-            session,
-            config,
-        })
+        let pool = SessionPool::new(session, model_path, probe_config);
+        Ok(Self { tokenizer, pool, config })
     }
 
     pub fn score_pairs(&self, pairs: &[(String, String)], apply_sigmoid: bool) -> Result<Vec<f32>> {
@@ -111,7 +107,8 @@ impl Reranker {
         apply_sigmoid: bool,
     ) -> Result<Vec<f32>> {
         let input_tensors = InputTensors::from_tokenized(tokenized, self.config.with_attention_mask)?;
-        let outputs = self.session.run(input_tensors.inputs)?;
+        let mut session = self.pool.acquire()?;
+        let outputs = session.run(input_tensors.inputs).map_err(|e| GteError::Ort(e.to_string()))?;
         let array = extract_output_tensor(&outputs, self.config.output_tensor.as_str())?;
 
         let mut scores = match array.ndim() {

data/ext/gte/src/ruby_embedder.rs

@@ -28,21 +28,24 @@ pub struct RbTensor {
     data: Vec<f32>,
 }
 
+// ---------------------------------------------------------------------------
+// GVL-release helpers
+// ---------------------------------------------------------------------------
+
 struct InferArgs {
     embedder: *const Embedder,
     texts: *const Vec<String>,
     normalize: bool,
-    result: Option<Result<ndarray::Array2<f32>, GteError>>,
+    result: Option<crate::error::Result<ndarray::Array2<f32>>>,
 }
 
 unsafe impl Send for InferArgs {}
 
 struct ScoreArgs {
     reranker: *const Reranker,
-    query: *const String,
-    candidates: *const Vec<String>,
+    pairs: *const Vec<(String, String)>,
     apply_sigmoid: bool,
-    result: Option<Result<Vec<f32>, GteError>>,
+    result: Option<crate::error::Result<Vec<f32>>>,
 }
 
 unsafe impl Send for ScoreArgs {}
@@ -57,6 +60,38 @@ fn panic_payload_to_string(payload: Box<dyn std::any::Any + Send>) -> String {
     }
 }
 
+unsafe extern "C" fn run_without_gvl(ptr: *mut c_void) -> *mut c_void {
+    let args = &mut *(ptr as *mut InferArgs);
+    let run_result = catch_unwind(AssertUnwindSafe(|| {
+        let tokenized = (*args.embedder).tokenize(&*args.texts)?;
+        let embeddings = (*args.embedder).run(&tokenized)?;
+        if args.normalize { Ok(normalize_l2(embeddings)) } else { Ok(embeddings) }
+    }));
+    args.result = Some(match run_result {
+        Ok(result) => result,
+        Err(payload) => Err(GteError::Inference(format!(
+            "panic during inference: {}",
+            panic_payload_to_string(payload),
+        ))),
+    });
+    std::ptr::null_mut()
+}
+
+unsafe extern "C" fn run_score_without_gvl(ptr: *mut c_void) -> *mut c_void {
+    let args = &mut *(ptr as *mut ScoreArgs);
+    let run_result = catch_unwind(AssertUnwindSafe(|| {
+        (*args.reranker).score_pairs(&*args.pairs, args.apply_sigmoid)
+    }));
+    args.result = Some(match run_result {
+        Ok(result) => result,
+        Err(payload) => Err(GteError::Inference(format!(
+            "panic during reranking: {}",
+            panic_payload_to_string(payload),
+        ))),
+    });
+    std::ptr::null_mut()
+}
+
 fn infer_without_gvl(
     embedder: &Arc<Embedder>,
     normalize: bool,
@@ -87,15 +122,13 @@ fn infer_without_gvl(
 
 fn score_without_gvl(
     reranker: &Arc<Reranker>,
-    query: String,
-    candidates: Vec<String>,
+    pairs: Vec<(String, String)>,
     apply_sigmoid: bool,
 ) -> Result<Vec<f32>, Error> {
     let scores = unsafe {
         let mut args = ScoreArgs {
             reranker: Arc::as_ptr(reranker),
-            query: &query as *const String,
-            candidates: &candidates as *const Vec<String>,
+            pairs: &pairs as *const Vec<(String, String)>,
             apply_sigmoid,
             result: None,
         };
@@ -115,41 +148,7 @@ fn score_without_gvl(
     Ok(scores)
 }
 
-unsafe extern "C" fn run_without_gvl(ptr: *mut c_void) -> *mut c_void {
-    let args = &mut *(ptr as *mut InferArgs);
-    let run_result = catch_unwind(AssertUnwindSafe(|| {
-        let tokenized = (*args.embedder).tokenize(&*args.texts)?;
-        let embeddings = (*args.embedder).run(&tokenized)?;
-        if args.normalize {
-            Ok(normalize_l2(embeddings))
-        } else {
-            Ok(embeddings)
-        }
-    }));
-    args.result = Some(match run_result {
-        Ok(result) => result,
-        Err(payload) => Err(GteError::Inference(format!(
-            "panic during inference: {}",
-            panic_payload_to_string(payload),
-        ))),
-    });
-    std::ptr::null_mut()
-}
-
-unsafe extern "C" fn run_score_without_gvl(ptr: *mut c_void) -> *mut c_void {
-    let args = &mut *(ptr as *mut ScoreArgs);
-    let run_result = catch_unwind(AssertUnwindSafe(|| {
-        (*args.reranker).score(&*args.query, &*args.candidates, args.apply_sigmoid)
-    }));
-    args.result = Some(match run_result {
-        Ok(result) => result,
-        Err(payload) => Err(GteError::Inference(format!(
-            "panic during reranking: {}",
-            panic_payload_to_string(payload),
-        ))),
-    });
-    std::ptr::null_mut()
-}
+// ---------------------------------------------------------------------------
 
 fn tensor_from_array(embeddings: ndarray::Array2<f32>) -> Result<RbTensor, Error> {
     let rows = embeddings.nrows();
@@ -177,31 +176,11 @@ impl RbEmbedder {
         padding: String,
         execution_providers: String,
     ) -> 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())
-        };
-        let max_length_override = if max_length == 0 {
-            None
-        } else {
-            Some(max_length)
-        };
-        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 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()) };
+        let max_length_override = if max_length == 0 { None } else { Some(max_length) };
+        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 overrides = ModelLoadOverrides {
             model_name: name,
             output_tensor: output_override,
@@ -209,17 +188,9 @@ impl RbEmbedder {
             padding: padding_override,
             execution_providers: execution_providers_override,
         };
-        let embedder = Embedder::from_dir(
-            &dir_path,
-            num_threads,
-            optimization_level,
-            overrides,
-        )
-        .map_err(magnus::Error::from)?;
-        Ok(RbEmbedder {
-            inner: Arc::new(embedder),
-            normalize,
-        })
+        let embedder = Embedder::from_dir(&dir_path, num_threads, optimization_level, overrides)
+            .map_err(magnus::Error::from)?;
+        Ok(RbEmbedder { inner: Arc::new(embedder), normalize })
     }
 
     pub fn rb_embed(_ruby: &Ruby, rb_self: &Self, texts: RArray) -> Result<RbTensor, Error> {
@@ -247,31 +218,11 @@ impl RbReranker {
         padding: String,
         execution_providers: String,
     ) -> 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())
-        };
-        let max_length_override = if max_length == 0 {
-            None
-        } else {
-            Some(max_length)
-        };
-        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 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()) };
+        let max_length_override = if max_length == 0 { None } else { Some(max_length) };
+        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 overrides = ModelLoadOverrides {
             model_name: name,
             output_tensor: output_override,
@@ -279,17 +230,9 @@ impl RbReranker {
             padding: padding_override,
             execution_providers: execution_providers_override,
         };
-        let reranker = Reranker::from_dir(
-            &dir_path,
-            num_threads,
-            optimization_level,
-            overrides,
-        )
-        .map_err(magnus::Error::from)?;
-        Ok(RbReranker {
-            inner: Arc::new(reranker),
-            sigmoid,
-        })
+        let reranker = Reranker::from_dir(&dir_path, num_threads, optimization_level, overrides)
+            .map_err(magnus::Error::from)?;
+        Ok(RbReranker { inner: Arc::new(reranker), sigmoid })
     }
 
     pub fn rb_score(
@@ -299,8 +242,8 @@ impl RbReranker {
         candidates: RArray,
     ) -> Result<RArray, Error> {
         let candidates: Vec<String> = candidates.to_vec()?;
-        let scores = score_without_gvl(&rb_self.inner, query, candidates, rb_self.sigmoid)?;
-
+        let pairs: Vec<(String, String)> = candidates.into_iter().map(|c| (query.clone(), c)).collect();
+        let scores = score_without_gvl(&rb_self.inner, pairs, rb_self.sigmoid)?;
         let out = ruby.ary_new_capa(scores.len());
         for score in scores {
             out.push(score)?;
@@ -336,7 +279,6 @@ impl RbTensor {
                 index, rb_self.rows
             ))));
         }
-
         let start = index * rb_self.cols;
         let end = start + rb_self.cols;
         let out = ruby.ary_new_capa(rb_self.cols);
@@ -361,7 +303,6 @@ impl RbTensor {
                 index, rb_self.rows
             ))));
         }
-
         let start = index * rb_self.cols;
         let end = start + rb_self.cols;
         let bytes = unsafe {

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/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)
@@ -23,13 +26,11 @@ module GTE
         )
       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,

data/lib/gte/reranker.rb

@@ -19,7 +19,7 @@ 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,

data/lib/gte.rb

@@ -19,17 +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,
-        padding: nil,
-        execution_providers: nil
-      )
+      cfg = Embedder.default_config(model_dir)
 
       cfg = yield(cfg) if block_given?
 

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: gte
 version: !ruby/object:Gem::Version
-  version: 0.0.7
+  version: 0.0.8
 platform: ruby
 authors:
 - elcuervo
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2026-04-16 00:00:00.000000000 Z
+date: 2026-04-28 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rake
@@ -42,16 +42,16 @@ dependencies:
   name: rb_sys
   requirement: !ruby/object:Gem::Requirement
     requirements:
-    - - ">="
+    - - '='
       - !ruby/object:Gem::Version
-        version: '0'
+        version: 0.9.126
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
-    - - ">="
+    - - '='
       - !ruby/object:Gem::Version
-        version: '0'
+        version: 0.9.126
 - !ruby/object:Gem::Dependency
   name: rspec
   requirement: !ruby/object:Gem::Requirement