ner-ruby 0.1.1 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: c1a6a8375878db28b788d78ae9c92ac2628321d1aa65a5368daf6fcd1701b569
-  data.tar.gz: 922ace5c18cf9949d5f92998b31d2c7bdf8097589ffc355c0f7b3a1b7151bf05
+  metadata.gz: 6c55a89d813528d2927bdf620869fa10c5660155aad6b15033f806426142c171
+  data.tar.gz: e0e58c4ac730b697b8e3b24c4228ca64bc2be048c346df63c8a9713d58fa0da1
 SHA512:
-  metadata.gz: 11f6b9feee7c04a79ed35c0eba290ce7d74b680d12dde7fddd403542e567180539328e1b85b80d6687984f51438ef17e5b26dfa9dcadbcf0265834ec076b006b
-  data.tar.gz: ac6dcd6c2292c65750d34cd76a379cf0597597e9733b2142529427f60ad6f9bb27beef42e32fddbdc09cf8a96ae58cc38c6b26f32d0fbd9bb5891c26acd0f38d
+  metadata.gz: 328af572e8556ebf0527288cab60bebb9e11441c29bdee5a68eeb835324d76d1d8f70cc70b2ff7b5190e569c7ab848690304492492183cbf41dc93b62a9d27fe
+  data.tar.gz: 603134f836d15da0af2068a26d1c99e47c311ada80a80167257b30841b3e4cf730e5ac062ce8090d32b5c180bfc6aa4c3027094b3ecf77299228a7405acb59ae

data/README.md

@@ -1,11 +1,11 @@
 # ner-ruby
 
-Named Entity Recognition for Ruby using ONNX models.
+Named Entity Recognition for Ruby. Extract entities (people, places, organizations) from text using ONNX models or API backends.
 
 ## Installation
 
 ```ruby
-gem "ner-ruby", "~> 0.1"
+gem "ner-ruby"
 ```
 
 ## Usage
@@ -13,19 +13,34 @@ gem "ner-ruby", "~> 0.1"
 ```ruby
 require "ner_ruby"
 
-ner = NerRuby::Recognizer.new(
-  model: "path/to/ner.onnx",
-  tokenizer: "path/to/tokenizer.json"
+# ONNX backend
+recognizer = NerRuby::Recognizer.new(
+  model_path: "path/to/model.onnx",
+  labels: [:PER, :LOC, :ORG, :MISC]
 )
 
-entities = ner.recognize("Jokowi visited Jakarta on Monday")
-# => [Entity(text: "Jokowi", label: :PER), Entity(text: "Jakarta", label: :LOC)]
+entities = recognizer.recognize("John works at Google in Mountain View")
+entities.each do |e|
+  puts "#{e.text} (#{e.label}) [#{e.start_offset}:#{e.end_offset}] score=#{e.score}"
+end
 
-entities = ner.recognize(text, labels: [:PER, :ORG])
-
-results = ner.recognize_batch(["Text one", "Text two"])
+# API backend
+recognizer = NerRuby::Recognizer.new(
+  backend: :api,
+  provider: :openai,
+  api_key: ENV["OPENAI_API_KEY"]
+)
 ```
 
+## Features
+
+- ONNX Runtime inference with auto label map from config.json
+- API backend support (OpenAI, etc.)
+- IOB/BIO tag decoding with wordpiece token merging
+- Character span offsets (start_offset, end_offset)
+- Numerically stable softmax
+- Empty/nil text guards
+
 ## License
 
 MIT

data/lib/ner_ruby/configuration.rb

@@ -3,7 +3,9 @@
 module NerRuby
   class Configuration
     attr_accessor :default_model_path, :default_tokenizer_path,
-                  :default_labels, :min_score, :batch_size
+                  :default_labels, :min_score, :batch_size,
+                  :min_scores_per_type, :enable_cache,
+                  :max_length, :stride, :merge_adjacent
 
     def initialize
       @default_model_path = nil
@@ -11,6 +13,19 @@ module NerRuby
       @default_labels = nil
       @min_score = 0.5
       @batch_size = 32
+      @min_scores_per_type = {}
+      @enable_cache = true
+      @max_length = 512
+      @stride = 128
+      @merge_adjacent = true
+    end
+
+    def model_registry
+      @model_registry ||= ModelRegistry.new
+    end
+
+    def register_model(name, **opts)
+      model_registry.register(name, **opts)
     end
   end
 end

data/lib/ner_ruby/model_cache.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+module NerRuby
+  class ModelCache
+    def initialize
+      @cache = {}
+      @mutex = Mutex.new
+    end
+
+    def get(key)
+      @mutex.synchronize { @cache[key] }
+    end
+
+    def set(key, value)
+      @mutex.synchronize { @cache[key] = value }
+    end
+
+    def has?(key)
+      @mutex.synchronize { @cache.key?(key) }
+    end
+
+    def clear
+      @mutex.synchronize { @cache.clear }
+    end
+
+    def size
+      @mutex.synchronize { @cache.size }
+    end
+  end
+end

data/lib/ner_ruby/model_registry.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+module NerRuby
+  class ModelRegistry
+    BUILT_IN = {
+      english: {
+        repo_id: "dslim/bert-base-NER",
+        model_file: "model.onnx",
+        tokenizer: "dslim/bert-base-NER",
+        label_map: { 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" }
+      },
+      indonesian: {
+        repo_id: "cahya/bert-base-indonesian-NER",
+        model_file: "model.onnx",
+        tokenizer: "cahya/bert-base-indonesian-NER",
+        label_map: { 0 => "O", 1 => "B-PER", 2 => "I-PER", 3 => "B-LOC",
+                     4 => "I-LOC", 5 => "B-ORG", 6 => "I-ORG" }
+      },
+      multilingual: {
+        repo_id: "Davlan/bert-base-multilingual-cased-ner-hrl",
+        model_file: "model.onnx",
+        tokenizer: "Davlan/bert-base-multilingual-cased-ner-hrl",
+        label_map: { 0 => "O", 1 => "B-PER", 2 => "I-PER", 3 => "B-ORG",
+                     4 => "I-ORG", 5 => "B-LOC", 6 => "I-LOC", 7 => "B-DATE", 8 => "I-DATE" }
+      }
+    }.freeze
+
+    def initialize
+      @custom = {}
+    end
+
+    def register(name, repo_id: nil, model_path: nil, tokenizer_path: nil, label_map: nil)
+      @custom[name.to_sym] = {
+        repo_id: repo_id,
+        model_path: model_path,
+        tokenizer_path: tokenizer_path,
+        label_map: label_map
+      }
+    end
+
+    def get(name)
+      @custom[name.to_sym] || BUILT_IN[name.to_sym]
+    end
+
+    def available
+      (BUILT_IN.keys + @custom.keys).uniq
+    end
+  end
+end

data/lib/ner_ruby/recognizer.rb

@@ -2,6 +2,8 @@
 
 module NerRuby
   class Recognizer
+    @@cache = ModelCache.new
+
     def initialize(model: nil, tokenizer: nil, label_map: nil, backend: nil, provider: nil, api_key: nil)
       config = NerRuby.configuration
 
@@ -14,8 +16,16 @@ module NerRuby
         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)
+          if config.enable_cache && @@cache.has?(model_path)
+            cached = @@cache.get(model_path)
+            @model = cached[:model]
+            @tokenizer = cached[:tokenizer]
+          else
+            @model = Models::Onnx.new(model_path: model_path)
+            @tokenizer = load_tokenizer(tokenizer_path)
+            @@cache.set(model_path, { model: @model, tokenizer: @tokenizer }) if config.enable_cache
+          end
+
           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)
@@ -23,6 +33,24 @@ module NerRuby
       end
     end
 
+    # Load a recognizer from a registered model name
+    def self.from_pretrained(name)
+      config = NerRuby.configuration
+      model_info = config.model_registry.get(name)
+      raise Error, "Unknown model: #{name}. Available: #{config.model_registry.available.join(', ')}" unless model_info
+
+      if model_info[:model_path] && model_info[:tokenizer_path]
+        new(
+          model: model_info[:model_path],
+          tokenizer: model_info[:tokenizer_path],
+          label_map: model_info[:label_map]
+        )
+      else
+        # API-based fallback
+        new(backend: :api, provider: :huggingface)
+      end
+    end
+
     def recognize(text, labels: nil)
       return [] if text.nil? || text.strip.empty?
       validate_labels!(labels) if labels
@@ -34,19 +62,29 @@ module NerRuby
         entities = @pipeline.call(text)
       end
 
+      # Merge adjacent entities of the same type
+      if NerRuby.configuration.merge_adjacent
+        entities = merge_adjacent_entities(entities)
+      end
+
+      # Filter by labels
       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 }
+      # Filter by per-type or global min_score
+      filter_by_score(entities)
     end
 
     def recognize_batch(texts, labels: nil)
       texts.map { |text| recognize(text, labels: labels) }
     end
 
+    def self.clear_cache
+      @@cache.clear
+    end
+
     private
 
     def load_tokenizer(path)
@@ -59,5 +97,43 @@ module NerRuby
         raise ValidationError, "labels must be an array of symbols or strings"
       end
     end
+
+    def merge_adjacent_entities(entities)
+      return entities if entities.empty?
+
+      merged = [entities.first]
+      entities[1..].each do |entity|
+        prev = merged.last
+        if prev.label == entity.label && adjacent?(prev, entity)
+          # Merge into a new entity
+          merged[-1] = Entity.new(
+            text: "#{prev.text} #{entity.text}",
+            label: prev.label,
+            start_offset: prev.start_offset,
+            end_offset: entity.end_offset,
+            score: ((prev.score + entity.score) / 2.0).round(4)
+          )
+        else
+          merged << entity
+        end
+      end
+      merged
+    end
+
+    def adjacent?(a, b)
+      return true if a.end_offset && b.start_offset && (b.start_offset - a.end_offset).abs <= 1
+      false
+    end
+
+    def filter_by_score(entities)
+      config = NerRuby.configuration
+      per_type = config.min_scores_per_type
+      global_min = config.min_score
+
+      entities.select do |e|
+        threshold = per_type[e.label] || global_min
+        e.score >= threshold
+      end
+    end
   end
 end

data/lib/ner_ruby/sliding_window.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+module NerRuby
+  class SlidingWindow
+    DEFAULT_MAX_LENGTH = 512
+    DEFAULT_STRIDE = 128
+
+    def initialize(max_length: DEFAULT_MAX_LENGTH, stride: DEFAULT_STRIDE)
+      @max_length = max_length
+      @stride = stride
+    end
+
+    # Split tokens into overlapping windows
+    def split(tokens, ids)
+      return [{ tokens: tokens, ids: ids, offset: 0 }] if tokens.length <= @max_length
+
+      windows = []
+      start = 0
+
+      while start < tokens.length
+        window_end = [start + @max_length, tokens.length].min
+        windows << {
+          tokens: tokens[start...window_end],
+          ids: ids[start...window_end],
+          offset: start
+        }
+        break if window_end >= tokens.length
+        start += @max_length - @stride
+      end
+
+      windows
+    end
+
+    # Merge entities from overlapping windows, preferring higher scores
+    def merge_entities(window_results)
+      all_entities = []
+
+      window_results.each do |entities|
+        entities.each do |entity|
+          existing = all_entities.find { |e| overlaps?(e, entity) }
+          if existing
+            # Keep the one with higher score
+            if entity.score > existing.score
+              all_entities.delete(existing)
+              all_entities << entity
+            end
+          else
+            all_entities << entity
+          end
+        end
+      end
+
+      all_entities.sort_by { |e| e.start_offset || 0 }
+    end
+
+    private
+
+    def overlaps?(a, b)
+      return false unless a.start_offset && b.start_offset && a.end_offset && b.end_offset
+      a.start_offset < b.end_offset && b.start_offset < a.end_offset
+    end
+  end
+end

data/lib/ner_ruby/version.rb

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

data/lib/ner_ruby.rb

@@ -9,6 +9,9 @@ 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/model_registry"
+require_relative "ner_ruby/model_cache"
+require_relative "ner_ruby/sliding_window"
 require_relative "ner_ruby/recognizer"
 
 module NerRuby

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: ner-ruby
 version: !ruby/object:Gem::Version
-  version: 0.1.1
+  version: 0.2.0
 platform: ruby
 authors:
 - Johannes Dwi Cahyo
@@ -96,11 +96,14 @@ files:
 - lib/ner_ruby/decoder.rb
 - lib/ner_ruby/entity.rb
 - lib/ner_ruby/error.rb
+- lib/ner_ruby/model_cache.rb
+- lib/ner_ruby/model_registry.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/sliding_window.rb
 - lib/ner_ruby/version.rb
 - ner-ruby.gemspec
 homepage: https://github.com/johannesdwicahyo/ner-ruby