onnx-ruby 0.1.0

data/lib/onnx_ruby/hub.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+require "fileutils"
+require "net/http"
+require "uri"
+require "json"
+
+module OnnxRuby
+  module Hub
+    DEFAULT_CACHE_DIR = File.join(Dir.home, ".cache", "onnx_ruby", "models")
+
+    # Download a model from Hugging Face Hub
+    # @param repo_id [String] e.g. "sentence-transformers/all-MiniLM-L6-v2"
+    # @param filename [String] ONNX file to download (default: "model.onnx")
+    # @param cache_dir [String] local cache directory
+    # @return [String] path to the downloaded model file
+    def self.download(repo_id, filename: "model.onnx", cache_dir: DEFAULT_CACHE_DIR, revision: "main")
+      model_dir = File.join(cache_dir, repo_id.tr("/", "--"), revision)
+      model_path = File.join(model_dir, filename)
+
+      return model_path if File.exist?(model_path)
+
+      FileUtils.mkdir_p(model_dir)
+
+      url = "https://huggingface.co/#{repo_id}/resolve/#{revision}/#{filename}"
+      download_file(url, model_path)
+
+      model_path
+    end
+
+    # List cached models
+    # @param cache_dir [String] cache directory to search
+    # @return [Array<String>] list of cached model paths
+    def self.cached_models(cache_dir: DEFAULT_CACHE_DIR)
+      return [] unless Dir.exist?(cache_dir)
+
+      Dir.glob(File.join(cache_dir, "**", "*.onnx"))
+    end
+
+    # Clear the model cache
+    # @param cache_dir [String] cache directory to clear
+    def self.clear_cache(cache_dir: DEFAULT_CACHE_DIR)
+      FileUtils.rm_rf(cache_dir) if Dir.exist?(cache_dir)
+    end
+
+    class << self
+      private
+
+      def download_file(url, dest)
+        uri = URI(url)
+        max_redirects = 5
+
+        max_redirects.times do
+          response = Net::HTTP.start(uri.host, uri.port, use_ssl: uri.scheme == "https") do |http|
+            http.request(Net::HTTP::Get.new(uri))
+          end
+
+          case response
+          when Net::HTTPSuccess
+            File.binwrite(dest, response.body)
+            return
+          when Net::HTTPRedirection
+            uri = URI(response["location"])
+          else
+            raise ModelError, "failed to download #{url}: #{response.code} #{response.message}"
+          end
+        end
+
+        raise ModelError, "too many redirects downloading #{url}"
+      end
+    end
+  end
+end

data/lib/onnx_ruby/lazy_session.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module OnnxRuby
+  class LazySession
+    def initialize(model_path, **opts)
+      @model_path = model_path
+      @opts = opts
+      @session = nil
+      @mutex = Mutex.new
+    end
+
+    def inputs
+      load_session.inputs
+    end
+
+    def outputs
+      load_session.outputs
+    end
+
+    def run(inputs, **kwargs)
+      load_session.run(inputs, **kwargs)
+    end
+
+    def loaded?
+      !@session.nil?
+    end
+
+    private
+
+    def load_session
+      return @session if @session
+
+      @mutex.synchronize do
+        @session ||= Session.new(@model_path, **@opts)
+      end
+    end
+  end
+end

data/lib/onnx_ruby/model.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+module OnnxRuby
+  # ActiveModel-style mixin for embedding generation.
+  #
+  # Usage:
+  #   class Document
+  #     include OnnxRuby::Model
+  #
+  #     onnx_model "embeddings.onnx"
+  #     onnx_input ->(doc) { { "input_ids" => doc.token_ids, "attention_mask" => doc.mask } }
+  #     onnx_output "embeddings"
+  #   end
+  #
+  #   doc = Document.new
+  #   doc.onnx_predict  # => [0.123, -0.456, ...]
+  module Model
+    def self.included(base)
+      base.extend(ClassMethods)
+    end
+
+    module ClassMethods
+      def onnx_model(path = nil, **opts)
+        if path
+          @onnx_model_path = path
+          @onnx_session_opts = opts
+        end
+        @onnx_model_path
+      end
+
+      def onnx_input(callable = nil, &block)
+        @onnx_input_fn = callable || block if callable || block
+        @onnx_input_fn
+      end
+
+      def onnx_output(name = nil)
+        @onnx_output_name = name if name
+        @onnx_output_name
+      end
+
+      def onnx_session
+        @onnx_session ||= begin
+          path = resolve_model_path(@onnx_model_path)
+          LazySession.new(path, **(@onnx_session_opts || {}))
+        end
+      end
+
+      private
+
+      def resolve_model_path(path)
+        return path if File.absolute_path?(path) && File.exist?(path)
+
+        full = File.join(OnnxRuby.configuration.models_path, path)
+        return full if File.exist?(full)
+
+        path
+      end
+    end
+
+    def onnx_predict(**run_opts)
+      input_fn = self.class.onnx_input
+      raise Error, "onnx_input not defined on #{self.class}" unless input_fn
+
+      inputs = input_fn.call(self)
+      result = self.class.onnx_session.run(inputs, **run_opts)
+
+      output_name = self.class.onnx_output
+      output_name ? result[output_name] : result.values.first
+    end
+  end
+end

data/lib/onnx_ruby/reranker.rb

@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+module OnnxRuby
+  class Reranker
+    attr_reader :session
+
+    def initialize(model_path, tokenizer: nil, **session_opts)
+      @session = Session.new(model_path, **session_opts)
+      @tokenizer = resolve_tokenizer(tokenizer)
+    end
+
+    # Rerank documents by relevance to a query
+    # @param query [String] the query text (requires tokenizer)
+    # @param documents [Array<String>] documents to rerank
+    # @return [Array<Hash>] sorted array of { document:, score:, index: }
+    def rerank(query, documents)
+      raise Error, "tokenizer is required for reranking" unless @tokenizer
+
+      pairs = documents.map { |doc| [query, doc] }
+      scores = score_pairs(pairs)
+
+      documents.each_with_index.map do |doc, i|
+        { document: doc, score: scores[i], index: i }
+      end.sort_by { |r| -r[:score] }
+    end
+
+    # Score query-document pairs with pre-tokenized inputs
+    # @param input_ids [Array<Array<Integer>>] batch of token ID sequences
+    # @param attention_mask [Array<Array<Integer>>] batch of attention masks
+    # @return [Array<Float>] relevance scores
+    def score(input_ids:, attention_mask:)
+      feed = build_feed(input_ids, attention_mask)
+      result = @session.run(feed)
+      raw_scores = find_output(result, %w[scores logits output])
+      raw_scores.map { |row| row.is_a?(Array) ? row.first : row }
+    end
+
+    private
+
+    def resolve_tokenizer(tokenizer)
+      return nil if tokenizer.nil?
+
+      if tokenizer.respond_to?(:encode)
+        tokenizer
+      else
+        begin
+          require "tokenizers"
+          Tokenizers::Tokenizer.from_pretrained(tokenizer.to_s)
+        rescue LoadError
+          raise Error, "tokenizer-ruby gem is required for text tokenization. " \
+                       "Install with: gem install tokenizers"
+        end
+      end
+    end
+
+    def score_pairs(pairs)
+      if @tokenizer.respond_to?(:encode_batch)
+        encodings = @tokenizer.encode_batch(pairs)
+        ids = encodings.map(&:ids)
+        masks = encodings.map(&:attention_mask)
+      else
+        encodings = pairs.map { |pair| @tokenizer.encode(*pair) }
+        ids = encodings.map(&:ids)
+        masks = encodings.map(&:attention_mask)
+      end
+
+      max_len = ids.map(&:length).max
+      ids = ids.map { |row| row + Array.new(max_len - row.length, 0) }
+      masks = masks.map { |row| row + Array.new(max_len - row.length, 0) }
+
+      feed = build_feed(ids, masks)
+      result = @session.run(feed)
+      raw_scores = find_output(result, %w[scores logits output])
+      raw_scores.map { |row| row.is_a?(Array) ? row.first : row }
+    end
+
+    def build_feed(ids, masks)
+      input_names = @session.inputs.map { |i| i[:name] }
+      feed = {}
+      feed[input_names.find { |n| n.include?("input_id") } || input_names[0]] = ids
+      mask_name = input_names.find { |n| n.include?("mask") || n.include?("attention") }
+      feed[mask_name] = masks if mask_name
+      feed
+    end
+
+    def find_output(result, candidate_names)
+      candidate_names.each { |name| return result[name] if result.key?(name) }
+      result.values.first
+    end
+  end
+end

data/lib/onnx_ruby/session.rb

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+module OnnxRuby
+  class Session
+    VALID_PROVIDERS = %i[cpu coreml cuda tensorrt].freeze
+
+    def initialize(model_path, providers: [:cpu], inter_threads: nil, intra_threads: nil,
+                   log_level: :warning, optimization_level: :all, memory_pattern: true,
+                   cpu_mem_arena: true, execution_mode: :sequential)
+      model_path = File.expand_path(model_path)
+      raise ModelError, "model file not found: #{model_path}" unless File.exist?(model_path)
+
+      provider_strs = Array(providers).map do |p|
+        p = p.to_sym
+        raise Error, "unknown provider: #{p}. Valid: #{VALID_PROVIDERS.join(", ")}" unless VALID_PROVIDERS.include?(p)
+        p.to_s
+      end
+
+      @session = Ext::SessionWrapper.new(
+        model_path,
+        log_level_to_int(log_level),
+        intra_threads || 0,
+        inter_threads || 0,
+        optimization_level.to_s,
+        memory_pattern,
+        cpu_mem_arena,
+        execution_mode.to_s,
+        provider_strs
+      )
+    end
+
+    def inputs
+      @session.input_info
+    end
+
+    def outputs
+      @session.output_info
+    end
+
+    def run(inputs, output_names: nil)
+      input_values = inputs.map do |name, data|
+        if data.is_a?(Tensor)
+          { name: name, data: data.flat_data, shape: data.shape, dtype: data.dtype.to_s }
+        else
+          flat = data.flatten
+          shape = infer_shape(data)
+          dtype = infer_dtype(flat)
+          { name: name, data: flat, shape: shape, dtype: dtype }
+        end
+      end
+
+      @session.run(input_values, output_names || [])
+    end
+
+    private
+
+    def log_level_to_int(level)
+      case level
+      when :verbose then 0
+      when :info then 1
+      when :warning then 2
+      when :error then 3
+      when :fatal then 4
+      else 2
+      end
+    end
+
+    def infer_shape(data)
+      shape = []
+      current = data
+      while current.is_a?(Array)
+        shape << current.length
+        current = current.first
+      end
+      shape
+    end
+
+    def infer_dtype(flat)
+      sample = flat.find { |v| !v.nil? }
+      case sample
+      when Float then "float"
+      when Integer then "int64"
+      when String then "string"
+      when true, false then "bool"
+      else "float"
+      end
+    end
+  end
+end

data/lib/onnx_ruby/session_pool.rb

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+module OnnxRuby
+  class SessionPool
+    class TimeoutError < Error; end
+
+    def initialize(model_path, size: nil, timeout: nil, **session_opts)
+      @model_path = model_path
+      @session_opts = session_opts
+      @size = size || OnnxRuby.configuration.pool_size
+      @timeout = timeout || OnnxRuby.configuration.pool_timeout
+      @pool = []
+      @mutex = Mutex.new
+      @condition = ConditionVariable.new
+      @created = 0
+    end
+
+    # Check out a session, yield it, then check it back in
+    def with_session(&block)
+      session = checkout
+      begin
+        yield session
+      ensure
+        checkin(session)
+      end
+    end
+
+    # Run inference using a pooled session
+    def run(inputs, **kwargs)
+      with_session { |s| s.run(inputs, **kwargs) }
+    end
+
+    # Current pool stats
+    def size
+      @mutex.synchronize { @created }
+    end
+
+    def available
+      @mutex.synchronize { @pool.size }
+    end
+
+    private
+
+    def checkout
+      @mutex.synchronize do
+        loop do
+          # Return an available session
+          return @pool.pop unless @pool.empty?
+
+          # Create a new one if under limit
+          if @created < @size
+            @created += 1
+            return create_session
+          end
+
+          # Wait for one to be returned
+          unless @condition.wait(@mutex, @timeout)
+            raise TimeoutError, "timed out waiting for session (pool size: #{@size})"
+          end
+        end
+      end
+    end
+
+    def checkin(session)
+      @mutex.synchronize do
+        @pool.push(session)
+        @condition.signal
+      end
+    end
+
+    def create_session
+      Session.new(@model_path, **@session_opts)
+    end
+  end
+end

data/lib/onnx_ruby/tensor.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+module OnnxRuby
+  class Tensor
+    DTYPE_MAP = {
+      float32: :float,
+      float: :float,
+      float64: :double,
+      double: :double,
+      int32: :int32,
+      int: :int32,
+      int64: :int64,
+      bool: :bool,
+      string: :string
+    }.freeze
+
+    attr_reader :shape, :dtype
+
+    def initialize(data, shape: nil, dtype: nil)
+      @data = data.flatten
+      @shape = shape || infer_shape(data)
+      @dtype = normalize_dtype(dtype || infer_dtype(@data))
+
+      validate!
+    end
+
+    def to_a
+      reshape(@data.dup, @shape)
+    end
+
+    def flat_data
+      @data
+    end
+
+    def self.float(data, shape: nil)
+      new(data, shape: shape, dtype: :float)
+    end
+
+    def self.int64(data, shape: nil)
+      new(data, shape: shape, dtype: :int64)
+    end
+
+    def self.int32(data, shape: nil)
+      new(data, shape: shape, dtype: :int32)
+    end
+
+    def self.double(data, shape: nil)
+      new(data, shape: shape, dtype: :double)
+    end
+
+    private
+
+    def normalize_dtype(dtype)
+      DTYPE_MAP.fetch(dtype) { raise TensorError, "unsupported dtype: #{dtype}" }
+    end
+
+    def infer_shape(data)
+      shape = []
+      current = data
+      while current.is_a?(Array)
+        shape << current.length
+        current = current.first
+      end
+      shape
+    end
+
+    def infer_dtype(flat)
+      sample = flat.find { |v| !v.nil? }
+      case sample
+      when Float then :float
+      when Integer then :int64
+      when String then :string
+      when true, false then :bool
+      else :float
+      end
+    end
+
+    def validate!
+      expected_size = @shape.reduce(1, :*)
+      if @data.length != expected_size
+        raise TensorError,
+              "data size #{@data.length} does not match shape #{@shape} (expected #{expected_size})"
+      end
+    end
+
+    def reshape(flat, dims)
+      return flat if dims.length <= 1
+      size = dims[1..].reduce(1, :*)
+      flat.each_slice(size).map { |slice| reshape(slice, dims[1..]) }
+    end
+  end
+end

data/lib/onnx_ruby/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module OnnxRuby
+  VERSION = "0.1.0"
+end

data/lib/onnx_ruby.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+require_relative "onnx_ruby/version"
+
+module OnnxRuby
+  class Error < StandardError; end
+  class ModelError < Error; end
+  class InferenceError < Error; end
+  class TensorError < Error; end
+end
+
+require_relative "onnx_ruby/onnx_ruby_ext"
+require_relative "onnx_ruby/tensor"
+require_relative "onnx_ruby/session"
+require_relative "onnx_ruby/embedder"
+require_relative "onnx_ruby/classifier"
+require_relative "onnx_ruby/reranker"
+require_relative "onnx_ruby/hub"
+require_relative "onnx_ruby/configuration"
+require_relative "onnx_ruby/lazy_session"
+require_relative "onnx_ruby/session_pool"
+require_relative "onnx_ruby/model"
+
+module OnnxRuby
+  class << self
+    def configuration
+      @configuration ||= Configuration.new
+    end
+
+    def configure
+      yield configuration
+    end
+  end
+
+  def self.optimize(input_path, output_path, level: :all)
+    input_path = File.expand_path(input_path)
+    raise ModelError, "model file not found: #{input_path}" unless File.exist?(input_path)
+
+    Ext.optimize_model(input_path, output_path, level.to_s)
+  end
+
+  def self.available_providers
+    Ext.available_providers
+  end
+end

data/onnx-ruby.gemspec

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+require_relative "lib/onnx_ruby/version"
+
+Gem::Specification.new do |spec|
+  spec.name = "onnx-ruby"
+  spec.version = OnnxRuby::VERSION
+  spec.authors = ["Johannes Dwi Cahyo"]
+  spec.email = ["johannesdwicahyo@gmail.com"]
+
+  spec.summary = "Ruby bindings for ONNX Runtime"
+  spec.description = "High-performance ONNX Runtime bindings for Ruby using Rice. " \
+                     "Run ONNX models locally for embeddings, classification, NER, and more."
+  spec.homepage = "https://github.com/johannesdwicahyo/onnx-ruby"
+  spec.license = "MIT"
+  spec.required_ruby_version = ">= 3.1.0"
+
+  spec.metadata["homepage_uri"] = spec.homepage
+  spec.metadata["source_code_uri"] = spec.homepage
+  spec.metadata["changelog_uri"] = "#{spec.homepage}/blob/main/CHANGELOG.md"
+
+  spec.files = Dir.chdir(__dir__) do
+    `git ls-files -z`.split("\x0").reject do |f|
+      (File.expand_path(f) == __FILE__) ||
+        f.start_with?("test/", "spec/", "features/", ".git", ".github", "script/")
+    end
+  end
+
+  spec.require_paths = ["lib"]
+  spec.extensions = ["ext/onnx_ruby/extconf.rb"]
+
+  spec.add_dependency "rice", ">= 4.0"
+
+  spec.add_development_dependency "rake", "~> 13.0"
+  spec.add_development_dependency "rake-compiler", "~> 1.2"
+  spec.add_development_dependency "minitest", "~> 5.0"
+end