From d6ec5155138fc3cfd0a4b8bb6114348d126bc09a Mon Sep 17 00:00:00 2001
From: conggguan <congguan@amazon.com>
Date: Tue, 6 Aug 2024 17:11:11 +0800
Subject: [PATCH] [Feature] Add a workflow parameter that model uploader can
 specific a customize prefix.

Signed-off-by: conggguan <congguan@amazon.com>
---
 .github/workflows/model_uploader.yml | 11 ++++++++++-
 CHANGELOG.md                         |  1 +
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/model_uploader.yml b/.github/workflows/model_uploader.yml
index 1c7362b2..cdfc3c28 100644
--- a/.github/workflows/model_uploader.yml
+++ b/.github/workflows/model_uploader.yml
@@ -24,6 +24,10 @@ on:
         - "BOTH"
         - "TORCH_SCRIPT"
         - "ONNX"
+      upload_prefix:
+        description: "Specifies the model prefix for uploading. For example, transforming the default path from '.../sentence-transformers/msmarco-distilbert-base-tas-b' to '.../{prefix}/msmarco-distilbert-base-tas-b'."
+        required: false
+        type: string
       model_type:
         description: "Model type for auto-tracing (SentenceTransformer/Sparse)"
         required: true
@@ -74,7 +78,12 @@ jobs:
       run: |
         model_id=${{ github.event.inputs.model_id }}
         echo "model_folder=ml-models/${{github.event.inputs.model_source}}/${model_id}" >> $GITHUB_OUTPUT
-        echo "model_prefix_folder=ml-models/${{github.event.inputs.model_source}}/${model_id%%/*}/" >> $GITHUB_OUTPUT
+        if [[ -n "${{ github.event.inputs.upload_prefix }}" ]]; then
+          model_prefix="ml-models/${{ github.event.inputs.model_source }}/${{ github.event.inputs.upload_prefix }}"
+        else
+          model_prefix="ml-models/${{ github.event.inputs.model_source }}/${model_id%%/*}"
+        fi
+        echo "model_prefix_folder=$model_prefix" >> $GITHUB_OUTPUT
     - name: Initiate workflow_info
       id: init_workflow_info
       run: |
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f8168762..49a52003 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - Add workflows and scripts for sparse encoding model tracing and uploading process by @conggguan in ([#394](https://github.com/opensearch-project/opensearch-py-ml/pull/394))
 
 ### Changed
+- Add a parameter for customize the upload folder prefix ([#398](https://github.com/opensearch-project/opensearch-py-ml/pull/398))
 - Modify ml-models.JenkinsFile so that it takes model format into account and can be triggered with generic webhook by @thanawan-atc in ([#211](https://github.com/opensearch-project/opensearch-py-ml/pull/211))
 - Update demo_tracing_model_torchscript_onnx.ipynb to use make_model_config_json by @thanawan-atc in ([#220](https://github.com/opensearch-project/opensearch-py-ml/pull/220))
 - Bump torch from 1.13.1 to 2.0.1 and add onnx dependency by @thanawan-atc ([#237](https://github.com/opensearch-project/opensearch-py-ml/pull/237))