Repository: deepjavalibrary/djl
Pull request #1869 (2022-08-04 13:41:56)

Bulk Batch creation
This pull requests optimizes Batch creation for ArrayDataset when using StackBatchifier .
My DJL application learns the data now in 58550ms  rather than 79843ms so 36% faster!
For applications not using range-based indexing but indices-based indexing, my DJL application would run in 66341ms , so still 20% faster. This is true for PyTorch, i get a test failure for MxNet, so this solution maybe restricted to PyTorch

codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1869?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
> Merging [#1869](https://codecov.io/gh/deepjavalibrary/djl/pull/1869?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (b4ad024) 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.36%`.
> The diff coverage is `64.73%`.

```diff
@@             Coverage Diff              @@
##             master    #1869      +/-   ##
============================================
- Coverage     72.08%   69.72%   -2.37%     
- Complexity     5126     5573     +447     
============================================
  Files           473      528      +55     
  Lines         21970    24642    +2672     
  Branches       2351     2686     +335     
============================================
+ Hits          15838    17181    +1343     
- Misses         4925     6145    +1220     
- Partials       1207     1316     +109     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1869?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/1869/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/1869/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/1869/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/1869/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/1869/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/1869/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/1869/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/1869/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/1869/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/1869/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 [426 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1869/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).

patins1: The problems with MxNet is fixed by using "take" operation instead of "pick".
TrainMnist executes now 4 times faster for MxNet .
lanking520: @KexinFeng can you help to check on the NDIndex part changes?
KexinFeng: @patins1 It looks like `BulkDataIterable` is not covered by any unit test. You mentioned that using `BulkDataIterable` is more efficient. Could you add a unit test that covers this class you added?
KexinFeng: > The problems with MxNet is fixed by using "take" operation instead of "pick".
> TrainMnist executes now 4 times faster for MxNet .

Could you integrate it with the existing `take` function in [PR](https://github.com/deepjavalibrary/djl/pull/1649)? 
KexinFeng: @patins1 Thanks for dealing with this issue in such a timely manner! But it looks like the unittest that covers the new file api/src/main/java/ai/djl/training/dataset/BulkDataIterable.java is still missing. You mentioned you have tested it, which display efficiency increase. Could you add it into the proper unit test file? 

Also you mentioned 
> The problems with MxNet is fixed by using "take" operation instead of "pick".
> TrainMnist executes now 4 times faster for MxNet .

It'd be better to have unit test for this too, which will prevent future edition from breaking this fix.

Thanks!

patins1: Tests added
KexinFeng: @siddvenk Hi Siddarth, in this pr, we have changed the definition of get(NDArray index) from pick to take, and have given the warning. This will affect the result you mentioned in https://github.com/deepjavalibrary/djl/issues/1800. See the test in commit "add index test and clean code" 8ffac00e6584c96500e6ac96ba35cc45513f1f4e. Pick can still be used though by addPickDim().
This change is for the purpose of making it consistent with numpy and pytorch engine, also here it is mor efficient.
Should we update the relevant part mentioned in https://github.com/deepjavalibrary/djl/issues/1800?

Patch:
diff --git a/api/src/main/java/ai/djl/DeviceType.java b/api/src/main/java/ai/djl/DeviceType.java
deleted file mode 100644
index 45cc4fad4..000000000
--- a/api/src/main/java/ai/djl/DeviceType.java
+++ /dev/null
@@ -1,37 +0,0 @@
-/*
- * Copyright 2019 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;
-
-/** {@code DeviceType} is a class used to map the Device name to their corresponding type number. */
-public interface DeviceType {
-
-    /**
-     * Map device to its type number.
-     *
-     * @param device {@link Device} to map from
-     * @return the number specified by engine
-     */
-    static int toDeviceType(Device device) {
-        return 0;
-    }
-
-    /**
-     * Map device to its type number.
-     *
-     * @param deviceType the number specified by engine
-     * @return {@link Device} to map to
-     */
-    static String fromDeviceType(int deviceType) {
-        return null;
-    }
-}
diff --git a/api/src/main/java/ai/djl/modality/audio/Audio.java b/api/src/main/java/ai/djl/modality/audio/Audio.java
new file mode 100644
index 000000000..4ad949845
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/audio/Audio.java
@@ -0,0 +1,73 @@
+/*
+ * 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.modality.audio;
+
+/**
+ * {@code Audio} is a container of an audio in DJL. The raw data of the audio is wrapped in a float
+ * array.
+ */
+public class Audio {
+
+    private float[] data;
+    private float sampleRate;
+    private int channels;
+
+    /**
+     * Constructs a new {@code Audio} instance.
+     *
+     * @param data the wrapped float array data
+     */
+    public Audio(float[] data) {
+        this.data = data;
+    }
+
+    /**
+     * Constructs a new {@code Audio} instance.
+     *
+     * @param data the wrapped float array data
+     * @param sampleRate the sample rate
+     * @param channels number of channels
+     */
+    public Audio(float[] data, float sampleRate, int channels) {
+        this.data = data;
+        this.sampleRate = sampleRate;
+        this.channels = channels;
+    }
+
+    /**
+     * Returns the float array data.
+     *
+     * @return The float array data.
+     */
+    public float[] getData() {
+        return data;
+    }
+
+    /**
+     * Returns the sample rate.
+     *
+     * @return sample rate.
+     */
+    public float getSampleRate() {
+        return sampleRate;
+    }
+
+    /**
+     * Get 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
new file mode 100644
index 000000000..310134c57
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/audio/AudioFactory.java
@@ -0,0 +1,124 @@
+/*
+ * 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.modality.audio;
+
+import ai.djl.ndarray.NDArray;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.net.URI;
+import java.net.URL;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+
+/**
+ * {@code AudioFactory} contains audio creation mechanism on top of different platforms like PC and
+ * Android. System will choose appropriate Factory based on the supported audio type.
+ */
+public abstract class AudioFactory {
+
+    private static final Logger logger = LoggerFactory.getLogger(AudioFactory.class);
+
+    private static final String[] FACTORIES = {
+        "ai.djl.audio.FFmpegAudioFactory", "ai.djl.modality.audio.SampledAudioFactory"
+    };
+
+    private static AudioFactory factory = newInstance();
+
+    private static AudioFactory newInstance() {
+        for (String f : FACTORIES) {
+            try {
+                Class<? extends AudioFactory> clazz =
+                        Class.forName(f).asSubclass(AudioFactory.class);
+                return clazz.getDeclaredConstructor().newInstance();
+            } catch (ReflectiveOperationException e) {
+                logger.trace("", e);
+            }
+        }
+        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.
+     *
+     * @param path the path to the audio
+     * @return {@link Audio}
+     * @throws IOException Audio not found or not readable
+     */
+    public abstract Audio fromFile(Path path) throws IOException;
+
+    /**
+     * Returns {@link Audio} from URL.
+     *
+     * @param url the URL to load from
+     * @return {@link Audio}
+     * @throws IOException URL is not valid.
+     */
+    public Audio fromUrl(URL url) throws IOException {
+        try (InputStream is = url.openStream()) {
+            return fromInputStream(is);
+        }
+    }
+
+    /**
+     * Returns {@link Audio} from URL.
+     *
+     * @param url the String represent URL to load from
+     * @return {@link Audio}
+     * @throws IOException URL is not valid.
+     */
+    public Audio fromUrl(String url) throws IOException {
+        URI uri = URI.create(url);
+        if (uri.isAbsolute()) {
+            return fromUrl(uri.toURL());
+        }
+        return fromFile(Paths.get(url));
+    }
+
+    /**
+     * Returns {@link Audio} from {@link InputStream}.
+     *
+     * @param is {@link InputStream}
+     * @return {@link Audio}
+     * @throws IOException image cannot be read from input stream.
+     */
+    public abstract Audio fromInputStream(InputStream is) throws IOException;
+
+    /**
+     * Returns {@link Audio} from raw data.
+     *
+     * @param data the raw data in float array form.
+     * @return {@link Audio}
+     */
+    public abstract Audio fromData(float[] data);
+
+    /**
+     * Returns {@link Audio} from {@link NDArray}.
+     *
+     * @param array the NDArray with CHW format
+     * @return {@link Audio}
+     */
+    public abstract Audio fromNDArray(NDArray array);
+}
diff --git a/api/src/main/java/ai/djl/modality/audio/SampledAudioFactory.java b/api/src/main/java/ai/djl/modality/audio/SampledAudioFactory.java
new file mode 100644
index 000000000..cdbeaccf2
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/audio/SampledAudioFactory.java
@@ -0,0 +1,117 @@
+/*
+ * 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.modality.audio;
+
+import ai.djl.modality.cv.ImageFactory;
+import ai.djl.ndarray.NDArray;
+
+import java.io.BufferedInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.ShortBuffer;
+import java.nio.file.Path;
+
+import javax.sound.sampled.AudioFormat;
+import javax.sound.sampled.AudioInputStream;
+import javax.sound.sampled.AudioSystem;
+import javax.sound.sampled.UnsupportedAudioFileException;
+
+/**
+ * {@code SampledAudioFactory} is an implementation of {@link ImageFactory} using the Java Sampled
+ * Package.
+ *
+ * @see <a
+ *     href="http://google.com">https://docs.oracle.com/javase/tutorial/sound/sampled-overview.html</a>
+ */
+public class SampledAudioFactory extends AudioFactory {
+
+    /** {@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);
+            return new Audio(floats, format.getSampleRate(), format.getChannels());
+        } catch (UnsupportedAudioFileException e) {
+            throw new IOException("Unsupported Audio file", e);
+        }
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Audio fromInputStream(InputStream is) throws IOException {
+        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);
+            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();
+
+        // Some audio formats may have unspecified frame size
+        if (frameSize == AudioSystem.NOT_SPECIFIED) {
+            frameSize = 1;
+        }
+
+        int size = (int) ais.getFrameLength() * frameSize;
+        byte[] ret = new byte[size];
+        byte[] buf = new byte[1024];
+        int offset = 0;
+        int read;
+        while ((read = ais.read(buf)) != -1) {
+            System.arraycopy(buf, 0, ret, offset, read);
+            offset += read;
+        }
+        return ret;
+    }
+
+    private float[] bytesToFloats(byte[] bytes, ByteOrder order) {
+        ShortBuffer buffer = ByteBuffer.wrap(bytes).order(order).asShortBuffer();
+        short[] shorts = new short[buffer.capacity()];
+        buffer.get(shorts);
+
+        // Feed in float values between -1.0f and 1.0f.
+        float[] floats = new float[shorts.length];
+        for (int i = 0; i < shorts.length; i++) {
+            floats[i] = ((float) shorts[i]) / (float) Short.MAX_VALUE;
+        }
+        return floats;
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/audio/package-info.java b/api/src/main/java/ai/djl/modality/audio/package-info.java
new file mode 100644
index 000000000..d9d2c2646
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/audio/package-info.java
@@ -0,0 +1,15 @@
+/*
+ * 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.
+ */
+
+/** Contains utility classes for audio processing. */
+package ai.djl.modality.audio;
diff --git a/api/src/main/java/ai/djl/modality/audio/translator/SpeechRecognitionTranslator.java b/api/src/main/java/ai/djl/modality/audio/translator/SpeechRecognitionTranslator.java
new file mode 100644
index 000000000..6d53de694
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/audio/translator/SpeechRecognitionTranslator.java
@@ -0,0 +1,43 @@
+/*
+ * 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.modality.audio.translator;
+
+import ai.djl.modality.audio.Audio;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.translate.NoBatchifyTranslator;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorContext;
+
+/**
+ * A {@link Translator} that post-process the {@link Audio} into {@link String} to get a text
+ * translation of the audio.
+ */
+public class SpeechRecognitionTranslator implements NoBatchifyTranslator<Audio, String> {
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, Audio input) throws Exception {
+        float[] data = input.getData();
+        NDArray array = ctx.getNDManager().create(data, new Shape(1, data.length));
+        return new NDList(array);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public String processOutput(TranslatorContext ctx, NDList list) throws Exception {
+        return list.get(0).toStringArray()[0];
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/audio/translator/SpeechRecognitionTranslatorFactory.java b/api/src/main/java/ai/djl/modality/audio/translator/SpeechRecognitionTranslatorFactory.java
new file mode 100644
index 000000000..f09a415e6
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/audio/translator/SpeechRecognitionTranslatorFactory.java
@@ -0,0 +1,52 @@
+/*
+ * 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.modality.audio.translator;
+
+import ai.djl.Model;
+import ai.djl.modality.audio.Audio;
+import ai.djl.translate.TranslateException;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorFactory;
+import ai.djl.util.Pair;
+
+import java.lang.reflect.Type;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+/** A {@link TranslatorFactory} that creates a {@link SpeechRecognitionTranslator} instance. */
+public class SpeechRecognitionTranslatorFactory implements TranslatorFactory {
+
+    private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>();
+
+    static {
+        SUPPORTED_TYPES.add(new Pair<>(Audio.class, String.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 {
+        if (input == Audio.class && output == String.class) {
+            return new SpeechRecognitionTranslator();
+        }
+        throw new IllegalArgumentException("Unsupported input/output types.");
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/audio/translator/package-info.java b/api/src/main/java/ai/djl/modality/audio/translator/package-info.java
new file mode 100644
index 000000000..661af2334
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/audio/translator/package-info.java
@@ -0,0 +1,15 @@
+/*
+ * 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.
+ */
+
+/** Contains utility classes for speech recognition processing. */
+package ai.djl.modality.audio.translator;
diff --git a/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index 0b5494840..8658f7389 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -587,7 +587,19 @@ public interface NDArray extends NDResource, BytesSupplier {
      * @param index picks the elements of an NDArray and output to the same entry as in index
      * @return the partial {@code NDArray} of the same shape as index
      */
-    NDArray take(NDArray index);
+    default NDArray take(NDArray index) {
+        return take(this.getManager(), index);
+    }
+
+    /**
+     * Returns a partial {@code NDArray} pointed by index according to linear indexing, and the of
+     * output is of the same shape as index.
+     *
+     * @param manager the manager used to create the arrays
+     * @param index picks the elements of an NDArray and output to the same entry as in index
+     * @return the partial {@code NDArray} of the same shape as index
+     */
+    NDArray take(NDManager manager, NDArray index);
 
     /**
      * Set the entries of {@code NDArray} pointed by index according to linear indexing, to be the
diff --git a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
index 519e1cafc..940ad0de8 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
@@ -193,7 +193,7 @@ public abstract class NDArrayAdapter implements NDArray {
 
     /** {@inheritDoc} */
     @Override
-    public NDArray take(NDArray index) {
+    public NDArray take(NDManager manager, NDArray index) {
         throw new UnsupportedOperationException(UNSUPPORTED_MSG);
     }
 
diff --git a/api/src/main/java/ai/djl/ndarray/index/NDArrayIndexer.java b/api/src/main/java/ai/djl/ndarray/index/NDArrayIndexer.java
index 3c7f67d4f..fe9fa7e2c 100644
--- a/api/src/main/java/ai/djl/ndarray/index/NDArrayIndexer.java
+++ b/api/src/main/java/ai/djl/ndarray/index/NDArrayIndexer.java
@@ -18,6 +18,10 @@ import ai.djl.ndarray.index.dim.NDIndexBooleans;
 import ai.djl.ndarray.index.dim.NDIndexElement;
 import ai.djl.ndarray.index.full.NDIndexFullPick;
 import ai.djl.ndarray.index.full.NDIndexFullSlice;
+import ai.djl.ndarray.index.full.NDIndexFullTake;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.List;
 import java.util.Optional;
@@ -34,6 +38,15 @@ public abstract class NDArrayIndexer {
      */
     public abstract NDArray get(NDArray array, NDIndexFullPick fullPick);
 
+    /**
+     * Returns a subarray by taken the elements from one axis.
+     *
+     * @param array the array to get from
+     * @param fullTake the elements to pick
+     * @return the subArray
+     */
+    public abstract NDArray get(NDArray array, NDIndexFullTake fullTake);
+
     /**
      * Returns a subarray at the slice.
      *
@@ -65,6 +78,16 @@ public abstract class NDArrayIndexer {
             return array.booleanMask(((NDIndexBooleans) indices.get(0)).getIndex());
         }
 
+        Optional<NDIndexFullTake> fullTake = NDIndexFullTake.fromIndex(index, array.getShape());
+        if (fullTake.isPresent()) {
+            Logger logger = LoggerFactory.getLogger(NDArrayIndexer.class);
+            logger.warn(
+                    "The definition of the getter by array NDIndex: get(NDIndex array) has changed"
+                        + " from pick to take.If you still want to use array index as pick, then do"
+                        + " it explicitly by get(new NDIndex().addPickDim(array));");
+            return get(array, fullTake.get());
+        }
+
         Optional<NDIndexFullPick> fullPick = NDIndexFullPick.fromIndex(index, array.getShape());
         if (fullPick.isPresent()) {
             return get(array, fullPick.get());
diff --git a/api/src/main/java/ai/djl/ndarray/index/full/NDIndexFullPick.java b/api/src/main/java/ai/djl/ndarray/index/full/NDIndexFullPick.java
index 7171ecec7..414d1c86a 100644
--- a/api/src/main/java/ai/djl/ndarray/index/full/NDIndexFullPick.java
+++ b/api/src/main/java/ai/djl/ndarray/index/full/NDIndexFullPick.java
@@ -17,7 +17,6 @@ import ai.djl.ndarray.index.NDIndex;
 import ai.djl.ndarray.index.dim.NDIndexAll;
 import ai.djl.ndarray.index.dim.NDIndexElement;
 import ai.djl.ndarray.index.dim.NDIndexPick;
-import ai.djl.ndarray.index.dim.NDIndexTake;
 import ai.djl.ndarray.types.Shape;
 
 import java.util.Optional;
@@ -60,18 +59,6 @@ public final class NDIndexFullPick {
                 }
                 NDArray indexElem = ((NDIndexPick) el).getIndex();
                 fullPick = new NDIndexFullPick(indexElem, axis);
-            } else if (el instanceof NDIndexTake) {
-                if (fullPick != null) {
-                    // Don't support multiple picks
-                    throw new UnsupportedOperationException(
-                            "Only one pick per get is currently supported");
-                }
-                NDArray indexElem = ((NDIndexTake) el).getIndex();
-                if (!indexElem.getShape().isRankOne()) {
-                    throw new UnsupportedOperationException(
-                            "Only rank-1 indexing array is supported for pick");
-                }
-                fullPick = new NDIndexFullPick(indexElem, axis);
             } else {
                 // Invalid dim for fullPick
                 return Optional.empty();
diff --git a/api/src/main/java/ai/djl/ndarray/index/full/NDIndexFullTake.java b/api/src/main/java/ai/djl/ndarray/index/full/NDIndexFullTake.java
new file mode 100644
index 000000000..22b269179
--- /dev/null
+++ b/api/src/main/java/ai/djl/ndarray/index/full/NDIndexFullTake.java
@@ -0,0 +1,91 @@
+/*
+ * Copyright 2020 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.ndarray.index.full;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.index.NDIndex;
+import ai.djl.ndarray.index.dim.NDIndexAll;
+import ai.djl.ndarray.index.dim.NDIndexElement;
+import ai.djl.ndarray.index.dim.NDIndexTake;
+import ai.djl.ndarray.types.Shape;
+
+import java.util.Optional;
+
+/** A simplified representation of a take-based {@link NDIndex}. */
+public final class NDIndexFullTake {
+
+    private NDArray indices;
+    private int axis;
+
+    /**
+     * Constructs a new {@link NDIndexFullTake}.
+     *
+     * @param indices the indices to take
+     * @param axis the axis to take at
+     */
+    private NDIndexFullTake(NDArray indices, int axis) {
+        this.indices = indices;
+        this.axis = axis;
+    }
+
+    /**
+     * Returns (if possible) the {@link NDIndexFullTake} representation of an {@link NDIndex}.
+     *
+     * @param index the index to represent
+     * @param target the shape of the array to index
+     * @return the full take representation or nothing if it can't represent the index
+     */
+    public static Optional<NDIndexFullTake> fromIndex(NDIndex index, Shape target) {
+        int axis = 0;
+        NDIndexFullTake fullTake = null;
+        for (NDIndexElement el : index.getIndices()) {
+            if (el instanceof NDIndexAll) {
+                axis++;
+            } else if (el instanceof NDIndexTake) {
+                if (fullTake != null) {
+                    // Don't support multiple takes
+                    throw new UnsupportedOperationException(
+                            "Only one take per get is currently supported");
+                }
+                NDArray indexElem = ((NDIndexTake) el).getIndex();
+                if (!indexElem.getShape().isRankOne()) {
+                    throw new UnsupportedOperationException(
+                            "Only rank-1 indexing array is supported for take");
+                }
+                fullTake = new NDIndexFullTake(indexElem, axis);
+            } else {
+                // Invalid dim for fullTake
+                return Optional.empty();
+            }
+        }
+        return Optional.ofNullable(fullTake);
+    }
+
+    /**
+     * Returns the indices to take.
+     *
+     * @return the indices to take
+     */
+    public NDArray getIndices() {
+        return indices;
+    }
+
+    /**
+     * Returns the axis to take.
+     *
+     * @return the axis to take
+     */
+    public int getAxis() {
+        return axis;
+    }
+}
diff --git a/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java b/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java
index 4a14cb7c5..8abb4ca44 100644
--- a/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java
+++ b/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java
@@ -33,6 +33,8 @@ import java.util.Map;
 import java.util.ServiceLoader;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
 class RepositoryFactoryImpl implements RepositoryFactory {
 
@@ -40,6 +42,7 @@ class RepositoryFactoryImpl implements RepositoryFactory {
 
     private static final RepositoryFactory FACTORY = new RepositoryFactoryImpl();
     private static final Map<String, RepositoryFactory> REGISTRY = init();
+    private static final Pattern PATTERN = Pattern.compile("(.+)/([\\d.]+)(/(.*))?");
 
     static RepositoryFactory getFactory() {
         return FACTORY;
@@ -217,17 +220,21 @@ class RepositoryFactoryImpl implements RepositoryFactory {
             if (groupId == null) {
                 throw new IllegalArgumentException("Invalid djl URL: " + uri);
             }
-            Path path = parseFilePath(uri);
-            int size = path.getNameCount();
-            if (size == 0) {
-                throw new IllegalArgumentException("Invalid djl URL: " + uri);
+            String artifactId = parseFilePath(uri).toString();
+            if (artifactId.startsWith("/")) {
+                artifactId = artifactId.substring(1);
             }
-            String artifactId = path.getName(0).toString();
             if (artifactId.isEmpty()) {
                 throw new IllegalArgumentException("Invalid djl URL: " + uri);
             }
-            String version = size > 1 ? path.getName(1).toString() : null;
-            String artifactName = size > 2 ? path.getName(2).toString() : null;
+            String version = null;
+            String artifactName = null;
+            Matcher m = PATTERN.matcher(artifactId);
+            if (m.matches()) {
+                artifactId = m.group(1);
+                version = m.group(2);
+                artifactName = m.group(4);
+            }
 
             ModelZoo zoo = ModelZoo.getModelZoo(groupId);
             if (zoo == null) {
diff --git a/api/src/main/java/ai/djl/repository/zoo/DefaultModelZoo.java b/api/src/main/java/ai/djl/repository/zoo/DefaultModelZoo.java
index 17c259dd2..527871067 100644
--- a/api/src/main/java/ai/djl/repository/zoo/DefaultModelZoo.java
+++ b/api/src/main/java/ai/djl/repository/zoo/DefaultModelZoo.java
@@ -19,9 +19,6 @@ import ai.djl.repository.Repository;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
 import java.util.Set;
 
 /** A {@link ModelZoo} that contains models in specified locations. */
@@ -31,10 +28,13 @@ public class DefaultModelZoo extends ModelZoo {
 
     private static final Logger logger = LoggerFactory.getLogger(DefaultModelZoo.class);
 
-    private List<ModelLoader> modelLoaders;
-
     /** Constructs a new {@code LocalModelZoo} instance. */
-    public DefaultModelZoo() {}
+    public DefaultModelZoo() {
+        String locations = System.getProperty("ai.djl.repository.zoo.location");
+        if (locations != null) {
+            parseLocation(locations);
+        }
+    }
 
     /**
      * Constructs a new {@code LocalModelZoo} instance from the given search locations.
@@ -42,20 +42,7 @@ public class DefaultModelZoo extends ModelZoo {
      * @param locations a comma separated urls where the models to be loaded from
      */
     public DefaultModelZoo(String locations) {
-        modelLoaders = parseLocation(locations);
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public List<ModelLoader> getModelLoaders() {
-        if (modelLoaders != null) {
-            return modelLoaders;
-        }
-        String locations = System.getProperty("ai.djl.repository.zoo.location");
-        if (locations != null) {
-            return parseLocation(locations);
-        }
-        return Collections.emptyList();
+        parseLocation(locations);
     }
 
     /** {@inheritDoc} */
@@ -70,21 +57,18 @@ public class DefaultModelZoo extends ModelZoo {
         return Engine.getAllEngines();
     }
 
-    private List<ModelLoader> parseLocation(String locations) {
+    private void parseLocation(String locations) {
         String[] urls = locations.split("\\s*,\\s*");
-        List<ModelLoader> list = new ArrayList<>(urls.length);
         for (String url : urls) {
             if (!url.isEmpty()) {
                 Repository repo = Repository.newInstance(url, url);
                 logger.debug("Scanning models in repo: {}, {}", repo.getClass(), url);
-                List<MRL> mrls = repo.getResources();
-                for (MRL mrl : mrls) {
-                    list.add(new BaseModelLoader(mrl));
+                for (MRL mrl : repo.getResources()) {
+                    addModel(mrl);
                 }
             } else {
                 logger.warn("Model location is empty.");
             }
         }
-        return list;
     }
 }
diff --git a/api/src/main/java/ai/djl/repository/zoo/ModelZoo.java b/api/src/main/java/ai/djl/repository/zoo/ModelZoo.java
index 0674c9d00..50b219be5 100644
--- a/api/src/main/java/ai/djl/repository/zoo/ModelZoo.java
+++ b/api/src/main/java/ai/djl/repository/zoo/ModelZoo.java
@@ -15,6 +15,7 @@ package ai.djl.repository.zoo;
 import ai.djl.Application;
 import ai.djl.MalformedModelException;
 import ai.djl.repository.Artifact;
+import ai.djl.repository.MRL;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -32,6 +33,8 @@ public abstract class ModelZoo {
 
     private static final Map<String, ModelZoo> MODEL_ZOO_MAP = new ConcurrentHashMap<>();
 
+    private Map<String, ModelLoader> modelLoaders = new ConcurrentHashMap<>();
+
     static {
         ServiceLoader<ZooProvider> providers = ServiceLoader.load(ZooProvider.class);
         for (ZooProvider provider : providers) {
@@ -54,7 +57,9 @@ public abstract class ModelZoo {
      *
      * @return the list of all available model families
      */
-    public abstract List<ModelLoader> getModelLoaders();
+    public Collection<ModelLoader> getModelLoaders() {
+        return modelLoaders.values();
+    }
 
     /**
      * Returns the {@link ModelLoader} based on the model name.
@@ -63,12 +68,7 @@ public abstract class ModelZoo {
      * @return the {@link ModelLoader} of the model
      */
     public ModelLoader getModelLoader(String name) {
-        for (ModelLoader loader : getModelLoaders()) {
-            if (name.equals(loader.getArtifactId())) {
-                return loader;
-            }
-        }
-        return null;
+        return modelLoaders.get(name);
     }
 
     /**
@@ -78,6 +78,14 @@ public abstract class ModelZoo {
      */
     public abstract Set<String> getSupportedEngines();
 
+    protected final void addModel(MRL mrl) {
+        modelLoaders.put(mrl.getArtifactId(), new BaseModelLoader(mrl));
+    }
+
+    protected final void addModel(ModelLoader loader) {
+        modelLoaders.put(loader.getArtifactId(), loader);
+    }
+
     /**
      * Refreshes model zoo.
      *
@@ -85,7 +93,7 @@ public abstract class ModelZoo {
      */
     public static void registerModelZoo(ZooProvider provider) {
         ModelZoo zoo = provider.getModelZoo();
-        MODEL_ZOO_MAP.putIfAbsent(zoo.getGroupId(), zoo);
+        MODEL_ZOO_MAP.put(zoo.getGroupId(), zoo);
     }
 
     /**
@@ -174,8 +182,7 @@ public abstract class ModelZoo {
                     continue;
                 }
             }
-            List<ModelLoader> list = zoo.getModelLoaders();
-            for (ModelLoader loader : list) {
+            for (ModelLoader loader : zoo.getModelLoaders()) {
                 Application app = loader.getApplication();
                 String loaderArtifactId = loader.getArtifactId();
                 if (artifactId != null && !artifactId.equals(loaderArtifactId)) {
diff --git a/api/src/main/java/ai/djl/training/dataset/ArrayDataset.java b/api/src/main/java/ai/djl/training/dataset/ArrayDataset.java
index 3dee4e607..c658ae18c 100644
--- a/api/src/main/java/ai/djl/training/dataset/ArrayDataset.java
+++ b/api/src/main/java/ai/djl/training/dataset/ArrayDataset.java
@@ -15,9 +15,14 @@ package ai.djl.training.dataset;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.index.NDIndex;
+import ai.djl.translate.Batchifier;
+import ai.djl.translate.TranslateException;
 import ai.djl.util.Progress;
 
 import java.io.IOException;
+import java.util.List;
+import java.util.concurrent.ExecutorService;
 import java.util.stream.Stream;
 
 /**
@@ -73,6 +78,8 @@ public class ArrayDataset extends RandomAccessDataset {
         }
     }
 
+    ArrayDataset() {}
+
     /** {@inheritDoc} */
     @Override
     protected long availableSize() {
@@ -95,6 +102,96 @@ public class ArrayDataset extends RandomAccessDataset {
         return new Record(datum, label);
     }
 
+    /**
+     * Gets the {@link Record} for the given indices from the dataset.
+     *
+     * @param manager the manager used to create the arrays
+     * @param indices indices of the requested data items
+     * @return a {@link Record} that contains the data and label of the requested data items
+     */
+    public Record getByIndices(NDManager manager, long... indices) {
+        try (NDArray ndIndices = manager.create(indices)) {
+            NDIndex index = new NDIndex("{}", ndIndices);
+            NDList datum = new NDList();
+            NDList label = new NDList();
+            for (NDArray array : data) {
+                datum.add(array.get(manager, index));
+            }
+            if (labels != null) {
+                for (NDArray array : labels) {
+                    label.add(array.get(manager, index));
+                }
+            }
+            return new Record(datum, label);
+        }
+    }
+
+    /**
+     * Gets the {@link Record} for the given range from the dataset.
+     *
+     * @param manager the manager used to create the arrays
+     * @param fromIndex low endpoint (inclusive) of the dataset
+     * @param toIndex high endpoint (exclusive) of the dataset
+     * @return a {@link Record} that contains the data and label of the requested data items
+     */
+    public Record getByRange(NDManager manager, long fromIndex, long toIndex) {
+        NDIndex index = new NDIndex().addSliceDim(fromIndex, toIndex);
+        NDList datum = new NDList();
+        NDList label = new NDList();
+        for (NDArray array : data) {
+            datum.add(array.get(manager, index));
+        }
+        if (labels != null) {
+            for (NDArray array : labels) {
+                label.add(array.get(manager, index));
+            }
+        }
+        return new Record(datum, label);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    protected RandomAccessDataset newSubDataset(int[] indices, int from, int to) {
+        return new SubDataset(this, indices, from, to);
+    }
+
+    @Override
+    protected RandomAccessDataset newSubDataset(List<Long> subIndices) {
+        return new SubDatasetByIndices(this, subIndices);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Iterable<Batch> getData(
+            NDManager manager, Sampler sampler, ExecutorService executorService)
+            throws IOException, TranslateException {
+        prepare();
+        if (dataBatchifier == Batchifier.STACK && labelBatchifier == Batchifier.STACK) {
+            return new BulkDataIterable(
+                    this,
+                    manager,
+                    sampler,
+                    dataBatchifier,
+                    labelBatchifier,
+                    pipeline,
+                    targetPipeline,
+                    executorService,
+                    prefetchNumber,
+                    device);
+        }
+        return new DataIterable(
+                this,
+                manager,
+                sampler,
+                dataBatchifier,
+                labelBatchifier,
+                pipeline,
+                targetPipeline,
+                executorService,
+                prefetchNumber,
+                device);
+    }
+
     /** {@inheritDoc} */
     @Override
     public void prepare(Progress progress) throws IOException {}
@@ -145,4 +242,122 @@ public class ArrayDataset extends RandomAccessDataset {
             return new ArrayDataset(this);
         }
     }
+
+    private static final class SubDataset extends ArrayDataset {
+
+        private ArrayDataset dataset;
+        private int[] indices;
+        private int from;
+        private int to;
+
+        public SubDataset(ArrayDataset dataset, int[] indices, int from, int to) {
+            this.dataset = dataset;
+            this.indices = indices;
+            this.from = from;
+            this.to = to;
+            this.sampler = dataset.sampler;
+            this.dataBatchifier = dataset.dataBatchifier;
+            this.labelBatchifier = dataset.labelBatchifier;
+            this.pipeline = dataset.pipeline;
+            this.targetPipeline = dataset.targetPipeline;
+            this.prefetchNumber = dataset.prefetchNumber;
+            this.device = dataset.device;
+
+            limit = Long.MAX_VALUE;
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public Record get(NDManager manager, long index) {
+            if (index >= size()) {
+                throw new IndexOutOfBoundsException("index(" + index + ") > size(" + size() + ").");
+            }
+            return dataset.get(manager, indices[Math.toIntExact(index) + from]);
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public Record getByIndices(NDManager manager, long... indices) {
+            long[] resolvedIndices = new long[indices.length];
+            int i = 0;
+            for (long index : indices) {
+                resolvedIndices[i++] = this.indices[Math.toIntExact(index) + from];
+            }
+            return dataset.getByIndices(manager, resolvedIndices);
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public Record getByRange(NDManager manager, long fromIndex, long toIndex) {
+            return dataset.getByRange(manager, fromIndex + from, toIndex + from);
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        protected long availableSize() {
+            return to - from;
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public void prepare(Progress progress) {}
+    }
+
+    private static final class SubDatasetByIndices extends ArrayDataset {
+
+        private ArrayDataset dataset;
+        private List<Long> subIndices;
+
+        public SubDatasetByIndices(ArrayDataset dataset, List<Long> subIndices) {
+            this.dataset = dataset;
+            this.subIndices = subIndices;
+            this.sampler = dataset.sampler;
+            this.dataBatchifier = dataset.dataBatchifier;
+            this.labelBatchifier = dataset.labelBatchifier;
+            this.pipeline = dataset.pipeline;
+            this.targetPipeline = dataset.targetPipeline;
+            this.prefetchNumber = dataset.prefetchNumber;
+            this.device = dataset.device;
+
+            limit = Long.MAX_VALUE;
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public Record get(NDManager manager, long index) {
+            return dataset.get(manager, subIndices.get(Math.toIntExact(index)));
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public Record getByIndices(NDManager manager, long... indices) {
+            long[] resolvedIndices = new long[indices.length];
+            int i = 0;
+            for (long index : indices) {
+                resolvedIndices[i++] = subIndices.get(Math.toIntExact(index));
+            }
+            return dataset.getByIndices(manager, resolvedIndices);
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public Record getByRange(NDManager manager, long fromIndex, long toIndex) {
+            long[] resolvedIndices = new long[(int) (toIndex - fromIndex)];
+            int i = 0;
+            for (long index = fromIndex; index < toIndex; index++) {
+                resolvedIndices[i++] = subIndices.get(Math.toIntExact(index));
+            }
+            return dataset.getByIndices(manager, resolvedIndices);
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        protected long availableSize() {
+            return subIndices.size();
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public void prepare(Progress progress) {}
+    }
 }
diff --git a/api/src/main/java/ai/djl/training/dataset/BulkDataIterable.java b/api/src/main/java/ai/djl/training/dataset/BulkDataIterable.java
new file mode 100644
index 000000000..521697779
--- /dev/null
+++ b/api/src/main/java/ai/djl/training/dataset/BulkDataIterable.java
@@ -0,0 +1,134 @@
+/*
+ * 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.training.dataset;
+
+import ai.djl.Device;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.translate.Batchifier;
+import ai.djl.translate.Pipeline;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.concurrent.ExecutorService;
+
+/**
+ * BulkDataIterable specializes DataIterable in using {@link ArrayDataset#getByRange(NDManager,
+ * long, long)} or {@link ArrayDataset#getByIndices(NDManager, long...)} to create {@link Batch}
+ * instances more efficiently.
+ */
+public class BulkDataIterable extends DataIterable {
+
+    /**
+     * Creates a new instance of {@code BulkDataIterable} with the given parameters.
+     *
+     * @param dataset the dataset to iterate on
+     * @param manager the manager to create the arrays
+     * @param sampler a sampler to sample data with
+     * @param dataBatchifier a batchifier for data
+     * @param labelBatchifier a batchifier for labels
+     * @param pipeline the pipeline of transforms to apply on the data
+     * @param targetPipeline the pipeline of transforms to apply on the labels
+     * @param executor an {@link ExecutorService}
+     * @param preFetchNumber the number of samples to prefetch
+     * @param device the {@link Device}
+     */
+    public BulkDataIterable(
+            ArrayDataset dataset,
+            NDManager manager,
+            Sampler sampler,
+            Batchifier dataBatchifier,
+            Batchifier labelBatchifier,
+            Pipeline pipeline,
+            Pipeline targetPipeline,
+            ExecutorService executor,
+            int preFetchNumber,
+            Device device) {
+        super(
+                dataset,
+                manager,
+                sampler,
+                dataBatchifier,
+                labelBatchifier,
+                pipeline,
+                targetPipeline,
+                executor,
+                preFetchNumber,
+                device);
+    }
+
+    @Override
+    protected Batch fetch(List<Long> indices, int progress) throws IOException {
+        NDManager subManager = manager.newSubManager();
+        subManager.setName("dataIter fetch");
+        int batchSize = indices.size();
+
+        Record record;
+        if (isRange(indices)) {
+            long fromIndex = indices.get(0);
+            long toIndex = fromIndex + indices.size();
+            record = ((ArrayDataset) dataset).getByRange(subManager, fromIndex, toIndex);
+        } else {
+            long[] indicesArr = indices.stream().mapToLong(Long::longValue).toArray();
+            record = ((ArrayDataset) dataset).getByIndices(subManager, indicesArr);
+        }
+
+        NDList batchData = record.getData();
+        // apply transform
+        if (pipeline != null) {
+            batchData = pipeline.transform(batchData);
+        }
+
+        NDList batchLabels = record.getLabels();
+
+        // apply label transform
+        if (targetPipeline != null) {
+            batchLabels = targetPipeline.transform(batchLabels);
+        }
+        // pin to a specific device
+        if (device != null) {
+            batchData = batchData.toDevice(device, false);
+            batchLabels = batchLabels.toDevice(device, false);
+        }
+        return new Batch(
+                subManager,
+                batchData,
+                batchLabels,
+                batchSize,
+                dataBatchifier,
+                labelBatchifier,
+                progress,
+                dataset.size(),
+                indices);
+    }
+
+    /**
+     * Checks whether the given indices actually represents a range.
+     *
+     * @param indices the indices to examine
+     * @return whether the given indices are sorted in ascending order with no gap and has at least
+     *     one element
+     */
+    public static boolean isRange(List<Long> indices) {
+        if (indices.isEmpty()) {
+            return false;
+        }
+        long from = indices.get(0);
+        for (long index : indices) {
+            if (index != from++) {
+                return false;
+            }
+        }
+        return true;
+    }
+}
diff --git a/api/src/main/java/ai/djl/training/dataset/DataIterable.java b/api/src/main/java/ai/djl/training/dataset/DataIterable.java
index 0f478eea1..aa5660ba2 100644
--- a/api/src/main/java/ai/djl/training/dataset/DataIterable.java
+++ b/api/src/main/java/ai/djl/training/dataset/DataIterable.java
@@ -44,14 +44,14 @@ public class DataIterable implements Iterable<Batch>, Iterator<Batch> {
 
     private static final Logger logger = LoggerFactory.getLogger(DataIterable.class);
 
-    private RandomAccessDataset dataset;
-    private NDManager manager;
-    private Batchifier dataBatchifier;
-    private Batchifier labelBatchifier;
-    private Pipeline pipeline;
-    private Pipeline targetPipeline;
+    protected RandomAccessDataset dataset;
+    protected NDManager manager;
+    protected Batchifier dataBatchifier;
+    protected Batchifier labelBatchifier;
+    protected Pipeline pipeline;
+    protected Pipeline targetPipeline;
     private ExecutorService executor;
-    private Device device;
+    protected Device device;
 
     private Iterator<List<Long>> sample;
     // for multithreading
@@ -160,7 +160,7 @@ public class DataIterable implements Iterable<Batch>, Iterator<Batch> {
         }
     }
 
-    private Batch fetch(List<Long> indices, int progress) throws IOException {
+    protected Batch fetch(List<Long> indices, int progress) throws IOException {
         NDManager subManager = manager.newSubManager();
         subManager.setName("dataIter fetch");
         int batchSize = indices.size();
diff --git a/api/src/main/java/ai/djl/training/dataset/RandomAccessDataset.java b/api/src/main/java/ai/djl/training/dataset/RandomAccessDataset.java
index 0ee0cc727..0d3e6a45c 100644
--- a/api/src/main/java/ai/djl/training/dataset/RandomAccessDataset.java
+++ b/api/src/main/java/ai/djl/training/dataset/RandomAccessDataset.java
@@ -26,8 +26,12 @@ import ai.djl.util.Progress;
 import ai.djl.util.RandomUtils;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -84,36 +88,14 @@ public abstract class RandomAccessDataset implements Dataset {
     /** {@inheritDoc} */
     @Override
     public Iterable<Batch> getData(NDManager manager) throws IOException, TranslateException {
-        prepare();
-        return new DataIterable(
-                this,
-                manager,
-                sampler,
-                dataBatchifier,
-                labelBatchifier,
-                pipeline,
-                targetPipeline,
-                null,
-                prefetchNumber,
-                device);
+        return getData(manager, sampler, null);
     }
 
     /** {@inheritDoc} */
     @Override
     public Iterable<Batch> getData(NDManager manager, ExecutorService executorService)
             throws IOException, TranslateException {
-        prepare();
-        return new DataIterable(
-                this,
-                manager,
-                sampler,
-                dataBatchifier,
-                labelBatchifier,
-                pipeline,
-                targetPipeline,
-                executorService,
-                prefetchNumber,
-                device);
+        return getData(manager, sampler, executorService);
     }
 
     /**
@@ -127,18 +109,7 @@ public abstract class RandomAccessDataset implements Dataset {
      */
     public Iterable<Batch> getData(NDManager manager, Sampler sampler)
             throws IOException, TranslateException {
-        prepare();
-        return new DataIterable(
-                this,
-                manager,
-                sampler,
-                dataBatchifier,
-                labelBatchifier,
-                pipeline,
-                targetPipeline,
-                null,
-                prefetchNumber,
-                device);
+        return getData(manager, sampler, null);
     }
 
     /**
@@ -209,10 +180,10 @@ public abstract class RandomAccessDataset implements Dataset {
         int from = 0;
         for (int i = 0; i < ratio.length - 1; ++i) {
             int to = from + (int) (ratio[i] / sum * size);
-            ret[i] = new SubDataset(this, indices, from, to);
+            ret[i] = newSubDataset(indices, from, to);
             from = to;
         }
-        ret[ratio.length - 1] = new SubDataset(this, indices, from, size);
+        ret[ratio.length - 1] = newSubDataset(indices, from, size);
         return ret;
     }
 
@@ -227,7 +198,94 @@ public abstract class RandomAccessDataset implements Dataset {
     public RandomAccessDataset subDataset(int fromIndex, int toIndex) {
         int size = Math.toIntExact(size());
         int[] indices = IntStream.range(0, size).toArray();
-        return new SubDataset(this, indices, fromIndex, toIndex);
+        return newSubDataset(indices, fromIndex, toIndex);
+    }
+
+    /**
+     * Returns a view of the portion of this data for the specified {@code subIndices}.
+     *
+     * @param subIndices sub-set of indices of this dataset
+     * @return a view of the specified indices within this dataset
+     */
+    public RandomAccessDataset subDataset(List<Long> subIndices) {
+        if (BulkDataIterable.isRange(subIndices)) {
+            int size = Math.toIntExact(size());
+            int[] indices = IntStream.range(0, size).toArray();
+            long fromIndex = subIndices.get(0);
+            long toIndex = subIndices.get(0) + subIndices.size();
+            return newSubDataset(indices, Math.toIntExact(fromIndex), Math.toIntExact(toIndex));
+        }
+        return newSubDataset(subIndices);
+    }
+
+    /**
+     * Returns a view of the portion of this data for the specified record keys. Assuming that the
+     * records of this database are represented by the keys in <code>recordKeys</code>, then <code>
+     * subRecordKeys</code> defines the view on the corresponding records of the database.
+     *
+     * @param recordKeys unique keys for all records of this dataset.
+     * @param subRecordKeys keys to define the view on the dataset. All keys in <code>subRecordKeys
+     *     </code> must be contained in <code>recordKeys</code> but may occur more than once.
+     * @param <K> the record key type.
+     * @return a view of the specified records within this dataset
+     */
+    public <K> RandomAccessDataset subDataset(List<K> recordKeys, List<K> subRecordKeys) {
+        if (this.size() != recordKeys.size()) {
+            throw new IllegalArgumentException(
+                    "Requires as many record keys as there are records in the dataset.");
+        }
+        Map<K, Long> indicesOfRecordKeys = new ConcurrentHashMap<>(recordKeys.size());
+        for (int index = 0; index < recordKeys.size(); index++) {
+            Long prevIndex = indicesOfRecordKeys.put(recordKeys.get(index), (long) index);
+            if (prevIndex != null) {
+                throw new IllegalArgumentException(
+                        "At least two keys at position "
+                                + prevIndex
+                                + " and "
+                                + index
+                                + " are equal!");
+            }
+        }
+        return subDataset(indicesOfRecordKeys, subRecordKeys);
+    }
+
+    /**
+     * Returns a view of the portion of this data for the specified record keys. Assuming that the
+     * records of this database are represented by the keys in <code>indicesOfRecordKeys</code>,
+     * then <code>
+     * subRecordKeys</code> defines the view on the corresponding records of the database.
+     *
+     * @param indicesOfRecordKeys Map for keys of the records in this dataset to their index
+     *     position within this dataset. While this map typically maps all records, technically it
+     *     just needs to map the ones occuring in <code>subRecordKeys</code>.
+     * @param subRecordKeys Keys to define the view on the dataset. All keys in <code>subRecordKeys
+     *     </code> must be contained in <code>indicesOfRecordKeys</code> but may occur more than
+     *     once.
+     * @param <K> the record key type.
+     * @return a view of the records identified by the specified keys of this dataset
+     */
+    public <K> RandomAccessDataset subDataset(
+            Map<K, Long> indicesOfRecordKeys, List<K> subRecordKeys) {
+        List<Long> subIndices = new ArrayList<>(subRecordKeys.size());
+        for (K recordKey : subRecordKeys) {
+            Long index = indicesOfRecordKeys.get(recordKey);
+            if (index == null) {
+                throw new IllegalArgumentException(
+                        "The key of subRecordKeys at position "
+                                + subRecordKeys.indexOf(recordKey)
+                                + " is not contained in recordKeys!");
+            }
+            subIndices.add(index);
+        }
+        return subDataset(subIndices);
+    }
+
+    protected RandomAccessDataset newSubDataset(int[] indices, int from, int to) {
+        return new SubDataset(this, indices, from, to);
+    }
+
+    protected RandomAccessDataset newSubDataset(List<Long> subIndices) {
+        return new SubDatasetByIndices(this, subIndices);
     }
 
     /**
@@ -498,4 +556,40 @@ public abstract class RandomAccessDataset implements Dataset {
         @Override
         public void prepare(Progress progress) {}
     }
+
+    private static final class SubDatasetByIndices extends RandomAccessDataset {
+
+        private RandomAccessDataset dataset;
+        private List<Long> subIndices;
+
+        public SubDatasetByIndices(RandomAccessDataset dataset, List<Long> subIndices) {
+            this.dataset = dataset;
+            this.subIndices = subIndices;
+            this.sampler = dataset.sampler;
+            this.dataBatchifier = dataset.dataBatchifier;
+            this.labelBatchifier = dataset.labelBatchifier;
+            this.pipeline = dataset.pipeline;
+            this.targetPipeline = dataset.targetPipeline;
+            this.prefetchNumber = dataset.prefetchNumber;
+            this.device = dataset.device;
+
+            limit = Long.MAX_VALUE;
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public Record get(NDManager manager, long index) throws IOException {
+            return dataset.get(manager, subIndices.get(Math.toIntExact(index)));
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        protected long availableSize() {
+            return subIndices.size();
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public void prepare(Progress progress) {}
+    }
 }
diff --git a/engines/dlr/dlr-engine/src/main/java/ai/djl/dlr/zoo/DlrModelZoo.java b/engines/dlr/dlr-engine/src/main/java/ai/djl/dlr/zoo/DlrModelZoo.java
index 897a04dbf..a4b8f0944 100644
--- a/engines/dlr/dlr-engine/src/main/java/ai/djl/dlr/zoo/DlrModelZoo.java
+++ b/engines/dlr/dlr-engine/src/main/java/ai/djl/dlr/zoo/DlrModelZoo.java
@@ -14,15 +14,10 @@ package ai.djl.dlr.zoo;
 
 import ai.djl.Application.CV;
 import ai.djl.dlr.engine.DlrEngine;
-import ai.djl.repository.MRL;
 import ai.djl.repository.Repository;
-import ai.djl.repository.zoo.BaseModelLoader;
-import ai.djl.repository.zoo.ModelLoader;
 import ai.djl.repository.zoo.ModelZoo;
 
-import java.util.ArrayList;
 import java.util.Collections;
-import java.util.List;
 import java.util.Set;
 
 /** DlrModelZoo is a repository that contains all dlr models for DJL. */
@@ -32,17 +27,8 @@ public class DlrModelZoo extends ModelZoo {
     private static final Repository REPOSITORY = Repository.newInstance("Dlr", DJL_REPO_URL);
     public static final String GROUP_ID = "ai.djl.dlr";
 
-    private static final List<ModelLoader> MODEL_LOADERS = new ArrayList<>();
-
-    static {
-        MRL resnet = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "resnet", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(resnet));
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public List<ModelLoader> getModelLoaders() {
-        return MODEL_LOADERS;
+    DlrModelZoo() {
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "resnet", "0.0.1"));
     }
 
     /** {@inheritDoc} */
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxDeviceType.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxDeviceType.java
index 7d4e754a8..6b38c2f40 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxDeviceType.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxDeviceType.java
@@ -13,10 +13,9 @@
 package ai.djl.mxnet.engine;
 
 import ai.djl.Device;
-import ai.djl.DeviceType;
 
 /** DeviceType is the MXNet equivalent of the types in {@link Device}. */
-public final class MxDeviceType implements DeviceType {
+public final class MxDeviceType {
 
     private static final String CPU_PINNED = "cpu_pinned";
 
@@ -38,7 +37,7 @@ public final class MxDeviceType implements DeviceType {
         } else if (CPU_PINNED.equals(deviceType)) {
             return 3;
         } else {
-            throw new IllegalArgumentException("Unsupported device: " + device.toString());
+            throw new IllegalArgumentException("Unsupported device: " + device);
         }
     }
 
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
index 0974876b2..78b041ba4 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
@@ -324,7 +324,7 @@ public class MxNDArray extends NativeResource<Pointer> implements LazyNDArray {
 
     /** {@inheritDoc} */
     @Override
-    public NDArray take(NDArray index) {
+    public NDArray take(NDManager manager, NDArray index) {
         MxOpParams params = new MxOpParams();
         params.add("mode", "wrap");
         return manager.invoke("take", new NDList(this.flatten(), index), params).singletonOrThrow();
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArrayIndexer.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArrayIndexer.java
index 76ab1b1aa..6812aa4a1 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArrayIndexer.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArrayIndexer.java
@@ -17,6 +17,7 @@ import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.index.NDArrayIndexer;
 import ai.djl.ndarray.index.full.NDIndexFullPick;
 import ai.djl.ndarray.index.full.NDIndexFullSlice;
+import ai.djl.ndarray.index.full.NDIndexFullTake;
 import ai.djl.ndarray.types.Shape;
 
 import java.util.Stack;
@@ -42,6 +43,17 @@ public class MxNDArrayIndexer extends NDArrayIndexer {
                 .singletonOrThrow();
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray get(NDArray array, NDIndexFullTake fullTake) {
+        array = manager.from(array);
+        MxOpParams params = new MxOpParams();
+        params.addParam("axis", fullTake.getAxis());
+        params.add("mode", "wrap");
+        return manager.invoke("take", new NDList(array, fullTake.getIndices()), params)
+                .singletonOrThrow();
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray get(NDArray array, NDIndexFullSlice fullSlice) {
diff --git a/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/MxModelZoo.java b/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/MxModelZoo.java
index 162d4b994..30f7884e0 100644
--- a/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/MxModelZoo.java
+++ b/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/MxModelZoo.java
@@ -15,15 +15,10 @@ package ai.djl.mxnet.zoo;
 import ai.djl.Application.CV;
 import ai.djl.Application.NLP;
 import ai.djl.mxnet.engine.MxEngine;
-import ai.djl.repository.MRL;
 import ai.djl.repository.Repository;
-import ai.djl.repository.zoo.BaseModelLoader;
-import ai.djl.repository.zoo.ModelLoader;
 import ai.djl.repository.zoo.ModelZoo;
 
-import java.util.ArrayList;
 import java.util.Collections;
-import java.util.List;
 import java.util.Set;
 
 /**
@@ -36,79 +31,28 @@ public class MxModelZoo extends ModelZoo {
     private static final Repository REPOSITORY = Repository.newInstance("MXNet", DJL_REPO_URL);
     public static final String GROUP_ID = "ai.djl.mxnet";
 
-    private static final List<ModelLoader> MODEL_LOADERS = new ArrayList<>();
-
-    static {
-        MRL ssd = REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "ssd", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(ssd));
-
-        MRL yolo = REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "yolo", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(yolo));
-
-        MRL alexnet = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "alexnet", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(alexnet));
-
-        MRL darknet = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "darknet", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(darknet));
-
-        MRL densenet = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "densenet", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(densenet));
-
-        MRL googlenet = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "googlenet", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(googlenet));
-
-        MRL inceptionv3 =
-                REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "inceptionv3", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(inceptionv3));
-
-        MRL mlp = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "mlp", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(mlp));
-
-        MRL mobilenet = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "mobilenet", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(mobilenet));
-
-        MRL resnest = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "resnest", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(resnest));
-
-        MRL resnet = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "resnet", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(resnet));
-
-        MRL senet = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "senet", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(senet));
-
-        MRL seresnext = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "se_resnext", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(seresnext));
-
-        MRL squeezenet = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "squeezenet", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(squeezenet));
-
-        MRL vgg = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "vgg", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(vgg));
-
-        MRL xception = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "xception", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(xception));
-
-        MRL simplePose = REPOSITORY.model(CV.POSE_ESTIMATION, GROUP_ID, "simple_pose", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(simplePose));
-
-        MRL maskrcnn = REPOSITORY.model(CV.INSTANCE_SEGMENTATION, GROUP_ID, "mask_rcnn", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(maskrcnn));
-
-        MRL actionRecognition =
-                REPOSITORY.model(CV.ACTION_RECOGNITION, GROUP_ID, "action_recognition", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(actionRecognition));
-
-        MRL bertQa = REPOSITORY.model(NLP.QUESTION_ANSWER, GROUP_ID, "bertqa", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(bertQa));
-
-        MRL glove = REPOSITORY.model(NLP.WORD_EMBEDDING, GROUP_ID, "glove", "0.0.2");
-        MODEL_LOADERS.add(new BaseModelLoader(glove));
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public List<ModelLoader> getModelLoaders() {
-        return MODEL_LOADERS;
+    MxModelZoo() {
+        addModel(REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "ssd", "0.0.1"));
+        addModel(REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "yolo", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "alexnet", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "darknet", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "densenet", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "googlenet", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "inceptionv3", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "mlp", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "mobilenet", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "resnest", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "resnet", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "senet", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "se_resnext", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "squeezenet", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "vgg", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "xception", "0.0.1"));
+        addModel(REPOSITORY.model(CV.POSE_ESTIMATION, GROUP_ID, "simple_pose", "0.0.1"));
+        addModel(REPOSITORY.model(CV.INSTANCE_SEGMENTATION, GROUP_ID, "mask_rcnn", "0.0.1"));
+        addModel(REPOSITORY.model(CV.ACTION_RECOGNITION, GROUP_ID, "action_recognition", "0.0.1"));
+        addModel(REPOSITORY.model(NLP.QUESTION_ANSWER, GROUP_ID, "bertqa", "0.0.1"));
+        addModel(REPOSITORY.model(NLP.WORD_EMBEDDING, GROUP_ID, "glove", "0.0.2"));
     }
 
     /** {@inheritDoc} */
diff --git a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/zoo/OrtModelZoo.java b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/zoo/OrtModelZoo.java
index 0421a5f96..d4d27b9ed 100644
--- a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/zoo/OrtModelZoo.java
+++ b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/zoo/OrtModelZoo.java
@@ -15,15 +15,10 @@ package ai.djl.onnxruntime.zoo;
 import ai.djl.Application.CV;
 import ai.djl.Application.Tabular;
 import ai.djl.onnxruntime.engine.OrtEngine;
-import ai.djl.repository.MRL;
 import ai.djl.repository.Repository;
-import ai.djl.repository.zoo.BaseModelLoader;
-import ai.djl.repository.zoo.ModelLoader;
 import ai.djl.repository.zoo.ModelZoo;
 
-import java.util.ArrayList;
 import java.util.Collections;
-import java.util.List;
 import java.util.Set;
 
 /** OrtModelZoo is a repository that contains all Onnx models for DJL. */
@@ -33,21 +28,9 @@ public class OrtModelZoo extends ModelZoo {
     private static final Repository REPOSITORY = Repository.newInstance("Ort", DJL_REPO_URL);
     public static final String GROUP_ID = "ai.djl.onnxruntime";
 
-    private static final List<ModelLoader> MODEL_LOADERS = new ArrayList<>();
-
-    static {
-        MRL resnet = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "resnet", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(resnet));
-
-        MRL irisFlower =
-                REPOSITORY.model(Tabular.SOFTMAX_REGRESSION, GROUP_ID, "iris_flowers", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(irisFlower));
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public List<ModelLoader> getModelLoaders() {
-        return MODEL_LOADERS;
+    OrtModelZoo() {
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "resnet", "0.0.1"));
+        addModel(REPOSITORY.model(Tabular.SOFTMAX_REGRESSION, GROUP_ID, "iris_flowers", "0.0.1"));
     }
 
     /** {@inheritDoc} */
diff --git a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/PpModelZoo.java b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/PpModelZoo.java
index 22af1f070..55d392b57 100644
--- a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/PpModelZoo.java
+++ b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/PpModelZoo.java
@@ -14,15 +14,10 @@ package ai.djl.paddlepaddle.zoo;
 
 import ai.djl.Application.CV;
 import ai.djl.paddlepaddle.engine.PpEngine;
-import ai.djl.repository.MRL;
 import ai.djl.repository.Repository;
-import ai.djl.repository.zoo.BaseModelLoader;
-import ai.djl.repository.zoo.ModelLoader;
 import ai.djl.repository.zoo.ModelZoo;
 
-import java.util.ArrayList;
 import java.util.Collections;
-import java.util.List;
 import java.util.Set;
 
 /** PpModelZoo is a repository that contains all PaddlePaddle models for DJL. */
@@ -32,34 +27,14 @@ public class PpModelZoo extends ModelZoo {
     private static final Repository REPOSITORY = Repository.newInstance("Paddle", DJL_REPO_URL);
     public static final String GROUP_ID = "ai.djl.paddlepaddle";
 
-    private static final List<ModelLoader> MODEL_LOADERS = new ArrayList<>();
-
-    static {
-        MRL maskDetection =
-                REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "mask_classification", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(maskDetection));
-
-        MRL wordRotation =
-                REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "word_rotation", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(wordRotation));
-
-        MRL faceDetection =
-                REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "face_detection", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(faceDetection));
-
-        MRL wordDetection =
-                REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "word_detection", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(wordDetection));
-
-        MRL wordRecognition =
-                REPOSITORY.model(CV.WORD_RECOGNITION, GROUP_ID, "word_recognition", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(wordRecognition));
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public List<ModelLoader> getModelLoaders() {
-        return MODEL_LOADERS;
+    PpModelZoo() {
+        addModel(
+                REPOSITORY.model(
+                        CV.IMAGE_CLASSIFICATION, GROUP_ID, "mask_classification", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "word_rotation", "0.0.1"));
+        addModel(REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "face_detection", "0.0.1"));
+        addModel(REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "word_detection", "0.0.1"));
+        addModel(REPOSITORY.model(CV.WORD_RECOGNITION, GROUP_ID, "word_recognition", "0.0.1"));
     }
 
     /** {@inheritDoc} */
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtDeviceType.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtDeviceType.java
index c898ddfd8..78b620345 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtDeviceType.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtDeviceType.java
@@ -13,10 +13,9 @@
 package ai.djl.pytorch.engine;
 
 import ai.djl.Device;
-import ai.djl.DeviceType;
 
 /** DeviceType is the PyTorch equivalent of the types in {@link Device}. */
-public final class PtDeviceType implements DeviceType {
+public final class PtDeviceType {
 
     private PtDeviceType() {}
 
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
index d61ecad3f..2829fee3f 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
@@ -268,11 +268,11 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
 
     /** {@inheritDoc} */
     @Override
-    public NDArray take(NDArray index) {
+    public NDArray take(NDManager manager, NDArray index) {
         if (!(index instanceof PtNDArray)) {
             throw new IllegalArgumentException("Only PtNDArray is supported.");
         }
-        return JniUtils.take(this, (PtNDArray) index);
+        return JniUtils.take(this, (PtNDArray) index, (PtNDManager) manager);
     }
 
     /** {@inheritDoc} */
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayIndexer.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayIndexer.java
index fcc8958d0..b38a038d1 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayIndexer.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayIndexer.java
@@ -18,6 +18,7 @@ import ai.djl.ndarray.index.NDIndex;
 import ai.djl.ndarray.index.dim.NDIndexBooleans;
 import ai.djl.ndarray.index.full.NDIndexFullPick;
 import ai.djl.ndarray.index.full.NDIndexFullSlice;
+import ai.djl.ndarray.index.full.NDIndexFullTake;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.pytorch.jni.JniUtils;
 
@@ -39,6 +40,12 @@ public class PtNDArrayIndexer extends NDArrayIndexer {
                 manager.from(array), manager.from(fullPick.getIndices()), fullPick.getAxis());
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray get(NDArray array, NDIndexFullTake fullTake) {
+        return JniUtils.take(manager.from(array), manager.from(fullTake.getIndices()), manager);
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray get(NDArray array, NDIndexFullSlice fullSlice) {
@@ -63,12 +70,12 @@ public class PtNDArrayIndexer extends NDArrayIndexer {
             index.addAllDim();
         }
 
-        if (array == null || array instanceof PtNDArray && array.getManager() == manager) {
-            return JniUtils.indexAdv((PtNDArray) array, index);
+        if (array == null || array instanceof PtNDArray) {
+            return JniUtils.indexAdv((PtNDArray) array, index, manager);
         } else {
             PtNDArray arrayNew =
                     manager.create(array.toByteBuffer(), array.getShape(), array.getDataType());
-            return JniUtils.indexAdv(arrayNew, index);
+            return JniUtils.indexAdv(arrayNew, index, manager);
         }
     }
 
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
index 61e21a0c8..5ae9e0494 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
@@ -358,7 +358,7 @@ public final class JniUtils {
     }
 
     @SuppressWarnings("OptionalGetWithoutIsPresent")
-    public static PtNDArray indexAdv(PtNDArray ndArray, NDIndex index) {
+    public static PtNDArray indexAdv(PtNDArray ndArray, NDIndex index, PtNDManager manager) {
         if (ndArray == null) {
             return ndArray;
         }
@@ -405,10 +405,7 @@ public final class JniUtils {
                 // Backward compatible
                 NDIndexFullPick fullPick =
                         NDIndexFullPick.fromIndex(index, ndArray.getShape()).get();
-                return pick(
-                        ndArray,
-                        ndArray.getManager().from(fullPick.getIndices()),
-                        fullPick.getAxis());
+                return pick(ndArray, manager.from(fullPick.getIndices()), fullPick.getAxis());
             }
         }
         if (indices.size() == index.getEllipsisIndex()) {
@@ -416,7 +413,7 @@ public final class JniUtils {
         }
 
         return new PtNDArray(
-                ndArray.getManager(),
+                manager,
                 PyTorchLibrary.LIB.torchIndexAdvGet(ndArray.getHandle(), torchIndexHandle));
     }
 
@@ -506,13 +503,12 @@ public final class JniUtils {
                 PyTorchLibrary.LIB.torchGather(ndArray.getHandle(), index.getHandle(), dim, false));
     }
 
-    public static PtNDArray take(PtNDArray ndArray, PtNDArray index) {
+    public static PtNDArray take(PtNDArray ndArray, PtNDArray index, PtNDManager manager) {
         if (index.getDataType() != DataType.INT64) {
             index = index.toType(DataType.INT64, true);
         }
         return new PtNDArray(
-                ndArray.getManager(),
-                PyTorchLibrary.LIB.torchTake(ndArray.getHandle(), index.getHandle()));
+                manager, PyTorchLibrary.LIB.torchTake(ndArray.getHandle(), index.getHandle()));
     }
 
     public static PtNDArray put(PtNDArray ndArray, PtNDArray index, PtNDArray data) {
diff --git a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/PtModelZoo.java b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/PtModelZoo.java
index e2cb7f790..c53aeedd5 100644
--- a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/PtModelZoo.java
+++ b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/PtModelZoo.java
@@ -15,15 +15,10 @@ package ai.djl.pytorch.zoo;
 import ai.djl.Application.CV;
 import ai.djl.Application.NLP;
 import ai.djl.pytorch.engine.PtEngine;
-import ai.djl.repository.MRL;
 import ai.djl.repository.Repository;
-import ai.djl.repository.zoo.BaseModelLoader;
-import ai.djl.repository.zoo.ModelLoader;
 import ai.djl.repository.zoo.ModelZoo;
 
-import java.util.ArrayList;
 import java.util.Collections;
-import java.util.List;
 import java.util.Set;
 
 /**
@@ -36,36 +31,14 @@ public class PtModelZoo extends ModelZoo {
     private static final Repository REPOSITORY = Repository.newInstance("PyTorch", DJL_REPO_URL);
     public static final String GROUP_ID = "ai.djl.pytorch";
 
-    private static final List<ModelLoader> MODEL_LOADERS = new ArrayList<>();
-
-    static {
-        MRL resnet = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "resnet", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(resnet));
-
-        MRL ssd = REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "ssd", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(ssd));
-
-        MRL bertQa = REPOSITORY.model(NLP.QUESTION_ANSWER, GROUP_ID, "bertqa", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(bertQa));
-
-        MRL sentimentAnalysis =
-                REPOSITORY.model(NLP.SENTIMENT_ANALYSIS, GROUP_ID, "distilbert", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(sentimentAnalysis));
-
-        MRL bigGan = REPOSITORY.model(CV.IMAGE_GENERATION, GROUP_ID, "biggan-deep", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(bigGan));
-
-        MRL cyclegan = REPOSITORY.model(CV.IMAGE_GENERATION, GROUP_ID, "cyclegan", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(cyclegan));
-
-        MRL deeplabv3 = REPOSITORY.model(CV.SEMANTIC_SEGMENTATION, GROUP_ID, "deeplabv3", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(deeplabv3));
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public List<ModelLoader> getModelLoaders() {
-        return MODEL_LOADERS;
+    PtModelZoo() {
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "resnet", "0.0.1"));
+        addModel(REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "ssd", "0.0.1"));
+        addModel(REPOSITORY.model(NLP.QUESTION_ANSWER, GROUP_ID, "bertqa", "0.0.1"));
+        addModel(REPOSITORY.model(NLP.SENTIMENT_ANALYSIS, GROUP_ID, "distilbert", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_GENERATION, GROUP_ID, "biggan-deep", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_GENERATION, GROUP_ID, "cyclegan", "0.0.1"));
+        addModel(REPOSITORY.model(CV.SEMANTIC_SEGMENTATION, GROUP_ID, "deeplabv3", "0.0.1"));
     }
 
     /** {@inheritDoc} */
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
index 89814cc86..915438789 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
@@ -210,7 +210,7 @@ public class TfNDArray extends NativeResource<TFE_TensorHandle> implements NDArr
 
     /** {@inheritDoc} */
     @Override
-    public NDArray take(NDArray index) {
+    public NDArray take(NDManager manager, NDArray index) {
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArrayIndexer.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArrayIndexer.java
index 73b3a4c5d..80f5e533c 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArrayIndexer.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArrayIndexer.java
@@ -16,6 +16,7 @@ import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.index.NDArrayIndexer;
 import ai.djl.ndarray.index.full.NDIndexFullPick;
 import ai.djl.ndarray.index.full.NDIndexFullSlice;
+import ai.djl.ndarray.index.full.NDIndexFullTake;
 
 /** The {@link NDArrayIndexer} used by the {@link TfNDArray}. */
 public class TfNDArrayIndexer extends NDArrayIndexer {
@@ -32,6 +33,12 @@ public class TfNDArrayIndexer extends NDArrayIndexer {
         throw new UnsupportedOperationException("Not implemented");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray get(NDArray array, NDIndexFullTake fullTake) {
+        throw new UnsupportedOperationException("Not implemented");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray get(NDArray array, NDIndexFullSlice fullSlice) {
diff --git a/engines/tensorflow/tensorflow-model-zoo/src/main/java/ai/djl/tensorflow/zoo/TfModelZoo.java b/engines/tensorflow/tensorflow-model-zoo/src/main/java/ai/djl/tensorflow/zoo/TfModelZoo.java
index ac27a9e3b..4a54be6c0 100644
--- a/engines/tensorflow/tensorflow-model-zoo/src/main/java/ai/djl/tensorflow/zoo/TfModelZoo.java
+++ b/engines/tensorflow/tensorflow-model-zoo/src/main/java/ai/djl/tensorflow/zoo/TfModelZoo.java
@@ -14,16 +14,11 @@
 package ai.djl.tensorflow.zoo;
 
 import ai.djl.Application.CV;
-import ai.djl.repository.MRL;
 import ai.djl.repository.Repository;
-import ai.djl.repository.zoo.BaseModelLoader;
-import ai.djl.repository.zoo.ModelLoader;
 import ai.djl.repository.zoo.ModelZoo;
 import ai.djl.tensorflow.engine.TfEngine;
 
-import java.util.ArrayList;
 import java.util.Collections;
-import java.util.List;
 import java.util.Set;
 
 /** TfModelZoo is a repository that contains the TensorFlow models for DJL. */
@@ -33,23 +28,10 @@ public class TfModelZoo extends ModelZoo {
     private static final Repository REPOSITORY = Repository.newInstance("TensorFlow", DJL_REPO_URL);
     public static final String GROUP_ID = "ai.djl.tensorflow";
 
-    private static final List<ModelLoader> MODEL_LOADERS = new ArrayList<>();
-
-    static {
-        MRL resnet = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "resnet", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(resnet));
-
-        MRL mobilenet = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "mobilenet", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(mobilenet));
-
-        MRL ssd = REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "ssd", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(ssd));
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public List<ModelLoader> getModelLoaders() {
-        return MODEL_LOADERS;
+    TfModelZoo() {
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "resnet", "0.0.1"));
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "mobilenet", "0.0.1"));
+        addModel(REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "ssd", "0.0.1"));
     }
 
     /** {@inheritDoc} */
diff --git a/engines/tflite/tflite-engine/src/main/java/ai/djl/tflite/zoo/TfLiteModelZoo.java b/engines/tflite/tflite-engine/src/main/java/ai/djl/tflite/zoo/TfLiteModelZoo.java
index 6c1ec7875..48614b915 100644
--- a/engines/tflite/tflite-engine/src/main/java/ai/djl/tflite/zoo/TfLiteModelZoo.java
+++ b/engines/tflite/tflite-engine/src/main/java/ai/djl/tflite/zoo/TfLiteModelZoo.java
@@ -13,16 +13,11 @@
 package ai.djl.tflite.zoo;
 
 import ai.djl.Application.CV;
-import ai.djl.repository.MRL;
 import ai.djl.repository.Repository;
-import ai.djl.repository.zoo.BaseModelLoader;
-import ai.djl.repository.zoo.ModelLoader;
 import ai.djl.repository.zoo.ModelZoo;
 import ai.djl.tflite.engine.TfLiteEngine;
 
-import java.util.ArrayList;
 import java.util.Collections;
-import java.util.List;
 import java.util.Set;
 
 /** TfLiteModelZoo is a repository that contains all TFLite models for DJL. */
@@ -32,17 +27,8 @@ public class TfLiteModelZoo extends ModelZoo {
     private static final Repository REPOSITORY = Repository.newInstance("TFLite", DJL_REPO_URL);
     public static final String GROUP_ID = "ai.djl.tflite";
 
-    private static final List<ModelLoader> MODEL_LOADERS = new ArrayList<>();
-
-    static {
-        MRL mobilenet = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "mobilenet", "0.0.1");
-        MODEL_LOADERS.add(new BaseModelLoader(mobilenet));
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public List<ModelLoader> getModelLoaders() {
-        return MODEL_LOADERS;
+    TfLiteModelZoo() {
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "mobilenet", "0.0.1"));
     }
 
     /** {@inheritDoc} */
diff --git a/examples/docs/speech_recognition.md b/examples/docs/speech_recognition.md
new file mode 100644
index 000000000..0018455dd
--- /dev/null
+++ b/examples/docs/speech_recognition.md
@@ -0,0 +1,28 @@
+# Speech Recognition Example
+
+In this example, you learn how to use Speech Recognition using PyTorch.
+You can provide the model with a wav input file. 
+You can find the source code in [SpeechRecognition.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/SpeechRecognition.java).
+
+Example:
+
+Input audio: https://resources.djl.ai/audios/speech.wav
+
+Result:
+
+```text
+THE NEAREST SAID THE DISTRICT DOCTOR IS A GOOD ITALIAN ABBE WHO LIVES NEXT DOOR TO YOU SHALL I CALL ON HIM AS I PASS 
+```
+
+## Setup guide
+
+Follow [setup](../../docs/development/setup.md) to configure your development environment.
+
+### Run Inference
+
+```
+cd examples
+./gradlew run -Dmain=ai.djl.examples.inference.SpeechRecognition
+
+[INFO ] - Result: THE NEAREST SAID THE DISTRICT DOCTOR IS A GOOD ITALIAN ABBE WHO LIVES NEXT DOOR TO YOU SHALL I CALL ON HIM AS I PASS 
+```
diff --git a/examples/src/main/java/ai/djl/examples/inference/SpeechRecognition.java b/examples/src/main/java/ai/djl/examples/inference/SpeechRecognition.java
new file mode 100644
index 000000000..1561bab1b
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/inference/SpeechRecognition.java
@@ -0,0 +1,73 @@
+/*
+ * 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.examples.inference;
+
+import ai.djl.ModelException;
+import ai.djl.inference.Predictor;
+import ai.djl.modality.audio.Audio;
+import ai.djl.modality.audio.AudioFactory;
+import ai.djl.modality.audio.translator.SpeechRecognitionTranslatorFactory;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.ZooModel;
+import ai.djl.translate.TranslateException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+
+import javax.sound.sampled.UnsupportedAudioFileException;
+
+/**
+ * An example of inference using a speech recognition model.
+ *
+ * <p>See this <a
+ * href="https://github.com/deepjavalibrary/djl/blob/master/examples/docs/speech_recognition.md">doc</a>
+ * for information about this example.
+ */
+public final class SpeechRecognition {
+
+    public static final Logger logger = LoggerFactory.getLogger((SpeechRecognition.class));
+
+    private SpeechRecognition() {}
+
+    public static void main(String[] args)
+            throws UnsupportedAudioFileException, IOException, TranslateException, ModelException {
+        logger.info("Result: {}", predict());
+    }
+
+    public static String predict()
+            throws UnsupportedAudioFileException, 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.
+        String url = "https://resources.djl.ai/test-models/pytorch/wav2vec2.zip";
+        Criteria<Audio, String> criteria =
+                Criteria.builder()
+                        .setTypes(Audio.class, String.class)
+                        .optModelUrls(url)
+                        .optTranslatorFactory(new SpeechRecognitionTranslatorFactory())
+                        .optModelName("wav2vec2.ptl")
+                        .optEngine("PyTorch")
+                        .build();
+
+        // Read in audio file
+        String wave = "https://resources.djl.ai/audios/speech.wav";
+        Audio audio = AudioFactory.getInstance().fromUrl(wave);
+        try (ZooModel<Audio, String> model = criteria.loadModel();
+                Predictor<Audio, String> predictor = model.newPredictor()) {
+            return predictor.predict(audio);
+        }
+    }
+}
diff --git a/extensions/audio/src/main/java/ai/djl/audio/FFmpegAudioFactory.java b/extensions/audio/src/main/java/ai/djl/audio/FFmpegAudioFactory.java
new file mode 100644
index 000000000..ad6ffa071
--- /dev/null
+++ b/extensions/audio/src/main/java/ai/djl/audio/FFmpegAudioFactory.java
@@ -0,0 +1,94 @@
+/*
+ * 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.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;
+import org.bytedeco.javacv.FrameGrabber;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.ShortBuffer;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * {@code FFmpegAudioFactory} is a high performance implementation of {@link AudioFactory} using
+ * FFmpeg.
+ */
+public class FFmpegAudioFactory extends AudioFactory {
+
+    /** {@inheritDoc} */
+    @Override
+    public Audio fromFile(Path path) throws IOException {
+        try (FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(path.toFile())) {
+            grabber.start();
+            float[] floats = grab(grabber);
+            return new Audio(floats, grabber.getSampleRate(), grabber.getAudioChannels());
+        } catch (FrameGrabber.Exception e) {
+            throw new IOException("Unsupported Audio file", e);
+        }
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Audio fromInputStream(InputStream is) throws IOException {
+        try (FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(is)) {
+            grabber.start();
+            float[] floats = grab(grabber);
+            return new Audio(floats, grabber.getSampleRate(), grabber.getAudioChannels());
+        } catch (FrameGrabber.Exception 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!");
+    }
+
+    /**
+     * Grab frames from the audio using {@link FFmpegFrameGrabber}.
+     *
+     * @param grabber the {@link FFmpegFrameGrabber}.
+     * @return the float array read from the audio.
+     * @throws FFmpegFrameGrabber.Exception if error occurs
+     */
+    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);
+            }
+        }
+        float[] ret = new float[list.size()];
+        for (int i = 0; i < list.size(); i++) {
+            ret[i] = list.get(i);
+        }
+        return ret;
+    }
+}
diff --git a/extensions/audio/src/main/java/ai/djl/audio/dataset/AudioData.java b/extensions/audio/src/main/java/ai/djl/audio/dataset/AudioData.java
index dcd0d2cd5..60f580c54 100644
--- a/extensions/audio/src/main/java/ai/djl/audio/dataset/AudioData.java
+++ b/extensions/audio/src/main/java/ai/djl/audio/dataset/AudioData.java
@@ -18,7 +18,6 @@ import ai.djl.audio.processor.AudioProcessor;
 import ai.djl.audio.processor.LinearSpecgram;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDManager;
-import ai.djl.repository.LocalRepository;
 
 import org.bytedeco.javacv.FFmpegFrameGrabber;
 import org.bytedeco.javacv.Frame;
@@ -44,7 +43,7 @@ import java.util.List;
  */
 public class AudioData {
 
-    private static final Logger logger = LoggerFactory.getLogger(LocalRepository.class);
+    private static final Logger logger = LoggerFactory.getLogger(AudioData.class);
 
     private int sampleRate;
     private int audioChannels;
@@ -63,13 +62,11 @@ public class AudioData {
     }
 
     /**
-     * Returns a good default {@link AudioData.Configuration} to use for the constructor with
-     * defaults.
+     * Returns a good default {@link Configuration} to use for the constructor with defaults.
      *
-     * @return a good default {@link AudioData.Configuration} to use for the constructor with
-     *     defaults.
+     * @return a good default {@link Configuration} to use for the constructor with defaults.
      */
-    public static AudioData.Configuration getDefaultConfiguration() {
+    public static Configuration getDefaultConfiguration() {
         float targetDb = -20f;
         float strideMs = 10f;
         float windowsMs = 20f;
@@ -78,9 +75,7 @@ public class AudioData {
                 Arrays.asList(
                         new AudioNormalizer(targetDb),
                         new LinearSpecgram(strideMs, windowsMs, sampleRate));
-        return new AudioData.Configuration()
-                .setProcessorList(defaultProcessors)
-                .setSampleRate(sampleRate);
+        return new Configuration().setProcessorList(defaultProcessors).setSampleRate(sampleRate);
     }
 
     /**
@@ -177,10 +172,7 @@ public class AudioData {
         return audioPaths.size();
     }
 
-    /**
-     * The configuration for creating a {@link AudioData} value in a {@link
-     * ai.djl.audio.dataset.AudioData}.
-     */
+    /** The configuration for creating a {@code AudioData} value in a {@code AudioData}. */
     public static final class Configuration {
 
         private int sampleRate;
@@ -211,13 +203,12 @@ public class AudioData {
         }
 
         /**
-         * Updates this {@link AudioData.Configuration} with the non-null values from another
-         * configuration.
+         * Updates this {@link Configuration} with the non-null values from another configuration.
          *
          * @param other the other configuration to use to update this
          * @return this configuration after updating
          */
-        public AudioData.Configuration update(AudioData.Configuration other) {
+        public Configuration update(AudioData.Configuration other) {
             processorList = other.processorList;
             return this;
         }
diff --git a/extensions/audio/src/main/java/ai/djl/audio/processor/AudioNormalizer.java b/extensions/audio/src/main/java/ai/djl/audio/processor/AudioNormalizer.java
index 55c4de4ed..5d76a5115 100644
--- a/extensions/audio/src/main/java/ai/djl/audio/processor/AudioNormalizer.java
+++ b/extensions/audio/src/main/java/ai/djl/audio/processor/AudioNormalizer.java
@@ -12,7 +12,7 @@
  */
 package ai.djl.audio.processor;
 
-import ai.djl.audio.AudioUtils;
+import ai.djl.audio.util.AudioUtils;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDManager;
 
diff --git a/extensions/audio/src/main/java/ai/djl/audio/AudioUtils.java b/extensions/audio/src/main/java/ai/djl/audio/util/AudioUtils.java
similarity index 97%
rename from extensions/audio/src/main/java/ai/djl/audio/AudioUtils.java
rename to extensions/audio/src/main/java/ai/djl/audio/util/AudioUtils.java
index e2a2bcb7d..2e0603023 100644
--- a/extensions/audio/src/main/java/ai/djl/audio/AudioUtils.java
+++ b/extensions/audio/src/main/java/ai/djl/audio/util/AudioUtils.java
@@ -10,7 +10,7 @@
  * 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;
+package ai.djl.audio.util;
 
 import ai.djl.ndarray.NDArray;
 
diff --git a/extensions/audio/src/main/java/ai/djl/audio/util/package-info.java b/extensions/audio/src/main/java/ai/djl/audio/util/package-info.java
new file mode 100644
index 000000000..ec145a739
--- /dev/null
+++ b/extensions/audio/src/main/java/ai/djl/audio/util/package-info.java
@@ -0,0 +1,15 @@
+/*
+ * 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.
+ */
+
+/** Contains audio utils for djl. */
+package ai.djl.audio.util;
diff --git a/extensions/fasttext/src/main/java/ai/djl/fasttext/zoo/FtModelZoo.java b/extensions/fasttext/src/main/java/ai/djl/fasttext/zoo/FtModelZoo.java
index a44ba6700..29e6077d3 100644
--- a/extensions/fasttext/src/main/java/ai/djl/fasttext/zoo/FtModelZoo.java
+++ b/extensions/fasttext/src/main/java/ai/djl/fasttext/zoo/FtModelZoo.java
@@ -15,12 +15,9 @@ package ai.djl.fasttext.zoo;
 import ai.djl.engine.Engine;
 import ai.djl.fasttext.zoo.nlp.textclassification.TextClassificationModelLoader;
 import ai.djl.repository.Repository;
-import ai.djl.repository.zoo.ModelLoader;
 import ai.djl.repository.zoo.ModelZoo;
 
-import java.util.ArrayList;
 import java.util.Collections;
-import java.util.List;
 import java.util.Set;
 
 /** FtModelZoo is a repository that contains all fastText models for DJL. */
@@ -30,16 +27,8 @@ public class FtModelZoo extends ModelZoo {
     private static final Repository REPOSITORY = Repository.newInstance("Fasttext", DJL_REPO_URL);
     public static final String GROUP_ID = "ai.djl.fasttext";
 
-    private static final List<ModelLoader> MODEL_LOADERS = new ArrayList<>();
-
-    static {
-        MODEL_LOADERS.add(new TextClassificationModelLoader(REPOSITORY));
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public List<ModelLoader> getModelLoaders() {
-        return MODEL_LOADERS;
+    FtModelZoo() {
+        addModel(new TextClassificationModelLoader(REPOSITORY));
     }
 
     /** {@inheritDoc} */
diff --git a/extensions/tokenizers/README.md b/extensions/tokenizers/README.md
index 38cc8529a..d10ea87a0 100644
--- a/extensions/tokenizers/README.md
+++ b/extensions/tokenizers/README.md
@@ -24,3 +24,83 @@ You can pull the module from the central Maven repository by including the follo
     <version>0.18.0</version>
 </dependency>
 ```
+
+## Usage
+
+### Use DJL HuggingFace model converter (experimental)
+
+If you are trying to convert a complete HuggingFace (transformers) model,
+you can try to use our all-in-one conversion solution to convert to Java:
+
+```bash
+python3 -m pip install -r src/main/python/requirements.txt
+python3 src/main/python/model_zoo_importer.py -m deepset/bert-base-cased-squad2
+```
+
+This will generate a zip file into your local folder:
+
+```
+model/nlp/question_answer/ai/djl/huggingface/pytorch/deepset/bert-base-cased-squad2/0.0.1/bert-base-cased-squad2.zip
+```
+
+Then, all you need to do, is to load this model in DJL:
+
+```java
+Criteria<QAInput, String> criteria = Criteria.builder()
+    .setTypes(QAInput.class, String.class)
+    .optModelPath(Paths.get("model/nlp/question_answer/ai/djl/huggingface/pytorch/deepset/bert-base-cased-squad2/0.0.1/bert-base-cased-squad2.zip"))
+    .optProgress(new ProgressBar()).build();
+```
+
+Currently, this converter support:
+
+- fill-mask
+- question-answering
+- sentence-similarity
+- text-classification
+- token-classification
+
+### From HuggingFace AutoTokenizer
+In most of the cases, you can easily use a pre-existing tokenizer in DJL:
+
+Python
+
+```python
+from transformers import AutoTokenizer
+tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/msmarco-distilbert-dot-v5")
+```
+
+Java
+
+```java
+HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.newInstance("sentence-transformers/msmarco-distilbert-dot-v5")
+```
+
+This way requires network connection to huggingface repo.
+The way to determine if you can use this way is through looking into the "Files and versions" in [HuggingFace model tab](https://huggingface.co/sentence-transformers/msmarco-distilbert-dot-v5)
+and see if there is a `tokenizer.json`. 
+
+If there is a `tokenizer.json`, you can get it directly through DJL. Otherwise, use the other way below to obtain a `tokenizer.json`.
+
+### From HuggingFace Pipeline
+
+If you are trying to get tokenizer from a HuggingFace pipeline,
+you can use the followings to extract `tokenizer.json` file.
+
+Python
+
+```python
+ pipeline.tokenizer.save_pretrained("./")
+```
+
+From your local directory, you will find a `tokenizer.json` file.
+
+Java
+
+```java
+HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.newInstance(Paths.get("./tokenizer.json"))
+```
+
+### From pretrained json file
+
+Same as above step, just save your tokenizer into `tokenizer.json` (done by huggingface).
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslator.java
index e0fe288ab..5e2db2a3a 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslator.java
@@ -73,8 +73,8 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
         NDArray endLogits = list.get(1).duplicate();
 
         // exclude <CLS>, TODO: exclude impossible ids properly and handle max answer length
-        startLogits.set(new NDIndex(0), 0);
-        endLogits.set(new NDIndex(0), 0);
+        startLogits.set(new NDIndex(0), -100000);
+        endLogits.set(new NDIndex(0), -100000);
         int startIdx = (int) startLogits.argMax().getLong();
         int endIdx = (int) endLogits.argMax().getLong();
         if (startIdx > endIdx) {
diff --git a/extensions/tokenizers/src/main/python/fill_mask_converter.py b/extensions/tokenizers/src/main/python/fill_mask_converter.py
index 7b06552ff..6d9e10deb 100644
--- a/extensions/tokenizers/src/main/python/fill_mask_converter.py
+++ b/extensions/tokenizers/src/main/python/fill_mask_converter.py
@@ -47,7 +47,9 @@ class FillMaskConverter(HuggingfaceConverter):
             pipeline_output = hf_pipeline(text)
 
             if prediction not in [o["token_str"] for o in pipeline_output]:
-                return False, f"Unexpected inference result: {prediction}"
+                logging.error(f"Unexpected inference result: {prediction}")
+                return False, "Unexpected inference result"
+
             logging.warning(
                 f"pipeline output differs from expected: {pipeline_output}")
 
diff --git a/extensions/tokenizers/src/main/python/huggingface_converter.py b/extensions/tokenizers/src/main/python/huggingface_converter.py
index 747bba7cf..1eb521b40 100644
--- a/extensions/tokenizers/src/main/python/huggingface_converter.py
+++ b/extensions/tokenizers/src/main/python/huggingface_converter.py
@@ -49,16 +49,28 @@ class HuggingfaceConverter:
         shutil.copyfile(config, os.path.join(temp_dir, "config.json"))
 
         # Save jit traced .pt file to temp dir
-        model_file = self.jit_trace_model(hf_pipeline, model_id, temp_dir)
+        include_types = False
+        model_file = self.jit_trace_model(hf_pipeline, model_id, temp_dir,
+                                          include_types)
         if not model_file:
             return False, "Failed to trace model", -1
 
-        result, reason = self.verify_jit_model(hf_pipeline, model_file)
+        result, reason = self.verify_jit_model(hf_pipeline, model_file,
+                                               include_types)
         if not result:
-            return False, reason, -1
+            include_types = True
+            model_file = self.jit_trace_model(hf_pipeline, model_id, temp_dir,
+                                              include_types)
+            if not model_file:
+                return False, reason, -1
+
+            result, reason = self.verify_jit_model(hf_pipeline, model_file,
+                                                   include_types)
+            if not result:
+                return False, reason, -1
 
         size = self.save_to_model_zoo(model_info, output_dir, temp_dir,
-                                      hf_pipeline)
+                                      hf_pipeline, include_types)
 
         return True, None, size
 
@@ -70,43 +82,47 @@ class HuggingfaceConverter:
             if path != "tokenizer.json":
                 os.remove(os.path.join(temp_dir, path))
 
-    def jit_trace_model(self, hf_pipeline, model_id: str, temp_dir: str):
-        logging.info(f"Tracing model: {model_id} ...")
+    def jit_trace_model(self, hf_pipeline, model_id: str, temp_dir: str,
+                        include_types: bool):
+        logging.info(
+            f"Tracing model: {model_id} include_token_types={include_types} ..."
+        )
         encoding = self.encode_inputs(hf_pipeline.tokenizer)
         input_ids = encoding["input_ids"]
         attention_mask = encoding["attention_mask"]
+        token_type_ids = encoding.get("token_type_ids")
+        if include_types and token_type_ids is None:
+            return None
 
         # noinspection PyBroadException
         try:
-            script_module = torch.jit.trace(hf_pipeline.model,
-                                            (input_ids, attention_mask),
-                                            strict=False)
-        except Exception as ex:
-            logging.info(
-                f"Trace with attention_mask failed, model: {model_id}.")
-            logging.debug(ex, exc_info=False)
-            try:
+            if include_types:
+                script_module = torch.jit.trace(
+                    hf_pipeline.model,
+                    (input_ids, attention_mask, token_type_ids),
+                    strict=False)
+            else:
                 script_module = torch.jit.trace(hf_pipeline.model,
-                                                input_ids,
+                                                (input_ids, attention_mask),
                                                 strict=False)
-            except Exception as e:
-                logging.warning(f"Failed to trace model: {model_id}.")
-                logging.warning(e, exc_info=True)
-                return None
 
-        model_name = model_id.split("/")[-1]
-        logging.info(f"Saving torchscript model: {model_name}.pt ...")
-        model_file = os.path.join(temp_dir, f"{model_name}.pt")
-        script_module.save(model_file)
+            model_name = model_id.split("/")[-1]
+            logging.info(f"Saving torchscript model: {model_name}.pt ...")
+            model_file = os.path.join(temp_dir, f"{model_name}.pt")
+            script_module.save(model_file)
+        except RuntimeError as e:
+            logging.warning(f"Failed to trace model: {model_id}.")
+            logging.warning(e, exc_info=True)
+            return None
+
         return model_file
 
     def save_to_model_zoo(self, model_info, output_dir: str, temp_dir: str,
-                          hf_pipeline):
+                          hf_pipeline, include_types: bool):
         model_id = model_info.modelId
-        artifact_ids = model_id.split("/")
-        model_name = artifact_ids[-1]
+        model_name = model_id.split("/")[-1]
 
-        repo_dir = f"{output_dir}/model/{self.application}/ai/djl/huggingface/{model_id}"
+        repo_dir = f"{output_dir}/model/{self.application}/ai/djl/huggingface/pytorch/{model_id}"
         model_dir = f"{repo_dir}/0.0.1"
         if not os.path.exists(model_dir):
             os.makedirs(model_dir)
@@ -119,6 +135,9 @@ class HuggingfaceConverter:
                     f"option.modelName={model_name}\n"
                     f"option.mapLocation=true\n"
                     f"translatorFactory={self.translator}\n")
+            if include_types:
+                f.write(f"includeTokenTypes={include_types}\n")
+
             for k, v in arguments.items():
                 f.write(f"{k}={v}\n")
 
@@ -130,34 +149,46 @@ class HuggingfaceConverter:
         # Save metadata.json
         sha1 = sha1_sum(zip_file)
         file_size = os.path.getsize(zip_file)
-        metadata = HuggingfaceMetadata(model_info, artifact_ids,
-                                       self.application, self.translator, sha1,
-                                       file_size)
+        metadata = HuggingfaceMetadata(model_info, self.application,
+                                       self.translator, sha1, file_size)
         metadata_file = os.path.join(repo_dir, "metadata.json")
         metadata.save_metadata(metadata_file)
 
         return file_size
 
-    def verify_jit_model(self, hf_pipeline, model_file: str):
-        logging.info(f"Verifying torchscript model: {model_file} ...")
+    def verify_jit_model(self, hf_pipeline, model_file: str,
+                         include_types: bool):
+        logging.info(
+            f"Verifying torchscript model(include_token_types={include_types}): {model_file} ..."
+        )
 
         tokenizer = hf_pipeline.tokenizer
         encoding = self.encode_inputs(tokenizer)
 
         input_ids = encoding["input_ids"]
         attention_mask = encoding["attention_mask"]
+        token_type_ids = encoding.get("token_type_ids")
         if torch.cuda.is_available():
             traced_model = torch.jit.load(model_file, map_location='cuda:0')
             traced_model.to(self.device)
             input_ids = input_ids.to(self.device)
             attention_mask = attention_mask.to(self.device)
+            if token_type_ids is not None:
+                token_type_ids = token_type_ids.to(self.device)
         else:
             traced_model = torch.jit.load(model_file)
 
         traced_model.eval()
 
-        # test traced model
-        out = traced_model(input_ids, attention_mask)
+        try:
+            # test traced model
+            if include_types:
+                out = traced_model(input_ids, attention_mask, token_type_ids)
+            else:
+                out = traced_model(input_ids, attention_mask)
+        except RuntimeError as e:
+            logging.warning(e, exc_info=True)
+            return False, "Failed to run inference on jit model"
 
         return self.verify_jit_output(hf_pipeline, encoding, out)
 
diff --git a/extensions/tokenizers/src/main/python/huggingface_models.py b/extensions/tokenizers/src/main/python/huggingface_models.py
index 7dcf3db06..abdc4ede9 100644
--- a/extensions/tokenizers/src/main/python/huggingface_models.py
+++ b/extensions/tokenizers/src/main/python/huggingface_models.py
@@ -13,12 +13,12 @@
 import json
 import logging
 import os
+from argparse import Namespace
 from typing import List
 
 from huggingface_hub import HfApi, ModelSearchArguments
 from huggingface_hub import hf_hub_download
 from huggingface_hub.hf_api import ModelInfo
-from argparse import Namespace
 
 ARCHITECTURES_2_TASK = {
     "ForQuestionAnswering": "question-answering",
@@ -87,11 +87,19 @@ class HuggingfaceModels:
                 logging.warning(f"Skip non-English model: {model_id}.")
                 continue
 
-            if self.processed_models.get(model_id):
+            existing_model = self.processed_models.get(model_id)
+            if existing_model:
+                existing_model["downloads"] = model_info.downloads
                 logging.info(f"Skip converted mode: {model_id}.")
                 continue
 
-            config = hf_hub_download(repo_id=model_id, filename="config.json")
+            try:
+                config = hf_hub_download(repo_id=model_id,
+                                         filename="config.json")
+            except EnvironmentError:
+                logging.info(f"Skip {model_id}, no config.json found.")
+                continue
+
             with open(config) as f:
                 config = json.load(f)
 
@@ -101,7 +109,8 @@ class HuggingfaceModels:
                 task, architecture = self.to_supported_task(config)
                 if not task:
                     logging.info(
-                        f"Unsupported model architecture: {architecture}")
+                        f"Unsupported model architecture: {architecture} for {model_id}."
+                    )
                     continue
 
             model = {
@@ -119,13 +128,14 @@ class HuggingfaceModels:
             "result": "success" if result else "failed",
             "application": application,
             "sha1": model_info.sha,
-            "size": size
+            "size": size,
+            "downloads": model_info.downloads,
         }
         if reason:
             status["reason"] = reason
         self.processed_models[model_info.modelId] = status
 
-        dict_file = os.path.join(self.output_dir, "processed_models.json")
+        dict_file = os.path.join(self.output_dir, "models.json")
         with open(dict_file, 'w') as f:
             json.dump(self.processed_models,
                       f,
@@ -135,7 +145,11 @@ class HuggingfaceModels:
 
     @staticmethod
     def to_supported_task(config: dict):
-        architecture = config.get("architectures")[0]
+        architectures = config.get("architectures")
+        if not architectures:
+            return None, "No architectures found"
+
+        architecture = architectures[0]
         for arch in ARCHITECTURES_2_TASK:
             if architecture.endswith(arch):
                 return ARCHITECTURES_2_TASK[arch], architecture
diff --git a/extensions/tokenizers/src/main/python/metadata.py b/extensions/tokenizers/src/main/python/metadata.py
index 90aa1830f..b4b61f115 100644
--- a/extensions/tokenizers/src/main/python/metadata.py
+++ b/extensions/tokenizers/src/main/python/metadata.py
@@ -16,12 +16,11 @@ from huggingface_models import get_lang_tags
 
 class HuggingfaceMetadata:
 
-    def __init__(self, model_info, artifact_ids: list[str], application: str,
-                 translator: str, sha1: str, file_size: int):
+    def __init__(self, model_info, application: str, translator: str,
+                 sha1: str, file_size: int):
         self.model_info = model_info
-        self.model_name = artifact_ids[-1]
-        artifact_ids[0:-1].insert(0, "ai.djl.huggingface")
-        self.group_id = ".".join(artifact_ids)
+        self.artifact_id = model_info.modelId
+        self.model_name = model_info.modelId.split("/")[-1]
         self.application = application
         self.translator = translator
         self.sha1 = sha1
@@ -31,28 +30,35 @@ class HuggingfaceMetadata:
         properties = get_lang_tags(self.model_info)
 
         metadata = {
-            "metadataVersion": "0.2",
-            "resourceType": "model",
-            "application": self.application,
-            "groupId": self.group_id,
-            "artifactId": self.model_name,
-            "name": self.model_name,
+            "metadataVersion":
+            "0.2",
+            "resourceType":
+            "model",
+            "application":
+            self.application,
+            "groupId":
+            "ai.djl.huggingface.pytorch",
+            "artifactId":
+            self.artifact_id,
+            "name":
+            self.model_name,
             "description":
             f"Huggingface transformers model: {self.model_name}",
-            "website": "http://www.djl.ai/extensions/tokenizers",
+            "website":
+            "http://www.djl.ai/extensions/tokenizers",
             "licenses": {
                 "license": {
                     "name": "The Apache License, Version 2.0",
                     "url": "https://www.apache.org/licenses/LICENSE-2.0"
                 }
             },
-            "artifacts": {
+            "artifacts": [{
                 "version": "0.0.1",
                 "snapshot": False,
                 "name": self.model_name,
-                "downloads": self.model_info.downloads,
                 "properties": properties,
                 "arguments": {
+                    "engine": "PyTorch",
                     "translatorFactory": self.translator
                 },
                 "options": {
@@ -66,7 +72,7 @@ class HuggingfaceMetadata:
                         "size": self.file_size
                     }
                 }
-            }
+            }]
         }
         with open(metadata_file, 'w') as f:
             json.dump(metadata,
diff --git a/extensions/tokenizers/src/main/python/model_zoo_importer.py b/extensions/tokenizers/src/main/python/model_zoo_importer.py
index a8c7e6437..a4e2c6d86 100644
--- a/extensions/tokenizers/src/main/python/model_zoo_importer.py
+++ b/extensions/tokenizers/src/main/python/model_zoo_importer.py
@@ -18,14 +18,16 @@ from arg_parser import converter_args
 from fill_mask_converter import FillMaskConverter
 from huggingface_models import HuggingfaceModels
 from question_answering_converter import QuestionAnsweringConverter
-from token_classification_converter import TokenClassificationConverter
 from sentence_similarity_converter import SentenceSimilarityConverter
+from text_classification_converter import TextClassificationConverter
+from token_classification_converter import TokenClassificationConverter
 
 SUPPORTED_TASK = {
     "fill-mask": FillMaskConverter(),
     "question-answering": QuestionAnsweringConverter(),
+    "sentence-similarity": SentenceSimilarityConverter(),
+    "text-classification": TextClassificationConverter(),
     "token-classification": TokenClassificationConverter(),
-    "sentence-similarity": SentenceSimilarityConverter()
 }
 
 
@@ -48,7 +50,7 @@ def main():
         result, reason, size = converter.save_model(model_info,
                                                     args.output_dir, temp_dir)
         if not result:
-            logging.error(reason)
+            logging.error(f"{model_info.modelId}: {reason}")
 
         huggingface_models.update_progress(model_info, converter.application,
                                            result, reason, size)
diff --git a/extensions/tokenizers/src/main/python/question_answering_converter.py b/extensions/tokenizers/src/main/python/question_answering_converter.py
index e69ce02d8..935c21186 100644
--- a/extensions/tokenizers/src/main/python/question_answering_converter.py
+++ b/extensions/tokenizers/src/main/python/question_answering_converter.py
@@ -39,8 +39,8 @@ class QuestionAnsweringConverter(HuggingfaceConverter):
 
         start_ = out["start_logits"]
         end_ = out["end_logits"]
-        start_[0, 0] = 0
-        end_[0, 0] = 0
+        start_[0, 0] = -10000
+        end_[0, 0] = -10000
         answer_start = torch.argmax(start_)
         answer_end = torch.argmax(end_) + 1
 
@@ -50,8 +50,10 @@ class QuestionAnsweringConverter(HuggingfaceConverter):
 
         if prediction.lower() != self.outputs:
             pipeline_output = hf_pipeline(self.inputs)
-            if pipeline_output != prediction:
-                return False, f"Unexpected inference result: {prediction}"
+            if pipeline_output["answer"].strip() != prediction:
+                logging.error(f"Unexpected inference result: {prediction}")
+                return False, "Unexpected inference result"
+
             logging.warning(
                 f"pipeline output differs from expected: {pipeline_output}")
 
diff --git a/extensions/tokenizers/src/main/python/requirements.txt b/extensions/tokenizers/src/main/python/requirements.txt
index c88c9b742..30990ae90 100644
--- a/extensions/tokenizers/src/main/python/requirements.txt
+++ b/extensions/tokenizers/src/main/python/requirements.txt
@@ -1,4 +1,4 @@
 huggingface_hub
 transformers
-torch
+torch==1.11.0
 protobuf==3.20.0
diff --git a/extensions/tokenizers/src/main/python/sentence_similarity_converter.py b/extensions/tokenizers/src/main/python/sentence_similarity_converter.py
index dbcb587a2..6e372b706 100644
--- a/extensions/tokenizers/src/main/python/sentence_similarity_converter.py
+++ b/extensions/tokenizers/src/main/python/sentence_similarity_converter.py
@@ -43,12 +43,13 @@ class SentenceSimilarityConverter(HuggingfaceConverter):
         return PipelineHolder(tokenizer, model)
 
     def verify_jit_output(self, hf_pipeline, encoding, out):
-        last_hidden_state = out["last_hidden_state"]
+        last_hidden_state = out["last_hidden_state"].to("cpu")
 
         pipeline_output = hf_pipeline.model(**encoding)
-        if not torch.allclose(pipeline_output.last_hidden_state,
-                              last_hidden_state):
-            return False, f"Unexpected inference result: {last_hidden_state}"
+        expected = pipeline_output.last_hidden_state
+        if not torch.allclose(
+                expected, last_hidden_state, atol=1e-05, rtol=1e-03):
+            return False, "Unexpected inference result"
 
         return True, None
 
diff --git a/extensions/tokenizers/src/main/python/text_classification_converter.py b/extensions/tokenizers/src/main/python/text_classification_converter.py
new file mode 100644
index 000000000..8b1dba4ef
--- /dev/null
+++ b/extensions/tokenizers/src/main/python/text_classification_converter.py
@@ -0,0 +1,56 @@
+#!/usr/bin/env python
+#
+# 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.txt" file accompanying this file. This file is distributed on an "AS IS"
+# BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for
+# the specific language governing permissions and limitations under the License.
+import logging
+import math
+
+import torch
+
+from huggingface_converter import HuggingfaceConverter
+
+
+class TextClassificationConverter(HuggingfaceConverter):
+
+    def __init__(self):
+        super().__init__()
+        self.task = "text-classification"
+        self.application = "nlp/text_classification"
+        self.translator = "ai.djl.huggingface.translator.TextClassificationTranslatorFactory"
+        self.inputs = "DJL is the best."
+        self.outputs = None
+
+    def verify_jit_output(self, hf_pipeline, encoding, out):
+        config = hf_pipeline.model.config
+        logits = out['logits'][0]
+
+        if config.problem_type == "multi_label_classification" or config.num_labels == 1:
+            logits = torch.sigmoid(logits)
+        elif config.problem_type == "single_label_classification" or config.num_labels > 1:
+            logits = torch.softmax(logits, dim=0)
+        elif hasattr(config, "function_to_apply"):
+            logging.error(
+                f"Customized function not supported: {config.function_to_apply}"
+            )
+            return False, "Customized function not supported"
+
+        index = logits.argmax().item()
+        label = config.id2label[index]
+        score = logits[index]
+        pipeline_output = hf_pipeline(self.inputs)
+
+        for item in pipeline_output:
+            if item["label"] == label:
+                if math.isclose(item["score"], score, abs_tol=1e-3):
+                    return True, None
+                break
+
+        return False, f"Unexpected inference result"
diff --git a/extensions/tokenizers/src/main/python/token_classification_converter.py b/extensions/tokenizers/src/main/python/token_classification_converter.py
index 5855bcbcf..fead8d247 100644
--- a/extensions/tokenizers/src/main/python/token_classification_converter.py
+++ b/extensions/tokenizers/src/main/python/token_classification_converter.py
@@ -65,12 +65,16 @@ class TokenClassificationConverter(HuggingfaceConverter):
             return False, "TokenClassification returns with empty result"
 
         pipeline_output = hf_pipeline(self.inputs)
-        entity = next(x for x in pipeline_output
-                      if x["word"] == entities[0]["word"])
+        for e in pipeline_output:
+            if e["word"] == entities[0]["word"]:
+                if e["entity"] == entities[0]["entity"]:
+                    logging.warning(
+                        f"pipeline output differs from expected: {pipeline_output}"
+                    )
+                    return True, None
+                else:
+                    break
 
-        if entities[0]["entity"] != entity["entity"]:
-            return False, f"Unexpected inference result: {entities}"
-        logging.warning(
-            f"pipeline output differs from expected: {pipeline_output}")
+        logging.error(f"Unexpected inference result: {entities[0]}")
 
-        return True, None
+        return False, "Unexpected inference result"
diff --git a/model-zoo/src/main/java/ai/djl/basicmodelzoo/BasicModelZoo.java b/model-zoo/src/main/java/ai/djl/basicmodelzoo/BasicModelZoo.java
index 34b2ab8a6..7cdfc040c 100644
--- a/model-zoo/src/main/java/ai/djl/basicmodelzoo/BasicModelZoo.java
+++ b/model-zoo/src/main/java/ai/djl/basicmodelzoo/BasicModelZoo.java
@@ -13,15 +13,10 @@
 package ai.djl.basicmodelzoo;
 
 import ai.djl.Application.CV;
-import ai.djl.repository.MRL;
 import ai.djl.repository.Repository;
-import ai.djl.repository.zoo.BaseModelLoader;
-import ai.djl.repository.zoo.ModelLoader;
 import ai.djl.repository.zoo.ModelZoo;
 
-import java.util.ArrayList;
 import java.util.HashSet;
-import java.util.List;
 import java.util.Set;
 
 /** Model Zoo is a repository that contains all models for DJL. */
@@ -31,23 +26,10 @@ public class BasicModelZoo extends ModelZoo {
     private static final Repository REPOSITORY = Repository.newInstance("zoo", REPO_URL);
     public static final String GROUP_ID = "ai.djl.zoo";
 
-    private static final List<ModelLoader> MODEL_LOADERS = new ArrayList<>();
-
-    static {
-        MRL mlp = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "mlp", "0.0.3");
-        MODEL_LOADERS.add(new BaseModelLoader(mlp));
-
-        MRL resnet = REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "resnet", "0.0.2");
-        MODEL_LOADERS.add(new BaseModelLoader(resnet));
-
-        MRL ssd = REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "ssd", "0.0.2");
-        MODEL_LOADERS.add(new BaseModelLoader(ssd));
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public List<ModelLoader> getModelLoaders() {
-        return MODEL_LOADERS;
+    BasicModelZoo() {
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "mlp", "0.0.3"));
+        addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "resnet", "0.0.2"));
+        addModel(REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "ssd", "0.0.2"));
     }
 
     /** {@inheritDoc} */


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
new file mode 100644
index 000000000..40ec26d27
--- /dev/null
+++ b/api/src/test/java/ai/djl/modality/audio/AudioFactoryTest.java
@@ -0,0 +1,81 @@
+/*
+ * 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.modality.audio;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.training.util.DownloadUtils;
+
+import org.testng.Assert;
+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 {
+
+    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");
+    }
+
+    @Test
+    public void testFromFile() throws IOException {
+        Audio audio = AudioFactory.getInstance().fromFile(Paths.get("build/test/test_01.wav"));
+        Assert.assertEquals(audio.getSampleRate(), 16000f);
+        Assert.assertEquals(audio.getChannels(), 1);
+    }
+
+    @Test
+    public void testFromUrl() throws IOException {
+        Audio audio = AudioFactory.getInstance().fromUrl("build/test/test_01.wav");
+        Assert.assertEquals(audio.getSampleRate(), 16000f);
+        Assert.assertEquals(audio.getChannels(), 1);
+
+        audio = AudioFactory.getInstance().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);
+        Assert.assertEquals(audio.getData(), data);
+        Assert.assertEquals(audio.getSampleRate(), 0);
+        Assert.assertEquals(audio.getChannels(), 0);
+    }
+
+    @Test(expectedExceptions = UnsupportedOperationException.class)
+    public void testFromNDArray() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray array = manager.zeros(new Shape(1));
+            AudioFactory.getInstance().fromNDArray(array);
+        }
+    }
+}
diff --git a/api/src/test/java/ai/djl/modality/audio/package-info.java b/api/src/test/java/ai/djl/modality/audio/package-info.java
new file mode 100644
index 000000000..29a1b0cbc
--- /dev/null
+++ b/api/src/test/java/ai/djl/modality/audio/package-info.java
@@ -0,0 +1,15 @@
+/*
+ * Copyright 2019 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 tests for the audio api. */
+package ai.djl.modality.audio;
diff --git a/basicdataset/src/test/java/ai/djl/basicdataset/MnistTest.java b/basicdataset/src/test/java/ai/djl/basicdataset/MnistTest.java
index be531b46e..8d0e609cc 100644
--- a/basicdataset/src/test/java/ai/djl/basicdataset/MnistTest.java
+++ b/basicdataset/src/test/java/ai/djl/basicdataset/MnistTest.java
@@ -15,20 +15,27 @@ package ai.djl.basicdataset;
 import ai.djl.Model;
 import ai.djl.basicdataset.cv.classification.Mnist;
 import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.Blocks;
 import ai.djl.repository.Repository;
 import ai.djl.training.DefaultTrainingConfig;
 import ai.djl.training.Trainer;
 import ai.djl.training.TrainingConfig;
 import ai.djl.training.dataset.Batch;
+import ai.djl.training.dataset.BulkDataIterable;
 import ai.djl.training.dataset.Dataset;
 import ai.djl.training.loss.Loss;
+import ai.djl.translate.StackBatchifier;
 import ai.djl.translate.TranslateException;
 
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
 
 public class MnistTest {
 
@@ -81,4 +88,96 @@ public class MnistTest {
             }
         }
     }
+
+    @SuppressWarnings("unchecked")
+    @Test
+    public void testRandomSamplingCoversAllIndices() throws IOException, TranslateException {
+        TrainingConfig config = new DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss());
+
+        try (Model model = Model.newInstance("model")) {
+            model.setBlock(Blocks.identityBlock());
+
+            NDManager manager = model.getNDManager();
+            Repository repository = Repository.newInstance("test", "src/test/resources/mlrepo");
+            Mnist mnist =
+                    Mnist.builder()
+                            .optManager(manager)
+                            .optUsage(Dataset.Usage.TEST)
+                            .optRepository(repository)
+                            .setSampling(32, true)
+                            .build();
+
+            List<Long> randomizedIndices = new ArrayList<>();
+            try (Trainer trainer = model.newTrainer(config)) {
+                for (Batch batch : trainer.iterateDataset(mnist)) {
+                    randomizedIndices.addAll((List<Long>) batch.getIndices());
+                    Assert.assertEquals(batch.getData().size(), 1);
+                    Assert.assertEquals(batch.getLabels().size(), 1);
+                    if (randomizedIndices.size() == 10000) {
+                        // 312 * 32 + 16 = 10000, so last batch has only 16 items left
+                        Assert.assertEquals(
+                                new Shape(16, 1, 28, 28), batch.getData().get(0).getShape());
+                    } else {
+                        Assert.assertEquals(
+                                new Shape(32, 1, 28, 28), batch.getData().get(0).getShape());
+                    }
+                    batch.close();
+                }
+            }
+            Assert.assertEquals(10000, randomizedIndices.size());
+            Assert.assertFalse(BulkDataIterable.isRange(randomizedIndices));
+            Collections.sort(randomizedIndices);
+            Assert.assertTrue(BulkDataIterable.isRange(randomizedIndices));
+        }
+    }
+
+    @Test
+    public void testBulkEqualsNonBulk() throws IOException, TranslateException {
+        TrainingConfig config = new DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss());
+
+        try (Model model = Model.newInstance("model")) {
+            model.setBlock(Blocks.identityBlock());
+
+            NDManager manager = model.getNDManager();
+            Repository repository = Repository.newInstance("test", "src/test/resources/mlrepo");
+
+            // use BulkDataIterable
+            Mnist mnistBulk =
+                    Mnist.builder()
+                            .optManager(manager)
+                            .optUsage(Dataset.Usage.TEST)
+                            .optRepository(repository)
+                            .setSampling(32, false)
+                            .build();
+
+            // use only DataIterable (as not using Batchifier.STACK)
+            Mnist mnist =
+                    Mnist.builder()
+                            .optManager(manager)
+                            .optUsage(Dataset.Usage.TEST)
+                            .optRepository(repository)
+                            .setSampling(32, false)
+                            .optLabelBatchifier(new StackBatchifier() {})
+                            .build();
+
+            try (Trainer trainer = model.newTrainer(config)) {
+                Iterator<Batch> iterable = trainer.iterateDataset(mnist).iterator();
+                Iterator<Batch> bulkIterable = trainer.iterateDataset(mnistBulk).iterator();
+                Assert.assertFalse(iterable instanceof BulkDataIterable);
+                Assert.assertTrue(bulkIterable instanceof BulkDataIterable);
+                while (iterable.hasNext()) {
+                    Batch batch = iterable.next();
+                    Assert.assertTrue(bulkIterable.hasNext());
+                    Batch batchBulk = bulkIterable.next();
+                    Assert.assertEquals(batch.getIndices(), batchBulk.getIndices());
+                    Assert.assertEquals(batch.getSize(), batchBulk.getSize());
+                    Assert.assertEquals(batch.getData(), batchBulk.getData());
+                    Assert.assertEquals(batch.getLabels(), batchBulk.getLabels());
+                    batch.close();
+                    batchBulk.close();
+                }
+                Assert.assertFalse(bulkIterable.hasNext());
+            }
+        }
+    }
 }
diff --git a/basicdataset/src/test/java/ai/djl/basicdataset/RandomAccessDatasetTest.java b/basicdataset/src/test/java/ai/djl/basicdataset/RandomAccessDatasetTest.java
index 01247a376..b4cb4c045 100644
--- a/basicdataset/src/test/java/ai/djl/basicdataset/RandomAccessDatasetTest.java
+++ b/basicdataset/src/test/java/ai/djl/basicdataset/RandomAccessDatasetTest.java
@@ -26,6 +26,10 @@ import org.testng.Assert;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.LongStream;
 
 public class RandomAccessDatasetTest {
 
@@ -53,6 +57,29 @@ public class RandomAccessDatasetTest {
             Assert.assertEquals(15, subset.size());
             Record record = subset.get(manager, 0);
             Assert.assertEquals(record.getData().head().getLong(), 10);
+
+            /** Tests {@link #ArrayDataset.getByRange()} */
+            List<Long> indices =
+                    LongStream.range(10, 25).mapToObj(n -> n).collect(Collectors.toList());
+            subset = dataset.subDataset(indices);
+            Assert.assertEquals(15, subset.size());
+            record = subset.get(manager, 0);
+            Assert.assertEquals(record.getData().head().getLong(), 10);
+
+            /** Tests {@link #ArrayDataset.getByIndices()} */
+            Collections.reverse(indices);
+            subset = dataset.subDataset(indices);
+            record = subset.get(manager, 0);
+            Assert.assertEquals(record.getData().head().getLong(), 24);
+
+            /** Tests {@link #RandomAccessDataset.subDataset(recordKeys, subRecordKeys)} */
+            List<String> subRecordKeys =
+                    indices.stream().map(k -> "key" + k).collect(Collectors.toList());
+            List<String> recordKeys =
+                    LongStream.range(0, 100).mapToObj(k -> "key" + k).collect(Collectors.toList());
+            subset = dataset.subDataset(recordKeys, subRecordKeys);
+            record = subset.get(manager, 0);
+            Assert.assertEquals(record.getData().head().getLong(), 24);
         }
     }
 }
diff --git a/engines/mxnet/mxnet-model-zoo/src/test/java/ai/djl/mxnet/integration/modality/nlp/WordEmbeddingTest.java b/engines/mxnet/mxnet-model-zoo/src/test/java/ai/djl/mxnet/integration/modality/nlp/WordEmbeddingTest.java
index 726374423..7f6822359 100644
--- a/engines/mxnet/mxnet-model-zoo/src/test/java/ai/djl/mxnet/integration/modality/nlp/WordEmbeddingTest.java
+++ b/engines/mxnet/mxnet-model-zoo/src/test/java/ai/djl/mxnet/integration/modality/nlp/WordEmbeddingTest.java
@@ -32,6 +32,7 @@ public class WordEmbeddingTest {
     @Test
     public void testGlove() throws IOException, ModelException, EmbeddingException {
         TestRequirements.notArm();
+        TestRequirements.engine("MXNet");
 
         Criteria<String, NDList> criteria =
                 Criteria.builder()
diff --git a/examples/src/test/java/ai/djl/examples/inference/InstanceSegmentationTest.java b/examples/src/test/java/ai/djl/examples/inference/InstanceSegmentationTest.java
index 68a4b0c96..33812b409 100644
--- a/examples/src/test/java/ai/djl/examples/inference/InstanceSegmentationTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/InstanceSegmentationTest.java
@@ -28,6 +28,7 @@ public class InstanceSegmentationTest {
     @Test
     public void testInstanceSegmentation() throws ModelException, TranslateException, IOException {
         TestRequirements.engine("MXNet");
+        TestRequirements.nightly();
 
         DetectedObjects result = InstanceSegmentation.predict();
         Classifications.Classification best = result.best();
diff --git a/examples/src/test/java/ai/djl/examples/inference/ListModelsTest.java b/examples/src/test/java/ai/djl/examples/inference/ListModelsTest.java
index a8266c2b1..2a0c58723 100644
--- a/examples/src/test/java/ai/djl/examples/inference/ListModelsTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/ListModelsTest.java
@@ -18,6 +18,7 @@ import ai.djl.Application.NLP;
 import ai.djl.ModelException;
 import ai.djl.repository.Artifact;
 import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.DefaultZooProvider;
 import ai.djl.repository.zoo.ModelZoo;
 
 import org.testng.Assert;
@@ -36,6 +37,7 @@ public class ListModelsTest {
     @AfterClass
     public void tearDown() {
         System.clearProperty("ai.djl.repository.zoo.location");
+        ModelZoo.registerModelZoo(new DefaultZooProvider());
     }
 
     @Test
@@ -43,6 +45,7 @@ public class ListModelsTest {
         Path path = Paths.get("../model-zoo/src/test/resources/mlrepo");
         String repoUrl = path.toRealPath().toAbsolutePath().toUri().toURL().toExternalForm();
         System.setProperty("ai.djl.repository.zoo.location", "src/test/resources," + repoUrl);
+        ModelZoo.registerModelZoo(new DefaultZooProvider());
         Map<Application, List<Artifact>> models = ModelZoo.listModels();
         System.out.println(Arrays.toString(models.keySet().toArray()));
         List<Artifact> artifacts = models.get(Application.UNDEFINED);
@@ -54,6 +57,7 @@ public class ListModelsTest {
         Path path = Paths.get("../model-zoo/src/test/resources/mlrepo");
         String repoUrl = path.toRealPath().toAbsolutePath().toUri().toURL().toExternalForm();
         System.setProperty("ai.djl.repository.zoo.location", "src/test/resources," + repoUrl);
+        ModelZoo.registerModelZoo(new DefaultZooProvider());
         Criteria<?, ?> criteria = Criteria.builder().optApplication(NLP.ANY).build();
         Map<Application, List<Artifact>> models = ModelZoo.listModels(criteria);
 
diff --git a/examples/src/test/java/ai/djl/examples/inference/SpeechRecognitionTest.java b/examples/src/test/java/ai/djl/examples/inference/SpeechRecognitionTest.java
new file mode 100644
index 000000000..a26d4dd7e
--- /dev/null
+++ b/examples/src/test/java/ai/djl/examples/inference/SpeechRecognitionTest.java
@@ -0,0 +1,40 @@
+/*
+ * 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.examples.inference;
+
+import ai.djl.ModelException;
+import ai.djl.testing.TestRequirements;
+import ai.djl.translate.TranslateException;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+
+import javax.sound.sampled.UnsupportedAudioFileException;
+
+public class SpeechRecognitionTest {
+
+    @Test
+    public void testSpeechRecognition()
+            throws ModelException, TranslateException, IOException, UnsupportedAudioFileException {
+        TestRequirements.nightly();
+        TestRequirements.engine("PyTorch");
+
+        String result = SpeechRecognition.predict();
+        Assert.assertEquals(
+                result,
+                "THE NEAREST SAID THE DISTRICT DOCTOR IS A GOOD ITALIAN ABBE WHO LIVES NEXT DOOR TO"
+                        + " YOU SHALL I CALL ON HIM AS I PASS ");
+    }
+}
diff --git a/extensions/audio/src/test/java/ai/djl/audio/FFmpegAudioFactoryTest.java b/extensions/audio/src/test/java/ai/djl/audio/FFmpegAudioFactoryTest.java
new file mode 100644
index 000000000..4539fedd7
--- /dev/null
+++ b/extensions/audio/src/test/java/ai/djl/audio/FFmpegAudioFactoryTest.java
@@ -0,0 +1,68 @@
+/*
+ * 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.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.testing.TestRequirements;
+import ai.djl.training.util.DownloadUtils;
+
+import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+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 {
+        TestRequirements.notArm();
+        DownloadUtils.download(URL, "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);
+
+        audio = AudioFactory.getInstance().fromUrl("build/test/test_01.wav");
+        Assert.assertEquals(audio.getSampleRate(), 16000f);
+        Assert.assertEquals(audio.getChannels(), 1);
+
+        audio = AudioFactory.getInstance().fromUrl(URL);
+        Assert.assertEquals(audio.getSampleRate(), 16000f);
+        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);
+    }
+
+    @Test(expectedExceptions = UnsupportedOperationException.class)
+    public void testFromNDArray() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray array = manager.zeros(new Shape(1));
+            AudioFactory.getInstance().fromNDArray(array);
+        }
+    }
+}
diff --git a/extensions/audio/src/test/java/ai/djl/audio/package-info.java b/extensions/audio/src/test/java/ai/djl/audio/package-info.java
new file mode 100644
index 000000000..24d79f6c5
--- /dev/null
+++ b/extensions/audio/src/test/java/ai/djl/audio/package-info.java
@@ -0,0 +1,13 @@
+/*
+ * 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.audio;
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 0c5894dd5..df9fa801f 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
@@ -14,8 +14,8 @@
 package ai.djl.audio.processor;
 
 import ai.djl.Device;
-import ai.djl.audio.AudioUtils;
 import ai.djl.audio.dataset.AudioData;
+import ai.djl.audio.util.AudioUtils;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.Shape;
diff --git a/extensions/timeseries/src/test/java/ai/djl/timeseries/transform/TransformTest.java b/extensions/timeseries/src/test/java/ai/djl/timeseries/transform/TransformTest.java
index 8e521f048..0a840fcc4 100644
--- a/extensions/timeseries/src/test/java/ai/djl/timeseries/transform/TransformTest.java
+++ b/extensions/timeseries/src/test/java/ai/djl/timeseries/transform/TransformTest.java
@@ -18,6 +18,7 @@ import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.index.NDIndex;
 import ai.djl.ndarray.types.Shape;
+import ai.djl.testing.TestRequirements;
 import ai.djl.timeseries.TimeSeriesData;
 import ai.djl.timeseries.dataset.FieldName;
 import ai.djl.timeseries.timefeature.TimeFeature;
@@ -27,6 +28,7 @@ import ai.djl.timeseries.transform.field.Field;
 import ai.djl.timeseries.transform.split.Split;
 
 import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 import java.time.LocalDateTime;
@@ -36,6 +38,11 @@ import java.util.function.BiFunction;
 
 public class TransformTest {
 
+    @BeforeClass
+    public void setUp() {
+        TestRequirements.notArm();
+    }
+
     @Test
     public void testVstackFeatures() {
         try (NDManager manager = NDManager.newBaseManager(Device.cpu())) {
diff --git a/extensions/timeseries/src/test/java/ai/djl/timeseries/translator/DeepARTranslatorTest.java b/extensions/timeseries/src/test/java/ai/djl/timeseries/translator/DeepARTranslatorTest.java
index 8e64746ba..4c442dde1 100644
--- a/extensions/timeseries/src/test/java/ai/djl/timeseries/translator/DeepARTranslatorTest.java
+++ b/extensions/timeseries/src/test/java/ai/djl/timeseries/translator/DeepARTranslatorTest.java
@@ -19,6 +19,7 @@ import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ZooModel;
+import ai.djl.testing.TestRequirements;
 import ai.djl.timeseries.Forecast;
 import ai.djl.timeseries.TimeSeriesData;
 import ai.djl.timeseries.dataset.FieldName;
@@ -37,6 +38,8 @@ public class DeepARTranslatorTest {
 
     @Test
     public void testDeepARTranslator() throws IOException, TranslateException, ModelException {
+        TestRequirements.notArm();
+
         String modelUrl = "https://resources.djl.ai/test-models/mxnet/timeseries/deepar.zip";
         Map<String, Object> arguments = new ConcurrentHashMap<>();
         arguments.put("prediction_length", 28);
diff --git a/extensions/timeseries/src/test/java/ai/djl/timeseries/translator/TransformerTranslatorTest.java b/extensions/timeseries/src/test/java/ai/djl/timeseries/translator/TransformerTranslatorTest.java
index e12af5150..49295b4f9 100644
--- a/extensions/timeseries/src/test/java/ai/djl/timeseries/translator/TransformerTranslatorTest.java
+++ b/extensions/timeseries/src/test/java/ai/djl/timeseries/translator/TransformerTranslatorTest.java
@@ -19,6 +19,7 @@ import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ZooModel;
+import ai.djl.testing.TestRequirements;
 import ai.djl.timeseries.Forecast;
 import ai.djl.timeseries.TimeSeriesData;
 import ai.djl.timeseries.dataset.FieldName;
@@ -37,6 +38,8 @@ public class TransformerTranslatorTest {
 
     @Test
     public void testTransformerTranslator() throws IOException, TranslateException, ModelException {
+        TestRequirements.notArm();
+
         String modelUrl = "https://resources.djl.ai/test-models/mxnet/timeseries/transformer.zip";
         Map<String, Object> arguments = new ConcurrentHashMap<>();
         arguments.put("prediction_length", 28);
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 88c72d41c..c37743565 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
@@ -33,7 +33,6 @@ import ai.djl.nn.LambdaBlock;
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ZooModel;
 import ai.djl.testing.Assertions;
-import ai.djl.testing.TestRequirements;
 import ai.djl.translate.TranslateException;
 import ai.djl.util.JsonUtils;
 import ai.djl.util.Utils;
@@ -59,8 +58,6 @@ public class TranslatorTest {
 
     @Test
     public void testQATranslator() throws ModelException, IOException, TranslateException {
-        TestRequirements.notArm();
-
         String question = "When did BBC Japan start broadcasting?";
         String paragraph =
                 "BBC Japan was a general entertainment Channel. "
@@ -146,8 +143,6 @@ public class TranslatorTest {
 
     @Test
     public void testFillMaskTranslator() throws ModelException, IOException, TranslateException {
-        TestRequirements.notArm();
-
         String text = "Hello I'm a [MASK] model.";
 
         Block block =
@@ -235,8 +230,6 @@ public class TranslatorTest {
     @Test
     public void testTokenClassificationTranslator()
             throws ModelException, IOException, TranslateException {
-        TestRequirements.notArm();
-
         String text = "My name is Wolfgang and I live in Berlin.";
 
         Block block =
@@ -331,8 +324,6 @@ public class TranslatorTest {
     @Test
     public void testTextEmbeddingTranslator()
             throws ModelException, IOException, TranslateException {
-        TestRequirements.notArm();
-
         String text = "This is an example sentence";
 
         Block block =
diff --git a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDIndexTest.java b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDIndexTest.java
index 5c59092f8..6a6a3fe1f 100644
--- a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDIndexTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDIndexTest.java
@@ -50,6 +50,15 @@ public class NDIndexTest {
                     original.get(
                             new NDIndex().addAllDim().addPickDim(manager.create(new int[] {0, 1})));
             Assert.assertEquals(actual, expected);
+
+            // The difference between take and pick used combined with addAllDim()
+            NDArray yHat = manager.create(new float[][] {{0.1f, 0.3f, 0.6f}, {0.3f, 0.2f, 0.5f}});
+            NDArray yGet = yHat.get(new NDIndex(":, {}", manager.create(new int[] {0, 2})));
+            NDArray yPick =
+                    yHat.get(
+                            new NDIndex().addAllDim().addPickDim(manager.create(new int[] {0, 2})));
+            Assert.assertEquals(yGet, manager.create(new float[][] {{0.1f, 0.6f}, {0.3f, 0.5f}}));
+            Assert.assertEquals(yPick, manager.create(new float[][] {{0.1f}, {0.5f}}));
         }
     }
 
diff --git a/integration/src/main/java/ai/djl/integration/tests/nn/BlockCoreTest.java b/integration/src/main/java/ai/djl/integration/tests/nn/BlockCoreTest.java
index f2dccd988..5bfd245c5 100644
--- a/integration/src/main/java/ai/djl/integration/tests/nn/BlockCoreTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/nn/BlockCoreTest.java
@@ -198,7 +198,7 @@ public class BlockCoreTest {
         // number of linear transformations modeled by LinearCollection
         int splitSize = 6;
 
-        // following ndarrays will be scoped to modelLinearCollection's ndmanager so they can be
+        // following ndarrays will be scoped to <code>sharedManager</code> so they can be
         // used later in the for-loop
         NDArray expectedWeightForLinearCollection;
         NDArray expectedBiasForLinearCollection;
@@ -256,21 +256,13 @@ public class BlockCoreTest {
                         // copy initial weight/bias values from LinearCollection
                         Shape inputShape = new Shape(batchSize, featureSize);
                         block.setInitializer(
-                                (m, s, t) -> {
-                                    NDArray w =
-                                            expectedWeightForLinearCollection
-                                                    .get(splitIndex)
-                                                    .transpose();
-                                    w.attach(m);
-                                    return w;
-                                },
+                                (m, s, t) ->
+                                        expectedWeightForLinearCollection
+                                                .get(m, splitIndex)
+                                                .transpose(),
                                 Parameter.Type.WEIGHT);
                         block.setInitializer(
-                                (m, s, t) -> {
-                                    NDArray w = expectedBiasForLinearCollection.get(splitIndex);
-                                    w.attach(m);
-                                    return w;
-                                },
+                                (m, s, t) -> expectedBiasForLinearCollection.get(m, splitIndex),
                                 Parameter.Type.BIAS);
                         trainer.initialize(inputShape);