From 784cd4db5acefb85ce6ff9b65d8ead9cbfc784f1 Mon Sep 17 00:00:00 2001 From: Yaliang Wu Date: Fri, 3 Nov 2023 00:26:55 -0700 Subject: [PATCH] fine tune processed doc to always respect step size Signed-off-by: Yaliang Wu --- .../engine/algorithms/remote/RemoteConnectorExecutor.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteConnectorExecutor.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteConnectorExecutor.java index 14841bc3a9..4e312bd16e 100644 --- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteConnectorExecutor.java +++ b/ml-algorithms/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 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); }