gte 0.0.1 → 0.0.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 566cb32a193255c0cf3d087a1a907cdbf0a96292ee58d4676fd18745d55ec1b2
-  data.tar.gz: 56479bb218282bf189ace46852a42766d4c2211230be635b6a79414cf3eb82c8
+  metadata.gz: 46dc6316dd04fd0a15bd0f0be71ecb1dca4f177612da7e2f7992a3030aa51e88
+  data.tar.gz: d05a7dc0230a9b2c670cd3001afad0724ba8543ec9de16f82d5d1abb4df73853
 SHA512:
-  metadata.gz: f873b83b16e4cf2685f26b84d26f4c3b8abd90a9d32d5e60046d1aad577d409704838d3da1329f741535fef1f1a90f1edeeeb6d4fbab559569117869ae42677e
-  data.tar.gz: d5d5b49b8f51cbf3222409b941fa39b3f00074bf34aafdfaf74d4a3fd37ab99ec8a90bba9be641efb42176b182fe191925042ae09d324c851324b27bd62031ce
+  metadata.gz: acdc81039070f6307548ea7438aeea69df4d9d5d6884f2046be1c62424d4529fa7ca10c7de7b6c60ca150f388dcda88936ec5a04bb34f67e022bf96a78682b73
+  data.tar.gz: fc4052f0c7d4b99c5c9bfa9d04f44c9cecd52e32b496e176dc8a72305b856683e5081a40146b7cbc9cd9352ae7e42a173360f7165711994f594802dbe03022ee

data/VERSION

@@ -1 +1 @@
-0.0.1
+0.0.2

data/ext/gte/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "gte"
-version = "0.0.1"
+version = "0.0.2"
 edition = "2021"
 authors = ["elcuervo <elcuervo@elcuervo.net>"]
 license = "MIT"

data/ext/gte/src/embedder.rs

@@ -40,10 +40,14 @@ impl Embedder {
         dir: P,
         num_threads: usize,
         optimization_level: u8,
+        model_name: Option<&str>,
     ) -> Result<Self> {
         let dir = dir.as_ref();
         let tokenizer_path = dir.join("tokenizer.json");
-        let model_path = resolve_model_path(dir)?;
+        let model_path = match model_name.filter(|s| !s.is_empty()) {
+            Some(name) => resolve_named_model(dir, name)?,
+            None => resolve_model_path(dir)?,
+        };
 
         if !tokenizer_path.exists() {
             return Err(GteError::Tokenizer(format!(
@@ -179,6 +183,20 @@ fn host_parallelism() -> usize {
         .unwrap_or(1)
 }
 
+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()
+    )))
+}
+
 fn resolve_model_path(dir: &Path) -> Result<PathBuf> {
     let candidates = [
         dir.join("onnx").join("text_model.onnx"),

data/ext/gte/src/ruby_embedder.rs

@@ -96,8 +96,14 @@ impl RbEmbedder {
         dir_path: String,
         num_threads: usize,
         optimization_level: u8,
+        model_name: String,
     ) -> Result<Self, Error> {
-        let embedder = Embedder::from_dir(&dir_path, num_threads, optimization_level)
+        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)?;
         Ok(RbEmbedder {
             inner: Arc::new(embedder),
@@ -202,7 +208,7 @@ 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, 3))?;
+    embedder_class.define_singleton_method("new", function!(RbEmbedder::rb_new, 4))?;
     embedder_class.define_method("embed", method!(RbEmbedder::rb_embed, 1))?;
     embedder_class.define_method("embed_one", method!(RbEmbedder::rb_embed_one, 1))?;
 

data/ext/gte/src/session.rs

@@ -36,13 +36,8 @@ pub fn build_session<P: AsRef<Path>>(model_path: P, config: &ModelConfig) -> Res
 }
 
 fn preferred_execution_providers() -> Vec<ExecutionProviderDispatch> {
-    let default_providers = if cfg!(all(target_os = "macos", target_arch = "aarch64")) {
-        "xnnpack,coreml"
-    } else {
-        "xnnpack"
-    };
     let order = std::env::var("GTE_EXECUTION_PROVIDERS")
-        .unwrap_or_else(|_| default_providers.to_string())
+        .unwrap_or_else(|_| "xnnpack".to_string())
         .to_ascii_lowercase();
 
     let mut providers = Vec::new();

data/ext/gte/tests/inference_integration_test.rs

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

data/lib/gte.rb

@@ -1,13 +1,17 @@
 # frozen_string_literal: true
 
-require 'gte/gte'
+begin
+  require "gte/#{RUBY_VERSION.to_f}/gte"
+rescue LoadError
+  require 'gte/gte'
+end
 
 module GTE
   VERSION = File.read(File.expand_path('../VERSION', __dir__)).strip
 
   class Model
-    def initialize(dir, num_threads: 0, optimization_level: 3)
-      @embedder = GTE::Embedder.new(dir, num_threads, optimization_level)
+    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
 
     def embed(texts)
@@ -26,7 +30,7 @@ module GTE
     end
   end
 
-  def self.new(dir, num_threads: 0, optimization_level: 3)
-    Model.new(dir, num_threads: num_threads, optimization_level: optimization_level)
+  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)
   end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: gte
 version: !ruby/object:Gem::Version
-  version: 0.0.1
+  version: 0.0.2
 platform: ruby
 authors:
 - elcuervo