-
Notifications
You must be signed in to change notification settings - Fork 171
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[Pipeline Refactor][Text-Generation] Refactor transformers
helpers functions
#1394
Changes from all commits
f18d5f3
2c4d231
672ca20
304eb35
71515ac
6f1b175
a508342
cbb0e86
2541581
8c8989d
f8d75e3
fd1e466
64c0552
913665a
e15521f
379481e
a90a20a
c0c4240
4f248dd
98f7a6d
d1683b4
51c4ee6
a4f6f19
71f4c6d
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -17,24 +17,26 @@ | |
""" | ||
|
||
|
||
import logging | ||
import os | ||
import re | ||
from pathlib import Path | ||
from tempfile import NamedTemporaryFile | ||
from typing import List, Optional, Tuple, Union | ||
from typing import Any, Dict, List, Optional, Tuple, Union | ||
|
||
import numpy | ||
import onnx | ||
import transformers | ||
from onnx import ModelProto | ||
|
||
from deepsparse.log import get_main_logger | ||
from deepsparse.utils.onnx import _MODEL_DIR_ONNX_NAME, truncate_onnx_model | ||
from deepsparse.utils.onnx import MODEL_ONNX_NAME, truncate_onnx_model | ||
from sparsezoo import Model | ||
from sparsezoo.utils import save_onnx | ||
|
||
|
||
__all__ = [ | ||
"get_deployment_path", | ||
"setup_transformers_pipeline", | ||
"overwrite_transformer_onnx_model_inputs", | ||
"fix_numpy_types", | ||
"get_transformer_layer_init_names", | ||
|
@@ -44,7 +46,94 @@ | |
_LOGGER = get_main_logger() | ||
|
||
|
||
def get_deployment_path(model_path: str) -> Tuple[str, str]: | ||
def setup_transformers_pipeline( | ||
model_path: str, | ||
sequence_length: int, | ||
tokenizer_padding_side: str = "left", | ||
engine_kwargs: Optional[Dict] = None, | ||
onnx_model_name: Optional[str] = None, | ||
) -> Tuple[ | ||
str, transformers.PretrainedConfig, transformers.PreTrainedTokenizer, Dict[str, Any] | ||
]: | ||
""" | ||
A helper function that sets up the model path, config, tokenizer, | ||
and engine kwargs for a transformers model. | ||
:param model_path: The path to the model to load | ||
:param sequence_length: The sequence length to use for the model | ||
:param tokenizer_padding_side: The side to pad on for the tokenizer, | ||
either "left" or "right" | ||
:param engine_kwargs: The kwargs to pass to the engine | ||
:param onnx_model_name: The name of the onnx model to be loaded. | ||
If not specified, defaults are used (see setup_onnx_file_path) | ||
:return The model path, config, tokenizer, and engine kwargs | ||
""" | ||
model_path, config, tokenizer = setup_onnx_file_path( | ||
model_path, sequence_length, onnx_model_name | ||
) | ||
|
||
tokenizer.padding_side = tokenizer_padding_side | ||
if not tokenizer.pad_token: | ||
tokenizer.pad_token = tokenizer.eos_token | ||
|
||
engine_kwargs = engine_kwargs or {} | ||
if engine_kwargs.get("model_path"): | ||
raise ValueError( | ||
"The engine kwargs already specify " | ||
f"a model path: {engine_kwargs['model_path']}, " | ||
f"but a model path was also provided: {model_path}. " | ||
"Please only provide one." | ||
) | ||
engine_kwargs["model_path"] = model_path | ||
dsikka marked this conversation as resolved.
Show resolved
Hide resolved
|
||
return model_path, config, tokenizer, engine_kwargs | ||
|
||
|
||
def setup_onnx_file_path( | ||
model_path: str, | ||
sequence_length: int, | ||
onnx_model_name: Optional[str] = None, | ||
task: Optional[str] = None, | ||
) -> Tuple[str, transformers.PretrainedConfig, transformers.PreTrainedTokenizer]: | ||
""" | ||
Parses ONNX model from the `model_path` provided. It additionally | ||
creates config and tokenizer objects from the `deployment path`, | ||
derived from the `model_path` provided. | ||
:param model_path: path to the model to be parsed | ||
:param sequence_length: maximum sequence length of the model | ||
:param onnx_model_name: optionally, the precise name of the ONNX model | ||
of interest may be specified. If not specified, the default ONNX model | ||
name will be used (refer to `get_deployment_path` for details) | ||
:return: file path to the processed ONNX file for the engine to compile | ||
""" | ||
deployment_path, onnx_path = get_deployment_path(model_path, onnx_model_name) | ||
|
||
hf_logger = logging.getLogger("transformers") | ||
hf_logger_level = hf_logger.level | ||
hf_logger.setLevel(logging.ERROR) | ||
|
||
config = transformers.PretrainedConfig.from_pretrained( | ||
deployment_path, finetuning_task=task | ||
) | ||
hf_logger.setLevel(hf_logger_level) | ||
|
||
trust_remote_code = False | ||
tokenizer = transformers.AutoTokenizer.from_pretrained( | ||
deployment_path, | ||
trust_remote_code=trust_remote_code, | ||
model_max_length=sequence_length, | ||
) | ||
|
||
if not config or not 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, config, tokenizer | ||
|
||
|
||
def get_deployment_path( | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We talked about a bug/fix for this in stand-up: #1396 There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Let's not make this PR take on too many hats. The problem that @bfineran was fixing was particular to the few transformer tasks. More of those are bound to happen once we swap v1 transformers pipelines in favor of v2 transformers pipelines. We will deal with them in their own time, this PR seems fairly orthogonal. |
||
model_path: str, onnx_model_name: Optional[str] = None | ||
) -> Tuple[str, str]: | ||
""" | ||
Returns the path to the deployment directory | ||
for the given model path and the path to the mandatory | ||
|
@@ -53,36 +142,39 @@ def get_deployment_path(model_path: str) -> Tuple[str, str]: | |
for running the transformers model in the deepsparse pipeline | ||
|
||
:param model_path: path to model directory, sparsezoo stub, or ONNX file | ||
:param onnx_model_name: name of the ONNX file to look for in the deployment | ||
directory. Defaults to MODEL_ONNX_NAME | ||
:return: path to the deployment directory and path to the ONNX file inside | ||
the deployment directory | ||
""" | ||
onnx_model_name = onnx_model_name or MODEL_ONNX_NAME | ||
if os.path.isfile(model_path): | ||
# return the parent directory of the ONNX file | ||
return os.path.dirname(model_path), model_path | ||
|
||
if os.path.isdir(model_path): | ||
model_files = os.listdir(model_path) | ||
|
||
if _MODEL_DIR_ONNX_NAME not in model_files: | ||
if onnx_model_name not in model_files: | ||
raise ValueError( | ||
f"{_MODEL_DIR_ONNX_NAME} not found in transformers model directory " | ||
f"{onnx_model_name} not found in transformers model directory " | ||
f"{model_path}. Be sure that an export of the model is written to " | ||
f"{os.path.join(model_path, _MODEL_DIR_ONNX_NAME)}" | ||
f"{os.path.join(model_path, onnx_model_name)}" | ||
) | ||
return model_path, os.path.join(model_path, _MODEL_DIR_ONNX_NAME) | ||
return model_path, os.path.join(model_path, onnx_model_name) | ||
|
||
elif model_path.startswith("zoo:"): | ||
zoo_model = Model(model_path) | ||
deployment_path = zoo_model.deployment_directory_path | ||
return deployment_path, os.path.join(deployment_path, _MODEL_DIR_ONNX_NAME) | ||
return deployment_path, os.path.join(deployment_path, onnx_model_name) | ||
elif model_path.startswith("hf:"): | ||
from huggingface_hub import snapshot_download | ||
|
||
deployment_path = snapshot_download(repo_id=model_path.replace("hf:", "", 1)) | ||
onnx_path = os.path.join(deployment_path, _MODEL_DIR_ONNX_NAME) | ||
onnx_path = os.path.join(deployment_path, onnx_model_name) | ||
if not os.path.isfile(onnx_path): | ||
raise ValueError( | ||
f"{_MODEL_DIR_ONNX_NAME} not found in transformers model directory " | ||
f"{onnx_model_name} not found in transformers model directory " | ||
f"{deployment_path}. Be sure that an export of the model is written to " | ||
f"{onnx_path}" | ||
) | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Could we also get the v1 pipeline to use these? Fear is that we'll have two places to update going forward.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sure, good idea.