|
17 | 17 | """
|
18 | 18 |
|
19 | 19 |
|
| 20 | +import logging |
20 | 21 | import os
|
21 | 22 | import re
|
22 | 23 | from pathlib import Path
|
23 | 24 | from tempfile import NamedTemporaryFile
|
24 |
| -from typing import List, Optional, Tuple, Union |
| 25 | +from typing import Any, Dict, List, Optional, Tuple, Union |
25 | 26 |
|
26 | 27 | import numpy
|
27 | 28 | import onnx
|
| 29 | +import transformers |
28 | 30 | from onnx import ModelProto
|
29 | 31 |
|
30 | 32 | from deepsparse.log import get_main_logger
|
|
38 | 40 |
|
39 | 41 | __all__ = [
|
40 | 42 | "get_deployment_path",
|
| 43 | + "setup_transformers_pipeline", |
41 | 44 | "overwrite_transformer_onnx_model_inputs",
|
42 | 45 | "fix_numpy_types",
|
43 | 46 | "get_transformer_layer_init_names",
|
|
47 | 50 | _LOGGER = get_main_logger()
|
48 | 51 |
|
49 | 52 |
|
| 53 | +def setup_transformers_pipeline( |
| 54 | + model_path: str, |
| 55 | + sequence_length: int, |
| 56 | + tokenizer_padding_side: str = "left", |
| 57 | + engine_kwargs: Optional[Dict] = None, |
| 58 | +) -> Tuple[ |
| 59 | + str, transformers.PretrainedConfig, transformers.PreTrainedTokenizer, Dict[str, Any] |
| 60 | +]: |
| 61 | + """ |
| 62 | + A helper function that sets up the model path, config, tokenizer, |
| 63 | + and engine kwargs for a transformers model. |
| 64 | + :param model_path: The path to the model to load |
| 65 | + :param sequence_length: The sequence length to use for the model |
| 66 | + :param tokenizer_padding_side: The side to pad on for the tokenizer, |
| 67 | + either "left" or "right" |
| 68 | + :param engine_kwargs: The kwargs to pass to the engine |
| 69 | + :return The model path, config, tokenizer, and engine kwargs |
| 70 | + """ |
| 71 | + model_path, config, tokenizer = fetch_onnx_file_path(model_path, sequence_length) |
| 72 | + |
| 73 | + tokenizer.padding_side = tokenizer_padding_side |
| 74 | + if not tokenizer.pad_token: |
| 75 | + tokenizer.pad_token = tokenizer.eos_token |
| 76 | + |
| 77 | + engine_kwargs = engine_kwargs or {} |
| 78 | + if engine_kwargs.get("model_path"): |
| 79 | + raise ValueError( |
| 80 | + "The engine kwargs already specify " |
| 81 | + f"a model path: {engine_kwargs['model_path']}, " |
| 82 | + f"but a model path was also provided: {model_path}. " |
| 83 | + "Please only provide one." |
| 84 | + ) |
| 85 | + engine_kwargs["model_path"] = model_path |
| 86 | + return model_path, config, tokenizer, engine_kwargs |
| 87 | + |
| 88 | + |
| 89 | +def fetch_onnx_file_path( |
| 90 | + model_path: str, |
| 91 | + sequence_length: int, |
| 92 | + task: Optional[str] = None, |
| 93 | +) -> Tuple[str, transformers.PretrainedConfig, transformers.PreTrainedTokenizer]: |
| 94 | + """ |
| 95 | + Parses ONNX model from the `model_path` provided. It additionally |
| 96 | + creates config and tokenizer objects from the `deployment path`, |
| 97 | + derived from the `model_path` provided. |
| 98 | + :param model_path: path to the model to be parsed |
| 99 | + :param sequence_length: maximum sequence length of the model |
| 100 | + :return: file path to the processed ONNX file for the engine to compile |
| 101 | + """ |
| 102 | + deployment_path, onnx_path = get_deployment_path(model_path) |
| 103 | + |
| 104 | + hf_logger = logging.getLogger("transformers") |
| 105 | + hf_logger_level = hf_logger.level |
| 106 | + hf_logger.setLevel(logging.ERROR) |
| 107 | + |
| 108 | + config = transformers.PretrainedConfig.from_pretrained( |
| 109 | + deployment_path, finetuning_task=task |
| 110 | + ) |
| 111 | + hf_logger.setLevel(hf_logger_level) |
| 112 | + |
| 113 | + trust_remote_code = False |
| 114 | + tokenizer = transformers.AutoTokenizer.from_pretrained( |
| 115 | + deployment_path, |
| 116 | + trust_remote_code=trust_remote_code, |
| 117 | + model_max_length=sequence_length, |
| 118 | + ) |
| 119 | + |
| 120 | + if not config or not tokenizer: |
| 121 | + raise RuntimeError( |
| 122 | + "Invalid config or tokenizer provided. Please provide " |
| 123 | + "paths to the files or ensure they exist in the `model_path` provided. " |
| 124 | + "See `tokenizer` and `config` arguments for details." |
| 125 | + ) |
| 126 | + return onnx_path, config, tokenizer |
| 127 | + |
| 128 | + |
50 | 129 | def get_deployment_path(model_path: str) -> Tuple[str, str]:
|
51 | 130 | """
|
52 | 131 | Returns the path to the deployment directory
|
|
0 commit comments