mitie 0.1.5 → 0.1.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: e0ceaa2d4609a2a1b3b4056d67d88a1b0f55616ac8fd1a0509f070c352a96ea1
-  data.tar.gz: 0b90eba1027ca5a46a405a97411d2b5fe193d20666ee8c8bad347bb2c79c225b
+  metadata.gz: 1e36228752078ce915d9d6cd49f3d8087c62263b0e41bebb5bcfd5a478398fe6
+  data.tar.gz: 54881a7c47cd855b678cf71a5baa7486b22acdbac802e15eda3984fc2215838d
 SHA512:
-  metadata.gz: 8281e51659e08157d305535f3cd242082d4173368c36fa167716771b008a82233d02e43d1015c39e27c4c704ae0ec53935834eb0fefb71dac3b515962987d7eb
-  data.tar.gz: 96d3564684f8197651f93238876f2bbeb69f860a8a373f36dcafb578db698555153b320f2bd0c799cd8d5d939ad05c3746c59e907d4ad62cfa4917fbb7ec2313
+  metadata.gz: 573e0e14ecd284cb9754abe95f2c3cd2ebbd9e2dbebaa5a0f33fad847abe6ebc3d0875c916d05f95fc1fdd9208c12d906315f0f5a557dfc6b3db63b620aa3f6f
+  data.tar.gz: 272335c83fc924aaf696382959d6cd93dd667139f282fef71b78b29a6784bda16f221cdd36691e885db6075312a00d463772f4408ee130ea5dde6a7a49b77d55

data/CHANGELOG.md

@@ -1,3 +1,8 @@
+## 0.1.6 (2022-03-20)
+
+- Added support for training NER models
+- Improved ARM detection
+
 ## 0.1.5 (2021-01-29)
 
 - Fixed issue with multibyte characters

data/README.md

@@ -1,18 +1,18 @@
-# MITIE
+# MITIE Ruby
 
 [MITIE](https://github.com/mit-nlp/MITIE) - named-entity recognition and binary relation detection - for Ruby
 
 - Finds people, organizations, and locations in text
 - Detects relationships between entities, like `PERSON` was born in `LOCATION`
 
-[![Build Status](https://github.com/ankane/mitie/workflows/build/badge.svg?branch=master)](https://github.com/ankane/mitie/actions)
+[![Build Status](https://github.com/ankane/mitie-ruby/workflows/build/badge.svg?branch=master)](https://github.com/ankane/mitie-ruby/actions)
 
 ## Installation
 
 Add this line to your application’s Gemfile:
 
 ```ruby
-gem 'mitie'
+gem "mitie"
 ```
 
 And download the pre-trained model for your language:
@@ -69,6 +69,41 @@ Get all tags for a model
 model.tags
 ```
 
+## Training
+
+Load an NER model into a trainer
+
+```ruby
+trainer = Mitie::NERTrainer.new("total_word_feature_extractor.dat")
+```
+
+Create training instances
+
+```ruby
+tokens = ["You", "can", "do", "machine", "learning", "in", "Ruby", "!"]
+instance = Mitie::NERTrainingInstance.new(tokens)
+instance.add_entity(3..4, "topic")    # machine learning
+instance.add_entity(6..6, "language") # Ruby
+```
+
+Add the training instances to the trainer
+
+```ruby
+trainer.add(instance)
+```
+
+Train the model
+
+```ruby
+model = trainer.train
+```
+
+Save the model
+
+```ruby
+model.save_to_disk("ner_model.dat")
+```
+
 ## Binary Relation Detection
 
 Detect relationships betweens two entities, like:
@@ -105,22 +140,22 @@ This returns
 
 ## History
 
-View the [changelog](https://github.com/ankane/mitie/blob/master/CHANGELOG.md)
+View the [changelog](https://github.com/ankane/mitie-ruby/blob/master/CHANGELOG.md)
 
 ## Contributing
 
 Everyone is encouraged to help improve this project. Here are a few ways you can help:
 
-- [Report bugs](https://github.com/ankane/mitie/issues)
-- Fix bugs and [submit pull requests](https://github.com/ankane/mitie/pulls)
+- [Report bugs](https://github.com/ankane/mitie-ruby/issues)
+- Fix bugs and [submit pull requests](https://github.com/ankane/mitie-ruby/pulls)
 - Write, clarify, or fix documentation
 - Suggest or add new features
 
 To get started with development:
 
 ```sh
-git clone https://github.com/ankane/mitie.git
-cd mitie
+git clone https://github.com/ankane/mitie-ruby.git
+cd mitie-ruby
 bundle install
 bundle exec rake vendor:all
 

data/lib/mitie/binary_relation_detector.rb

@@ -35,6 +35,7 @@ module Mitie
 
         score_ptr = Fiddle::Pointer.malloc(Fiddle::SIZEOF_DOUBLE)
         status = FFI.mitie_classify_binary_relation(pointer, relation, score_ptr)
+        # TODO make Mitie::Error
         raise "Bad status: #{status}" if status != 0
         score = score_ptr.to_s(Fiddle::SIZEOF_DOUBLE).unpack1("d")
         if score > 0

data/lib/mitie/document.rb

@@ -84,11 +84,7 @@ module Mitie
     def tokenize
       @tokenize ||= begin
         if text.is_a?(Array)
-          # malloc uses memset to set all bytes to 0
-          tokens_ptr = Fiddle::Pointer.malloc(Fiddle::SIZEOF_VOIDP * (text.size + 1))
-          text.size.times do |i|
-            tokens_ptr[i * Fiddle::SIZEOF_VOIDP, Fiddle::SIZEOF_VOIDP] = Fiddle::Pointer.to_ptr(text[i]).ref
-          end
+          tokens_ptr = Utils.array_to_pointer(text)
           [tokens_ptr, nil]
         else
           offsets_ptr = Fiddle::Pointer.malloc(Fiddle::SIZEOF_VOIDP)

data/lib/mitie/ffi.rb

@@ -10,9 +10,12 @@ module Mitie
       raise e
     end
 
+    # https://github.com/mit-nlp/MITIE/blob/master/mitielib/include/mitie.h
+
     extern "void mitie_free(void* object)"
     extern "char** mitie_tokenize(const char* text)"
     extern "char** mitie_tokenize_with_offsets(const char* text, unsigned long** token_offsets)"
+    extern "int mitie_save_named_entity_extractor(const char* filename, const mitie_named_entity_extractor* ner)"
 
     extern "mitie_named_entity_extractor* mitie_load_named_entity_extractor(const char* filename)"
     extern "unsigned long mitie_get_num_possible_ner_tags(const mitie_named_entity_extractor* ner)"
@@ -26,6 +29,21 @@ module Mitie
     extern "const char* mitie_ner_get_detection_tagstr(const mitie_named_entity_detections* dets, unsigned long idx)"
     extern "double mitie_ner_get_detection_score(const mitie_named_entity_detections* dets, unsigned long idx)"
 
+    extern "mitie_ner_training_instance* mitie_create_ner_training_instance(char** tokens)"
+    extern "unsigned long mitie_ner_training_instance_num_entities(const mitie_ner_training_instance* instance)"
+    extern "unsigned long mitie_ner_training_instance_num_tokens(const mitie_ner_training_instance* instance)"
+    extern "int mitie_overlaps_any_entity(mitie_ner_training_instance* instance, unsigned long start, unsigned long length)"
+    extern "int mitie_add_ner_training_entity(mitie_ner_training_instance* instance, unsigned long start, unsigned long length, const char* label)"
+
+    extern "mitie_ner_trainer* mitie_create_ner_trainer(const char* filename)"
+    extern "unsigned long mitie_ner_trainer_size(const mitie_ner_trainer* trainer)"
+    extern "int mitie_add_ner_training_instance(mitie_ner_trainer* trainer, const mitie_ner_training_instance* instance)"
+    extern "void mitie_ner_trainer_set_beta(mitie_ner_trainer* trainer, double beta)"
+    extern "double mitie_ner_trainer_get_beta(const mitie_ner_trainer* trainer)"
+    extern "void mitie_ner_trainer_set_num_threads(mitie_ner_trainer* trainer, unsigned long num_threads)"
+    extern "unsigned long mitie_ner_trainer_get_num_threads(const mitie_ner_trainer* trainer)"
+    extern "mitie_named_entity_extractor* mitie_train_named_entity_extractor(const mitie_ner_trainer* trainer)"
+
     extern "mitie_binary_relation_detector* mitie_load_binary_relation_detector(const char* filename)"
     extern "const char* mitie_binary_relation_detector_name_string(const mitie_binary_relation_detector* detector)"
     extern "int mitie_entities_overlap(unsigned long arg1_start, unsigned long arg1_length, unsigned long arg2_start, unsigned long arg2_length)"

data/lib/mitie/ner.rb

@@ -2,11 +2,18 @@ module Mitie
   class NER
     attr_reader :pointer
 
-    def initialize(path)
-      # better error message
-      raise ArgumentError, "File does not exist" unless File.exist?(path)
-      @pointer = FFI.mitie_load_named_entity_extractor(path)
-      ObjectSpace.define_finalizer(self, self.class.finalize(pointer))
+    def initialize(path = nil, pointer: nil)
+      if path
+        # better error message
+        raise ArgumentError, "File does not exist" unless File.exist?(path)
+        @pointer = FFI.mitie_load_named_entity_extractor(path)
+      elsif pointer
+        @pointer = pointer
+      else
+        raise ArgumentError, "Must pass either a path or a pointer"
+      end
+
+      ObjectSpace.define_finalizer(self, self.class.finalize(@pointer))
     end
 
     def tags
@@ -23,6 +30,13 @@ module Mitie
       doc(text).entities
     end
 
+    def save_to_disk(filename)
+      if FFI.mitie_save_named_entity_extractor(filename, pointer) != 0
+        raise Error, "Unable to save model"
+      end
+      nil
+    end
+
     def tokens(text)
       doc(text).tokens
     end

data/lib/mitie/ner_trainer.rb

@@ -0,0 +1,51 @@
+module Mitie
+  class NERTrainer
+    def initialize(filename)
+      raise ArgumentError, "File does not exist" unless File.exist?(filename)
+      @pointer = FFI.mitie_create_ner_trainer(filename)
+
+      ObjectSpace.define_finalizer(self, self.class.finalize(@pointer))
+    end
+
+    def add(instance)
+      FFI.mitie_add_ner_training_instance(@pointer, instance.pointer)
+    end
+
+    def beta
+      FFI.mitie_ner_trainer_get_beta(@pointer)
+    end
+
+    def beta=(value)
+      raise ArgumentError, "beta must be greater than or equal to zero" unless value >= 0
+
+      FFI.mitie_ner_trainer_set_beta(@pointer, value)
+    end
+
+    def num_threads
+      FFI.mitie_ner_trainer_get_num_threads(@pointer)
+    end
+
+    def num_threads=(value)
+      FFI.mitie_ner_trainer_set_num_threads(@pointer, value)
+    end
+
+    def size
+      FFI.mitie_ner_trainer_size(@pointer)
+    end
+
+    def train
+      raise Error, "You can't call train() on an empty trainer" if size.zero?
+
+      extractor = FFI.mitie_train_named_entity_extractor(@pointer)
+
+      raise Error, "Unable to create named entity extractor. Probably ran out of RAM." if extractor.null?
+
+      Mitie::NER.new(pointer: extractor)
+    end
+
+    def self.finalize(pointer)
+      # must use proc instead of stabby lambda
+      proc { FFI.mitie_free(pointer) }
+    end
+  end
+end

data/lib/mitie/ner_training_instance.rb

@@ -0,0 +1,49 @@
+module Mitie
+  class NERTrainingInstance
+    attr_reader :pointer
+
+    def initialize(tokens)
+      tokens_pointer = Utils.array_to_pointer(tokens)
+
+      @pointer = FFI.mitie_create_ner_training_instance(tokens_pointer)
+      raise Error, "Unable to create training instance. Probably ran out of RAM." if @pointer.null?
+
+      ObjectSpace.define_finalizer(self, self.class.finalize(@pointer))
+    end
+
+    def add_entity(range, label)
+      if range.none? || range.end >= num_tokens || range.begin < 0
+        raise ArgumentError, "Invalid range"
+      end
+
+      raise ArgumentError, "Range overlaps existing entity" if overlaps_any_entity?(range)
+
+      unless FFI.mitie_add_ner_training_entity(@pointer, range.begin, range.size, label).zero?
+        raise Error, "Unable to add entity to training instance. Probably ran out of RAM."
+      end
+
+      nil
+    end
+
+    def num_entities
+      FFI.mitie_ner_training_instance_num_entities(@pointer)
+    end
+
+    def num_tokens
+      FFI.mitie_ner_training_instance_num_tokens(@pointer)
+    end
+
+    def overlaps_any_entity?(range)
+      if range.none? || range.max >= num_tokens
+        raise ArgumentError, "Invalid range"
+      end
+
+      FFI.mitie_overlaps_any_entity(@pointer, range.begin, range.size) == 1
+    end
+
+    def self.finalize(pointer)
+      # must use proc instead of stabby lambda
+      proc { FFI.mitie_free(pointer) }
+    end
+  end
+end

data/lib/mitie/utils.rb

@@ -0,0 +1,12 @@
+module Mitie
+  module Utils
+    def self.array_to_pointer(text)
+      # malloc uses memset to set all bytes to 0
+      tokens_ptr = Fiddle::Pointer.malloc(Fiddle::SIZEOF_VOIDP * (text.size + 1))
+      text.size.times do |i|
+        tokens_ptr[i * Fiddle::SIZEOF_VOIDP, Fiddle::SIZEOF_VOIDP] = Fiddle::Pointer.to_ptr(text[i]).ref
+      end
+      tokens_ptr
+    end
+  end
+end

data/lib/mitie/version.rb

@@ -1,3 +1,3 @@
 module Mitie
-  VERSION = "0.1.5"
+  VERSION = "0.1.6"
 end

data/lib/mitie.rb

@@ -5,6 +5,9 @@ require "fiddle/import"
 require "mitie/binary_relation_detector"
 require "mitie/document"
 require "mitie/ner"
+require "mitie/ner_training_instance"
+require "mitie/ner_trainer"
+require "mitie/utils"
 require "mitie/version"
 
 module Mitie
@@ -16,10 +19,12 @@ module Mitie
   lib_name =
     if Gem.win_platform?
       "mitie.dll"
-    elsif RbConfig::CONFIG["arch"] =~ /arm64-darwin/i
-      "libmitie.arm64.dylib"
     elsif RbConfig::CONFIG["host_os"] =~ /darwin/i
-      "libmitie.dylib"
+      if RbConfig::CONFIG["host_cpu"] =~ /arm|aarch64/i
+        "libmitie.arm64.dylib"
+      else
+        "libmitie.dylib"
+      end
     else
       "libmitie.so"
     end

metadata

@@ -1,17 +1,17 @@
 --- !ruby/object:Gem::Specification
 name: mitie
 version: !ruby/object:Gem::Version
-  version: 0.1.5
+  version: 0.1.6
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2021-01-30 00:00:00.000000000 Z
+date: 2022-03-21 00:00:00.000000000 Z
 dependencies: []
 description:
-email: andrew@chartkick.com
+email: andrew@ankane.org
 executables: []
 extensions: []
 extra_rdoc_files: []
@@ -24,13 +24,16 @@ files:
 - lib/mitie/document.rb
 - lib/mitie/ffi.rb
 - lib/mitie/ner.rb
+- lib/mitie/ner_trainer.rb
+- lib/mitie/ner_training_instance.rb
+- lib/mitie/utils.rb
 - lib/mitie/version.rb
 - vendor/LICENSE.txt
 - vendor/libmitie.arm64.dylib
 - vendor/libmitie.dylib
 - vendor/libmitie.so
 - vendor/mitie.dll
-homepage: https://github.com/ankane/mitie
+homepage: https://github.com/ankane/mitie-ruby
 licenses:
 - BSL-1.0
 metadata: {}
@@ -49,7 +52,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.2.3
+rubygems_version: 3.3.7
 signing_key:
 specification_version: 4
 summary: Named-entity recognition for Ruby