fix test

Signed-off-by: Yaliang Wu <[email protected]>
zane-neo · Jun 12, 2024 · e82d89e · e82d89e
1 parent 38d8ddf
commit e82d89e
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/...earch/ml/common/connector/functions/preprocess/MultiModalConnectorPreProcessFunction.java b/...earch/ml/common/connector/functions/preprocess/MultiModalConnectorPreProcessFunction.java
@@ -11,6 +11,7 @@
 import org.opensearch.ml.common.dataset.remote.RemoteInferenceInputDataSet;
 import org.opensearch.ml.common.input.MLInput;
 
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 

diff --git a/plugin/src/test/java/org/opensearch/ml/rest/RestBedRockInferenceIT.java b/plugin/src/test/java/org/opensearch/ml/rest/RestBedRockInferenceIT.java
@@ -73,7 +73,7 @@ public void test_bedrock_multimodal_model() throws Exception {
 
             TextDocsInputDataSet inputDataSet = TextDocsInputDataSet.builder().docs(List.of("hello", imageBase64)).build();
             MLInput mlInput = MLInput.builder().inputDataset(inputDataSet).algorithm(FunctionName.TEXT_EMBEDDING).build();
-            Map inferenceResult = predictTextEmbedding(modelId, mlInput);
+            Map inferenceResult = predictTextEmbeddingModel(modelId, mlInput);
             assertTrue(errorMsg, inferenceResult.containsKey("inference_results"));
             List output = (List) inferenceResult.get("inference_results");
             assertEquals(errorMsg, 1, output.size());