Add check to multi-modal non image case

Signed-off-by: zane-neo <zaniu@amazon.com>
zane-neo · May 8, 2024 · 410f42b · 410f42b
1 parent 284bb98
commit 410f42b
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/...earch/ml/common/connector/functions/preprocess/MultiModalEmbeddingPreProcessFunction.java b/...earch/ml/common/connector/functions/preprocess/MultiModalEmbeddingPreProcessFunction.java
@@ -27,7 +27,10 @@ public void validate(MLInput mlInput) {
     @Override
     public RemoteInferenceInputDataSet process(MLInput mlInput) {
         TextDocsInputDataSet inputData = (TextDocsInputDataSet) mlInput.getInputDataset();
-        Map<String, Object> processedResult = Map.of("parameters", Map.of("inputText", inputData.getDocs().get(0), "inputImage", inputData.getDocs().get(1)));
-        return RemoteInferenceInputDataSet.builder().parameters(convertScriptStringToJsonString(processedResult)).build();
+        if (inputData.getDocs().size() == 1) {
+            return RemoteInferenceInputDataSet.builder().parameters(convertScriptStringToJsonString(Map.of("inputText", inputData.getDocs().get(0)))).build();
+        } else {
+            return RemoteInferenceInputDataSet.builder().parameters(convertScriptStringToJsonString(Map.of("inputText", inputData.getDocs().get(0), "inputImage", inputData.getDocs().get(1)))).build();
+        }
     }
 }