gte 0.0.15 → 0.0.16

data/ext/gte/tests/padding_regression_test.rs

@@ -1,40 +1,24 @@
-// Regression tests for the fixed-padding performance bug.
-//
-// Root cause: PaddingMode::Auto silently read "padding.strategy.Fixed: N" from
-// tokenizer.json and applied it, padding every input to max_length tokens.
-// A query like "cat" (1 token) was padded to 64 tokens for Siglip2, making
-// inference ~6x slower (44ms vs 7ms measured on Heroku).
-//
-// These tests use tests/fixtures/minimal/tokenizer.json which has
-// "padding.strategy.Fixed: 64" baked in — exactly the condition that triggered
-// the regression in production models like Siglip2.
-
 use gte::model_config::PaddingMode;
 use gte::tokenizer::Tokenizer;
 
 const TOKENIZER: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/minimal/tokenizer.json");
 
-// Short input tokenizes to 1 token with this vocabulary.
 const SHORT_INPUT: &str = "cat";
 const MAX_LENGTH: usize = 64;
 
 #[test]
 fn auto_padding_uses_batch_longest_regardless_of_tokenizer_json() {
-    // fixed_padding_length: Some(MAX_LENGTH) simulates what model_profile::read_tokenizer_profile
-    // returns when tokenizer.json has "padding.strategy.Fixed: 64".
     let tokenizer = Tokenizer::new(TOKENIZER, MAX_LENGTH, false, PaddingMode::Auto, Some(MAX_LENGTH))
         .expect("tokenizer should load");
 
     let tokenized = tokenizer.tokenize(&[SHORT_INPUT.to_string()]).expect("tokenize should succeed");
 
-    // Old behavior: cols == 64 (silently padded to max_length)
-    // New behavior: cols == actual token count (1 for "cat")
     assert!(
-        tokenized.cols < MAX_LENGTH,
+        tokenized.input_ids.ncols() < MAX_LENGTH,
         "Auto padding should use batch_longest, got cols={} (expected < {}). \
          This is the Siglip2 regression: short queries were padded to max_length, \
          making inference ~6x slower.",
-        tokenized.cols,
+        tokenized.input_ids.ncols(),
         MAX_LENGTH
     );
 }
@@ -46,7 +30,7 @@ fn fixed_padding_mode_pads_to_max_length() {
 
     let tokenized = tokenizer.tokenize(&[SHORT_INPUT.to_string()]).expect("tokenize should succeed");
 
-    assert_eq!(tokenized.cols, MAX_LENGTH, "Fixed mode should pad to max_length");
+    assert_eq!(tokenized.input_ids.ncols(), MAX_LENGTH, "Fixed mode should pad to max_length");
     assert_eq!(tokenized.input_ids.len(), MAX_LENGTH);
     assert_eq!(tokenized.attn_masks.len(), MAX_LENGTH);
 }
@@ -56,26 +40,24 @@ fn batch_longest_padding_uses_longest_sequence_in_batch() {
     let tokenizer =
         Tokenizer::new(TOKENIZER, MAX_LENGTH, false, PaddingMode::BatchLongest, None).expect("tokenizer should load");
 
-    // "cat" = 1 token, "hello world" = 2 tokens — batch pads to 2, not 64
     let tokenized =
         tokenizer.tokenize(&["cat".to_string(), "hello world".to_string()]).expect("tokenize should succeed");
 
-    assert_eq!(tokenized.rows, 2);
+    assert_eq!(tokenized.input_ids.nrows(), 2);
     assert!(
-        tokenized.cols < MAX_LENGTH,
+        tokenized.input_ids.ncols() < MAX_LENGTH,
         "BatchLongest should pad to longest in batch (2 tokens), not max_length ({}). Got cols={}",
         MAX_LENGTH,
-        tokenized.cols
+        tokenized.input_ids.ncols()
     );
 }
 
 #[test]
 fn auto_padding_with_no_fixed_hint_also_uses_batch_longest() {
-    // Sanity check: Auto with fixed_padding_length=None also uses BatchLongest
     let tokenizer =
         Tokenizer::new(TOKENIZER, MAX_LENGTH, false, PaddingMode::Auto, None).expect("tokenizer should load");
 
     let tokenized = tokenizer.tokenize(&[SHORT_INPUT.to_string()]).expect("tokenize should succeed");
 
-    assert!(tokenized.cols < MAX_LENGTH);
+    assert!(tokenized.input_ids.ncols() < MAX_LENGTH);
 }

data/ext/gte/tests/tokenizer_unit_test.rs

@@ -12,13 +12,13 @@ fn test_e5_tokenizer_output_shape() {
 
     let tokenized = tokenizer.tokenize(&texts).expect("tokenize should succeed");
 
-    assert_eq!(tokenized.rows, 2, "batch size should be 2");
-    assert!(tokenized.cols > 0, "sequence length should be non-zero");
-    assert_eq!(tokenized.input_ids.len(), tokenized.rows * tokenized.cols);
-    assert_eq!(tokenized.attn_masks.len(), tokenized.rows * tokenized.cols);
+    assert_eq!(tokenized.input_ids.nrows(), 2, "batch size should be 2");
+    assert!(tokenized.input_ids.ncols() > 0, "sequence length should be non-zero");
+    assert_eq!(tokenized.input_ids.len(), tokenized.input_ids.nrows() * tokenized.input_ids.ncols());
+    assert_eq!(tokenized.attn_masks.len(), tokenized.attn_masks.nrows() * tokenized.attn_masks.ncols());
 
     let type_ids = tokenized.type_ids.as_ref().expect("type_ids should exist");
-    assert_eq!(type_ids.len(), tokenized.rows * tokenized.cols);
+    assert_eq!(type_ids.len(), type_ids.nrows() * type_ids.ncols());
 }
 
 #[test]
@@ -31,6 +31,6 @@ fn test_e5_truncation_at_max_length() {
     let long_text = "word ".repeat(200);
     let tokenized = tokenizer.tokenize(&[long_text]).expect("tokenize should not error on long input");
 
-    assert_eq!(tokenized.rows, 1);
-    assert_eq!(tokenized.cols, 16, "sequence length should be truncated to max_length");
+    assert_eq!(tokenized.input_ids.nrows(), 1);
+    assert_eq!(tokenized.input_ids.ncols(), 16, "sequence length should be truncated to max_length");
 }

data/lib/gte/config.rb

@@ -4,8 +4,7 @@ module GTE
   module Config
     Text = Data.define(
       :model_dir, :optimization_level,
-      :model_name, :normalize, :output_tensor, :max_length, :padding, :execution_providers,
-      :lowercase_input, :max_input_chars
+      :model_name, :output_tensor, :max_length, :padding, :execution_providers
     )
 
     Reranker = Data.define(

data/lib/gte/embedder.rb

@@ -5,24 +5,15 @@ module GTE
     DEFAULT_OPTIMIZATION_LEVEL = 3
 
     class << self
-      def config(model_dir)
-        cfg = default_config(model_dir)
-        cfg = yield(cfg) if block_given?
-        from_config(cfg)
-      end
-
       def from_config(config)
         new(
           config.model_dir,
           config.optimization_level,
           config.model_name.to_s,
-          config.normalize,
           config.output_tensor.to_s,
           config.max_length || 0,
           config.padding.to_s,
-          config.execution_providers.to_s,
-          config.lowercase_input ? true : false,
-          config.max_input_chars || 0
+          config.execution_providers.to_s
         )
       end
 
@@ -31,13 +22,10 @@ module GTE
           model_dir: File.expand_path(model_dir),
           optimization_level: DEFAULT_OPTIMIZATION_LEVEL,
           model_name: nil,
-          normalize: true,
           output_tensor: nil,
           max_length: nil,
           padding: nil,
-          execution_providers: nil,
-          lowercase_input: false,
-          max_input_chars: nil
+          execution_providers: nil
         )
       end
     end

data/lib/gte/model.rb

@@ -17,13 +17,6 @@ module GTE
       @embedder.embed(Array(texts))
     end
 
-    def [](input)
-      case input
-      when String then embed(input).row(0)
-      when Array then embed(input)
-      end
-    end
-
     def embed_binary(text)
       embed(text).row_binary_f32(0)
     end

data/lib/gte/reranker.rb

@@ -3,15 +3,21 @@
 module GTE
   class Reranker
     class << self
-      def config(model_dir)
-        cfg = default_config(model_dir)
-
-        if block_given?
-          yielded = yield(cfg)
-          cfg = yielded if yielded.is_a?(Config::Reranker)
-        end
+      alias native_new new
 
-        build(cfg)
+      def new(model_dir, &block)
+        cfg = default_config(model_dir)
+        cfg = block.call(cfg) if block
+        native_new(
+          cfg.model_dir,
+          cfg.optimization_level,
+          cfg.model_name.to_s,
+          cfg.sigmoid,
+          cfg.output_tensor.to_s,
+          cfg.max_length || 0,
+          cfg.padding.to_s,
+          cfg.execution_providers.to_s
+        )
       end
 
       private
@@ -28,31 +34,6 @@ module GTE
           execution_providers: nil
         )
       end
-
-      def build(cfg)
-        new(
-          cfg.model_dir,
-          cfg.optimization_level,
-          cfg.model_name.to_s,
-          cfg.sigmoid,
-          cfg.output_tensor.to_s,
-          cfg.max_length || 0,
-          cfg.padding.to_s,
-          cfg.execution_providers.to_s,
-          false, # lowercase_input
-          0 # max_input_chars
-        )
-      end
-    end
-
-    def rerank(query:, candidates:)
-      rows = Array(candidates).map(&:to_s)
-      scores = score(query.to_s, rows)
-
-      rows
-        .each_with_index
-        .map { |text, idx| { index: idx, score: scores[idx], text: text } }
-        .sort_by { |row| -row[:score] }
     end
   end
 end

data/lib/gte.rb

@@ -14,30 +14,9 @@ require 'gte/model'
 require 'gte/reranker'
 
 module GTE
-  @model_cache_mutex = Mutex.new
-  @model_cache = {}
-
-  class << self
-    def config(model_dir)
-      cfg = Embedder.default_config(model_dir)
-
-      cfg = yield(cfg) if block_given?
-
-      @model_cache_mutex.synchronize do
-        @model_cache[cache_key(cfg)] ||= Model.new(cfg)
-      end
-    end
-
-    def warmup(runner, threads:)
-      threads.times.map do
-        Thread.new { runner.embed('warmup') }
-      end.each(&:join)
-    end
-
-    private
-
-    def cache_key(cfg)
-      cfg.to_h
-    end
+  def self.config(model_dir, &block)
+    cfg = Embedder.default_config(model_dir)
+    cfg = block.call(cfg) if block
+    Model.new(cfg)
   end
 end

metadata

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