gte 0.0.13 → 0.0.14

data/ext/gte/src/tokenizer.rs

@@ -24,69 +24,51 @@ impl Tokenizer {
         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()))?;
-
-        let truncation = TruncationParams {
-            max_length,
-            ..Default::default()
-        };
-        tokenizer
-            .with_truncation(Some(truncation))
-            .map_err(|e| GteError::Tokenizer(e.to_string()))?;
-
-        let padding = PaddingParams {
-            strategy: resolve_padding_strategy(padding_mode, max_length, fixed_padding_length),
-            ..Default::default()
-        };
-        tokenizer.with_padding(Some(padding));
-
-        Ok(Self {
-            tokenizer,
-            with_type_ids,
-        })
+        #[allow(unused_results)]
+        {
+            let mut tokenizer =
+                tokenizers::Tokenizer::from_file(tokenizer_path).map_err(|e| GteError::Tokenizer(e.to_string()))?;
+
+            let truncation = TruncationParams { max_length, ..Default::default() };
+            let padding = PaddingParams {
+                strategy: resolve_padding_strategy(padding_mode, max_length, fixed_padding_length),
+                ..Default::default()
+            };
+            tokenizer.with_truncation(Some(truncation)).map_err(|e| GteError::Tokenizer(e.to_string()))?;
+            tokenizer.with_padding(Some(padding));
+
+            Ok(Self { tokenizer, with_type_ids })
+        }
     }
 
     pub fn tokenize(&self, texts: &[String]) -> Result<Tokenized> {
         if texts.len() == 1 {
-            let encoding = self
-                .tokenizer
-                .encode_fast(texts[0].as_str(), true)
-                .map_err(|e| GteError::Tokenizer(e.to_string()))?;
-            return build_tokenized_single(&encoding, self.with_type_ids);
+            let encoding =
+                self.tokenizer.encode_fast(texts[0].as_str(), true).map_err(|e| GteError::Tokenizer(e.to_string()))?;
+            return Ok(build_tokenized_single(&encoding, self.with_type_ids));
         }
 
         let encode_inputs: Vec<&str> = texts.iter().map(String::as_str).collect();
-        let encodings = self
-            .tokenizer
-            .encode_batch_fast(encode_inputs, true)
-            .map_err(|e| GteError::Tokenizer(e.to_string()))?;
+        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)
+        Ok(build_tokenized(&encodings, self.with_type_ids))
     }
 
     pub fn tokenize_pairs(&self, pairs: &[(String, String)]) -> Result<Tokenized> {
-        let encode_inputs: Vec<tokenizers::EncodeInput<'_>> = pairs
-            .iter()
-            .map(|(left, right)| (left.as_str(), right.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)
+        let encode_inputs: Vec<tokenizers::EncodeInput<'_>> =
+            pairs.iter().map(|(left, right)| (left.as_str(), right.as_str()).into()).collect();
+        let encodings =
+            self.tokenizer.encode_batch_fast(encode_inputs, true).map_err(|e| GteError::Tokenizer(e.to_string()))?;
+        Ok(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)
+        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()))?;
+        Ok(build_tokenized(&encodings, self.with_type_ids))
     }
 }
 
@@ -102,8 +84,7 @@ pub fn parse_padding_mode_override(value: Option<&str>) -> Result<Option<Padding
         "fixed" => PaddingMode::Fixed,
         _ => {
             return Err(GteError::Inference(format!(
-                "invalid padding mode '{}'; expected one of: auto, batch_longest, fixed",
-                raw
+                "invalid padding mode '{raw}'; expected one of: auto, batch_longest, fixed"
             )))
         }
     };
@@ -121,45 +102,20 @@ fn resolve_padding_strategy(
     }
 }
 
-fn build_tokenized_single(
-    encoding: &tokenizers::Encoding,
-    with_type_ids: bool,
-) -> Result<Tokenized> {
+fn build_tokenized_single(encoding: &tokenizers::Encoding, with_type_ids: bool) -> Tokenized {
     let cols = encoding.len();
 
-    let input_ids: Vec<i64> = encoding
-        .get_ids()
-        .iter()
-        .map(|&value| i64::from(value))
-        .collect();
-    let attn_masks: Vec<i64> = encoding
-        .get_attention_mask()
-        .iter()
-        .map(|&value| i64::from(value))
-        .collect();
-    let type_ids: Option<Vec<i64>> = with_type_ids.then(|| {
-        encoding
-            .get_type_ids()
-            .iter()
-            .map(|&value| i64::from(value))
-            .collect()
-    });
-
-    Ok(Tokenized {
-        rows: 1,
-        cols,
-        input_ids,
-        attn_masks,
-        type_ids,
-    })
+    let input_ids: Vec<i64> = encoding.get_ids().iter().map(|&v| i64::from(v)).collect();
+    let attn_masks: Vec<i64> = encoding.get_attention_mask().iter().map(|&v| i64::from(v)).collect();
+    let type_ids: Option<Vec<i64>> =
+        with_type_ids.then(|| encoding.get_type_ids().iter().map(|&v| i64::from(v)).collect());
+
+    Tokenized { rows: 1, cols, input_ids, attn_masks, type_ids }
 }
 
-fn build_tokenized(encodings: &[tokenizers::Encoding], with_type_ids: bool) -> Result<Tokenized> {
+fn build_tokenized(encodings: &[tokenizers::Encoding], with_type_ids: bool) -> Tokenized {
     let rows = encodings.len();
-    let cols = encodings
-        .first()
-        .map(|encoding| encoding.len())
-        .unwrap_or(0);
+    let cols = encodings.first().map_or(0, tokenizers::Encoding::len);
     let len = rows * cols;
 
     let mut input_ids = Vec::with_capacity(len);
@@ -167,27 +123,15 @@ 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 {
-        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));
-        }
+        input_ids.extend(encoding.get_ids().iter().map(|&v| i64::from(v)));
+        attn_masks.extend(encoding.get_attention_mask().iter().map(|&v| i64::from(v)));
 
         if let Some(type_ids) = type_ids.as_mut() {
-            for &value in encoding.get_type_ids() {
-                type_ids.push(i64::from(value));
-            }
+            type_ids.extend(encoding.get_type_ids().iter().map(|&v| i64::from(v)));
         }
     }
 
-    Ok(Tokenized {
-        rows,
-        cols,
-        input_ids,
-        attn_masks,
-        type_ids,
-    })
+    Tokenized { rows, cols, input_ids, attn_masks, type_ids }
 }
 
 #[cfg(test)]
@@ -198,18 +142,9 @@ mod tests {
 
     #[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)
-        );
+        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]
@@ -222,21 +157,12 @@ mod tests {
         // Auto ignores fixed_padding_length from tokenizer.json — BatchLongest is
         // always faster for inference and correct for variable-length inputs.
         // Use PaddingMode::Fixed explicitly when fixed-length padding is required.
-        assert!(matches!(
-            resolve_padding_strategy(PaddingMode::Auto, 64, Some(64)),
-            PaddingStrategy::BatchLongest
-        ));
-        assert!(matches!(
-            resolve_padding_strategy(PaddingMode::Auto, 512, None),
-            PaddingStrategy::BatchLongest
-        ));
+        assert!(matches!(resolve_padding_strategy(PaddingMode::Auto, 64, Some(64)), PaddingStrategy::BatchLongest));
+        assert!(matches!(resolve_padding_strategy(PaddingMode::Auto, 512, None), PaddingStrategy::BatchLongest));
     }
 
     #[test]
     fn resolve_padding_strategy_fixed_uses_max_length() {
-        assert!(matches!(
-            resolve_padding_strategy(PaddingMode::Fixed, 64, None),
-            PaddingStrategy::Fixed(64)
-        ));
+        assert!(matches!(resolve_padding_strategy(PaddingMode::Fixed, 64, None), PaddingStrategy::Fixed(64)));
     }
 }

data/ext/gte/tests/inference_integration_test.rs

@@ -8,11 +8,8 @@ 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, ModelLoadOverrides::default())
-        .expect("embedder should initialize");
-    let result = embedder
-        .embed(vec!["query: Hello world".to_string()])
-        .expect("embed should succeed");
+    let embedder = Embedder::from_dir(&dir, 0, ModelLoadOverrides::default()).expect("embedder should initialize");
+    let result = embedder.embed(&["query: Hello world".to_string()]).expect("embed should succeed");
 
     assert_eq!(result.shape()[0], 1);
     assert!(result.shape()[1] > 0);
@@ -21,11 +18,8 @@ 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, ModelLoadOverrides::default())
-        .expect("embedder should initialize");
-    let result = embedder
-        .embed(vec!["a photo of a cat".to_string()])
-        .expect("embed should succeed");
+    let embedder = Embedder::from_dir(&dir, 0, ModelLoadOverrides::default()).expect("embedder should initialize");
+    let result = embedder.embed(&["a photo of a cat".to_string()]).expect("embed should succeed");
 
     assert_eq!(result.shape()[0], 1);
     assert!(result.shape()[1] > 0);
@@ -34,15 +28,14 @@ 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, ModelLoadOverrides::default())
-        .expect("embedder should initialize");
+    let embedder = Embedder::from_dir(&dir, 0, ModelLoadOverrides::default()).expect("embedder should initialize");
     let texts = vec![
         "query: first sentence".to_string(),
         "query: second sentence".to_string(),
         "query: third sentence for batch".to_string(),
     ];
 
-    let result = embedder.embed(texts).expect("batch embed should succeed");
+    let result = embedder.embed(&texts).expect("batch embed should succeed");
 
     assert_eq!(result.shape()[0], 3);
     assert!(result.shape()[1] > 0);
@@ -51,12 +44,9 @@ 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, ModelLoadOverrides::default())
-        .expect("embedder should initialize");
+    let embedder = Embedder::from_dir(&dir, 0, ModelLoadOverrides::default()).expect("embedder should initialize");
     let very_long_text = "word ".repeat(1000);
-    let result = embedder
-        .embed(vec![very_long_text])
-        .expect("long input should be truncated without error");
+    let result = embedder.embed(&[very_long_text]).expect("long input should be truncated without error");
 
     assert_eq!(result.shape()[0], 1);
     assert!(result.shape()[1] > 0);

data/ext/gte/tests/padding_regression_test.rs

@@ -12,10 +12,7 @@
 use gte::model_config::PaddingMode;
 use gte::tokenizer::Tokenizer;
 
-const TOKENIZER: &str = concat!(
-    env!("CARGO_MANIFEST_DIR"),
-    "/tests/fixtures/minimal/tokenizer.json"
-);
+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";
@@ -28,9 +25,7 @@ fn auto_padding_uses_batch_longest_regardless_of_tokenizer_json() {
     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");
+    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")
@@ -46,30 +41,24 @@ fn auto_padding_uses_batch_longest_regardless_of_tokenizer_json() {
 
 #[test]
 fn fixed_padding_mode_pads_to_max_length() {
-    let tokenizer = Tokenizer::new(TOKENIZER, MAX_LENGTH, false, PaddingMode::Fixed, None)
-        .expect("tokenizer should load");
+    let tokenizer =
+        Tokenizer::new(TOKENIZER, MAX_LENGTH, false, PaddingMode::Fixed, None).expect("tokenizer should load");
 
-    let tokenized = tokenizer
-        .tokenize(&[SHORT_INPUT.to_string()])
-        .expect("tokenize should succeed");
+    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.cols, 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);
 }
 
 #[test]
 fn batch_longest_padding_uses_longest_sequence_in_batch() {
-    let tokenizer = Tokenizer::new(TOKENIZER, MAX_LENGTH, false, PaddingMode::BatchLongest, None)
-        .expect("tokenizer should load");
+    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");
+    let tokenized =
+        tokenizer.tokenize(&["cat".to_string(), "hello world".to_string()]).expect("tokenize should succeed");
 
     assert_eq!(tokenized.rows, 2);
     assert!(
@@ -83,12 +72,10 @@ fn batch_longest_padding_uses_longest_sequence_in_batch() {
 #[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 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");
+    let tokenized = tokenizer.tokenize(&[SHORT_INPUT.to_string()]).expect("tokenize should succeed");
 
     assert!(tokenized.cols < MAX_LENGTH);
 }

data/ext/gte/tests/tokenizer_unit_test.rs

@@ -4,17 +4,11 @@ use gte::tokenizer::Tokenizer;
 #[test]
 #[ignore = "requires ext/gte/tests/fixtures/e5/tokenizer.json"]
 fn test_e5_tokenizer_output_shape() {
-    const TOKENIZER: &str = concat!(
-        env!("CARGO_MANIFEST_DIR"),
-        "/tests/fixtures/e5/tokenizer.json"
-    );
-
-    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(),
-    ];
+    const TOKENIZER: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/e5/tokenizer.json");
+
+    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()];
 
     let tokenized = tokenizer.tokenize(&texts).expect("tokenize should succeed");
 
@@ -30,21 +24,13 @@ fn test_e5_tokenizer_output_shape() {
 #[test]
 #[ignore = "requires ext/gte/tests/fixtures/e5/tokenizer.json"]
 fn test_e5_truncation_at_max_length() {
-    const TOKENIZER: &str = concat!(
-        env!("CARGO_MANIFEST_DIR"),
-        "/tests/fixtures/e5/tokenizer.json"
-    );
+    const TOKENIZER: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/e5/tokenizer.json");
 
-    let tokenizer = Tokenizer::new(TOKENIZER, 16, false, PaddingMode::BatchLongest, None)
-        .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])
-        .expect("tokenize should not error on long input");
+    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.cols, 16, "sequence length should be truncated to max_length");
 }

data/lib/gte/config.rb

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

data/lib/gte/embedder.rb

@@ -20,7 +20,9 @@ module GTE
           config.output_tensor.to_s,
           config.max_length || 0,
           config.padding.to_s,
-          config.execution_providers.to_s
+          config.execution_providers.to_s,
+          config.lowercase_input ? true : false,
+          config.max_input_chars || 0
         )
       end
 
@@ -33,7 +35,9 @@ module GTE
           output_tensor: nil,
           max_length: nil,
           padding: nil,
-          execution_providers: nil
+          execution_providers: nil,
+          lowercase_input: false,
+          max_input_chars: nil
         )
       end
     end

data/lib/gte/reranker.rb

@@ -38,7 +38,9 @@ module GTE
           cfg.output_tensor.to_s,
           cfg.max_length || 0,
           cfg.padding.to_s,
-          cfg.execution_providers.to_s
+          cfg.execution_providers.to_s,
+          false, # lowercase_input
+          0 # max_input_chars
         )
       end
     end

data/lib/gte.rb

@@ -28,6 +28,12 @@ module GTE
       end
     end
 
+    def warmup(runner, threads:)
+      threads.times.map do
+        Thread.new { runner.embed('warmup') }
+      end.each(&:join)
+    end
+
     private
 
     def cache_key(cfg)

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: gte
 version: !ruby/object:Gem::Version
-  version: 0.0.13
+  version: 0.0.14
 platform: ruby
 authors:
 - elcuervo
@@ -95,6 +95,7 @@ files:
 - ext/gte/benches/hot_path.rs
 - ext/gte/build.rs
 - ext/gte/extconf.rb
+- ext/gte/rustfmt.toml
 - ext/gte/src/embedder.rs
 - ext/gte/src/error.rs
 - ext/gte/src/lib.rs