Add tokenizer and sparse encoding

common/build.gradle

@@ -19,8 +19,8 @@ dependencies {
     testImplementation group: 'org.mockito', name: 'mockito-core', version: '4.4.0'
 
     compileOnly group: 'org.apache.commons', name: 'commons-text', version: '1.10.0'
-    compileOnly group: 'com.google.code.gson', name: 'gson', version: '2.10.1'
-    compileOnly group: 'org.json', name: 'json', version: '20230227'
+    implementation group: 'com.google.code.gson', name: 'gson', version: '2.10.1'
+    implementation group: 'org.json', name: 'json', version: '20230227'
 }
 
 lombok {

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,
+    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 == TOKENIZE) {
             return true;
         }
         return false;

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

@@ -104,9 +104,9 @@ public <T> void parseResponse(T response, List<ModelTensor> modelTensors, boolea
             Map<String, Object> data = StringUtils.fromJson((String) response, "response");
             modelTensors.add(ModelTensor.builder().name("response").dataAsMap(data).build());
         } else {
-            Map<String, Object> map = new HashMap<>();
-            map.put("response", response);
-            modelTensors.add(ModelTensor.builder().name("response").dataAsMap(map).build());
+        Map<String, Object> map = new HashMap<>();
+        map.put("response", response);
+        modelTensors.add(ModelTensor.builder().name("response").dataAsMap(map).build());
         }
     }
 

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.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.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

@@ -104,7 +104,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.TOKENIZE && functionName != FunctionName.SPARSE_ENCODING) {
                 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.TOKENIZE && functionName != FunctionName.SPARSE_ENCODING) {
             throw new IllegalArgumentException("model config is null");
         }
         if (totalChunks == null) {

common/src/main/java/org/opensearch/ml/common/utils/StringUtils.java

@@ -17,6 +17,7 @@
 import java.security.AccessController;
 import java.security.PrivilegedActionException;
 import java.security.PrivilegedExceptionAction;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;

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.TOKENIZE);
+        testClassLoader_MLInput_DlModel(FunctionName.SPARSE_ENCODING);
+    }
+
     public enum TestEnum {
         TEST
     }

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

@@ -110,19 +110,19 @@ public void parse_LinearRegression() throws IOException {
         });
     }
 
-    @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 -> {
+
+        TextDocsInputDataSet inputDataset = TextDocsInputDataSet.builder().docs(Arrays.asList(sentence)).resultFilter(resultFilter).build();
+        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());
@@ -134,19 +134,33 @@ 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.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));
         });
     }
 
-    private void testParse(FunctionName algorithm, MLInputDataset inputDataset, String expectedInputStr,
-            Consumer<MLInput> verify) throws IOException {
+
+    @Test
+    public void parse_NLPRelated_NullResultFilter() throws IOException {
+        parse_NLPModel_NullResultFilter(FunctionName.TEXT_EMBEDDING);
+        parse_NLPModel_NullResultFilter(FunctionName.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 = MediaTypeRegistry.contentBuilder(XContentType.JSON);
         input.toXContent(builder, ToXContent.EMPTY_PARAMS);

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

@@ -193,7 +193,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);
@@ -202,7 +202,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);
@@ -224,7 +224,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;
@@ -239,7 +239,7 @@ public void verifyModelZipFile(MLModelFormat modelFormat, String modelZipFilePat
                 }
             }
         }
-        if (!hasPtFile && !hasOnnxFile) {
+        if (!hasPtFile && !hasOnnxFile && functionName != FunctionName.TOKENIZE) {
             throw new IllegalArgumentException("Can't find model file");
         }
         if (!hasTokenizerFile) {