gte 0.0.3 → 0.0.5

data/ext/gte/src/model_profile.rs

@@ -0,0 +1,179 @@
+use crate::error::{GteError, Result};
+use crate::model_config::ExtractorMode;
+use ort::session::Session;
+use std::path::{Path, PathBuf};
+
+const SUPPORTED_INPUTS: [&str; 3] = ["input_ids", "attention_mask", "token_type_ids"];
+
+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()
+        )));
+    }
+    Ok(tokenizer_path)
+}
+
+pub fn resolve_named_model(dir: &Path, name: &str) -> Result<PathBuf> {
+    let candidates = [dir.join("onnx").join(name), dir.join(name)];
+    for path in &candidates {
+        if path.exists() {
+            return Ok(path.clone());
+        }
+    }
+    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> {
+    let candidates = [
+        dir.join("onnx").join("text_model.onnx"),
+        dir.join("text_model.onnx"),
+        dir.join("onnx").join("model.onnx"),
+        dir.join("model.onnx"),
+    ];
+    for path in &candidates {
+        if path.exists() {
+            return Ok(path.clone());
+        }
+    }
+    Err(GteError::Inference(format!(
+        "no ONNX model found in {} (checked text_model.onnx and model.onnx)",
+        dir.display()
+    )))
+}
+
+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 validate_supported_text_inputs(session: &Session, api_label: &str) -> Result<()> {
+    let unsupported: Vec<String> = session
+        .inputs
+        .iter()
+        .filter(|i| !SUPPORTED_INPUTS.contains(&i.name.as_str()))
+        .map(|i| i.name.clone())
+        .collect();
+
+    if unsupported.is_empty() {
+        return Ok(());
+    }
+
+    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).",
+        );
+    } else {
+        message.push_str(". Supported inputs are: input_ids, attention_mask, token_type_ids.");
+    }
+    Err(GteError::Inference(message))
+}
+
+pub fn has_input(session: &Session, name: &str) -> bool {
+    session.inputs.iter().any(|input| input.name == name)
+}
+
+fn output_name_matches(name: &str, preferred: &str) -> bool {
+    let lower = name.to_ascii_lowercase();
+    lower == preferred || lower.ends_with(&format!("/{}", preferred))
+}
+
+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.as_str(), requested_name))
+        {
+            return Ok(output.name.clone());
+        }
+        let available = session
+            .outputs
+            .iter()
+            .map(|o| o.name.as_str())
+            .collect::<Vec<_>>()
+            .join(", ");
+        return Err(GteError::Inference(format!(
+            "requested output tensor '{}' not found in model outputs: {}",
+            requested_name, available
+        )));
+    }
+
+    for preferred in preferred_outputs {
+        if let Some(output) = session
+            .outputs
+            .iter()
+            .find(|o| output_name_matches(o.name.as_str(), preferred))
+        {
+            return Ok(output.name.clone());
+        }
+    }
+
+    session
+        .outputs
+        .first()
+        .map(|o| o.name.clone())
+        .ok_or_else(|| GteError::Inference("model has no outputs".into()))
+}
+
+fn output_basename(name: &str) -> &str {
+    name.rsplit('/').next().unwrap_or(name)
+}
+
+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 ndims = match &output.output_type {
+        ort::value::ValueType::Tensor { dimensions, .. } => dimensions.len(),
+        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),
+        (_, 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
+        ))),
+    }
+}

data/ext/gte/src/pipeline.rs

@@ -0,0 +1,60 @@
+use crate::error::{GteError, Result};
+use crate::tokenizer::Tokenized;
+use ndarray::ArrayView2;
+use ort::session::SessionInputValue;
+use ort::value::Value;
+
+pub struct InputTensors<'a> {
+    pub inputs: Vec<(&'static str, SessionInputValue<'a>)>,
+    pub attention_mask: ArrayView2<'a, i64>,
+}
+
+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 mut inputs = Vec::with_capacity(2 + usize::from(tokenized.type_ids.is_some()));
+        inputs.push((
+            "input_ids",
+            SessionInputValue::from(Value::from_array(input_ids_view)?),
+        ));
+
+        if with_attention_mask {
+            inputs.push((
+                "attention_mask",
+                SessionInputValue::from(Value::from_array(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(Value::from_array(type_ids_view)?),
+            ));
+        }
+
+        Ok(Self {
+            inputs,
+            attention_mask,
+        })
+    }
+}
+
+pub fn extract_output_tensor<'a>(
+    outputs: &'a ort::session::SessionOutputs<'a, 'a>,
+    output_name: &str,
+) -> Result<ndarray::CowArray<'a, f32, ndarray::IxDyn>> {
+    let tensor_value = outputs.get(output_name).ok_or_else(|| {
+        GteError::Inference(format!("output tensor '{}' not found in model outputs", output_name))
+    })?;
+    Ok(tensor_value.try_extract_tensor::<f32>()?.into())
+}

data/ext/gte/src/postprocess.rs

@@ -75,6 +75,12 @@ pub fn normalize_l2(mut embeddings: Array2<f32>) -> Array2<f32> {
     embeddings
 }
 
+pub fn sigmoid_scores(mut scores: ndarray::ArrayViewMut1<'_, f32>) {
+    scores.map_inplace(|value| {
+        *value = 1.0 / (1.0 + (-*value).exp());
+    });
+}
+
 fn mean_pool_contiguous(
     hidden: &[f32],
     attention_mask: &[i64],
@@ -87,10 +93,27 @@ fn mean_pool_contiguous(
         let mask_base = batch_index * seq;
         let hidden_base = batch_index * seq * dim;
         let output_row = &mut output[batch_index * dim..(batch_index + 1) * dim];
+        let mask_row = &attention_mask[mask_base..mask_base + seq];
+
+        if mask_row.iter().all(|&weight| weight == 1) {
+            for token_index in 0..seq {
+                let token_base = hidden_base + token_index * dim;
+                for dim_index in 0..dim {
+                    output_row[dim_index] += hidden[token_base + dim_index];
+                }
+            }
+
+            let inverse = (seq as f32).recip();
+            for value in output_row {
+                *value *= inverse;
+            }
+            continue;
+        }
+
         let mut weight_sum = 0.0f32;
 
-        for token_index in 0..seq {
-            let weight = attention_mask[mask_base + token_index];
+        for (token_index, &weight_raw) in mask_row.iter().enumerate() {
+            let weight = weight_raw;
             if weight <= 0 {
                 continue;
             }

data/ext/gte/src/reranker.rs

@@ -0,0 +1,120 @@
+use crate::error::{GteError, Result};
+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,
+};
+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;
+
+#[derive(Debug, Clone)]
+struct RerankerConfig {
+    max_length: usize,
+    output_tensor: String,
+    with_type_ids: bool,
+    with_attention_mask: bool,
+}
+
+pub struct Reranker {
+    tokenizer: Tokenizer,
+    session: Session,
+    config: RerankerConfig,
+}
+
+impl Reranker {
+    pub fn from_dir<P: AsRef<Path>>(
+        dir: P,
+        num_threads: usize,
+        optimization_level: u8,
+        model_name: Option<&str>,
+        output_tensor_override: Option<&str>,
+        max_length_override: Option<usize>,
+    ) -> 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()) {
+            Some(name) => resolve_named_model(dir, name)?,
+            None => resolve_default_text_model(dir)?,
+        };
+
+        let max_length = if let Some(override_value) = max_length_override {
+            if override_value == 0 {
+                return Err(GteError::Inference(
+                    "max_length override must be greater than 0".to_string(),
+                ));
+            }
+            override_value
+        } else {
+            read_max_length(dir)
+        };
+
+        let probe_config = crate::model_config::ModelConfig {
+            max_length,
+            output_tensor: String::new(),
+            mode: crate::model_config::ExtractorMode::Raw,
+            with_type_ids: false,
+            with_attention_mask: true,
+            num_threads,
+            optimization_level,
+        };
+        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 config = RerankerConfig {
+            max_length,
+            output_tensor,
+            with_type_ids,
+            with_attention_mask,
+        };
+
+        let tokenizer = Tokenizer::new(&tokenizer_path, config.max_length, config.with_type_ids)?;
+
+        Ok(Self {
+            tokenizer,
+            session,
+            config,
+        })
+    }
+
+    pub fn score_pairs(&self, pairs: &[(String, String)], apply_sigmoid: bool) -> Result<Array1<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)?;
+        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(),
+            2 => {
+                let shape = array.shape();
+                if shape[1] == 0 {
+                    return Err(GteError::Inference(format!(
+                        "reranker output '{}' has invalid shape {:?}",
+                        self.config.output_tensor, shape
+                    )));
+                }
+                array.slice(ndarray::s![.., 0]).into_owned()
+            }
+            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(scores.view_mut());
+        }
+
+        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::reranker::Reranker;
 use magnus::{function, method, prelude::*, wrap, Error, RArray, Ruby};
 use std::os::raw::c_void;
 use std::panic::{catch_unwind, AssertUnwindSafe};
@@ -10,6 +11,13 @@ use std::sync::Arc;
 #[wrap(class = "GTE::Embedder", free_immediately, size)]
 pub struct RbEmbedder {
     inner: Arc<Embedder>,
+    normalize: bool,
+}
+
+#[wrap(class = "GTE::Reranker", free_immediately, size)]
+pub struct RbReranker {
+    inner: Arc<Reranker>,
+    sigmoid: bool,
 }
 
 #[wrap(class = "GTE::Tensor", free_immediately, size)]
@@ -22,11 +30,21 @@ pub struct RbTensor {
 struct InferArgs {
     embedder: *const Embedder,
     texts: *const Vec<String>,
+    normalize: bool,
     result: Option<Result<ndarray::Array2<f32>, GteError>>,
 }
 
 unsafe impl Send for InferArgs {}
 
+struct ScoreArgs {
+    reranker: *const Reranker,
+    pairs: *const Vec<(String, String)>,
+    apply_sigmoid: bool,
+    result: Option<Result<Vec<f32>, GteError>>,
+}
+
+unsafe impl Send for ScoreArgs {}
+
 fn panic_payload_to_string(payload: Box<dyn std::any::Any + Send>) -> String {
     if let Some(msg) = payload.downcast_ref::<&str>() {
         (*msg).to_string()
@@ -37,11 +55,16 @@ fn panic_payload_to_string(payload: Box<dyn std::any::Any + Send>) -> String {
     }
 }
 
-fn infer_without_gvl(embedder: &Arc<Embedder>, texts: Vec<String>) -> Result<ndarray::Array2<f32>, Error> {
+fn infer_without_gvl(
+    embedder: &Arc<Embedder>,
+    normalize: bool,
+    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,
         };
         rb_sys::rb_thread_call_without_gvl(
@@ -60,12 +83,44 @@ fn infer_without_gvl(embedder: &Arc<Embedder>, texts: Vec<String>) -> Result<nda
     Ok(embeddings)
 }
 
+fn score_without_gvl(
+    reranker: &Arc<Reranker>,
+    pairs: Vec<(String, 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)>,
+            apply_sigmoid,
+            result: None,
+        };
+        rb_sys::rb_thread_call_without_gvl(
+            Some(run_score_without_gvl),
+            &mut args as *mut ScoreArgs as *mut c_void,
+            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(),
+            ))
+        })?;
+        result.map_err(magnus::Error::from)?
+    };
+    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)?;
-        Ok(normalize_l2(embeddings))
+        if args.normalize {
+            Ok(normalize_l2(embeddings))
+        } else {
+            Ok(embeddings)
+        }
     }));
     args.result = Some(match run_result {
         Ok(result) => result,
@@ -77,6 +132,22 @@ unsafe extern "C" fn run_without_gvl(ptr: *mut c_void) -> *mut c_void {
     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();
     let cols = embeddings.ncols();
@@ -97,31 +168,114 @@ impl RbEmbedder {
         num_threads: usize,
         optimization_level: u8,
         model_name: String,
+        normalize: bool,
+        output_tensor: String,
+        max_length: usize,
     ) -> Result<Self, Error> {
         let name = if model_name.is_empty() {
             None
         } else {
             Some(model_name.as_str())
         };
-        let embedder = Embedder::from_dir(&dir_path, num_threads, optimization_level, name)
-            .map_err(magnus::Error::from)?;
+        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 embedder = Embedder::from_dir(
+            &dir_path,
+            num_threads,
+            optimization_level,
+            name,
+            output_override,
+            max_length_override,
+        )
+        .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> {
         let texts: Vec<String> = texts.to_vec()?;
-        let embeddings = infer_without_gvl(&rb_self.inner, texts)?;
+        let embeddings = infer_without_gvl(&rb_self.inner, rb_self.normalize, texts)?;
         tensor_from_array(embeddings)
     }
 
     pub fn rb_embed_one(_ruby: &Ruby, rb_self: &Self, text: String) -> Result<RbTensor, Error> {
-        let embeddings = infer_without_gvl(&rb_self.inner, vec![text])?;
+        let embeddings = infer_without_gvl(&rb_self.inner, rb_self.normalize, vec![text])?;
         tensor_from_array(embeddings)
     }
 }
 
+impl RbReranker {
+    pub fn rb_new(
+        _ruby: &Ruby,
+        dir_path: String,
+        num_threads: usize,
+        optimization_level: u8,
+        model_name: String,
+        sigmoid: bool,
+        output_tensor: String,
+        max_length: 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 reranker = Reranker::from_dir(
+            &dir_path,
+            num_threads,
+            optimization_level,
+            name,
+            output_override,
+            max_length_override,
+        )
+        .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> {
+        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 out = ruby.ary_new_capa(scores.len());
+        for score in scores {
+            out.push(score)?;
+        }
+        Ok(out)
+    }
+}
+
 impl RbTensor {
     pub fn len(&self) -> usize {
         self.rows
@@ -208,10 +362,14 @@ 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, 4))?;
+    embedder_class.define_singleton_method("new", function!(RbEmbedder::rb_new, 7))?;
     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, 7))?;
+    reranker_class.define_method("score", method!(RbReranker::rb_score, 2))?;
+
     let tensor_class = module.define_class("Tensor", ruby.class_object())?;
     tensor_class.define_method("rows", method!(RbTensor::rows, 0))?;
     tensor_class.define_method("size", method!(RbTensor::len, 0))?;