From 28c666cb1ba96eb8ae58b555f4b13ea096001df2 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Thu, 2 Nov 2023 12:16:40 -0600 Subject: [PATCH] Add hf: stub support to model_to_path (#1378) * Add hf: stub support to model_to_path * Refactor to use model_to_path in transformers/helpers --- src/deepsparse/transformers/helpers.py | 26 ++++++++------------------ src/deepsparse/utils/onnx.py | 15 +++++++++++++++ 2 files changed, 23 insertions(+), 18 deletions(-) diff --git a/src/deepsparse/transformers/helpers.py b/src/deepsparse/transformers/helpers.py index d7acc71a99..cb3f43e484 100644 --- a/src/deepsparse/transformers/helpers.py +++ b/src/deepsparse/transformers/helpers.py @@ -28,8 +28,11 @@ from onnx import ModelProto from deepsparse.log import get_main_logger -from deepsparse.utils.onnx import _MODEL_DIR_ONNX_NAME, truncate_onnx_model -from sparsezoo import Model +from deepsparse.utils.onnx import ( + _MODEL_DIR_ONNX_NAME, + model_to_path, + truncate_onnx_model, +) from sparsezoo.utils import save_onnx @@ -71,22 +74,9 @@ def get_deployment_path(model_path: str) -> Tuple[str, str]: ) return model_path, os.path.join(model_path, _MODEL_DIR_ONNX_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) - 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) - if not os.path.isfile(onnx_path): - raise ValueError( - f"{_MODEL_DIR_ONNX_NAME} not found in transformers model directory " - f"{deployment_path}. Be sure that an export of the model is written to " - f"{onnx_path}" - ) - return deployment_path, onnx_path + elif model_path.startswith("zoo:") or model_path.startswith("hf:"): + onnx_model_path = model_to_path(model_path) + return os.path.dirname(onnx_model_path), onnx_model_path else: raise ValueError( f"model_path {model_path} is not a valid file, directory, or zoo stub" diff --git a/src/deepsparse/utils/onnx.py b/src/deepsparse/utils/onnx.py index 86fcb1cef6..a3358c8f41 100644 --- a/src/deepsparse/utils/onnx.py +++ b/src/deepsparse/utils/onnx.py @@ -138,6 +138,21 @@ def model_to_path(model: Union[str, Model, File]) -> str: # get the downloaded_path -- will auto download if not on local system model = model.path + if isinstance(model, str) and model.startswith("hf:"): + # load Hugging Face model from stub + from huggingface_hub import snapshot_download + + deployment_path = snapshot_download(repo_id=model.replace("hf:", "", 1)) + onnx_path = os.path.join(deployment_path, _MODEL_DIR_ONNX_NAME) + if not os.path.isfile(onnx_path): + raise ValueError( + f"Could not find the ONNX model file '{_MODEL_DIR_ONNX_NAME}' in the " + f"Hugging Face Hub repository located at {deployment_path}. Please " + f"ensure the model has been correctly exported to ONNX format and " + f"exists in the repository." + ) + return onnx_path + if not isinstance(model, str): raise ValueError("unsupported type for model: {}".format(type(model)))