gte 0.0.11 → 0.0.13

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 8123419de1a0fe86fae1de8808354318ca5b8d575514c7452ac45f71b657252e
-  data.tar.gz: da92e0ba5cab358dde743f4b444c5500f88757268084138571bd578bc505f1f3
+  metadata.gz: 278028df09fbcdd14fd583f0af5e1a8c9553adb28fe7aa0bc67b67666dbbdccd
+  data.tar.gz: ce994e3f505200ed4654ca8f87f585ff88919201fe82dd79007622f07a3d1ea0
 SHA512:
-  metadata.gz: 1fd3dc5a8a7e005d797f352c3fefea348a6badc529eca0169b9d349ffda0aa707f0327897112e019340091b2613e6948210966676ddc6717e72146f506c93fae
-  data.tar.gz: ea12d947a02133c69990f24c104a2003df465149fa8beb3ec857ced623b3ba39310f160d956c3da40636a58e92e50720e37c3845905db74dc772ca8b4a6d0c85
+  metadata.gz: 742f1830ff2b83f89726be527c4323a81649b04f341b7adc0544a9000373f6a097c0b4b4ba211ead5912ba45d876565fbaab6d723ef8f06c488ab7827323f827
+  data.tar.gz: 75e91b3d4c3980b166268c6468b96bebe4b74db999e0cee433a295e57d89bec95c7614b004c61e8b3ed88cff30f02f3b6aff74de710d3dd3bb34552f36fb3422

data/README.md

@@ -33,10 +33,6 @@ raw_model = GTE.config(ENV.fetch("GTE_MODEL_DIR")) do |config|
   config.with(normalize: false)
 end
 
-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|
   config.with(
     output_tensor: "last_hidden_state",
@@ -50,7 +46,6 @@ end
 Config fields and defaults:
 
 - `model_dir`: absolute path to model directory
-- `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)
@@ -64,11 +59,48 @@ Notes:
 - Return a `Config::Text` from the block (for example, `config.with(...)`).
 - Model instances are cached by full config key; different config values create different cached instances.
 
+Common model presets:
+
+```ruby
+e5 = GTE.config(ENV.fetch("GTE_MODEL_DIR")) do |config|
+  config.with(
+    model_name: "model.onnx",
+    output_tensor: "last_hidden_state",
+    max_length: 512,
+    execution_providers: "cpu"
+  )
+end
+
+siglip2 = GTE.config(ENV.fetch("GTE_SIGLIP2_DIR")) do |config|
+  config.with(
+    model_name: "text_model_int8.onnx",
+    output_tensor: "pooler_output",
+    max_length: 64,
+    execution_providers: "cpu"
+  )
+end
+
+clip = GTE.config(ENV.fetch("GTE_CLIP_DIR")) do |config|
+  config.with(
+    output_tensor: "sentence_embedding",
+    max_length: 512,
+    execution_providers: "cpu"
+  )
+end
+```
+
+Picking a specific layer:
+
+- Use `output_tensor:` to request a named model output.
+- `last_hidden_state` gives token-level hidden states and is mean-pooled by `gte` when the tensor is rank 3.
+- `pooler_output`, `sentence_embedding`, and similar 2D tensors are returned directly and then L2-normalized by default.
+- If the requested tensor is not present in the model, `gte` raises an error instead of silently falling back.
+
 Low-level embedder setup (without model cache):
 
 ```ruby
 embedder = GTE::Embedder.config(ENV.fetch("GTE_MODEL_DIR")) do |config|
-  config.with(threads: 1, execution_providers: "cpu")
+  config.with(execution_providers: "cpu")
 end
 ```
 
@@ -78,7 +110,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: 1)
+  config.with(sigmoid: true)
 end
 
 query = "how to train a neural network?"
@@ -102,7 +134,6 @@ ranked = reranker.rerank(query: query, candidates: candidates)
 Reranker config fields and defaults:
 
 - `model_dir`: absolute path to model directory
-- `threads`: `1`
 - `optimization_level`: `3`
 - `model_name`: `nil`
 - `sigmoid`: `false` (set `true` if you want bounded [0,1] style scores)
@@ -111,6 +142,11 @@ Reranker config fields and defaults:
 - `padding`: `nil` (auto; accepts `auto`, `batch_longest`, `fixed`)
 - `execution_providers`: `nil`
 
+Session pool sizing:
+
+- `GTE_SESSION_POOL_CAP`: optional positive integer cap for internal ONNX session pool size.
+- Unset by default; runtime uses available CPU parallelism.
+
 ## Runtime + Result Examples
 
 Process-local reuse (recommended for Puma/web servers):
@@ -185,7 +221,7 @@ 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
+- `rake bench:matrix_sweep`: GTE provider 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:

data/Rakefile

@@ -74,7 +74,7 @@ namespace :bench do
     )
   end
 
-  desc 'Sweep execution-provider and thread settings for Puma-like benchmark'
+  desc 'Sweep execution-provider settings for Puma-like benchmark'
   task :matrix_sweep do
     run_in_nix(
       'bundle', 'exec', 'ruby', 'bench/puma_matrix_sweep.rb',

data/VERSION

@@ -1 +1 @@
-0.0.11
+0.0.13

data/ext/gte/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "gte"
-version = "0.0.11"
+version = "0.0.13"
 edition = "2021"
 authors = ["elcuervo <elcuervo@elcuervo.net>"]
 license = "MIT"

data/ext/gte/benches/hot_path.rs

@@ -90,7 +90,7 @@ fn bench_padding_impact(c: &mut Criterion) {
 //   GTE_BENCH_E5_DIR       — sentence-transformers / E5-style text model dir
 //   GTE_BENCH_SIGLIP2_DIR  — siglip2 text encoder dir
 //   GTE_BENCH_CLIP_DIR     — clip text encoder dir
-// Sweeps threads ∈ {0 (auto/all-cores), 1, 2} to validate DEFAULT_THREADS=0.
+// Sweeps execution providers for quick local comparison.
 fn bench_embedding_e2e(c: &mut Criterion) {
     let cases = [
         ("e5", "GTE_BENCH_E5_DIR", "query: cat", "query: ".to_string() + &"the quick brown fox jumps over the lazy dog ".repeat(20)),
@@ -106,17 +106,21 @@ fn bench_embedding_e2e(c: &mut Criterion) {
             continue;
         };
 
-        for &threads in &[0usize, 1, 2] {
-            let embedder = match Embedder::from_dir(&dir, threads, 3, ModelLoadOverrides::default()) {
+        for provider in ["cpu", "xnnpack"] {
+            let overrides = ModelLoadOverrides {
+                execution_providers: Some(provider),
+                ..ModelLoadOverrides::default()
+            };
+            let embedder = match Embedder::from_dir(&dir, 3, overrides) {
                 Ok(e) => e,
                 Err(err) => {
-                    eprintln!("skip {model_label} threads={threads}: {err}");
+                    eprintln!("skip {model_label} provider={provider}: {err}");
                     continue;
                 }
             };
 
             for (input_label, input) in [("short", short_input.to_string()), ("long", long_input.clone())] {
-                let id = BenchmarkId::from_parameter(format!("{model_label}/threads_{threads}/{input_label}"));
+                let id = BenchmarkId::from_parameter(format!("{model_label}/{provider}/{input_label}"));
                 group.bench_with_input(id, &input, |b, text| {
                     b.iter(|| {
                         embedder

data/ext/gte/src/embedder.rs

@@ -37,7 +37,6 @@ impl Embedder {
 
     pub fn from_dir<P: AsRef<Path>>(
         dir: P,
-        num_threads: usize,
         optimization_level: u8,
         overrides: ModelLoadOverrides<'_>,
     ) -> Result<Self> {
@@ -76,7 +75,6 @@ impl Embedder {
             mode: ExtractorMode::Raw,
             with_type_ids: false,
             with_attention_mask: true,
-            num_threads,
             optimization_level,
             execution_providers: overrides.execution_providers.map(str::to_string),
         };
@@ -101,7 +99,6 @@ impl Embedder {
             mode,
             with_type_ids,
             with_attention_mask,
-            num_threads,
             optimization_level,
             execution_providers: overrides.execution_providers.map(str::to_string),
         };
@@ -119,7 +116,11 @@ impl Embedder {
     }
 
     pub fn embed(&self, texts: Vec<String>) -> Result<Array2<f32>> {
-        let tokenized = self.tokenize(&texts)?;
+        self.embed_ref(&texts)
+    }
+
+    pub fn embed_ref(&self, texts: &[String]) -> Result<Array2<f32>> {
+        let tokenized = self.tokenize(texts)?;
         self.run(&tokenized)
     }
 

data/ext/gte/src/model_config.rs

@@ -21,7 +21,6 @@ pub struct ModelConfig {
     pub mode: ExtractorMode,
     pub with_type_ids: bool,
     pub with_attention_mask: bool,
-    pub num_threads: usize,
     pub optimization_level: u8,
     pub execution_providers: Option<String>,
 }

data/ext/gte/src/model_profile.rs

@@ -201,10 +201,14 @@ pub fn select_output_tensor(
         }
     }
 
-    session
-        .outputs()
-        .first()
-        .map(|o| o.name().to_owned())
+    let outputs = session.outputs();
+    let best = outputs
+        .iter()
+        .find(|o| {
+            matches!(o.dtype(), ort::value::ValueType::Tensor { shape, .. } if shape.len() == 2)
+        })
+        .or_else(|| outputs.first());
+    best.map(|o| o.name().to_owned())
         .ok_or_else(|| GteError::Inference("model has no outputs".into()))
 }
 

data/ext/gte/src/reranker.rs

@@ -28,7 +28,6 @@ pub struct Reranker {
 impl Reranker {
     pub fn from_dir<P: AsRef<Path>>(
         dir: P,
-        num_threads: usize,
         optimization_level: u8,
         overrides: ModelLoadOverrides<'_>,
     ) -> Result<Self> {
@@ -60,7 +59,6 @@ impl Reranker {
             mode: crate::model_config::ExtractorMode::Raw,
             with_type_ids: false,
             with_attention_mask: true,
-            num_threads,
             optimization_level,
             execution_providers: overrides.execution_providers.map(str::to_string),
         };

data/ext/gte/src/ruby_embedder.rs

@@ -4,7 +4,6 @@ use crate::embedder::{normalize_l2, Embedder};
 use crate::error::GteError;
 use crate::model_config::ModelLoadOverrides;
 use crate::reranker::Reranker;
-use crate::tokenizer::Tokenized;
 use magnus::{function, method, prelude::*, wrap, Error, RArray, Ruby};
 use std::os::raw::c_void;
 use std::panic::{catch_unwind, AssertUnwindSafe};
@@ -33,10 +32,9 @@ pub struct RbTensor {
 // GVL-release helpers
 // ---------------------------------------------------------------------------
 
-// Tokenized holds only Vec<i64> fields — safe to send across threads.
 struct InferArgs {
     embedder: *const Embedder,
-    tokenized: *const Tokenized,
+    texts: *const Vec<String>,
     normalize: bool,
     result: Option<crate::error::Result<ndarray::Array2<f32>>>,
 }
@@ -45,7 +43,8 @@ unsafe impl Send for InferArgs {}
 
 struct ScoreArgs {
     reranker: *const Reranker,
-    pairs: *const Vec<(String, String)>,
+    query: *const String,
+    candidates: *const Vec<String>,
     apply_sigmoid: bool,
     result: Option<crate::error::Result<Vec<f32>>>,
 }
@@ -62,12 +61,11 @@ 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 {
+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(|| {
-        // Tokenization happens before GVL release (in rb_embed / rb_embed_one).
-        // Only ONNX inference runs here without the GVL.
-        let embeddings = (*args.embedder).run(&*args.tokenized)?;
+        // 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) }
     }));
     args.result = Some(match run_result {
@@ -83,7 +81,7 @@ unsafe extern "C" fn run_without_gvl(ptr: *mut c_void) -> *mut c_void {
 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.reranker).score(&*args.query, &*args.candidates, args.apply_sigmoid)
     }));
     args.result = Some(match run_result {
         Ok(result) => result,
@@ -98,17 +96,17 @@ unsafe extern "C" fn run_score_without_gvl(ptr: *mut c_void) -> *mut c_void {
 fn infer_without_gvl(
     embedder: &Arc<Embedder>,
     normalize: bool,
-    tokenized: &Tokenized,
+    texts: Vec<String>,
 ) -> Result<ndarray::Array2<f32>, Error> {
     let embeddings = unsafe {
         let mut args = InferArgs {
             embedder: Arc::as_ptr(embedder),
-            tokenized: tokenized as *const Tokenized,
+            texts: &texts as *const Vec<String>,
             normalize,
             result: None,
         };
         rb_sys::rb_thread_call_without_gvl(
-            Some(run_without_gvl),
+            Some(run_embed_without_gvl),
             &mut args as *mut InferArgs as *mut c_void,
             None,
             std::ptr::null_mut(),
@@ -125,13 +123,15 @@ fn infer_without_gvl(
 
 fn score_without_gvl(
     reranker: &Arc<Reranker>,
-    pairs: Vec<(String, String)>,
+    query: String,
+    candidates: Vec<String>,
     apply_sigmoid: bool,
 ) -> Result<Vec<f32>, Error> {
     let scores = unsafe {
         let mut args = ScoreArgs {
             reranker: Arc::as_ptr(reranker),
-            pairs: &pairs as *const Vec<(String, String)>,
+            query: &query as *const String,
+            candidates: &candidates as *const Vec<String>,
             apply_sigmoid,
             result: None,
         };
@@ -170,7 +170,6 @@ impl RbEmbedder {
     pub fn rb_new(
         _ruby: &Ruby,
         dir_path: String,
-        num_threads: usize,
         optimization_level: u8,
         model_name: String,
         normalize: bool,
@@ -191,21 +190,19 @@ impl RbEmbedder {
             padding: padding_override,
             execution_providers: execution_providers_override,
         };
-        let embedder = Embedder::from_dir(&dir_path, num_threads, optimization_level, overrides)
+        let embedder = Embedder::from_dir(&dir_path, 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> {
         let texts: Vec<String> = texts.to_vec()?;
-        let tokenized = rb_self.inner.tokenize(&texts).map_err(magnus::Error::from)?;
-        let embeddings = infer_without_gvl(&rb_self.inner, rb_self.normalize, &tokenized)?;
+        let embeddings = infer_without_gvl(&rb_self.inner, rb_self.normalize, texts)?;
         tensor_from_array(embeddings)
     }
 
     pub fn rb_embed_one(_ruby: &Ruby, rb_self: &Self, text: String) -> Result<RbTensor, Error> {
-        let tokenized = rb_self.inner.tokenize(&[text]).map_err(magnus::Error::from)?;
-        let embeddings = infer_without_gvl(&rb_self.inner, rb_self.normalize, &tokenized)?;
+        let embeddings = infer_without_gvl(&rb_self.inner, rb_self.normalize, vec![text])?;
         tensor_from_array(embeddings)
     }
 }
@@ -214,7 +211,6 @@ impl RbReranker {
     pub fn rb_new(
         _ruby: &Ruby,
         dir_path: String,
-        num_threads: usize,
         optimization_level: u8,
         model_name: String,
         sigmoid: bool,
@@ -235,7 +231,7 @@ impl RbReranker {
             padding: padding_override,
             execution_providers: execution_providers_override,
         };
-        let reranker = Reranker::from_dir(&dir_path, num_threads, optimization_level, overrides)
+        let reranker = Reranker::from_dir(&dir_path, optimization_level, overrides)
             .map_err(magnus::Error::from)?;
         Ok(RbReranker { inner: Arc::new(reranker), sigmoid })
     }
@@ -247,8 +243,7 @@ impl RbReranker {
         candidates: RArray,
     ) -> Result<RArray, Error> {
         let candidates: Vec<String> = candidates.to_vec()?;
-        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 scores = score_without_gvl(&rb_self.inner, query, candidates, rb_self.sigmoid)?;
         let out = ruby.ary_new_capa(scores.len());
         for score in scores {
             out.push(score)?;
@@ -345,12 +340,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, 9))?;
+    embedder_class.define_singleton_method("new", function!(RbEmbedder::rb_new, 8))?;
     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, 9))?;
+    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

@@ -7,7 +7,7 @@ use ndarray::{Array2, ArrayView2, ArrayViewD, Ix2};
 use ort::execution_providers::{
     CoreMLExecutionProvider, ExecutionProviderDispatch, XNNPACKExecutionProvider,
 };
-use ort::session::Session;
+use ort::session::{OutputSelector, RunOptions, Session};
 use std::path::{Path, PathBuf};
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Condvar, Mutex};
@@ -27,8 +27,6 @@ pub fn build_session<P: AsRef<Path>>(model_path: P, config: &ModelConfig) -> Res
     let mut builder = Session::builder()
         .map_err(ort_err)?
         .with_optimization_level(opt_level)
-        .map_err(ort_err)?
-        .with_memory_pattern(false)
         .map_err(ort_err)?;
 
     let providers = preferred_execution_providers(config.execution_providers.as_deref());
@@ -38,15 +36,6 @@ pub fn build_session<P: AsRef<Path>>(model_path: P, config: &ModelConfig) -> Res
             .map_err(ort_err)?;
     }
 
-    if config.num_threads > 0 {
-        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)
 }
 
@@ -54,25 +43,17 @@ pub fn build_session<P: AsRef<Path>>(model_path: P, config: &ModelConfig) -> Res
 // Session pool
 // ---------------------------------------------------------------------------
 
-fn pool_capacity(num_threads: usize) -> usize {
-    let available_parallelism = std::thread::available_parallelism()
+fn pool_capacity() -> usize {
+    let available = std::thread::available_parallelism()
         .map(|n| n.get())
         .unwrap_or(1);
-    pool_capacity_with_parallelism(num_threads, available_parallelism)
+    parse_pool_capacity_override().map_or(available, |cap| cap.min(available).max(1))
 }
 
-fn pool_capacity_with_parallelism(num_threads: usize, available_parallelism: usize) -> usize {
-    if available_parallelism == 0 {
-        return 1;
-    }
-
-    // Auto-thread mode: ORT grabs all cores per session. One session avoids
-    // N² intra-op oversubscription when multiple Ruby threads call concurrently.
-    if num_threads == 0 {
-        return 1;
-    }
-
-    available_parallelism.div_ceil(num_threads).max(1)
+fn parse_pool_capacity_override() -> Option<usize> {
+    let raw = std::env::var("GTE_SESSION_POOL_CAP").ok()?;
+    let parsed = raw.trim().parse::<usize>().ok()?;
+    (parsed > 0).then_some(parsed)
 }
 
 pub struct SessionPool {
@@ -86,7 +67,7 @@ pub struct SessionPool {
 
 impl SessionPool {
     pub fn new(initial: Session, model_path: PathBuf, build_config: ModelConfig) -> Self {
-        let capacity = pool_capacity(build_config.num_threads);
+        let capacity = pool_capacity();
         Self {
             sessions: Mutex::new(vec![initial]),
             available: Condvar::new(),
@@ -235,8 +216,11 @@ pub fn run_session(
     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()))?
+        .with_outputs(OutputSelector::no_default().with(config.output_tensor.as_str()));
     let outputs = session
-        .run(input_tensors.inputs)
+        .run_with_options(input_tensors.inputs, &run_opts)
         .map_err(|e| GteError::Ort(e.to_string()))?;
     let array = extract_output_tensor(&outputs, config.output_tensor.as_str())?;
 
@@ -282,7 +266,7 @@ mod tests {
     use ndarray::{array, ArrayView2};
 
     use super::{
-        extract_embeddings, parse_provider_registrations, pool_capacity_with_parallelism,
+        extract_embeddings, parse_pool_capacity_override, parse_provider_registrations,
         resolve_provider_order_with_env,
     };
 
@@ -294,7 +278,6 @@ mod tests {
             mode,
             with_type_ids: false,
             with_attention_mask: true,
-            num_threads: 1,
             optimization_level: 3,
             execution_providers: None,
         }
@@ -343,21 +326,30 @@ mod tests {
     }
 
     #[test]
-    fn pool_capacity_uses_single_session_for_auto_thread_mode() {
-        // Auto-thread = ORT uses all cores per session. Pool=1 avoids N²
-        // intra-op oversubscription under concurrent Ruby threads.
-        assert_eq!(pool_capacity_with_parallelism(0, 1), 1);
-        assert_eq!(pool_capacity_with_parallelism(0, 4), 1);
-        assert_eq!(pool_capacity_with_parallelism(0, 8), 1);
-    }
+    fn parse_pool_capacity_override_uses_positive_integer_only() {
+        unsafe {
+            std::env::remove_var("GTE_SESSION_POOL_CAP");
+        }
+        assert_eq!(parse_pool_capacity_override(), None);
 
-    #[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);
+        unsafe {
+            std::env::set_var("GTE_SESSION_POOL_CAP", "0");
+        }
+        assert_eq!(parse_pool_capacity_override(), None);
+
+        unsafe {
+            std::env::set_var("GTE_SESSION_POOL_CAP", "4");
+        }
+        assert_eq!(parse_pool_capacity_override(), Some(4));
+
+        unsafe {
+            std::env::set_var("GTE_SESSION_POOL_CAP", "abc");
+        }
+        assert_eq!(parse_pool_capacity_override(), None);
+
+        unsafe {
+            std::env::remove_var("GTE_SESSION_POOL_CAP");
+        }
     }
 
     #[test]

data/ext/gte/tests/inference_integration_test.rs

@@ -8,7 +8,7 @@ fn model_dir(env_var: &str) -> Option<String> {
 #[test]
 fn test_e5_single_embedding_shape() {
     let Some(dir) = model_dir("GTE_BENCH_E5_DIR") else { return };
-    let embedder = Embedder::from_dir(&dir, 0, 3, ModelLoadOverrides::default())
+    let embedder = Embedder::from_dir(&dir, 0, ModelLoadOverrides::default())
         .expect("embedder should initialize");
     let result = embedder
         .embed(vec!["query: Hello world".to_string()])
@@ -21,7 +21,7 @@ fn test_e5_single_embedding_shape() {
 #[test]
 fn test_clip_single_embedding_shape() {
     let Some(dir) = model_dir("GTE_BENCH_CLIP_DIR") else { return };
-    let embedder = Embedder::from_dir(&dir, 0, 3, ModelLoadOverrides::default())
+    let embedder = Embedder::from_dir(&dir, 0, ModelLoadOverrides::default())
         .expect("embedder should initialize");
     let result = embedder
         .embed(vec!["a photo of a cat".to_string()])
@@ -34,7 +34,7 @@ fn test_clip_single_embedding_shape() {
 #[test]
 fn test_e5_batch_embedding_shape() {
     let Some(dir) = model_dir("GTE_BENCH_E5_DIR") else { return };
-    let embedder = Embedder::from_dir(&dir, 0, 3, ModelLoadOverrides::default())
+    let embedder = Embedder::from_dir(&dir, 0, ModelLoadOverrides::default())
         .expect("embedder should initialize");
     let texts = vec![
         "query: first sentence".to_string(),
@@ -51,7 +51,7 @@ fn test_e5_batch_embedding_shape() {
 #[test]
 fn test_e5_long_input_truncation_no_error() {
     let Some(dir) = model_dir("GTE_BENCH_E5_DIR") else { return };
-    let embedder = Embedder::from_dir(&dir, 0, 3, ModelLoadOverrides::default())
+    let embedder = Embedder::from_dir(&dir, 0, ModelLoadOverrides::default())
         .expect("embedder should initialize");
     let very_long_text = "word ".repeat(1000);
     let result = embedder

data/lib/gte/config.rb

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

data/lib/gte/embedder.rb

@@ -2,7 +2,6 @@
 
 module GTE
   class Embedder
-    DEFAULT_THREADS = 0
     DEFAULT_OPTIMIZATION_LEVEL = 3
 
     class << self
@@ -15,7 +14,6 @@ module GTE
       def from_config(config)
         new(
           config.model_dir,
-          config.threads,
           config.optimization_level,
           config.model_name.to_s,
           config.normalize,
@@ -29,7 +27,6 @@ module GTE
       def default_config(model_dir)
         Config::Text.new(
           model_dir: File.expand_path(model_dir),
-          threads: DEFAULT_THREADS,
           optimization_level: DEFAULT_OPTIMIZATION_LEVEL,
           model_name: nil,
           normalize: true,

data/lib/gte/reranker.rb

@@ -19,7 +19,6 @@ module GTE
       def default_config(model_dir)
         Config::Reranker.new(
           model_dir: File.expand_path(model_dir),
-          threads: 1,
           optimization_level: 3,
           model_name: nil,
           sigmoid: false,
@@ -33,7 +32,6 @@ module GTE
       def build(cfg)
         new(
           cfg.model_dir,
-          cfg.threads,
           cfg.optimization_level,
           cfg.model_name.to_s,
           cfg.sigmoid,

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: gte
 version: !ruby/object:Gem::Version
-  version: 0.0.11
+  version: 0.0.13
 platform: ruby
 authors:
 - elcuervo