gte 0.0.13 → 0.0.14

data/ext/gte/src/model_profile.rs

@@ -18,10 +18,7 @@ pub struct TokenizerProfile {
 pub fn resolve_tokenizer_path(dir: &Path) -> Result<PathBuf> {
     let tokenizer_path = dir.join("tokenizer.json");
     if !tokenizer_path.exists() {
-        return Err(GteError::Tokenizer(format!(
-            "tokenizer.json not found in {}",
-            dir.display()
-        )));
+        return Err(GteError::Tokenizer(format!("tokenizer.json not found in {}", dir.display())));
     }
     Ok(tokenizer_path)
 }
@@ -33,11 +30,7 @@ pub fn resolve_named_model(dir: &Path, name: &str) -> Result<PathBuf> {
             return Ok(path.clone());
         }
     }
-    Err(GteError::Inference(format!(
-        "model '{}' not found in {} (checked onnx/{0} and {0})",
-        name,
-        dir.display()
-    )))
+    Err(GteError::Inference(format!("model '{}' not found in {} (checked onnx/{0} and {0})", name, dir.display())))
 }
 
 pub fn resolve_default_text_model(dir: &Path) -> Result<PathBuf> {
@@ -62,9 +55,7 @@ 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 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() {
@@ -90,19 +81,10 @@ pub fn read_tokenizer_profile(dir: &Path) -> TokenizerProfile {
         candidates.push(v.min(MAX_SUPPORTED_LENGTH));
     }
 
-    let default_max_length = candidates
-        .iter()
-        .copied()
-        .min()
-        .unwrap_or(DEFAULT_MAX_LENGTH)
-        .max(1);
+    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,
-    }
+    TokenizerProfile { default_max_length, safe_max_length, fixed_padding_length }
 }
 
 fn read_json(path: PathBuf) -> Option<Value> {
@@ -113,12 +95,7 @@ fn read_json(path: PathBuf) -> Option<Value> {
 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_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)
@@ -133,7 +110,9 @@ fn parse_fixed_padding_length_from_tokenizer_json(tokenizer_json: &Value) -> Opt
 }
 
 pub fn validate_supported_text_inputs(session: &Session, api_label: &str) -> Result<()> {
-    let unsupported: Vec<String> = session.inputs().iter()
+    let unsupported: Vec<String> = session
+        .inputs()
+        .iter()
         .filter(|i| !SUPPORTED_INPUTS.contains(&i.name()))
         .map(|i| i.name().to_owned())
         .collect();
@@ -142,11 +121,7 @@ pub fn validate_supported_text_inputs(session: &Session, api_label: &str) -> Res
         return Ok(());
     }
 
-    let mut message = format!(
-        "unsupported model inputs for {} API: {}",
-        api_label,
-        unsupported.join(", ")
-    );
+    let mut message = format!("unsupported model inputs for {} API: {}", api_label, unsupported.join(", "));
     if unsupported.iter().any(|n| n == "pixel_values") {
         message.push_str(
             ". This looks like a multimodal graph. Provide a text-only export (for example onnx/text_model.onnx).",
@@ -163,40 +138,23 @@ pub fn has_input(session: &Session, name: &str) -> bool {
 
 fn output_name_matches(name: &str, preferred: &str) -> bool {
     let lower = name.to_ascii_lowercase();
-    lower == preferred || lower.ends_with(&format!("/{}", preferred))
+    lower == preferred || lower.ends_with(&format!("/{preferred}"))
 }
 
-pub fn select_output_tensor(
-    session: &Session,
-    requested: Option<&str>,
-    preferred_outputs: &[&str],
-) -> Result<String> {
+#[allow(clippy::redundant_closure_for_method_calls)]
+pub fn select_output_tensor(session: &Session, requested: Option<&str>, preferred_outputs: &[&str]) -> Result<String> {
     if let Some(requested_name) = requested.map(str::trim).filter(|name| !name.is_empty()) {
-        if let Some(output) = session
-            .outputs()
-            .iter()
-            .find(|o| output_name_matches(o.name(), requested_name))
-        {
+        if let Some(output) = session.outputs().iter().find(|o| output_name_matches(o.name(), requested_name)) {
             return Ok(output.name().to_owned());
         }
-        let available = session
-            .outputs()
-            .iter()
-            .map(|o| o.name())
-            .collect::<Vec<_>>()
-            .join(", ");
+        let available = session.outputs().iter().map(|o| o.name()).collect::<Vec<_>>().join(", ");
         return Err(GteError::Inference(format!(
-            "requested output tensor '{}' not found in model outputs: {}",
-            requested_name, available
+            "requested output tensor '{requested_name}' not found in model outputs: {available}"
         )));
     }
 
     for preferred in preferred_outputs {
-        if let Some(output) = session
-            .outputs()
-            .iter()
-            .find(|o| output_name_matches(o.name(), preferred))
-        {
+        if let Some(output) = session.outputs().iter().find(|o| output_name_matches(o.name(), preferred)) {
             return Ok(output.name().to_owned());
         }
     }
@@ -204,12 +162,9 @@ pub fn select_output_tensor(
     let outputs = session.outputs();
     let best = outputs
         .iter()
-        .find(|o| {
-            matches!(o.dtype(), ort::value::ValueType::Tensor { shape, .. } if shape.len() == 2)
-        })
+        .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()))
+    best.map(|o| o.name().to_owned()).ok_or_else(|| GteError::Inference("model has no outputs".into()))
 }
 
 fn output_basename(name: &str) -> &str {
@@ -217,34 +172,21 @@ fn output_basename(name: &str) -> &str {
 }
 
 pub fn infer_extraction_mode(session: &Session, output_tensor: &str) -> Result<ExtractorMode> {
-    let output = session
-        .outputs()
-        .iter()
-        .find(|o| o.name() == output_tensor)
-        .ok_or_else(|| {
-            GteError::Inference(format!(
-                "output tensor '{}' not found in model outputs",
-                output_tensor
-            ))
+    let output =
+        session.outputs().iter().find(|o| o.name() == output_tensor).ok_or_else(|| {
+            GteError::Inference(format!("output tensor '{output_tensor}' not found in model outputs"))
         })?;
 
     let ndims = match output.dtype() {
         ort::value::ValueType::Tensor { shape, .. } => shape.len(),
-        other => {
-            return Err(GteError::Inference(format!(
-                "output is not a tensor: {:?}",
-                other
-            )))
-        }
+        other => return Err(GteError::Inference(format!("output is not a tensor: {other:?}"))),
     };
 
     match (output_basename(output_tensor), ndims) {
-        ("last_hidden_state", 3) => Ok(ExtractorMode::MeanPool),
+        ("last_hidden_state" | _, 3) => Ok(ExtractorMode::MeanPool),
         (_, 2) => Ok(ExtractorMode::Raw),
-        (_, 3) => Ok(ExtractorMode::MeanPool),
         (_, n) => Err(GteError::Inference(format!(
-            "unexpected output tensor rank {} for '{}': expected 2 (Raw) or 3 (MeanPool)",
-            n, output_tensor
+            "unexpected output tensor rank {n} for '{output_tensor}': expected 2 (Raw) or 3 (MeanPool)"
         ))),
     }
 }
@@ -271,9 +213,6 @@ mod tests {
                 }
             }
         });
-        assert_eq!(
-            parse_fixed_padding_length_from_tokenizer_json(&tokenizer_json),
-            Some(64)
-        );
+        assert_eq!(parse_fixed_padding_length_from_tokenizer_json(&tokenizer_json), Some(64));
     }
 }

data/ext/gte/src/pipeline.rs

@@ -11,41 +11,25 @@ pub struct InputTensors<'a> {
 
 impl<'a> InputTensors<'a> {
     pub fn from_tokenized(tokenized: &'a Tokenized, with_attention_mask: bool) -> Result<Self> {
-        let input_ids_view: ArrayView2<'_, i64> = ArrayView2::from_shape(
-            (tokenized.rows, tokenized.cols),
-            tokenized.input_ids.as_slice(),
-        )?;
-        let attention_mask: ArrayView2<'_, i64> = ArrayView2::from_shape(
-            (tokenized.rows, tokenized.cols),
-            tokenized.attn_masks.as_slice(),
-        )?;
+        let input_ids_view: ArrayView2<'_, i64> =
+            ArrayView2::from_shape((tokenized.rows, tokenized.cols), tokenized.input_ids.as_slice())?;
+        let attention_mask: ArrayView2<'_, i64> =
+            ArrayView2::from_shape((tokenized.rows, tokenized.cols), tokenized.attn_masks.as_slice())?;
 
         let mut inputs = Vec::with_capacity(2 + usize::from(tokenized.type_ids.is_some()));
-        inputs.push((
-            "input_ids",
-            SessionInputValue::from(TensorRef::from_array_view(input_ids_view)?),
-        ));
+        inputs.push(("input_ids", SessionInputValue::from(TensorRef::from_array_view(input_ids_view)?)));
 
         if with_attention_mask {
-            inputs.push((
-                "attention_mask",
-                SessionInputValue::from(TensorRef::from_array_view(attention_mask)?),
-            ));
+            inputs.push(("attention_mask", SessionInputValue::from(TensorRef::from_array_view(attention_mask)?)));
         }
 
         if let Some(type_ids) = tokenized.type_ids.as_deref() {
             let type_ids_view: ArrayView2<'_, i64> =
                 ArrayView2::from_shape((tokenized.rows, tokenized.cols), type_ids)?;
-            inputs.push((
-                "token_type_ids",
-                SessionInputValue::from(TensorRef::from_array_view(type_ids_view)?),
-            ));
+            inputs.push(("token_type_ids", SessionInputValue::from(TensorRef::from_array_view(type_ids_view)?)));
         }
 
-        Ok(Self {
-            inputs,
-            attention_mask,
-        })
+        Ok(Self { inputs, attention_mask })
     }
 }
 
@@ -53,11 +37,8 @@ pub fn extract_output_tensor<'a>(
     outputs: &'a ort::session::SessionOutputs<'_>,
     output_name: &str,
 ) -> Result<ArrayViewD<'a, f32>> {
-    let tensor_value = outputs.get(output_name).ok_or_else(|| {
-        GteError::Inference(format!(
-            "output tensor '{}' not found in model outputs",
-            output_name
-        ))
-    })?;
+    let tensor_value = outputs
+        .get(output_name)
+        .ok_or_else(|| GteError::Inference(format!("output tensor '{output_name}' not found in model outputs")))?;
     Ok(tensor_value.try_extract_array::<f32>()?)
 }

data/ext/gte/src/postprocess.rs

@@ -1,10 +1,7 @@
 use crate::error::{GteError, Result};
 use ndarray::{Array2, ArrayView2, ArrayView3};
 
-pub fn mean_pool(
-    hidden_states: ArrayView3<'_, f32>,
-    attention_mask: ArrayView2<'_, i64>,
-) -> Result<Array2<f32>> {
+pub fn mean_pool(hidden_states: ArrayView3<'_, f32>, attention_mask: ArrayView2<'_, i64>) -> Result<Array2<f32>> {
     let (batch, seq, dim) = hidden_states.dim();
     if attention_mask.dim() != (batch, seq) {
         return Err(GteError::Inference(format!(
@@ -34,17 +31,14 @@ pub fn mean_pool(
 
             let weight = weight as f32;
             for dim_index in 0..dim {
-                pooled[[batch_index, dim_index]] +=
-                    hidden_states[[batch_index, token_index, dim_index]] * weight;
+                pooled[[batch_index, dim_index]] += hidden_states[[batch_index, token_index, dim_index]] * weight;
             }
             weight_sum += weight;
         }
 
         if weight_sum > 0.0 {
             let inverse = weight_sum.recip();
-            pooled
-                .row_mut(batch_index)
-                .map_inplace(|value| *value *= inverse);
+            pooled.row_mut(batch_index).map_inplace(|value| *value *= inverse);
         }
     }
 
@@ -89,6 +83,8 @@ fn mean_pool_contiguous(
     seq: usize,
     dim: usize,
 ) {
+    let seq_inverse = (seq as f32).recip();
+
     for batch_index in 0..batch {
         let mask_base = batch_index * seq;
         let hidden_base = batch_index * seq * dim;
@@ -103,9 +99,8 @@ fn mean_pool_contiguous(
                 }
             }
 
-            let inverse = (seq as f32).recip();
             for value in output_row {
-                *value *= inverse;
+                *value *= seq_inverse;
             }
             continue;
         }
@@ -113,12 +108,11 @@ fn mean_pool_contiguous(
         let mut weight_sum = 0.0f32;
 
         for (token_index, &weight_raw) in mask_row.iter().enumerate() {
-            let weight = weight_raw;
-            if weight <= 0 {
+            if weight_raw <= 0 {
                 continue;
             }
 
-            let weight = weight as f32;
+            let weight = weight_raw as f32;
             let token_base = hidden_base + token_index * dim;
             for dim_index in 0..dim {
                 output_row[dim_index] += hidden[token_base + dim_index] * weight;

data/ext/gte/src/reranker.rs

@@ -1,8 +1,8 @@
 use crate::error::{GteError, Result};
 use crate::model_config::{ModelLoadOverrides, PaddingMode};
 use crate::model_profile::{
-    has_input, read_tokenizer_profile, 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;
@@ -26,11 +26,7 @@ pub struct Reranker {
 }
 
 impl Reranker {
-    pub fn from_dir<P: AsRef<Path>>(
-        dir: P,
-        optimization_level: u8,
-        overrides: ModelLoadOverrides<'_>,
-    ) -> Result<Self> {
+    pub fn from_dir<P: AsRef<Path>>(dir: P, optimization_level: u8, overrides: ModelLoadOverrides<'_>) -> Result<Self> {
         let dir = dir.as_ref();
         let tokenizer_path = resolve_tokenizer_path(dir)?;
         let model_path: PathBuf = match overrides.model_name.filter(|s| !s.is_empty()) {
@@ -41,16 +37,13 @@ impl Reranker {
         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(),
-                ));
+                return Err(GteError::Inference("max_length override must be greater than 0".to_string()));
             }
             override_value.min(tokenizer_profile.safe_max_length)
         } else {
             tokenizer_profile.default_max_length
         };
-        let padding_mode =
-            parse_padding_mode_override(overrides.padding)?.unwrap_or(PaddingMode::Auto);
+        let padding_mode = parse_padding_mode_override(overrides.padding)?.unwrap_or(PaddingMode::Auto);
 
         let probe_config = crate::model_config::ModelConfig {
             max_length,
@@ -61,6 +54,8 @@ impl Reranker {
             with_attention_mask: true,
             optimization_level,
             execution_providers: overrides.execution_providers.map(str::to_string),
+            lowercase_input: false,
+            max_input_chars: None,
         };
         let session = build_session(&model_path, &probe_config)?;
 
@@ -69,13 +64,7 @@ impl Reranker {
         let with_attention_mask = has_input(&session, "attention_mask");
         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 config = RerankerConfig { max_length, padding_mode, output_tensor, with_type_ids, with_attention_mask };
 
         let tokenizer = Tokenizer::new(
             &tokenizer_path,
@@ -85,7 +74,19 @@ impl Reranker {
             tokenizer_profile.fixed_padding_length,
         )?;
 
-        let pool = SessionPool::new(session, model_path, probe_config);
+        let model_config = crate::model_config::ModelConfig {
+            max_length,
+            padding_mode,
+            output_tensor: config.output_tensor.clone(),
+            mode: crate::model_config::ExtractorMode::Raw,
+            with_type_ids: config.with_type_ids,
+            with_attention_mask: config.with_attention_mask,
+            optimization_level,
+            execution_providers: None,
+            lowercase_input: false,
+            max_input_chars: None,
+        };
+        let pool = SessionPool::new(session, &model_path, &model_config)?;
         Ok(Self { tokenizer, pool, config })
     }
 
@@ -99,40 +100,39 @@ impl Reranker {
         self.score_tokenized(&tokenized, apply_sigmoid)
     }
 
-    fn score_tokenized(
-        &self,
-        tokenized: &crate::tokenizer::Tokenized,
-        apply_sigmoid: bool,
-    ) -> Result<Vec<f32>> {
+    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 mut session = self.pool.acquire()?;
-        let outputs = session.run(input_tensors.inputs).map_err(|e| GteError::Ort(e.to_string()))?;
-        let array = extract_output_tensor(&outputs, self.config.output_tensor.as_str())?;
-
-        let mut scores = match array.ndim() {
-            1 => array.into_dimensionality::<ndarray::Ix1>()?.to_vec(),
-            2 => {
-                let shape = array.shape();
-                if shape[1] == 0 {
+        let output_name = self.config.output_tensor.clone();
+        let inputs = input_tensors.inputs;
+
+        self.pool.with_session(|session| {
+            let outputs = session.run(inputs).map_err(|e| GteError::Ort(e.to_string()))?;
+
+            let array = extract_output_tensor(&outputs, output_name.as_str())?;
+
+            let mut scores = match array.ndim() {
+                1 => array.into_dimensionality::<ndarray::Ix1>()?.to_vec(),
+                2 => {
+                    let shape = array.shape();
+                    if shape[1] == 0 {
+                        return Err(GteError::Inference(format!(
+                            "reranker output '{output_name}' has invalid shape {shape:?}"
+                        )));
+                    }
+                    array.slice(ndarray::s![.., 0]).to_vec()
+                }
+                n => {
                     return Err(GteError::Inference(format!(
-                        "reranker output '{}' has invalid shape {:?}",
-                        self.config.output_tensor, shape
-                    )));
+                        "reranker output '{output_name}' rank {n} is unsupported; expected rank 1 or 2"
+                    )))
                 }
-                array.slice(ndarray::s![.., 0]).to_vec()
-            }
-            n => {
-                return Err(GteError::Inference(format!(
-                    "reranker output '{}' rank {} is unsupported; expected rank 1 or 2",
-                    self.config.output_tensor, n
-                )))
-            }
-        };
+            };
 
-        if apply_sigmoid {
-            sigmoid_scores(ndarray::ArrayViewMut1::from(scores.as_mut_slice()));
-        }
+            if apply_sigmoid {
+                sigmoid_scores(ndarray::ArrayViewMut1::from(scores.as_mut_slice()));
+            }
 
-        Ok(scores)
+            Ok(scores)
+        })
     }
 }