ner-ruby 0.1.1

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: c1a6a8375878db28b788d78ae9c92ac2628321d1aa65a5368daf6fcd1701b569
+  data.tar.gz: 922ace5c18cf9949d5f92998b31d2c7bdf8097589ffc355c0f7b3a1b7151bf05
+SHA512:
+  metadata.gz: 11f6b9feee7c04a79ed35c0eba290ce7d74b680d12dde7fddd403542e567180539328e1b85b80d6687984f51438ef17e5b26dfa9dcadbcf0265834ec076b006b
+  data.tar.gz: ac6dcd6c2292c65750d34cd76a379cf0597597e9733b2142529427f60ad6f9bb27beef42e32fddbdc09cf8a96ae58cc38c6b26f32d0fbd9bb5891c26acd0f38d

data/CHANGELOG.md

@@ -0,0 +1,24 @@
+# Changelog
+
+## 0.1.1 (2026-03-09)
+
+### Fixed
+- Wire API backend to Recognizer via `backend: :api` option
+- Softmax numerical stability (subtract max before exp, clamp output to 0.0..1.0)
+- Empty/nil text input returns empty array instead of crashing
+
+### Added
+- Model config auto-detection: reads `config.json` alongside ONNX model for label map
+- Entity character span offsets (`start_offset`, `end_offset`) in original text
+- Input validation for model path, API key, and labels parameter
+- `ValidationError` and `ConfigurationError` error classes
+
+## 0.1.0 (2026-03-09)
+
+- Initial release
+- Named Entity Recognition using ONNX models
+- IOB/BIO tag decoder with sub-token merging
+- Entity result objects with type predicates
+- Support for PER, LOC, ORG, MISC entity types
+- API backend for OpenAI and HuggingFace
+- Configuration DSL

data/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Johannes Dwi Cahyo
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

data/README.md

@@ -0,0 +1,31 @@
+# ner-ruby
+
+Named Entity Recognition for Ruby using ONNX models.
+
+## Installation
+
+```ruby
+gem "ner-ruby", "~> 0.1"
+```
+
+## Usage
+
+```ruby
+require "ner_ruby"
+
+ner = NerRuby::Recognizer.new(
+  model: "path/to/ner.onnx",
+  tokenizer: "path/to/tokenizer.json"
+)
+
+entities = ner.recognize("Jokowi visited Jakarta on Monday")
+# => [Entity(text: "Jokowi", label: :PER), Entity(text: "Jakarta", label: :LOC)]
+
+entities = ner.recognize(text, labels: [:PER, :ORG])
+
+results = ner.recognize_batch(["Text one", "Text two"])
+```
+
+## License
+
+MIT

data/Rakefile

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+require "rake/testtask"
+
+Rake::TestTask.new(:test) do |t|
+  t.libs << "test"
+  t.libs << "lib"
+  t.test_files = FileList["test/**/test_*.rb"]
+end
+
+task default: :test

data/lib/ner_ruby/configuration.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+module NerRuby
+  class Configuration
+    attr_accessor :default_model_path, :default_tokenizer_path,
+                  :default_labels, :min_score, :batch_size
+
+    def initialize
+      @default_model_path = nil
+      @default_tokenizer_path = nil
+      @default_labels = nil
+      @min_score = 0.5
+      @batch_size = 32
+    end
+  end
+end

data/lib/ner_ruby/decoder.rb

@@ -0,0 +1,108 @@
+# frozen_string_literal: true
+
+module NerRuby
+  class Decoder
+    LABEL_MAPS = {
+      "bert-base-NER" => {
+        0 => "O",
+        1 => "B-MISC",
+        2 => "I-MISC",
+        3 => "B-PER",
+        4 => "I-PER",
+        5 => "B-ORG",
+        6 => "I-ORG",
+        7 => "B-LOC",
+        8 => "I-LOC"
+      }
+    }.freeze
+
+    def initialize(label_map: nil)
+      @label_map = label_map || LABEL_MAPS["bert-base-NER"]
+    end
+
+    def decode(tokens, predictions, scores: nil, original_text: nil)
+      entities = []
+      current_entity = nil
+
+      tokens.each_with_index do |token, i|
+        next if special_token?(token)
+
+        label = @label_map[predictions[i]] || "O"
+        score = scores ? scores[i] : 1.0
+
+        if label.start_with?("B-")
+          entities << build_entity(current_entity, original_text) if current_entity
+          entity_type = label.sub("B-", "")
+          current_entity = { raw_tokens: [token], label: entity_type, scores: [score] }
+        elsif label.start_with?("I-") && current_entity
+          entity_type = label.sub("I-", "")
+          if entity_type == current_entity[:label]
+            current_entity[:raw_tokens] << token
+            current_entity[:scores] << score
+          else
+            entities << build_entity(current_entity, original_text)
+            current_entity = nil
+          end
+        else
+          entities << build_entity(current_entity, original_text) if current_entity
+          current_entity = nil
+        end
+      end
+
+      entities << build_entity(current_entity, original_text) if current_entity
+      entities
+    end
+
+    private
+
+    def special_token?(token)
+      token == "[CLS]" || token == "[SEP]" || token == "[PAD]" ||
+        token == "<s>" || token == "</s>" || token == "<pad>"
+    end
+
+    def clean_token(token)
+      token.sub(/^##/, "")
+    end
+
+    def build_entity(entity_data, original_text = nil)
+      text = merge_tokens(entity_data[:raw_tokens])
+      avg_score = entity_data[:scores].sum / entity_data[:scores].size
+      avg_score = avg_score.clamp(0.0, 1.0)
+
+      start_offset = nil
+      end_offset = nil
+
+      if original_text
+        idx = original_text.index(text)
+        if idx
+          start_offset = idx
+          end_offset = idx + text.length
+        end
+      end
+
+      Entity.new(
+        text: text,
+        label: entity_data[:label],
+        start_offset: start_offset,
+        end_offset: end_offset,
+        score: avg_score.round(4)
+      )
+    end
+
+    def merge_tokens(raw_tokens)
+      result = clean_token(raw_tokens.first) || ""
+      raw_tokens[1..].each do |token|
+        if wordpiece?(token)
+          result += clean_token(token)
+        else
+          result += " #{token}"
+        end
+      end
+      result
+    end
+
+    def wordpiece?(token)
+      token.start_with?("##")
+    end
+  end
+end

data/lib/ner_ruby/entity.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module NerRuby
+  class Entity
+    attr_reader :text, :label, :start_offset, :end_offset, :score
+
+    LABELS = %i[PER LOC ORG MISC DATE TIME MONEY PERCENT QUANTITY].freeze
+
+    def initialize(text:, label:, start_offset: nil, end_offset: nil, score: 0.0)
+      @text = text
+      @label = label.to_sym
+      @start_offset = start_offset
+      @end_offset = end_offset
+      @score = score
+    end
+
+    def person?
+      label == :PER
+    end
+
+    def location?
+      label == :LOC
+    end
+
+    def organization?
+      label == :ORG
+    end
+
+    def to_h
+      {
+        text: @text,
+        label: @label,
+        start_offset: @start_offset,
+        end_offset: @end_offset,
+        score: @score
+      }
+    end
+
+    def to_s
+      "#{@text} [#{@label}] (#{(@score * 100).round(1)}%)"
+    end
+  end
+end

data/lib/ner_ruby/error.rb

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+module NerRuby
+  class Error < StandardError; end
+  class ModelNotFoundError < Error; end
+  class TokenizerError < Error; end
+  class InferenceError < Error; end
+  class ConfigurationError < Error; end
+  class ValidationError < Error; end
+end

data/lib/ner_ruby/models/api.rb

@@ -0,0 +1,102 @@
+# frozen_string_literal: true
+
+require "net/http"
+require "uri"
+require "json"
+
+module NerRuby
+  module Models
+    class Api < Base
+      PROVIDERS = {
+        openai: {
+          url: "https://api.openai.com/v1/chat/completions",
+          model: "gpt-4o"
+        },
+        huggingface: {
+          url: "https://api-inference.huggingface.co/models/",
+          model: "dslim/bert-base-NER"
+        }
+      }.freeze
+
+      def initialize(provider: :openai, api_key: nil, model: nil)
+        @provider = provider
+        @api_key = api_key || ENV["#{provider.to_s.upcase}_API_KEY"]
+        @model = model || PROVIDERS.dig(provider, :model)
+
+        raise Error, "API key is required for #{provider}" unless @api_key
+      end
+
+      def recognize(text, labels: nil)
+        return [] if text.nil? || text.strip.empty?
+
+        case @provider
+        when :openai then recognize_openai(text, labels: labels)
+        when :huggingface then recognize_huggingface(text)
+        else raise Error, "Unknown provider: #{@provider}"
+        end
+      end
+
+      private
+
+      def recognize_openai(text, labels: nil)
+        label_str = (labels || Entity::LABELS).map(&:to_s).join(", ")
+        prompt = <<~PROMPT
+          Extract named entities from the following text. Return JSON array with objects having keys: text, label, score.
+          Labels: #{label_str}
+          Text: #{text}
+        PROMPT
+
+        body = {
+          model: @model,
+          messages: [{ role: "user", content: prompt }],
+          response_format: { type: "json_object" }
+        }
+
+        response = post_json(PROVIDERS[:openai][:url], body, {
+          "Authorization" => "Bearer #{@api_key}"
+        })
+
+        content = response.dig("choices", 0, "message", "content")
+        parsed = JSON.parse(content)
+        entities = parsed["entities"] || parsed
+
+        entities.map do |e|
+          Entity.new(
+            text: e["text"],
+            label: e["label"],
+            score: e["score"] || 0.9
+          )
+        end
+      end
+
+      def recognize_huggingface(text)
+        url = "#{PROVIDERS[:huggingface][:url]}#{@model}"
+        response = post_json(url, { inputs: text }, {
+          "Authorization" => "Bearer #{@api_key}"
+        })
+
+        response.map do |e|
+          Entity.new(
+            text: e["word"],
+            label: e["entity_group"] || e["entity"],
+            score: e["score"] || 0.0
+          )
+        end
+      end
+
+      def post_json(url, body, headers)
+        uri = URI.parse(url)
+        http = Net::HTTP.new(uri.host, uri.port)
+        http.use_ssl = true
+
+        req = Net::HTTP::Post.new(uri.request_uri)
+        req["Content-Type"] = "application/json"
+        headers.each { |k, v| req[k] = v }
+        req.body = JSON.generate(body)
+
+        response = http.request(req)
+        JSON.parse(response.body)
+      end
+    end
+  end
+end

data/lib/ner_ruby/models/base.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module NerRuby
+  module Models
+    class Base
+      def predict(input_ids)
+        raise NotImplementedError, "#{self.class}#predict not implemented"
+      end
+    end
+  end
+end

data/lib/ner_ruby/models/onnx.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+require "json"
+
+module NerRuby
+  module Models
+    class Onnx < Base
+      attr_reader :label_map
+
+      def initialize(model_path:)
+        require "onnx_ruby"
+        @model_path = model_path
+        raise ModelNotFoundError, "Model not found: #{model_path}" unless File.exist?(model_path)
+
+        @session = OnnxRuby::Session.new(model_path)
+        @label_map = load_config_label_map
+      end
+
+      def predict(input_ids)
+        attention_mask = Array.new(input_ids.length, 1)
+        token_type_ids = Array.new(input_ids.length, 0)
+
+        outputs = @session.run(
+          input_ids: [input_ids],
+          attention_mask: [attention_mask],
+          token_type_ids: [token_type_ids]
+        )
+
+        outputs[0][0]
+      end
+
+      private
+
+      def load_config_label_map
+        config_path = File.join(File.dirname(@model_path), "config.json")
+        return nil unless File.exist?(config_path)
+
+        config = JSON.parse(File.read(config_path))
+        id2label = config["id2label"]
+        return nil unless id2label.is_a?(Hash)
+
+        id2label.each_with_object({}) do |(k, v), map|
+          map[k.to_i] = v.to_s
+        end
+      end
+    end
+  end
+end

data/lib/ner_ruby/pipeline.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module NerRuby
+  class Pipeline
+    def initialize(model:, tokenizer:, decoder: nil)
+      @model = model
+      @tokenizer = tokenizer
+      @decoder = decoder || Decoder.new
+    end
+
+    def call(text)
+      encoding = @tokenizer.encode(text)
+      tokens = encoding[:tokens] || encoding["tokens"]
+      input_ids = encoding[:ids] || encoding["ids"]
+
+      logits = @model.predict(input_ids)
+
+      predictions = logits.map { |row| row.each_with_index.max_by { |v, _| v }.last }
+      scores = logits.map { |row| softmax(row).max }
+
+      @decoder.decode(tokens, predictions, scores: scores, original_text: text)
+    end
+
+    private
+
+    def softmax(logits)
+      max = logits.max
+      exps = logits.map { |x| Math.exp(x - max) }
+      sum = exps.sum
+      exps.map { |x| (x / sum).clamp(0.0, 1.0) }
+    end
+  end
+end

data/lib/ner_ruby/recognizer.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+module NerRuby
+  class Recognizer
+    def initialize(model: nil, tokenizer: nil, label_map: nil, backend: nil, provider: nil, api_key: nil)
+      config = NerRuby.configuration
+
+      if backend == :api
+        @api_model = Models::Api.new(provider: provider || :openai, api_key: api_key)
+      else
+        model_path = model || config.default_model_path
+        tokenizer_path = tokenizer || config.default_tokenizer_path
+
+        if model_path && tokenizer_path
+          raise ModelNotFoundError, "Model not found: #{model_path}" unless File.exist?(model_path)
+
+          @model = Models::Onnx.new(model_path: model_path)
+          @tokenizer = load_tokenizer(tokenizer_path)
+          detected_label_map = label_map || @model.label_map
+          @decoder = Decoder.new(label_map: detected_label_map)
+          @pipeline = Pipeline.new(model: @model, tokenizer: @tokenizer, decoder: @decoder)
+        end
+      end
+    end
+
+    def recognize(text, labels: nil)
+      return [] if text.nil? || text.strip.empty?
+      validate_labels!(labels) if labels
+
+      if @api_model
+        entities = @api_model.recognize(text, labels: labels)
+      else
+        raise Error, "No model loaded. Provide model and tokenizer paths." unless @pipeline
+        entities = @pipeline.call(text)
+      end
+
+      if labels
+        label_syms = labels.map(&:to_sym)
+        entities = entities.select { |e| label_syms.include?(e.label) }
+      end
+
+      min = NerRuby.configuration.min_score
+      entities.select { |e| e.score >= min }
+    end
+
+    def recognize_batch(texts, labels: nil)
+      texts.map { |text| recognize(text, labels: labels) }
+    end
+
+    private
+
+    def load_tokenizer(path)
+      require "tokenizer_ruby"
+      TokenizerRuby::Tokenizer.new(path)
+    end
+
+    def validate_labels!(labels)
+      unless labels.is_a?(Array) && labels.all? { |l| l.is_a?(Symbol) || l.is_a?(String) }
+        raise ValidationError, "labels must be an array of symbols or strings"
+      end
+    end
+  end
+end

data/lib/ner_ruby/version.rb

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

data/lib/ner_ruby.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+require_relative "ner_ruby/version"
+require_relative "ner_ruby/error"
+require_relative "ner_ruby/configuration"
+require_relative "ner_ruby/entity"
+require_relative "ner_ruby/decoder"
+require_relative "ner_ruby/pipeline"
+require_relative "ner_ruby/models/base"
+require_relative "ner_ruby/models/onnx"
+require_relative "ner_ruby/models/api"
+require_relative "ner_ruby/recognizer"
+
+module NerRuby
+  class << self
+    def configuration
+      @configuration ||= Configuration.new
+    end
+
+    def configure
+      yield(configuration)
+    end
+
+    def reset_configuration!
+      @configuration = Configuration.new
+    end
+  end
+end

data/ner-ruby.gemspec

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+require_relative "lib/ner_ruby/version"
+
+Gem::Specification.new do |spec|
+  spec.name = "ner-ruby"
+  spec.version = NerRuby::VERSION
+  spec.authors = ["Johannes Dwi Cahyo"]
+  spec.email = ["johannes@example.com"]
+  spec.summary = "Named Entity Recognition for Ruby using ONNX models"
+  spec.description = "NER using ONNX models via onnx-ruby and tokenizer-ruby. Extracts people, places, organizations, and other entities from text."
+  spec.homepage = "https://github.com/johannesdwicahyo/ner-ruby"
+  spec.license = "MIT"
+  spec.required_ruby_version = ">= 3.0.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[
+    "lib/**/*.rb",
+    "README.md",
+    "LICENSE",
+    "CHANGELOG.md",
+    "Rakefile",
+    "ner-ruby.gemspec"
+  ]
+  spec.require_paths = ["lib"]
+
+  spec.add_dependency "onnx-ruby", "~> 0.1"
+  spec.add_dependency "tokenizer-ruby", "~> 0.1"
+
+  spec.add_development_dependency "minitest", "~> 5.0"
+  spec.add_development_dependency "rake", "~> 13.0"
+  spec.add_development_dependency "webmock", "~> 3.0"
+end

metadata

@@ -0,0 +1,130 @@
+--- !ruby/object:Gem::Specification
+name: ner-ruby
+version: !ruby/object:Gem::Version
+  version: 0.1.1
+platform: ruby
+authors:
+- Johannes Dwi Cahyo
+bindir: bin
+cert_chain: []
+date: 1980-01-02 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: onnx-ruby
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.1'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.1'
+- !ruby/object:Gem::Dependency
+  name: tokenizer-ruby
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.1'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.1'
+- !ruby/object:Gem::Dependency
+  name: minitest
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '5.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '5.0'
+- !ruby/object:Gem::Dependency
+  name: rake
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '13.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '13.0'
+- !ruby/object:Gem::Dependency
+  name: webmock
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.0'
+description: NER using ONNX models via onnx-ruby and tokenizer-ruby. Extracts people,
+  places, organizations, and other entities from text.
+email:
+- johannes@example.com
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- CHANGELOG.md
+- LICENSE
+- README.md
+- Rakefile
+- lib/ner_ruby.rb
+- lib/ner_ruby/configuration.rb
+- lib/ner_ruby/decoder.rb
+- lib/ner_ruby/entity.rb
+- lib/ner_ruby/error.rb
+- lib/ner_ruby/models/api.rb
+- lib/ner_ruby/models/base.rb
+- lib/ner_ruby/models/onnx.rb
+- lib/ner_ruby/pipeline.rb
+- lib/ner_ruby/recognizer.rb
+- lib/ner_ruby/version.rb
+- ner-ruby.gemspec
+homepage: https://github.com/johannesdwicahyo/ner-ruby
+licenses:
+- MIT
+metadata:
+  homepage_uri: https://github.com/johannesdwicahyo/ner-ruby
+  source_code_uri: https://github.com/johannesdwicahyo/ner-ruby
+  changelog_uri: https://github.com/johannesdwicahyo/ner-ruby/blob/main/CHANGELOG.md
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: 3.0.0
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.6.9
+specification_version: 4
+summary: Named Entity Recognition for Ruby using ONNX models
+test_files: []