gte 0.0.6 → 0.0.8

data/ext/gte/src/reranker.rs

@@ -1,19 +1,19 @@
 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 ort::session::Session;
-use std::path::Path;
+use crate::session::{build_session, SessionPool};
+use crate::tokenizer::{parse_padding_mode_override, Tokenizer};
+use std::path::{Path, PathBuf};
 
 #[derive(Debug, Clone)]
 struct RerankerConfig {
     max_length: usize,
+    padding_mode: PaddingMode,
     output_tensor: String,
     with_type_ids: bool,
     with_attention_mask: bool,
@@ -21,7 +21,7 @@ struct RerankerConfig {
 
 pub struct Reranker {
     tokenizer: Tokenizer,
-    session: Session,
+    pool: SessionPool,
     config: RerankerConfig,
 }
 
@@ -30,70 +30,89 @@ 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: PathBuf = 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,
-            session,
-            config,
-        })
+        let pool = SessionPool::new(session, model_path, probe_config);
+        Ok(Self { tokenizer, pool, config })
     }
 
-    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)?;
-        let outputs = self.session.run(input_tensors.inputs)?;
+        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 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>()?.into_owned(),
+            1 => array.into_dimensionality::<ndarray::Ix1>()?.to_vec(),
             2 => {
                 let shape = array.shape();
                 if shape[1] == 0 {
@@ -102,7 +121,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 +132,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;
@@ -27,11 +28,15 @@ pub struct RbTensor {
     data: Vec<f32>,
 }
 
+// ---------------------------------------------------------------------------
+// GVL-release helpers
+// ---------------------------------------------------------------------------
+
 struct InferArgs {
     embedder: *const Embedder,
     texts: *const Vec<String>,
     normalize: bool,
-    result: Option<Result<ndarray::Array2<f32>, GteError>>,
+    result: Option<crate::error::Result<ndarray::Array2<f32>>>,
 }
 
 unsafe impl Send for InferArgs {}
@@ -40,7 +45,7 @@ struct ScoreArgs {
     reranker: *const Reranker,
     pairs: *const Vec<(String, String)>,
     apply_sigmoid: bool,
-    result: Option<Result<Vec<f32>, GteError>>,
+    result: Option<crate::error::Result<Vec<f32>>>,
 }
 
 unsafe impl Send for ScoreArgs {}
@@ -55,6 +60,38 @@ fn panic_payload_to_string(payload: Box<dyn std::any::Any + Send>) -> String {
     }
 }
 
+unsafe extern "C" fn run_without_gvl(ptr: *mut c_void) -> *mut c_void {
+    let args = &mut *(ptr as *mut InferArgs);
+    let run_result = catch_unwind(AssertUnwindSafe(|| {
+        let tokenized = (*args.embedder).tokenize(&*args.texts)?;
+        let embeddings = (*args.embedder).run(&tokenized)?;
+        if args.normalize { Ok(normalize_l2(embeddings)) } else { Ok(embeddings) }
+    }));
+    args.result = Some(match run_result {
+        Ok(result) => result,
+        Err(payload) => Err(GteError::Inference(format!(
+            "panic during inference: {}",
+            panic_payload_to_string(payload),
+        ))),
+    });
+    std::ptr::null_mut()
+}
+
+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(|| {
+        (*args.reranker).score_pairs(&*args.pairs, args.apply_sigmoid)
+    }));
+    args.result = Some(match run_result {
+        Ok(result) => result,
+        Err(payload) => Err(GteError::Inference(format!(
+            "panic during reranking: {}",
+            panic_payload_to_string(payload),
+        ))),
+    });
+    std::ptr::null_mut()
+}
+
 fn infer_without_gvl(
     embedder: &Arc<Embedder>,
     normalize: bool,
@@ -111,42 +148,7 @@ fn score_without_gvl(
     Ok(scores)
 }
 
-unsafe extern "C" fn run_without_gvl(ptr: *mut c_void) -> *mut c_void {
-    let args = &mut *(ptr as *mut InferArgs);
-    let run_result = catch_unwind(AssertUnwindSafe(|| {
-        let tokenized = (*args.embedder).tokenize(&*args.texts)?;
-        let embeddings = (*args.embedder).run(&tokenized)?;
-        if args.normalize {
-            Ok(normalize_l2(embeddings))
-        } else {
-            Ok(embeddings)
-        }
-    }));
-    args.result = Some(match run_result {
-        Ok(result) => result,
-        Err(payload) => Err(GteError::Inference(format!(
-            "panic during inference: {}",
-            panic_payload_to_string(payload),
-        ))),
-    });
-    std::ptr::null_mut()
-}
-
-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.result = Some(match run_result {
-        Ok(result) => result,
-        Err(payload) => Err(GteError::Inference(format!(
-            "panic during reranking: {}",
-            panic_payload_to_string(payload),
-        ))),
-    });
-    std::ptr::null_mut()
-}
+// ---------------------------------------------------------------------------
 
 fn tensor_from_array(embeddings: ndarray::Array2<f32>) -> Result<RbTensor, Error> {
     let rows = embeddings.nrows();
@@ -171,42 +173,24 @@ 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() {
-            None
-        } else {
-            Some(model_name.as_str())
-        };
-        let output_override = if output_tensor.is_empty() {
-            None
-        } else {
-            Some(output_tensor.as_str())
-        };
-        let max_length_override = if max_length == 0 {
-            None
-        } else {
-            Some(max_length)
+        let name = if model_name.is_empty() { None } else { Some(model_name.as_str()) };
+        let output_override = if output_tensor.is_empty() { None } else { Some(output_tensor.as_str()) };
+        let max_length_override = if max_length == 0 { None } else { Some(max_length) };
+        let execution_providers_override = if execution_providers.is_empty() { None } 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 execution_providers_override = if execution_providers.is_empty() {
-            None
-        } else {
-            Some(execution_providers.as_str())
-        };
-        let embedder = Embedder::from_dir(
-            &dir_path,
-            num_threads,
-            optimization_level,
-            name,
-            output_override,
-            max_length_override,
-            execution_providers_override,
-        )
-        .map_err(magnus::Error::from)?;
-        Ok(RbEmbedder {
-            inner: Arc::new(embedder),
-            normalize,
-        })
+        let embedder = Embedder::from_dir(&dir_path, num_threads, optimization_level, overrides)
+            .map_err(magnus::Error::from)?;
+        Ok(RbEmbedder { inner: Arc::new(embedder), normalize })
     }
 
     pub fn rb_embed(_ruby: &Ruby, rb_self: &Self, texts: RArray) -> Result<RbTensor, Error> {
@@ -231,42 +215,24 @@ 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() {
-            None
-        } else {
-            Some(model_name.as_str())
-        };
-        let output_override = if output_tensor.is_empty() {
-            None
-        } else {
-            Some(output_tensor.as_str())
-        };
-        let max_length_override = if max_length == 0 {
-            None
-        } else {
-            Some(max_length)
+        let name = if model_name.is_empty() { None } else { Some(model_name.as_str()) };
+        let output_override = if output_tensor.is_empty() { None } else { Some(output_tensor.as_str()) };
+        let max_length_override = if max_length == 0 { None } else { Some(max_length) };
+        let execution_providers_override = if execution_providers.is_empty() { None } 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 execution_providers_override = if execution_providers.is_empty() {
-            None
-        } else {
-            Some(execution_providers.as_str())
-        };
-        let reranker = Reranker::from_dir(
-            &dir_path,
-            num_threads,
-            optimization_level,
-            name,
-            output_override,
-            max_length_override,
-            execution_providers_override,
-        )
-        .map_err(magnus::Error::from)?;
-        Ok(RbReranker {
-            inner: Arc::new(reranker),
-            sigmoid,
-        })
+        let reranker = Reranker::from_dir(&dir_path, num_threads, optimization_level, overrides)
+            .map_err(magnus::Error::from)?;
+        Ok(RbReranker { inner: Arc::new(reranker), sigmoid })
     }
 
     pub fn rb_score(
@@ -276,12 +242,8 @@ 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 pairs: Vec<(String, String)> = candidates.into_iter().map(|c| (query.clone(), c)).collect();
         let scores = score_without_gvl(&rb_self.inner, pairs, rb_self.sigmoid)?;
-
         let out = ruby.ary_new_capa(scores.len());
         for score in scores {
             out.push(score)?;
@@ -317,7 +279,6 @@ impl RbTensor {
                 index, rb_self.rows
             ))));
         }
-
         let start = index * rb_self.cols;
         let end = start + rb_self.cols;
         let out = ruby.ary_new_capa(rb_self.cols);
@@ -342,7 +303,6 @@ impl RbTensor {
                 index, rb_self.rows
             ))));
         }
-
         let start = index * rb_self.cols;
         let end = start + rb_self.cols;
         let bytes = unsafe {
@@ -376,12 +336,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())?;