[TransformersPipeline] Add in and refactor TransformersPipeline args (#…

…1218)

* add in and refactor TransformersPipeline args

* debug

* clean-up

* add warning, fix bug in helper, add additional check

src/deepsparse/transformers/helpers.py

@@ -19,6 +19,7 @@
 
 import os
 import re
+import warnings
 from pathlib import Path
 from tempfile import NamedTemporaryFile
 from typing import List, Optional, Tuple, Union
@@ -34,7 +35,8 @@
 
 
 __all__ = [
-    "get_onnx_path_and_configs",
+    "get_hugging_face_configs",
+    "get_onnx_path",
     "overwrite_transformer_onnx_model_inputs",
     "fix_numpy_types",
     "get_transformer_layer_init_names",
@@ -49,28 +51,9 @@
 _MODEL_DIR_TOKENIZER_CONFIG_NAME = "tokenizer_config.json"
 
 
-def get_onnx_path_and_configs(
-    model_path: str,
-    require_configs: bool = False,
-) -> Tuple[str, Optional[str], Optional[str]]:
-    """
-    :param model_path: path to onnx file, transformers sparsezoo stub,
-        or directory containing `model.onnx`, `config.json`, and/or
-        `tokenizer.json` files. If no `model.onnx` file is found in
-        a model directory, an exception will be raised
-    :param require_configs: if True, model_path must be a directory containing
-        `model.onnx`, `config.json`, and `tokenizer.json` files. Will raise
-        an exception otherwise
-    :return: tuple of ONNX file path, parent directory of config file
-        if it exists, and parent directory of tokenizer config file if it
-        exists. (Parent directories returned instead of absolute path
-        for compatibility with transformers .from_pretrained() method)
-    """
-    if os.path.isfile(model_path) and not require_configs:
-        return model_path, None, None
-
-    config_path = None
-    tokenizer_path = None
+def get_onnx_path(model_path: str) -> str:
+    if os.path.isfile(model_path):
+        return model_path
 
     if os.path.isdir(model_path):
         model_files = os.listdir(model_path)
@@ -83,6 +66,32 @@ def get_onnx_path_and_configs(
             )
         onnx_path = os.path.join(model_path, _MODEL_DIR_ONNX_NAME)
 
+    elif model_path.startswith("zoo:"):
+        zoo_model = Model(model_path)
+        onnx_path = zoo_model.onnx_model.path
+    else:
+        raise ValueError(
+            f"model_path {model_path} is not a valid file, directory, or zoo stub"
+        )
+
+    return onnx_path
+
+
+def get_hugging_face_configs(model_path: str) -> Tuple[str, str]:
+    """
+    :param model_path: path to model directory, transformers sparsezoo stub,
+        or directory containing `config.json`, and `tokenizer.json` files.
+        If the json files are not found, an exception will be raised.
+    :return: tuple of ONNX file path, parent directory of config file
+        if it exists, and parent directory of tokenizer config file if it
+        exists. (Parent directories returned instead of absolute path
+        for compatibility with transformers .from_pretrained() method)
+    """
+    config_path = None
+    tokenizer_path = None
+
+    if os.path.isdir(model_path):
+        model_files = os.listdir(model_path)
         # attempt to read config and tokenizer from sparsezoo-like framework directory
         framework_dir = None
         if "framework" in model_files:
@@ -102,12 +111,14 @@ def get_onnx_path_and_configs(
         # prefer config and tokenizer files in same directory as model.onnx
         if _MODEL_DIR_CONFIG_NAME in model_files:
             config_path = model_path
-        if _MODEL_DIR_TOKENIZER_NAME or _MODEL_DIR_TOKENIZER_CONFIG_NAME in model_files:
+        if (
+            _MODEL_DIR_TOKENIZER_NAME in model_files
+            or _MODEL_DIR_TOKENIZER_CONFIG_NAME in model_files
+        ):
             tokenizer_path = model_path
 
     elif model_path.startswith("zoo:"):
         zoo_model = Model(model_path)
-        onnx_path = zoo_model.onnx_model.path
         config_path = _get_file_parent(
             zoo_model.deployment.default.get_file(_MODEL_DIR_CONFIG_NAME).path
         )
@@ -119,19 +130,21 @@ def get_onnx_path_and_configs(
         )
         if tokenizer_config_path is not None:
             tokenizer_config_path.path  # trigger download of tokenizer_config
-    elif require_configs and (config_path is None or tokenizer_path is None):
-        raise RuntimeError(
-            f"Unable to find model and tokenizer config for model_path {model_path}. "
-            f"model_path must be a directory containing model.onnx, config.json, and "
-            f"tokenizer.json files. Found config and tokenizer paths: {config_path}, "
-            f"{tokenizer_path}"
-        )
     else:
         raise ValueError(
-            f"model_path {model_path} is not a valid file, directory, or zoo stub"
+            f"model_path {model_path} is not a valid directory or zoo stub"
+        )
+
+    if config_path is None or tokenizer_path is None:
+        warnings.warn(
+            f"Unable to find model or tokenizer configs for model_path {model_path}. "
+            f"model_path must be a directory containing config.json, and/or "
+            f"tokenizer.json files. Found config and tokenizer paths: {config_path}, "
+            f"{tokenizer_path}. If not given, set the `tokenizer` and `config` args "
+            "for the Pipeline."
         )
 
-    return onnx_path, config_path, tokenizer_path
+    return config_path, tokenizer_path
 
 
 def overwrite_transformer_onnx_model_inputs(

src/deepsparse/transformers/pipelines/pipeline.py

@@ -16,16 +16,20 @@
 Base Pipeline class for transformers inference pipeline
 """
 
+import json
 import os
 import warnings
-from typing import Any, List, Mapping, Optional, Union
+from pathlib import Path
+from typing import Any, Dict, List, Mapping, Optional, Union
 
 import numpy
+import transformers
 from transformers.models.auto import AutoConfig, AutoTokenizer
 
 from deepsparse import Bucketable, Pipeline
 from deepsparse.transformers.helpers import (
-    get_onnx_path_and_configs,
+    get_hugging_face_configs,
+    get_onnx_path,
     overwrite_transformer_onnx_model_inputs,
 )
 
@@ -70,21 +74,33 @@ class TransformersPipeline(Pipeline, Bucketable):
         should only be set to `True` for repositories you trust and in which
         you have read the code, as it will execute possibly unsafe code
         on your local machine. Default is False
+    :param config: hugging face transformers model config. Can be a path to the config,
+        a dictionary with the config values, a transformers.PretrainedConfig, or None.
+        If a directory is provided, it is assumed that the file is named config.json.
+        If None, an attempt is made to read the config file from the the model_path
+        directory provided. Default is None
+    :param tokenizer: hugging face transfromers tokenizer. Can be a path to the
+        a directory with the relevant tokenizer files, a
+        transformers.PreTrainedTokenizerBase, or None. If None, an attempt is made to
+        read the json file from the model_path directory provided. Default is None.
     """
 
     def __init__(
         self,
         *,
         sequence_length: Union[int, List[int]] = 128,
         trust_remote_code: bool = False,
+        config: Union[str, Path, Dict, transformers.PretrainedConfig] = None,
+        tokenizer: Union[str, Path, transformers.PreTrainedTokenizerBase] = None,
         **kwargs,
     ):
 
         self._sequence_length = sequence_length
         self._trust_remote_code = trust_remote_code
 
-        self.config = None
-        self.tokenizer = None
+        self.config = config
+        self.tokenizer = tokenizer
+
         self.config_path = None
         self.tokenizer_config_path = None  # path to 'tokenizer.json'
         self.onnx_input_names = None
@@ -104,30 +120,51 @@ def sequence_length(self) -> Union[int, List[int]]:
 
     def setup_onnx_file_path(self) -> str:
         """
-        Parses ONNX, tokenizer, and config file paths from the given `model_path`.
-        Supports sparsezoo stubs
+        Parses ONNX model from the `model_path` provided.
+        For tokenizers and model configs, supports paths, dictionaries,
+        or transformers.PretrainedConfig/transformes.PreTrainedTokenizerBase types. Also
+        supports the default None, in which case the config and tokenizer are read from
+        the provided `model_path`.
 
-        :param delay_overwriting_inputs: if True, do not overwrite the ONNX model
-            inputs to the given sequence length. Default is False
+        Supports sparsezoo stubs
 
         :return: file path to the processed ONNX file for the engine to compile
         """
-        onnx_path, config_path, tokenizer_path = get_onnx_path_and_configs(
-            self.model_path, require_configs=True
-        )
+        onnx_path = get_onnx_path(self.model_path)
+
+        if not self.config or not self.tokenizer:
+            config_found, tokenizer_found = get_hugging_face_configs(self.model_path)
+            if config_found:
+                self.config = config_found
+            if tokenizer_found:
+                self.tokenizer = tokenizer_found
+
+        if isinstance(self.config, dict):
+            local_config_path = os.path.join(self.model_path, "config.json")
+            with open(local_config_path, "w") as f:
+                json.dump(self.config, f)
+            self.config = local_config_path
+
+        if isinstance(self.config, (str, Path)):
+            if str(self.config).endswith(".json"):
+                self.config_path = self.config
+            else:
+                self.config_path = os.path.join(self.config, "config.json")
+
+            self.config = AutoConfig.from_pretrained(
+                self.config,
+                trust_remote_code=self._trust_remote_code,
+                finetuning_task=self.task if hasattr(self, "task") else None,
+            )
 
-        self.config = AutoConfig.from_pretrained(
-            config_path,
-            trust_remote_code=self._trust_remote_code,
-            finetuning_task=self.task if hasattr(self, "task") else None,
-        )
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            tokenizer_path,
-            trust_remote_code=self._trust_remote_code,
-            model_max_length=self.sequence_length,
-        )
-        self.config_path = os.path.join(config_path, "config.json")
-        self.tokenizer_config_path = os.path.join(tokenizer_path, "tokenizer.json")
+        if isinstance(self.tokenizer, (str, Path)):
+            self.tokenizer_config_path = os.path.join(self.tokenizer, "tokenizer.json")
+
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.tokenizer,
+                trust_remote_code=self._trust_remote_code,
+                model_max_length=self.sequence_length,
+            )
 
         if not self._delay_overwriting_inputs:
             # overwrite onnx graph to given required input shape
@@ -139,6 +176,12 @@ def setup_onnx_file_path(self) -> str:
                 onnx_path, max_length=self.sequence_length
             )
 
+        if not self.config or not self.tokenizer:
+            raise RuntimeError(
+                "Invalid config or tokenizer provided. Please provide "
+                "paths to the files or ensure they exist in the `model_path` provided. "
+                "See `tokenizer` and `config` arguments for details."
+            )
         return onnx_path
 
     def tokens_to_engine_input(

tests/deepsparse/transformers/test_helpers.py

@@ -18,7 +18,8 @@
 
 import pytest
 from deepsparse.transformers.helpers import (
-    get_onnx_path_and_configs,
+    get_hugging_face_configs,
+    get_onnx_path,
     get_transformer_layer_init_names,
     truncate_transformer_onnx_model,
 )
@@ -35,7 +36,8 @@
     ],
 )
 def test_get_onnx_path_and_configs_from_stub(stub):
-    onnx_path, config_dir, tokenizer_dir = get_onnx_path_and_configs(stub)
+    onnx_path = get_onnx_path(stub)
+    config_dir, tokenizer_dir = get_hugging_face_configs(stub)
 
     assert onnx_path.endswith("model.onnx")
     assert os.path.exists(onnx_path)