transformers-rb 0.1.0

data/lib/transformers/pipelines/token_classification.rb

@@ -0,0 +1,282 @@
+module Transformers
+  class TokenClassificationArgumentHandler < ArgumentHandler
+  end
+
+  class AggregationStrategy < ExplicitEnum
+    NONE = "none"
+    SIMPLE = "simple"
+    FIRST = "first"
+    AVERAGE = "average"
+    MAX = "max"
+  end
+
+  class TokenClassificationPipeline < ChunkPipeline
+    extend ClassAttribute
+
+    class_attribute :default_input_names, "sequences"
+
+    def initialize(*args, args_parser: TokenClassificationArgumentHandler.new, **kwargs)
+      super(*args, **kwargs)
+      check_model_type(MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES)
+
+      @basic_tokenizer = Bert::BertTokenizer::BasicTokenizer.new(do_lower_case: false)
+      @args_parser = args_parser
+    end
+
+    def _sanitize_parameters(
+      ignore_labels: nil,
+      grouped_entities: nil,
+      ignore_subwords: nil,
+      aggregation_strategy: nil,
+      offset_mapping: nil,
+      stride: nil
+    )
+      preprocess_params = {}
+      if !offset_mapping.nil?
+        preprocess_params[:offset_mapping] = offset_mapping
+      end
+
+      postprocess_params = {}
+      if !grouped_entities.nil? || !ignore_subwords.nil?
+        if grouped_entities && ignore_subwords
+          aggregation_strategy = AggregationStrategy::FIRST
+        elsif grouped_entities && !ignore_subwords
+          aggregation_strategy = AggregationStrategy::SIMPLE
+        else
+          aggregation_strategy = AggregationStrategy::NONE
+        end
+
+        if !grouped_entities.nil?
+          warn(
+            "`grouped_entities` is deprecated and will be removed in version v5.0.0, defaulted to" +
+            " `aggregation_strategy=\"#{aggregation_strategy}\"` instead."
+          )
+        end
+        if !ignore_subwords.nil?
+          warn(
+            "`ignore_subwords` is deprecated and will be removed in version v5.0.0, defaulted to" +
+            " `aggregation_strategy=\"#{aggregation_strategy}\"` instead."
+          )
+        end
+      end
+
+      if !aggregation_strategy.nil?
+        if aggregation_strategy.is_a?(String)
+          aggregation_strategy = AggregationStrategy.new(aggregation_strategy.upcase).to_s
+        end
+        if (
+          [AggregationStrategy::FIRST, AggregationStrategy::MAX, AggregationStrategy::AVERAGE].include?(aggregation_strategy) &&
+          !@tokenizer.is_fast
+        )
+          raise ArgumentError,
+            "Slow tokenizers cannot handle subwords. Please set the `aggregation_strategy` option" +
+            ' to `"simple"` or use a fast tokenizer.'
+        end
+        postprocess_params[:aggregation_strategy] = aggregation_strategy
+      end
+      if !ignore_labels.nil?
+        postprocess_params[:ignore_labels] = ignore_labels
+      end
+      if !stride.nil?
+        if stride >= @tokenizer.model_max_length
+          raise ArgumentError,
+            "`stride` must be less than `tokenizer.model_max_length` (or even lower if the tokenizer adds special tokens)"
+        end
+        if aggregation_strategy == AggregationStrategy::NONE
+          raise ArgumentError,
+            "`stride` was provided to process all the text but `aggregation_strategy=" +
+            "\"#{aggregation_strategy}\"`, please select another one instead."
+        else
+          if @tokenizer.is_fast
+            tokenizer_params = {
+              return_overflowing_tokens: true,
+              padding: true,
+              stride: stride
+            }
+            preprocess_params[:tokenizer_params] = tokenizer_params
+          else
+            raise ArgumentError,
+              "`stride` was provided to process all the text but you're using a slow tokenizer." +
+              " Please use a fast tokenizer."
+          end
+        end
+      end
+      [preprocess_params, {}, postprocess_params]
+    end
+
+    def preprocess(sentence, offset_mapping: nil, **preprocess_params)
+      tokenizer_params = preprocess_params.delete(:tokenizer_params) { {} }
+      truncation = @tokenizer.model_max_length && @tokenizer.model_max_length > 0
+      inputs = @tokenizer.(
+        sentence,
+        return_tensors: @framework,
+        truncation: truncation,
+        return_special_tokens_mask: true,
+        return_offsets_mapping: @tokenizer.is_fast,
+        **tokenizer_params
+      )
+      inputs.delete(:overflow_to_sample_mapping)
+      num_chunks = inputs[:input_ids].length
+
+      num_chunks.times do |i|
+        if @framework == "tf"
+          raise Todo
+        else
+          model_inputs = inputs.to_h { |k, v| [k, v[i].unsqueeze(0)] }
+        end
+        if !@offset_mapping.nil?
+          model_inputs[:offset_mapping] = offset_mapping
+        end
+        model_inputs[:sentence] = i == 0 ? sentence : nil
+        model_inputs[:is_last] = (i == num_chunks - 1)
+
+        yield model_inputs
+      end
+    end
+
+    def _forward(model_inputs)
+      # Forward
+      special_tokens_mask = model_inputs.delete(:special_tokens_mask)
+      offset_mapping = model_inputs.delete(:offset_mapping)
+      sentence = model_inputs.delete(:sentence)
+      is_last = model_inputs.delete(:is_last)
+      if @framework == "tf"
+        logits = @model.(**model_inputs)[0]
+      else
+        output = @model.(**model_inputs)
+        logits = output.is_a?(Hash) ? output[:logits] : output[0]
+      end
+
+      {
+        logits: logits,
+        special_tokens_mask: special_tokens_mask,
+        offset_mapping: offset_mapping,
+        sentence: sentence,
+        is_last: is_last,
+        **model_inputs
+      }
+    end
+
+    def postprocess(all_outputs, aggregation_strategy: AggregationStrategy::NONE, ignore_labels: nil)
+      if ignore_labels.nil?
+        ignore_labels = ["O"]
+      end
+      all_entities = []
+      all_outputs.each do |model_outputs|
+        logits = model_outputs[:logits][0].numo
+        sentence = all_outputs[0][:sentence]
+        input_ids = model_outputs[:input_ids][0]
+        offset_mapping = (
+          !model_outputs[:offset_mapping].nil? ? model_outputs[:offset_mapping][0] : nil
+        )
+        special_tokens_mask = model_outputs[:special_tokens_mask][0].numo
+
+        maxes = logits.max(axis: -1).expand_dims(-1)
+        shifted_exp = Numo::NMath.exp(logits - maxes)
+        scores = shifted_exp / shifted_exp.sum(axis: -1).expand_dims(-1)
+
+        if @framework == "tf"
+          raise Todo
+        end
+
+        pre_entities = gather_pre_entities(
+          sentence, input_ids, scores, offset_mapping, special_tokens_mask, aggregation_strategy
+        )
+        grouped_entities = aggregate(pre_entities, aggregation_strategy)
+        # Filter anything that is in self.ignore_labels
+        entities =
+          grouped_entities.select do |entity|
+            !ignore_labels.include?(entity[:entity]) && !ignore_labels.include?(entity[:entity_group])
+          end
+        all_entities.concat(entities)
+      end
+      num_chunks = all_outputs.length
+      if num_chunks > 1
+        all_entities = aggregate_overlapping_entities(all_entities)
+      end
+      all_entities
+    end
+
+    def gather_pre_entities(
+      sentence,
+      input_ids,
+      scores,
+      offset_mapping,
+      special_tokens_mask,
+      aggregation_strategy
+    )
+      pre_entities = []
+      scores.each_over_axis(0).with_index do |token_scores, idx|
+        # Filter special_tokens
+        if special_tokens_mask[idx] != 0
+          next
+        end
+
+        word = @tokenizer.convert_ids_to_tokens(input_ids[idx].to_i)
+        if !offset_mapping.nil?
+          start_ind, end_ind = offset_mapping[idx].to_a
+          if !start_ind.is_a?(Integer)
+            if @framework == "pt"
+              start_ind = start_ind.item
+              end_ind = end_ind.item
+            end
+          end
+          word_ref = sentence[start_ind...end_ind]
+          if @tokenizer.instance_variable_get(:@tokenizer).respond_to?(:continuing_subword_prefix)
+            # This is a BPE, word aware tokenizer, there is a correct way
+            # to fuse tokens
+            is_subword = word.length != word_ref.length
+          else
+            is_subword = start_ind > 0 && !sentence[(start_ind - 1)...(start_ind + 1)].include?(" ")
+          end
+
+          if input_ids[idx].to_i == @tokenizer.unk_token_id
+            word = word_ref
+            is_subword = false
+          end
+        else
+          start_ind = nil
+          end_ind = nil
+          is_subword = nil
+        end
+
+        pre_entity = {
+          word: word,
+          scores: token_scores,
+          start: start_ind,
+          end: end_ind,
+          index: idx,
+          is_subword: is_subword
+        }
+        pre_entities << pre_entity
+      end
+      pre_entities
+    end
+
+    def aggregate(pre_entities, aggregation_strategy)
+      if [AggregationStrategy::NONE, AggregationStrategy::SIMPLE].include?(aggregation_strategy)
+        entities = []
+        pre_entities.each do |pre_entity|
+          entity_idx = pre_entity[:scores].argmax
+          score = pre_entity[:scores][entity_idx]
+          entity = {
+            entity: @model.config.id2label[entity_idx],
+            score: score,
+            index: pre_entity[:index],
+            word: pre_entity[:word],
+            start: pre_entity[:start],
+            end: pre_entity[:end]
+          }
+          entities << entity
+        end
+      else
+        entities = aggregate_words(pre_entities, aggregation_strategy)
+      end
+
+      if aggregation_strategy == AggregationStrategy::NONE
+        return entities
+      end
+      group_entities(entities)
+    end
+  end
+end

data/lib/transformers/ruby_utils.rb

@@ -0,0 +1,33 @@
+module Transformers
+  module ClassAttribute
+    def class_attribute(name, default = nil)
+      singleton_class.attr_writer name
+      var = "@#{name}"
+      instance_variable_set(var, default)
+      singleton_class.define_method(name) do
+        # ancestors includes current module
+        ancestors.find { |c| c.instance_variable_defined?(var) }.instance_variable_get(var)
+      end
+      define_method(name) do
+        self.class.send(name)
+      end
+    end
+  end
+
+  module Copy
+    def self.deepcopy(value, memo = {})
+      key = value.object_id
+      if !memo.key?(key)
+        copy = value.dup
+        memo[key] = copy
+        if value.is_a?(Hash)
+          copy.transform_keys! { |k| deepcopy(k, memo) }
+          copy.transform_values! { |v| deepcopy(v, memo) }
+        elsif value.is_a?(Array)
+          copy.map! { |v| deepcopy(v, memo) }
+        end
+      end
+      memo[key]
+    end
+  end
+end

data/lib/transformers/sentence_transformer.rb

@@ -0,0 +1,37 @@
+module Transformers
+  class SentenceTransformer
+    def initialize(model_id)
+      @model_id = model_id
+      @tokenizer = Transformers::AutoTokenizer.from_pretrained(model_id)
+      @model = Transformers::AutoModel.from_pretrained(model_id)
+    end
+
+    def encode(sentences)
+      singular = sentences.is_a?(String)
+      sentences = [sentences] if singular
+
+      input = @tokenizer.(sentences, padding: true, truncation: true, return_tensors: "pt")
+      output = Torch.no_grad { @model.(**input) }[0]
+
+      # TODO check modules.json
+      if [
+        "sentence-transformers/all-MiniLM-L6-v2",
+        "sentence-transformers/multi-qa-MiniLM-L6-cos-v1"
+      ].include?(@model_id)
+        output = mean_pooling(output, input[:attention_mask])
+        output = Torch::NN::Functional.normalize(output, p: 2, dim: 1).to_a
+      else
+        output = output[0.., 0].to_a
+      end
+
+      singular ? output[0] : output
+    end
+
+    private
+
+    def mean_pooling(output, attention_mask)
+      input_mask_expanded = attention_mask.unsqueeze(-1).expand(output.size).float
+      Torch.sum(output * input_mask_expanded, 1) / Torch.clamp(input_mask_expanded.sum(1), min: 1e-9)
+    end
+  end
+end

data/lib/transformers/tokenization_utils.rb

@@ -0,0 +1,152 @@
+# Copyright 2020 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+module Transformers
+  class PreTrainedTokenizer < PreTrainedTokenizerBase
+    def initialize(**kwargs)
+
+      # 2. init `_added_tokens_decoder` if child class did not
+      if !instance_variable_defined?(:@added_tokens_decoder)
+        @added_tokens_decoder = {}
+      end
+
+      # 3. if a `added_tokens_decoder` is passed, we are loading from a saved tokenizer, we overwrite
+      @added_tokens_decoder.merge!(kwargs.delete(:added_tokens_decoder) { {} })
+      @added_tokens_encoder = @added_tokens_decoder.to_h { |k, v| [k.content, v] }
+
+      # 4 init the parent class
+      super(**kwargs)
+    end
+
+    def is_fast
+      false
+    end
+
+    def vocab_size
+      raise NotImplementedError
+    end
+
+    def tokenize(text, **kwargs)
+      raise Todo
+    end
+
+    def _encode_plus(
+      text:,
+      text_pair: nil,
+      add_special_tokens: true,
+      padding_strategy: PaddingStrategy::DO_NOT_PAD,
+      truncation_strategy: TruncationStrategy::DO_NOT_TRUNCATE,
+      max_length: nil,
+      stride: 0,
+      is_split_into_words: false,
+      pad_to_multiple_of: nil,
+      return_tensors: nil,
+      return_token_type_ids: nil,
+      return_attention_mask: nil,
+      return_overflowing_tokens: false,
+      return_special_tokens_mask: false,
+      return_offsets_mapping: false,
+      return_length: false,
+      verbose: true,
+      **kwargs
+    )
+      get_input_ids = lambda do |text|
+        if text.is_a?(String)
+          tokens = tokenize(text, **kwargs)
+          convert_tokens_to_ids(tokens)
+        elsif text.is_a?(Array) && text.length > 0 && text[0].is_a?(String)
+          if is_split_into_words
+            raise Todo
+          else
+            convert_tokens_to_ids(text)
+          end
+        elsif text.is_a?(Array) && text.length > 0 && text[0].is_a?(Integer)
+          text
+        else
+          if is_split_into_words
+            raise ArgumentError,
+              "Input #{text} is not valid. Should be a string or a list/tuple of strings when" +
+              " `is_split_into_words=True`."
+          else
+            raise ArgumentError,
+              "Input #{text} is not valid. Should be a string, a list/tuple of strings or a list/tuple of" +
+              " integers."
+          end
+        end
+      end
+
+      if return_offsets_mapping
+        raise RuntimeError,
+          "return_offset_mapping is not available when using Ruby tokenizers. " +
+          "To use this feature, change your tokenizer to one deriving from " +
+          "Transformers::PreTrainedTokenizerFast. " +
+          "More information on available tokenizers at " +
+          "https://github.com/huggingface/transformers/pull/2674"
+      end
+
+      first_ids = get_input_ids.(text)
+      second_ids = !text_pair.nil? ? get_input_ids.(text_pair) : nil
+
+      prepare_for_model(
+        first_ids,
+        pair_ids: second_ids,
+        add_special_tokens: add_special_tokens,
+        padding: padding_strategy,
+        truncation: truncation_strategy,
+        max_length: max_length,
+        stride: stride,
+        pad_to_multiple_of: pad_to_multiple_of,
+        return_tensors: return_tensors,
+        prepend_batch_axis: true,
+        return_attention_mask: return_attention_mask,
+        return_token_type_ids: return_token_type_ids,
+        return_overflowing_tokens: return_overflowing_tokens,
+        return_special_tokens_mask: return_special_tokens_mask,
+        return_length: return_length,
+        verbose: verbose
+      )
+    end
+
+    def convert_tokens_to_ids(tokens)
+      if tokens.nil?
+        return nil
+      end
+
+      if tokens.is_a?(String)
+        return _convert_token_to_id_with_added_voc(tokens)
+      end
+
+      ids = []
+      tokens.each do |token|
+        ids << _convert_token_to_id_with_added_voc(token)
+      end
+      ids
+    end
+
+    def _convert_token_to_id_with_added_voc(token)
+      if token.nil?
+        return nil
+      end
+
+      if @added_tokens_encoder.include?(token)
+        return @added_tokens_encoder[token]
+      end
+      _convert_token_to_id(token)
+    end
+
+    def _convert_token_to_id(token)
+      raise NotImplementedError
+    end
+  end
+end