Add model_type

Signed-off-by: Thanawan Atchariyachanvanit <[email protected]>
thanawan-atc · Aug 16, 2023 · 5c7b4a9 · 5c7b4a9
1 parent 6f69ba2
commit 5c7b4a9
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git 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
@@ -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