gte 0.0.3 → 0.0.5

data/ext/gte/src/session.rs

@@ -1,14 +1,13 @@
 use crate::error::{GteError, Result};
 use crate::model_config::{ExtractorMode, ModelConfig};
+use crate::pipeline::{extract_output_tensor, InputTensors};
 use crate::postprocess::mean_pool;
 use crate::tokenizer::Tokenized;
-use ndarray::{Array2, ArrayView2, Ix2};
+use ndarray::{Array2, Ix2};
 use ort::execution_providers::{
     CoreMLExecutionProvider, ExecutionProviderDispatch, XNNPACKExecutionProvider,
 };
 use ort::session::Session;
-use ort::session::SessionInputValue;
-use ort::value::Value;
 use std::path::Path;
 
 pub fn build_session<P: AsRef<Path>>(model_path: P, config: &ModelConfig) -> Result<Session> {
@@ -43,7 +42,9 @@ fn preferred_execution_providers() -> Vec<ExecutionProviderDispatch> {
     let mut providers = Vec::new();
     for provider in order.split(',').map(str::trim).filter(|p| !p.is_empty()) {
         match provider {
-            "xnnpack" => providers.push(XNNPACKExecutionProvider::default().build().fail_silently()),
+            "xnnpack" => {
+                providers.push(XNNPACKExecutionProvider::default().build().fail_silently())
+            }
             "coreml" => providers.push(CoreMLExecutionProvider::default().build().fail_silently()),
             "none" => {}
             _ => {}
@@ -57,40 +58,9 @@ pub fn run_session(
     tokenized: &Tokenized,
     config: &ModelConfig,
 ) -> Result<Array2<f32>> {
-    let input_ids_view: ArrayView2<'_, i64> =
-        ArrayView2::from_shape((tokenized.rows, tokenized.cols), tokenized.input_ids.as_slice())?;
-    let attn_masks_view: 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 config.with_attention_mask {
-        inputs.push((
-            "attention_mask",
-            SessionInputValue::from(Value::from_array(attn_masks_view)?),
-        ));
-    }
-    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)?),
-        ));
-    }
-
-    let outputs = session.run(inputs)?;
-    let tensor_value = outputs.get(config.output_tensor.as_str()).ok_or_else(|| {
-        GteError::Inference(format!(
-            "output tensor '{}' not found in model outputs",
-            &config.output_tensor
-        ))
-    })?;
-
-    let array = tensor_value.try_extract_tensor::<f32>()?;
+    let input_tensors = InputTensors::from_tokenized(tokenized, config.with_attention_mask)?;
+    let outputs = session.run(input_tensors.inputs)?;
+    let array = extract_output_tensor(&outputs, config.output_tensor.as_str())?;
 
     match config.mode {
         ExtractorMode::Token(idx) => {
@@ -111,7 +81,7 @@ pub fn run_session(
                     ndim
                 ))
             })?;
-            mean_pool(hidden_states, attn_masks_view)
+            mean_pool(hidden_states.view(), input_tensors.attention_mask)
         }
         ExtractorMode::Raw => Ok(array.into_dimensionality::<Ix2>()?.into_owned()),
     }

data/ext/gte/src/tokenizer.rs

@@ -61,12 +61,31 @@ impl Tokenizer {
 
         build_tokenized(&encodings, self.with_type_ids)
     }
+
+    pub fn tokenize_pairs(&self, pairs: &[(String, String)]) -> Result<Tokenized> {
+        let encode_inputs: Vec<tokenizers::EncodeInput<'_>> = pairs
+            .iter()
+            .map(|(left, right)| (left.as_str(), right.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)
+    }
 }
 
-fn build_tokenized_single(encoding: &tokenizers::Encoding, with_type_ids: bool) -> Result<Tokenized> {
+fn build_tokenized_single(
+    encoding: &tokenizers::Encoding,
+    with_type_ids: bool,
+) -> Result<Tokenized> {
     let cols = encoding.len();
 
-    let input_ids: Vec<i64> = encoding.get_ids().iter().map(|&value| i64::from(value)).collect();
+    let input_ids: Vec<i64> = encoding
+        .get_ids()
+        .iter()
+        .map(|&value| i64::from(value))
+        .collect();
     let attn_masks: Vec<i64> = encoding
         .get_attention_mask()
         .iter()

data/ext/gte/tests/inference_integration_test.rs

@@ -5,7 +5,8 @@ use gte::embedder::Embedder;
 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).expect("embedder should initialize");
+    let embedder =
+        Embedder::from_dir(DIR, 0, 3, None, None, None).expect("embedder should initialize");
     let result = embedder
         .embed(vec!["query: Hello world".to_string()])
         .expect("embed should succeed");
@@ -19,7 +20,8 @@ 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).expect("embedder should initialize");
+    let embedder =
+        Embedder::from_dir(DIR, 0, 3, None, None, None).expect("embedder should initialize");
     let result = embedder
         .embed(vec!["a photo of a cat".to_string()])
         .expect("embed should succeed");
@@ -33,7 +35,8 @@ 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).expect("embedder should initialize");
+    let embedder =
+        Embedder::from_dir(DIR, 0, 3, None, None, None).expect("embedder should initialize");
     let texts = vec![
         "query: first sentence".to_string(),
         "query: second sentence".to_string(),
@@ -51,7 +54,8 @@ 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).expect("embedder should initialize");
+    let embedder =
+        Embedder::from_dir(DIR, 0, 3, None, None, None).expect("embedder should initialize");
     let very_long_text = "word ".repeat(1000);
     let result = embedder
         .embed(vec![very_long_text])

data/ext/gte/tests/postprocess_unit_test.rs

@@ -0,0 +1,17 @@
+use gte::postprocess::sigmoid_scores;
+use ndarray::array;
+
+#[test]
+fn test_sigmoid_scores_monotonic_and_bounded() {
+    let mut scores = array![-10.0f32, -1.0, 0.0, 1.0, 10.0];
+    sigmoid_scores(scores.view_mut());
+
+    assert!(scores[0] < scores[1]);
+    assert!(scores[1] < scores[2]);
+    assert!(scores[2] < scores[3]);
+    assert!(scores[3] < scores[4]);
+
+    for score in scores.iter() {
+        assert!((*score >= 0.0) && (*score <= 1.0));
+    }
+}

data/ext/gte/tests/tokenizer_unit_test.rs

@@ -40,5 +40,8 @@ fn test_e5_truncation_at_max_length() {
         .expect("tokenize should not error on long input");
 
     assert_eq!(tokenized.rows, 1);
-    assert_eq!(tokenized.cols, 16, "sequence length should be truncated to max_length");
+    assert_eq!(
+        tokenized.cols, 16,
+        "sequence length should be truncated to max_length"
+    );
 }

data/lib/gte/config.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module GTE
+  module Config
+    Text = Data.define(
+      :model_dir, :threads, :optimization_level,
+      :model_name, :normalize, :output_tensor, :max_length
+    )
+
+    Reranker = Data.define(
+      :model_dir, :threads, :optimization_level,
+      :model_name, :sigmoid, :output_tensor, :max_length
+    )
+  end
+end

data/lib/gte/model.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module GTE
+  class Model
+    attr_reader :config
+
+    def initialize(config)
+      raise ArgumentError, 'config must be a GTE::Config::Text' unless config.is_a?(Config::Text)
+
+      @config = config
+      @embedder = GTE::Embedder.new(
+        config.model_dir,
+        config.threads,
+        config.optimization_level,
+        config.model_name.to_s,
+        config.normalize,
+        config.output_tensor.to_s,
+        config.max_length || 0
+      )
+    end
+
+    def embed(texts)
+      return @embedder.embed_one(texts) if texts.is_a?(String)
+
+      @embedder.embed(Array(texts))
+    end
+
+    def [](input)
+      case input
+      when String then embed(input).row(0)
+      when Array then embed(input)
+      end
+    end
+  end
+end

data/lib/gte/reranker.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+module GTE
+  class Reranker
+    class << self
+      def config(model_dir)
+        cfg = default_config(model_dir)
+
+        if block_given?
+          yielded = yield(cfg)
+          cfg = yielded if yielded.is_a?(Config::Reranker)
+        end
+
+        build(cfg)
+      end
+
+      private
+
+      def default_config(model_dir)
+        Config::Reranker.new(
+          model_dir: File.expand_path(model_dir),
+          threads: 3,
+          optimization_level: 3,
+          model_name: nil,
+          sigmoid: false,
+          output_tensor: nil,
+          max_length: nil
+        )
+      end
+
+      def build(cfg)
+        new(
+          cfg.model_dir,
+          cfg.threads,
+          cfg.optimization_level,
+          cfg.model_name.to_s,
+          cfg.sigmoid,
+          cfg.output_tensor.to_s,
+          cfg.max_length || 0
+        )
+      end
+    end
+
+    def rerank(query:, candidates:)
+      rows = Array(candidates).map(&:to_s)
+      scores = score(query.to_s, rows)
+
+      rows
+        .each_with_index
+        .map { |text, idx| { index: idx, score: scores[idx], text: text } }
+        .sort_by { |row| -row[:score] }
+    end
+  end
+end

data/lib/gte/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module GTE
+  VERSION = File.read(File.expand_path('../../VERSION', __dir__)).strip
+end

data/lib/gte.rb

@@ -1,36 +1,44 @@
 # frozen_string_literal: true
 
+require 'gte/version'
+
 begin
   require "gte/#{RUBY_VERSION.to_f}/gte"
 rescue LoadError
   require 'gte/gte'
 end
 
+require 'gte/config'
+require 'gte/model'
+require 'gte/reranker'
+
 module GTE
-  VERSION = File.read(File.expand_path('../VERSION', __dir__)).strip
+  @model_cache_mutex = Mutex.new
+  @model_cache = {}
 
-  class Model
-    def initialize(dir, num_threads: 0, optimization_level: 3, model_name: nil)
-      @embedder = GTE::Embedder.new(dir, num_threads, optimization_level, model_name.to_s)
-    end
+  class << self
+    def config(model_dir)
+      cfg = Config::Text.new(
+        model_dir: File.expand_path(model_dir),
+        threads: 3,
+        optimization_level: 3,
+        model_name: nil,
+        normalize: true,
+        output_tensor: nil,
+        max_length: nil
+      )
 
-    def embed(texts)
-      if texts.is_a?(String)
-        @embedder.embed_one(texts)
-      else
-        @embedder.embed(Array(texts))
-      end
-    end
+      cfg = yield(cfg) if block_given?
 
-    def [](input)
-      case input
-      when String then embed(input).row(0)
-      when Array  then embed(input)
+      @model_cache_mutex.synchronize do
+        @model_cache[cache_key(cfg)] ||= Model.new(cfg)
       end
     end
-  end
 
-  def self.new(dir, num_threads: 0, optimization_level: 3, model_name: nil)
-    Model.new(dir, num_threads: num_threads, optimization_level: optimization_level, model_name: model_name)
+    private
+
+    def cache_key(cfg)
+      cfg.to_h
+    end
   end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: gte
 version: !ruby/object:Gem::Version
-  version: 0.0.3
+  version: 0.0.5
 platform: ruby
 authors:
 - elcuervo
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2026-04-10 00:00:00.000000000 Z
+date: 2026-04-15 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rake
@@ -101,14 +101,22 @@ files:
 - ext/gte/src/error.rs
 - ext/gte/src/lib.rs
 - ext/gte/src/model_config.rs
+- ext/gte/src/model_profile.rs
+- ext/gte/src/pipeline.rs
 - ext/gte/src/postprocess.rs
+- ext/gte/src/reranker.rs
 - ext/gte/src/ruby_embedder.rs
 - ext/gte/src/session.rs
 - ext/gte/src/tokenizer.rs
 - ext/gte/tests/embedder_unit_test.rs
 - ext/gte/tests/inference_integration_test.rs
+- ext/gte/tests/postprocess_unit_test.rs
 - ext/gte/tests/tokenizer_unit_test.rs
 - lib/gte.rb
+- lib/gte/config.rb
+- lib/gte/model.rb
+- lib/gte/reranker.rb
+- lib/gte/version.rb
 homepage: https://github.com/elcuervo/gte
 licenses:
 - MIT