mitie 0.1.5 → 0.2.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: e0ceaa2d4609a2a1b3b4056d67d88a1b0f55616ac8fd1a0509f070c352a96ea1
-  data.tar.gz: 0b90eba1027ca5a46a405a97411d2b5fe193d20666ee8c8bad347bb2c79c225b
+  metadata.gz: dc2b7bdba2fba6b335ab9750efab8766190e618c4b8e2542ff6409f6727ec8b2
+  data.tar.gz: 03d85b928082a04b46209694c8e1b294e5c4205057a26b5cc65e695bbe3564a9
 SHA512:
-  metadata.gz: 8281e51659e08157d305535f3cd242082d4173368c36fa167716771b008a82233d02e43d1015c39e27c4c704ae0ec53935834eb0fefb71dac3b515962987d7eb
-  data.tar.gz: 96d3564684f8197651f93238876f2bbeb69f860a8a373f36dcafb578db698555153b320f2bd0c799cd8d5d939ad05c3746c59e907d4ad62cfa4917fbb7ec2313
+  metadata.gz: 4092a2dc005bb76527429454c301179c63f0eeee4913a3a8f56190e13d7f7551ee3bbfcd1098a6c335cf3b886ec0dee4e52f5c062975a4a058d117229f45b340
+  data.tar.gz: 066e5800a520b16002388088fc367939a26738bd638631d3b86ff9074d508cf5cd051154fbbf7be129f07f82c97a851caf7a6486e6fff2fb7f6a4b252d426ec0

data/CHANGELOG.md

@@ -1,3 +1,19 @@
+## 0.2.1 (2022-06-12)
+
+- Added `tokenize` and `tokenize_file` methods
+- Added support for untokenized text to text categorization
+
+## 0.2.0 (2022-06-01)
+
+- Added support for text categorization
+- Added support for training binary relation detectors
+- Dropped support for Ruby < 2.7
+
+## 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,21 +1,21 @@
-# MITIE
+# MITIE Ruby
 
-[MITIE](https://github.com/mit-nlp/MITIE) - named-entity recognition and binary relation detection - for Ruby
+[MITIE](https://github.com/mit-nlp/MITIE) - named-entity recognition, binary relation detection, and text categorization - 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:
+And download the pre-trained models for your language:
 
 - [English](https://github.com/mit-nlp/MITIE/releases/download/v0.4/MITIE-models-v0.2.tar.bz2)
 - [Spanish](https://github.com/mit-nlp/MITIE/releases/download/v0.4/MITIE-models-v0.2-Spanish.zip)
@@ -23,6 +23,12 @@ And download the pre-trained model for your language:
 
 ## Getting Started
 
+- [Named Entity Recognition](#named-entity-recognition)
+- [Binary Relation Detection](#binary-relation-detection)
+- [Text Categorization](#text-categorization)
+
+## Named Entity Recognition
+
 Load an NER model
 
 ```ruby
@@ -69,6 +75,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:
@@ -103,24 +144,98 @@ This returns
 [{first: "Shopify", second: "Ottawa", score: 0.17649169745814464}]
 ```
 
+### Training
+
+Load an NER model into a trainer
+
+```ruby
+trainer = Mitie::BinaryRelationTrainer.new(model)
+```
+
+Add positive and negative examples to the trainer
+
+```ruby
+tokens = ["Shopify", "was", "founded", "in", "Ottawa"]
+trainer.add_positive_binary_relation(tokens, 0..0, 4..4)
+trainer.add_negative_binary_relation(tokens, 4..4, 0..0)
+```
+
+Train the detector
+
+```ruby
+detector = trainer.train
+```
+
+Save the detector
+
+```ruby
+detector.save_to_disk("binary_relation_detector.svm")
+```
+
+## Text Categorization
+
+Load a model into a trainer
+
+```ruby
+trainer = Mitie::TextCategorizerTrainer.new("total_word_feature_extractor.dat")
+```
+
+Add labeled text to the trainer
+
+```ruby
+trainer.add("This is super cool", "positive")
+```
+
+Train the model
+
+```ruby
+model = trainer.train
+```
+
+Save the model
+
+```ruby
+model.save_to_disk("text_categorization_model.dat")
+```
+
+Load a saved model
+
+```ruby
+model = Mitie::TextCategorizer.new("text_categorization_model.dat")
+```
+
+Categorize text
+
+```ruby
+model.categorize("What a super nice day")
+```
+
+## Deployment
+
+Check out [Trove](https://github.com/ankane/trove) for deploying models.
+
+```sh
+trove push ner_model.dat
+```
+
 ## 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

@@ -1,9 +1,16 @@
 module Mitie
   class BinaryRelationDetector
-    def initialize(path)
-      # better error message
-      raise ArgumentError, "File does not exist" unless File.exist?(path)
-      @pointer = FFI.mitie_load_binary_relation_detector(path)
+    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_binary_relation_detector(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
 
@@ -23,37 +30,52 @@ module Mitie
 
       relations = []
       combinations.each do |entity1, entity2|
-        relation =
-          FFI.mitie_extract_binary_relation(
-            doc.model.pointer,
-            doc.send(:tokens_ptr),
-            entity1[:token_index],
-            entity1[:token_length],
-            entity2[:token_index],
-            entity2[:token_length]
-          )
-
-        score_ptr = Fiddle::Pointer.malloc(Fiddle::SIZEOF_DOUBLE)
-        status = FFI.mitie_classify_binary_relation(pointer, relation, score_ptr)
-        raise "Bad status: #{status}" if status != 0
-        score = score_ptr.to_s(Fiddle::SIZEOF_DOUBLE).unpack1("d")
-        if score > 0
-          relations << {
-            first: entity1[:text],
-            second: entity2[:text],
-            score: score
-          }
-        end
+        relation = extract_relation(doc, entity1, entity2)
+        relations << relation if relation
       end
       relations
     end
 
+    def save_to_disk(filename)
+      if FFI.mitie_save_binary_relation_detector(filename, pointer) != 0
+        raise Error, "Unable to save detector"
+      end
+      nil
+    end
+
     private
 
     def pointer
       @pointer
     end
 
+    def extract_relation(doc, entity1, entity2)
+      relation =
+        FFI.mitie_extract_binary_relation(
+          doc.model.pointer,
+          doc.send(:tokens_ptr),
+          entity1[:token_index],
+          entity1[:token_length],
+          entity2[:token_index],
+          entity2[:token_length]
+        )
+
+      score_ptr = Fiddle::Pointer.malloc(Fiddle::SIZEOF_DOUBLE)
+      status = FFI.mitie_classify_binary_relation(pointer, relation, score_ptr)
+      raise Error, "Bad status: #{status}" if status != 0
+
+      score = Utils.read_double(score_ptr)
+      if score > 0
+        {
+          first: entity1[:text],
+          second: entity2[:text],
+          score: score
+        }
+      end
+    ensure
+      FFI.mitie_free(relation) if relation
+    end
+
     def self.finalize(pointer)
       # must use proc instead of stabby lambda
       proc { FFI.mitie_free(pointer) }

data/lib/mitie/binary_relation_trainer.rb

@@ -0,0 +1,87 @@
+module Mitie
+  class BinaryRelationTrainer
+    def initialize(ner, name: "")
+      @pointer = FFI.mitie_create_binary_relation_trainer(name, ner.pointer)
+
+      ObjectSpace.define_finalizer(self, self.class.finalize(@pointer))
+    end
+
+    def add_positive_binary_relation(tokens, range1, range2)
+      check_add(tokens, range1, range2)
+
+      tokens_pointer = Utils.array_to_pointer(tokens)
+      status = FFI.mitie_add_positive_binary_relation(@pointer, tokens_pointer, range1.begin, range1.size, range2.begin, range2.size)
+      if status != 0
+        raise Error, "Unable to add binary relation"
+      end
+    end
+
+    def add_negative_binary_relation(tokens, range1, range2)
+      check_add(tokens, range1, range2)
+
+      tokens_pointer = Utils.array_to_pointer(tokens)
+      status = FFI.mitie_add_negative_binary_relation(@pointer, tokens_pointer, range1.begin, range1.size, range2.begin, range2.size)
+      if status != 0
+        raise Error, "Unable to add binary relation"
+      end
+    end
+
+    def beta
+      FFI.mitie_binary_relation_trainer_get_beta(@pointer)
+    end
+
+    def beta=(value)
+      raise ArgumentError, "beta must be greater than or equal to zero" unless value >= 0
+
+      FFI.mitie_binary_relation_trainer_set_beta(@pointer, value)
+    end
+
+    def num_threads
+      FFI.mitie_binary_relation_trainer_get_num_threads(@pointer)
+    end
+
+    def num_threads=(value)
+      FFI.mitie_binary_relation_trainer_set_num_threads(@pointer, value)
+    end
+
+    def num_positive_examples
+      FFI.mitie_binary_relation_trainer_num_positive_examples(@pointer)
+    end
+
+    def num_negative_examples
+      FFI.mitie_binary_relation_trainer_num_negative_examples(@pointer)
+    end
+
+    def train
+      if num_positive_examples + num_negative_examples == 0
+        raise Error, "You can't call train() on an empty trainer"
+      end
+
+      detector = FFI.mitie_train_binary_relation_detector(@pointer)
+
+      raise Error, "Unable to create binary relation detector. Probably ran out of RAM." if detector.null?
+
+      Mitie::BinaryRelationDetector.new(pointer: detector)
+    end
+
+    private
+
+    def check_add(tokens, range1, range2)
+      Utils.check_range(range1, tokens.size)
+      Utils.check_range(range2, tokens.size)
+
+      if entities_overlap?(range1, range2)
+        raise ArgumentError, "Entities overlap"
+      end
+    end
+
+    def entities_overlap?(range1, range2)
+      FFI.mitie_entities_overlap(range1.begin, range1.size, range2.begin, range2.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/document.rb

@@ -33,37 +33,35 @@ module Mitie
 
     def entities
       @entities ||= begin
-        begin
-          entities = []
-          tokens = tokens_with_offset
-          detections = FFI.mitie_extract_entities(pointer, tokens_ptr)
-          num_detections = FFI.mitie_ner_get_num_detections(detections)
-          num_detections.times do |i|
-            pos = FFI.mitie_ner_get_detection_position(detections, i)
-            len = FFI.mitie_ner_get_detection_length(detections, i)
-            tag = FFI.mitie_ner_get_detection_tagstr(detections, i).to_s
-            score = FFI.mitie_ner_get_detection_score(detections, i)
-            tok = tokens[pos, len]
-            offset = tok[0][1]
+        entities = []
+        tokens = tokens_with_offset
+        detections = FFI.mitie_extract_entities(pointer, tokens_ptr)
+        num_detections = FFI.mitie_ner_get_num_detections(detections)
+        num_detections.times do |i|
+          pos = FFI.mitie_ner_get_detection_position(detections, i)
+          len = FFI.mitie_ner_get_detection_length(detections, i)
+          tag = FFI.mitie_ner_get_detection_tagstr(detections, i).to_s
+          score = FFI.mitie_ner_get_detection_score(detections, i)
+          tok = tokens[pos, len]
+          offset = tok[0][1]
 
-            entity = {}
-            if offset
-              finish = tok[-1][1] + tok[-1][0].bytesize
-              entity[:text] = text.byteslice(offset...finish)
-            else
-              entity[:text] = tok.map(&:first)
-            end
-            entity[:tag] = tag
-            entity[:score] = score
-            entity[:offset] = offset if offset
-            entity[:token_index] = pos
-            entity[:token_length] = len
-            entities << entity
+          entity = {}
+          if offset
+            finish = tok[-1][1] + tok[-1][0].bytesize
+            entity[:text] = text.byteslice(offset...finish)
+          else
+            entity[:text] = tok.map(&:first)
           end
-          entities
-        ensure
-          FFI.mitie_free(detections) if detections
+          entity[:tag] = tag
+          entity[:score] = score
+          entity[:offset] = offset if offset
+          entity[:token_index] = pos
+          entity[:token_length] = len
+          entities << entity
         end
+        entities
+      ensure
+        FFI.mitie_free(detections) if detections
       end
     end
 
@@ -84,11 +82,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,14 +10,17 @@ 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_file(const char* filename)"
     extern "char** mitie_tokenize_with_offsets(const char* text, unsigned long** token_offsets)"
 
+    # 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)"
     extern "const char* mitie_get_named_entity_tagstr(const mitie_named_entity_extractor* ner, unsigned long idx)"
-
     extern "mitie_named_entity_detections* mitie_extract_entities(const mitie_named_entity_extractor* ner, char** tokens)"
     extern "unsigned long mitie_ner_get_num_detections(const mitie_named_entity_detections* dets)"
     extern "unsigned long mitie_ner_get_detection_position(const mitie_named_entity_detections* dets, unsigned long idx)"
@@ -26,10 +29,57 @@ 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)"
 
+    # binary relation detector
     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)"
     extern "mitie_binary_relation* mitie_extract_binary_relation(const mitie_named_entity_extractor* ner, char** tokens, unsigned long arg1_start, unsigned long arg1_length, unsigned long arg2_start, unsigned long arg2_length)"
     extern "int mitie_classify_binary_relation(const mitie_binary_relation_detector* detector, const mitie_binary_relation* relation, double* score)"
+
+    # text categorizer
+    extern "mitie_text_categorizer* mitie_load_text_categorizer(const char* filename)"
+    extern "int mitie_categorize_text(const mitie_text_categorizer* tcat, const char** tokens, char** text_tag, double* text_score)"
+
+    # save
+    extern "int mitie_save_named_entity_extractor(const char* filename, const mitie_named_entity_extractor* ner)"
+    extern "int mitie_save_binary_relation_detector(const char* filename, const mitie_binary_relation_detector* detector)"
+    extern "int mitie_save_text_categorizer(const char* filename, const mitie_text_categorizer* tcat)"
+
+    # ner trainer
+    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)"
+
+    # binary relation trainer
+    extern "mitie_binary_relation_trainer* mitie_create_binary_relation_trainer(const char* relation_name, const mitie_named_entity_extractor* ner)"
+    extern "unsigned long mitie_binary_relation_trainer_num_positive_examples(const mitie_binary_relation_trainer* trainer)"
+    extern "unsigned long mitie_binary_relation_trainer_num_negative_examples(const mitie_binary_relation_trainer* trainer)"
+    extern "int mitie_add_positive_binary_relation(mitie_binary_relation_trainer* trainer, char** tokens, unsigned long arg1_start, unsigned long arg1_length, unsigned long arg2_start, unsigned long arg2_length)"
+    extern "int mitie_add_negative_binary_relation(mitie_binary_relation_trainer* trainer, char** tokens, unsigned long arg1_start, unsigned long arg1_length, unsigned long arg2_start, unsigned long arg2_length)"
+    extern "void mitie_binary_relation_trainer_set_beta(mitie_binary_relation_trainer* trainer, double beta)"
+    extern "double mitie_binary_relation_trainer_get_beta(const mitie_binary_relation_trainer* trainer)"
+    extern "void mitie_binary_relation_trainer_set_num_threads(mitie_binary_relation_trainer* trainer, unsigned long num_threads)"
+    extern "unsigned long mitie_binary_relation_trainer_get_num_threads(const mitie_binary_relation_trainer* trainer)"
+    extern "mitie_binary_relation_detector* mitie_train_binary_relation_detector(const mitie_binary_relation_trainer* trainer)"
+
+    # text categorizer trainer
+    extern "mitie_text_categorizer_trainer* mitie_create_text_categorizer_trainer(const char* filename)"
+    extern "unsigned long mitie_text_categorizer_trainer_size(const mitie_text_categorizer_trainer* trainer)"
+    extern "void mitie_text_categorizer_trainer_set_beta(mitie_text_categorizer_trainer* trainer, double beta)"
+    extern "double mitie_text_categorizer_trainer_get_beta(const mitie_text_categorizer_trainer* trainer)"
+    extern "void mitie_text_categorizer_trainer_set_num_threads(mitie_text_categorizer_trainer* trainer, unsigned long num_threads)"
+    extern "unsigned long mitie_text_categorizer_trainer_get_num_threads(const mitie_text_categorizer_trainer* trainer)"
+    extern "int mitie_add_text_categorizer_labeled_text(mitie_text_categorizer_trainer* trainer, const char** tokens, const char* label)"
+    extern "mitie_text_categorizer* mitie_train_text_categorizer(const mitie_text_categorizer_trainer* trainer)"
   end
 end

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,45 @@
+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)
+      Utils.check_range(range, num_tokens)
+
+      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)
+      Utils.check_range(range, num_tokens)
+
+      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/text_categorizer.rb

@@ -0,0 +1,48 @@
+module Mitie
+  class TextCategorizer
+    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_text_categorizer(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 categorize(text)
+      tokens = text.is_a?(Array) ? text : Mitie.tokenize(text)
+      tokens_pointer = Utils.array_to_pointer(tokens)
+      text_tag = Fiddle::Pointer.malloc(Fiddle::SIZEOF_VOIDP)
+      text_score = Fiddle::Pointer.malloc(Fiddle::SIZEOF_DOUBLE)
+
+      if FFI.mitie_categorize_text(@pointer, tokens_pointer, text_tag, text_score) != 0
+        raise Error, "Unable to categorize"
+      end
+
+      {
+        tag: text_tag.ptr.to_s,
+        score: Utils.read_double(text_score)
+      }
+    ensure
+      # text_tag must be freed
+      FFI.mitie_free(text_tag.ptr) if text_tag
+    end
+
+    def save_to_disk(filename)
+      if FFI.mitie_save_text_categorizer(filename, @pointer) != 0
+        raise Error, "Unable to save model"
+      end
+      nil
+    end
+
+    def self.finalize(pointer)
+      # must use proc instead of stabby lambda
+      proc { FFI.mitie_free(pointer) }
+    end
+  end
+end

data/lib/mitie/text_categorizer_trainer.rb

@@ -0,0 +1,53 @@
+module Mitie
+  class TextCategorizerTrainer
+    def initialize(filename)
+      raise ArgumentError, "File does not exist" unless File.exist?(filename)
+      @pointer = FFI.mitie_create_text_categorizer_trainer(filename)
+
+      ObjectSpace.define_finalizer(self, self.class.finalize(@pointer))
+    end
+
+    def add(text, label)
+      tokens = text.is_a?(Array) ? text : Mitie.tokenize(text)
+      tokens_pointer = Utils.array_to_pointer(tokens)
+      FFI.mitie_add_text_categorizer_labeled_text(@pointer, tokens_pointer, label)
+    end
+
+    def beta
+      FFI.mitie_text_categorizer_trainer_get_beta(@pointer)
+    end
+
+    def beta=(value)
+      raise ArgumentError, "beta must be greater than or equal to zero" unless value >= 0
+
+      FFI.mitie_text_categorizer_trainer_set_beta(@pointer, value)
+    end
+
+    def num_threads
+      FFI.mitie_text_categorizer_trainer_get_num_threads(@pointer)
+    end
+
+    def num_threads=(value)
+      FFI.mitie_text_categorizer_trainer_set_num_threads(@pointer, value)
+    end
+
+    def size
+      FFI.mitie_text_categorizer_trainer_size(@pointer)
+    end
+
+    def train
+      raise Error, "You can't call train() on an empty trainer" if size.zero?
+
+      categorizer = FFI.mitie_train_text_categorizer(@pointer)
+
+      raise Error, "Unable to create text categorizer. Probably ran out of RAM." if categorizer.null?
+
+      Mitie::TextCategorizer.new(pointer: categorizer)
+    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,22 @@
+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
+
+    def self.check_range(range, num_tokens)
+      if range.none? || !(0..(num_tokens - 1)).cover?(range)
+        raise ArgumentError, "Invalid range"
+      end
+    end
+
+    def self.read_double(ptr)
+      ptr.to_s(Fiddle::SIZEOF_DOUBLE).unpack1("d")
+    end
+  end
+end

data/lib/mitie/version.rb

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

data/lib/mitie.rb

@@ -3,8 +3,14 @@ require "fiddle/import"
 
 # modules
 require "mitie/binary_relation_detector"
+require "mitie/binary_relation_trainer"
 require "mitie/document"
 require "mitie/ner"
+require "mitie/ner_training_instance"
+require "mitie/ner_trainer"
+require "mitie/text_categorizer"
+require "mitie/text_categorizer_trainer"
+require "mitie/utils"
 require "mitie/version"
 
 module Mitie
@@ -16,10 +22,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
@@ -28,4 +36,36 @@ module Mitie
 
   # friendlier error message
   autoload :FFI, "mitie/ffi"
+
+  class << self
+    def tokenize(text)
+      tokens_ptr = FFI.mitie_tokenize(text)
+      tokens = read_tokens(tokens_ptr)
+      tokens.each { |t| t.force_encoding(text.encoding) }
+      tokens
+    ensure
+      FFI.mitie_free(tokens_ptr) if tokens_ptr
+    end
+
+    def tokenize_file(filename)
+      tokens_ptr = FFI.mitie_tokenize_file(filename)
+      read_tokens(tokens_ptr)
+    ensure
+      FFI.mitie_free(tokens_ptr) if tokens_ptr
+    end
+
+    private
+
+    def read_tokens(tokens_ptr)
+      i = 0
+      tokens = []
+      loop do
+        token = (tokens_ptr + i * Fiddle::SIZEOF_VOIDP).ptr
+        break if token.null?
+        tokens << token.to_s
+        i += 1
+      end
+      tokens
+    end
+  end
 end

metadata

@@ -1,17 +1,17 @@
 --- !ruby/object:Gem::Specification
 name: mitie
 version: !ruby/object:Gem::Version
-  version: 0.1.5
+  version: 0.2.1
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2021-01-30 00:00:00.000000000 Z
+date: 2022-06-12 00:00:00.000000000 Z
 dependencies: []
 description:
-email: andrew@chartkick.com
+email: andrew@ankane.org
 executables: []
 extensions: []
 extra_rdoc_files: []
@@ -21,16 +21,22 @@ files:
 - README.md
 - lib/mitie.rb
 - lib/mitie/binary_relation_detector.rb
+- lib/mitie/binary_relation_trainer.rb
 - 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/text_categorizer.rb
+- lib/mitie/text_categorizer_trainer.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: {}
@@ -42,14 +48,14 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: '2.5'
+      version: '2.7'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !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