gte 0.0.6 → 0.0.7

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: fc149108c647dc5b14154bfbdc4975b53670b9ed3cf7d80760cc2b415c935a48
-  data.tar.gz: 32a682a95d56c8fab8d0d64a7ada0c0347ae796b6aefe6191f9aca8fc96426c2
+  metadata.gz: 29659e3ab6072d858b1710a779c3d2e5981f7749782182d141ccd5e9790a1fbb
+  data.tar.gz: c42d51cfa1a2ba6a2e83249e8a725c978b11c7ef80c6d69f09a64e884be42031
 SHA512:
-  metadata.gz: f5c69d954f51a51521b143b576942a9c0505ad60574c1727f963dd79e0b6c22cacc4e6d9af75394ae06f451521dbc788af51f1e79397a5cc66a41b4ce1b31933
-  data.tar.gz: 9e75fdbc9b5c8cfdd9d0e377a7e4a944057ec604e38ab23d960c4ed75ec6a72ce1dd27c2dd1bb2802721387babdabe0996e0c42be34d17d98253e0582b375de1
+  metadata.gz: ff2c2b1450a6e82c07aacd2ec98437f03678d56eef9c5516f904021a54f59b2ba5c42b8f6af22b5c4b2dacea98615b99bc54d2c7cdc4e8fbccc1abc195fe9975
+  data.tar.gz: 04ca056458d40e2ba7fabcdbcab415a087d54802fb3bd86748dc901c2cf0ecb44072fd1820a73e3dcaca097f165df3e70bab747b38340cd738876af5f0ea7645

data/README.md

@@ -41,6 +41,7 @@ custom = GTE.config(ENV.fetch("GTE_MODEL_DIR")) do |config|
   config.with(
     output_tensor: "last_hidden_state",
     max_length: 256,
+    padding: "batch_longest",
     optimization_level: 3
   )
 end
@@ -55,6 +56,7 @@ Config fields and defaults:
 - `normalize`: `true` (L2 normalization at Ruby-facing API)
 - `output_tensor`: `nil` (auto-select output tensor)
 - `max_length`: `nil` (uses tokenizer/model defaults)
+- `padding`: `nil` (auto; accepts `auto`, `batch_longest`, `fixed`)
 - `execution_providers`: `nil` (falls back to `GTE_EXECUTION_PROVIDERS` / CPU default)
 
 Notes:
@@ -106,6 +108,7 @@ Reranker config fields and defaults:
 - `sigmoid`: `false` (set `true` if you want bounded [0,1] style scores)
 - `output_tensor`: `nil`
 - `max_length`: `nil`
+- `padding`: `nil` (auto; accepts `auto`, `batch_longest`, `fixed`)
 - `execution_providers`: `nil`
 
 ## Runtime + Result Examples

data/VERSION

@@ -1 +1 @@
-0.0.6
+0.0.7

data/ext/gte/Cargo.toml

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

data/ext/gte/src/embedder.rs

@@ -1,12 +1,12 @@
 use crate::error::{GteError, Result};
-use crate::model_config::{ExtractorMode, ModelConfig};
+use crate::model_config::{ExtractorMode, ModelConfig, ModelLoadOverrides, PaddingMode};
 use crate::model_profile::{
-    has_input, infer_extraction_mode, read_max_length, resolve_default_text_model, resolve_named_model,
-    resolve_tokenizer_path, select_output_tensor, validate_supported_text_inputs,
+    has_input, infer_extraction_mode, read_tokenizer_profile, resolve_default_text_model,
+    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::tokenizer::{Tokenized, Tokenizer};
+use crate::tokenizer::{parse_padding_mode_override, Tokenized, Tokenizer};
 use ndarray::Array2;
 use ort::session::Session;
 use std::path::Path;
@@ -23,7 +23,13 @@ impl Embedder {
         P1: AsRef<Path>,
         P2: AsRef<Path>,
     {
-        let tokenizer = Tokenizer::new(tokenizer_path, config.max_length, config.with_type_ids)?;
+        let tokenizer = Tokenizer::new(
+            tokenizer_path,
+            config.max_length,
+            config.with_type_ids,
+            config.padding_mode,
+            None,
+        )?;
         let session = build_session(model_path, &config)?;
         Ok(Self {
             tokenizer,
@@ -36,10 +42,7 @@ impl Embedder {
         dir: P,
         num_threads: usize,
         optimization_level: u8,
-        model_name: Option<&str>,
-        output_tensor_override: Option<&str>,
-        max_length_override: Option<usize>,
-        execution_providers_override: Option<&str>,
+        overrides: ModelLoadOverrides<'_>,
     ) -> Result<Self> {
         const PREFERRED_EMBEDDING_OUTPUTS: [&str; 4] = [
             "pooler_output",
@@ -50,31 +53,35 @@ impl Embedder {
 
         let dir = dir.as_ref();
         let tokenizer_path = resolve_tokenizer_path(dir)?;
-        let model_path = match model_name.filter(|s| !s.is_empty()) {
+        let model_path = match overrides.model_name.filter(|s| !s.is_empty()) {
             Some(name) => resolve_named_model(dir, name)?,
             None => resolve_default_text_model(dir)?,
         };
 
-        let max_length = if let Some(override_value) = max_length_override {
+        let tokenizer_profile = read_tokenizer_profile(dir);
+        let max_length = if let Some(override_value) = overrides.max_length {
             if override_value == 0 {
                 return Err(GteError::Inference(
                     "max_length override must be greater than 0".to_string(),
                 ));
             }
-            override_value
+            override_value.min(tokenizer_profile.safe_max_length)
         } else {
-            read_max_length(dir)
+            tokenizer_profile.default_max_length
         };
+        let padding_mode =
+            parse_padding_mode_override(overrides.padding)?.unwrap_or(PaddingMode::Auto);
 
         let session_config = ModelConfig {
             max_length,
+            padding_mode,
             output_tensor: String::new(),
             mode: ExtractorMode::Raw,
             with_type_ids: false,
             with_attention_mask: true,
             num_threads,
             optimization_level,
-            execution_providers: execution_providers_override.map(str::to_string),
+            execution_providers: overrides.execution_providers.map(str::to_string),
         };
         let session = build_session(&model_path, &session_config)?;
 
@@ -82,7 +89,7 @@ impl Embedder {
         let with_type_ids = has_input(&session, "token_type_ids");
         let with_attention_mask = has_input(&session, "attention_mask");
         let output_tensor =
-            select_output_tensor(&session, output_tensor_override, &PREFERRED_EMBEDDING_OUTPUTS)?;
+            select_output_tensor(&session, overrides.output_tensor, &PREFERRED_EMBEDDING_OUTPUTS)?;
         let mode = infer_extraction_mode(&session, output_tensor.as_str())?;
         if matches!(mode, ExtractorMode::MeanPool) && !with_attention_mask {
             return Err(GteError::Inference(
@@ -92,16 +99,23 @@ impl Embedder {
 
         let config = ModelConfig {
             max_length,
+            padding_mode,
             output_tensor,
             mode,
             with_type_ids,
             with_attention_mask,
             num_threads,
             optimization_level,
-            execution_providers: execution_providers_override.map(str::to_string),
+            execution_providers: overrides.execution_providers.map(str::to_string),
         };
 
-        let tokenizer = Tokenizer::new(&tokenizer_path, config.max_length, config.with_type_ids)?;
+        let tokenizer = Tokenizer::new(
+            &tokenizer_path,
+            config.max_length,
+            config.with_type_ids,
+            config.padding_mode,
+            tokenizer_profile.fixed_padding_length,
+        )?;
 
         Ok(Self {
             tokenizer,

data/ext/gte/src/model_config.rs

@@ -5,9 +5,18 @@ pub enum ExtractorMode {
     Raw,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
+pub enum PaddingMode {
+    #[default]
+    Auto,
+    BatchLongest,
+    Fixed,
+}
+
 #[derive(Debug, Clone)]
 pub struct ModelConfig {
     pub max_length: usize,
+    pub padding_mode: PaddingMode,
     pub output_tensor: String,
     pub mode: ExtractorMode,
     pub with_type_ids: bool,
@@ -16,3 +25,12 @@ pub struct ModelConfig {
     pub optimization_level: u8,
     pub execution_providers: Option<String>,
 }
+
+#[derive(Debug, Clone, Copy, Default)]
+pub struct ModelLoadOverrides<'a> {
+    pub model_name: Option<&'a str>,
+    pub output_tensor: Option<&'a str>,
+    pub max_length: Option<usize>,
+    pub padding: Option<&'a str>,
+    pub execution_providers: Option<&'a str>,
+}

data/ext/gte/src/model_profile.rs

@@ -1,9 +1,19 @@
 use crate::error::{GteError, Result};
 use crate::model_config::ExtractorMode;
 use ort::session::Session;
+use serde_json::Value;
 use std::path::{Path, PathBuf};
 
 const SUPPORTED_INPUTS: [&str; 3] = ["input_ids", "attention_mask", "token_type_ids"];
+const DEFAULT_MAX_LENGTH: usize = 512;
+const MAX_SUPPORTED_LENGTH: usize = 8192;
+
+#[derive(Debug, Clone, Copy)]
+pub struct TokenizerProfile {
+    pub default_max_length: usize,
+    pub safe_max_length: usize,
+    pub fixed_padding_length: Option<usize>,
+}
 
 pub fn resolve_tokenizer_path(dir: &Path) -> Result<PathBuf> {
     let tokenizer_path = dir.join("tokenizer.json");
@@ -48,19 +58,78 @@ pub fn resolve_default_text_model(dir: &Path) -> Result<PathBuf> {
     )))
 }
 
-pub fn read_max_length(dir: &Path) -> usize {
-    (|| -> Option<usize> {
-        let contents = std::fs::read_to_string(dir.join("tokenizer_config.json")).ok()?;
-        let json: serde_json::Value = serde_json::from_str(&contents).ok()?;
-        let v = json.get("model_max_length")?;
-        let n = v.as_u64().or_else(|| {
-            v.as_f64()
-                .filter(|&f| f > 0.0 && f < 1e15)
-                .map(|f| f as u64)
-        })?;
-        Some((n as usize).min(8192))
-    })()
-    .unwrap_or(512)
+pub fn read_tokenizer_profile(dir: &Path) -> TokenizerProfile {
+    let tokenizer_config = read_json(dir.join("tokenizer_config.json"));
+    let tokenizer_json = read_json(dir.join("tokenizer.json"));
+
+    let fixed_padding_length = tokenizer_json
+        .as_ref()
+        .and_then(parse_fixed_padding_length_from_tokenizer_json);
+
+    let mut candidates = Vec::new();
+    if let Some(config) = tokenizer_config.as_ref() {
+        if let Some(v) = config.get("max_length").and_then(parse_positive_usize) {
+            candidates.push(v.min(MAX_SUPPORTED_LENGTH));
+        }
+        if let Some(v) = config.get("model_max_length").and_then(parse_positive_usize) {
+            candidates.push(v.min(MAX_SUPPORTED_LENGTH));
+        }
+    }
+
+    if let Some(tokenizer) = tokenizer_json.as_ref() {
+        if let Some(v) = tokenizer
+            .get("truncation")
+            .and_then(|truncation| truncation.get("max_length"))
+            .and_then(parse_positive_usize)
+        {
+            candidates.push(v.min(MAX_SUPPORTED_LENGTH));
+        }
+    }
+
+    if let Some(v) = fixed_padding_length {
+        candidates.push(v.min(MAX_SUPPORTED_LENGTH));
+    }
+
+    let default_max_length = candidates
+        .iter()
+        .copied()
+        .min()
+        .unwrap_or(DEFAULT_MAX_LENGTH)
+        .max(1);
+    let safe_max_length = fixed_padding_length.unwrap_or(default_max_length).max(1);
+
+    TokenizerProfile {
+        default_max_length,
+        safe_max_length,
+        fixed_padding_length,
+    }
+}
+
+fn read_json(path: PathBuf) -> Option<Value> {
+    let contents = std::fs::read_to_string(path).ok()?;
+    serde_json::from_str(&contents).ok()
+}
+
+fn parse_positive_usize(value: &Value) -> Option<usize> {
+    let raw = value
+        .as_u64()
+        .or_else(|| {
+            value
+                .as_f64()
+                .filter(|&v| v.is_finite() && v > 0.0)
+                .map(|v| v as u64)
+        })
+        .or_else(|| value.as_str().and_then(|s| s.parse::<u64>().ok()))?;
+    let parsed = usize::try_from(raw).ok()?;
+    (parsed > 0).then_some(parsed)
+}
+
+fn parse_fixed_padding_length_from_tokenizer_json(tokenizer_json: &Value) -> Option<usize> {
+    tokenizer_json
+        .get("padding")
+        .and_then(|padding| padding.get("strategy"))
+        .and_then(|strategy| strategy.get("Fixed"))
+        .and_then(parse_positive_usize)
 }
 
 pub fn validate_supported_text_inputs(session: &Session, api_label: &str) -> Result<()> {
@@ -177,3 +246,32 @@ pub fn infer_extraction_mode(session: &Session, output_tensor: &str) -> Result<E
         ))),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{parse_fixed_padding_length_from_tokenizer_json, parse_positive_usize};
+    use serde_json::json;
+
+    #[test]
+    fn parse_positive_usize_handles_integer_float_and_string() {
+        assert_eq!(parse_positive_usize(&json!(64)), Some(64));
+        assert_eq!(parse_positive_usize(&json!(64.0)), Some(64));
+        assert_eq!(parse_positive_usize(&json!("64")), Some(64));
+        assert_eq!(parse_positive_usize(&json!(0)), None);
+    }
+
+    #[test]
+    fn parse_fixed_padding_length_reads_fixed_padding_strategy() {
+        let tokenizer_json = json!({
+            "padding": {
+                "strategy": {
+                    "Fixed": 64
+                }
+            }
+        });
+        assert_eq!(
+            parse_fixed_padding_length_from_tokenizer_json(&tokenizer_json),
+            Some(64)
+        );
+    }
+}

data/ext/gte/src/reranker.rs

@@ -1,19 +1,20 @@
 use crate::error::{GteError, Result};
+use crate::model_config::{ModelLoadOverrides, PaddingMode};
 use crate::model_profile::{
-    has_input, read_max_length, resolve_default_text_model, resolve_named_model, resolve_tokenizer_path,
-    select_output_tensor, validate_supported_text_inputs,
+    has_input, read_tokenizer_profile, resolve_default_text_model, resolve_named_model,
+    resolve_tokenizer_path, select_output_tensor, validate_supported_text_inputs,
 };
 use crate::pipeline::{extract_output_tensor, InputTensors};
 use crate::postprocess::sigmoid_scores;
 use crate::session::build_session;
-use crate::tokenizer::Tokenizer;
-use ndarray::Array1;
+use crate::tokenizer::{parse_padding_mode_override, Tokenizer};
 use ort::session::Session;
 use std::path::Path;
 
 #[derive(Debug, Clone)]
 struct RerankerConfig {
     max_length: usize,
+    padding_mode: PaddingMode,
     output_tensor: String,
     with_type_ids: bool,
     with_attention_mask: bool,
@@ -30,54 +31,62 @@ impl Reranker {
         dir: P,
         num_threads: usize,
         optimization_level: u8,
-        model_name: Option<&str>,
-        output_tensor_override: Option<&str>,
-        max_length_override: Option<usize>,
-        execution_providers_override: Option<&str>,
+        overrides: ModelLoadOverrides<'_>,
     ) -> Result<Self> {
         let dir = dir.as_ref();
         let tokenizer_path = resolve_tokenizer_path(dir)?;
-        let model_path = match model_name.filter(|s| !s.is_empty()) {
+        let model_path = match overrides.model_name.filter(|s| !s.is_empty()) {
             Some(name) => resolve_named_model(dir, name)?,
             None => resolve_default_text_model(dir)?,
         };
 
-        let max_length = if let Some(override_value) = max_length_override {
+        let tokenizer_profile = read_tokenizer_profile(dir);
+        let max_length = if let Some(override_value) = overrides.max_length {
             if override_value == 0 {
                 return Err(GteError::Inference(
                     "max_length override must be greater than 0".to_string(),
                 ));
             }
-            override_value
+            override_value.min(tokenizer_profile.safe_max_length)
         } else {
-            read_max_length(dir)
+            tokenizer_profile.default_max_length
         };
+        let padding_mode =
+            parse_padding_mode_override(overrides.padding)?.unwrap_or(PaddingMode::Auto);
 
         let probe_config = crate::model_config::ModelConfig {
             max_length,
+            padding_mode,
             output_tensor: String::new(),
             mode: crate::model_config::ExtractorMode::Raw,
             with_type_ids: false,
             with_attention_mask: true,
             num_threads,
             optimization_level,
-            execution_providers: execution_providers_override.map(str::to_string),
+            execution_providers: overrides.execution_providers.map(str::to_string),
         };
         let session = build_session(&model_path, &probe_config)?;
 
         validate_supported_text_inputs(&session, "text reranking")?;
         let with_type_ids = has_input(&session, "token_type_ids");
         let with_attention_mask = has_input(&session, "attention_mask");
-        let output_tensor = select_output_tensor(&session, output_tensor_override, &["logits"])?;
+        let output_tensor = select_output_tensor(&session, overrides.output_tensor, &["logits"])?;
 
         let config = RerankerConfig {
             max_length,
+            padding_mode,
             output_tensor,
             with_type_ids,
             with_attention_mask,
         };
 
-        let tokenizer = Tokenizer::new(&tokenizer_path, config.max_length, config.with_type_ids)?;
+        let tokenizer = Tokenizer::new(
+            &tokenizer_path,
+            config.max_length,
+            config.with_type_ids,
+            config.padding_mode,
+            tokenizer_profile.fixed_padding_length,
+        )?;
 
         Ok(Self {
             tokenizer,
@@ -86,14 +95,27 @@ impl Reranker {
         })
     }
 
-    pub fn score_pairs(&self, pairs: &[(String, String)], apply_sigmoid: bool) -> Result<Array1<f32>> {
+    pub fn score_pairs(&self, pairs: &[(String, String)], apply_sigmoid: bool) -> Result<Vec<f32>> {
         let tokenized = self.tokenizer.tokenize_pairs(pairs)?;
-        let input_tensors = InputTensors::from_tokenized(&tokenized, self.config.with_attention_mask)?;
+        self.score_tokenized(&tokenized, apply_sigmoid)
+    }
+
+    pub fn score(&self, query: &str, candidates: &[String], apply_sigmoid: bool) -> Result<Vec<f32>> {
+        let tokenized = self.tokenizer.tokenize_query_candidates(query, candidates)?;
+        self.score_tokenized(&tokenized, apply_sigmoid)
+    }
+
+    fn score_tokenized(
+        &self,
+        tokenized: &crate::tokenizer::Tokenized,
+        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 array = extract_output_tensor(&outputs, self.config.output_tensor.as_str())?;
 
         let mut scores = match array.ndim() {
-            1 => array.into_dimensionality::<ndarray::Ix1>()?.into_owned(),
+            1 => array.into_dimensionality::<ndarray::Ix1>()?.to_vec(),
             2 => {
                 let shape = array.shape();
                 if shape[1] == 0 {
@@ -102,7 +124,7 @@ impl Reranker {
                         self.config.output_tensor, shape
                     )));
                 }
-                array.slice(ndarray::s![.., 0]).into_owned()
+                array.slice(ndarray::s![.., 0]).to_vec()
             }
             n => {
                 return Err(GteError::Inference(format!(
@@ -113,10 +135,9 @@ impl Reranker {
         };
 
         if apply_sigmoid {
-            sigmoid_scores(scores.view_mut());
+            sigmoid_scores(ndarray::ArrayViewMut1::from(scores.as_mut_slice()));
         }
 
         Ok(scores)
     }
-
 }

data/ext/gte/src/ruby_embedder.rs

@@ -2,6 +2,7 @@
 
 use crate::embedder::{normalize_l2, Embedder};
 use crate::error::GteError;
+use crate::model_config::ModelLoadOverrides;
 use crate::reranker::Reranker;
 use magnus::{function, method, prelude::*, wrap, Error, RArray, Ruby};
 use std::os::raw::c_void;
@@ -38,7 +39,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<Result<Vec<f32>, GteError>>,
 }
@@ -85,13 +87,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,
         };
@@ -135,8 +139,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(|| {
-        let scores = (*args.reranker).score_pairs(&*args.pairs, args.apply_sigmoid)?;
-        Ok(scores.to_vec())
+        (*args.reranker).score(&*args.query, &*args.candidates, args.apply_sigmoid)
     }));
     args.result = Some(match run_result {
         Ok(result) => result,
@@ -171,6 +174,7 @@ impl RbEmbedder {
         normalize: bool,
         output_tensor: String,
         max_length: usize,
+        padding: String,
         execution_providers: String,
     ) -> Result<Self, Error> {
         let name = if model_name.is_empty() {
@@ -193,14 +197,23 @@ impl RbEmbedder {
         } 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,
+            max_length: max_length_override,
+            padding: padding_override,
+            execution_providers: execution_providers_override,
+        };
         let embedder = Embedder::from_dir(
             &dir_path,
             num_threads,
             optimization_level,
-            name,
-            output_override,
-            max_length_override,
-            execution_providers_override,
+            overrides,
         )
         .map_err(magnus::Error::from)?;
         Ok(RbEmbedder {
@@ -231,6 +244,7 @@ impl RbReranker {
         sigmoid: bool,
         output_tensor: String,
         max_length: usize,
+        padding: String,
         execution_providers: String,
     ) -> Result<Self, Error> {
         let name = if model_name.is_empty() {
@@ -253,14 +267,23 @@ impl RbReranker {
         } 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,
+            max_length: max_length_override,
+            padding: padding_override,
+            execution_providers: execution_providers_override,
+        };
         let reranker = Reranker::from_dir(
             &dir_path,
             num_threads,
             optimization_level,
-            name,
-            output_override,
-            max_length_override,
-            execution_providers_override,
+            overrides,
         )
         .map_err(magnus::Error::from)?;
         Ok(RbReranker {
@@ -276,11 +299,7 @@ impl RbReranker {
         candidates: RArray,
     ) -> Result<RArray, Error> {
         let candidates: Vec<String> = candidates.to_vec()?;
-        let pairs: Vec<(String, String)> = candidates
-            .into_iter()
-            .map(|candidate| (query.clone(), candidate))
-            .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 {
@@ -376,12 +395,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, 8))?;
+    embedder_class.define_singleton_method("new", function!(RbEmbedder::rb_new, 9))?;
     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, 8))?;
+    reranker_class.define_singleton_method("new", function!(RbReranker::rb_new, 9))?;
     reranker_class.define_method("score", method!(RbReranker::rb_score, 2))?;
 
     let tensor_class = module.define_class("Tensor", ruby.class_object())?;

data/ext/gte/src/tokenizer.rs

@@ -1,4 +1,5 @@
 use crate::error::{GteError, Result};
+use crate::model_config::PaddingMode;
 use std::path::Path;
 use tokenizers::{PaddingParams, PaddingStrategy, TruncationParams};
 
@@ -20,6 +21,8 @@ impl Tokenizer {
         tokenizer_path: P,
         max_length: usize,
         with_type_ids: bool,
+        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()))?;
@@ -33,7 +36,7 @@ impl Tokenizer {
             .map_err(|e| GteError::Tokenizer(e.to_string()))?;
 
         let padding = PaddingParams {
-            strategy: PaddingStrategy::BatchLongest,
+            strategy: resolve_padding_strategy(padding_mode, max_length, fixed_padding_length),
             ..Default::default()
         };
         tokenizer.with_padding(Some(padding));
@@ -73,6 +76,56 @@ impl Tokenizer {
             .map_err(|e| GteError::Tokenizer(e.to_string()))?;
         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)
+    }
+}
+
+pub fn parse_padding_mode_override(value: Option<&str>) -> Result<Option<PaddingMode>> {
+    let Some(raw) = value.map(str::trim).filter(|v| !v.is_empty()) else {
+        return Ok(None);
+    };
+
+    let normalized = raw.to_ascii_lowercase().replace('-', "_");
+    let parsed = match normalized.as_str() {
+        "auto" => PaddingMode::Auto,
+        "batch_longest" | "batchlongest" => PaddingMode::BatchLongest,
+        "fixed" => PaddingMode::Fixed,
+        _ => {
+            return Err(GteError::Inference(format!(
+                "invalid padding mode '{}'; expected one of: auto, batch_longest, fixed",
+                raw
+            )))
+        }
+    };
+    Ok(Some(parsed))
+}
+
+fn resolve_padding_strategy(
+    padding_mode: PaddingMode,
+    max_length: usize,
+    fixed_padding_length: Option<usize>,
+) -> PaddingStrategy {
+    match padding_mode {
+        PaddingMode::BatchLongest => PaddingStrategy::BatchLongest,
+        PaddingMode::Fixed => PaddingStrategy::Fixed(max_length),
+        PaddingMode::Auto => {
+            if fixed_padding_length.is_some() {
+                PaddingStrategy::Fixed(max_length)
+            } else {
+                PaddingStrategy::BatchLongest
+            }
+        }
+    }
 }
 
 fn build_tokenized_single(
@@ -121,21 +174,17 @@ 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 {
-        input_ids.extend(encoding.get_ids().iter().map(|&value| i64::from(value)));
-        attn_masks.extend(
-            encoding
-                .get_attention_mask()
-                .iter()
-                .map(|&value| i64::from(value)),
-        );
+        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));
+        }
 
         if let Some(type_ids) = type_ids.as_mut() {
-            type_ids.extend(
-                encoding
-                    .get_type_ids()
-                    .iter()
-                    .map(|&value| i64::from(value)),
-            );
+            for &value in encoding.get_type_ids() {
+                type_ids.push(i64::from(value));
+            }
         }
     }
 
@@ -147,3 +196,39 @@ fn build_tokenized(encodings: &[tokenizers::Encoding], with_type_ids: bool) -> R
         type_ids,
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{parse_padding_mode_override, resolve_padding_strategy};
+    use crate::model_config::PaddingMode;
+    use tokenizers::PaddingStrategy;
+
+    #[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)
+        );
+    }
+
+    #[test]
+    fn parse_padding_mode_override_rejects_invalid_values() {
+        assert!(parse_padding_mode_override(Some("unknown")).is_err());
+    }
+
+    #[test]
+    fn resolve_padding_strategy_uses_fixed_for_auto_when_model_has_fixed_padding() {
+        match resolve_padding_strategy(PaddingMode::Auto, 64, Some(64)) {
+            PaddingStrategy::Fixed(64) => {}
+            other => panic!("expected Fixed(64), got {:?}", other),
+        }
+    }
+}

data/ext/gte/tests/inference_integration_test.rs

@@ -1,11 +1,12 @@
 use gte::embedder::Embedder;
+use gte::model_config::ModelLoadOverrides;
 
 #[test]
 #[ignore = "requires ext/gte/tests/fixtures/e5/tokenizer.json and model.onnx"]
 fn test_e5_single_embedding_shape() {
     const DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/e5");
 
-    let embedder = Embedder::from_dir(DIR, 0, 3, None, None, None, None)
+    let embedder = Embedder::from_dir(DIR, 0, 3, ModelLoadOverrides::default())
         .expect("embedder should initialize");
     let result = embedder
         .embed(vec!["query: Hello world".to_string()])
@@ -20,7 +21,7 @@ fn test_e5_single_embedding_shape() {
 fn test_clip_single_embedding_shape() {
     const DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/clip");
 
-    let embedder = Embedder::from_dir(DIR, 0, 3, None, None, None, None)
+    let embedder = Embedder::from_dir(DIR, 0, 3, ModelLoadOverrides::default())
         .expect("embedder should initialize");
     let result = embedder
         .embed(vec!["a photo of a cat".to_string()])
@@ -35,7 +36,7 @@ fn test_clip_single_embedding_shape() {
 fn test_e5_batch_embedding_shape() {
     const DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/e5");
 
-    let embedder = Embedder::from_dir(DIR, 0, 3, None, None, None, None)
+    let embedder = Embedder::from_dir(DIR, 0, 3, ModelLoadOverrides::default())
         .expect("embedder should initialize");
     let texts = vec![
         "query: first sentence".to_string(),
@@ -54,7 +55,7 @@ fn test_e5_batch_embedding_shape() {
 fn test_e5_long_input_truncation_no_error() {
     const DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/e5");
 
-    let embedder = Embedder::from_dir(DIR, 0, 3, None, None, None, None)
+    let embedder = Embedder::from_dir(DIR, 0, 3, ModelLoadOverrides::default())
         .expect("embedder should initialize");
     let very_long_text = "word ".repeat(1000);
     let result = embedder

data/ext/gte/tests/tokenizer_unit_test.rs

@@ -1,3 +1,4 @@
+use gte::model_config::PaddingMode;
 use gte::tokenizer::Tokenizer;
 
 #[test]
@@ -8,7 +9,8 @@ fn test_e5_tokenizer_output_shape() {
         "/tests/fixtures/e5/tokenizer.json"
     );
 
-    let tokenizer = Tokenizer::new(TOKENIZER, 512, true).expect("tokenizer should load");
+    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(),
@@ -33,7 +35,8 @@ fn test_e5_truncation_at_max_length() {
         "/tests/fixtures/e5/tokenizer.json"
     );
 
-    let tokenizer = Tokenizer::new(TOKENIZER, 16, false).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])

data/lib/gte/config.rb

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

data/lib/gte/embedder.rb

@@ -18,6 +18,7 @@ module GTE
           config.normalize,
           config.output_tensor.to_s,
           config.max_length || 0,
+          config.padding.to_s,
           config.execution_providers.to_s
         )
       end
@@ -33,6 +34,7 @@ module GTE
           normalize: true,
           output_tensor: nil,
           max_length: nil,
+          padding: nil,
           execution_providers: nil
         )
       end

data/lib/gte/reranker.rb

@@ -25,6 +25,7 @@ module GTE
           sigmoid: false,
           output_tensor: nil,
           max_length: nil,
+          padding: nil,
           execution_providers: nil
         )
       end
@@ -38,6 +39,7 @@ module GTE
           cfg.sigmoid,
           cfg.output_tensor.to_s,
           cfg.max_length || 0,
+          cfg.padding.to_s,
           cfg.execution_providers.to_s
         )
       end

data/lib/gte.rb

@@ -27,6 +27,7 @@ module GTE
         normalize: true,
         output_tensor: nil,
         max_length: nil,
+        padding: nil,
         execution_providers: nil
       )
 

metadata

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