fine tune processed doc to always respect step size

Signed-off-by: Yaliang Wu <[email protected]>
opensearch-project · Nov 3, 2023 · 784cd4d · 784cd4d
1 parent 0a3febc
commit 784cd4d
Showing 1 changed file with 4 additions and 4 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
@@ -45,11 +45,11 @@ default ModelTensorOutput executePredict(MLInput mlInput) {
                 // This is to support some model which takes N text docs and embedding size is less than N-1.
                 // We need to tell executor what's the step size for each model run.
                 Map<String, String> parameters = getConnector().getParameters();
-                int stepSize = 1;
-                if (parameters != null) {
-                    stepSize = Integer.parseInt(Optional.ofNullable(parameters.get("input_docs_processed_step_size")).orElse("1"));
+                if (parameters != null && parameters.containsKey("input_docs_processed_step_size")) {
+                    processedDocs += Integer.parseInt(parameters.get("input_docs_processed_step_size"));
+                } else {
+                    processedDocs += Math.max(tensorCount, 1);
                 }
-                processedDocs += Math.max(tensorCount, stepSize);
                 tensorOutputs.addAll(tempTensorOutputs);
             }