fix text docs input unescaped error; enable deploy remote model (open…

…search-project#1407) Signed-off-by: Yaliang Wu <ylwu@amazon.com>
rbhavna · Sep 29, 2023 · 6e0d949 · 6e0d949
1 parent cda89ff
commit 6e0d949
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/ConnectorUtils.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/ConnectorUtils.java
@@ -81,17 +81,24 @@ public static RemoteInferenceInputDataSet processInput(MLInput mlInput, Connecto
         return inputData;
     }
     private static RemoteInferenceInputDataSet processTextDocsInput(TextDocsInputDataSet inputDataSet, Connector connector, Map<String, String> parameters, ScriptService scriptService) {
-        List<String> docs = new ArrayList<>(inputDataSet.getDocs());
         Optional<ConnectorAction> predictAction = connector.findPredictAction();
         if (predictAction.isEmpty()) {
             throw new IllegalArgumentException("no predict action found");
         }
         String preProcessFunction = predictAction.get().getPreProcessFunction();
         preProcessFunction = preProcessFunction == null ? MLPreProcessFunction.TEXT_DOCS_TO_DEFAULT_EMBEDDING_INPUT : preProcessFunction;
         if (MLPreProcessFunction.contains(preProcessFunction)) {
-            Map<String, Object> buildInFunctionResult = MLPreProcessFunction.get(preProcessFunction).apply(docs);
+            Map<String, Object> buildInFunctionResult = MLPreProcessFunction.get(preProcessFunction).apply(inputDataSet.getDocs());
             return RemoteInferenceInputDataSet.builder().parameters(convertScriptStringToJsonString(buildInFunctionResult)).build();
         } else {
+            List<String> docs = new ArrayList<>();
+            for (String doc : inputDataSet.getDocs()) {
+                if (doc != null) {
+                    docs.add(gson.toJson(doc));
+                } else {
+                    docs.add(null);
+                }
+            }
             if (preProcessFunction.contains("${parameters")) {
                 StringSubstitutor substitutor = new StringSubstitutor(parameters, "${parameters.", "}");
                 preProcessFunction = substitutor.replace(preProcessFunction);

diff --git a/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java b/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java
@@ -519,9 +519,9 @@ private void indexRemoteModel(
                     mlTask.setModelId(modelId);
                     log.info("create new model meta doc {} for upload task {}", modelId, taskId);
                     mlTaskManager.updateMLTask(taskId, ImmutableMap.of(MODEL_ID_FIELD, modelId, STATE_FIELD, COMPLETED), 5000, true);
-                    // if (registerModelInput.isDeployModel()) {
-                    // deployModelAfterRegistering(registerModelInput, modelId);
-                    // }
+                    if (registerModelInput.isDeployModel()) {
+                        deployModelAfterRegistering(registerModelInput, modelId);
+                    }
                     listener.onResponse(new MLRegisterModelResponse(taskId, MLTaskState.CREATED.name(), modelId));
                 }, e -> {
                     log.error("Failed to index model meta doc", e);