Repository: deepjavalibrary/djl
Pull request #1995 (2022-09-06 18:38:43)

Return Batches with ArrayDataset helpers
Fixes minor issue with #1869


Patch:
diff --git a/api/src/main/java/ai/djl/Application.java b/api/src/main/java/ai/djl/Application.java
index 81e98b13f..cfd8b41b3 100644
--- a/api/src/main/java/ai/djl/Application.java
+++ b/api/src/main/java/ai/djl/Application.java
@@ -70,20 +70,24 @@ public class Application {
                 return CV.IMAGE_ENHANCEMENT;
             case "nlp":
                 return NLP.ANY;
-            case "nlp/question_answer":
-                return NLP.QUESTION_ANSWER;
-            case "nlp/text_classification":
-                return NLP.TEXT_CLASSIFICATION;
-            case "nlp/sentiment_analysis":
-                return NLP.SENTIMENT_ANALYSIS;
-            case "nlp/word_embedding":
-                return NLP.WORD_EMBEDDING;
+            case "nlp/fill_mask":
+                return NLP.FILL_MASK;
             case "nlp/machine_translation":
                 return NLP.MACHINE_TRANSLATION;
             case "nlp/multiple_choice":
                 return NLP.MULTIPLE_CHOICE;
+            case "nlp/question_answer":
+                return NLP.QUESTION_ANSWER;
+            case "nlp/sentiment_analysis":
+                return NLP.SENTIMENT_ANALYSIS;
+            case "nlp/text_classification":
+                return NLP.TEXT_CLASSIFICATION;
             case "nlp/text_embedding":
                 return NLP.TEXT_EMBEDDING;
+            case "nlp/token_classification":
+                return NLP.TOKEN_CLASSIFICATION;
+            case "nlp/word_embedding":
+                return NLP.WORD_EMBEDDING;
             case "tabular":
                 return Tabular.ANY;
             case "tabular/linear_regression":
@@ -210,6 +214,12 @@ public class Application {
         /** Any NLP application, including those in {@link NLP}. */
         Application ANY = new Application("nlp");
 
+        /**
+         * An application that masking some words in a sentence and predicting which words should
+         * replace those masks.
+         */
+        Application FILL_MASK = new Application("nlp/fill_mask");
+
         /**
          * An application that a reference document and a question about the document and returns
          * text answering the question.
@@ -228,11 +238,17 @@ public class Application {
         Application TEXT_CLASSIFICATION = new Application("nlp/text_classification");
 
         /**
-         * An application that classifies text into positive or negative, an specific case of {@link
+         * An application that classifies text into positive or negative, a specific case of {@link
          * #TEXT_CLASSIFICATION}.
          */
         Application SENTIMENT_ANALYSIS = new Application("nlp/sentiment_analysis");
 
+        /**
+         * A natural language understanding application that assigns a label to some tokens in a
+         * text.
+         */
+        Application TOKEN_CLASSIFICATION = new Application("nlp/token_classification");
+
         /**
          * An application that takes a word and returns a feature vector that represents the word.
          *
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
index f09a415e6..e4dba51b4 100644
--- a/api/src/main/java/ai/djl/modality/audio/translator/SpeechRecognitionTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/audio/translator/SpeechRecognitionTranslatorFactory.java
@@ -14,7 +14,6 @@ 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;
@@ -41,11 +40,11 @@ public class SpeechRecognitionTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
-            throws TranslateException {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
         if (input == Audio.class && output == String.class) {
-            return new SpeechRecognitionTranslator();
+            return (Translator<I, O>) new SpeechRecognitionTranslator();
         }
         throw new IllegalArgumentException("Unsupported input/output types.");
     }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/BigGANTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/BigGANTranslatorFactory.java
index e2319dbcd..ef738250e 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/BigGANTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/BigGANTranslatorFactory.java
@@ -15,7 +15,6 @@ package ai.djl.modality.cv.translator;
 import ai.djl.Model;
 import ai.djl.modality.cv.Image;
 import ai.djl.translate.ArgumentsUtil;
-import ai.djl.translate.TranslateException;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 import ai.djl.util.Pair;
@@ -36,13 +35,13 @@ public class BigGANTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
-            throws TranslateException {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
         if (!isSupported(input, output)) {
             throw new IllegalArgumentException("Unsupported input/output types.");
         }
         float truncation = ArgumentsUtil.floatValue(arguments, "truncation", 0.5f);
-        return new BigGANTranslator(truncation);
+        return (Translator<I, O>) new BigGANTranslator(truncation);
     }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/ImageClassificationTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/ImageClassificationTranslatorFactory.java
index a963e856f..0cdcdea64 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/ImageClassificationTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/ImageClassificationTranslatorFactory.java
@@ -53,20 +53,21 @@ public class ImageClassificationTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments) {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
+        ImageClassificationTranslator translator =
+                ImageClassificationTranslator.builder(arguments).build();
         if (input == Image.class && output == Classifications.class) {
-            return ImageClassificationTranslator.builder(arguments).build();
+            return (Translator<I, O>) translator;
         } else if (input == Path.class && output == Classifications.class) {
-            return new FileTranslator<>(ImageClassificationTranslator.builder(arguments).build());
+            return (Translator<I, O>) new FileTranslator<>(translator);
         } else if (input == URL.class && output == Classifications.class) {
-            return new UrlTranslator<>(ImageClassificationTranslator.builder(arguments).build());
+            return (Translator<I, O>) new UrlTranslator<>(translator);
         } else if (input == InputStream.class && output == Classifications.class) {
-            return new InputStreamTranslator<>(
-                    ImageClassificationTranslator.builder(arguments).build());
+            return (Translator<I, O>) new InputStreamTranslator<>(translator);
         } else if (input == Input.class && output == Output.class) {
-            return new ImageServingTranslator(
-                    ImageClassificationTranslator.builder(arguments).build());
+            return (Translator<I, O>) new ImageServingTranslator(translator);
         }
         throw new IllegalArgumentException("Unsupported input/output types.");
     }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslatorFactory.java
index 613a09832..7394356cb 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslatorFactory.java
@@ -33,20 +33,21 @@ public class InstanceSegmentationTranslatorFactory extends ObjectDetectionTransl
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments) {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
+        InstanceSegmentationTranslator translator =
+                InstanceSegmentationTranslator.builder(arguments).build();
         if (input == Image.class && output == DetectedObjects.class) {
-            return InstanceSegmentationTranslator.builder(arguments).build();
+            return (Translator<I, O>) translator;
         } else if (input == Path.class && output == DetectedObjects.class) {
-            return new FileTranslator<>(InstanceSegmentationTranslator.builder(arguments).build());
+            return (Translator<I, O>) new FileTranslator<>(translator);
         } else if (input == URL.class && output == DetectedObjects.class) {
-            return new UrlTranslator<>(InstanceSegmentationTranslator.builder(arguments).build());
+            return (Translator<I, O>) new UrlTranslator<>(translator);
         } else if (input == InputStream.class && output == DetectedObjects.class) {
-            return new InputStreamTranslator<>(
-                    InstanceSegmentationTranslator.builder(arguments).build());
+            return (Translator<I, O>) new InputStreamTranslator<>(translator);
         } else if (input == Input.class && output == Output.class) {
-            return new ImageServingTranslator(
-                    InstanceSegmentationTranslator.builder(arguments).build());
+            return (Translator<I, O>) new ImageServingTranslator(translator);
         }
         throw new IllegalArgumentException("Unsupported input/output types.");
     }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactory.java
index bbd3613c0..ec69dabdd 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactory.java
@@ -42,20 +42,21 @@ public class SemanticSegmentationTranslatorFactory implements TranslatorFactory
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments) {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
+        SemanticSegmentationTranslator translator =
+                SemanticSegmentationTranslator.builder(arguments).build();
         if (input == Image.class && output == Image.class) {
-            return SemanticSegmentationTranslator.builder(arguments).build();
+            return (Translator<I, O>) translator;
         } else if (input == Path.class && output == Image.class) {
-            return new FileTranslator<>(SemanticSegmentationTranslator.builder(arguments).build());
+            return (Translator<I, O>) new FileTranslator<>(translator);
         } else if (input == URL.class && output == Image.class) {
-            return new UrlTranslator<>(SemanticSegmentationTranslator.builder(arguments).build());
+            return (Translator<I, O>) new UrlTranslator<>(translator);
         } else if (input == InputStream.class && output == Image.class) {
-            return new InputStreamTranslator<>(
-                    SemanticSegmentationTranslator.builder(arguments).build());
+            return (Translator<I, O>) new InputStreamTranslator<>(translator);
         } else if (input == Input.class && output == Output.class) {
-            return new ImageServingTranslator(
-                    SemanticSegmentationTranslator.builder(arguments).build());
+            return (Translator<I, O>) new ImageServingTranslator(translator);
         }
         throw new IllegalArgumentException("Unsupported input/output types.");
     }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/SimplePoseTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/SimplePoseTranslatorFactory.java
index ad2bd9b6f..33c0ee858 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/SimplePoseTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/SimplePoseTranslatorFactory.java
@@ -44,14 +44,15 @@ public class SimplePoseTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments) {
-        if (!isSupported(input, output)) {
-            throw new IllegalArgumentException("Unsupported input/output types.");
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
+        SimplePoseTranslator translator = SimplePoseTranslator.builder(arguments).build();
+        if (input == Image.class && output == Joints.class) {
+            return (Translator<I, O>) translator;
+        } else if (input == Input.class && output == Output.class) {
+            return (Translator<I, O>) new ImageServingTranslator(translator);
         }
-        if (input == Input.class && output == Output.class) {
-            return new ImageServingTranslator(SimplePoseTranslator.builder(arguments).build());
-        }
-        return SimplePoseTranslator.builder(arguments).build();
+        throw new IllegalArgumentException("Unsupported input/output types.");
     }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/SingleShotDetectionTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/SingleShotDetectionTranslatorFactory.java
index aeb874737..d4f794ad7 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/SingleShotDetectionTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/SingleShotDetectionTranslatorFactory.java
@@ -33,20 +33,21 @@ public class SingleShotDetectionTranslatorFactory extends ObjectDetectionTransla
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments) {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
+        SingleShotDetectionTranslator translator =
+                SingleShotDetectionTranslator.builder(arguments).build();
         if (input == Image.class && output == DetectedObjects.class) {
-            return SingleShotDetectionTranslator.builder(arguments).build();
+            return (Translator<I, O>) translator;
         } else if (input == Path.class && output == DetectedObjects.class) {
-            return new FileTranslator<>(SingleShotDetectionTranslator.builder(arguments).build());
+            return (Translator<I, O>) new FileTranslator<>(translator);
         } else if (input == URL.class && output == DetectedObjects.class) {
-            return new UrlTranslator<>(SingleShotDetectionTranslator.builder(arguments).build());
+            return (Translator<I, O>) new UrlTranslator<>(translator);
         } else if (input == InputStream.class && output == DetectedObjects.class) {
-            return new InputStreamTranslator<>(
-                    SingleShotDetectionTranslator.builder(arguments).build());
+            return (Translator<I, O>) new InputStreamTranslator<>(translator);
         } else if (input == Input.class && output == Output.class) {
-            return new ImageServingTranslator(
-                    SingleShotDetectionTranslator.builder(arguments).build());
+            return (Translator<I, O>) new ImageServingTranslator(translator);
         }
         throw new IllegalArgumentException("Unsupported input/output types.");
     }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/StyleTransferTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/StyleTransferTranslatorFactory.java
index 68af21ff8..4017bdfe4 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/StyleTransferTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/StyleTransferTranslatorFactory.java
@@ -14,7 +14,6 @@ package ai.djl.modality.cv.translator;
 
 import ai.djl.Model;
 import ai.djl.modality.cv.Image;
-import ai.djl.translate.TranslateException;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 import ai.djl.util.Pair;
@@ -35,12 +34,12 @@ public class StyleTransferTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
-            throws TranslateException {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
         if (!isSupported(input, output)) {
             throw new IllegalArgumentException("Unsupported input/output types.");
         }
-        return new StyleTransferTranslator();
+        return (Translator<I, O>) new StyleTransferTranslator();
     }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/YoloTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/YoloTranslatorFactory.java
index 9b57a2f31..b474e15fc 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/YoloTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/YoloTranslatorFactory.java
@@ -20,7 +20,6 @@ import ai.djl.modality.cv.output.DetectedObjects;
 import ai.djl.modality.cv.translator.wrapper.FileTranslator;
 import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
 import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
-import ai.djl.translate.TranslateException;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 
@@ -34,19 +33,20 @@ public class YoloTranslatorFactory extends ObjectDetectionTranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
-            throws TranslateException {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
+        YoloTranslator translator = YoloTranslator.builder(arguments).build();
         if (input == Image.class && output == DetectedObjects.class) {
-            return YoloTranslator.builder(arguments).build();
+            return (Translator<I, O>) translator;
         } else if (input == Path.class && output == DetectedObjects.class) {
-            return new FileTranslator<>(YoloTranslator.builder(arguments).build());
+            return (Translator<I, O>) new FileTranslator<>(translator);
         } else if (input == URL.class && output == DetectedObjects.class) {
-            return new UrlTranslator<>(YoloTranslator.builder(arguments).build());
+            return (Translator<I, O>) new UrlTranslator<>(translator);
         } else if (input == InputStream.class && output == DetectedObjects.class) {
-            return new InputStreamTranslator<>(YoloTranslator.builder(arguments).build());
+            return (Translator<I, O>) new InputStreamTranslator<>(translator);
         } else if (input == Input.class && output == Output.class) {
-            return new ImageServingTranslator(YoloTranslator.builder(arguments).build());
+            return (Translator<I, O>) new ImageServingTranslator(translator);
         }
         throw new IllegalArgumentException("Unsupported input/output types.");
     }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/YoloV5TranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/YoloV5TranslatorFactory.java
index da26aae6b..c18051aa0 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/YoloV5TranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/YoloV5TranslatorFactory.java
@@ -20,7 +20,6 @@ import ai.djl.modality.cv.output.DetectedObjects;
 import ai.djl.modality.cv.translator.wrapper.FileTranslator;
 import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
 import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
-import ai.djl.translate.TranslateException;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 
@@ -34,19 +33,20 @@ public class YoloV5TranslatorFactory extends ObjectDetectionTranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
-            throws TranslateException {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
+        YoloV5Translator translator = YoloV5Translator.builder(arguments).build();
         if (input == Image.class && output == DetectedObjects.class) {
-            return YoloV5Translator.builder(arguments).build();
+            return (Translator<I, O>) translator;
         } else if (input == Path.class && output == DetectedObjects.class) {
-            return new FileTranslator<>(YoloV5Translator.builder(arguments).build());
+            return (Translator<I, O>) new FileTranslator<>(translator);
         } else if (input == URL.class && output == DetectedObjects.class) {
-            return new UrlTranslator<>(YoloV5Translator.builder(arguments).build());
+            return (Translator<I, O>) new UrlTranslator<>(translator);
         } else if (input == InputStream.class && output == DetectedObjects.class) {
-            return new InputStreamTranslator<>(YoloV5Translator.builder(arguments).build());
+            return (Translator<I, O>) new InputStreamTranslator<>(translator);
         } else if (input == Input.class && output == Output.class) {
-            return new ImageServingTranslator(YoloV5Translator.builder(arguments).build());
+            return (Translator<I, O>) new ImageServingTranslator(translator);
         }
         throw new IllegalArgumentException("Unsupported input/output types.");
     }
diff --git a/api/src/main/java/ai/djl/repository/zoo/BaseModelLoader.java b/api/src/main/java/ai/djl/repository/zoo/BaseModelLoader.java
index 584efefc4..6e15103ca 100644
--- a/api/src/main/java/ai/djl/repository/zoo/BaseModelLoader.java
+++ b/api/src/main/java/ai/djl/repository/zoo/BaseModelLoader.java
@@ -69,7 +69,6 @@ public class BaseModelLoader implements ModelLoader {
 
     /** {@inheritDoc} */
     @Override
-    @SuppressWarnings("unchecked")
     public <I, O> ZooModel<I, O> loadModel(Criteria<I, O> criteria)
             throws IOException, ModelNotFoundException, MalformedModelException {
         Artifact artifact = mrl.match(criteria.getFilters());
@@ -158,8 +157,7 @@ public class BaseModelLoader implements ModelLoader {
                             arguments,
                             engine);
             model.load(modelPath, null, options);
-            Translator<I, O> translator =
-                    (Translator<I, O>) factory.newInstance(input, output, model, arguments);
+            Translator<I, O> translator = factory.newInstance(input, output, model, arguments);
             return new ZooModel<>(model, translator);
         } catch (TranslateException e) {
             throw new ModelNotFoundException("No matching translator found", e);
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 c658ae18c..cbeb45ee9 100644
--- a/api/src/main/java/ai/djl/training/dataset/ArrayDataset.java
+++ b/api/src/main/java/ai/djl/training/dataset/ArrayDataset.java
@@ -103,13 +103,13 @@ public class ArrayDataset extends RandomAccessDataset {
     }
 
     /**
-     * Gets the {@link Record} for the given indices from the dataset.
+     * Gets the {@link Batch} 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
+     * @return a {@link Batch} that contains the data and label of the requested data items
      */
-    public Record getByIndices(NDManager manager, long... indices) {
+    public Batch getByIndices(NDManager manager, long... indices) {
         try (NDArray ndIndices = manager.create(indices)) {
             NDIndex index = new NDIndex("{}", ndIndices);
             NDList datum = new NDList();
@@ -122,19 +122,27 @@ public class ArrayDataset extends RandomAccessDataset {
                     label.add(array.get(manager, index));
                 }
             }
-            return new Record(datum, label);
+            return new Batch(
+                    manager,
+                    datum,
+                    label,
+                    indices.length,
+                    Batchifier.STACK,
+                    Batchifier.STACK,
+                    -1,
+                    -1);
         }
     }
 
     /**
-     * Gets the {@link Record} for the given range from the dataset.
+     * Gets the {@link Batch} 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
+     * @return a {@link Batch} that contains the data and label of the requested data items
      */
-    public Record getByRange(NDManager manager, long fromIndex, long toIndex) {
+    public Batch getByRange(NDManager manager, long fromIndex, long toIndex) {
         NDIndex index = new NDIndex().addSliceDim(fromIndex, toIndex);
         NDList datum = new NDList();
         NDList label = new NDList();
@@ -146,7 +154,8 @@ public class ArrayDataset extends RandomAccessDataset {
                 label.add(array.get(manager, index));
             }
         }
-        return new Record(datum, label);
+        int size = Math.toIntExact(toIndex - fromIndex);
+        return new Batch(manager, datum, label, size, Batchifier.STACK, Batchifier.STACK, -1, -1);
     }
 
     /** {@inheritDoc} */
@@ -277,7 +286,7 @@ public class ArrayDataset extends RandomAccessDataset {
 
         /** {@inheritDoc} */
         @Override
-        public Record getByIndices(NDManager manager, long... indices) {
+        public Batch getByIndices(NDManager manager, long... indices) {
             long[] resolvedIndices = new long[indices.length];
             int i = 0;
             for (long index : indices) {
@@ -288,7 +297,7 @@ public class ArrayDataset extends RandomAccessDataset {
 
         /** {@inheritDoc} */
         @Override
-        public Record getByRange(NDManager manager, long fromIndex, long toIndex) {
+        public Batch getByRange(NDManager manager, long fromIndex, long toIndex) {
             return dataset.getByRange(manager, fromIndex + from, toIndex + from);
         }
 
@@ -330,7 +339,7 @@ public class ArrayDataset extends RandomAccessDataset {
 
         /** {@inheritDoc} */
         @Override
-        public Record getByIndices(NDManager manager, long... indices) {
+        public Batch getByIndices(NDManager manager, long... indices) {
             long[] resolvedIndices = new long[indices.length];
             int i = 0;
             for (long index : indices) {
@@ -341,7 +350,7 @@ public class ArrayDataset extends RandomAccessDataset {
 
         /** {@inheritDoc} */
         @Override
-        public Record getByRange(NDManager manager, long fromIndex, long toIndex) {
+        public Batch getByRange(NDManager manager, long fromIndex, long toIndex) {
             long[] resolvedIndices = new long[(int) (toIndex - fromIndex)];
             int i = 0;
             for (long index = fromIndex; index < toIndex; index++) {
diff --git a/api/src/main/java/ai/djl/training/dataset/BulkDataIterable.java b/api/src/main/java/ai/djl/training/dataset/BulkDataIterable.java
index 521697779..0aa47566d 100644
--- a/api/src/main/java/ai/djl/training/dataset/BulkDataIterable.java
+++ b/api/src/main/java/ai/djl/training/dataset/BulkDataIterable.java
@@ -73,23 +73,23 @@ public class BulkDataIterable extends DataIterable {
         subManager.setName("dataIter fetch");
         int batchSize = indices.size();
 
-        Record record;
+        Batch raw;
         if (isRange(indices)) {
             long fromIndex = indices.get(0);
             long toIndex = fromIndex + indices.size();
-            record = ((ArrayDataset) dataset).getByRange(subManager, fromIndex, toIndex);
+            raw = ((ArrayDataset) dataset).getByRange(subManager, fromIndex, toIndex);
         } else {
             long[] indicesArr = indices.stream().mapToLong(Long::longValue).toArray();
-            record = ((ArrayDataset) dataset).getByIndices(subManager, indicesArr);
+            raw = ((ArrayDataset) dataset).getByIndices(subManager, indicesArr);
         }
 
-        NDList batchData = record.getData();
+        NDList batchData = raw.getData();
         // apply transform
         if (pipeline != null) {
             batchData = pipeline.transform(batchData);
         }
 
-        NDList batchLabels = record.getLabels();
+        NDList batchLabels = raw.getLabels();
 
         // apply label transform
         if (targetPipeline != null) {
diff --git a/api/src/main/java/ai/djl/translate/DefaultTranslatorFactory.java b/api/src/main/java/ai/djl/translate/DefaultTranslatorFactory.java
index 80c103b15..3502ecc78 100644
--- a/api/src/main/java/ai/djl/translate/DefaultTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/translate/DefaultTranslatorFactory.java
@@ -74,18 +74,19 @@ public class DefaultTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments)
             throws TranslateException {
         if (translators != null) {
             Translator<?, ?> translator = translators.get(new Pair<Type, Type>(input, output));
             if (translator != null) {
-                return translator;
+                return (Translator<I, O>) translator;
             }
         }
 
         if (input == NDList.class && output == NDList.class) {
-            return new NoopTranslator();
+            return (Translator<I, O>) new NoopTranslator();
         }
         if (servingTranslatorFactory.isSupported(input, output)) {
             return servingTranslatorFactory.newInstance(input, output, model, arguments);
diff --git a/api/src/main/java/ai/djl/translate/DeferredTranslatorFactory.java b/api/src/main/java/ai/djl/translate/DeferredTranslatorFactory.java
index c500cebf3..5f01d3720 100644
--- a/api/src/main/java/ai/djl/translate/DeferredTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/translate/DeferredTranslatorFactory.java
@@ -51,8 +51,8 @@ public class DeferredTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments)
             throws TranslateException {
         String factoryClass = ArgumentsUtil.stringValue(arguments, "translatorFactory");
         if (factoryClass == null || factoryClass.isEmpty()) {
diff --git a/api/src/main/java/ai/djl/translate/NoopServingTranslatorFactory.java b/api/src/main/java/ai/djl/translate/NoopServingTranslatorFactory.java
index c23f0efea..0935abdd4 100644
--- a/api/src/main/java/ai/djl/translate/NoopServingTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/translate/NoopServingTranslatorFactory.java
@@ -35,14 +35,14 @@ public class NoopServingTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
-            throws TranslateException {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
         if (!isSupported(input, output)) {
             throw new IllegalArgumentException("Unsupported input/output types.");
         }
         String batchifier = ArgumentsUtil.stringValue(arguments, "batchifier", "none");
-        return new NoopServingTranslator(Batchifier.fromString(batchifier));
+        return (Translator<I, O>) new NoopServingTranslator(Batchifier.fromString(batchifier));
     }
 
     static final class NoopServingTranslator implements Translator<Input, Output> {
@@ -77,9 +77,11 @@ public class NoopServingTranslatorFactory implements TranslatorFactory {
         public Output processOutput(TranslatorContext ctx, NDList list) {
             Map<String, String> prop = (Map<String, String>) ctx.getAttachment("properties");
             String contentType = prop.get("Content-Type");
+            String accept = prop.get("Accept");
 
             Output output = new Output();
-            if ("tensor/npz".equalsIgnoreCase(contentType)) {
+            if ("tensor/npz".equalsIgnoreCase(accept)
+                    || "tensor/npz".equalsIgnoreCase(contentType)) {
                 output.add(list.encode(true));
                 output.addProperty("Content-Type", "tensor/npz");
             } else {
diff --git a/api/src/main/java/ai/djl/translate/ServingTranslatorFactory.java b/api/src/main/java/ai/djl/translate/ServingTranslatorFactory.java
index 0755a9688..dfaa514f8 100644
--- a/api/src/main/java/ai/djl/translate/ServingTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/translate/ServingTranslatorFactory.java
@@ -18,8 +18,6 @@ import ai.djl.modality.Input;
 import ai.djl.modality.Output;
 import ai.djl.modality.cv.translator.ImageClassificationTranslator;
 import ai.djl.modality.cv.translator.ImageServingTranslator;
-import ai.djl.ndarray.NDList;
-import ai.djl.ndarray.NDManager;
 import ai.djl.util.ClassLoaderUtils;
 import ai.djl.util.Pair;
 
@@ -62,8 +60,9 @@ public class ServingTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments)
             throws TranslateException {
         if (!isSupported(input, output)) {
             throw new IllegalArgumentException("Unsupported input/output types.");
@@ -87,16 +86,16 @@ public class ServingTranslatorFactory implements TranslatorFactory {
         if (!Files.isDirectory(libPath)) {
             libPath = modelDir.resolve("lib");
             if (!Files.isDirectory(libPath) && className == null) {
-                return loadDefaultTranslator(arguments);
+                return (Translator<I, O>) loadDefaultTranslator(arguments);
             }
         }
         ServingTranslator translator = findTranslator(libPath, className);
         if (translator != null) {
             translator.setArguments(arguments);
             logger.info("Using translator: {}", translator.getClass().getName());
-            return translator;
+            return (Translator<I, O>) translator;
         }
-        return loadDefaultTranslator(arguments);
+        return (Translator<I, O>) loadDefaultTranslator(arguments);
     }
 
     private ServingTranslator findTranslator(Path path, String className) {
@@ -222,8 +221,8 @@ public class ServingTranslatorFactory implements TranslatorFactory {
                 return getImageClassificationTranslator(arguments);
             }
         }
-        String batchifier = ArgumentsUtil.stringValue(arguments, "batchifier", "none");
-        return new RawTranslator(Batchifier.fromString(batchifier));
+        NoopServingTranslatorFactory factory = new NoopServingTranslatorFactory();
+        return factory.newInstance(Input.class, Output.class, null, arguments);
     }
 
     private Translator<Input, Output> getImageClassificationTranslator(Map<String, ?> arguments) {
@@ -251,49 +250,4 @@ public class ServingTranslatorFactory implements TranslatorFactory {
             logger.warn("Failed to compile bundled java file", e);
         }
     }
-
-    private static final class RawTranslator implements Translator<Input, Output> {
-
-        private Batchifier batchifier;
-
-        RawTranslator(Batchifier batchifier) {
-            this.batchifier = batchifier;
-        }
-
-        /** {@inheritDoc} */
-        @Override
-        public Batchifier getBatchifier() {
-            return batchifier;
-        }
-
-        /** {@inheritDoc} */
-        @Override
-        public NDList processInput(TranslatorContext ctx, Input input) throws TranslateException {
-            NDManager manager = ctx.getNDManager();
-            try {
-                ctx.setAttachment("properties", input.getProperties());
-                return input.getDataAsNDList(manager);
-            } catch (IllegalArgumentException e) {
-                throw new TranslateException("Input is not a NDList data type", e);
-            }
-        }
-
-        /** {@inheritDoc} */
-        @Override
-        @SuppressWarnings("unchecked")
-        public Output processOutput(TranslatorContext ctx, NDList list) {
-            Map<String, String> prop = (Map<String, String>) ctx.getAttachment("properties");
-            String contentType = prop.get("Content-Type");
-
-            Output output = new Output();
-            if ("tensor/npz".equalsIgnoreCase(contentType)) {
-                output.add(list.encode(true));
-                output.addProperty("Content-Type", "tensor/npz");
-            } else {
-                output.add(list.encode(false));
-                output.addProperty("Content-Type", "tensor/ndlist");
-            }
-            return output;
-        }
-    }
 }
diff --git a/api/src/main/java/ai/djl/translate/TranslatorFactory.java b/api/src/main/java/ai/djl/translate/TranslatorFactory.java
index fc86106a6..ad17f2f8d 100644
--- a/api/src/main/java/ai/djl/translate/TranslatorFactory.java
+++ b/api/src/main/java/ai/djl/translate/TranslatorFactory.java
@@ -43,6 +43,8 @@ public interface TranslatorFactory {
     /**
      * Returns a new instance of the {@link Translator} class.
      *
+     * @param <I> the input data type
+     * @param <O> the output data type
      * @param input the input class
      * @param output the output class
      * @param model the {@link Model} that uses the {@link Translator}
@@ -50,7 +52,7 @@ public interface TranslatorFactory {
      * @return a new instance of the {@code Translator} class
      * @throws TranslateException if failed to create Translator instance
      */
-    Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
+    <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments)
             throws TranslateException;
 }
diff --git a/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/nlp/embedding/GloveWordEmbeddingTranslatorFactory.java b/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/nlp/embedding/GloveWordEmbeddingTranslatorFactory.java
index c8eda44fd..de9cb4ffc 100644
--- a/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/nlp/embedding/GloveWordEmbeddingTranslatorFactory.java
+++ b/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/nlp/embedding/GloveWordEmbeddingTranslatorFactory.java
@@ -16,7 +16,6 @@ import ai.djl.Model;
 import ai.djl.ndarray.NDList;
 import ai.djl.nn.core.Embedding;
 import ai.djl.translate.ArgumentsUtil;
-import ai.djl.translate.TranslateException;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorContext;
 import ai.djl.translate.TranslatorFactory;
@@ -38,14 +37,14 @@ public class GloveWordEmbeddingTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
-            throws TranslateException {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
         if (!isSupported(input, output)) {
             throw new IllegalArgumentException("Unsupported input/output types.");
         }
         String unknownToken = ArgumentsUtil.stringValue(arguments, "unknownToken");
-        return new GloveWordEmbeddingTranslator(unknownToken);
+        return (Translator<I, O>) new GloveWordEmbeddingTranslator(unknownToken);
     }
 
     private static final class GloveWordEmbeddingTranslator implements Translator<String, NDList> {
diff --git a/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/nlp/qa/MxBertQATranslatorFactory.java b/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/nlp/qa/MxBertQATranslatorFactory.java
index 9a1797f77..002b4fa37 100644
--- a/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/nlp/qa/MxBertQATranslatorFactory.java
+++ b/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/nlp/qa/MxBertQATranslatorFactory.java
@@ -18,7 +18,6 @@ import ai.djl.modality.Output;
 import ai.djl.modality.nlp.qa.QAInput;
 import ai.djl.modality.nlp.translator.QATranslator;
 import ai.djl.modality.nlp.translator.QaServingTranslator;
-import ai.djl.translate.TranslateException;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 import ai.djl.util.Pair;
@@ -46,16 +45,16 @@ public class MxBertQATranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
-            throws TranslateException {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
         if (!isSupported(input, output)) {
             throw new IllegalArgumentException("Unsupported input/output types.");
         }
         QATranslator translator = MxBertQATranslator.builder(arguments).build();
         if (input == Input.class && output == Output.class) {
-            return new QaServingTranslator(translator);
+            return (Translator<I, O>) new QaServingTranslator(translator);
         }
-        return translator;
+        return (Translator<I, O>) translator;
     }
 }
diff --git a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/zoo/tabular/softmax_regression/IrisClassificationTranslatorFactory.java b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/zoo/tabular/softmax_regression/IrisClassificationTranslatorFactory.java
index 6cfa713a9..34281ec42 100644
--- a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/zoo/tabular/softmax_regression/IrisClassificationTranslatorFactory.java
+++ b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/zoo/tabular/softmax_regression/IrisClassificationTranslatorFactory.java
@@ -42,12 +42,13 @@ public class IrisClassificationTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments) {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
         if (!isSupported(input, output)) {
             throw new IllegalArgumentException("Unsupported input/output types.");
         }
-        return new IrisTranslator();
+        return (Translator<I, O>) new IrisTranslator();
     }
 
     private static final class IrisTranslator
diff --git a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/imageclassification/PpImageClassificationTranslatorFactory.java b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/imageclassification/PpImageClassificationTranslatorFactory.java
index c37cf8964..9ce71339a 100644
--- a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/imageclassification/PpImageClassificationTranslatorFactory.java
+++ b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/imageclassification/PpImageClassificationTranslatorFactory.java
@@ -42,8 +42,9 @@ public class PpImageClassificationTranslatorFactory extends ImageClassificationT
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments) {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
         ImageClassificationTranslator translator =
                 ImageClassificationTranslator.builder()
                         .addTransform(new Resize(128, 128))
@@ -55,15 +56,15 @@ public class PpImageClassificationTranslatorFactory extends ImageClassificationT
                         .addTransform(nd -> nd.flip(0)) // RGB -> GBR
                         .build();
         if (input == Image.class && output == Classifications.class) {
-            return translator;
+            return (Translator<I, O>) translator;
         } else if (input == Path.class && output == Classifications.class) {
-            return new FileTranslator<>(translator);
+            return (Translator<I, O>) new FileTranslator<>(translator);
         } else if (input == URL.class && output == Classifications.class) {
-            return new UrlTranslator<>(translator);
+            return (Translator<I, O>) new UrlTranslator<>(translator);
         } else if (input == InputStream.class && output == Classifications.class) {
-            return new InputStreamTranslator<>(translator);
+            return (Translator<I, O>) new InputStreamTranslator<>(translator);
         } else if (input == Input.class && output == Output.class) {
-            return new ImageServingTranslator(translator);
+            return (Translator<I, O>) new ImageServingTranslator(translator);
         }
         throw new IllegalArgumentException("Unsupported input/output types.");
     }
diff --git a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/imageclassification/PpWordRotateTranslatorFactory.java b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/imageclassification/PpWordRotateTranslatorFactory.java
index f242c1904..ca4b83a3b 100644
--- a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/imageclassification/PpWordRotateTranslatorFactory.java
+++ b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/imageclassification/PpWordRotateTranslatorFactory.java
@@ -35,18 +35,19 @@ public class PpWordRotateTranslatorFactory extends ImageClassificationTranslator
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments) {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
         if (input == Image.class && output == Classifications.class) {
-            return new PpWordRotateTranslator();
+            return (Translator<I, O>) new PpWordRotateTranslator();
         } else if (input == Path.class && output == Classifications.class) {
-            return new FileTranslator<>(new PpWordRotateTranslator());
+            return (Translator<I, O>) new FileTranslator<>(new PpWordRotateTranslator());
         } else if (input == URL.class && output == Classifications.class) {
-            return new UrlTranslator<>(new PpWordRotateTranslator());
+            return (Translator<I, O>) new UrlTranslator<>(new PpWordRotateTranslator());
         } else if (input == InputStream.class && output == Classifications.class) {
-            return new InputStreamTranslator<>(new PpWordRotateTranslator());
+            return (Translator<I, O>) new InputStreamTranslator<>(new PpWordRotateTranslator());
         } else if (input == Input.class && output == Output.class) {
-            return new ImageServingTranslator(new PpWordRotateTranslator());
+            return (Translator<I, O>) new ImageServingTranslator(new PpWordRotateTranslator());
         }
         throw new IllegalArgumentException("Unsupported input/output types.");
     }
diff --git a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpFaceDetectionTranslatorFactory.java b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpFaceDetectionTranslatorFactory.java
index 026d9c93f..c55987644 100644
--- a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpFaceDetectionTranslatorFactory.java
+++ b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpFaceDetectionTranslatorFactory.java
@@ -35,18 +35,20 @@ public class PpFaceDetectionTranslatorFactory extends ObjectDetectionTranslatorF
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments) {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
+        PpFaceDetectionTranslator translator = new PpFaceDetectionTranslator(arguments);
         if (input == Image.class && output == DetectedObjects.class) {
-            return new PpFaceDetectionTranslator(arguments);
+            return (Translator<I, O>) translator;
         } else if (input == Path.class && output == DetectedObjects.class) {
-            return new FileTranslator<>(new PpFaceDetectionTranslator(arguments));
+            return (Translator<I, O>) new FileTranslator<>(translator);
         } else if (input == URL.class && output == DetectedObjects.class) {
-            return new UrlTranslator<>(new PpFaceDetectionTranslator(arguments));
+            return (Translator<I, O>) new UrlTranslator<>(translator);
         } else if (input == InputStream.class && output == DetectedObjects.class) {
-            return new InputStreamTranslator<>(new PpFaceDetectionTranslator(arguments));
+            return (Translator<I, O>) new InputStreamTranslator<>(translator);
         } else if (input == Input.class && output == Output.class) {
-            return new ImageServingTranslator(new PpFaceDetectionTranslator(arguments));
+            return (Translator<I, O>) new ImageServingTranslator(translator);
         }
         throw new IllegalArgumentException("Unsupported input/output types.");
     }
diff --git a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpWordDetectionTranslatorFactory.java b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpWordDetectionTranslatorFactory.java
index dfaf52bea..2abbbca0e 100644
--- a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpWordDetectionTranslatorFactory.java
+++ b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpWordDetectionTranslatorFactory.java
@@ -35,18 +35,20 @@ public class PpWordDetectionTranslatorFactory extends ObjectDetectionTranslatorF
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments) {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
+        PpWordDetectionTranslator translator = new PpWordDetectionTranslator(arguments);
         if (input == Image.class && output == DetectedObjects.class) {
-            return new PpWordDetectionTranslator(arguments);
+            return (Translator<I, O>) translator;
         } else if (input == Path.class && output == DetectedObjects.class) {
-            return new FileTranslator<>(new PpWordDetectionTranslator(arguments));
+            return (Translator<I, O>) new FileTranslator<>(translator);
         } else if (input == URL.class && output == DetectedObjects.class) {
-            return new UrlTranslator<>(new PpWordDetectionTranslator(arguments));
+            return (Translator<I, O>) new UrlTranslator<>(translator);
         } else if (input == InputStream.class && output == DetectedObjects.class) {
-            return new InputStreamTranslator<>(new PpWordDetectionTranslator(arguments));
+            return (Translator<I, O>) new InputStreamTranslator<>(translator);
         } else if (input == Input.class && output == Output.class) {
-            return new ImageServingTranslator(new PpWordDetectionTranslator(arguments));
+            return (Translator<I, O>) new ImageServingTranslator(translator);
         }
         throw new IllegalArgumentException("Unsupported input/output types.");
     }
diff --git a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/wordrecognition/PpWordRecognitionTranslatorFactory.java b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/wordrecognition/PpWordRecognitionTranslatorFactory.java
index e26b46344..5223e1385 100644
--- a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/wordrecognition/PpWordRecognitionTranslatorFactory.java
+++ b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/wordrecognition/PpWordRecognitionTranslatorFactory.java
@@ -52,18 +52,20 @@ public class PpWordRecognitionTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments) {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
+        PpWordRecognitionTranslator translator = new PpWordRecognitionTranslator();
         if (input == Image.class && output == String.class) {
-            return new PpWordRecognitionTranslator();
+            return (Translator<I, O>) translator;
         } else if (input == Path.class && output == String.class) {
-            return new FileTranslator<>(new PpWordRecognitionTranslator());
+            return (Translator<I, O>) new FileTranslator<>(translator);
         } else if (input == URL.class && output == String.class) {
-            return new UrlTranslator<>(new PpWordRecognitionTranslator());
+            return (Translator<I, O>) new UrlTranslator<>(translator);
         } else if (input == InputStream.class && output == String.class) {
-            return new InputStreamTranslator<>(new PpWordRecognitionTranslator());
+            return (Translator<I, O>) new InputStreamTranslator<>(translator);
         } else if (input == Input.class && output == Output.class) {
-            return new ImageServingTranslator(new PpWordRecognitionTranslator());
+            return (Translator<I, O>) new ImageServingTranslator(translator);
         }
         throw new IllegalArgumentException("Unsupported input/output types.");
     }
diff --git a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/cv/objectdetection/PtSsdTranslatorFactory.java b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/cv/objectdetection/PtSsdTranslatorFactory.java
index 2bdafe197..c0ac429fa 100644
--- a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/cv/objectdetection/PtSsdTranslatorFactory.java
+++ b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/cv/objectdetection/PtSsdTranslatorFactory.java
@@ -35,18 +35,20 @@ public class PtSsdTranslatorFactory extends ObjectDetectionTranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments) {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
+        PtSsdTranslator translator = PtSsdTranslator.builder(arguments).build();
         if (input == Image.class && output == DetectedObjects.class) {
-            return PtSsdTranslator.builder(arguments).build();
+            return (Translator<I, O>) translator;
         } else if (input == Path.class && output == DetectedObjects.class) {
-            return new FileTranslator<>(PtSsdTranslator.builder(arguments).build());
+            return (Translator<I, O>) new FileTranslator<>(translator);
         } else if (input == URL.class && output == DetectedObjects.class) {
-            return new UrlTranslator<>(PtSsdTranslator.builder(arguments).build());
+            return (Translator<I, O>) new UrlTranslator<>(translator);
         } else if (input == InputStream.class && output == DetectedObjects.class) {
-            return new InputStreamTranslator<>(PtSsdTranslator.builder(arguments).build());
+            return (Translator<I, O>) new InputStreamTranslator<>(translator);
         } else if (input == Input.class && output == Output.class) {
-            return new ImageServingTranslator(PtSsdTranslator.builder(arguments).build());
+            return (Translator<I, O>) new ImageServingTranslator(translator);
         }
         throw new IllegalArgumentException("Unsupported input/output types.");
     }
diff --git a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/qa/PtBertQATranslatorFactory.java b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/qa/PtBertQATranslatorFactory.java
index ea46a6432..75532fc58 100644
--- a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/qa/PtBertQATranslatorFactory.java
+++ b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/qa/PtBertQATranslatorFactory.java
@@ -18,7 +18,6 @@ import ai.djl.modality.Output;
 import ai.djl.modality.nlp.qa.QAInput;
 import ai.djl.modality.nlp.translator.QATranslator;
 import ai.djl.modality.nlp.translator.QaServingTranslator;
-import ai.djl.translate.TranslateException;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 import ai.djl.util.Pair;
@@ -46,16 +45,16 @@ public class PtBertQATranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
-            throws TranslateException {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
         if (!isSupported(input, output)) {
             throw new IllegalArgumentException("Unsupported input/output types.");
         }
         QATranslator translator = PtBertQATranslator.builder(arguments).build();
         if (input == Input.class && output == Output.class) {
-            return new QaServingTranslator(translator);
+            return (Translator<I, O>) new QaServingTranslator(translator);
         }
-        return translator;
+        return (Translator<I, O>) translator;
     }
 }
diff --git a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/sentimentanalysis/PtDistilBertTranslatorFactory.java b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/sentimentanalysis/PtDistilBertTranslatorFactory.java
index 8e4c4761c..634243c76 100644
--- a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/sentimentanalysis/PtDistilBertTranslatorFactory.java
+++ b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/sentimentanalysis/PtDistilBertTranslatorFactory.java
@@ -17,7 +17,6 @@ import ai.djl.modality.Classifications;
 import ai.djl.modality.Input;
 import ai.djl.modality.Output;
 import ai.djl.modality.nlp.translator.TextClassificationServingTranslator;
-import ai.djl.translate.TranslateException;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 import ai.djl.util.Pair;
@@ -45,16 +44,16 @@ public class PtDistilBertTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
-            throws TranslateException {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
         if (!isSupported(input, output)) {
             throw new IllegalArgumentException("Unsupported input/output types.");
         }
         Translator<String, Classifications> translator = new PtDistilBertTranslator();
         if (input == Input.class && output == Output.class) {
-            return new TextClassificationServingTranslator(translator);
+            return (Translator<I, O>) new TextClassificationServingTranslator(translator);
         }
-        return translator;
+        return (Translator<I, O>) translator;
     }
 }
diff --git a/engines/tensorflow/tensorflow-model-zoo/src/main/java/ai/djl/tensorflow/zoo/cv/objectdetction/TfSsdTranslatorFactory.java b/engines/tensorflow/tensorflow-model-zoo/src/main/java/ai/djl/tensorflow/zoo/cv/objectdetction/TfSsdTranslatorFactory.java
index aa54457d0..2b0b7e52e 100644
--- a/engines/tensorflow/tensorflow-model-zoo/src/main/java/ai/djl/tensorflow/zoo/cv/objectdetction/TfSsdTranslatorFactory.java
+++ b/engines/tensorflow/tensorflow-model-zoo/src/main/java/ai/djl/tensorflow/zoo/cv/objectdetction/TfSsdTranslatorFactory.java
@@ -35,18 +35,20 @@ public class TfSsdTranslatorFactory extends ObjectDetectionTranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments) {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
+        TfSsdTranslator translator = TfSsdTranslator.builder(arguments).build();
         if (input == Image.class && output == DetectedObjects.class) {
-            return TfSsdTranslator.builder(arguments).build();
+            return (Translator<I, O>) translator;
         } else if (input == Path.class && output == DetectedObjects.class) {
-            return new FileTranslator<>(TfSsdTranslator.builder(arguments).build());
+            return (Translator<I, O>) new FileTranslator<>(translator);
         } else if (input == URL.class && output == DetectedObjects.class) {
-            return new UrlTranslator<>(TfSsdTranslator.builder(arguments).build());
+            return (Translator<I, O>) new UrlTranslator<>(translator);
         } else if (input == InputStream.class && output == DetectedObjects.class) {
-            return new InputStreamTranslator<>(TfSsdTranslator.builder(arguments).build());
+            return (Translator<I, O>) new InputStreamTranslator<>(translator);
         } else if (input == Input.class && output == Output.class) {
-            return new ImageServingTranslator(TfSsdTranslator.builder(arguments).build());
+            return (Translator<I, O>) new ImageServingTranslator(translator);
         }
         throw new IllegalArgumentException("Unsupported input/output types.");
     }
diff --git a/extensions/timeseries/src/main/java/ai/djl/timeseries/translator/DeepARTranslatorFactory.java b/extensions/timeseries/src/main/java/ai/djl/timeseries/translator/DeepARTranslatorFactory.java
index 09d19fdf0..5eacba22e 100644
--- a/extensions/timeseries/src/main/java/ai/djl/timeseries/translator/DeepARTranslatorFactory.java
+++ b/extensions/timeseries/src/main/java/ai/djl/timeseries/translator/DeepARTranslatorFactory.java
@@ -15,7 +15,6 @@ package ai.djl.timeseries.translator;
 import ai.djl.Model;
 import ai.djl.timeseries.Forecast;
 import ai.djl.timeseries.TimeSeriesData;
-import ai.djl.translate.TranslateException;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 import ai.djl.util.Pair;
@@ -36,12 +35,12 @@ public class DeepARTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
-            throws TranslateException {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
         if (!isSupported(input, output)) {
             throw new IllegalArgumentException("Unsupported input/output types.");
         }
-        return DeepARTranslator.builder(arguments).build();
+        return (Translator<I, O>) DeepARTranslator.builder(arguments).build();
     }
 }
diff --git a/extensions/timeseries/src/main/java/ai/djl/timeseries/translator/TransformerTranslatorFactory.java b/extensions/timeseries/src/main/java/ai/djl/timeseries/translator/TransformerTranslatorFactory.java
index e80a8f59b..9ee1b4280 100644
--- a/extensions/timeseries/src/main/java/ai/djl/timeseries/translator/TransformerTranslatorFactory.java
+++ b/extensions/timeseries/src/main/java/ai/djl/timeseries/translator/TransformerTranslatorFactory.java
@@ -15,7 +15,6 @@ package ai.djl.timeseries.translator;
 import ai.djl.Model;
 import ai.djl.timeseries.Forecast;
 import ai.djl.timeseries.TimeSeriesData;
-import ai.djl.translate.TranslateException;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 import ai.djl.util.Pair;
@@ -36,12 +35,12 @@ public class TransformerTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
-            throws TranslateException {
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
         if (!isSupported(input, output)) {
             throw new IllegalArgumentException("Unsupported input/output types.");
         }
-        return TransformerTranslator.builder(arguments).build();
+        return (Translator<I, O>) TransformerTranslator.builder(arguments).build();
     }
 }
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslatorFactory.java
index 71daf9488..5c23f3114 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslatorFactory.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslatorFactory.java
@@ -48,8 +48,9 @@ public class FillMaskTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments)
             throws TranslateException {
         Path modelPath = model.getModelPath();
         try {
@@ -61,9 +62,9 @@ public class FillMaskTranslatorFactory implements TranslatorFactory {
             FillMaskTranslator translator =
                     FillMaskTranslator.builder(tokenizer, arguments).build();
             if (input == String.class && output == Classifications.class) {
-                return translator;
+                return (Translator<I, O>) translator;
             } else if (input == Input.class && output == Output.class) {
-                return new TextClassificationServingTranslator(translator);
+                return (Translator<I, O>) new TextClassificationServingTranslator(translator);
             }
             throw new IllegalArgumentException("Unsupported input/output types.");
         } catch (IOException e) {
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslatorFactory.java
index 87fbe5a20..2f0e1fbbb 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslatorFactory.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslatorFactory.java
@@ -48,8 +48,9 @@ public class QuestionAnsweringTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments)
             throws TranslateException {
         Path modelPath = model.getModelPath();
         try {
@@ -61,9 +62,9 @@ public class QuestionAnsweringTranslatorFactory implements TranslatorFactory {
             QuestionAnsweringTranslator translator =
                     QuestionAnsweringTranslator.builder(tokenizer, arguments).build();
             if (input == QAInput.class && output == String.class) {
-                return translator;
+                return (Translator<I, O>) translator;
             } else if (input == Input.class && output == Output.class) {
-                return new QaServingTranslator(translator);
+                return (Translator<I, O>) new QaServingTranslator(translator);
             }
             throw new IllegalArgumentException("Unsupported input/output types.");
         } catch (IOException e) {
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslatorFactory.java
index a20031cac..a66a3a745 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslatorFactory.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslatorFactory.java
@@ -48,8 +48,9 @@ public class TextClassificationTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments)
             throws TranslateException {
         Path modelPath = model.getModelPath();
         try {
@@ -61,9 +62,9 @@ public class TextClassificationTranslatorFactory implements TranslatorFactory {
             TextClassificationTranslator translator =
                     TextClassificationTranslator.builder(tokenizer, arguments).build();
             if (input == String.class && output == Classifications.class) {
-                return translator;
+                return (Translator<I, O>) translator;
             } else if (input == Input.class && output == Output.class) {
-                return new TextClassificationServingTranslator(translator);
+                return (Translator<I, O>) new TextClassificationServingTranslator(translator);
             }
             throw new IllegalArgumentException("Unsupported input/output types.");
         } catch (IOException e) {
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java
index 87d692abe..43bc4f03b 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java
@@ -47,8 +47,9 @@ public class TextEmbeddingTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments)
             throws TranslateException {
         Path modelPath = model.getModelPath();
         try {
@@ -60,9 +61,9 @@ public class TextEmbeddingTranslatorFactory implements TranslatorFactory {
             TextEmbeddingTranslator translator =
                     TextEmbeddingTranslator.builder(tokenizer, arguments).build();
             if (input == String.class && output == float[].class) {
-                return translator;
+                return (Translator<I, O>) translator;
             } else if (input == Input.class && output == Output.class) {
-                return new TextEmbeddingServingTranslator(translator);
+                return (Translator<I, O>) new TextEmbeddingServingTranslator(translator);
             }
             throw new IllegalArgumentException("Unsupported input/output types.");
         } catch (IOException e) {
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslatorFactory.java
index e590bf7f7..efe14f391 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslatorFactory.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslatorFactory.java
@@ -48,8 +48,9 @@ public class TokenClassificationTranslatorFactory implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    public Translator<?, ?> newInstance(
-            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments)
             throws TranslateException {
         Path modelPath = model.getModelPath();
         try {
@@ -61,9 +62,9 @@ public class TokenClassificationTranslatorFactory implements TranslatorFactory {
             TokenClassificationTranslator translator =
                     TokenClassificationTranslator.builder(tokenizer, arguments).build();
             if (input == String.class && output == NamedEntity[].class) {
-                return translator;
+                return (Translator<I, O>) translator;
             } else if (input == Input.class && output == Output.class) {
-                return new TokenClassificationServingTranslator(translator);
+                return (Translator<I, O>) new TokenClassificationServingTranslator(translator);
             }
             throw new IllegalArgumentException("Unsupported input/output types.");
         } catch (IOException e) {
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/zoo/HfModelZoo.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/zoo/HfModelZoo.java
new file mode 100644
index 000000000..15fed4572
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/zoo/HfModelZoo.java
@@ -0,0 +1,136 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.huggingface.zoo;
+
+import ai.djl.Application;
+import ai.djl.Application.NLP;
+import ai.djl.engine.Engine;
+import ai.djl.repository.Repository;
+import ai.djl.repository.Version;
+import ai.djl.repository.VersionRange;
+import ai.djl.repository.zoo.ModelZoo;
+import ai.djl.util.JsonUtils;
+import ai.djl.util.Utils;
+
+import com.google.gson.reflect.TypeToken;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.io.Reader;
+import java.io.Writer;
+import java.lang.reflect.Type;
+import java.net.URL;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.time.Duration;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Set;
+import java.util.zip.GZIPInputStream;
+
+/** HfModelZoo is a repository that contains HuggingFace models. */
+public class HfModelZoo extends ModelZoo {
+
+    private static final Logger logger = LoggerFactory.getLogger(HfModelZoo.class);
+
+    private static final String REPO = "https://mlrepo.djl.ai/";
+    private static final Repository REPOSITORY = Repository.newInstance("Huggingface", REPO);
+    private static final String GROUP_ID = "ai.djl.huggingface.pytorch";
+
+    private static final long ONE_DAY = Duration.ofDays(1).toMillis();
+
+    HfModelZoo() {
+        Version version = new Version(Engine.class.getPackage().getSpecificationVersion());
+        addModels(NLP.FILL_MASK, version);
+        addModels(NLP.QUESTION_ANSWER, version);
+        addModels(NLP.TEXT_CLASSIFICATION, version);
+        addModels(NLP.TEXT_EMBEDDING, version);
+        addModels(NLP.TOKEN_CLASSIFICATION, version);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public String getGroupId() {
+        return GROUP_ID;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Set<String> getSupportedEngines() {
+        return Collections.singleton("PyTorch");
+    }
+
+    private void addModels(Application app, Version version) {
+        Map<String, Map<String, Object>> map = listModels(app);
+        for (Map.Entry<String, Map<String, Object>> entry : map.entrySet()) {
+            Map<String, Object> model = entry.getValue();
+            if ("failed".equals(model.get("result"))) {
+                continue;
+            }
+            String requires = (String) model.get("requires");
+            if (requires != null) {
+                // the model requires specific DJL version
+                VersionRange range = VersionRange.parse(requires);
+                if (!range.contains(version)) {
+                    continue;
+                }
+            }
+            String artifactId = entry.getKey();
+            addModel(REPOSITORY.model(app, GROUP_ID, artifactId, "0.0.1"));
+        }
+    }
+
+    private Map<String, Map<String, Object>> listModels(Application app) {
+        try {
+            String path = "model/" + app.getPath() + "/ai/djl/huggingface/pytorch/";
+            Path dir = Utils.getCacheDir().resolve("cache/repo/" + path);
+            if (Files.notExists(dir)) {
+                Files.createDirectories(dir);
+            } else if (!Files.isDirectory(dir)) {
+                logger.warn("Failed initialize cache directory: " + dir);
+                return Collections.emptyMap();
+            }
+            Type type = new TypeToken<Map<String, Map<String, Object>>>() {}.getType();
+
+            Path file = dir.resolve("models.json");
+            if (Files.exists(file)) {
+                long lastModified = Files.getLastModifiedTime(file).toMillis();
+                if (Boolean.getBoolean("offline")
+                        || System.currentTimeMillis() - lastModified < ONE_DAY) {
+                    try (Reader reader = Files.newBufferedReader(file)) {
+                        return JsonUtils.GSON.fromJson(reader, type);
+                    }
+                }
+            }
+
+            String url = REPO + path + "models.json.gz";
+            Path tmp = Files.createTempFile(dir, "models", ".tmp");
+            try (GZIPInputStream gis = new GZIPInputStream(new URL(url).openStream())) {
+                String json = Utils.toString(gis);
+                try (Writer writer = Files.newBufferedWriter(tmp)) {
+                    writer.write(json);
+                }
+                Utils.moveQuietly(tmp, file);
+                return JsonUtils.GSON.fromJson(json, type);
+            } finally {
+                Utils.deleteQuietly(tmp);
+            }
+        } catch (IOException e) {
+            logger.warn("Failed load index of models: " + app, e);
+        }
+
+        return Collections.emptyMap();
+    }
+}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/zoo/HfZooProvider.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/zoo/HfZooProvider.java
new file mode 100644
index 000000000..1f31f714e
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/zoo/HfZooProvider.java
@@ -0,0 +1,29 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.huggingface.zoo;
+
+import ai.djl.repository.zoo.ModelZoo;
+import ai.djl.repository.zoo.ZooProvider;
+
+/**
+ * An Huggingface model zoo provider implements the {@link ai.djl.repository.zoo.ZooProvider}
+ * interface.
+ */
+public class HfZooProvider implements ZooProvider {
+
+    /** {@inheritDoc} */
+    @Override
+    public ModelZoo getModelZoo() {
+        return new HfModelZoo();
+    }
+}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/zoo/package-info.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/zoo/package-info.java
new file mode 100644
index 000000000..20149d707
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/zoo/package-info.java
@@ -0,0 +1,14 @@
+/*
+ * 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 the built-in {@link ai.djl.huggingface.zoo.HfModelZoo}. */
+package ai.djl.huggingface.zoo;
diff --git a/extensions/tokenizers/src/main/resources/.empty b/extensions/tokenizers/src/main/resources/.empty
deleted file mode 100644
index e69de29bb..000000000
diff --git a/extensions/tokenizers/src/main/resources/META-INF/services/ai.djl.repository.zoo.ZooProvider b/extensions/tokenizers/src/main/resources/META-INF/services/ai.djl.repository.zoo.ZooProvider
new file mode 100644
index 000000000..956e00dcf
--- /dev/null
+++ b/extensions/tokenizers/src/main/resources/META-INF/services/ai.djl.repository.zoo.ZooProvider
@@ -0,0 +1,1 @@
+ai.djl.huggingface.zoo.HfZooProvider


Test patch:
diff --git a/api/src/test/java/ai/djl/translate/NoopServingTranslatorFactoryTest.java b/api/src/test/java/ai/djl/translate/NoopServingTranslatorFactoryTest.java
index 69c11b368..6a3496f81 100644
--- a/api/src/test/java/ai/djl/translate/NoopServingTranslatorFactoryTest.java
+++ b/api/src/test/java/ai/djl/translate/NoopServingTranslatorFactoryTest.java
@@ -29,7 +29,7 @@ import java.nio.file.Paths;
 public class NoopServingTranslatorFactoryTest {
 
     @Test
-    public void testRawTranslatorFactory() throws ModelException, IOException {
+    public void testNoopTranslatorFactory() throws ModelException, IOException {
         NoopServingTranslatorFactory factory = new NoopServingTranslatorFactory();
         Assert.assertEquals(factory.getSupportedTypes().size(), 1);
 
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/zoo/ModelZooTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/zoo/ModelZooTest.java
new file mode 100644
index 000000000..ea9a5ca1b
--- /dev/null
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/zoo/ModelZooTest.java
@@ -0,0 +1,92 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.huggingface.zoo;
+
+import ai.djl.Application.NLP;
+import ai.djl.ModelException;
+import ai.djl.inference.Predictor;
+import ai.djl.modality.nlp.qa.QAInput;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.ZooModel;
+import ai.djl.testing.TestRequirements;
+import ai.djl.translate.TranslateException;
+import ai.djl.util.JsonUtils;
+import ai.djl.util.Utils;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+import java.io.Writer;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+public class ModelZooTest {
+
+    @Test
+    public void testModelZoo() throws ModelException, IOException, TranslateException {
+        TestRequirements.nightly();
+
+        String question = "When did BBC Japan start broadcasting?";
+        String paragraph =
+                "BBC Japan was a general entertainment Channel. "
+                        + "Which operated between December 2004 and April 2006. "
+                        + "It ceased operations after its Japanese distributor folded.";
+
+        String url = "djl://ai.djl.huggingface.pytorch/deepset/minilm-uncased-squad2";
+        Criteria<QAInput, String> criteria =
+                Criteria.builder().setTypes(QAInput.class, String.class).optModelUrls(url).build();
+
+        try (ZooModel<QAInput, String> model = criteria.loadModel();
+                Predictor<QAInput, String> predictor = model.newPredictor()) {
+            QAInput input = new QAInput(question, paragraph);
+            String res = predictor.predict(input);
+            Assert.assertEquals(res, "december 2004");
+        }
+    }
+
+    @Test
+    public void testFutureVersion() throws IOException {
+        Map<String, Map<String, Object>> map = new ConcurrentHashMap<>();
+        Map<String, Object> model = new ConcurrentHashMap<>();
+        model.put("result", "failed");
+        map.put("model1", model);
+
+        model = new ConcurrentHashMap<>();
+        model.put("requires", "10.100.0+");
+        map.put("model2", model);
+
+        model = new ConcurrentHashMap<>();
+        model.put("requires", "0.19.0+");
+        map.put("model3", model);
+        map.put("model4", new ConcurrentHashMap<>());
+
+        String path = "model/" + NLP.QUESTION_ANSWER.getPath() + "/ai/djl/huggingface/pytorch/";
+        Path dir = Utils.getCacheDir().resolve("cache/repo/" + path);
+        Files.createDirectories(dir);
+        Path file = dir.resolve("models.json");
+        try (Writer writer = Files.newBufferedWriter(file)) {
+            writer.write(JsonUtils.GSON_PRETTY.toJson(map));
+        }
+        HfModelZoo zoo = new HfModelZoo();
+
+        Assert.assertNull(zoo.getModelLoader("model1"));
+        Assert.assertNull(zoo.getModelLoader("model2"));
+        Assert.assertNull(zoo.getModelLoader("model3"));
+        Assert.assertNotNull(zoo.getModelLoader("model4"));
+
+        Utils.deleteQuietly(file);
+    }
+}
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/zoo/package-info.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/zoo/package-info.java
new file mode 100644
index 000000000..85d08046c
--- /dev/null
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/zoo/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.huggingface.zoo;