gte 0.0.3-arm64-darwin → 0.0.4-arm64-darwin

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: e86f2a3e303b2ad2dd726f2a14b04127f07591d94f220aae65a39e922818756f
-  data.tar.gz: cccb448f456718089b73080ff29541e1c525fa28c0fc4e27674dede3bd59389f
+  metadata.gz: 86bd8582b693e5d6b1b1d880dec05b21912b3cc45de1a09e143210fcdb76dc5d
+  data.tar.gz: 8785600d70d86d0f9de669631054f279105801f47ca6f6d227b1e93fe9e6ff9f
 SHA512:
-  metadata.gz: 72ec78ebb4c400bebf8cd52dc44aa034cf8ddbf5911e81547eb8120087ae5f72b32873f04be515aeb7170374eab938d1031d80ae6a680b1f37f43441c4acbaa4
-  data.tar.gz: bc00deb09690233e212526e1cd5115af1a813d3e2f0e55445bbf90d94e169a25cef2a0c07c64d8ac6d98224eaa58166bdbdb26a6c715a584d7a634598e1e6ae8
+  metadata.gz: 7ac75c64e28ca6c227be56439f15bea76355c775cec5ba203c95666b8f603d370b492647a4b9c242b84d9fdfa36f6f34fa2436895596f5b887863b193bf4e523
+  data.tar.gz: 2b9f12675b6e0fd461f611f199e7ee66f366e3bbdf93c15bd018a0b45eb2eef001b9adb9670b02a07f58589a915ab7391762c5fabb54c14254193ff11d1867ec

data/README.md

@@ -13,6 +13,12 @@ model = GTE.new(ENV.fetch("GTE_MODEL_DIR"))
 vector = model["query: hello world"]
 ```
 
+For Puma or other thread pools, prefer process-local reuse:
+
+```ruby
+MODEL = GTE.new(ENV.fetch("GTE_MODEL_DIR"))
+```
+
 ## Model Directory
 
 A model directory must include `tokenizer.json` and one ONNX model, resolved in this order:
@@ -40,6 +46,7 @@ The repo includes two benchmark paths:
 bundle exec rake bench:pure_compare
 bundle exec rake bench:puma_compare
 bundle exec rake bench:matrix_sweep
+bundle exec ruby bench/memory_probe.rb --compare-pure
 ```
 
 For release tracking and regression detection, record a run entry in `RUNS.md`:

data/Rakefile

@@ -48,6 +48,14 @@ namespace :bench do
     )
   end
 
+  desc 'Run memory probe for single-instance vs duplicate-instance behavior'
+  task :memory_probe do
+    run_in_nix(
+      'bundle', 'exec', 'ruby', 'bench/memory_probe.rb',
+      '--compare-pure'
+    )
+  end
+
   desc 'Run Puma benchmark, append RUNS.md entry, and enforce goal/regression checks'
   task :record_run do
     run_in_nix(

data/VERSION

@@ -1 +1 @@
-0.0.3
+0.0.4

data/ext/gte/Cargo.toml

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

data/ext/gte/src/embedder.rs

@@ -57,16 +57,17 @@ impl Embedder {
         }
 
         let max_length = read_max_length(dir);
+        let probe_num_threads = if num_threads == 0 { 1 } else { num_threads };
         let temp_config = ModelConfig {
             max_length,
             output_tensor: String::new(),
             mode: ExtractorMode::Raw,
             with_type_ids: false,
             with_attention_mask: true,
-            num_threads,
+            num_threads: probe_num_threads,
             optimization_level,
         };
-        let session = build_session(&model_path, &temp_config)?;
+        let mut session = build_session(&model_path, &temp_config)?;
 
         validate_supported_inputs(&session)?;
         let with_type_ids = session.inputs.iter().any(|i| i.name == "token_type_ids");
@@ -97,11 +98,11 @@ impl Embedder {
             optimization_level,
         };
 
-        let session = if tuned_num_threads != num_threads {
-            build_session(&model_path, &config)?
-        } else {
-            session
-        };
+        if tuned_num_threads != probe_num_threads {
+            // Release probe session before rebuilding to minimize transient peak RSS.
+            drop(session);
+            session = build_session(&model_path, &config)?;
+        }
 
         let tokenizer = Tokenizer::new(&tokenizer_path, config.max_length, config.with_type_ids)?;
 
@@ -138,16 +139,13 @@ fn tune_num_threads(
     }
 
     let family = infer_model_family(with_attention_mask, with_type_ids, output_name);
-    let target_concurrency = puma_target_concurrency();
-    let host_cores = host_parallelism();
-    let budgeted_threads = (host_cores / target_concurrency).max(1);
 
     match family {
         // Puma-like workloads typically run many concurrent single-item requests where
         // one intra-op thread per request gives the best tail behavior.
-        ModelFamily::E5Like | ModelFamily::ClipLike | ModelFamily::SiglipLike => {
-            budgeted_threads.min(1)
-        }
+        ModelFamily::E5Like | ModelFamily::ClipLike => 1,
+        // Siglip2 text path benefits from a small intra-op pool under concurrency.
+        ModelFamily::SiglipLike => 3,
         ModelFamily::Other => 0,
     }
 }
@@ -169,20 +167,6 @@ fn infer_model_family(
     ModelFamily::Other
 }
 
-fn puma_target_concurrency() -> usize {
-    std::env::var("GTE_PUMA_CONCURRENCY")
-        .ok()
-        .and_then(|raw| raw.parse::<usize>().ok())
-        .filter(|value| *value > 0)
-        .unwrap_or(16)
-}
-
-fn host_parallelism() -> usize {
-    std::thread::available_parallelism()
-        .map(|n| n.get())
-        .unwrap_or(1)
-}
-
 fn resolve_named_model(dir: &Path, name: &str) -> Result<PathBuf> {
     let candidates = [dir.join("onnx").join(name), dir.join(name)];
     for path in &candidates {

data/ext/gte/src/postprocess.rs

@@ -87,10 +87,27 @@ fn mean_pool_contiguous(
         let mask_base = batch_index * seq;
         let hidden_base = batch_index * seq * dim;
         let output_row = &mut output[batch_index * dim..(batch_index + 1) * dim];
+        let mask_row = &attention_mask[mask_base..mask_base + seq];
+
+        if mask_row.iter().all(|&weight| weight == 1) {
+            for token_index in 0..seq {
+                let token_base = hidden_base + token_index * dim;
+                for dim_index in 0..dim {
+                    output_row[dim_index] += hidden[token_base + dim_index];
+                }
+            }
+
+            let inverse = (seq as f32).recip();
+            for value in output_row {
+                *value *= inverse;
+            }
+            continue;
+        }
+
         let mut weight_sum = 0.0f32;
 
-        for token_index in 0..seq {
-            let weight = attention_mask[mask_base + token_index];
+        for (token_index, &weight_raw) in mask_row.iter().enumerate() {
+            let weight = weight_raw;
             if weight <= 0 {
                 continue;
             }

data/lib/gte.rb

@@ -9,6 +9,9 @@ end
 module GTE
   VERSION = File.read(File.expand_path('../VERSION', __dir__)).strip
 
+  @model_cache_mutex = Mutex.new
+  @model_cache = {}
+
   class Model
     def initialize(dir, num_threads: 0, optimization_level: 3, model_name: nil)
       @embedder = GTE::Embedder.new(dir, num_threads, optimization_level, model_name.to_s)
@@ -30,7 +33,23 @@ module GTE
     end
   end
 
-  def self.new(dir, num_threads: 0, optimization_level: 3, model_name: nil)
-    Model.new(dir, num_threads: num_threads, optimization_level: optimization_level, model_name: model_name)
+  def self.new(dir, threads: 0, optimization: 3, model_name: nil)
+    key = [
+      File.expand_path(dir),
+      Integer(threads),
+      Integer(optimization),
+      model_name.to_s
+    ].freeze
+
+    @model_cache_mutex.synchronize do
+      @model_cache[key] ||= Model.new(
+        key[0],
+        num_threads: key[1],
+        optimization_level: key[2],
+        model_name: key[3].empty? ? nil : key[3]
+      )
+    end
   end
+
+  def self.fetch(*) = new(*)
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: gte
 version: !ruby/object:Gem::Version
-  version: 0.0.3
+  version: 0.0.4
 platform: arm64-darwin
 authors:
 - elcuervo
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2026-04-10 00:00:00.000000000 Z
+date: 2026-04-13 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rake