onnx-ruby 0.1.1 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: ec86be6d12455b53b31aeb0b680db0239f0fc6d2b0d93a57efcacf847728af43
-  data.tar.gz: e228293936cffd177c42055a5dcf34536ab1c30670797d231bc3a9885eba0a8e
+  metadata.gz: 16caafd6c79615589d18c10f16be77816875f823f0024b78d32357b025645d30
+  data.tar.gz: 915a8bbf683808063b752e81da2a5629157baf73959f99bee7a61d1c1871fad3
 SHA512:
-  metadata.gz: 609dc877ac507261a573b4284d249e779f07d6e45661ee894e29dd1c65e906cd3d89349f13d1931ab69c8a8f71765fc0c698e95bd59a114c9cef7db3a86ee39f
-  data.tar.gz: 86cd66b410257e644fad7a4e2148912fba30a03a1aa39e9112c5c087903fc311f8d7626416d52ccb5810380bf6316ac88ba2fb4f4a856791e0315e548f5e947b
+  metadata.gz: b3fcc3eb5eb0605ff0866af424b0436f2bd294675a5eb32c0348182ad65bea9b64997d19c3bb026c0b7b8290c8fe59038693f3737ca7984ffb4712cd0e88c2d3
+  data.tar.gz: d22068ed40991670d3dee5780339cea9232bc7627a07f796d9b0a7bc870b648250622e66b8eb3337a186eb9ff7b092e7e6bc6dd1da7c1ac4cfc99840b2b4a9ae

data/MILESTONES.md

@@ -0,0 +1,12 @@
+# Milestones
+
+## v0.1.1 (2026-03-10)
+
+### Changes
+- Empty input guards
+- TokenizerSupport extraction
+- config validation
+- session mutex fix
+
+## v0.1.0 (Initial release)
+- Initial release

data/ext/onnx_ruby/onnx_ruby_ext.cpp

@@ -94,12 +94,63 @@ static Rice::Object tensor_to_ruby(const Ort::Value& tensor) {
       break;
     }
     case ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL: {
-      const bool* data = tensor.GetTensorData<bool>();
+      const uint8_t* data = reinterpret_cast<const uint8_t*>(tensor.GetTensorData<bool>());
       for (size_t i = 0; i < total; i++) {
         flat.push(Rice::Object(data[i] ? Qtrue : Qfalse));
       }
       break;
     }
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: {
+      const uint8_t* data = tensor.GetTensorData<uint8_t>();
+      for (size_t i = 0; i < total; i++) {
+        flat.push(Rice::Object(INT2NUM(data[i])));
+      }
+      break;
+    }
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8: {
+      const int8_t* data = tensor.GetTensorData<int8_t>();
+      for (size_t i = 0; i < total; i++) {
+        flat.push(Rice::Object(INT2NUM(data[i])));
+      }
+      break;
+    }
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16: {
+      const uint16_t* data = tensor.GetTensorData<uint16_t>();
+      for (size_t i = 0; i < total; i++) {
+        flat.push(Rice::Object(INT2NUM(data[i])));
+      }
+      break;
+    }
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16: {
+      const int16_t* data = tensor.GetTensorData<int16_t>();
+      for (size_t i = 0; i < total; i++) {
+        flat.push(Rice::Object(INT2NUM(data[i])));
+      }
+      break;
+    }
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16: {
+      // float16 is stored as uint16_t; convert to float for Ruby
+      const uint16_t* data = tensor.GetTensorData<uint16_t>();
+      for (size_t i = 0; i < total; i++) {
+        // IEEE 754 half-precision to single-precision conversion
+        uint16_t h = data[i];
+        uint32_t sign = (h & 0x8000u) << 16;
+        uint32_t exponent = (h >> 10) & 0x1F;
+        uint32_t mantissa = h & 0x03FF;
+        uint32_t f;
+        if (exponent == 0) {
+          f = sign; // zero or subnormal (treat as zero for simplicity)
+        } else if (exponent == 31) {
+          f = sign | 0x7F800000u | (mantissa << 13); // inf or nan
+        } else {
+          f = sign | ((exponent + 112) << 23) | (mantissa << 13);
+        }
+        float val;
+        memcpy(&val, &f, sizeof(float));
+        flat.push(Rice::Object(rb_float_new(static_cast<double>(val))));
+      }
+      break;
+    }
     case ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING: {
       size_t count = total;
       std::vector<std::string> strings(count);
@@ -317,7 +368,13 @@ public:
       }
 
       size_t total_elements = 1;
-      for (auto dim : shape) total_elements *= dim;
+      for (auto dim : shape) {
+        if (dim < 0) throw std::runtime_error("Negative shape dimension: " + std::to_string(dim));
+        if (dim > 0 && total_elements > SIZE_MAX / static_cast<size_t>(dim)) {
+          throw std::runtime_error("Shape dimension overflow");
+        }
+        total_elements *= static_cast<size_t>(dim);
+      }
 
       if (dtype == "float") {
         float_buffers.emplace_back(total_elements);

data/lib/onnx_ruby/classifier.rb

@@ -86,9 +86,15 @@ module OnnxRuby
     end
 
     def softmax(logits)
+      # Clamp extreme values to prevent overflow
       max_val = logits.max
-      exps = logits.map { |v| Math.exp(v - max_val) }
+      exps = logits.map do |v|
+        clamped = v - max_val
+        clamped = -500.0 if clamped < -500.0
+        Math.exp(clamped)
+      end
       sum = exps.sum
+      sum = Float::MIN if sum.zero?
       exps.map { |v| v / sum }
     end
   end

data/lib/onnx_ruby/session.rb

@@ -42,8 +42,13 @@ module OnnxRuby
         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)
+          flat = data.flatten
+          expected_size = shape.reduce(1, :*)
+          if flat.length != expected_size
+            raise TensorError,
+                  "input '#{name}' data size #{flat.length} does not match shape #{shape} (expected #{expected_size})"
+          end
           dtype = infer_dtype(flat)
           { name: name, data: flat, shape: shape, dtype: dtype }
         end
@@ -70,6 +75,13 @@ module OnnxRuby
       current = data
       while current.is_a?(Array)
         shape << current.length
+        if current.length > 1 && current.all? { |el| el.is_a?(Array) }
+          lengths = current.map(&:length).uniq
+          if lengths.size > 1
+            raise TensorError,
+                  "jagged array detected: sub-arrays have lengths #{lengths.sort.join(', ')} at dimension #{shape.size - 1}"
+          end
+        end
         current = current.first
       end
       shape

data/lib/onnx_ruby/session_pool.rb

@@ -64,12 +64,15 @@ module OnnxRuby
     def checkin(session)
       @mutex.synchronize do
         @pool.push(session)
-        @condition.signal
+        @condition.broadcast
       end
     end
 
     def create_session
       Session.new(@model_path, **@session_opts)
+    rescue => e
+      @created -= 1
+      raise
     end
   end
 end

data/lib/onnx_ruby/tensor.rb

@@ -59,6 +59,14 @@ module OnnxRuby
       current = data
       while current.is_a?(Array)
         shape << current.length
+        # Check for jagged arrays: all sub-arrays at this level must have the same length
+        if current.length > 1 && current.all? { |el| el.is_a?(Array) }
+          lengths = current.map(&:length).uniq
+          if lengths.size > 1
+            raise TensorError,
+                  "jagged array detected: sub-arrays have lengths #{lengths.sort.join(', ')} at dimension #{shape.size - 1}"
+          end
+        end
         current = current.first
       end
       shape

data/lib/onnx_ruby/version.rb

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

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: onnx-ruby
 version: !ruby/object:Gem::Version
-  version: 0.1.1
+  version: 0.2.0
 platform: ruby
 authors:
 - Johannes Dwi Cahyo
@@ -74,9 +74,9 @@ extensions:
 - ext/onnx_ruby/extconf.rb
 extra_rdoc_files: []
 files:
-- CLAUDE.md
 - Gemfile
 - LICENSE
+- MILESTONES.md
 - README.md
 - Rakefile
 - examples/classification.rb

data/CLAUDE.md

@@ -1,334 +0,0 @@
-# onnx-ruby
-
-## Project Overview
-
-Ruby bindings for [ONNX Runtime](https://github.com/microsoft/onnxruntime), Microsoft's high-performance inference engine for ONNX models. This gem wraps the ONNX Runtime C++ API using **Rice** (same approach as zvec-ruby) to give Ruby developers fast local model inference.
-
-This unlocks: local embeddings, text classification, named entity recognition, sentiment analysis, reranking, and any other ML model exported to ONNX format — all without Python or API calls.
-
-## Author
-
-- Name: Johannes Dwi Cahyo
-- GitHub: johannesdwicahyo
-- Repo: git@github.com:johannesdwicahyo/onnx-ruby.git
-
-## Technical Approach
-
-### Binding Strategy: Rice 4.x (C++ → Ruby)
-
-ONNX Runtime has a C++ API (`onnxruntime_cxx_api.h`). We wrap it using Rice, exactly like zvec-ruby.
-
-**Important lessons from zvec-ruby to apply here:**
-- Use `require "mkmf-rice"` (not `require "rice/extconf"`) for Rice 4.x
-- Use `define_module_under()` and `define_enum_under()` for Rice 4.x
-- Wrap all raw `VALUE` returns in `Rice::Object()` when pushing to Arrays
-- Use `std::make_shared` when C++ API expects shared_ptr
-- Extract results to Ruby Hashes/Arrays in C++ before returning (avoid dangling pointers)
-- Ship precompiled gems — ONNX Runtime is a large C++ library, nobody wants to build it
-- On macOS use `-force_load` for static archives with static initializers
-- Default to safe options (like mmap=true in zvec-ruby)
-
-### ONNX Runtime Linking
-
-ONNX Runtime provides **prebuilt shared libraries** (`.so`/`.dylib`) for all platforms. Unlike zvec (which required building from source), we can download the official release and link against it. This is much simpler.
-
-Download from: https://github.com/microsoft/onnxruntime/releases
-
-The `extconf.rb` should:
-1. Check for `ONNX_RUNTIME_DIR` env var
-2. Check for system-installed onnxruntime via pkg-config
-3. Auto-download the correct prebuilt release if neither found
-
-### Precompiled Gems
-
-For precompiled gems, statically link or bundle the ONNX Runtime `.dylib`/`.so` inside the gem. The gem will be ~50-80MB but users get zero-install experience.
-
-## Core API Design
-
-```ruby
-require "onnx_ruby"
-
-# --- Session (model loading) ---
-
-# Load a model
-session = OnnxRuby::Session.new("model.onnx")
-
-# With options
-session = OnnxRuby::Session.new("model.onnx",
-  providers: [:cpu],              # :cpu, :cuda, :coreml, :tensorrt
-  inter_threads: 4,
-  intra_threads: 2,
-  log_level: :warning
-)
-
-# Model info
-session.inputs   # => [{ name: "input_ids", type: :int64, shape: [-1, 512] }]
-session.outputs  # => [{ name: "embeddings", type: :float, shape: [-1, 384] }]
-
-# --- Inference ---
-
-# Run inference
-result = session.run(
-  { "input_ids" => [[101, 2023, 2003, 1037, 3231, 102]] },
-)
-result["embeddings"]  # => [[0.0123, -0.0456, ...]]
-
-# With output names
-result = session.run(inputs, output_names: ["embeddings"])
-
-# --- Tensor ---
-
-# Create tensors explicitly
-tensor = OnnxRuby::Tensor.new([1, 2, 3, 4], shape: [2, 2], dtype: :int64)
-tensor.to_a     # => [[1, 2], [3, 4]]
-tensor.shape    # => [2, 2]
-tensor.dtype    # => :int64
-
-# From flat array
-tensor = OnnxRuby::Tensor.float([0.1, 0.2, 0.3], shape: [1, 3])
-
-# --- High-Level Helpers ---
-
-# Embedding model (wraps session with pre/post processing)
-embedder = OnnxRuby::Embedder.new("all-MiniLM-L6-v2.onnx",
-  tokenizer: "sentence-transformers/all-MiniLM-L6-v2"  # requires tokenizer-ruby
-)
-embeddings = embedder.embed("Hello world")         # => [0.0123, ...]
-embeddings = embedder.embed_batch(["Hello", "World"])  # => [[...], [...]]
-
-# Classifier
-classifier = OnnxRuby::Classifier.new("intent_model.onnx",
-  tokenizer: "bert-base-uncased",
-  labels: ["greeting", "farewell", "question", "command"]
-)
-classifier.predict("Hello there!")  # => { label: "greeting", score: 0.95 }
-```
-
-## Features to Implement
-
-### Phase 1 — Core (MVP)
-- [ ] `Session.new(path, options)` — load ONNX model
-- [ ] `session.run(inputs)` — run inference, return outputs
-- [ ] `session.inputs` / `session.outputs` — model metadata
-- [ ] `Tensor` class — create and manipulate tensors
-- [ ] Support dtypes: float32, float64, int32, int64, string, bool
-- [ ] Support shapes: 1D, 2D, 3D, 4D tensors
-- [ ] CPU execution provider
-
-### Phase 2 — Providers & Options
-- [ ] CoreML provider (macOS acceleration)
-- [ ] CUDA provider (NVIDIA GPU)
-- [ ] Session options: threading, memory, optimization level
-- [ ] Model optimization: `OnnxRuby.optimize("model.onnx", "optimized.onnx")`
-- [ ] Dynamic shapes (batching)
-
-### Phase 3 — High-Level API
-- [ ] `Embedder` — embedding model wrapper (tokenize → infer → normalize)
-- [ ] `Classifier` — text classification wrapper
-- [ ] `Reranker` — cross-encoder reranking wrapper
-- [ ] Integration with tokenizer-ruby for text preprocessing
-- [ ] Model hub: `OnnxRuby::Hub.download("sentence-transformers/all-MiniLM-L6-v2")`
-
-### Phase 4 — Rails Integration
-- [ ] `OnnxRuby.configure { |c| c.models_path = "app/models/onnx" }`
-- [ ] Lazy model loading (load on first inference)
-- [ ] Connection pool for thread-safe concurrent inference
-- [ ] ActiveModel integration for embedding generation
-
-## Project Structure
-
-```
-onnx-ruby/
-├── CLAUDE.md
-├── Gemfile
-├── Rakefile
-├── LICENSE                  # MIT
-├── README.md
-├── onnx-ruby.gemspec
-├── lib/
-│   ├── onnx_ruby.rb
-│   └── onnx_ruby/
-│       ├── version.rb
-│       ├── session.rb
-│       ├── tensor.rb
-│       ├── embedder.rb
-│       ├── classifier.rb
-│       └── reranker.rb
-├── ext/
-│   └── onnx_ruby/
-│       ├── extconf.rb
-│       └── onnx_ruby_ext.cpp
-├── test/
-│   ├── test_helper.rb
-│   ├── test_session.rb
-│   ├── test_tensor.rb
-│   ├── test_inference.rb
-│   └── models/           # small test ONNX models
-│       └── .gitkeep
-├── script/
-│   ├── download_onnxruntime.sh
-│   └── package_native_gem.rb
-└── examples/
-    ├── embedding.rb
-    ├── classification.rb
-    └── with_zvec.rb       # full RAG example with zvec-ruby
-```
-
-## Dependencies
-
-### Runtime
-- `rice` (>= 4.0) — C++ to Ruby bindings
-- ONNX Runtime shared library (bundled in precompiled gems)
-
-### Optional
-- `tokenizer-ruby` — for Embedder/Classifier text preprocessing
-
-### Development
-- `rake-compiler` for building native extensions
-- `rake-compiler-dock` for cross-compilation
-- `minitest` for testing
-- `rake` for tasks
-
-## Key C++ Binding Details
-
-### ONNX Runtime C++ API Structure
-
-```cpp
-#include <onnxruntime_cxx_api.h>
-
-// Key classes to wrap:
-Ort::Env           // Runtime environment (singleton)
-Ort::Session       // Model session
-Ort::SessionOptions
-Ort::Value         // Tensor (input/output)
-Ort::MemoryInfo    // Memory allocation info
-Ort::TypeInfo      // Model input/output type info
-Ort::TensorTypeAndShapeInfo
-```
-
-### extconf.rb approach
-
-```ruby
-require "mkmf-rice"
-
-# Try to find ONNX Runtime
-ort_dir = ENV["ONNX_RUNTIME_DIR"]
-
-unless ort_dir
-  # Auto-download prebuilt ONNX Runtime for the current platform
-  # from https://github.com/microsoft/onnxruntime/releases
-  ort_dir = download_onnxruntime()  # helper function
-end
-
-dir_config("onnxruntime", "#{ort_dir}/include", "#{ort_dir}/lib")
-$INCFLAGS << " -I#{ort_dir}/include"
-$LDFLAGS << " -L#{ort_dir}/lib"
-$libs << " -lonnxruntime"
-
-have_header("onnxruntime_cxx_api.h") or
-  abort "Cannot find ONNX Runtime headers"
-
-create_makefile("onnx_ruby/onnx_ruby_ext")
-```
-
-### C++ Extension Skeleton
-
-```cpp
-#include <rice/rice.hpp>
-#include <onnxruntime_cxx_api.h>
-
-using namespace Rice;
-
-// Global ORT environment (initialized once)
-static Ort::Env& get_env() {
-  static Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "onnx_ruby");
-  return env;
-}
-
-// Wrap Ort::Session
-// Wrap Ort::Value (Tensor)
-// Handle type conversion: Ruby Array ↔ ORT Tensor
-// Map ORT errors to Ruby exceptions
-
-void Init_onnx_ruby_ext() {
-  Module rb_mOnnxRuby = define_module("OnnxRuby");
-  Module rb_mExt = define_module_under(rb_mOnnxRuby, "Ext");
-
-  // Define Session, Tensor, etc.
-}
-```
-
-### Critical: Tensor ↔ Ruby Array Conversion
-
-The most complex part. Need to handle:
-- Ruby Array of floats → ORT float32 tensor (most common for embeddings)
-- Ruby Array of integers → ORT int64 tensor (for token IDs)
-- Nested Ruby Arrays → multi-dimensional tensors
-- ORT output tensors → Ruby Arrays (with proper Float wrapping via `Rice::Object(rb_float_new())`)
-
-```cpp
-// Ruby Array → ORT Tensor
-Ort::Value array_to_tensor(Rice::Array arr, const std::vector<int64_t>& shape) {
-  // Flatten nested arrays
-  // Detect dtype from Ruby values
-  // Create ORT tensor with proper memory allocation
-}
-
-// ORT Tensor → Ruby Array
-Rice::Object tensor_to_array(const Ort::Value& tensor) {
-  // Read shape
-  // Read dtype
-  // Copy data to Ruby Array (with Rice::Object wrapping!)
-}
-```
-
-## Testing Strategy
-
-- Test with small ONNX models (generate test models with Python: `torch.onnx.export`)
-- Test model loading and metadata inspection
-- Test inference with known inputs/outputs
-- Test all supported dtypes
-- Test batch inference
-- Test error handling (invalid model, wrong input shape, etc.)
-- Benchmark against Python's onnxruntime for correctness
-
-### Create test models (Python script to include):
-```python
-# script/create_test_models.py
-import torch
-import torch.nn as nn
-
-# Simple linear model for testing
-class SimpleModel(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.linear = nn.Linear(4, 3)
-    def forward(self, x):
-        return self.linear(x)
-
-model = SimpleModel()
-dummy = torch.randn(1, 4)
-torch.onnx.export(model, dummy, "test/models/simple.onnx",
-                  input_names=["input"], output_names=["output"])
-```
-
-## Publishing
-
-- RubyGems.org: `GEM_HOST_API_KEY=rubygems_5d46e91ceb51fb455e98a7f491a2321bb6879f9be35d6842 gem push onnx-ruby-*.gem`
-- gem.coop: `GEM_HOST_API_KEY=hjncPswY8PbGDfLPw4RMj928 gem push onnx-ruby-*.gem --host https://beta.gem.coop/@johannesdwicahyo`
-
-## Notes from zvec-ruby Experience
-
-- **Rice 4.x API**: `define_module_under()`, `define_enum_under()`, not the 3.x syntax
-- **Rice::Object wrapping**: ALWAYS wrap `rb_float_new()`, `Qtrue`, `Qnil` in `Rice::Object()` when pushing to Arrays
-- **shared_ptr**: Use `make_shared` when C++ expects `shared_ptr`, accept by `const T&` in bindings
-- **Extract results in C++**: Don't try to push C++ objects directly into Ruby arrays. Extract to Hashes/Arrays first.
-- **Precompiled gems**: Essential. Use `script/package_native_gem.rb` for macOS, `rake-compiler-dock` for Linux.
-- **ONNX Runtime ships prebuilt binaries**: Much easier than zvec. Download from GitHub releases, link against `.dylib`/`.so`.
-- **Static initializers**: May need `-force_load` if ONNX Runtime uses static registration patterns.
-- **mmap/memory**: ONNX Runtime manages its own memory via allocators. Let it handle memory, don't fight it.
-
-## Existing Ruby ONNX Solutions
-
-- `onnxruntime` gem by ankane — exists but is FFI-based and limited. We can provide better performance and API with Rice + additional high-level features (Embedder, Classifier, Reranker).
-- Differentiate by: better API, precompiled gems, high-level wrappers, tokenizer-ruby integration