Added support for models with external data - #10

Author: ankane

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 1.2.0 (unreleased)
+
+- Added support for models with external data
+
 ## 1.1.1 (2024-10-14)
 
 - Added `audio-classification` pipeline

lib/informers/models.rb

@@ -178,7 +178,15 @@ def self.construct_session(pretrained_model_name_or_path, file_name, **options)
       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)
+      begin
+        OnnxRuntime::InferenceSession.new(path)
+      rescue OnnxRuntime::Error => e
+        raise e unless e.message.include?(".onnx_data")
+
+        Utils::Hub.get_model_file(pretrained_model_name_or_path, "#{model_file_name}_data", true, **options)
+
+        OnnxRuntime::InferenceSession.new(path)
+      end
     end
 
     def call(model_inputs, **kwargs)

lib/informers/utils/hub.rb

@@ -81,12 +81,14 @@ def match(request)
           file if file.exists
         end
 
-        def put(request, buffer)
+        def put(request, response)
           output_path = resolve_path(request)
 
           begin
             FileUtils.mkdir_p(File.dirname(output_path))
-            File.binwrite(output_path, buffer)
+            File.open(output_path, "wb") do |f|
+              f.write(response.read(1024 * 1024)) until response.eof?
+            end
           rescue => e
             warn "An error occurred while writing the file to cache: #{e}"
           end
@@ -189,10 +191,8 @@ def self.get_model_file(path_or_repo_id, filename, fatal = true, **options)
           to_cache_response = cache && !response.is_a?(FileResponse) && response.status[0] == "200"
         end
 
-        buffer = response.read
-
         if to_cache_response && cache_key && cache.match(cache_key).nil?
-          cache.put(cache_key, buffer)
+          cache.put(cache_key, response)
         end
 
         Utils.dispatch_callback(options[:progress_callback], {

test/model_test.rb

@@ -174,6 +174,14 @@ def test_all_mpnet
     assert_elements_in_delta [0.04170236, 0.00109747, -0.01553415], embeddings[1][..2]
   end
 
+  # https://huggingface.co/BAAI/bge-m3
+  def test_bge_m3
+    sentences = ["This is an example sentence", "Each sentence is converted"]
+
+    model = Informers.pipeline("embedding", "BAAI/bge-m3")
+    model.(sentences, model_output: "token_embeddings")
+  end
+
   # https://huggingface.co/mixedbread-ai/mxbai-rerank-base-v1
   def test_mxbai_rerank
     query = "How many people live in London?"