[Backport to 2.x]Add tokenizer and sparse encoding (#1301) (#1393)

common/src/main/java/org/opensearch/ml/common/CommonValue.java

@@ -42,6 +42,7 @@ public class CommonValue {
     public static final Integer ML_CONNECTOR_SCHEMA_VERSION = 2;
     public static final String ML_CONFIG_INDEX = ".plugins-ml-config";
     public static final Integer ML_CONFIG_INDEX_SCHEMA_VERSION = 2;
+    public static final String ML_MAP_RESPONSE_KEY = "response";
     public static final String USER_FIELD_MAPPING = "      \""
             + CommonValue.USER
             + "\": {\n"

common/src/main/java/org/opensearch/ml/common/FunctionName.java

@@ -17,6 +17,8 @@ public enum FunctionName {
     RCF_SUMMARIZE,
     LOGISTIC_REGRESSION,
     TEXT_EMBEDDING,
+    SPARSE_ENCODING,
+    SPARSE_TOKENIZE,
     METRICS_CORRELATION,
     REMOTE;
 
@@ -33,7 +35,7 @@ public static FunctionName from(String value) {
      * @return true for deep learning model.
      */
     public static boolean isDLModel(FunctionName functionName) {
-        if (functionName == TEXT_EMBEDDING) {
+        if (functionName == TEXT_EMBEDDING || functionName == SPARSE_ENCODING || functionName == SPARSE_TOKENIZE) {
             return true;
         }
         return false;

common/src/main/java/org/opensearch/ml/common/connector/AbstractConnector.java

@@ -23,6 +23,7 @@
 import java.util.Map;
 import java.util.Optional;
 
+import static org.opensearch.ml.common.CommonValue.ML_MAP_RESPONSE_KEY;
 import static org.opensearch.ml.common.utils.StringUtils.isJson;
 
 @Getter
@@ -101,7 +102,7 @@ public <T> void parseResponse(T response, List<ModelTensor> modelTensors, boolea
             return;
         }
         if (response instanceof String && isJson((String)response)) {
-            Map<String, Object> data = StringUtils.fromJson((String) response, "response");
+            Map<String, Object> data = StringUtils.fromJson((String) response, ML_MAP_RESPONSE_KEY);
             modelTensors.add(ModelTensor.builder().name("response").dataAsMap(data).build());
         } else {
             Map<String, Object> map = new HashMap<>();

common/src/main/java/org/opensearch/ml/common/input/MLInput.java

@@ -239,7 +239,7 @@ public static MLInput parse(XContentParser parser, String inputAlgoName) throws
             }
         }
         MLInputDataset inputDataSet = null;
-        if (algorithm == FunctionName.TEXT_EMBEDDING) {
+        if (algorithm == FunctionName.TEXT_EMBEDDING || algorithm == FunctionName.SPARSE_ENCODING || algorithm == FunctionName.SPARSE_TOKENIZE) {
             ModelResultFilter filter = new ModelResultFilter(returnBytes, returnNumber, targetResponse, targetResponsePositions);
             inputDataSet = new TextDocsInputDataSet(textDocs, filter);
         }

common/src/main/java/org/opensearch/ml/common/input/nlp/TextDocsMLInput.java

@@ -25,7 +25,7 @@
  * ML input class which supports a list fo text docs.
  * This class can be used for TEXT_EMBEDDING model.
  */
-@org.opensearch.ml.common.annotation.MLInput(functionNames = {FunctionName.TEXT_EMBEDDING})
+@org.opensearch.ml.common.annotation.MLInput(functionNames = {FunctionName.TEXT_EMBEDDING, FunctionName.SPARSE_ENCODING, FunctionName.SPARSE_TOKENIZE})
 public class TextDocsMLInput extends MLInput {
     public static final String TEXT_DOCS_FIELD = "text_docs";
     public static final String RESULT_FILTER_FIELD = "result_filter";

common/src/main/java/org/opensearch/ml/common/transport/register/MLRegisterModelInput.java

@@ -103,7 +103,7 @@ public MLRegisterModelInput(FunctionName functionName,
             if (modelFormat == null) {
                 throw new IllegalArgumentException("model format is null");
             }
-            if (url != null && modelConfig == null) {
+            if (url != null && modelConfig == null && functionName != FunctionName.SPARSE_TOKENIZE && functionName != FunctionName.SPARSE_ENCODING) { // The tokenize model doesn't require a model configuration. Currently, we only support one type of sparse model, which is pretrained, and it doesn't necessitate a model configuration.
                 throw new IllegalArgumentException("model config is null");
             }
         }

common/src/main/java/org/opensearch/ml/common/transport/upload_chunk/MLRegisterModelMetaInput.java

@@ -84,7 +84,7 @@ public MLRegisterModelMetaInput(String name, FunctionName functionName, String m
         if (modelContentHashValue == null) {
             throw new IllegalArgumentException("model content hash value is null");
         }
-        if (modelConfig == null) {
+        if (modelConfig == null && functionName != FunctionName.SPARSE_TOKENIZE && functionName != FunctionName.SPARSE_ENCODING) { // The tokenize model doesn't require a model configuration. Currently, we only support one type of sparse model, which is pretrained, and it doesn't necessitate a model configuration.
             throw new IllegalArgumentException("model config is null");
         }
         if (totalChunks == null) {

common/src/test/java/org/opensearch/ml/common/MLCommonsClassLoaderTests.java

@@ -149,21 +149,27 @@ public void testClassLoader_ExecuteOutputMCorr() throws IOException {
         assertArrayEquals(new long[]{1, 2}, metrics);
     }
 
-    @Test
-    public void testClassLoader_MLInput() throws IOException {
-        assertTrue(MLCommonsClassLoader.canInitMLInput(FunctionName.TEXT_EMBEDDING));
+    private void testClassLoader_MLInput_DlModel(FunctionName functionName) throws IOException {
+        assertTrue(MLCommonsClassLoader.canInitMLInput(functionName));
 
         String jsonStr = "{\"text_docs\":[\"doc1\",\"doc2\"],\"result_filter\":{\"return_bytes\":true,\"return_number\":true,\"target_response\":[\"field1\"], \"target_response_positions\": [2]}}";
         XContentParser parser = XContentType.JSON.xContent().createParser(new NamedXContentRegistry(new SearchModule(Settings.EMPTY,
                 Collections.emptyList()).getNamedXContents()), null, jsonStr);
         parser.nextToken();
 
-        TextDocsMLInput mlInput = MLCommonsClassLoader.initMLInput(FunctionName.TEXT_EMBEDDING, new Object[]{parser, FunctionName.TEXT_EMBEDDING}, XContentParser.class, FunctionName.class);
+        TextDocsMLInput mlInput = MLCommonsClassLoader.initMLInput(functionName, new Object[]{parser, functionName}, XContentParser.class, FunctionName.class);
         assertNotNull(mlInput);
-        assertEquals(FunctionName.TEXT_EMBEDDING, mlInput.getFunctionName());
+        assertEquals(functionName, mlInput.getFunctionName());
         assertEquals(2, ((TextDocsInputDataSet)mlInput.getInputDataset()).getDocs().size());
     }
 
+    @Test
+    public void testClassLoader_MLInput() throws IOException {
+        testClassLoader_MLInput_DlModel(FunctionName.TEXT_EMBEDDING);
+        testClassLoader_MLInput_DlModel(FunctionName.SPARSE_TOKENIZE);
+        testClassLoader_MLInput_DlModel(FunctionName.SPARSE_ENCODING);
+    }
+
     public enum TestEnum {
         TEST
     }

common/src/test/java/org/opensearch/ml/common/input/MLInputTest.java

@@ -13,8 +13,8 @@
 import org.opensearch.common.io.stream.BytesStreamOutput;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.common.settings.Settings;
-import org.opensearch.common.xcontent.XContentFactory;
 import org.opensearch.common.xcontent.XContentType;
+import org.opensearch.core.xcontent.MediaTypeRegistry;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.core.xcontent.XContentBuilder;
@@ -68,11 +68,11 @@ public class MLInputTest {
 
     @Before
     public void setUp() throws Exception {
-        final ColumnMeta[] columnMetas = new ColumnMeta[]{new ColumnMeta("test", ColumnType.DOUBLE)};
+        final ColumnMeta[] columnMetas = new ColumnMeta[] { new ColumnMeta("test", ColumnType.DOUBLE) };
         List<Row> rows = new ArrayList<>();
-        rows.add(new Row(new ColumnValue[]{new DoubleValue(1.0)}));
-        rows.add(new Row(new ColumnValue[]{new DoubleValue(2.0)}));
-        rows.add(new Row(new ColumnValue[]{new DoubleValue(3.0)}));
+        rows.add(new Row(new ColumnValue[] { new DoubleValue(1.0) }));
+        rows.add(new Row(new ColumnValue[] { new DoubleValue(2.0) }));
+        rows.add(new Row(new ColumnValue[] { new DoubleValue(3.0) }));
         DataFrame dataFrame = new DefaultDataFrame(columnMetas, rows);
         input = MLInput.builder()
                 .algorithm(algorithm)
@@ -96,35 +96,39 @@ public void parse_LinearRegression() throws IOException {
                 .searchSourceBuilder(new SearchSourceBuilder().query(new MatchAllQueryBuilder()).size(1))
                 .build();
         String expectedInputStr = "{\"algorithm\":\"LINEAR_REGRESSION\",\"input_index\":[\"index1\"],\"input_query\":{\"size\":1,\"query\":{\"match_all\":{\"boost\":1.0}}}}";
-        testParse(FunctionName.LINEAR_REGRESSION, inputDataset, expectedInputStr,  parsedInput -> {
+        testParse(FunctionName.LINEAR_REGRESSION, inputDataset, expectedInputStr, parsedInput -> {
             assertNotNull(parsedInput.getInputDataset());
-            assertEquals(1, ((SearchQueryInputDataset)parsedInput.getInputDataset()).getIndices().size());
-            assertEquals(indexName, ((SearchQueryInputDataset)parsedInput.getInputDataset()).getIndices().get(0));
+            assertEquals(1, ((SearchQueryInputDataset) parsedInput.getInputDataset()).getIndices().size());
+            assertEquals(indexName, ((SearchQueryInputDataset) parsedInput.getInputDataset()).getIndices().get(0));
         });
 
-        @NonNull DataFrame dataFrame = new DefaultDataFrame(new ColumnMeta[]{ColumnMeta.builder().name("value").columnType(ColumnType.FLOAT).build()});
-        dataFrame.appendRow(new Float[]{1.0f});
+        @NonNull
+        DataFrame dataFrame = new DefaultDataFrame(
+                new ColumnMeta[] { ColumnMeta.builder().name("value").columnType(ColumnType.FLOAT).build() });
+        dataFrame.appendRow(new Float[] { 1.0f });
         DataFrameInputDataset dataFrameInputDataset = DataFrameInputDataset.builder().dataFrame(dataFrame).build();
         expectedInputStr = "{\"algorithm\":\"LINEAR_REGRESSION\",\"input_data\":{\"column_metas\":[{\"name\":\"value\",\"column_type\":\"FLOAT\"}],\"rows\":[{\"values\":[{\"column_type\":\"FLOAT\",\"value\":1.0}]}]}}";
-        testParse(FunctionName.LINEAR_REGRESSION, dataFrameInputDataset, expectedInputStr,  parsedInput -> {
+        testParse(FunctionName.LINEAR_REGRESSION, dataFrameInputDataset, expectedInputStr, parsedInput -> {
             assertNotNull(parsedInput.getInputDataset());
-            assertEquals(1, ((DataFrameInputDataset)parsedInput.getInputDataset()).getDataFrame().size());
-            assertEquals(1.0f, ((DataFrameInputDataset)parsedInput.getInputDataset()).getDataFrame().getRow(0).getValue(0).floatValue(), 1e-5);
+            assertEquals(1, ((DataFrameInputDataset) parsedInput.getInputDataset()).getDataFrame().size());
+            assertEquals(1.0f, ((DataFrameInputDataset) parsedInput.getInputDataset()).getDataFrame().getRow(0)
+                    .getValue(0).floatValue(), 1e-5);
         });
     }
 
-    @Test
-    public void parse_TextEmbedding() throws IOException {
+    private void parse_NLPModel(FunctionName functionName) throws IOException {
         String sentence = "test sentence";
         String column = "column1";
         Integer position = 1;
         ModelResultFilter resultFilter = ModelResultFilter.builder()
                 .targetResponse(Arrays.asList(column))
                 .targetResponsePositions(Arrays.asList(position))
                 .build();
+
         TextDocsInputDataSet inputDataset = TextDocsInputDataSet.builder().docs(Arrays.asList(sentence)).resultFilter(resultFilter).build();
-        String expectedInputStr = "{\"algorithm\":\"TEXT_EMBEDDING\",\"text_docs\":[\"test sentence\"],\"return_bytes\":false,\"return_number\":false,\"target_response\":[\"column1\"],\"target_response_positions\":[1]}";
-        testParse(FunctionName.TEXT_EMBEDDING, inputDataset, expectedInputStr, parsedInput -> {
+        String expectedInputStr = "{\"algorithm\":\"functionName\",\"text_docs\":[\"test sentence\"],\"return_bytes\":false,\"return_number\":false,\"target_response\":[\"column1\"],\"target_response_positions\":[1]}";
+        expectedInputStr = expectedInputStr.replace("functionName", functionName.toString());
+        testParse(functionName, inputDataset, expectedInputStr, parsedInput -> {
             assertNotNull(parsedInput.getInputDataset());
             TextDocsInputDataSet parsedInputDataSet = (TextDocsInputDataSet) parsedInput.getInputDataset();
             assertEquals(1, parsedInputDataSet.getDocs().size());
@@ -136,35 +140,50 @@ public void parse_TextEmbedding() throws IOException {
     }
 
     @Test
-    public void parse_TextEmbedding_NullResultFilter() throws IOException {
+    public void parse_NLP_Related() throws IOException {
+        parse_NLPModel(FunctionName.TEXT_EMBEDDING);
+        parse_NLPModel(FunctionName.SPARSE_TOKENIZE);
+        parse_NLPModel(FunctionName.SPARSE_ENCODING);
+    }
+
+    private void parse_NLPModel_NullResultFilter(FunctionName functionName) throws IOException {
         String sentence = "test sentence";
         TextDocsInputDataSet inputDataset = TextDocsInputDataSet.builder().docs(Arrays.asList(sentence)).build();
-        String expectedInputStr = "{\"algorithm\":\"TEXT_EMBEDDING\",\"text_docs\":[\"test sentence\"]}";
-        testParse(FunctionName.TEXT_EMBEDDING, inputDataset, expectedInputStr, parsedInput -> {
+        String expectedInputStr = "{\"algorithm\":\"functionName\",\"text_docs\":[\"test sentence\"]}";
+        expectedInputStr = expectedInputStr.replace("functionName", functionName.toString());
+        testParse(functionName, inputDataset, expectedInputStr, parsedInput -> {
             assertNotNull(parsedInput.getInputDataset());
-            assertEquals(1, ((TextDocsInputDataSet)parsedInput.getInputDataset()).getDocs().size());
-            assertEquals(sentence, ((TextDocsInputDataSet)parsedInput.getInputDataset()).getDocs().get(0));
+            assertEquals(1, ((TextDocsInputDataSet) parsedInput.getInputDataset()).getDocs().size());
+            assertEquals(sentence, ((TextDocsInputDataSet) parsedInput.getInputDataset()).getDocs().get(0));
         });
     }
 
+
+    @Test
+    public void parse_NLPRelated_NullResultFilter() throws IOException {
+        parse_NLPModel_NullResultFilter(FunctionName.TEXT_EMBEDDING);
+        parse_NLPModel_NullResultFilter(FunctionName.SPARSE_TOKENIZE);
+        parse_NLPModel_NullResultFilter(FunctionName.SPARSE_ENCODING);
+    }
+
     private void testParse(FunctionName algorithm, MLInputDataset inputDataset, String expectedInputStr, Consumer<MLInput> verify) throws IOException {
         MLInput input = MLInput.builder().inputDataset(inputDataset).algorithm(algorithm).build();
-        XContentBuilder builder = XContentFactory.jsonBuilder();
+        XContentBuilder builder = MediaTypeRegistry.contentBuilder(XContentType.JSON);
         input.toXContent(builder, ToXContent.EMPTY_PARAMS);
         assertNotNull(builder);
         String jsonStr = builder.toString();
         assertEquals(expectedInputStr, jsonStr);
 
-        XContentParser parser = XContentType.JSON.xContent().createParser(new NamedXContentRegistry(new SearchModule(Settings.EMPTY,
-                Collections.emptyList()).getNamedXContents()), null, jsonStr);
+        XContentParser parser = XContentType.JSON.xContent()
+                .createParser(new NamedXContentRegistry(new SearchModule(Settings.EMPTY,
+                        Collections.emptyList()).getNamedXContents()), null, jsonStr);
         parser.nextToken();
         MLInput parsedInput = MLInput.parse(parser, algorithm.name());
         assertEquals(input.getFunctionName(), parsedInput.getFunctionName());
         assertEquals(input.getInputDataset().getInputDataType(), parsedInput.getInputDataset().getInputDataType());
         verify.accept(parsedInput);
     }
 
-
     @Test
     public void readInputStream_Success() throws IOException {
         readInputStream(input, parsedInput -> {

docs/model_serving_framework/text_embedding_model_examples.md

@@ -239,7 +239,7 @@ GET /_plugins/_ml/profile/models/zwla5YUB1qmVrJFlwzXJ
             "models": {
                 "zwla5YUB1qmVrJFlwzXJ": { # model id
                     "model_state": "LOADED",
-                    "predictor": "org.opensearch.ml.engine.algorithms.text_embedding.TextEmbeddingModel@1a0b0793",
+                    "predictor": "org.opensearch.ml.engine.algorithms.text_embedding.TextEmbeddingDenseModel@1a0b0793",
                     "target_worker_nodes": [ # plan to deploy model to these nodes
                         "0TLL4hHxRv6_G3n6y1l0BQ"
                     ],

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

@@ -190,7 +190,7 @@ public List downloadPrebuiltModelMetaList(String taskId, MLRegisterModelInput re
      * @param modelContentHash model content hash value
      * @param listener action listener
      */
-    public void downloadAndSplit(MLModelFormat modelFormat, String taskId, String modelName, String version, String url, String modelContentHash, ActionListener<Map<String, Object>> listener) {
+    public void downloadAndSplit(MLModelFormat modelFormat, String taskId, String modelName, String version, String url, String modelContentHash, FunctionName functionName, ActionListener<Map<String, Object>> listener) {
         try {
             AccessController.doPrivileged((PrivilegedExceptionAction<Void>) () -> {
                 Path registerModelPath = mlEngine.getRegisterModelPath(taskId, modelName, version);
@@ -199,7 +199,7 @@ public void downloadAndSplit(MLModelFormat modelFormat, String taskId, String mo
                 File modelZipFile = new File(modelPath);
                 log.debug("download model to file {}", modelZipFile.getAbsolutePath());
                 DownloadUtils.download(url, modelPath, new ProgressBar());
-                verifyModelZipFile(modelFormat, modelPath, modelName);
+                verifyModelZipFile(modelFormat, modelPath, modelName, functionName);
                 String hash = calculateFileHash(modelZipFile);
                 if (hash.equals(modelContentHash)) {
                     List<String> chunkFiles = splitFileIntoChunks(modelZipFile, modelPartsPath, CHUNK_SIZE);
@@ -221,7 +221,7 @@ public void downloadAndSplit(MLModelFormat modelFormat, String taskId, String mo
         }
     }
 
-    public void verifyModelZipFile(MLModelFormat modelFormat, String modelZipFilePath, String modelName) throws IOException {
+    public void verifyModelZipFile(MLModelFormat modelFormat, String modelZipFilePath, String modelName, FunctionName functionName) throws IOException {
         boolean hasPtFile = false;
         boolean hasOnnxFile = false;
         boolean hasTokenizerFile = false;
@@ -236,7 +236,7 @@ public void verifyModelZipFile(MLModelFormat modelFormat, String modelZipFilePat
                 }
             }
         }
-        if (!hasPtFile && !hasOnnxFile) {
+        if (!hasPtFile && !hasOnnxFile && functionName != FunctionName.SPARSE_TOKENIZE) { // sparse tokenizer model doesn't need model file.
             throw new IllegalArgumentException("Can't find model file");
         }
         if (!hasTokenizerFile) {