tune parameter name

Signed-off-by: Yaliang Wu <[email protected]>
ylwu-amzn · Nov 3, 2023 · 0a3febc · 0a3febc
1 parent fd436d8
commit 0a3febc
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/...hms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteConnectorExecutor.java b/...hms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteConnectorExecutor.java
@@ -47,7 +47,7 @@ default ModelTensorOutput executePredict(MLInput mlInput) {
                 Map<String, String> parameters = getConnector().getParameters();
                 int stepSize = 1;
                 if (parameters != null) {
-                    stepSize = Integer.parseInt(Optional.ofNullable(parameters.get("text_docs_processed_step_size")).orElse("1"));
+                    stepSize = Integer.parseInt(Optional.ofNullable(parameters.get("input_docs_processed_step_size")).orElse("1"));
                 }
                 processedDocs += Math.max(tensorCount, stepSize);
                 tensorOutputs.addAll(tempTensorOutputs);

diff --git a/...c/test/java/org/opensearch/ml/engine/algorithms/remote/HttpJsonConnectorExecutorTest.java b/...c/test/java/org/opensearch/ml/engine/algorithms/remote/HttpJsonConnectorExecutorTest.java
@@ -209,7 +209,7 @@ public void executePredict_TextDocsInput_LessEmbeddingThanInputDocs() throws IOE
                 .postProcessFunction(MLPostProcessFunction.OPENAI_EMBEDDING)
                 .requestBody("{\"input\": ${parameters.input}}")
                 .build();
-        Map<String, String> parameters = ImmutableMap.of("text_docs_processed_step_size", "2");
+        Map<String, String> parameters = ImmutableMap.of("input_docs_processed_step_size", "2");
         HttpConnector connector = HttpConnector.builder().name("test connector").version("1").protocol("http").parameters(parameters).actions(Arrays.asList(predictAction)).build();
         HttpJsonConnectorExecutor executor = spy(new HttpJsonConnectorExecutor(connector));
         executor.setScriptService(scriptService);