Repository: deepjavalibrary/djl
Pull request #1954 (2022-08-25 06:29:53)

[example] Refactor NMT example
Change-Id: Ib7703cafbcaace3a9a24b3ecad64c5d211754ba4

## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here


codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1954?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
> Merging [#1954](https://codecov.io/gh/deepjavalibrary/djl/pull/1954?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (cd6602b) into [master](https://codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (bb5073f) will **decrease** coverage by `2.11%`.
> The diff coverage is `67.18%`.

```diff
@@             Coverage Diff              @@
##             master    #1954      +/-   ##
============================================
- Coverage     72.08%   69.97%   -2.12%     
- Complexity     5126     5795     +669     
============================================
  Files           473      569      +96     
  Lines         21970    25721    +3751     
  Branches       2351     2777     +426     
============================================
+ Hits          15838    17997    +2159     
- Misses         4925     6359    +1434     
- Partials       1207     1365     +158     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1954?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [api/src/main/java/ai/djl/modality/cv/Image.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1954/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0ltYWdlLmphdmE=) | `69.23% <ø> (-4.11%)` | :arrow_down: |
| [...rc/main/java/ai/djl/modality/cv/MultiBoxPrior.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1954/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L011bHRpQm94UHJpb3IuamF2YQ==) | `76.00% <ø> (ø)` | |
| [...rc/main/java/ai/djl/modality/cv/output/Joints.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1954/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9Kb2ludHMuamF2YQ==) | `71.42% <ø> (ø)` | |
| [.../main/java/ai/djl/modality/cv/output/Landmark.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1954/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9MYW5kbWFyay5qYXZh) | `100.00% <ø> (ø)` | |
| [...main/java/ai/djl/modality/cv/output/Rectangle.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1954/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9SZWN0YW5nbGUuamF2YQ==) | `72.41% <0.00%> (ø)` | |
| [...i/djl/modality/cv/translator/BigGANTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1954/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvci5qYXZh) | `21.42% <0.00%> (-5.24%)` | :arrow_down: |
| [...odality/cv/translator/BigGANTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1954/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvckZhY3RvcnkuamF2YQ==) | `33.33% <0.00%> (+8.33%)` | :arrow_up: |
| [...nslator/InstanceSegmentationTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1954/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW5zdGFuY2VTZWdtZW50YXRpb25UcmFuc2xhdG9yRmFjdG9yeS5qYXZh) | `14.28% <0.00%> (-3.90%)` | :arrow_down: |
| [.../cv/translator/SemanticSegmentationTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1954/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvU2VtYW50aWNTZWdtZW50YXRpb25UcmFuc2xhdG9yLmphdmE=) | `0.00% <0.00%> (ø)` | |
| [.../cv/translator/StyleTransferTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1954/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvU3R5bGVUcmFuc2ZlclRyYW5zbGF0b3JGYWN0b3J5LmphdmE=) | `40.00% <ø> (ø)` | |
| ... and [470 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1954/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)


Patch:
diff --git a/api/src/main/java/ai/djl/modality/nlp/translator/TextEmbeddingServingTranslator.java b/api/src/main/java/ai/djl/modality/nlp/translator/TextEmbeddingServingTranslator.java
new file mode 100644
index 000000000..b83624fb8
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/translator/TextEmbeddingServingTranslator.java
@@ -0,0 +1,64 @@
+/*
+ * Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.nlp.translator;
+
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.ndarray.NDList;
+import ai.djl.translate.Batchifier;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorContext;
+import ai.djl.util.JsonUtils;
+
+/** A {@link Translator} that can handle generic text embedding {@link Input} and {@link Output}. */
+public class TextEmbeddingServingTranslator implements Translator<Input, Output> {
+
+    private Translator<String, float[]> translator;
+
+    /**
+     * Constructs a {@code TextEmbeddingServingTranslator} instance.
+     *
+     * @param translator a {@code Translator} processes text embedding input
+     */
+    public TextEmbeddingServingTranslator(Translator<String, float[]> translator) {
+        this.translator = translator;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Batchifier getBatchifier() {
+        return translator.getBatchifier();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void prepare(TranslatorContext ctx) throws Exception {
+        translator.prepare(ctx);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, Input input) throws Exception {
+        String text = input.getData().getAsString();
+        return translator.processInput(ctx, text);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Output processOutput(TranslatorContext ctx, NDList list) throws Exception {
+        float[] ret = translator.processOutput(ctx, list);
+        Output output = new Output();
+        output.add(JsonUtils.GSON_PRETTY.toJson(ret));
+        return output;
+    }
+}
diff --git a/bom/README.md b/bom/README.md
index ff04b2a35..a161972fe 100644
--- a/bom/README.md
+++ b/bom/README.md
@@ -57,11 +57,6 @@ will need to mention the type as pom and the scope as import) as the following:
         <groupId>ai.djl.mxnet</groupId>
         <artifactId>mxnet-engine</artifactId>
     </dependency>
-    <dependency>
-        <groupId>ai.djl.mxnet</groupId>
-        <artifactId>mxnet-native-auto</artifactId>
-        <scope>runtime</scope>
-    </dependency>
 </dependencies>
 ```
 
diff --git a/docs/development/dependency_management.md b/docs/development/dependency_management.md
index a53483d2f..cb7d064ef 100644
--- a/docs/development/dependency_management.md
+++ b/docs/development/dependency_management.md
@@ -17,7 +17,9 @@ Although DJL itself has many packages, but in a real production environment, the
 
 - ai.djl:api
 - one of the engine package (e.g. ai.djl.pytorch:pytorch-engine)
-- native library of selected engine (e.g. ai.djl.pytorch:pytorch-native-auto)
+- native library of selected engine for no internet use case 
+  (e.g. ai.djl.pytorch:pytorch-native-cpu:1.11.0:linux-x86_64 + ai.djl.pytorch:pytorch-jni).
+  See [How to distribute DJL application](https://github.com/deepjavalibrary/djl-demo/tree/master/development/fatjar#readme) for more detail.
 
 We don't recommend include more than one engine into your project unless you need both of them. DJL will load
 all engines into memory as long as they are in classpath, and those engines usually consume significant of memory.
diff --git a/engines/pytorch/pytorch-engine/README.md b/engines/pytorch/pytorch-engine/README.md
index 9c56491c6..9b022d88e 100644
--- a/engines/pytorch/pytorch-engine/README.md
+++ b/engines/pytorch/pytorch-engine/README.md
@@ -39,7 +39,7 @@ Since DJL 0.14.0, `pytorch-engine` can load older version of pytorch native libr
 ways to specify PyTorch version:
 
 1. Explicitly specify `pytorch-native-xxx` package version to override the version in the [BOM](../../../bom/README.md).
-2. Sets environment variable: `PYTORCH_VERSION` to override the package version defined by `pytorch-native-auto`.
+2. Sets environment variable: `PYTORCH_VERSION` to override the default package version.
 
 The following table illustrates which pytorch version that DJL supports:
 
diff --git a/examples/README.md b/examples/README.md
index b5c7a7c7e..4b756913d 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -95,14 +95,13 @@ mvn exec:java
 
 ## Engine selection
 
-DJL is engine agnostic, so it's capable of supporting different backends.
+DJL is engine agnostic, so it's capable of supporting different backends. DJL by default will select
+proper native library for you automatically and download those libraries from internet. If your
+production environment doesn't have network access, you can distribute DJL's offline native packages
+together with your application to avoid download engine native libraries at runtime.
 
 With Apache MXNet, PyTorch, TensorFlow and ONNX Runtime, you can choose different
 builds of the native library.
-We recommend the automatic engine selection which downloads the best engine for your
-platform and available hardware during the first runtime.
-Activate the automatic selection by adding `ai.djl.mxnet:mxnet-native-auto:1.8.0`
-for Apache MXNet, and `ai.djl.pytorch:pytorch-native-auto:1.9.1` for PyTorch as a dependency.
 You can also see:
 
 - [MXNet Engine](../engines/mxnet/mxnet-engine/README.md)
diff --git a/examples/src/main/java/ai/djl/examples/inference/NeuralMachineTranslation.java b/examples/src/main/java/ai/djl/examples/inference/NeuralMachineTranslation.java
index d24ebc888..8bed71c7a 100644
--- a/examples/src/main/java/ai/djl/examples/inference/NeuralMachineTranslation.java
+++ b/examples/src/main/java/ai/djl/examples/inference/NeuralMachineTranslation.java
@@ -14,6 +14,7 @@
 package ai.djl.examples.inference;
 
 import ai.djl.ModelException;
+import ai.djl.engine.Engine;
 import ai.djl.modality.nlp.preprocess.PunctuationSeparator;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
@@ -23,11 +24,10 @@ import ai.djl.nn.Block;
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ZooModel;
 import ai.djl.training.ParameterStore;
-import ai.djl.training.util.ProgressBar;
 import ai.djl.translate.TranslateException;
+import ai.djl.util.JsonUtils;
+import ai.djl.util.Utils;
 
-import com.google.gson.Gson;
-import com.google.gson.internal.LinkedTreeMap;
 import com.google.gson.reflect.TypeToken;
 
 import org.slf4j.Logger;
@@ -41,6 +41,7 @@ import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -60,7 +61,38 @@ public final class NeuralMachineTranslation {
     }
 
     public static void predict() throws ModelException, TranslateException, IOException {
-        String french = "trop tard";
+        Path path = Paths.get("src/test/resources/source_wrd2idx.json");
+        Map<String, Long> wrd2idx;
+        try (InputStream is = Files.newInputStream(path)) {
+            String json = Utils.toString(is);
+            Type mapType = new TypeToken<Map<String, Long>>() {}.getType();
+            wrd2idx = JsonUtils.GSON.fromJson(json, mapType);
+        }
+
+        path = Paths.get("src/test/resources/target_idx2wrd.json");
+        Map<String, String> idx2wrd;
+        try (InputStream is = Files.newInputStream(path)) {
+            String json = Utils.toString(is);
+            Type mapType = new TypeToken<Map<String, String>>() {}.getType();
+            idx2wrd = JsonUtils.GSON.fromJson(json, mapType);
+        }
+
+        Engine engine = Engine.getEngine("PyTorch");
+        try (NDManager manager = engine.newBaseManager()) {
+            try (ZooModel<NDList, NDList> encoder = getEncoderModel();
+                    ZooModel<NDList, NDList> decoder = getDecoderModel()) {
+
+                String french = "trop tard";
+                NDList toDecode = predictEncoder(french, encoder, wrd2idx, manager);
+                String english = predictDecoder(toDecode, decoder, idx2wrd, manager);
+
+                logger.info("French: {}", french);
+                logger.info("English: {}", english);
+            }
+        }
+    }
+
+    public static ZooModel<NDList, NDList> getEncoderModel() throws ModelException, IOException {
         String url =
                 "https://resources.djl.ai/demo/pytorch/android/neural_machine_translation/optimized_encoder_150k.zip";
 
@@ -70,39 +102,31 @@ public final class NeuralMachineTranslation {
                         .optModelUrls(url)
                         .optModelName("optimized_encoder_150k.ptl")
                         .optEngine("PyTorch")
-                        .optProgress(new ProgressBar())
                         .build();
+        return criteria.loadModel();
+    }
 
-        String english;
-        try (ZooModel<NDList, NDList> model = criteria.loadModel()) {
-            try (NDManager manager = model.getNDManager().newSubManager()) {
-                NDList toDecode = predictEncoder(french, manager, model);
-                english = predictDecoder(toDecode, manager);
-            }
-        }
+    public static ZooModel<NDList, NDList> getDecoderModel() throws ModelException, IOException {
+        String url =
+                "https://resources.djl.ai/demo/pytorch/android/neural_machine_translation/optimized_decoder_150k.zip";
 
-        logger.info("French: {}", french);
-        logger.info("English: {}", english);
+        Criteria<NDList, NDList> criteria =
+                Criteria.builder()
+                        .setTypes(NDList.class, NDList.class)
+                        .optModelUrls(url)
+                        .optModelName("optimized_decoder_150k.ptl")
+                        .optEngine("PyTorch")
+                        .build();
+        return criteria.loadModel();
     }
 
-    private static NDList predictEncoder(
-            String text, NDManager manager, ZooModel<NDList, NDList> model)
-            throws ModelException, IOException {
-
-        // reads in the french words to be mapped to
-        Path words = Paths.get("src/test/resources/source_wrd2idx.json");
-        InputStream is = Files.newInputStream(words);
-        int size = is.available();
-        byte[] buffer = new byte[size];
-        is.read(buffer);
-        is.close();
-        String json = new String(buffer, "UTF-8");
-        Type mapType = new TypeToken<Map<String, Long>>() {}.getType();
-        LinkedTreeMap<String, Long> wrd2idx = new Gson().fromJson(json, mapType);
-
+    public static NDList predictEncoder(
+            String text,
+            ZooModel<NDList, NDList> model,
+            Map<String, Long> wrd2idx,
+            NDManager manager) {
         // maps french input to id's from french file
-        List<String> list = new ArrayList<>();
-        list.add(text);
+        List<String> list = Collections.singletonList(text);
         PunctuationSeparator punc = new PunctuationSeparator();
         list = punc.preprocess(list);
         List<Long> inputs = new ArrayList<>();
@@ -112,12 +136,12 @@ public final class NeuralMachineTranslation {
             }
             Long id = wrd2idx.get(word.toLowerCase(Locale.FRENCH));
             if (id == null) {
-                throw new ModelException("Word \"" + word + "\" not found.");
+                throw new IllegalArgumentException("Word \"" + word + "\" not found.");
             }
             inputs.add(id);
         }
 
-        // for prediction
+        // for forwarding the model
         Shape inputShape = new Shape(1);
         Shape hiddenShape = new Shape(1, 1, 256);
         FloatBuffer fb = FloatBuffer.allocate(256);
@@ -140,76 +164,55 @@ public final class NeuralMachineTranslation {
         }
         outputTensorBuffer.rewind();
         NDArray outputsTensor = manager.create(outputTensorBuffer, new Shape(outputsShape));
+
         return new NDList(outputsTensor, hiddenTensor);
     }
 
-    private static String predictDecoder(NDList toDecode, NDManager manager)
-            throws ModelException, IOException {
-        String url =
-                "https://resources.djl.ai/demo/pytorch/android/neural_machine_translation/optimized_decoder_150k.zip";
-
-        Criteria<NDList, NDList> criteria =
-                Criteria.builder()
-                        .setTypes(NDList.class, NDList.class)
-                        .optModelUrls(url)
-                        .optModelName("optimized_decoder_150k.ptl")
-                        .optEngine("PyTorch")
-                        .optProgress(new ProgressBar())
-                        .build();
+    public static String predictDecoder(
+            NDList toDecode,
+            ZooModel<NDList, NDList> model,
+            Map<String, String> idx2wrd,
+            NDManager manager) {
+        // for forwarding the model
+        Shape decoderInputShape = new Shape(1, 1);
+        NDArray inputTensor = manager.create(new long[] {0}, decoderInputShape);
+        ArrayList<Integer> result = new ArrayList<>(MAX_LENGTH);
+        NDArray outputsTensor = toDecode.get(0);
+        NDArray hiddenTensor = toDecode.get(1);
 
-        StringBuilder english = new StringBuilder();
-        try (ZooModel<NDList, NDList> model = criteria.loadModel()) {
-            // for forwarding the model
-            Shape decoderInputShape = new Shape(1, 1);
-            NDArray inputTensor = manager.create(new long[] {0}, decoderInputShape);
-            ArrayList<Integer> result = new ArrayList<>(MAX_LENGTH);
-            NDArray outputsTensor = toDecode.get(0);
-            NDArray hiddenTensor = toDecode.get(1);
-
-            // for using the model
-            Block block = model.getBlock();
-            ParameterStore ps = new ParameterStore();
-
-            // loops through forwarding of each word
-            for (int i = 0; i < MAX_LENGTH; i++) {
-                NDList inputTensorList = new NDList(inputTensor, hiddenTensor, outputsTensor);
-                NDList outputs = block.forward(ps, inputTensorList, false);
-                NDArray outputTensor = outputs.get(0);
-                hiddenTensor = outputs.get(1);
-                float[] buf = outputTensor.toFloatArray();
-                int topIdx = 0;
-                double topVal = -Double.MAX_VALUE;
-                for (int j = 0; j < buf.length; j++) {
-                    if (buf[j] > topVal) {
-                        topVal = buf[j];
-                        topIdx = j;
-                    }
-                }
+        // for using the model
+        Block block = model.getBlock();
+        ParameterStore ps = new ParameterStore();
 
-                if (topIdx == EOS_TOKEN) {
-                    break;
+        // loops through forwarding of each word
+        for (int i = 0; i < MAX_LENGTH; i++) {
+            NDList inputTensorList = new NDList(inputTensor, hiddenTensor, outputsTensor);
+            NDList outputs = block.forward(ps, inputTensorList, false);
+            NDArray outputTensor = outputs.get(0);
+            hiddenTensor = outputs.get(1);
+            float[] buf = outputTensor.toFloatArray();
+            int topIdx = 0;
+            double topVal = -Double.MAX_VALUE;
+            for (int j = 0; j < buf.length; j++) {
+                if (buf[j] > topVal) {
+                    topVal = buf[j];
+                    topIdx = j;
                 }
+            }
 
-                result.add(topIdx);
-                inputTensor = manager.create(new long[] {topIdx}, decoderInputShape);
+            if (topIdx == EOS_TOKEN) {
+                break;
             }
 
-            // read in english words to be mapped to
-            Path words = Paths.get("src/test/resources/target_idx2wrd.json");
-            InputStream is = Files.newInputStream(words);
-            int size = is.available();
-            byte[] buffer = new byte[size];
-            is.read(buffer);
-            is.close();
-            String json = new String(buffer, "UTF-8");
-            Type mapType = new TypeToken<Map<String, String>>() {}.getType();
-            LinkedTreeMap<String, String> idx2wrd = new Gson().fromJson(json, mapType);
+            result.add(topIdx);
+            inputTensor = manager.create(new long[] {topIdx}, decoderInputShape);
+        }
 
-            // map english words and create output string
-            for (Integer word : result) {
-                english.append(' ').append(idx2wrd.get(word.toString()));
-            }
+        StringBuilder sb = new StringBuilder();
+        // map english words and create output string
+        for (Integer word : result) {
+            sb.append(idx2wrd.get(word.toString())).append(' ');
         }
-        return english.toString();
+        return sb.toString().trim();
     }
 }
diff --git a/examples/src/main/java/ai/djl/examples/inference/SemanticSegmentation.java b/examples/src/main/java/ai/djl/examples/inference/SemanticSegmentation.java
index 381cae625..92a677484 100644
--- a/examples/src/main/java/ai/djl/examples/inference/SemanticSegmentation.java
+++ b/examples/src/main/java/ai/djl/examples/inference/SemanticSegmentation.java
@@ -17,6 +17,8 @@ import ai.djl.ModelException;
 import ai.djl.inference.Predictor;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.ImageFactory;
+import ai.djl.modality.cv.transform.Normalize;
+import ai.djl.modality.cv.transform.ToTensor;
 import ai.djl.modality.cv.translator.SemanticSegmentationTranslator;
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ZooModel;
@@ -30,8 +32,6 @@ import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * An example of inference using a semantic segmentation model.
@@ -55,13 +55,16 @@ public final class SemanticSegmentation {
         Path imageFile = Paths.get("src/test/resources/dog_bike_car.jpg");
         Image img = ImageFactory.getInstance().fromFile(imageFile);
 
+        float[] mean = {0.485f, 0.456f, 0.406f};
+        float[] std = {0.229f, 0.224f, 0.225f};
+
         String url =
                 "https://mlrepo.djl.ai/model/cv/semantic_segmentation/ai/djl/pytorch/deeplabv3/0.0.1/deeplabv3.zip";
-        Map<String, String> arguments = new ConcurrentHashMap<>();
-        arguments.put("toTensor", "true");
-        arguments.put("normalize", "true");
         SemanticSegmentationTranslator translator =
-                SemanticSegmentationTranslator.builder(arguments).build();
+                SemanticSegmentationTranslator.builder()
+                        .addTransform(new ToTensor())
+                        .addTransform(new Normalize(mean, std))
+                        .build();
 
         Criteria<Image, Image> criteria =
                 Criteria.builder()
@@ -72,11 +75,10 @@ public final class SemanticSegmentation {
                         .optEngine("PyTorch")
                         .optProgress(new ProgressBar())
                         .build();
-        try (ZooModel<Image, Image> model = criteria.loadModel()) {
-            try (Predictor<Image, Image> predictor = model.newPredictor()) {
-                Image semanticImage = predictor.predict(img);
-                saveSemanticImage(semanticImage);
-            }
+        try (ZooModel<Image, Image> model = criteria.loadModel();
+                Predictor<Image, Image> predictor = model.newPredictor()) {
+            Image semanticImage = predictor.predict(img);
+            saveSemanticImage(semanticImage);
         }
     }
 
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
new file mode 100644
index 000000000..0187ae3fa
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
@@ -0,0 +1,147 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.huggingface.translator;
+
+import ai.djl.huggingface.tokenizers.Encoding;
+import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.translate.ArgumentsUtil;
+import ai.djl.translate.Batchifier;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorContext;
+
+import java.io.IOException;
+import java.util.Map;
+
+/** The translator for Huggingface text embedding model. */
+public class TextEmbeddingTranslator implements Translator<String, float[]> {
+
+    private static final int[] AXIS = {0};
+
+    private HuggingFaceTokenizer tokenizer;
+    private Batchifier batchifier;
+
+    TextEmbeddingTranslator(HuggingFaceTokenizer tokenizer, Batchifier batchifier) {
+        this.tokenizer = tokenizer;
+        this.batchifier = batchifier;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Batchifier getBatchifier() {
+        return batchifier;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, String input) {
+        NDManager manager = ctx.getNDManager();
+        Encoding encoding = tokenizer.encode(input);
+        ctx.setAttachment("encoding", encoding);
+        long[] indices = encoding.getIds();
+        long[] attentionMask = encoding.getAttentionMask();
+        NDList ndList = new NDList(2);
+        ndList.add(manager.create(indices));
+        ndList.add(manager.create(attentionMask));
+        return ndList;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public float[] processOutput(TranslatorContext ctx, NDList list) {
+        NDArray embeddings = list.get("last_hidden_state");
+        Encoding encoding = (Encoding) ctx.getAttachment("encoding");
+        long[] attentionMask = encoding.getAttentionMask();
+        NDManager manager = ctx.getNDManager();
+        NDArray inputAttentionMask = manager.create(attentionMask).toType(DataType.FLOAT32, true);
+        long[] shape = embeddings.getShape().getShape();
+        inputAttentionMask = inputAttentionMask.tile(shape[shape.length - 1]);
+        inputAttentionMask = inputAttentionMask.reshape(embeddings.getShape());
+        NDArray inputAttentionMaskSum = inputAttentionMask.sum(AXIS);
+        NDArray clamp = inputAttentionMaskSum.clip(1e-9, 1e12);
+        NDArray prod = embeddings.mul(inputAttentionMask);
+        NDArray sum = prod.sum(AXIS);
+        embeddings = sum.div(clamp).normalize(2, 0);
+
+        return embeddings.toFloatArray();
+    }
+
+    /**
+     * Creates a builder to build a {@code TextEmbeddingTranslator}.
+     *
+     * @param tokenizer the tokenizer
+     * @return a new builder
+     */
+    public static Builder builder(HuggingFaceTokenizer tokenizer) {
+        return new Builder(tokenizer);
+    }
+
+    /**
+     * Creates a builder to build a {@code TextEmbeddingTranslator}.
+     *
+     * @param tokenizer the tokenizer
+     * @param arguments the models' arguments
+     * @return a new builder
+     */
+    public static Builder builder(HuggingFaceTokenizer tokenizer, Map<String, ?> arguments) {
+        Builder builder = builder(tokenizer);
+        builder.configure(arguments);
+
+        return builder;
+    }
+
+    /** The builder for token classification translator. */
+    public static final class Builder {
+
+        private HuggingFaceTokenizer tokenizer;
+        private Batchifier batchifier = Batchifier.STACK;
+
+        Builder(HuggingFaceTokenizer tokenizer) {
+            this.tokenizer = tokenizer;
+        }
+
+        /**
+         * Sets the {@link Batchifier} for the {@link Translator}.
+         *
+         * @param batchifier true to include token types
+         * @return this builder
+         */
+        public TextEmbeddingTranslator.Builder optBatchifier(Batchifier batchifier) {
+            this.batchifier = batchifier;
+            return this;
+        }
+
+        /**
+         * Configures the builder with the model arguments.
+         *
+         * @param arguments the model arguments
+         */
+        public void configure(Map<String, ?> arguments) {
+            String batchifierStr = ArgumentsUtil.stringValue(arguments, "batchifier", "stack");
+            optBatchifier(Batchifier.fromString(batchifierStr));
+        }
+
+        /**
+         * Builds the translator.
+         *
+         * @return the new translator
+         * @throws IOException if I/O error occurs
+         */
+        public TextEmbeddingTranslator build() throws IOException {
+            return new TextEmbeddingTranslator(tokenizer, batchifier);
+        }
+    }
+}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java
new file mode 100644
index 000000000..87d692abe
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java
@@ -0,0 +1,72 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.huggingface.translator;
+
+import ai.djl.Model;
+import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.nlp.translator.TextEmbeddingServingTranslator;
+import ai.djl.translate.TranslateException;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorFactory;
+import ai.djl.util.Pair;
+
+import java.io.IOException;
+import java.lang.reflect.Type;
+import java.nio.file.Path;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+/** A {@link TranslatorFactory} that creates a {@link TextEmbeddingTranslator} instance. */
+public class TextEmbeddingTranslatorFactory implements TranslatorFactory {
+
+    private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>();
+
+    static {
+        SUPPORTED_TYPES.add(new Pair<>(String.class, float[].class));
+        SUPPORTED_TYPES.add(new Pair<>(Input.class, Output.class));
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Set<Pair<Type, Type>> getSupportedTypes() {
+        return SUPPORTED_TYPES;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Translator<?, ?> newInstance(
+            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
+            throws TranslateException {
+        Path modelPath = model.getModelPath();
+        try {
+            HuggingFaceTokenizer tokenizer =
+                    HuggingFaceTokenizer.builder(arguments)
+                            .optTokenizerPath(modelPath)
+                            .optManager(model.getNDManager())
+                            .build();
+            TextEmbeddingTranslator translator =
+                    TextEmbeddingTranslator.builder(tokenizer, arguments).build();
+            if (input == String.class && output == float[].class) {
+                return translator;
+            } else if (input == Input.class && output == Output.class) {
+                return new TextEmbeddingServingTranslator(translator);
+            }
+            throw new IllegalArgumentException("Unsupported input/output types.");
+        } catch (IOException e) {
+            throw new TranslateException("Failed to load tokenizer.", e);
+        }
+    }
+}
diff --git a/extensions/tokenizers/src/main/python/fill_mask_converter.py b/extensions/tokenizers/src/main/python/fill_mask_converter.py
index f9b1dadbd..7b06552ff 100644
--- a/extensions/tokenizers/src/main/python/fill_mask_converter.py
+++ b/extensions/tokenizers/src/main/python/fill_mask_converter.py
@@ -56,3 +56,6 @@ class FillMaskConverter(HuggingfaceConverter):
     def encode_inputs(self, tokenizer):
         text = self.inputs.replace("[MASK]", tokenizer.mask_token)
         return tokenizer.encode_plus(text, return_tensors='pt')
+
+    def get_extra_arguments(self, hf_pipeline) -> dict:
+        return {"maskToken": hf_pipeline.tokenizer.mask_token}
diff --git a/extensions/tokenizers/src/main/python/huggingface_converter.py b/extensions/tokenizers/src/main/python/huggingface_converter.py
index f965a7652..6fdd53f22 100644
--- a/extensions/tokenizers/src/main/python/huggingface_converter.py
+++ b/extensions/tokenizers/src/main/python/huggingface_converter.py
@@ -57,7 +57,7 @@ class HuggingfaceConverter:
             return False, reason, -1
 
         size = self.save_to_model_zoo(model_id, output_dir, temp_dir,
-                                      hf_pipeline.tokenizer.mask_token)
+                                      hf_pipeline)
 
         return True, None, size
 
@@ -100,7 +100,7 @@ class HuggingfaceConverter:
         return model_file
 
     def save_to_model_zoo(self, model_id: str, output_dir: str, temp_dir: str,
-                          mask_token: str):
+                          hf_pipeline):
         artifact_ids = model_id.split("/")
         model_name = artifact_ids[-1]
 
@@ -111,13 +111,14 @@ class HuggingfaceConverter:
 
         # Save serving.properties
         serving_file = os.path.join(temp_dir, "serving.properties")
+        arguments = self.get_extra_arguments(hf_pipeline)
         with open(serving_file, 'w') as f:
             f.write(f"engine=PyTorch\n"
                     f"option.modelName={model_name}\n"
                     f"option.mapLocation=true\n"
                     f"translatorFactory={self.translator}\n")
-            if mask_token:
-                f.write(f"maskToken={mask_token}\n")
+            for k, v in arguments.items():
+                f.write(f"{k}={v}\n")
 
         # Save model as .zip file
         logging.info(f"Saving DJL model as zip: {model_name}.zip ...")
@@ -157,6 +158,9 @@ class HuggingfaceConverter:
 
         return self.verify_jit_output(hf_pipeline, encoding, out)
 
+    def get_extra_arguments(self, hf_pipeline) -> dict:
+        return {}
+
     def verify_jit_output(self, hf_pipeline, encoding, out):
         if not hasattr(out, "last_hidden_layer"):
             return False, f"Unexpected inference result: {out}"
diff --git a/extensions/tokenizers/src/main/python/sentence_similarity_converter.py b/extensions/tokenizers/src/main/python/sentence_similarity_converter.py
index baa0382a1..dbcb587a2 100644
--- a/extensions/tokenizers/src/main/python/sentence_similarity_converter.py
+++ b/extensions/tokenizers/src/main/python/sentence_similarity_converter.py
@@ -51,3 +51,6 @@ class SentenceSimilarityConverter(HuggingfaceConverter):
             return False, f"Unexpected inference result: {last_hidden_state}"
 
         return True, None
+
+    def get_extra_arguments(self, hf_pipeline) -> dict:
+        return {"padding": "true"}


Test patch:
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TranslatorTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TranslatorTest.java
index 2cad1353b..88c72d41c 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TranslatorTest.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TranslatorTest.java
@@ -16,6 +16,7 @@ import ai.djl.Model;
 import ai.djl.ModelException;
 import ai.djl.huggingface.translator.FillMaskTranslatorFactory;
 import ai.djl.huggingface.translator.QuestionAnsweringTranslatorFactory;
+import ai.djl.huggingface.translator.TextEmbeddingTranslatorFactory;
 import ai.djl.huggingface.translator.TokenClassificationTranslatorFactory;
 import ai.djl.inference.Predictor;
 import ai.djl.modality.Classifications;
@@ -26,6 +27,7 @@ import ai.djl.modality.nlp.translator.NamedEntity;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.Block;
 import ai.djl.nn.LambdaBlock;
 import ai.djl.repository.zoo.Criteria;
@@ -325,4 +327,83 @@ public class TranslatorTest {
                     () -> factory.newInstance(String.class, Integer.class, model, arguments));
         }
     }
+
+    @Test
+    public void testTextEmbeddingTranslator()
+            throws ModelException, IOException, TranslateException {
+        TestRequirements.notArm();
+
+        String text = "This is an example sentence";
+
+        Block block =
+                new LambdaBlock(
+                        a -> {
+                            NDManager manager = a.getManager();
+                            NDArray arr = manager.ones(new Shape(1, 7, 384));
+                            arr.setName("last_hidden_state");
+                            return new NDList(arr);
+                        },
+                        "model");
+        Path modelDir = Paths.get("build/model");
+        Files.createDirectories(modelDir);
+
+        Criteria<String, float[]> criteria =
+                Criteria.builder()
+                        .setTypes(String.class, float[].class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-uncased")
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
+                        .build();
+
+        try (ZooModel<String, float[]> model = criteria.loadModel();
+                Predictor<String, float[]> predictor = model.newPredictor()) {
+            float[] res = predictor.predict(text);
+            Assert.assertEquals(res.length, 384);
+            Assertions.assertAlmostEquals(res[0], 0.05103);
+        }
+
+        Criteria<Input, Output> criteria2 =
+                Criteria.builder()
+                        .setTypes(Input.class, Output.class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-uncased")
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
+                        .build();
+
+        try (ZooModel<Input, Output> model = criteria2.loadModel();
+                Predictor<Input, Output> predictor = model.newPredictor()) {
+            Input input = new Input();
+            input.add(text);
+            Output out = predictor.predict(input);
+            float[] res = JsonUtils.GSON.fromJson(out.getAsString(0), float[].class);
+            Assert.assertEquals(res.length, 384);
+            Assertions.assertAlmostEquals(res[0], 0.05103);
+        }
+
+        try (Model model = Model.newInstance("test")) {
+            model.setBlock(block);
+            Map<String, String> options = new HashMap<>();
+            options.put("hasParameter", "false");
+            model.load(modelDir, "test", options);
+
+            TextEmbeddingTranslatorFactory factory = new TextEmbeddingTranslatorFactory();
+            Map<String, String> arguments = new HashMap<>();
+
+            Assert.assertThrows(
+                    TranslateException.class,
+                    () -> factory.newInstance(String.class, Integer.class, model, arguments));
+
+            arguments.put("tokenizer", "bert-base-uncased");
+
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> factory.newInstance(String.class, Integer.class, model, arguments));
+        }
+    }
 }