gte 0.0.13 → 0.0.14

data/ext/gte/src/ruby_embedder.rs

@@ -1,6 +1,9 @@
 #![cfg(feature = "ruby-ffi")]
+#![allow(unsafe_code)]
+#![allow(unused_results)]
+#![allow(unused_qualifications)]
 
-use crate::embedder::{normalize_l2, Embedder};
+use crate::embedder::{normalize_l2, output_name_suggests_normalized, Embedder};
 use crate::error::GteError;
 use crate::model_config::ModelLoadOverrides;
 use crate::reranker::Reranker;
@@ -66,29 +69,30 @@ unsafe extern "C" fn run_embed_without_gvl(ptr: *mut c_void) -> *mut c_void {
     let run_result = catch_unwind(AssertUnwindSafe(|| {
         // Full embedding path (tokenization + inference) runs without the GVL.
         let embeddings = (*args.embedder).embed_ref(&*args.texts)?;
-        if args.normalize { Ok(normalize_l2(embeddings)) } else { Ok(embeddings) }
+        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),
-        ))),
+        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(&*args.query, &*args.candidates, args.apply_sigmoid)
-    }));
+    let run_result =
+        catch_unwind(AssertUnwindSafe(|| (*args.reranker).score(&*args.query, &*args.candidates, 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),
-        ))),
+        Err(payload) => {
+            Err(GteError::Inference(format!("panic during reranking: {}", panic_payload_to_string(payload),)))
+        }
     });
     std::ptr::null_mut()
 }
@@ -99,23 +103,18 @@ fn infer_without_gvl(
     texts: Vec<String>,
 ) -> Result<ndarray::Array2<f32>, Error> {
     let embeddings = unsafe {
-        let mut args = InferArgs {
-            embedder: Arc::as_ptr(embedder),
-            texts: &texts as *const Vec<String>,
-            normalize,
-            result: None,
-        };
+        let mut args =
+            InferArgs { embedder: Arc::as_ptr(embedder), texts: &texts as *const Vec<String>, normalize, result: None };
         rb_sys::rb_thread_call_without_gvl(
             Some(run_embed_without_gvl),
             &mut args as *mut InferArgs as *mut c_void,
             None,
             std::ptr::null_mut(),
         );
-        let result = args.result.take().ok_or_else(|| {
-            magnus::Error::from(GteError::Inference(
-                "inference did not return a result".to_string(),
-            ))
-        })?;
+        let result = args
+            .result
+            .take()
+            .ok_or_else(|| magnus::Error::from(GteError::Inference("inference did not return a result".to_string())))?;
         result.map_err(magnus::Error::from)?
     };
     Ok(embeddings)
@@ -141,11 +140,10 @@ fn score_without_gvl(
             None,
             std::ptr::null_mut(),
         );
-        let result = args.result.take().ok_or_else(|| {
-            magnus::Error::from(GteError::Inference(
-                "reranking did not return a result".to_string(),
-            ))
-        })?;
+        let result = args
+            .result
+            .take()
+            .ok_or_else(|| magnus::Error::from(GteError::Inference("reranking did not return a result".to_string())))?;
         result.map_err(magnus::Error::from)?
     };
     Ok(scores)
@@ -158,10 +156,7 @@ fn tensor_from_array(embeddings: ndarray::Array2<f32>) -> Result<RbTensor, Error
     let cols = embeddings.ncols();
     let (data, offset) = embeddings.into_raw_vec_and_offset();
     if let Some(off) = offset.filter(|&o| o != 0) {
-        return Err(magnus::Error::from(GteError::Inference(format!(
-            "unexpected non-zero tensor offset: {}",
-            off
-        ))));
+        return Err(magnus::Error::from(GteError::Inference(format!("unexpected non-zero tensor offset: {}", off))));
     }
     Ok(RbTensor { rows, cols, data })
 }
@@ -177,22 +172,28 @@ impl RbEmbedder {
         max_length: usize,
         padding: String,
         execution_providers: String,
+        lowercase_input: bool,
+        max_input_chars: usize,
     ) -> 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 execution_providers_override = if execution_providers.is_empty() { None } else { Some(execution_providers.as_str()) };
+        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 max_input_chars_override = if max_input_chars == 0 { None } else { Some(max_input_chars) };
         let overrides = ModelLoadOverrides {
             model_name: name,
             output_tensor: output_override,
             max_length: max_length_override,
             padding: padding_override,
             execution_providers: execution_providers_override,
+            lowercase_input: Some(lowercase_input),
+            max_input_chars: max_input_chars_override,
         };
-        let embedder = Embedder::from_dir(&dir_path, optimization_level, overrides)
-            .map_err(magnus::Error::from)?;
-        Ok(RbEmbedder { inner: Arc::new(embedder), normalize })
+        let embedder = Embedder::from_dir(&dir_path, optimization_level, overrides).map_err(magnus::Error::from)?;
+        let skip_normalize = normalize && output_name_suggests_normalized(&embedder.config.output_tensor);
+        Ok(RbEmbedder { inner: Arc::new(embedder), normalize: normalize && !skip_normalize })
     }
 
     pub fn rb_embed(_ruby: &Ruby, rb_self: &Self, texts: RArray) -> Result<RbTensor, Error> {
@@ -218,11 +219,14 @@ impl RbReranker {
         max_length: usize,
         padding: String,
         execution_providers: String,
+        _lowercase_input: bool,
+        _max_input_chars: usize,
     ) -> 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 execution_providers_override = if execution_providers.is_empty() { None } else { Some(execution_providers.as_str()) };
+        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,
@@ -230,18 +234,13 @@ impl RbReranker {
             max_length: max_length_override,
             padding: padding_override,
             execution_providers: execution_providers_override,
+            ..ModelLoadOverrides::default()
         };
-        let reranker = Reranker::from_dir(&dir_path, optimization_level, overrides)
-            .map_err(magnus::Error::from)?;
+        let reranker = Reranker::from_dir(&dir_path, optimization_level, overrides).map_err(magnus::Error::from)?;
         Ok(RbReranker { inner: Arc::new(reranker), sigmoid })
     }
 
-    pub fn rb_score(
-        ruby: &Ruby,
-        rb_self: &Self,
-        query: String,
-        candidates: RArray,
-    ) -> Result<RArray, Error> {
+    pub fn rb_score(ruby: &Ruby, rb_self: &Self, query: String, candidates: RArray) -> Result<RArray, Error> {
         let candidates: Vec<String> = candidates.to_vec()?;
         let scores = score_without_gvl(&rb_self.inner, query, candidates, rb_self.sigmoid)?;
         let out = ruby.ary_new_capa(scores.len());
@@ -296,11 +295,7 @@ impl RbTensor {
         Self::row_binary_f32(ruby, rb_self, 0)
     }
 
-    pub fn row_binary_f32(
-        ruby: &Ruby,
-        rb_self: &Self,
-        index: usize,
-    ) -> Result<magnus::RString, Error> {
+    pub fn row_binary_f32(ruby: &Ruby, rb_self: &Self, index: usize) -> Result<magnus::RString, Error> {
         if index >= rb_self.rows {
             return Err(magnus::Error::from(GteError::Inference(format!(
                 "row index {} out of bounds for {} rows",
@@ -340,12 +335,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, 10))?;
     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, 10))?;
     reranker_class.define_method("score", method!(RbReranker::rb_score, 2))?;
 
     let tensor_class = module.define_class("Tensor", ruby.class_object())?;