Update common/src/main/java/org/opensearch/ml/common/connector/functi…

…ons/preprocess/MultiModalConnectorPreProcessFunction.java Co-authored-by: Yaliang Wu <[email protected]> Signed-off-by: zane-neo <[email protected]>
zane-neo · Jun 11, 2024 · 3591d25 · 3591d25
1 parent 8162dc8
commit 3591d25
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/...earch/ml/common/connector/functions/preprocess/MultiModalConnectorPreProcessFunction.java b/...earch/ml/common/connector/functions/preprocess/MultiModalConnectorPreProcessFunction.java
@@ -47,10 +47,12 @@ public void validate(MLInput mlInput) {
     @Override
     public RemoteInferenceInputDataSet process(MLInput mlInput) {
         TextDocsInputDataSet inputData = (TextDocsInputDataSet) mlInput.getInputDataset();
-        if (inputData.getDocs().size() == 1) {
-            return RemoteInferenceInputDataSet.builder().parameters(convertScriptStringToJsonString(Map.of("parameters", Map.of("inputText", inputData.getDocs().get(0))))).build();
-        } else {
-            return RemoteInferenceInputDataSet.builder().parameters(convertScriptStringToJsonString(Map.of("parameters", Map.of("inputText", inputData.getDocs().get(0), "inputImage", inputData.getDocs().get(1))))).build();
+        Map parametersMap = new HashMap<>();
+        parametersMap.put("inputText", inputData.getDocs().get(0));
+        if (inputData.getDocs().size() > 1) {
+            parametersMap.put("inputImage", inputData.getDocs().get(1));
         }
+        return RemoteInferenceInputDataSet.builder().parameters(convertScriptStringToJsonString(Map.of("parameters", parametersMap))).build();
+
     }
 }