fix multiple docs support (opensearch-project#1516)

Signed-off-by: Yaliang Wu <[email protected]>
rbhavna · Nov 16, 2023 · 0711f17 · 0711f17
1 parent bd02cfe
commit 0711f17
Show file tree

Hide file tree

Showing 2 changed files with 7 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
@@ -37,7 +37,11 @@ default ModelTensorOutput executePredict(MLInput mlInput) {
  List<String> textDocs = textDocsInputDataSet.getDocs().subList(processedDocs, textDocsInputDataSet.getDocs().size());
  List<ModelTensors> tempTensorOutputs = new ArrayList<>();
  preparePayloadAndInvokeRemoteModel(MLInput.builder().algorithm(FunctionName.TEXT_EMBEDDING).inputDataset(TextDocsInputDataSet.builder().docs(textDocs).build()).build(), tempTensorOutputs);
- processedDocs += Math.max(tempTensorOutputs.size(), 1);
+ int tensorCount = 0;
+ if (tempTensorOutputs.size() > 0 && tempTensorOutputs.get(0).getMlModelTensors() != null) {
+ tensorCount = tempTensorOutputs.get(0).getMlModelTensors().size();
+ }
+ processedDocs += Math.max(tensorCount, 1);
  tensorOutputs.addAll(tempTensorOutputs);
  }
  } else {

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
@@ -213,7 +213,8 @@ public void executePredict_TextDocsInput() throws IOException {
  when(executor.getHttpClient()).thenReturn(httpClient);
  MLInputDataset inputDataSet = TextDocsInputDataSet.builder().docs(Arrays.asList("test doc1", "test doc2")).build();
  ModelTensorOutput modelTensorOutput = executor.executePredict(MLInput.builder().algorithm(FunctionName.REMOTE).inputDataset(inputDataSet).build());
- Assert.assertEquals(2, modelTensorOutput.getMlModelOutputs().size());
+ Assert.assertEquals(1, modelTensorOutput.getMlModelOutputs().size());
+ Assert.assertEquals(2, modelTensorOutput.getMlModelOutputs().get(0).getMlModelTensors().size());
  Assert.assertEquals("sentence_embedding", modelTensorOutput.getMlModelOutputs().get(0).getMlModelTensors().get(0).getName());
  Assert
  .assertArrayEquals(