transformers-rb 0.1.0

data/lib/transformers/pipelines/_init.rb

@@ -0,0 +1,348 @@
+# Copyright 2018 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
+  TASK_ALIASES = {
+    "sentiment-analysis" => "text-classification",
+    "ner" => "token-classification"
+  }
+
+  SUPPORTED_TASKS = {
+    "feature-extraction" => {
+      "impl" => FeatureExtractionPipeline,
+      "pt" => [AutoModel],
+      "default" => {
+        "model" => {
+          "pt" => ["distilbert/distilbert-base-cased", "935ac13"]
+        }
+      },
+      "type" => "multimodal"
+    },
+    "text-classification" => {
+      "impl" => TextClassificationPipeline,
+      "pt" => [AutoModelForSequenceClassification],
+      "default" => {
+        "model" => {
+          "pt" => ["distilbert/distilbert-base-uncased-finetuned-sst-2-english", "af0f99b"]
+        }
+      },
+      "type" => "text"
+    },
+    "token-classification" => {
+      "impl" => TokenClassificationPipeline,
+      "pt" => [AutoModelForTokenClassification],
+      "default" => {
+        "model" => {
+          "pt" => ["dbmdz/bert-large-cased-finetuned-conll03-english", "f2482bf"]
+        }
+      },
+      "type" => "text"
+    },
+    "question-answering" => {
+      "impl" => QuestionAnsweringPipeline,
+      "pt" => [AutoModelForQuestionAnswering],
+      "default" => {
+        "model" => {
+          "pt" => ["distilbert/distilbert-base-cased-distilled-squad", "626af31"]
+        }
+      },
+      "type" => "text"
+    },
+    "image-classification" => {
+      "impl" => ImageClassificationPipeline,
+      "pt" => [AutoModelForImageClassification],
+      "default" => {
+        "model" => {
+          "pt" => ["google/vit-base-patch16-224", "5dca96d"]
+        }
+      },
+      "type" => "image"
+    },
+    "image-feature-extraction" => {
+      "impl" => ImageFeatureExtractionPipeline,
+      "pt" => [AutoModel],
+      "default" => {
+        "model" => {
+          "pt" => ["google/vit-base-patch16-224", "3f49326"]
+        }
+      },
+      "type" => "image"
+    }
+  }
+
+  PIPELINE_REGISTRY = PipelineRegistry.new(supported_tasks: SUPPORTED_TASKS, task_aliases: TASK_ALIASES)
+
+  class << self
+    def pipeline(
+      task,
+      model: nil,
+      config: nil,
+      tokenizer: nil,
+      feature_extractor: nil,
+      image_processor: nil,
+      framework: nil,
+      revision: nil,
+      use_fast: true,
+      token: nil,
+      device: nil,
+      device_map: nil,
+      torch_dtype: nil,
+      trust_remote_code: nil,
+      model_kwargs: nil,
+      pipeline_class: nil,
+      **kwargs
+    )
+      model_kwargs ||= {}
+      # Make sure we only pass use_auth_token once as a kwarg (it used to be possible to pass it in model_kwargs,
+      # this is to keep BC).
+      use_auth_token = model_kwargs.delete(:use_auth_token)
+      if !use_auth_token.nil?
+        raise Todo
+      end
+
+      code_revision = kwargs.delete(:code_revision)
+      commit_hash = kwargs.delete(:_commit_hash)
+
+      hub_kwargs = {
+        revision: revision,
+        token: token,
+        trust_remote_code: trust_remote_code,
+        _commit_hash: commit_hash
+      }
+
+      if task.nil? && model.nil?
+        raise RuntimeError,
+          "Impossible to instantiate a pipeline without either a task or a model " +
+          "being specified. " +
+          "Please provide a task class or a model"
+      end
+
+      if model.nil? && !tokenizer.nil?
+        raise RuntimeError,
+          "Impossible to instantiate a pipeline with tokenizer specified but not the model as the provided tokenizer" +
+          " may not be compatible with the default model. Please provide a PreTrainedModel class or a" +
+          " path/identifier to a pretrained model when providing tokenizer."
+      end
+      if model.nil? && !feature_extractor.nil?
+        raise RuntimeError,
+          "Impossible to instantiate a pipeline with feature_extractor specified but not the model as the provided" +
+          " feature_extractor may not be compatible with the default model. Please provide a PreTrainedModel class" +
+          " or a path/identifier to a pretrained model when providing feature_extractor."
+      end
+      if model.is_a?(Pathname)
+        model = model.to_s
+      end
+
+      if commit_hash.nil?
+        pretrained_model_name_or_path = nil
+        if config.is_a?(String)
+          pretrained_model_name_or_path = config
+        elsif config.nil? && model.is_a?(String)
+          pretrained_model_name_or_path = model
+        end
+
+        if !config.is_a?(PretrainedConfig) && !pretrained_model_name_or_path.nil?
+          # We make a call to the config file first (which may be absent) to get the commit hash as soon as possible
+          resolved_config_file = Utils::Hub.cached_file(
+            pretrained_model_name_or_path,
+            CONFIG_NAME,
+            _raise_exceptions_for_gated_repo: false,
+            _raise_exceptions_for_missing_entries: false,
+            _raise_exceptions_for_connection_errors: false,
+            cache_dir: model_kwargs[:cache_dir],
+            **hub_kwargs
+          )
+          hub_kwargs[:_commit_hash] = Utils::Hub.extract_commit_hash(resolved_config_file, commit_hash)
+        else
+          hub_kwargs[:_commit_hash] = nil # getattr(config, "_commit_hash", None)
+        end
+      end
+
+      # Config is the primordial information item.
+      # Instantiate config if needed
+      if config.is_a?(String)
+        raise Todo
+      elsif config.nil? && model.is_a?(String)
+        config = AutoConfig.from_pretrained(
+          model, _from_pipeline: task, code_revision: code_revision, **hub_kwargs, **model_kwargs
+        )
+        hub_kwargs[:_commit_hash] = config._commit_hash
+      end
+
+      custom_tasks = {}
+      if !config.nil? && (config.instance_variable_get(:@custom_pipelines) || {}).length > 0
+        raise Todo
+      end
+
+      if task.nil? && !model.nil?
+        raise Todo
+      end
+
+      # Retrieve the task
+      if custom_tasks.include?(task)
+        raise Todo
+      else
+        _normalized_task, targeted_task, task_options = check_task(task)
+        if pipeline_class.nil?
+          pipeline_class = targeted_task["impl"]
+        end
+      end
+
+      # Use default model/config/tokenizer for the task if no model is provided
+      if model.nil?
+        # At that point framework might still be undetermined
+        model, default_revision = Pipelines.get_default_model_and_revision(targeted_task, framework, task_options)
+        revision = !revision.nil? ? revision : default_revision
+        Transformers.logger.warn(
+          "No model was supplied, defaulted to #{model} and revision" +
+          " #{revision} (#{Utils::Hub::HUGGINGFACE_CO_RESOLVE_ENDPOINT}/#{model}).\n" +
+          "Using a pipeline without specifying a model name and revision in production is not recommended."
+        )
+        if config.nil? && model.is_a?(String)
+          config = AutoConfig.from_pretrained(model, _from_pipeline: task, **hub_kwargs, **model_kwargs)
+          hub_kwargs[:_commit_hash] = config._commit_hash
+        end
+      end
+
+      if !device_map.nil?
+        raise Todo
+      end
+      if !torch_dtype.nil?
+        raise Todo
+      end
+
+      model_name = model.is_a?(String) ? model : nil
+
+      # Load the correct model if possible
+      # Infer the framework from the model if not already defined
+      if model.is_a?(String) || framework.nil?
+        model_classes = {"tf" => targeted_task["tf"], "pt" => targeted_task["pt"]}
+        framework, model =
+          Pipelines.infer_framework_load_model(
+            model,
+            config,
+            model_classes: model_classes,
+            framework: framework,
+            task: task,
+            **hub_kwargs,
+            **model_kwargs
+          )
+      end
+
+      model_config = model.config
+      hub_kwargs[:_commit_hash] = model.config._commit_hash
+      model_config_type = model_config.class.name.split("::").last
+      load_tokenizer = TOKENIZER_MAPPING.include?(model_config_type) || !model_config.tokenizer_class.nil?
+      load_feature_extractor = FEATURE_EXTRACTOR_MAPPING.include?(model_config_type) || !feature_extractor.nil?
+      load_image_processor = IMAGE_PROCESSOR_MAPPING.include?(model_config_type) || !image_processor.nil?
+
+      if load_tokenizer
+        # Try to infer tokenizer from model or config name (if provided as str)
+        if tokenizer.nil?
+          if model_name.is_a?(String)
+            tokenizer = model_name
+          elsif config.is_a?(String)
+            tokenizer = config
+          else
+            # Impossible to guess what is the right tokenizer here
+            raise "Impossible to guess which tokenizer to use. Please provide a PreTrainedTokenizer class or a path/identifier to a pretrained tokenizer."
+          end
+        end
+
+        # Instantiate tokenizer if needed
+        if tokenizer.is_a?(String) || tokenizer.is_a?(Array)
+          if tokenizer.is_a?(Array)
+            # For array we have [tokenizer name, {kwargs}]
+            use_fast = tokenizer[1].delete(:use_fast) { use_fast }
+            tokenizer_identifier = tokenizer[0]
+            tokenizer_kwargs = tokenizer[1]
+          else
+            tokenizer_identifier = tokenizer
+            tokenizer_kwargs = model_kwargs.dup
+            tokenizer_kwargs.delete(:torch_dtype)
+          end
+
+          tokenizer =
+            AutoTokenizer.from_pretrained(
+              tokenizer_identifier, use_fast: use_fast, _from_pipeline: task, **hub_kwargs, **tokenizer_kwargs
+            )
+        end
+      end
+
+      if load_image_processor
+        # Try to infer image processor from model or config name (if provided as str)
+        if image_processor.nil?
+          if model_name.is_a?(String)
+            image_processor = model_name
+          elsif config.is_a?(String)
+            image_processor = config
+          # Backward compatibility, as `feature_extractor` used to be the name
+          # for `ImageProcessor`.
+          elsif !feature_extractor.nil? && feature_extractor.is_a?(BaseImageProcessor)
+            image_processor = feature_extractor
+          else
+            # Impossible to guess what is the right image_processor here
+            raise RuntimeError,
+              "Impossible to guess which image processor to use. " +
+              "Please provide a PreTrainedImageProcessor class or a path/identifier " +
+              "to a pretrained image processor."
+          end
+        end
+
+        # Instantiate image_processor if needed
+        if image_processor.is_a?(String) || image_processor.is_a?(Array)
+          image_processor = AutoImageProcessor.from_pretrained(
+            image_processor, _from_pipeline: task, **hub_kwargs, **model_kwargs
+          )
+        end
+      end
+
+      if load_feature_extractor
+        raise Todo
+      end
+
+      if task == "translation" && model.config.task_specific_params
+        raise Todo
+      end
+
+      if !tokenizer.nil?
+        kwargs[:tokenizer] = tokenizer
+      end
+
+      if !feature_extractor.nil?
+        kwargs[:feature_extractor] = feature_extractor
+      end
+
+      if !torch_dtype.nil?
+        kwargs[:torch_dtype] = torch_dtype
+      end
+
+      if !image_processor.nil?
+        kwargs[:image_processor] = image_processor
+      end
+
+      if !device.nil?
+        kwargs[:device] = device
+      end
+
+      pipeline_class.new(model, framework: framework, task: task, **kwargs)
+    end
+
+    private
+
+    def check_task(task)
+      PIPELINE_REGISTRY.check_task(task)
+    end
+  end
+end

data/lib/transformers/pipelines/base.rb

@@ -0,0 +1,301 @@
+# Copyright 2018 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
+  module Pipelines
+    def self.get_default_model_and_revision(targeted_task, framework, task_options)
+      defaults = targeted_task["default"]
+
+      if defaults.key?("model")
+        default_models = targeted_task["default"]["model"]
+      end
+
+      if framework.nil?
+        framework = "pt"
+      end
+
+      default_models[framework]
+    end
+
+    def self.infer_framework_load_model(
+      model,
+      config,
+      model_classes: nil,
+      task: nil,
+      framework: nil,
+      **model_kwargs
+    )
+      if model.is_a?(String)
+        model_kwargs[:_from_pipeline] = task
+        class_tuple = []
+        look_pt = true
+
+        if model_classes
+          if look_pt
+            class_tuple = class_tuple + model_classes.fetch("pt", AutoModel)
+          end
+        end
+        if config.architectures
+          classes = []
+          config.architectures.each do |architecture|
+            if look_pt
+              _class = Transformers.const_get(architecture)
+              if !_class.nil?
+                classes << _class
+              end
+            end
+          end
+          class_tuple = class_tuple + classes
+        end
+
+        if class_tuple.length == 0
+          raise ArgumentError, "Pipeline cannot infer suitable model classes from #{model}"
+        end
+
+        class_tuple.each do |model_class|
+          raise Error, "Invalid auto model class: #{model_class}" unless model_class < BaseAutoModelClass
+          kwargs = model_kwargs.dup
+
+          begin
+            model = model_class.from_pretrained(model, **kwargs)
+            if model.respond_to?(:eval)
+              model = model.eval
+            end
+            break
+          rescue
+            # TODO
+            raise
+          end
+        end
+      end
+
+      if framework.nil?
+        framework = Utils.infer_framework(model.class)
+      end
+      [framework, model]
+    end
+  end
+
+  class ArgumentHandler
+  end
+
+  class Pipeline
+    def initialize(
+      model,
+      tokenizer: nil,
+      feature_extractor: nil,
+      image_processor: nil,
+      modelcard: nil,
+      framework: nil,
+      task: "",
+      device: nil,
+      **kwargs
+    )
+      if framework.nil?
+        raise Todo
+      end
+
+      @task = task
+      @model = model
+      @tokenizer = tokenizer
+      @feature_extractor = feature_extractor
+      @image_processor = image_processor
+      @modelcard = modelcard
+      @framework = framework
+
+      if device.nil?
+        if Torch::CUDA.available? || Torch::Backends::MPS.available?
+          Transformers.logger.warn(
+            "Hardware accelerator e.g. GPU is available in the environment, but no `device` argument" +
+            " is passed to the `Pipeline` object. Model will be on CPU."
+          )
+        end
+      end
+
+      @call_count = 0
+      @batch_size = kwargs.delete(:batch_size)
+      @num_workers = kwargs.delete(:num_workers)
+      @preprocess_params, @forward_params, @postprocess_params = _sanitize_parameters(**kwargs)
+    end
+
+    def torch_dtype
+      @model.dtype
+    end
+
+    def check_model_type(supported_models)
+      if !supported_models.is_a?(Array)
+        supported_models_names = []
+        supported_models.each do |_, model_name|
+          # Mapping can now contain tuples of models for the same configuration.
+          if model_name.is_a?(Array)
+            supported_models_names.concat(model_name)
+          else
+            supported_models_names << model_name
+          end
+        end
+        supported_models = supported_models_names
+      end
+      if !supported_models.include?(@model.class.name.split("::").last)
+        Transformers.logger.error(
+          "The model '#{@model.class.name}' is not supported for #{@task}. Supported models are" +
+          " #{supported_models}."
+        )
+      end
+    end
+
+    def get_iterator(
+      inputs, num_workers, batch_size, preprocess_params, forward_params, postprocess_params
+    )
+      if inputs.respond_to?(:size)
+        dataset = PipelineDataset.new(inputs, method(:preprocess), preprocess_params)
+      else
+        if num_workers > 1
+          Transformers.logger.warn(
+            "For iterable dataset using num_workers>1 is likely to result" +
+            " in errors since everything is iterable, setting `num_workers: 1`" +
+            " to guarantee correctness."
+          )
+          num_workers = 1
+        end
+        dataset = PipelineIterator.new(inputs, method(:preprocess), preprocess_params)
+      end
+
+      # TODO hack by collating feature_extractor and image_processor
+      feature_extractor = !@feature_extractor.nil? ? @feature_extractor : @image_processor
+      collate_fn = batch_size == 1 ? method(:no_collate_fn) : pad_collate_fn(@tokenizer, feature_extractor)
+      dataloader = Torch::Utils::Data::DataLoader.new(dataset, batch_size: batch_size, collate_fn: collate_fn) # num_workers: num_workers,
+      model_iterator = PipelineIterator.new(dataloader, method(:forward), forward_params, loader_batch_size: batch_size)
+      final_iterator = PipelineIterator.new(model_iterator, method(:postprocess), postprocess_params)
+      final_iterator
+    end
+
+    def call(inputs, *args, num_workers: nil, batch_size: nil, **kwargs)
+      if args.any?
+        Transformers.logger.warn("Ignoring args : #{args}")
+      end
+
+      if num_workers.nil?
+        if @num_workers.nil?
+          num_workers = 0
+        else
+          num_workers = @num_workers
+        end
+      end
+      if batch_size.nil?
+        if @batch_size.nil?
+          batch_size = 1
+        else
+          batch_size = @batch_size
+        end
+      end
+
+      preprocess_params, forward_params, postprocess_params = _sanitize_parameters(**kwargs)
+
+      preprocess_params = @preprocess_params.merge(preprocess_params)
+      forward_params = @forward_params.merge(forward_params)
+      postprocess_params = @postprocess_params.merge(postprocess_params)
+
+      @call_count += 1
+      if @call_count > 10 && @framework == "pt" && @device.type == "cuda"
+        Transformers.logger.warn(
+          "You seem to be using the pipelines sequentially on GPU. In order to maximize efficiency please use a" +
+          " dataset"
+        )
+      end
+
+      is_dataset = inputs.is_a?(Torch::Utils::Data::Dataset)
+      is_generator = inputs.is_a?(Enumerable)
+      is_list = inputs.is_a?(Array)
+
+      _is_iterable = is_dataset || is_generator || is_list
+
+      # TODO make the get_iterator work also for `tf` (and `flax`).
+      can_use_iterator = @framework == "pt" && (is_dataset || is_generator || is_list)
+
+      if is_list
+        if can_use_iterator
+          final_iterator = get_iterator(
+            inputs, num_workers, batch_size, preprocess_params, forward_params, postprocess_params
+          )
+          outputs = final_iterator.to_a
+          outputs
+        else
+          run_multi(inputs, preprocess_params, forward_params, postprocess_params)
+        end
+      else
+        run_single(inputs, preprocess_params, forward_params, postprocess_params)
+      end
+    end
+
+    private
+
+    def _sanitize_parameters(**kwargs)
+      raise NotImplementedError, "_sanitize_parameters not implemented"
+    end
+
+    def forward(model_inputs, **forward_params)
+      _forward(model_inputs, **forward_params)
+    end
+
+    def run_single(inputs, preprocess_params, forward_params, postprocess_params)
+      model_inputs = preprocess(inputs, **preprocess_params)
+      model_outputs = forward(model_inputs, **forward_params)
+      outputs = postprocess(model_outputs, **postprocess_params)
+      outputs
+    end
+
+    def no_collate_fn(items)
+      if items.length != 1
+        raise ArgumentError, "This collate_fn is meant to be used with batch_size=1"
+      end
+      items[0]
+    end
+  end
+
+  class ChunkPipeline < Pipeline
+    def run_single(inputs, preprocess_params, forward_params, postprocess_params)
+      all_outputs = []
+      preprocess(inputs, **preprocess_params) do |model_inputs|
+        model_outputs = forward(model_inputs, **forward_params)
+        all_outputs << model_outputs
+      end
+      outputs = postprocess(all_outputs, **postprocess_params)
+      outputs
+    end
+  end
+
+  class PipelineRegistry
+    def initialize(supported_tasks:, task_aliases:)
+      @supported_tasks = supported_tasks
+      @task_aliases = task_aliases
+    end
+
+    def get_supported_tasks
+      supported_task = @supported_tasks.keys + @task_aliases.keys
+      supported_task.sort
+    end
+
+    def check_task(task)
+      if @task_aliases[task]
+        task = @task_aliases[task]
+      end
+      if @supported_tasks[task]
+        targeted_task = @supported_tasks[task]
+        return task, targeted_task, nil
+      end
+
+      raise KeyError, "Unknown task #{task}, available tasks are #{get_supported_tasks}"
+    end
+  end
+end

data/lib/transformers/pipelines/feature_extraction.rb

@@ -0,0 +1,47 @@
+module Transformers
+  class FeatureExtractionPipeline < Pipeline
+    def _sanitize_parameters(truncation: nil, tokenize_kwargs: nil, return_tensors: nil, **kwargs)
+      if tokenize_kwargs.nil?
+        tokenize_kwargs = {}
+      end
+
+      if !truncation.nil?
+        if tokenize_kwargs.include?(:truncation)
+          raise ArgumentError,
+            "truncation parameter defined twice (given as keyword argument as well as in tokenize_kwargs)"
+        end
+        tokenize_kwargs[:truncation] = truncation
+      end
+
+      preprocess_params = tokenize_kwargs
+
+      postprocess_params = {}
+      if !return_tensors.nil?
+        postprocess_params[:return_tensors] = return_tensors
+      end
+
+      [preprocess_params, {}, postprocess_params]
+    end
+
+    def preprocess(inputs, **tokenize_kwargs)
+      model_inputs = @tokenizer.(inputs, return_tensors: @framework, **tokenize_kwargs)
+      model_inputs
+    end
+
+    def _forward(model_inputs)
+      model_outputs = @model.(**model_inputs)
+      model_outputs
+    end
+
+    def postprocess(model_outputs, return_tensors: false)
+      # [0] is the first available tensor, logits or last_hidden_state.
+      if return_tensors
+        model_outputs[0]
+      elsif @framework == "pt"
+        model_outputs[0].to_a
+      elsif @framework == "tf"
+        raise Todo
+      end
+    end
+  end
+end