read function Name from pretrained model

Signed-off-by: xinyual <[email protected]>
xinyual · Oct 18, 2023 · e0b2487 · e0b2487
1 parent e9e3834
commit e0b2487
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/ModelHelper.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/ModelHelper.java
@@ -87,7 +87,8 @@ public void downloadPrebuiltModelConfig(String taskId, MLRegisterModelInput regi
                         .url(modelZipFileUrl)
                         .deployModel(deployModel)
                         .modelNodeIds(modelNodeIds)
-                        .modelGroupId(modelGroupId);
+                        .modelGroupId(modelGroupId)
+                        .functionName(FunctionName.from((String) config.get("model_task_type")));;
                 config.entrySet().forEach(entry -> {
                     switch (entry.getKey().toString()) {
                         case MLRegisterModelInput.MODEL_FORMAT_FIELD:

diff --git a/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java b/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java
@@ -609,7 +609,7 @@ private void uploadModel(MLRegisterModelInput registerModelInput, MLTask mlTask,
 
     private void registerModelFromUrl(MLRegisterModelInput registerModelInput, MLTask mlTask, String modelVersion) {
         String taskId = mlTask.getTaskId();
-        FunctionName functionName = mlTask.getFunctionName();
+        FunctionName functionName = registerModelInput.getFunctionName();
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
             String modelName = registerModelInput.getModelName();
             String version = modelVersion == null ? registerModelInput.getVersion() : modelVersion;