Repository: deepjavalibrary/djl
Pull request #2310 (2023-01-09 17:18:17)

[tokenizer] Adds batch support for TextEmbeddingTranslator
## 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/2310?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Base: **72.08**% // Head: **74.38**% // Increases project coverage by **`+2.29%`** :tada:
> Coverage data is based on head [(`9089e7e`)](https://codecov.io/gh/deepjavalibrary/djl/pull/2310?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) compared to base [(`bb5073f`)](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).
> Patch coverage: 75.03% of modified lines in pull request are covered.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #2310      +/-   ##
============================================
+ Coverage     72.08%   74.38%   +2.29%     
- Complexity     5126     6694    +1568     
============================================
  Files           473      657     +184     
  Lines         21970    29105    +7135     
  Branches       2351     3012     +661     
============================================
+ Hits          15838    21649    +5811     
- Misses         4925     5988    +1063     
- Partials       1207     1468     +261     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/2310?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/2310?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/2310?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/2310?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/2310?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/2310?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/2310?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: |
| [.../modality/cv/translator/ImageFeatureExtractor.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2310?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW1hZ2VGZWF0dXJlRXh0cmFjdG9yLmphdmE=) | `0.00% <0.00%> (ø)` | |
| [.../ai/djl/modality/cv/translator/YoloTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2310?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1RyYW5zbGF0b3IuamF2YQ==) | `27.77% <0.00%> (+18.95%)` | :arrow_up: |
| [...ain/java/ai/djl/modality/cv/util/NDImageUtils.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2310?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3V0aWwvTkRJbWFnZVV0aWxzLmphdmE=) | `67.10% <0.00%> (+7.89%)` | :arrow_up: |
| [api/src/main/java/ai/djl/modality/nlp/Decoder.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2310?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9EZWNvZGVyLmphdmE=) | `63.63% <ø> (ø)` | |
| ... and [600 more](https://codecov.io/gh/deepjavalibrary/djl/pull/2310?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)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2310?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-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/audio/Audio.java b/api/src/main/java/ai/djl/modality/audio/Audio.java
index 4ad949845..429359fc8 100644
--- a/api/src/main/java/ai/djl/modality/audio/Audio.java
+++ b/api/src/main/java/ai/djl/modality/audio/Audio.java
@@ -63,9 +63,9 @@ public class Audio {
     }
 
     /**
-     * Get the number of channels of an audio file.
+     * Returns the number of channels of an audio file.
      *
-     * @return The number of channels of an audio file.
+     * @return the number of channels of an audio file
      */
     public int getChannels() {
         return channels;
diff --git a/api/src/main/java/ai/djl/modality/audio/AudioFactory.java b/api/src/main/java/ai/djl/modality/audio/AudioFactory.java
index 310134c57..97d6446c7 100644
--- a/api/src/main/java/ai/djl/modality/audio/AudioFactory.java
+++ b/api/src/main/java/ai/djl/modality/audio/AudioFactory.java
@@ -36,9 +36,16 @@ public abstract class AudioFactory {
         "ai.djl.audio.FFmpegAudioFactory", "ai.djl.modality.audio.SampledAudioFactory"
     };
 
-    private static AudioFactory factory = newInstance();
+    protected int channels;
+    protected int sampleRate;
+    protected int sampleFormat;
 
-    private static AudioFactory newInstance() {
+    /**
+     * Constructs a new instance of {@code AudioFactory}.
+     *
+     * @return a new instance of {@code AudioFactory}
+     */
+    public static AudioFactory newInstance() {
         for (String f : FACTORIES) {
             try {
                 Class<? extends AudioFactory> clazz =
@@ -51,15 +58,6 @@ public abstract class AudioFactory {
         throw new IllegalStateException("Failed to create AudioFactory!");
     }
 
-    /**
-     * Gets new instance of Audio factory from the provided factory implementation.
-     *
-     * @return {@link AudioFactory}
-     */
-    public static AudioFactory getInstance() {
-        return factory;
-    }
-
     /**
      * Returns {@link Audio} from file.
      *
@@ -112,7 +110,9 @@ public abstract class AudioFactory {
      * @param data the raw data in float array form.
      * @return {@link Audio}
      */
-    public abstract Audio fromData(float[] data);
+    public Audio fromData(float[] data) {
+        return new Audio(data);
+    }
 
     /**
      * Returns {@link Audio} from {@link NDArray}.
@@ -120,5 +120,67 @@ public abstract class AudioFactory {
      * @param array the NDArray with CHW format
      * @return {@link Audio}
      */
-    public abstract Audio fromNDArray(NDArray array);
+    public Audio fromNDArray(NDArray array) {
+        throw new UnsupportedOperationException("Not supported!");
+    }
+
+    /**
+     * Sets the number of channels for {@link AudioFactory} to use.
+     *
+     * @param channels the number of channels for {@link AudioFactory} to use
+     * @return this factory
+     */
+    public AudioFactory setChannels(int channels) {
+        this.channels = channels;
+        return this;
+    }
+
+    /**
+     * Returns the channels of this factory.
+     *
+     * @return the channels of this factory
+     */
+    public int getChannels() {
+        return channels;
+    }
+
+    /**
+     * Sets the sampleRate for {@link AudioFactory} to use.
+     *
+     * @param sampleRate the sampleRate for {@link AudioFactory} to use
+     * @return this factory
+     */
+    public AudioFactory setSampleRate(int sampleRate) {
+        this.sampleRate = sampleRate;
+        return this;
+    }
+
+    /**
+     * Returns the sample rate.
+     *
+     * @return the sample rate in integer
+     */
+    public int getSampleRate() {
+        return sampleRate;
+    }
+
+    /**
+     * Sets the audio sample format for {@link AudioFactory} to use.
+     *
+     * @param sampleFormat the sample format
+     * @return this factory.
+     */
+    public AudioFactory setSampleFormat(int sampleFormat) {
+        this.sampleFormat = sampleFormat;
+        return this;
+    }
+
+    /**
+     * Returns the sample format name of the audio.
+     *
+     * @return the format name of the audio
+     */
+    public int getSampleFormat() {
+        return sampleFormat;
+    }
 }
diff --git a/api/src/main/java/ai/djl/modality/audio/SampledAudioFactory.java b/api/src/main/java/ai/djl/modality/audio/SampledAudioFactory.java
index cdbeaccf2..04ce1691f 100644
--- a/api/src/main/java/ai/djl/modality/audio/SampledAudioFactory.java
+++ b/api/src/main/java/ai/djl/modality/audio/SampledAudioFactory.java
@@ -13,7 +13,6 @@
 package ai.djl.modality.audio;
 
 import ai.djl.modality.cv.ImageFactory;
-import ai.djl.ndarray.NDArray;
 
 import java.io.BufferedInputStream;
 import java.io.IOException;
@@ -37,16 +36,34 @@ import javax.sound.sampled.UnsupportedAudioFileException;
  */
 public class SampledAudioFactory extends AudioFactory {
 
+    /** {@inheritDoc} */
+    @Override
+    public AudioFactory setChannels(int channel) {
+        if (channel != 0) {
+            throw new UnsupportedOperationException("SampledAudioFactory only support 1 channel.");
+        }
+        return this;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public AudioFactory setSampleRate(int sampleRate) {
+        throw new UnsupportedOperationException("Not supported.");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public AudioFactory setSampleFormat(int sampleFormat) {
+        throw new UnsupportedOperationException("Not supported.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public Audio fromFile(Path path) throws IOException {
         try (AudioInputStream ais = AudioSystem.getAudioInputStream(path.toFile())) {
             AudioFormat format = ais.getFormat();
             byte[] bytes = read(ais);
-            float[] floats =
-                    bytesToFloats(
-                            bytes,
-                            format.isBigEndian() ? ByteOrder.BIG_ENDIAN : ByteOrder.LITTLE_ENDIAN);
+            float[] floats = bytesToFloats(bytes, format.isBigEndian());
             return new Audio(floats, format.getSampleRate(), format.getChannels());
         } catch (UnsupportedAudioFileException e) {
             throw new IOException("Unsupported Audio file", e);
@@ -59,28 +76,13 @@ public class SampledAudioFactory extends AudioFactory {
         try (AudioInputStream ais = AudioSystem.getAudioInputStream(new BufferedInputStream(is))) {
             AudioFormat format = ais.getFormat();
             byte[] bytes = read(ais);
-            float[] floats =
-                    bytesToFloats(
-                            bytes,
-                            format.isBigEndian() ? ByteOrder.BIG_ENDIAN : ByteOrder.LITTLE_ENDIAN);
+            float[] floats = bytesToFloats(bytes, format.isBigEndian());
             return new Audio(floats, format.getSampleRate(), format.getChannels());
         } catch (UnsupportedAudioFileException e) {
             throw new IOException("Unsupported Audio file", e);
         }
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public Audio fromData(float[] data) {
-        return new Audio(data);
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public Audio fromNDArray(NDArray array) {
-        throw new UnsupportedOperationException("Not supported!");
-    }
-
     private byte[] read(AudioInputStream ais) throws IOException {
         AudioFormat format = ais.getFormat();
         int frameSize = format.getFrameSize();
@@ -102,7 +104,8 @@ public class SampledAudioFactory extends AudioFactory {
         return ret;
     }
 
-    private float[] bytesToFloats(byte[] bytes, ByteOrder order) {
+    private float[] bytesToFloats(byte[] bytes, boolean isBigEndian) {
+        ByteOrder order = isBigEndian ? ByteOrder.BIG_ENDIAN : ByteOrder.LITTLE_ENDIAN;
         ShortBuffer buffer = ByteBuffer.wrap(bytes).order(order).asShortBuffer();
         short[] shorts = new short[buffer.capacity()];
         buffer.get(shorts);
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
index b83624fb8..193a32afe 100644
--- a/api/src/main/java/ai/djl/modality/nlp/translator/TextEmbeddingServingTranslator.java
+++ b/api/src/main/java/ai/djl/modality/nlp/translator/TextEmbeddingServingTranslator.java
@@ -16,14 +16,16 @@ import ai.djl.modality.Input;
 import ai.djl.modality.Output;
 import ai.djl.ndarray.NDList;
 import ai.djl.translate.Batchifier;
+import ai.djl.translate.NoBatchifyTranslator;
 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> {
+public class TextEmbeddingServingTranslator implements NoBatchifyTranslator<Input, Output> {
 
     private Translator<String, float[]> translator;
+    private Translator<String[], float[][]> batchTranslator;
 
     /**
      * Constructs a {@code TextEmbeddingServingTranslator} instance.
@@ -32,12 +34,7 @@ public class TextEmbeddingServingTranslator implements Translator<Input, Output>
      */
     public TextEmbeddingServingTranslator(Translator<String, float[]> translator) {
         this.translator = translator;
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public Batchifier getBatchifier() {
-        return translator.getBatchifier();
+        this.batchTranslator = translator.toBatchTranslator();
     }
 
     /** {@inheritDoc} */
@@ -49,16 +46,35 @@ public class TextEmbeddingServingTranslator implements Translator<Input, Output>
     /** {@inheritDoc} */
     @Override
     public NDList processInput(TranslatorContext ctx, Input input) throws Exception {
+        String contentType = input.getProperty("Content-Type", null);
         String text = input.getData().getAsString();
-        return translator.processInput(ctx, text);
+        if ("application/json".equals(contentType)) {
+            ctx.setAttachment("batch", Boolean.TRUE);
+            String[] inputs = JsonUtils.GSON.fromJson(text, String[].class);
+            return batchTranslator.processInput(ctx, inputs);
+        }
+        NDList ret = translator.processInput(ctx, text);
+        Batchifier batchifier = translator.getBatchifier();
+        if (batchifier != null) {
+            NDList[] batch = {ret};
+            return batchifier.batchify(batch);
+        }
+        return ret;
     }
 
     /** {@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));
+        if (ctx.getAttachment("batch") != null) {
+            output.add(JsonUtils.GSON_PRETTY.toJson(batchTranslator.processOutput(ctx, list)));
+        } else {
+            Batchifier batchifier = translator.getBatchifier();
+            if (batchifier != null) {
+                list = batchifier.unbatchify(list)[0];
+            }
+            output.add(JsonUtils.GSON_PRETTY.toJson(translator.processOutput(ctx, list)));
+        }
         return output;
     }
 }
diff --git a/api/src/main/java/ai/djl/translate/Translator.java b/api/src/main/java/ai/djl/translate/Translator.java
index 43022c4b4..91268995a 100644
--- a/api/src/main/java/ai/djl/translate/Translator.java
+++ b/api/src/main/java/ai/djl/translate/Translator.java
@@ -13,6 +13,11 @@
 package ai.djl.translate;
 
 import ai.djl.inference.Predictor;
+import ai.djl.ndarray.NDList;
+
+import java.lang.reflect.Array;
+import java.util.ArrayList;
+import java.util.List;
 
 /**
  * The {@code Translator} interface provides model pre-processing and postprocessing functionality.
@@ -75,7 +80,7 @@ import ai.djl.inference.Predictor;
 public interface Translator<I, O> extends PreProcessor<I>, PostProcessor<O> {
 
     /**
-     * Gets the {@link Batchifier}.
+     * Returns the {@link Batchifier}.
      *
      * @return the {@link Batchifier}
      */
@@ -93,11 +98,66 @@ public interface Translator<I, O> extends PreProcessor<I>, PostProcessor<O> {
     default void prepare(TranslatorContext ctx) throws Exception {}
 
     /**
-     * Return possible {@link TranslatorOptions} that can be built using this {@link Translator}.
+     * Returns possible {@link TranslatorOptions} that can be built using this {@link Translator}.
      *
      * @return possible options or null if not defined
      */
     default TranslatorOptions getExpansions() {
         return null;
     }
+
+    /**
+     * Returns a batch translator.
+     *
+     * @return a batch translator
+     */
+    default Translator<I[], O[]> toBatchTranslator() {
+        return toBatchTranslator(getBatchifier());
+    }
+
+    /**
+     * Returns a batch translator.
+     *
+     * @param batchifier the {@link Batchifier} to use
+     * @return a batch translator
+     */
+    default Translator<I[], O[]> toBatchTranslator(Batchifier batchifier) {
+        if (batchifier == null) {
+            return null;
+        }
+
+        return new NoBatchifyTranslator<I[], O[]>() {
+
+            /** {@inheritDoc} */
+            @Override
+            @SuppressWarnings("PMD.SignatureDeclareThrowsException")
+            public void prepare(TranslatorContext ctx) throws Exception {
+                Translator.this.prepare(ctx);
+            }
+
+            /** {@inheritDoc} */
+            @Override
+            @SuppressWarnings("PMD.SignatureDeclareThrowsException")
+            public NDList processInput(TranslatorContext ctx, I[] inputs) throws Exception {
+                NDList[] preprocessed = new NDList[inputs.length];
+                for (int i = 0; i < inputs.length; ++i) {
+                    preprocessed[i] = Translator.this.processInput(ctx, inputs[i]);
+                }
+                return batchifier.batchify(preprocessed);
+            }
+
+            /** {@inheritDoc} */
+            @Override
+            @SuppressWarnings({"PMD.SignatureDeclareThrowsException", "unchecked"})
+            public O[] processOutput(TranslatorContext ctx, NDList list) throws Exception {
+                NDList[] unbatched = batchifier.unbatchify(list);
+                List<O> outputs = new ArrayList<>(unbatched.length);
+                for (NDList output : unbatched) {
+                    outputs.add(Translator.this.processOutput(ctx, output));
+                }
+                O[] type = (O[]) Array.newInstance(outputs.get(0).getClass(), 0);
+                return outputs.toArray(type);
+            }
+        };
+    }
 }
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/jna/LibUtils.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/jna/LibUtils.java
index 8dc9ca6b3..793f14691 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/jna/LibUtils.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/jna/LibUtils.java
@@ -296,6 +296,7 @@ public final class LibUtils {
     private static boolean supported(Platform platform) {
         // mxnet-native-cu102mkl:1.8.0: 3.0, 5.0, 6.0, 7.0, 7.5
         // mxnet-native-cu110mkl:1.8.0: 5.0, 6.0, 7.0, 8.0
+        // mxnet-native-cu112mkl:1.9.1: 5.0, 6.0, 7.0, 7.5, 8.0, 8.6
         String version = platform.getVersion();
         if (version.startsWith("1.8.") || version.startsWith("1.9.")) {
             String flavor = platform.getFlavor();
@@ -304,7 +305,7 @@ public final class LibUtils {
                 if (version.startsWith("1.8.")) {
                     return Arrays.asList("50", "60", "70", "80").contains(cudaArch);
                 }
-                return Arrays.asList("50", "60", "70", "75", "80").contains(cudaArch);
+                return Arrays.asList("50", "60", "70", "75", "80", "86").contains(cudaArch);
             } else if (flavor.startsWith("cu10")) {
                 return Arrays.asList("30", "50", "60", "70", "75").contains(cudaArch);
             }
diff --git a/examples/build.gradle b/examples/build.gradle
index 1489919fe..8c42e9649 100644
--- a/examples/build.gradle
+++ b/examples/build.gradle
@@ -9,6 +9,7 @@ dependencies {
     implementation project(":model-zoo")
     implementation project(":extensions:timeseries")
     implementation project(":extensions:tokenizers")
+    implementation project(":extensions:audio")
 
     runtimeOnly project(":engines:pytorch:pytorch-model-zoo")
     runtimeOnly project(":engines:tensorflow:tensorflow-model-zoo")
diff --git a/examples/docs/whisper_speech_text.md b/examples/docs/whisper_speech_text.md
new file mode 100644
index 000000000..245a5b8b1
--- /dev/null
+++ b/examples/docs/whisper_speech_text.md
@@ -0,0 +1,57 @@
+# OpenAI Whipser model in DJL
+
+[Whisper](https://github.com/openai/whisper) is an open source model released by OpenAI.
+It can do speech recognition and also machine translation within a single model. In this tutorial,
+we just convert the English portion of the model into Java.
+
+Our input is an audio file:
+
+```
+https://resources.djl.ai/audios/jfk.flac
+```
+
+## Run the example
+
+```
+cd examples
+./gradlew run -Dmain=ai.djl.examples.inference.whisper.SpeechToTextGeneration
+```
+
+Output:
+
+```
+
 
 
 
 And so my fellow Americans , ask not what your country can do for you , ask what you can do for your country . 

+```
+
+## Trace the model
+
+You can use the following script to get the traced model and also vocabulary files:
+
+```python
+from transformers import WhisperProcessor, WhisperForConditionalGeneration
+from datasets import load_dataset
+import torch
+import numpy as np
+
+processor = WhisperProcessor.from_pretrained("openai/whisper-base")
+processor.tokenizer.save_pretrained("whisper-tokenizer")
+model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-base", return_dict=False)
+ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+
+test = []
+for ele in ds:
+    test.append(ele["audio"]["array"])
+
+input_features = processor(np.concatenate(test), return_tensors="pt").input_features
+generated_ids = model.generate(inputs=input_features)
+transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+print("Original: " + transcription)
+
+# Start tracing
+traced_model = torch.jit.trace_module(model, {"generate": [input_features]})
+generated_ids = traced_model.generate(input_features)
+transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+print("Traced: " + transcription)
+
+torch.jit.save(traced_model, "whisper_en.pt")
+```
\ No newline at end of file
diff --git a/examples/src/main/java/ai/djl/examples/inference/SpeechRecognition.java b/examples/src/main/java/ai/djl/examples/inference/SpeechRecognition.java
index 1561bab1b..5d9a74fad 100644
--- a/examples/src/main/java/ai/djl/examples/inference/SpeechRecognition.java
+++ b/examples/src/main/java/ai/djl/examples/inference/SpeechRecognition.java
@@ -47,8 +47,7 @@ public final class SpeechRecognition {
         logger.info("Result: {}", predict());
     }
 
-    public static String predict()
-            throws UnsupportedAudioFileException, IOException, ModelException, TranslateException {
+    public static String predict() throws IOException, ModelException, TranslateException {
         // Load model.
         // Wav2Vec2 model is a speech model that accepts a float array corresponding to the raw
         // waveform of the speech signal.
@@ -64,7 +63,7 @@ public final class SpeechRecognition {
 
         // Read in audio file
         String wave = "https://resources.djl.ai/audios/speech.wav";
-        Audio audio = AudioFactory.getInstance().fromUrl(wave);
+        Audio audio = AudioFactory.newInstance().fromUrl(wave);
         try (ZooModel<Audio, String> model = criteria.loadModel();
                 Predictor<Audio, String> predictor = model.newPredictor()) {
             return predictor.predict(audio);
diff --git a/examples/src/main/java/ai/djl/examples/inference/whisper/SpeechToTextGeneration.java b/examples/src/main/java/ai/djl/examples/inference/whisper/SpeechToTextGeneration.java
new file mode 100644
index 000000000..de702cd5a
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/inference/whisper/SpeechToTextGeneration.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright 2023 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.examples.inference.whisper;
+
+import ai.djl.ModelException;
+import ai.djl.training.util.DownloadUtils;
+import ai.djl.translate.TranslateException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.nio.file.Paths;
+
+public final class SpeechToTextGeneration {
+
+    private static final Logger logger = LoggerFactory.getLogger(SpeechToTextGeneration.class);
+
+    private SpeechToTextGeneration() {}
+
+    public static void main(String[] args)
+            throws ModelException, IOException, TranslateException, InterruptedException {
+        DownloadUtils.download(
+                "https://resources.djl.ai/audios/jfk.flac", "build/example/jfk.flac");
+
+        try (WhisperModel model = new WhisperModel()) {
+            logger.info(model.speechToText(Paths.get("build/example/jfk.flac")));
+        }
+    }
+}
diff --git a/examples/src/main/java/ai/djl/examples/inference/whisper/WhisperModel.java b/examples/src/main/java/ai/djl/examples/inference/whisper/WhisperModel.java
new file mode 100644
index 000000000..2cb7cb7a8
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/inference/whisper/WhisperModel.java
@@ -0,0 +1,66 @@
+/*
+ * Copyright 2023 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.examples.inference.whisper;
+
+import ai.djl.ModelException;
+import ai.djl.inference.Predictor;
+import ai.djl.modality.audio.Audio;
+import ai.djl.modality.audio.AudioFactory;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.ZooModel;
+import ai.djl.translate.TranslateException;
+
+import org.bytedeco.ffmpeg.global.avutil;
+
+import java.io.IOException;
+import java.nio.file.Path;
+
+/** An example implementation of OpenAI Whisper Model. */
+public class WhisperModel implements AutoCloseable {
+
+    ZooModel<Audio, String> whisperModel;
+
+    public WhisperModel() throws ModelException, IOException {
+        Criteria<Audio, String> criteria =
+                Criteria.builder()
+                        .setTypes(Audio.class, String.class)
+                        .optModelUrls(
+                                "https://resources.djl.ai/demo/pytorch/whisper/whisper_en.zip")
+                        .optEngine("PyTorch")
+                        .optTranslator(new WhisperTranslator())
+                        .build();
+        whisperModel = criteria.loadModel();
+    }
+
+    public String speechToText(Audio speech) throws TranslateException {
+        try (Predictor<Audio, String> predictor = whisperModel.newPredictor()) {
+            return predictor.predict(speech);
+        }
+    }
+
+    public String speechToText(Path file) throws IOException, TranslateException {
+        Audio audio =
+                AudioFactory.newInstance()
+                        .setChannels(1)
+                        .setSampleRate(16000)
+                        .setSampleFormat(avutil.AV_SAMPLE_FMT_S16)
+                        .fromFile(file);
+        return speechToText(audio);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void close() {
+        whisperModel.close();
+    }
+}
diff --git a/examples/src/main/java/ai/djl/examples/inference/whisper/WhisperTranslator.java b/examples/src/main/java/ai/djl/examples/inference/whisper/WhisperTranslator.java
new file mode 100644
index 000000000..155cfc4af
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/inference/whisper/WhisperTranslator.java
@@ -0,0 +1,101 @@
+/*
+ * Copyright 2023 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.examples.inference.whisper;
+
+import ai.djl.audio.processor.AudioProcessor;
+import ai.djl.audio.processor.LogMelSpectrogram;
+import ai.djl.audio.processor.PadOrTrim;
+import ai.djl.modality.audio.Audio;
+import ai.djl.modality.nlp.DefaultVocabulary;
+import ai.djl.modality.nlp.Vocabulary;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.translate.NoBatchifyTranslator;
+import ai.djl.translate.TranslatorContext;
+import ai.djl.util.JsonUtils;
+
+import com.google.gson.reflect.TypeToken;
+
+import java.io.IOException;
+import java.io.Reader;
+import java.lang.reflect.Type;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+
+public class WhisperTranslator implements NoBatchifyTranslator<Audio, String> {
+
+    private List<AudioProcessor> processors;
+    private Vocabulary vocabulary;
+
+    public WhisperTranslator() {
+        processors = new ArrayList<>();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void prepare(TranslatorContext ctx) throws IOException {
+        NDManager manager = ctx.getNDManager();
+        Path path = ctx.getModel().getModelPath();
+        Path melFile = path.resolve("mel_80_filters.npz");
+
+        processors.add(new PadOrTrim(480000));
+        processors.add(LogMelSpectrogram.newInstance(melFile, 80, manager));
+
+        Map<String, Integer> vocab;
+        Map<String, Integer> added;
+        Type type = new TypeToken<Map<String, Integer>>() {}.getType();
+        try (Reader reader = Files.newBufferedReader(path.resolve("vocab.json"))) {
+            vocab = JsonUtils.GSON.fromJson(reader, type);
+        }
+        try (Reader reader = Files.newBufferedReader(path.resolve("added_tokens.json"))) {
+            added = JsonUtils.GSON.fromJson(reader, type);
+        }
+        String[] result = new String[vocab.size() + added.size()];
+        vocab.forEach((key, value) -> result[value] = key);
+        added.forEach((key, value) -> result[value] = key);
+        vocabulary = new DefaultVocabulary(Arrays.asList(result));
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, Audio input) throws Exception {
+        NDArray samples = ctx.getNDManager().create(input.getData());
+        for (AudioProcessor processor : processors) {
+            samples = processor.extractFeatures(samples.getManager(), samples);
+        }
+        samples = samples.expandDims(0);
+        NDArray placeholder = ctx.getNDManager().create("");
+        placeholder.setName("module_method:generate");
+        return new NDList(samples, placeholder);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public String processOutput(TranslatorContext ctx, NDList list) throws Exception {
+        NDArray result = list.singletonOrThrow();
+        List<String> sentence = new ArrayList<>();
+        for (long ele : result.toLongArray()) {
+            sentence.add(vocabulary.getToken(ele));
+            if ("
".equals(vocabulary.getToken(ele))) {
+                break;
+            }
+        }
+        String output = String.join(" ", sentence);
+        return output.replaceAll("[^a-zA-Z0-9<|> ,.!]", "");
+    }
+}
diff --git a/examples/src/main/java/ai/djl/examples/inference/whisper/package-info.java b/examples/src/main/java/ai/djl/examples/inference/whisper/package-info.java
new file mode 100644
index 000000000..8016afbf6
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/inference/whisper/package-info.java
@@ -0,0 +1,15 @@
+/*
+ * Copyright 2023 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.
+ */
+
+/** Contains examples of Speech to Text task. */
+package ai.djl.examples.inference.whisper;
diff --git a/extensions/audio/src/main/java/ai/djl/audio/FFmpegAudioFactory.java b/extensions/audio/src/main/java/ai/djl/audio/FFmpegAudioFactory.java
index ad6ffa071..07f238708 100644
--- a/extensions/audio/src/main/java/ai/djl/audio/FFmpegAudioFactory.java
+++ b/extensions/audio/src/main/java/ai/djl/audio/FFmpegAudioFactory.java
@@ -14,7 +14,6 @@ package ai.djl.audio;
 
 import ai.djl.modality.audio.Audio;
 import ai.djl.modality.audio.AudioFactory;
-import ai.djl.ndarray.NDArray;
 
 import org.bytedeco.javacv.FFmpegFrameGrabber;
 import org.bytedeco.javacv.Frame;
@@ -22,6 +21,8 @@ import org.bytedeco.javacv.FrameGrabber;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.Buffer;
+import java.nio.IntBuffer;
 import java.nio.ShortBuffer;
 import java.nio.file.Path;
 import java.util.ArrayList;
@@ -37,6 +38,7 @@ public class FFmpegAudioFactory extends AudioFactory {
     @Override
     public Audio fromFile(Path path) throws IOException {
         try (FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(path.toFile())) {
+            applyConfig(grabber);
             grabber.start();
             float[] floats = grab(grabber);
             return new Audio(floats, grabber.getSampleRate(), grabber.getAudioChannels());
@@ -49,6 +51,7 @@ public class FFmpegAudioFactory extends AudioFactory {
     @Override
     public Audio fromInputStream(InputStream is) throws IOException {
         try (FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(is)) {
+            applyConfig(grabber);
             grabber.start();
             float[] floats = grab(grabber);
             return new Audio(floats, grabber.getSampleRate(), grabber.getAudioChannels());
@@ -57,20 +60,22 @@ public class FFmpegAudioFactory extends AudioFactory {
         }
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public Audio fromData(float[] data) {
-        return new Audio(data);
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public Audio fromNDArray(NDArray array) {
-        throw new UnsupportedOperationException("Not supported!");
+    private void applyConfig(FFmpegFrameGrabber grabber) {
+        if (channels > 0) {
+            grabber.setAudioChannels(channels);
+        }
+        if (sampleRate > 0) {
+            grabber.setSampleRate(sampleRate);
+        }
+        if (sampleFormat > 0) {
+            grabber.setSampleFormat(sampleFormat);
+        }
     }
 
     /**
-     * Grab frames from the audio using {@link FFmpegFrameGrabber}.
+     * Grabs frames from the audio using {@link FFmpegFrameGrabber}.
+     *
+     * <p>The default channel to grab is 0.
      *
      * @param grabber the {@link FFmpegFrameGrabber}.
      * @return the float array read from the audio.
@@ -79,10 +84,21 @@ public class FFmpegAudioFactory extends AudioFactory {
     private float[] grab(FFmpegFrameGrabber grabber) throws FFmpegFrameGrabber.Exception {
         List<Float> list = new ArrayList<>();
         Frame frame;
-        while ((frame = grabber.grab()) != null) {
-            ShortBuffer buffer = (ShortBuffer) frame.samples[0];
-            for (int i = 0; i < buffer.limit(); i++) {
-                list.add(buffer.get() / (float) Short.MAX_VALUE);
+        while ((frame = grabber.grabFrame(true, false, true, false, false)) != null) {
+            Buffer buf = frame.samples[0];
+            if (buf instanceof ShortBuffer) {
+                ShortBuffer buffer = (ShortBuffer) buf;
+                for (int i = 0; i < buffer.limit(); i++) {
+                    list.add(buffer.get() / (float) Short.MAX_VALUE);
+                }
+            } else if (buf instanceof IntBuffer) {
+                IntBuffer buffer = (IntBuffer) buf;
+                for (int i = 0; i < buffer.limit(); i++) {
+                    list.add(buffer.get() / (float) Integer.MAX_VALUE);
+                }
+            } else {
+                throw new UnsupportedOperationException(
+                        "Unsupported sample format: " + sampleFormat);
             }
         }
         float[] ret = new float[list.size()];
diff --git a/extensions/audio/src/main/java/ai/djl/audio/processor/LogMelSpectrogram.java b/extensions/audio/src/main/java/ai/djl/audio/processor/LogMelSpectrogram.java
new file mode 100644
index 000000000..b76598a45
--- /dev/null
+++ b/extensions/audio/src/main/java/ai/djl/audio/processor/LogMelSpectrogram.java
@@ -0,0 +1,87 @@
+/*
+ * Copyright 2023 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.audio.processor;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+
+/** Apply Log Mel spectrogram to the given data. */
+public class LogMelSpectrogram implements AudioProcessor {
+
+    private static final int N_FFT = 400;
+    private static final int HOP_LENGTH = 160;
+
+    private NDArray melFilters;
+
+    /**
+     * Constructs a new instance of {@code LogMelSpectrogram}.
+     *
+     * @param melFilter the mel filter
+     */
+    public LogMelSpectrogram(NDArray melFilter) {
+        this.melFilters = melFilter;
+    }
+
+    /**
+     * Loads the mel filterbank matrix for projecting STFT into a Mel spectrogram.
+     *
+     * <p>Allows decoupling librosa dependency; saved using: np.savez_compressed( "mel_filters.npz",
+     * mel_80=librosa.filters.mel(sr=16000, n_fft=400, n_mels=80), )
+     *
+     * @param melFile the mdel file saved in .npz format
+     * @param numMel number of mel
+     * @param manager manager to set for the content
+     * @return a new instance of {@code LogMelSpectrogram}
+     * @throws IOException file not loadable
+     */
+    public static LogMelSpectrogram newInstance(Path melFile, int numMel, NDManager manager)
+            throws IOException {
+        try (InputStream is = Files.newInputStream(melFile)) {
+            return newInstance(is, numMel, manager);
+        }
+    }
+
+    /**
+     * Loads the mel filterbank matrix for projecting STFT into a Mel spectrogram.
+     *
+     * @param is the input stream
+     * @param numMel number of mel
+     * @param manager manager to set for the content
+     * @return a new instance of {@code LogMelSpectrogram}
+     * @throws IOException file not loadable
+     */
+    public static LogMelSpectrogram newInstance(InputStream is, int numMel, NDManager manager)
+            throws IOException {
+        return new LogMelSpectrogram(NDList.decode(manager, is).get("mel_" + numMel));
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray extractFeatures(NDManager manager, NDArray samples) {
+        NDArray window = manager.hanningWindow(N_FFT);
+        NDArray stft = samples.stft(N_FFT, HOP_LENGTH, true, window, true);
+        NDArray magnitudes = stft.get(":,:-1").abs().pow(2);
+        NDArray melSpec = melFilters.matMul(magnitudes);
+        melSpec.attach(manager);
+        NDArray logSpec = melSpec.clip(1e-10, Float.MAX_VALUE).log10();
+        logSpec = logSpec.maximum(logSpec.max().sub(8.0f));
+        logSpec = logSpec.add(4.0f).div(4.0f);
+        return logSpec;
+    }
+}
diff --git a/extensions/audio/src/main/java/ai/djl/audio/processor/PadOrTrim.java b/extensions/audio/src/main/java/ai/djl/audio/processor/PadOrTrim.java
new file mode 100644
index 000000000..368155e3b
--- /dev/null
+++ b/extensions/audio/src/main/java/ai/djl/audio/processor/PadOrTrim.java
@@ -0,0 +1,45 @@
+/*
+ * Copyright 2023 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.audio.processor;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.Shape;
+
+/** Pad or trim the samples to the desired numbers. */
+public class PadOrTrim implements AudioProcessor {
+
+    private int desiredSamples;
+
+    /**
+     * Pad or trim the samples to the fixed length.
+     *
+     * @param desiredSamples the desired sample points
+     */
+    public PadOrTrim(int desiredSamples) {
+        this.desiredSamples = desiredSamples;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray extractFeatures(NDManager manager, NDArray samples) {
+        if (samples.getShape().dimension() != 1) {
+            throw new UnsupportedOperationException("Batch samples not supported.");
+        }
+        long sampleLength = samples.getShape().get(0);
+        if (sampleLength > desiredSamples) {
+            return samples.get(":" + desiredSamples);
+        }
+        return samples.concat(manager.zeros(new Shape(desiredSamples - sampleLength)));
+    }
+}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingBatchTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingBatchTranslator.java
new file mode 100644
index 000000000..eb3fd1584
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingBatchTranslator.java
@@ -0,0 +1,81 @@
+/*
+ * Copyright 2023 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.translate.Batchifier;
+import ai.djl.translate.NoBatchifyTranslator;
+import ai.djl.translate.TranslatorContext;
+
+/** The translator for Huggingface text embedding model. */
+public class TextEmbeddingBatchTranslator implements NoBatchifyTranslator<String[], float[][]> {
+
+    private HuggingFaceTokenizer tokenizer;
+    private Batchifier batchifier;
+    private boolean normalize;
+    private String pooling;
+
+    TextEmbeddingBatchTranslator(
+            HuggingFaceTokenizer tokenizer,
+            Batchifier batchifier,
+            String pooling,
+            boolean normalize) {
+        this.tokenizer = tokenizer;
+        this.batchifier = batchifier;
+        this.pooling = pooling;
+        this.normalize = normalize;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, String[] input) {
+        NDManager manager = ctx.getNDManager();
+        Encoding[] encodings = tokenizer.batchEncode(input);
+        ctx.setAttachment("encoding", encodings);
+        NDList[] batch = new NDList[encodings.length];
+        for (int i = 0; i < encodings.length; ++i) {
+            long[] indices = encodings[i].getIds();
+            long[] attentionMask = encodings[i].getAttentionMask();
+            batch[i] = new NDList(2);
+            batch[i].add(manager.create(indices));
+            batch[i].add(manager.create(attentionMask));
+        }
+        return batchifier.batchify(batch);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public float[][] processOutput(TranslatorContext ctx, NDList list) {
+        NDArray embeddings = list.get("last_hidden_state");
+        NDList[] batch = batchifier.unbatchify(new NDList(embeddings));
+        Encoding[] encoding = (Encoding[]) ctx.getAttachment("encoding");
+        NDManager manager = ctx.getNDManager();
+        float[][] ret = new float[batch.length][];
+        for (int i = 0; i < batch.length; ++i) {
+            NDArray embedding = batch[i].singletonOrThrow();
+            NDArray array =
+                    TextEmbeddingTranslator.processEmbedding(
+                            manager, embedding, encoding[i], pooling);
+            if (normalize) {
+                array = array.normalize(2, 0);
+            }
+            ret[i] = array.toFloatArray();
+        }
+
+        return ret;
+    }
+}
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
index b611a90d0..544a13d70 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
@@ -72,36 +72,42 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
     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();
+        embeddings = processEmbedding(manager, embeddings, encoding, pooling);
+        if (normalize) {
+            embeddings = embeddings.normalize(2, 0);
+        }
+
+        return embeddings.toFloatArray();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public TextEmbeddingBatchTranslator toBatchTranslator(Batchifier batchifier) {
+        return new TextEmbeddingBatchTranslator(tokenizer, batchifier, pooling, normalize);
+    }
+
+    static NDArray processEmbedding(
+            NDManager manager, NDArray embeddings, Encoding encoding, String pooling) {
+        long[] attentionMask = encoding.getAttentionMask();
         NDArray inputAttentionMask = manager.create(attentionMask).toType(DataType.FLOAT32, true);
         switch (pooling) {
             case "mean_tokens":
-                embeddings = meanPool(embeddings, inputAttentionMask, false);
-                break;
+                return meanPool(embeddings, inputAttentionMask, false);
             case "mean_sqrt_len_tokens":
-                embeddings = meanPool(embeddings, inputAttentionMask, true);
-                break;
+                return meanPool(embeddings, inputAttentionMask, true);
             case "max_tokens":
-                embeddings = maxPool(embeddings, inputAttentionMask);
-                break;
+                return maxPool(embeddings, inputAttentionMask);
             case "weightedmean_tokens":
-                embeddings = weightedMeanPool(embeddings, inputAttentionMask);
-                break;
+                return weightedMeanPool(embeddings, inputAttentionMask);
             case "cls_token":
-                embeddings = embeddings.get(0);
-                break;
+                return embeddings.get(0);
             default:
                 throw new AssertionError("Unexpected pooling model: " + pooling);
         }
-        if (normalize) {
-            embeddings = embeddings.normalize(2, 0);
-        }
-
-        return embeddings.toFloatArray();
     }
 
-    private NDArray meanPool(NDArray embeddings, NDArray attentionMask, boolean sqrt) {
+    private static NDArray meanPool(NDArray embeddings, NDArray attentionMask, boolean sqrt) {
         long[] shape = embeddings.getShape().getShape();
         attentionMask = attentionMask.expandDims(-1).broadcast(shape);
         NDArray inputAttentionMaskSum = attentionMask.sum(AXIS);
@@ -114,7 +120,7 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
         return sum.div(clamp);
     }
 
-    private NDArray maxPool(NDArray embeddings, NDArray inputAttentionMask) {
+    private static NDArray maxPool(NDArray embeddings, NDArray inputAttentionMask) {
         long[] shape = embeddings.getShape().getShape();
         inputAttentionMask = inputAttentionMask.expandDims(-1).broadcast(shape);
         inputAttentionMask = inputAttentionMask.eq(0);
@@ -124,7 +130,7 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
         return embeddings.max(AXIS, true);
     }
 
-    private NDArray weightedMeanPool(NDArray embeddings, NDArray attentionMask) {
+    private static NDArray weightedMeanPool(NDArray embeddings, NDArray attentionMask) {
         long[] shape = embeddings.getShape().getShape();
         NDArray weight = embeddings.getManager().arange(1, shape[0] + 1);
         weight = weight.expandDims(-1).broadcast(shape);
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
index 43bc4f03b..49e436612 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java
@@ -36,6 +36,7 @@ public class TextEmbeddingTranslatorFactory implements TranslatorFactory {
 
     static {
         SUPPORTED_TYPES.add(new Pair<>(String.class, float[].class));
+        SUPPORTED_TYPES.add(new Pair<>(String[].class, float[][].class));
         SUPPORTED_TYPES.add(new Pair<>(Input.class, Output.class));
     }
 
@@ -62,6 +63,8 @@ public class TextEmbeddingTranslatorFactory implements TranslatorFactory {
                     TextEmbeddingTranslator.builder(tokenizer, arguments).build();
             if (input == String.class && output == float[].class) {
                 return (Translator<I, O>) translator;
+            } else if (input == String[].class && output == float[][].class) {
+                return (Translator<I, O>) translator.toBatchTranslator();
             } else if (input == Input.class && output == Output.class) {
                 return (Translator<I, O>) new TextEmbeddingServingTranslator(translator);
             }


Test patch:
diff --git a/api/src/test/java/ai/djl/modality/audio/AudioFactoryTest.java b/api/src/test/java/ai/djl/modality/audio/AudioFactoryTest.java
index 40ec26d27..c5b082c7a 100644
--- a/api/src/test/java/ai/djl/modality/audio/AudioFactoryTest.java
+++ b/api/src/test/java/ai/djl/modality/audio/AudioFactoryTest.java
@@ -22,8 +22,6 @@ import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
-import java.io.InputStream;
-import java.nio.file.Files;
 import java.nio.file.Paths;
 
 public class AudioFactoryTest {
@@ -37,37 +35,36 @@ public class AudioFactoryTest {
 
     @Test
     public void testFromFile() throws IOException {
-        Audio audio = AudioFactory.getInstance().fromFile(Paths.get("build/test/test_01.wav"));
+        Audio audio = AudioFactory.newInstance().fromFile(Paths.get("build/test/test_01.wav"));
         Assert.assertEquals(audio.getSampleRate(), 16000f);
         Assert.assertEquals(audio.getChannels(), 1);
+
+        AudioFactory factory = AudioFactory.newInstance().setChannels(0);
+        Assert.assertEquals(factory.getChannels(), 0);
+        Assert.assertEquals(factory.getSampleFormat(), 0);
+        Assert.assertEquals(factory.getSampleRate(), 0);
+        Assert.assertThrows(() -> factory.setChannels(1));
+        Assert.assertThrows(() -> factory.setSampleRate(1000));
+        Assert.assertThrows(() -> factory.setSampleFormat(1)); // AV_SAMPLE_FMT_S16
     }
 
     @Test
     public void testFromUrl() throws IOException {
-        Audio audio = AudioFactory.getInstance().fromUrl("build/test/test_01.wav");
+        Audio audio = AudioFactory.newInstance().fromUrl("build/test/test_01.wav");
         Assert.assertEquals(audio.getSampleRate(), 16000f);
         Assert.assertEquals(audio.getChannels(), 1);
 
-        audio = AudioFactory.getInstance().fromUrl(URL);
+        audio = AudioFactory.newInstance().fromUrl(URL);
         Assert.assertEquals(audio.getSampleRate(), 16000f);
         Assert.assertEquals(audio.getChannels(), 1);
     }
 
-    @Test
-    public void testFromInputStream() throws IOException {
-        try (InputStream is = Files.newInputStream(Paths.get("build/test/test_01.wav"))) {
-            Audio audio = AudioFactory.getInstance().fromInputStream(is);
-            Assert.assertEquals(audio.getSampleRate(), 16000f);
-            Assert.assertEquals(audio.getChannels(), 1);
-        }
-    }
-
     @Test
     public void testFromData() {
         float[] data = {0.001f, 0.002f, 0.003f};
-        Audio audio = AudioFactory.getInstance().fromData(data);
+        Audio audio = AudioFactory.newInstance().fromData(data);
         Assert.assertEquals(audio.getData(), data);
-        Assert.assertEquals(audio.getSampleRate(), 0);
+        Assert.assertEquals(audio.getSampleRate(), 0f);
         Assert.assertEquals(audio.getChannels(), 0);
     }
 
@@ -75,7 +72,7 @@ public class AudioFactoryTest {
     public void testFromNDArray() {
         try (NDManager manager = NDManager.newBaseManager()) {
             NDArray array = manager.zeros(new Shape(1));
-            AudioFactory.getInstance().fromNDArray(array);
+            AudioFactory.newInstance().fromNDArray(array);
         }
     }
 }
diff --git a/api/src/test/java/ai/djl/translate/TranslatorTest.java b/api/src/test/java/ai/djl/translate/TranslatorTest.java
new file mode 100644
index 000000000..5ce63faa2
--- /dev/null
+++ b/api/src/test/java/ai/djl/translate/TranslatorTest.java
@@ -0,0 +1,88 @@
+/*
+ * Copyright 2023 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.translate;
+
+import ai.djl.ModelException;
+import ai.djl.inference.Predictor;
+import ai.djl.modality.Classifications;
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.ImageFactory;
+import ai.djl.modality.cv.transform.CenterCrop;
+import ai.djl.modality.cv.transform.Resize;
+import ai.djl.modality.cv.transform.ToTensor;
+import ai.djl.modality.cv.translator.ImageClassificationTranslator;
+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;
+import ai.djl.repository.zoo.ZooModel;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Arrays;
+
+public class TranslatorTest {
+
+    @Test
+    public void testBatchTranslator() throws IOException, ModelException, TranslateException {
+        Block block =
+                new LambdaBlock(
+                        a -> {
+                            NDManager manager = a.getManager();
+                            NDArray arr = manager.ones(new Shape(2, 1000));
+                            return new NDList(arr);
+                        },
+                        "model");
+        Path modelDir = Paths.get("build/model");
+        Files.createDirectories(modelDir);
+
+        ImageFactory factory = ImageFactory.getInstance();
+        Image[] inputs = {
+            factory.fromFile(Paths.get("../examples/src/test/resources/kitten.jpg")),
+            factory.fromFile(Paths.get("../examples/src/test/resources/dog-cat.jpg"))
+        };
+        String[] classes = new String[1000];
+        Arrays.fill(classes, "something");
+
+        ImageClassificationTranslator translator =
+                ImageClassificationTranslator.builder()
+                        .addTransform(new Resize(256))
+                        .addTransform(new CenterCrop(224, 224))
+                        .addTransform(new ToTensor())
+                        .optSynset(Arrays.asList(classes))
+                        .build();
+
+        Criteria<Image[], Classifications[]> criteria =
+                Criteria.builder()
+                        .setTypes(Image[].class, Classifications[].class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optOption("hasParameter", "false")
+                        .optTranslator(translator.toBatchTranslator())
+                        .build();
+
+        try (ZooModel<Image[], Classifications[]> model = criteria.loadModel();
+                Predictor<Image[], Classifications[]> predictor = model.newPredictor()) {
+            Classifications[] res = predictor.predict(inputs);
+            Assert.assertEquals(res.length, 2);
+        }
+    }
+}
diff --git a/extensions/audio/src/test/java/ai/djl/audio/FFmpegAudioFactoryTest.java b/extensions/audio/src/test/java/ai/djl/audio/FFmpegAudioFactoryTest.java
index 7cfe2ca63..43da44d8b 100644
--- a/extensions/audio/src/test/java/ai/djl/audio/FFmpegAudioFactoryTest.java
+++ b/extensions/audio/src/test/java/ai/djl/audio/FFmpegAudioFactoryTest.java
@@ -14,53 +14,56 @@ package ai.djl.audio;
 
 import ai.djl.modality.audio.Audio;
 import ai.djl.modality.audio.AudioFactory;
-import ai.djl.ndarray.NDArray;
-import ai.djl.ndarray.NDManager;
-import ai.djl.ndarray.types.Shape;
 import ai.djl.training.util.DownloadUtils;
 
+import org.bytedeco.ffmpeg.global.avutil;
 import org.testng.Assert;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
+import java.net.URL;
 import java.nio.file.Paths;
 
 public class FFmpegAudioFactoryTest {
 
-    private static final String URL = "https://resources.djl.ai/audios/test_01.wav";
-
     @BeforeClass
     public void setUp() throws IOException {
-        DownloadUtils.download(URL, "build/test/test_01.wav");
+        DownloadUtils.download(
+                "https://resources.djl.ai/audios/test_01.wav", "build/test/test_01.wav");
     }
 
     @Test
     public void testFactory() throws IOException {
-        Audio audio = AudioFactory.getInstance().fromFile(Paths.get("build/test/test_01.wav"));
-        Assert.assertEquals(audio.getSampleRate(), 16000f);
-        Assert.assertEquals(audio.getChannels(), 1);
+        AudioFactory factory =
+                AudioFactory.newInstance()
+                        .setChannels(1)
+                        .setSampleRate(10000)
+                        .setSampleFormat(avutil.AV_SAMPLE_FMT_S16);
 
-        audio = AudioFactory.getInstance().fromUrl("build/test/test_01.wav");
-        Assert.assertEquals(audio.getSampleRate(), 16000f);
-        Assert.assertEquals(audio.getChannels(), 1);
+        Assert.assertEquals(factory.getChannels(), 1);
+        Assert.assertEquals(factory.getSampleRate(), 10000);
+        Assert.assertEquals(factory.getSampleFormat(), avutil.AV_SAMPLE_FMT_S16);
 
-        audio = AudioFactory.getInstance().fromUrl(URL);
-        Assert.assertEquals(audio.getSampleRate(), 16000f);
+        Audio audio = factory.fromUrl("build/test/test_01.wav");
+        Assert.assertEquals(audio.getSampleRate(), 10000f);
         Assert.assertEquals(audio.getChannels(), 1);
 
-        float[] data = {0.001f, 0.002f, 0.003f};
-        audio = AudioFactory.getInstance().fromData(data);
-        Assert.assertEquals(audio.getData(), data);
-        Assert.assertEquals(audio.getSampleRate(), 0);
-        Assert.assertEquals(audio.getChannels(), 0);
-    }
+        URL url = Paths.get("build/test/test_01.wav").toAbsolutePath().toUri().toURL();
+
+        audio =
+                AudioFactory.newInstance()
+                        .setSampleRate(8000)
+                        .setSampleFormat(avutil.AV_SAMPLE_FMT_S32)
+                        .fromUrl(url);
+        Assert.assertEquals(audio.getSampleRate(), 8000f);
+        Assert.assertEquals(audio.getChannels(), 1);
 
-    @Test(expectedExceptions = UnsupportedOperationException.class)
-    public void testFromNDArray() {
-        try (NDManager manager = NDManager.newBaseManager()) {
-            NDArray array = manager.zeros(new Shape(1));
-            AudioFactory.getInstance().fromNDArray(array);
-        }
+        Assert.assertThrows(
+                () -> {
+                    AudioFactory.newInstance()
+                            .setSampleFormat(avutil.AV_SAMPLE_FMT_DBL)
+                            .fromUrl(url);
+                });
     }
 }
diff --git a/extensions/audio/src/test/java/ai/djl/audio/processor/AudioProcessorTest.java b/extensions/audio/src/test/java/ai/djl/audio/processor/AudioProcessorTest.java
index 8070a1842..4b9517369 100644
--- a/extensions/audio/src/test/java/ai/djl/audio/processor/AudioProcessorTest.java
+++ b/extensions/audio/src/test/java/ai/djl/audio/processor/AudioProcessorTest.java
@@ -26,16 +26,19 @@ import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
+import java.nio.file.Paths;
 import java.util.Arrays;
 import java.util.Collections;
 
 public class AudioProcessorTest {
 
-    private static final String URL = "https://resources.djl.ai/audios/test_01.wav";
+    private static final String URL = "https://resources.djl.ai/";
 
     @BeforeClass
     public void setUp() throws IOException {
-        DownloadUtils.download(URL, "build/test/test_01.wav");
+        DownloadUtils.download(URL + "audios/test_01.wav", "build/test/test_01.wav");
+        DownloadUtils.download(
+                URL + "demo/pytorch/whisper/mel_80_filters.npz", "build/test/mel_80_filters.npz");
     }
 
     @Test
@@ -50,6 +53,31 @@ public class AudioProcessorTest {
         Assert.assertEquals(AudioUtils.rmsDb(samples), -20.0f, 1e-3);
     }
 
+    @Test
+    public void testPadOrTrim() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            AudioProcessor processor = new PadOrTrim(100);
+            NDArray test = manager.ones(new Shape(10));
+            NDArray result = processor.extractFeatures(manager, test);
+            Assert.assertEquals(result.getShape(), new Shape(100));
+            test = manager.ones(new Shape(1000));
+            result = processor.extractFeatures(manager, test);
+            Assert.assertEquals(result.getShape(), new Shape(100));
+        }
+    }
+
+    @Test
+    void testLogMelSpectrogram() throws IOException {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            AudioProcessor processor =
+                    LogMelSpectrogram.newInstance(
+                            Paths.get("build/test/mel_80_filters.npz"), 80, manager);
+            NDArray samples = manager.zeros(new Shape(480000));
+            NDArray result = processor.extractFeatures(manager, samples);
+            Assert.assertEquals(result.getShape(), new Shape(80, 3000));
+        }
+    }
+
     @Test
     public static void testLinearSpecgram() {
         NDManager manager = NDManager.newBaseManager(Device.cpu());
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 e8abc5d03..821afb117 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
@@ -460,6 +460,66 @@ public class TranslatorTest {
         }
     }
 
+    @Test
+    public void testTextEmbeddingBatchTranslator()
+            throws ModelException, IOException, TranslateException {
+        String[] text = {"This is an example sentence", "This is the second sentence"};
+
+        Block block =
+                new LambdaBlock(
+                        a -> {
+                            NDManager manager = a.getManager();
+                            NDArray arr = manager.ones(new Shape(2, 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")
+                        .optArgument("padding", "true")
+                        .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[0].length, 384);
+            Assertions.assertAlmostEquals(res[0][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")
+                        .optArgument("padding", "true")
+                        .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(JsonUtils.GSON.toJson(text));
+            input.addProperty("Content-Type", "application/json");
+            Output out = predictor.predict(input);
+            float[][] res = JsonUtils.GSON.fromJson(out.getAsString(0), float[][].class);
+            Assert.assertEquals(res[0].length, 384);
+            Assertions.assertAlmostEquals(res[0][0], 0.05103);
+        }
+    }
+
     @Test
     public void testTextClassificationTranslator()
             throws ModelException, IOException, TranslateException {