Text Similarity

common/src/main/java/org/opensearch/ml/common/settings/MLCommonsSettings.java

@@ -76,6 +76,15 @@ private MLCommonsSettings() {}
             Setting.Property.NodeScope,
             Setting.Property.Dynamic
         );
+    public static final Setting<Integer> ML_COMMONS_TEXT_SIMILARITY_BATCH_SIZE = Setting
+        .intSetting(
+            ML_PLUGIN_SETTING_PREFIX + "text_similarity_batch_size",
+            500,
+            1,
+            1000,
+            Setting.Property.NodeScope,
+            Setting.Property.Dynamic
+        );
     public static final Setting<Integer> ML_COMMONS_MAX_DEPLOY_MODEL_TASKS_PER_NODE = Setting
         .intSetting(
             ML_PLUGIN_SETTING_PREFIX + "max_deploy_model_tasks_per_node",

ml-algorithms/src/main/java/org/opensearch/ml/engine/MLEngine.java

@@ -12,6 +12,7 @@
 import java.util.Locale;
 import java.util.Map;
 
+import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.ml.common.FunctionName;
 import org.opensearch.ml.common.MLModel;
@@ -51,11 +52,24 @@ public class MLEngine {
 
     private Encryptor encryptor;
 
-    public MLEngine(Path opensearchDataFolder, Encryptor encryptor) {
+    @Getter
+    private ClusterService clusterService;
+
+    public MLEngine(Path opensearchDataFolder, Encryptor encryptor, ClusterService clusterService) {
         this.mlCachePath = opensearchDataFolder.resolve("ml_cache");
         this.mlModelsCachePath = mlCachePath.resolve("models_cache");
         this.mlConfigPath = mlCachePath.resolve("config");
         this.encryptor = encryptor;
+        this.clusterService = clusterService;
+    }
+
+    /**
+     * @deprecated Retained for backward compatibility. Scheduled for removal. <br/>
+     * Use {@link #MLEngine(Path, Encryptor, ClusterService)} instead.
+     */
+    @Deprecated(forRemoval = true)
+    public MLEngine(Path opensearchDataFolder, Encryptor encryptor) {
+        this(opensearchDataFolder, encryptor, null);
     }
 
     public String getPrebuiltModelMetaListPath() {

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/DLModel.java

@@ -23,6 +23,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 
 import org.apache.commons.io.FileUtils;
+import org.opensearch.common.settings.Settings;
 import org.opensearch.ml.common.FunctionName;
 import org.opensearch.ml.common.MLModel;
 import org.opensearch.ml.common.exception.MLException;
@@ -342,4 +343,10 @@ public ModelTensors parseModelTensorOutput(Output output, ModelResultFilter resu
         return tensorOutput;
     }
 
+    public Settings getClusterSettings() {
+        if (mlEngine.getClusterService() != null) {
+            return mlEngine.getClusterService().getSettings();
+        }
+        return Settings.EMPTY;
+    }
 }

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/text_similarity/TextSimilarityCrossEncoderModel.java

@@ -20,13 +20,15 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import org.opensearch.common.settings.Settings;
 import org.opensearch.ml.common.FunctionName;
 import org.opensearch.ml.common.dataset.MLInputDataset;
 import org.opensearch.ml.common.dataset.TextSimilarityInputDataSet;
 import org.opensearch.ml.common.input.MLInput;
 import org.opensearch.ml.common.model.MLModelConfig;
 import org.opensearch.ml.common.output.model.ModelTensorOutput;
 import org.opensearch.ml.common.output.model.ModelTensors;
+import org.opensearch.ml.common.settings.MLCommonsSettings;
 import org.opensearch.ml.engine.algorithms.DLModel;
 import org.opensearch.ml.engine.annotation.Function;
 
@@ -43,16 +45,31 @@ public class TextSimilarityCrossEncoderModel extends DLModel {
     public ModelTensorOutput predict(String modelId, MLInput mlInput) throws TranslateException {
         MLInputDataset inputDataSet = mlInput.getInputDataset();
         List<ModelTensors> tensorOutputs = new ArrayList<>();
-        Output output;
         TextSimilarityInputDataSet textSimInput = (TextSimilarityInputDataSet) inputDataSet;
         String queryText = textSimInput.getQueryText();
-        for (String doc : textSimInput.getTextDocs()) {
-            Input input = new Input();
-            input.add(queryText);
-            input.add(doc);
-            output = getPredictor().predict(input);
-            ModelTensors outputTensors = ModelTensors.fromBytes(output.getData().getAsBytes());
-            tensorOutputs.add(outputTensors);
+        List<String> textDocs = textSimInput.getTextDocs();
+
+        Settings clusterSettings = getClusterSettings();
+        final int batchSize = MLCommonsSettings.ML_COMMONS_TEXT_SIMILARITY_BATCH_SIZE.get(clusterSettings);
+
+        for (int i = 0; i < textDocs.size(); i += batchSize) {
+            int endIndex = Math.min(i + batchSize, textDocs.size());
+            List<String> batchDocs = textDocs.subList(i, endIndex);
+            List<Input> batchInputs = new ArrayList<>(batchDocs.size());
+
+            for (String doc : batchDocs) {
+                Input input = new Input();
+                input.add(queryText);
+                input.add(doc);
+                batchInputs.add(input);
+            }
+
+            List<Output> batchOutputs = getPredictor().batchPredict(batchInputs);
+
+            for (Output output : batchOutputs) {
+                ModelTensors outputTensors = ModelTensors.fromBytes(output.getData().getAsBytes());
+                tensorOutputs.add(outputTensors);
+            }
         }
         return new ModelTensorOutput(tensorOutputs);
     }

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/text_similarity/TextSimilarityTranslator.java

@@ -35,6 +35,7 @@
 import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.DataType;
 import ai.djl.translate.TranslatorContext;
+import ai.djl.util.PairList;
 
 public class TextSimilarityTranslator extends SentenceTransformerTranslator {
     public final String SIMILARITY_NAME = "similarity";
@@ -95,4 +96,73 @@ public Output processOutput(TranslatorContext ctx, NDList list) {
         return output;
     }
 
+    @Override
+    public NDList batchProcessInput(TranslatorContext ctx, List<Input> inputs) {
+        NDManager manager = ctx.getNDManager();
+        int batchSize = inputs.size();
+        List<String> sentences = new ArrayList<>(batchSize);
+        List<String> contexts = new ArrayList<>(batchSize);
+        for (Input input : inputs) {
+            String sentence = input.getAsString(0);
+            String context = input.getAsString(1);
+            sentences.add(sentence);
+            contexts.add(context);
+        }
+        // Tokenize in batches
+        Encoding[] encodings = tokenizer.batchEncode(new PairList<>(sentences, contexts));
+        int seqLen = encodings[0].getIds().length;
+        for (Encoding enc : encodings) {
+            seqLen = Math.max(seqLen, enc.getIds().length);
+        }
+        long[][] inputIds = new long[batchSize][seqLen];
+        long[][] attentionMasks = new long[batchSize][seqLen];
+        long[][] tokenTypeIds = new long[batchSize][seqLen];
+        for (int i = 0; i < batchSize; i++) {
+            inputIds[i] = encodings[i].getIds();
+            attentionMasks[i] = encodings[i].getAttentionMask();
+            tokenTypeIds[i] = encodings[i].getTypeIds();
+        }
+        NDArray inputIdsArray = manager.create(inputIds);
+        inputIdsArray.setName("input_ids");
+        NDArray attentionMaskArray = manager.create(attentionMasks);
+        attentionMaskArray.setName("attention_mask");
+        NDArray tokenTypeArray = manager.create(tokenTypeIds);
+        tokenTypeArray.setName("token_type_ids");
+        NDList ndList = new NDList();
+        ndList.add(inputIdsArray);
+        ndList.add(attentionMaskArray);
+        ndList.add(tokenTypeArray);
+        return ndList;
+    }
+
+    @Override
+    public List<Output> batchProcessOutput(TranslatorContext ctx, NDList list) {
+        NDArray batchArray = list.getFirst();
+        int batchSize = (int) batchArray.getShape().get(0);
+        List<Output> outputs = new ArrayList<>(batchSize);
+        for (int i = 0; i < batchSize; i++) {
+            NDArray itemArray = batchArray.get(i);
+
+            Number[] itemData = itemArray.toArray();
+            long[] itemShape = itemArray.getShape().getShape();
+            DataType dataType = itemArray.getDataType();
+            MLResultDataType mlResultDataType = MLResultDataType.valueOf(dataType.name());
+            ByteBuffer itemBuffer = itemArray.toByteBuffer();
+
+            ModelTensor tensor = ModelTensor
+                .builder()
+                .name(SIMILARITY_NAME)
+                .data(itemData)
+                .shape(itemShape)
+                .dataType(mlResultDataType)
+                .byteBuffer(itemBuffer)
+                .build();
+
+            ModelTensors modelTensorOutput = new ModelTensors(List.of(tensor));
+            Output output = new Output(200, "OK");
+            output.add(modelTensorOutput.toBytes());
+            outputs.add(output);
+        }
+        return outputs;
+    }
 }

ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/text_similarity/TextSimilarityCrossEncoderModelTest.java

@@ -29,8 +29,10 @@
 import java.io.File;
 import java.io.IOException;
 import java.net.URISyntaxException;
+import java.nio.ByteBuffer;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
@@ -173,6 +175,80 @@ public void test_TextSimilarity_Translator_ProcessOutput() throws URISyntaxExcep
         assertEquals(1, data.length);
     }
 
+    @Test
+    public void test_TextSimilarity_Translator_BatchProcessInput() throws URISyntaxException, IOException {
+        TextSimilarityTranslator textSimilarityTranslator = new TextSimilarityTranslator();
+        TranslatorContext translatorContext = mock(TranslatorContext.class);
+        Model mlModel = mock(Model.class);
+        when(translatorContext.getModel()).thenReturn(mlModel);
+        when(mlModel.getModelPath()).thenReturn(Paths.get(getClass().getResource("../tokenize/tokenizer.json").toURI()).getParent());
+        textSimilarityTranslator.prepare(translatorContext);
+
+        NDManager manager = mock(NDManager.class);
+        when(translatorContext.getNDManager()).thenReturn(manager);
+        Input input = mock(Input.class);
+        String testSentence = "hello world";
+        when(input.getAsString(0)).thenReturn(testSentence);
+        when(input.getAsString(1)).thenReturn(testSentence);
+        NDArray indiceNdArray = mock(NDArray.class);
+        when(indiceNdArray.toLongArray()).thenReturn(new long[] { 102l, 101l });
+        when(manager.create((long[][]) any())).thenReturn(indiceNdArray);
+        doNothing().when(indiceNdArray).setName(any());
+        List<Input> inputList = new ArrayList<>(1);
+        inputList.add(input);
+        NDList outputList = textSimilarityTranslator.batchProcessInput(translatorContext, inputList);
+        assertEquals(3, outputList.size());
+        Iterator<NDArray> iterator = outputList.iterator();
+        while (iterator.hasNext()) {
+            NDArray ndArray = iterator.next();
+            long[] output = ndArray.toLongArray();
+            assertEquals(2, output.length);
+        }
+    }
+
+    @Test
+    public void test_TextSimilarity_Translator_BatchProcessOutput() throws URISyntaxException, IOException {
+        TextSimilarityTranslator textSimilarityTranslator = new TextSimilarityTranslator();
+        TranslatorContext translatorContext = mock(TranslatorContext.class);
+        Model mlModel = mock(Model.class);
+        when(translatorContext.getModel()).thenReturn(mlModel);
+        when(mlModel.getModelPath()).thenReturn(Paths.get(getClass().getResource("../tokenize/tokenizer.json").toURI()).getParent());
+        textSimilarityTranslator.prepare(translatorContext);
+
+        NDArray batchArray = mock(NDArray.class);
+        Shape batchShape = mock(Shape.class);
+        when(batchArray.getShape()).thenReturn(batchShape);
+        when(batchShape.get(0)).thenReturn(2L);
+
+        NDArray itemArray1 = mock(NDArray.class);
+        NDArray itemArray2 = mock(NDArray.class);
+        Shape itemShape = mock(Shape.class);
+        when(itemShape.getShape()).thenReturn(new long[] { 1 });
+        when(itemArray1.getShape()).thenReturn(itemShape);
+        when(itemArray2.getShape()).thenReturn(itemShape);
+        when(itemArray1.toArray()).thenReturn(new Number[] { 1.0f });
+        when(itemArray2.toArray()).thenReturn(new Number[] { 2.0f });
+        when(itemArray1.getDataType()).thenReturn(DataType.FLOAT32);
+        when(itemArray2.getDataType()).thenReturn(DataType.FLOAT32);
+        when(itemArray1.toByteBuffer()).thenReturn(ByteBuffer.allocate(4));
+        when(itemArray2.toByteBuffer()).thenReturn(ByteBuffer.allocate(4));
+        when(batchArray.get(0)).thenReturn(itemArray1);
+        when(batchArray.get(1)).thenReturn(itemArray2);
+
+        NDList ndList = new NDList(batchArray);
+        List<Output> outputs = textSimilarityTranslator.batchProcessOutput(translatorContext, ndList);
+        assertEquals(2, outputs.size());
+        for (Output output : outputs) {
+            byte[] bytes = output.getData().getAsBytes();
+            ModelTensors tensorOutput = ModelTensors.fromBytes(bytes);
+            List<ModelTensor> modelTensorsList = tensorOutput.getMlModelTensors();
+            assertEquals(1, modelTensorsList.size());
+            ModelTensor modelTensor = modelTensorsList.get(0);
+            assertEquals("similarity", modelTensor.getName());
+            assertEquals(1, modelTensor.getData().length);
+        }
+    }
+
     @Test
     public void initModel_predict_TorchScript_CrossEncoder() throws URISyntaxException {
         textSimilarityCrossEncoderModel.initModel(model, params, encryptor);

plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java

@@ -682,7 +682,7 @@ public Collection<Object> createComponents(
 
         encryptor = new EncryptorImpl(clusterService, client, sdkClient, mlIndicesHandler);
 
-        mlEngine = new MLEngine(dataPath, encryptor);
+        mlEngine = new MLEngine(dataPath, encryptor, clusterService);
         nodeHelper = new DiscoveryNodeHelper(clusterService, settings);
         modelCacheHelper = new MLModelCacheHelper(clusterService, settings);
         cmHandler = new OpenSearchConversationalMemoryHandler(client, clusterService);
@@ -1364,7 +1364,8 @@ public List<Setting<?>> getSettings() {
                 MLCommonsSettings.ML_COMMONS_INDEX_INSIGHT_FEATURE_ENABLED,
                 MLCommonsSettings.REMOTE_METADATA_GLOBAL_TENANT_ID,
                 MLCommonsSettings.REMOTE_METADATA_GLOBAL_RESOURCE_CACHE_TTL,
-                MLCommonsSettings.ML_COMMONS_STREAM_ENABLED
+                MLCommonsSettings.ML_COMMONS_STREAM_ENABLED,
+                MLCommonsSettings.ML_COMMONS_TEXT_SIMILARITY_BATCH_SIZE
             );
         return settings;
     }