Improved model_file_name option

ankane · Aug 28, 2024 · 1cc07a3 · 1cc07a3
1 parent 66d6d4c
commit 1cc07a3
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 7 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 1.0.3 (unreleased)
+
+- Improved `model_file_name` option
+
 ## 1.0.2 (2024-08-28)
 
 - Added `embedding` pipeline

diff --git a/README.md b/README.md
@@ -30,7 +30,7 @@ Embedding
 - [intfloat/e5-base-v2](#intfloate5-base-v2)
 - [nomic-ai/nomic-embed-text-v1](#nomic-ainomic-embed-text-v1)
 - [BAAI/bge-base-en-v1.5](#baaibge-base-en-v15)
-- [Xenova/jina-embeddings-v2-base-en](#xenovajina-embeddings-v2-base-en)
+- [jinaai/jina-embeddings-v2-base-en](#jinaaijina-embeddings-v2-base-en)
 
 Reranking (experimental)
 
@@ -148,14 +148,14 @@ model = Informers.pipeline("embedding", "BAAI/bge-base-en-v1.5")
 embeddings = model.(input)
 ```
 
-### Xenova/jina-embeddings-v2-base-en
+### jinaai/jina-embeddings-v2-base-en
 
-[Docs](https://huggingface.co/Xenova/jina-embeddings-v2-base-en)
+[Docs](https://huggingface.co/jinaai/jina-embeddings-v2-base-en) [unreleased]
 
 ```ruby
 sentences = ["How is the weather today?", "What is the current weather like today?"]
 
-model = Informers.pipeline("embedding", "Xenova/jina-embeddings-v2-base-en")
+model = Informers.pipeline("embedding", "jinaai/jina-embeddings-v2-base-en", model_file_name: "../model")
 embeddings = model.(sentences)
 ```
 

diff --git a/lib/informers/models.rb b/lib/informers/models.rb
@@ -135,7 +135,15 @@ def self.from_pretrained(
     end
 
     def self.construct_session(pretrained_model_name_or_path, file_name, **options)
-      model_file_name = "onnx/#{file_name}#{options[:quantized] ? "_quantized" : ""}.onnx"
+      prefix = "onnx/"
+      if file_name.start_with?("../")
+        prefix = ""
+        file_name = file_name[3..]
+      elsif file_name.start_with?("/")
+        prefix = ""
+        file_name = file_name[1..]
+      end
+      model_file_name = "#{prefix}#{file_name}#{options[:quantized] ? "_quantized" : ""}.onnx"
       path = Utils::Hub.get_model_file(pretrained_model_name_or_path, model_file_name, true, **options)
 
       OnnxRuntime::InferenceSession.new(path)

diff --git a/test/model_test.rb b/test/model_test.rb
@@ -123,11 +123,11 @@ def test_bge_base
     assert_elements_in_delta [0.00029264, -0.0619305, -0.06199387], embeddings[-1][..2]
   end
 
-  # https://huggingface.co/Xenova/jina-embeddings-v2-base-en
+  # https://huggingface.co/jinaai/jina-embeddings-v2-base-en
   def test_jina_embeddings
     sentences = ["How is the weather today?", "What is the current weather like today?"]
 
-    model = Informers.pipeline("embedding", "Xenova/jina-embeddings-v2-base-en")
+    model = Informers.pipeline("embedding", "jinaai/jina-embeddings-v2-base-en", model_file_name: "../model")
     embeddings = model.(sentences)
 
     assert_elements_in_delta [-0.02488641, -0.0429398, 0.04303398], embeddings[0][..2]