From 5c7b4a9ea06b292fab189f629620c14c1f1d7368 Mon Sep 17 00:00:00 2001 From: Thanawan Atchariyachanvanit Date: Wed, 16 Aug 2023 13:04:40 -0700 Subject: [PATCH] Add model_type Signed-off-by: Thanawan Atchariyachanvanit --- utils/model_uploader/model_autotracing.py | 9 ++++----- utils/model_uploader/save_model_file_path_to_env.py | 4 ++-- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/utils/model_uploader/model_autotracing.py b/utils/model_uploader/model_autotracing.py index e99616b0..e5789486 100644 --- a/utils/model_uploader/model_autotracing.py +++ b/utils/model_uploader/model_autotracing.py @@ -40,9 +40,8 @@ ONNX_FORMAT = "ONNX" TEMP_MODEL_PATH = "temp_model_path" -ORIGINAL_FOLDER_PATH = "sentence-transformers-original/" -TORCHSCRIPT_FOLDER_PATH = "sentence-transformers-torchscript/" -ONNX_FOLDER_PATH = "sentence-transformers-onnx/" +TORCHSCRIPT_FOLDER_PATH = "model-torchscript/" +ONNX_FOLDER_PATH = "model-onnx/" UPLOAD_FOLDER_PATH = "upload/" MODEL_CONFIG_FILE_NAME = "ml-commons_model_config.json" OUTPUT_DIR = "trace_output/" @@ -305,7 +304,7 @@ def prepare_files_for_uploading( (path to model config json file) in the UPLOAD_FOLDER_PATH :rtype: Tuple[str, str] """ - model_name = str(model_id.split("/")[-1]) + model_type, model_name = model_id.split("/") model_format = model_format.lower() folder_to_delete = ( TORCHSCRIPT_FOLDER_PATH if model_format == "torch_script" else ONNX_FOLDER_PATH @@ -318,7 +317,7 @@ def prepare_files_for_uploading( ) os.makedirs(dst_model_dir, exist_ok=True) dst_model_filename = ( - f"sentence-transformers_{model_name}-{model_version}-{model_format}.zip" + f"{model_type}_{model_name}-{model_version}-{model_format}.zip" ) dst_model_path = dst_model_dir + "/" + dst_model_filename with ZipFile(src_model_path, "a") as zipObj: diff --git a/utils/model_uploader/save_model_file_path_to_env.py b/utils/model_uploader/save_model_file_path_to_env.py index 93576fff..c0fe9cb1 100644 --- a/utils/model_uploader/save_model_file_path_to_env.py +++ b/utils/model_uploader/save_model_file_path_to_env.py @@ -49,11 +49,11 @@ def get_model_file_path( :return: Expected model file path on model hub :rtype: string """ - model_name = str(model_id.split("/")[-1]) + model_type, model_name = model_id.split("/") model_format = model_format.lower() model_dirname = f"{model_folder}{model_name}/{model_version}/{model_format}" model_filename = ( - f"sentence-transformers_{model_name}-{model_version}-{model_format}.zip" + f"{model_type}_{model_name}-{model_version}-{model_format}.zip" ) model_file_path = model_dirname + "/" + model_filename return model_file_path