Repository: deepjavalibrary/djl
Pull request #3140 (2024-04-29 18:32:30)

[api] Adds OnesBlockFactory to make it easy for testing
## Description ##

Brief description of what this PR is about

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



Patch:
diff --git a/.gitignore b/.gitignore
index af083e92a..503d4c8c0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -15,6 +15,15 @@ libs
 *.dll
 *.class
 *.log
+*.onnx
+*.pt
+*.safetensors
+*.pb
+*.trt
+*.uff
+*.gguf
+*.xgb
+*.bin
 
 # Eclipse
 .settings
diff --git a/api/src/main/java/ai/djl/Application.java b/api/src/main/java/ai/djl/Application.java
index a0352e0d3..6fce7f0fe 100644
--- a/api/src/main/java/ai/djl/Application.java
+++ b/api/src/main/java/ai/djl/Application.java
@@ -51,44 +51,64 @@ public class Application {
             case "cv":
                 return CV.ANY;
             case "cv/image_classification":
+            case "image_classification":
                 return CV.IMAGE_CLASSIFICATION;
             case "cv/object_detection":
+            case "object_detection":
                 return CV.OBJECT_DETECTION;
             case "cv/semantic_segmentation":
+            case "semantic_segmentation":
                 return CV.SEMANTIC_SEGMENTATION;
             case "cv/instance_segmentation":
+            case "instance_segmentation":
                 return CV.INSTANCE_SEGMENTATION;
             case "cv/pose_estimation":
+            case "pose_estimation":
                 return CV.POSE_ESTIMATION;
             case "cv/action_recognition":
+            case "action_recognition":
                 return CV.ACTION_RECOGNITION;
             case "cv/word_recognition":
+            case "word_recognition":
                 return CV.WORD_RECOGNITION;
             case "cv/image_generation":
+            case "image_generation":
                 return CV.IMAGE_GENERATION;
             case "cv/image_enhancement":
+            case "image_enhancement":
                 return CV.IMAGE_ENHANCEMENT;
             case "nlp":
                 return NLP.ANY;
             case "nlp/fill_mask":
+            case "fill_mask":
                 return NLP.FILL_MASK;
             case "nlp/machine_translation":
+            case "machine_translation":
                 return NLP.MACHINE_TRANSLATION;
             case "nlp/multiple_choice":
+            case "multiple_choice":
                 return NLP.MULTIPLE_CHOICE;
             case "nlp/question_answer":
+            case "question_answering":
                 return NLP.QUESTION_ANSWER;
             case "nlp/sentiment_analysis":
+            case "sentiment_analysis":
                 return NLP.SENTIMENT_ANALYSIS;
             case "nlp/text_classification":
+            case "text_classification":
                 return NLP.TEXT_CLASSIFICATION;
             case "nlp/text_embedding":
+            case "text_embedding":
+            case "sentence_similarity":
                 return NLP.TEXT_EMBEDDING;
             case "nlp/token_classification":
+            case "token_classification":
                 return NLP.TOKEN_CLASSIFICATION;
             case "nlp/word_embedding":
+            case "word_embedding":
                 return NLP.WORD_EMBEDDING;
             case "nlp/text_generation":
+            case "text_generation":
                 return NLP.TEXT_GENERATION;
             case "tabular":
                 return Tabular.ANY;
diff --git a/api/src/main/java/ai/djl/ndarray/types/Shape.java b/api/src/main/java/ai/djl/ndarray/types/Shape.java
index e7751f7d5..51b4d2ad6 100644
--- a/api/src/main/java/ai/djl/ndarray/types/Shape.java
+++ b/api/src/main/java/ai/djl/ndarray/types/Shape.java
@@ -23,6 +23,8 @@ import java.util.Arrays;
 import java.util.List;
 import java.util.function.Function;
 import java.util.function.Predicate;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import java.util.stream.LongStream;
 import java.util.stream.Stream;
@@ -535,4 +537,64 @@ public class Shape {
         }
         return max == ans;
     }
+
+    /**
+     * Parses a string representation of shapes for NDList.
+     *
+     * @param value a string representation of shapes for NDList
+     * @return a list of Shape and datatype pairs
+     */
+    public static PairList<DataType, Shape> parseShapes(String value) {
+        PairList<DataType, Shape> inputShapes = new PairList<>();
+        if (value != null) {
+            if (value.contains("(")) {
+                Pattern pattern = Pattern.compile("\\((\\s*(\\d+)([,\\s]+\\d+)*\\s*)\\)(\\w?)");
+                Matcher matcher = pattern.matcher(value);
+                while (matcher.find()) {
+                    String[] tokens = matcher.group(1).split(",");
+                    long[] array = Arrays.stream(tokens).mapToLong(Long::parseLong).toArray();
+                    DataType dataType;
+                    String dataTypeStr = matcher.group(4);
+                    if (dataTypeStr == null || dataTypeStr.isEmpty()) {
+                        dataType = DataType.FLOAT32;
+                    } else {
+                        switch (dataTypeStr) {
+                            case "s":
+                                dataType = DataType.FLOAT16;
+                                break;
+                            case "d":
+                                dataType = DataType.FLOAT64;
+                                break;
+                            case "u":
+                                dataType = DataType.UINT8;
+                                break;
+                            case "b":
+                                dataType = DataType.INT8;
+                                break;
+                            case "i":
+                                dataType = DataType.INT32;
+                                break;
+                            case "l":
+                                dataType = DataType.INT64;
+                                break;
+                            case "B":
+                                dataType = DataType.BOOLEAN;
+                                break;
+                            case "f":
+                                dataType = DataType.FLOAT32;
+                                break;
+                            default:
+                                throw new IllegalArgumentException("Invalid input-shape: " + value);
+                        }
+                    }
+                    inputShapes.add(dataType, new Shape(array));
+                }
+            } else {
+                String[] tokens = value.split(",");
+                long[] shapes = Arrays.stream(tokens).mapToLong(Long::parseLong).toArray();
+                inputShapes.add(DataType.FLOAT32, new Shape(shapes));
+            }
+        }
+        return inputShapes;
+    }
 }
diff --git a/api/src/main/java/ai/djl/nn/Blocks.java b/api/src/main/java/ai/djl/nn/Blocks.java
index 47cd58435..136c9bd0f 100644
--- a/api/src/main/java/ai/djl/nn/Blocks.java
+++ b/api/src/main/java/ai/djl/nn/Blocks.java
@@ -14,8 +14,11 @@ package ai.djl.nn;
 
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.util.Pair;
+import ai.djl.util.PairList;
 
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -87,6 +90,31 @@ public final class Blocks {
         return new LambdaBlock(x -> x, "identity");
     }
 
+    /**
+     * Creates a {@link LambdaBlock} that return all-ones NDList.
+     *
+     * @return an all-ones {@link Block}
+     */
+    public static Block onesBlock(PairList<DataType, Shape> shapes, String[] names) {
+        return new LambdaBlock(
+                a -> {
+                    NDManager manager = a.getManager();
+                    NDList list = new NDList(shapes.size());
+                    int index = 0;
+                    for (Pair<DataType, Shape> pair : shapes) {
+                        DataType dataType = pair.getKey();
+                        Shape shape = pair.getValue();
+                        NDArray arr = manager.ones(shape, dataType);
+                        if (names.length == list.size()) {
+                            arr.setName(names[index++]);
+                        }
+                        list.add(arr);
+                    }
+                    return list;
+                },
+                "ones");
+    }
+
     /**
      * Returns a string representation of the passed {@link Block} describing the input axes, output
      * axes, and the block's children.
diff --git a/api/src/main/java/ai/djl/nn/OnesBlockFactory.java b/api/src/main/java/ai/djl/nn/OnesBlockFactory.java
new file mode 100644
index 000000000..55f17269f
--- /dev/null
+++ b/api/src/main/java/ai/djl/nn/OnesBlockFactory.java
@@ -0,0 +1,45 @@
+/*
+ * Copyright 2024 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.nn;
+
+import ai.djl.Model;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.translate.ArgumentsUtil;
+import ai.djl.util.PairList;
+import ai.djl.util.Utils;
+
+import java.nio.file.Path;
+import java.util.Map;
+
+/** A {@link BlockFactory} class that creates LambdaBlock. */
+public class OnesBlockFactory implements BlockFactory {
+
+    private static final long serialVersionUID = 1L;
+
+    /** {@inheritDoc} */
+    @Override
+    public Block newBlock(Model model, Path modelPath, Map<String, ?> arguments) {
+        String shapes = ArgumentsUtil.stringValue(arguments, "block_shapes");
+        String blockNames = ArgumentsUtil.stringValue(arguments, "block_names");
+        PairList<DataType, Shape> pairs = Shape.parseShapes(shapes);
+        String[] names;
+        if (blockNames != null) {
+            names = blockNames.split(",");
+        } else {
+            names = Utils.EMPTY_ARRAY;
+        }
+
+        return Blocks.onesBlock(pairs, names);
+    }
+}
diff --git a/api/src/main/java/ai/djl/translate/ServingTranslatorFactory.java b/api/src/main/java/ai/djl/translate/ServingTranslatorFactory.java
index 3ec161488..32680b084 100644
--- a/api/src/main/java/ai/djl/translate/ServingTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/translate/ServingTranslatorFactory.java
@@ -16,10 +16,9 @@ import ai.djl.Application;
 import ai.djl.Model;
 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.util.ClassLoaderUtils;
 import ai.djl.util.Pair;
+import ai.djl.util.Utils;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -29,6 +28,7 @@ import java.lang.reflect.Type;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Collections;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
 
@@ -55,32 +55,33 @@ public class ServingTranslatorFactory implements TranslatorFactory {
 
         Path modelDir = model.getModelPath();
         String factoryClass = ArgumentsUtil.stringValue(arguments, "translatorFactory");
-        if (factoryClass != null && !factoryClass.isEmpty()) {
-            TranslatorFactory factory = loadTranslatorFactory(factoryClass);
-            if (factory != null && factory.isSupported(input, output)) {
-                logger.info("Using TranslatorFactory: {}", factory.getClass().getName());
-                return factory.newInstance(input, output, model, arguments);
+        if (factoryClass != null) {
+            Translator<Input, Output> translator =
+                    getServingTranslator(factoryClass, model, arguments);
+            if (translator != null) {
+                return (Translator<I, O>) translator;
             }
+            throw new TranslateException("Failed to load translatorFactory: " + factoryClass);
         }
 
         String className = (String) arguments.get("translator");
-
         Path libPath = modelDir.resolve("libs");
         if (!Files.isDirectory(libPath)) {
             libPath = modelDir.resolve("lib");
             if (!Files.isDirectory(libPath) && className == null) {
-                return (Translator<I, O>) loadDefaultTranslator(arguments);
+                return (Translator<I, O>) loadDefaultTranslator(model, arguments);
             }
         }
-        ServingTranslator translator = findTranslator(libPath, className);
-        if (translator != null) {
-            translator.setArguments(arguments);
-            logger.info("Using translator: {}", translator.getClass().getName());
-            return (Translator<I, O>) translator;
+        ServingTranslator servingTranslator = findTranslator(libPath, className);
+        if (servingTranslator != null) {
+            servingTranslator.setArguments(arguments);
+            logger.info("Using translator: {}", servingTranslator.getClass().getName());
+            return (Translator<I, O>) servingTranslator;
         } else if (className != null) {
             throw new TranslateException("Failed to load translator: " + className);
         }
-        return (Translator<I, O>) loadDefaultTranslator(arguments);
+
+        return (Translator<I, O>) loadDefaultTranslator(model, arguments);
     }
 
     private ServingTranslator findTranslator(Path path, String className) {
@@ -101,19 +102,56 @@ public class ServingTranslatorFactory implements TranslatorFactory {
         return null;
     }
 
-    private Translator<Input, Output> loadDefaultTranslator(Map<String, ?> arguments) {
-        String appName = ArgumentsUtil.stringValue(arguments, "application");
-        if (appName != null) {
-            Application application = Application.of(appName);
-            if (application == Application.CV.IMAGE_CLASSIFICATION) {
-                return getImageClassificationTranslator(arguments);
-            }
+    private Translator<Input, Output> loadDefaultTranslator(Model model, Map<String, ?> arguments)
+            throws TranslateException {
+        String factoryClass = detectTranslatorFactory(arguments);
+        Translator<Input, Output> translator = getServingTranslator(factoryClass, model, arguments);
+        if (translator != null) {
+            return translator;
         }
+
         NoopServingTranslatorFactory factory = new NoopServingTranslatorFactory();
         return factory.newInstance(Input.class, Output.class, null, arguments);
     }
 
-    private Translator<Input, Output> getImageClassificationTranslator(Map<String, ?> arguments) {
-        return new ImageServingTranslator(ImageClassificationTranslator.builder(arguments).build());
+    private String detectTranslatorFactory(Map<String, ?> arguments) {
+        Application application;
+        String app = ArgumentsUtil.stringValue(arguments, "application");
+        if (app != null) {
+            application = Application.of(app);
+        } else {
+            String task = Utils.getEnvOrSystemProperty("HF_TASK");
+            task = ArgumentsUtil.stringValue(arguments, "task", task);
+            if (task != null) {
+                task = task.replace("-", "_").toLowerCase(Locale.ROOT);
+                application = Application.of(task);
+            } else {
+                application = Application.UNDEFINED;
+            }
+        }
+        if (application == Application.CV.IMAGE_CLASSIFICATION) {
+            return "ai.djl.modality.cv.translator.ImageClassificationTranslatorFactory";
+        } else if (application == Application.NLP.FILL_MASK) {
+            return "ai.djl.huggingface.translator.FillMaskTranslatorFactory";
+        } else if (application == Application.NLP.QUESTION_ANSWER) {
+            return "ai.djl.huggingface.translator.QuestionAnsweringTranslatorFactory";
+        } else if (application == Application.NLP.TEXT_CLASSIFICATION) {
+            return "ai.djl.huggingface.translator.TextClassificationTranslatorFactory";
+        } else if (application == Application.NLP.TEXT_EMBEDDING) {
+            return "ai.djl.huggingface.translator.TextEmbeddingTranslatorFactory";
+        } else if (application == Application.NLP.TOKEN_CLASSIFICATION) {
+            return "ai.djl.huggingface.translator.TokenClassificationTranslatorFactory";
+        }
+        return null;
+    }
+
+    private Translator<Input, Output> getServingTranslator(
+            String factoryClass, Model model, Map<String, ?> arguments) throws TranslateException {
+        TranslatorFactory factory = loadTranslatorFactory(factoryClass);
+        if (factory != null && factory.isSupported(Input.class, Output.class)) {
+            logger.info("Using TranslatorFactory: {}", factoryClass);
+            return factory.newInstance(Input.class, Output.class, model, arguments);
+        }
+        return null;
     }
 }
diff --git a/docs/development/development_guideline.md b/docs/development/development_guideline.md
index 55ddb3a83..b6a6201ba 100644
--- a/docs/development/development_guideline.md
+++ b/docs/development/development_guideline.md
@@ -134,7 +134,7 @@ Before you run any examples in IntelliJ, configure your Application template as
 3. Change the "Working directory:" value to: "$MODULE_WORKING_DIR$".
 4. Select "OK" to save the template.
 
-Navigate to the 'examples' module. Open the class that you want to execute (e.g. ai.djl.examples.inference.ObjectDetection).
+Navigate to the 'examples' module. Open the class that you want to execute (e.g. ai.djl.examples.inference.cv.ObjectDetection).
 Select the triangle at the class declaration line. A popup menu appears with 3 items:
 
 - Run 'ObjectDetection.main()'
diff --git a/examples/build.gradle b/examples/build.gradle
index 20daaae9b..6edf30527 100644
--- a/examples/build.gradle
+++ b/examples/build.gradle
@@ -22,7 +22,7 @@ dependencies {
 }
 
 application {
-    mainClass = System.getProperty("main", "ai.djl.examples.inference.ObjectDetection")
+    mainClass = System.getProperty("main", "ai.djl.examples.inference.cv.ObjectDetection")
 }
 
 run {
diff --git a/examples/docs/action_recognition.md b/examples/docs/action_recognition.md
index 72b585de3..c8f424364 100644
--- a/examples/docs/action_recognition.md
+++ b/examples/docs/action_recognition.md
@@ -4,7 +4,7 @@ Action recognition is a computer vision technique to infer human actions (presen
 
 In this example, you learn how to implement inference code with a [ModelZoo model](../../docs/model-zoo.md) to detect human actions in an image.
 
-The source code can be found at [ActionRecognition.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/ActionRecognition.java).
+The source code can be found at [ActionRecognition.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/cv/ActionRecognition.java).
 
 ## Setup Guide
 
@@ -22,7 +22,7 @@ Use the following command to run the project:
 
 ```sh
 cd examples
-./gradlew run -Dmain=ai.djl.examples.inference.ActionRecognition
+./gradlew run -Dmain=ai.djl.examples.inference.cv.ActionRecognition
 ```
 
 Your output should look like the following:
diff --git a/examples/docs/biggan.md b/examples/docs/biggan.md
index 646b9071a..f46b5f72c 100644
--- a/examples/docs/biggan.md
+++ b/examples/docs/biggan.md
@@ -5,7 +5,7 @@ They consist of 2 neural networks that act as adversaries, the Generator and the
 
 In this example, you will learn how to use a [BigGAN](https://deepmind.com/research/open-source/biggan) generator to create images, using the generator directly from the [ModelZoo](../../docs/model-zoo.md).
 
-The source code for this example can be found at [BigGAN.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/BigGAN.java).
+The source code for this example can be found at [BigGAN.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/cv/BigGAN.java).
 
 ## Setup guide
 
@@ -30,7 +30,7 @@ Use the following commands to run the project:
 
 ```sh
 cd examples
-./gradlew run -Dmain=ai.djl.examples.inference.BigGAN
+./gradlew run -Dmain=ai.djl.examples.inference.cv.BigGAN
 ```
 
 ### Output
diff --git a/examples/docs/image_classification.md b/examples/docs/image_classification.md
index c8f331320..7c3e00570 100644
--- a/examples/docs/image_classification.md
+++ b/examples/docs/image_classification.md
@@ -4,7 +4,7 @@ Image classification refers to the task of extracting information classes from a
 
 In this example, you learn how to implement inference code with Deep Java Library (DJL) to recognize handwritten digits from an image.
 
-The image classification example code can be found at [ImageClassification.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/ImageClassification.java).
+The image classification example code can be found at [ImageClassification.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/cv/ImageClassification.java).
 
 You can also use the [Jupyter notebook tutorial](http://docs.djl.ai/docs/demos/jupyter/tutorial/03_image_classification_with_your_model.html).
 The Jupyter notebook explains the key concepts in detail.
@@ -32,7 +32,7 @@ Run the project by using the following command:
 
 ```sh
 cd examples
-./gradlew run -Dmain=ai.djl.examples.inference.ImageClassification
+./gradlew run -Dmain=ai.djl.examples.inference.cv.ImageClassification
 ```
 
 Your output should look like the following:
diff --git a/examples/docs/instance_segmentation.md b/examples/docs/instance_segmentation.md
index 1c3559768..c9525984e 100644
--- a/examples/docs/instance_segmentation.md
+++ b/examples/docs/instance_segmentation.md
@@ -6,7 +6,7 @@ In this example, you learn how to implement inference code with Deep Java Librar
 
 The following is the instance segmentation example source code:
 
-[InstanceSegmentation.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/InstanceSegmentation.java).
+[InstanceSegmentation.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/cv/InstanceSegmentation.java).
 
 
 ## Setup guide
@@ -24,7 +24,7 @@ You can find the image used in this example in project test resource folder: `sr
 
 ```sh
 cd examples
-./gradlew run -Dmain=ai.djl.examples.inference.InstanceSegmentation
+./gradlew run -Dmain=ai.djl.examples.inference.cv.InstanceSegmentation
 ```
 
 This should produce the following output
diff --git a/examples/docs/mask_detection.md b/examples/docs/mask_detection.md
index 303ed5219..55bd56b4a 100644
--- a/examples/docs/mask_detection.md
+++ b/examples/docs/mask_detection.md
@@ -5,7 +5,7 @@ YOLOv5 is a powerful model for object detection tasks. With the transfer learnin
 In this example, we apply it on the [Face Mask Detection dataset](https://www.kaggle.com/datasets/andrewmvd/face-mask-detection?select=images). We first train the YOLOv5s model in Python, with the help of [ATLearn](https://github.com/awslabs/atlearn/blob/main/examples/docs/face_mask_detection.md), a python transfer learning toolkit.
 Then, the model is saved as an ONNX model, which is then imported into DJL for inference. We apply it on the mask wearing detection task. 
 
-The source code can be found at [MaskDetectionOnnx.java](../src/main/java/ai/djl/examples/inference/MaskDetection.java)
+The source code can be found at [MaskDetectionOnnx.java](../src/main/java/ai/djl/examples/inference/cv/MaskDetection.java)
 
 ## The training part in ATLearn
 
@@ -13,7 +13,7 @@ We initially attempted to import a pretrained YOLOv5 into DJL, and fine-tune it
 
 The retraining of YOLOv5 can be found in an example in ATLearn: `examples/docs/face_mask_detection.md`. In this example, the YOLOv5 layers near the input are frozen while those near the output are fine-tuned with the customized data. This follows the transfer learning idea.
 
-In this example, the trained model is first exported to ONNX file, eg. `mask.onnx` and then imported in  [MaskDetectionOnnx.java](../src/main/java/ai/djl/examples/inference/MaskDetection.java).
+In this example, the trained model is first exported to ONNX file, eg. `mask.onnx` and then imported in  [MaskDetectionOnnx.java](../src/main/java/ai/djl/examples/inference/cv/MaskDetection.java).
 This ONNX model file can also be used for inference in python, which will serve as a benchmark. See the [tutorial doc](https://github.com/awslabs/atlearn/blob/main/examples/docs/face_mask_detection.md). 
 
 ## Setup guide
@@ -32,7 +32,7 @@ Use the following command to run the project:
 
 ```sh
 cd examples
-./gradlew run -Dmain=ai.djl.examples.inference.MaskDetection
+./gradlew run -Dmain=ai.djl.examples.inference.cv.MaskDetection
 ```
 
 Your output should look like the following:
diff --git a/examples/docs/object_detection.md b/examples/docs/object_detection.md
index 84286fb6e..220883f62 100644
--- a/examples/docs/object_detection.md
+++ b/examples/docs/object_detection.md
@@ -5,7 +5,7 @@ for locating instances of objects in images or videos.
 
 In this example, you learn how to implement inference code with a [ModelZoo model](../../docs/model-zoo.md) to detect dogs in an image.
 
-The source code can be found at [ObjectDetection.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/ObjectDetection.java).
+The source code can be found at [ObjectDetection.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/cv/ObjectDetection.java).
 
 You can also use the [Jupyter notebook tutorial](http://docs.djl.ai/docs/demos/jupyter/object_detection_with_model_zoo.html).
 The Jupyter notebook explains the key concepts in detail.
@@ -26,7 +26,7 @@ Use the following command to run the project:
 
 ```sh
 cd examples
-./gradlew run -Dmain=ai.djl.examples.inference.ObjectDetection
+./gradlew run -Dmain=ai.djl.examples.inference.cv.ObjectDetection
 ```
 
 Your output should look like the following:
diff --git a/examples/docs/object_detection_with_tensorflow_saved_model.md b/examples/docs/object_detection_with_tensorflow_saved_model.md
index d02ed663c..ad56989db 100644
--- a/examples/docs/object_detection_with_tensorflow_saved_model.md
+++ b/examples/docs/object_detection_with_tensorflow_saved_model.md
@@ -7,7 +7,7 @@ In this example we will use pre-trained model from [tensorflow model zoo](https:
 The following code has been tested with EfficientDet, SSD MobileNet V2, Faster RCNN Inception Resnet V2,
 but should work with most of tensorflow object detection models.
 
-The source code can be found at [ObjectDetectionWithTensorflowSavedModel.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/ObjectDetectionWithTensorflowSavedModel.java).
+The source code can be found at [ObjectDetectionWithTensorflowSavedModel.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/cv/ObjectDetectionWithTensorflowSavedModel.java).
 
 ## Setup guide
 
@@ -45,7 +45,7 @@ Use the following command to run the project:
 
 ```sh
 cd examples
-./gradlew run -Dmain=ai.djl.examples.inference.ObjectDetectionWithTensorflowSavedModel
+./gradlew run -Dmain=ai.djl.examples.inference.cv.ObjectDetectionWithTensorflowSavedModel
 ```
 
 Your output should look like the following:
diff --git a/examples/docs/pose_estimation.md b/examples/docs/pose_estimation.md
index c8b4df2cd..86bfdc52b 100644
--- a/examples/docs/pose_estimation.md
+++ b/examples/docs/pose_estimation.md
@@ -4,7 +4,7 @@ Pose estimation is a computer vision technique for determining the pose of an ob
 
 In this example, you learn how to implement inference code with a [ModelZoo model](../../docs/model-zoo.md) to detect people and their joints in an image.
 
-The source code can be found at [PoseEstimation.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/PoseEstimation.java).
+The source code can be found at [PoseEstimation.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/cv/PoseEstimation.java).
 
 ## Setup guide
 
@@ -22,7 +22,7 @@ Use the following command to run the project:
 
 ```sh
 cd examples
-./gradlew run -Dmain=ai.djl.examples.inference.PoseEstimation
+./gradlew run -Dmain=ai.djl.examples.inference.cv.PoseEstimation
 ```
 
 Your output should look like the following:
diff --git a/examples/docs/semantic_segmentation.md b/examples/docs/semantic_segmentation.md
index 2e22d6d96..e6d989de1 100644
--- a/examples/docs/semantic_segmentation.md
+++ b/examples/docs/semantic_segmentation.md
@@ -6,7 +6,7 @@ In this example, you learn how to implement inference code with Deep Java Librar
 
 The following is the semantic segmentation example source code:
 
-[SemanticSegmentation.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/SemanticSegmentation.java).
+[SemanticSegmentation.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/cv/SemanticSegmentation.java).
 
 
 ## Setup guide
@@ -24,7 +24,7 @@ You can find the image used in this example in project test resource folder: `sr
 
 ```sh
 cd examples
-./gradlew run -Dmain=ai.djl.examples.inference.SemanticSegmentation
+./gradlew run -Dmain=ai.djl.examples.inference.cv.SemanticSegmentation
 ```
 
 This should produce the following output
@@ -57,7 +57,7 @@ Path imageFile = Paths.get("src/test/resources/dog_bike_car.jpg");
 
 ```sh
 cd examples
-./gradlew run -Dmain=ai.djl.examples.inference.InstanceSegmentation
+./gradlew run -Dmain=ai.djl.examples.inference.cv.SemanticSegmentation
 ```
 
 This should produce the following output
diff --git a/examples/docs/super_resolution.md b/examples/docs/super_resolution.md
index 1d9ea6b1d..483751217 100644
--- a/examples/docs/super_resolution.md
+++ b/examples/docs/super_resolution.md
@@ -3,7 +3,7 @@
 [Generative Adversarial Networks](https://en.wikipedia.org/wiki/Generative_adversarial_network) (GANs) are a branch of deep learning used for generative modeling. 
 They consist of 2 neural networks that act as adversaries, the Generator and the Discriminator. The Generator is assigned to generate fake images that look real, and the Discriminator needs to correctly identify the fake ones.
 
-In this example, you will learn how to use a [ESRGAN](https://deepmind.com/research/open-source/biggan) generator to upscale images x4 better than any upscaling algorithm like bicubic or bilinear, using the [generator](https://tfhub.dev/captain-pool/esrgan-tf2/1) available on TensorFlow Hub. ESRGAN is short for Enhanced Super-Resolution Generative Adversarial Networks.
+In this example, you will learn how to use a [ESRGAN](https://esrgan.readthedocs.io/en/latest/) generator to upscale images x4 better than any upscaling algorithm like bicubic or bilinear, using the [generator](https://tfhub.dev/captain-pool/esrgan-tf2/1) available on TensorFlow Hub. ESRGAN is short for Enhanced Super-Resolution Generative Adversarial Networks.
 
 The source code for this example can be found in the [examples/sr](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/sr/) package.
 
diff --git a/examples/pom.xml b/examples/pom.xml
index e6a099871..63f9f4662 100644
--- a/examples/pom.xml
+++ b/examples/pom.xml
@@ -11,7 +11,7 @@
         <maven.compiler.source>11</maven.compiler.source>
         <maven.compiler.target>11</maven.compiler.target>
         <djl.version>0.28.0-SNAPSHOT</djl.version>
-        <exec.mainClass>ai.djl.examples.inference.ObjectDetection</exec.mainClass>
+        <exec.mainClass>ai.djl.examples.inference.cv.ObjectDetection</exec.mainClass>
     </properties>
 
     <repositories>
diff --git a/examples/src/main/java/ai/djl/examples/inference/ActionRecognition.java b/examples/src/main/java/ai/djl/examples/inference/cv/ActionRecognition.java
similarity index 82%
rename from examples/src/main/java/ai/djl/examples/inference/ActionRecognition.java
rename to examples/src/main/java/ai/djl/examples/inference/cv/ActionRecognition.java
index c088c5df9..5530d44f1 100644
--- a/examples/src/main/java/ai/djl/examples/inference/ActionRecognition.java
+++ b/examples/src/main/java/ai/djl/examples/inference/cv/ActionRecognition.java
@@ -10,9 +10,8 @@
  * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.
  */
-package ai.djl.examples.inference;
+package ai.djl.examples.inference.cv;
 
-import ai.djl.Application;
 import ai.djl.ModelException;
 import ai.djl.inference.Predictor;
 import ai.djl.modality.Classifications;
@@ -44,7 +43,7 @@ public final class ActionRecognition {
     private ActionRecognition() {}
 
     public static void main(String[] args) throws IOException, ModelException, TranslateException {
-        Classifications classification = ActionRecognition.predict();
+        Classifications classification = predict();
         logger.info("{}", classification);
     }
 
@@ -52,20 +51,19 @@ public final class ActionRecognition {
         Path imageFile = Paths.get("src/test/resources/action_discus_throw.png");
         Image img = ImageFactory.getInstance().fromFile(imageFile);
 
+        // Use DJL MXNet model zoo model
         Criteria<Image, Classifications> criteria =
                 Criteria.builder()
-                        .optApplication(Application.CV.ACTION_RECOGNITION)
                         .setTypes(Image.class, Classifications.class)
-                        .optFilter("backbone", "inceptionv3")
-                        .optFilter("dataset", "ucf101")
+                        .optModelUrls(
+                                "djl://ai.djl.mxnet/action_recognition/0.0.1/inceptionv3_ucf101")
                         .optEngine("MXNet")
                         .optProgress(new ProgressBar())
                         .build();
 
-        try (ZooModel<Image, Classifications> inception = criteria.loadModel()) {
-            try (Predictor<Image, Classifications> action = inception.newPredictor()) {
-                return action.predict(img);
-            }
+        try (ZooModel<Image, Classifications> inception = criteria.loadModel();
+                Predictor<Image, Classifications> action = inception.newPredictor()) {
+            return action.predict(img);
         }
     }
 }
diff --git a/examples/src/main/java/ai/djl/examples/inference/BigGAN.java b/examples/src/main/java/ai/djl/examples/inference/cv/BigGAN.java
similarity index 98%
rename from examples/src/main/java/ai/djl/examples/inference/BigGAN.java
rename to examples/src/main/java/ai/djl/examples/inference/cv/BigGAN.java
index 3d215aadb..c3f9715bb 100644
--- a/examples/src/main/java/ai/djl/examples/inference/BigGAN.java
+++ b/examples/src/main/java/ai/djl/examples/inference/cv/BigGAN.java
@@ -10,7 +10,7 @@
  * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.
  */
-package ai.djl.examples.inference;
+package ai.djl.examples.inference.cv;
 
 import ai.djl.Application;
 import ai.djl.ModelException;
diff --git a/examples/src/main/java/ai/djl/examples/inference/ImageClassification.java b/examples/src/main/java/ai/djl/examples/inference/cv/ImageClassification.java
similarity index 96%
rename from examples/src/main/java/ai/djl/examples/inference/ImageClassification.java
rename to examples/src/main/java/ai/djl/examples/inference/cv/ImageClassification.java
index 386ea066c..d0ba39e43 100644
--- a/examples/src/main/java/ai/djl/examples/inference/ImageClassification.java
+++ b/examples/src/main/java/ai/djl/examples/inference/cv/ImageClassification.java
@@ -10,7 +10,7 @@
  * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.
  */
-package ai.djl.examples.inference;
+package ai.djl.examples.inference.cv;
 
 import ai.djl.Model;
 import ai.djl.ModelException;
@@ -48,7 +48,7 @@ public final class ImageClassification {
     private ImageClassification() {}
 
     public static void main(String[] args) throws IOException, ModelException, TranslateException {
-        Classifications classifications = ImageClassification.predict();
+        Classifications classifications = predict();
         logger.info("{}", classifications);
     }
 
diff --git a/examples/src/main/java/ai/djl/examples/inference/InstanceSegmentation.java b/examples/src/main/java/ai/djl/examples/inference/cv/InstanceSegmentation.java
similarity index 80%
rename from examples/src/main/java/ai/djl/examples/inference/InstanceSegmentation.java
rename to examples/src/main/java/ai/djl/examples/inference/cv/InstanceSegmentation.java
index 45890cf2b..057f15bf3 100644
--- a/examples/src/main/java/ai/djl/examples/inference/InstanceSegmentation.java
+++ b/examples/src/main/java/ai/djl/examples/inference/cv/InstanceSegmentation.java
@@ -10,9 +10,8 @@
  * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.
  */
-package ai.djl.examples.inference;
+package ai.djl.examples.inference.cv;
 
-import ai.djl.Application;
 import ai.djl.ModelException;
 import ai.djl.inference.Predictor;
 import ai.djl.modality.cv.Image;
@@ -45,7 +44,7 @@ public final class InstanceSegmentation {
     private InstanceSegmentation() {}
 
     public static void main(String[] args) throws IOException, ModelException, TranslateException {
-        DetectedObjects detection = InstanceSegmentation.predict();
+        DetectedObjects detection = predict();
         logger.info("{}", detection);
     }
 
@@ -55,20 +54,18 @@ public final class InstanceSegmentation {
 
         Criteria<Image, DetectedObjects> criteria =
                 Criteria.builder()
-                        .optApplication(Application.CV.INSTANCE_SEGMENTATION)
                         .setTypes(Image.class, DetectedObjects.class)
-                        .optFilter("backbone", "resnet18")
-                        .optFilter("flavor", "v1b")
-                        .optFilter("dataset", "coco")
+                        .optModelUrls(
+                                "djl://ai.djl.mxnet/mask_rcnn/0.0.1/mask_rcnn_resnet18_v1b_coco")
+                        .optEngine("MXNet")
                         .optProgress(new ProgressBar())
                         .build();
 
-        try (ZooModel<Image, DetectedObjects> model = criteria.loadModel()) {
-            try (Predictor<Image, DetectedObjects> predictor = model.newPredictor()) {
-                DetectedObjects detection = predictor.predict(img);
-                saveBoundingBoxImage(img, detection);
-                return detection;
-            }
+        try (ZooModel<Image, DetectedObjects> model = criteria.loadModel();
+                Predictor<Image, DetectedObjects> predictor = model.newPredictor()) {
+            DetectedObjects detection = predictor.predict(img);
+            saveBoundingBoxImage(img, detection);
+            return detection;
         }
     }
 
diff --git a/examples/src/main/java/ai/djl/examples/inference/MaskDetection.java b/examples/src/main/java/ai/djl/examples/inference/cv/MaskDetection.java
similarity index 99%
rename from examples/src/main/java/ai/djl/examples/inference/MaskDetection.java
rename to examples/src/main/java/ai/djl/examples/inference/cv/MaskDetection.java
index ab06b29ff..b9d71ed67 100644
--- a/examples/src/main/java/ai/djl/examples/inference/MaskDetection.java
+++ b/examples/src/main/java/ai/djl/examples/inference/cv/MaskDetection.java
@@ -10,7 +10,7 @@
  * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.
  */
-package ai.djl.examples.inference;
+package ai.djl.examples.inference.cv;
 
 import ai.djl.ModelException;
 import ai.djl.inference.Predictor;
diff --git a/examples/src/main/java/ai/djl/examples/inference/ObjectDetection.java b/examples/src/main/java/ai/djl/examples/inference/cv/ObjectDetection.java
similarity index 89%
rename from examples/src/main/java/ai/djl/examples/inference/ObjectDetection.java
rename to examples/src/main/java/ai/djl/examples/inference/cv/ObjectDetection.java
index 08643cb3b..f1844d10f 100644
--- a/examples/src/main/java/ai/djl/examples/inference/ObjectDetection.java
+++ b/examples/src/main/java/ai/djl/examples/inference/cv/ObjectDetection.java
@@ -10,7 +10,7 @@
  * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.
  */
-package ai.djl.examples.inference;
+package ai.djl.examples.inference.cv;
 
 import ai.djl.Application;
 import ai.djl.ModelException;
@@ -46,7 +46,7 @@ public final class ObjectDetection {
     private ObjectDetection() {}
 
     public static void main(String[] args) throws IOException, ModelException, TranslateException {
-        DetectedObjects detection = ObjectDetection.predict();
+        DetectedObjects detection = predict();
         logger.info("{}", detection);
     }
 
@@ -70,12 +70,11 @@ public final class ObjectDetection {
                         .optProgress(new ProgressBar())
                         .build();
 
-        try (ZooModel<Image, DetectedObjects> model = criteria.loadModel()) {
-            try (Predictor<Image, DetectedObjects> predictor = model.newPredictor()) {
-                DetectedObjects detection = predictor.predict(img);
-                saveBoundingBoxImage(img, detection);
-                return detection;
-            }
+        try (ZooModel<Image, DetectedObjects> model = criteria.loadModel();
+                Predictor<Image, DetectedObjects> predictor = model.newPredictor()) {
+            DetectedObjects detection = predictor.predict(img);
+            saveBoundingBoxImage(img, detection);
+            return detection;
         }
     }
 
diff --git a/examples/src/main/java/ai/djl/examples/inference/ObjectDetectionWithTensorflowSavedModel.java b/examples/src/main/java/ai/djl/examples/inference/cv/ObjectDetectionWithTensorflowSavedModel.java
similarity index 98%
rename from examples/src/main/java/ai/djl/examples/inference/ObjectDetectionWithTensorflowSavedModel.java
rename to examples/src/main/java/ai/djl/examples/inference/cv/ObjectDetectionWithTensorflowSavedModel.java
index ca2458bcc..bfff70077 100644
--- a/examples/src/main/java/ai/djl/examples/inference/ObjectDetectionWithTensorflowSavedModel.java
+++ b/examples/src/main/java/ai/djl/examples/inference/cv/ObjectDetectionWithTensorflowSavedModel.java
@@ -10,7 +10,7 @@
  * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.
  */
-package ai.djl.examples.inference;
+package ai.djl.examples.inference.cv;
 
 import ai.djl.Application;
 import ai.djl.ModelException;
@@ -71,7 +71,7 @@ public final class ObjectDetectionWithTensorflowSavedModel {
     private ObjectDetectionWithTensorflowSavedModel() {}
 
     public static void main(String[] args) throws IOException, ModelException, TranslateException {
-        DetectedObjects detection = ObjectDetectionWithTensorflowSavedModel.predict();
+        DetectedObjects detection = predict();
         logger.info("{}", detection);
     }
 
diff --git a/examples/src/main/java/ai/djl/examples/inference/PoseEstimation.java b/examples/src/main/java/ai/djl/examples/inference/cv/PoseEstimation.java
similarity index 86%
rename from examples/src/main/java/ai/djl/examples/inference/PoseEstimation.java
rename to examples/src/main/java/ai/djl/examples/inference/cv/PoseEstimation.java
index 5932fd3c1..1efc59477 100644
--- a/examples/src/main/java/ai/djl/examples/inference/PoseEstimation.java
+++ b/examples/src/main/java/ai/djl/examples/inference/cv/PoseEstimation.java
@@ -10,9 +10,8 @@
  * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.
  */
-package ai.djl.examples.inference;
+package ai.djl.examples.inference.cv;
 
-import ai.djl.Application;
 import ai.djl.MalformedModelException;
 import ai.djl.ModelException;
 import ai.djl.inference.Predictor;
@@ -52,7 +51,7 @@ public final class PoseEstimation {
     private PoseEstimation() {}
 
     public static void main(String[] args) throws IOException, ModelException, TranslateException {
-        List<Joints> joints = PoseEstimation.predict();
+        List<Joints> joints = predict();
         logger.info("{}", joints);
     }
 
@@ -76,21 +75,16 @@ public final class PoseEstimation {
 
         Criteria<Image, DetectedObjects> criteria =
                 Criteria.builder()
-                        .optApplication(Application.CV.OBJECT_DETECTION)
                         .setTypes(Image.class, DetectedObjects.class)
-                        .optFilter("size", "512")
-                        .optFilter("backbone", "resnet50")
-                        .optFilter("flavor", "v1")
-                        .optFilter("dataset", "voc")
+                        .optModelUrls("djl://ai.djl.mxnet/ssd/0.0.1/ssd_512_resnet50_v1_voc")
                         .optEngine("MXNet")
                         .optProgress(new ProgressBar())
                         .build();
 
         DetectedObjects detectedObjects;
-        try (ZooModel<Image, DetectedObjects> ssd = criteria.loadModel()) {
-            try (Predictor<Image, DetectedObjects> predictor = ssd.newPredictor()) {
-                detectedObjects = predictor.predict(img);
-            }
+        try (ZooModel<Image, DetectedObjects> ssd = criteria.loadModel();
+                Predictor<Image, DetectedObjects> predictor = ssd.newPredictor()) {
+            detectedObjects = predictor.predict(img);
         }
 
         List<DetectedObjects.DetectedObject> items = detectedObjects.items();
@@ -115,13 +109,14 @@ public final class PoseEstimation {
             throws MalformedModelException, ModelNotFoundException, IOException,
                     TranslateException {
 
+        // Use DJL MXNet model zoo model, model can be found:
+        // https://mlrepo.djl.ai/model/cv/pose_estimation/ai/djl/mxnet/simple_pose/0.0.1/simple_pose_resnet18_v1b-0000.params.gz
+        // https://mlrepo.djl.ai/model/cv/pose_estimation/ai/djl/mxnet/simple_pose/0.0.1/simple_pose_resnet18_v1b-symbol.json
         Criteria<Image, Joints> criteria =
                 Criteria.builder()
-                        .optApplication(Application.CV.POSE_ESTIMATION)
                         .setTypes(Image.class, Joints.class)
-                        .optFilter("backbone", "resnet18")
-                        .optFilter("flavor", "v1b")
-                        .optFilter("dataset", "imagenet")
+                        .optModelUrls(
+                                "djl://ai.djl.mxnet/simple_pose/0.0.1/simple_pose_resnet18_v1b")
                         .build();
 
         List<Joints> allJoints = new ArrayList<>();
diff --git a/examples/src/main/java/ai/djl/examples/inference/SemanticSegmentation.java b/examples/src/main/java/ai/djl/examples/inference/cv/SemanticSegmentation.java
similarity index 93%
rename from examples/src/main/java/ai/djl/examples/inference/SemanticSegmentation.java
rename to examples/src/main/java/ai/djl/examples/inference/cv/SemanticSegmentation.java
index ef281e2c0..deb2606c9 100644
--- a/examples/src/main/java/ai/djl/examples/inference/SemanticSegmentation.java
+++ b/examples/src/main/java/ai/djl/examples/inference/cv/SemanticSegmentation.java
@@ -10,7 +10,7 @@
  * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.
  */
-package ai.djl.examples.inference;
+package ai.djl.examples.inference.cv;
 
 import ai.djl.ModelException;
 import ai.djl.inference.Predictor;
@@ -54,13 +54,12 @@ public final class SemanticSegmentation {
         ImageFactory factory = ImageFactory.getInstance();
         Image img = factory.fromFile(imageFile);
 
-        String url =
-                "https://mlrepo.djl.ai/model/cv/semantic_segmentation/ai/djl/pytorch/deeplabv3/0.0.1/deeplabv3.zip";
-
+        // Use DJL PyTorch model zoo model, model can be found:
+        // https://mlrepo.djl.ai/model/cv/semantic_segmentation/ai/djl/pytorch/deeplabv3/0.0.1/deeplabv3.zip
         Criteria<Image, CategoryMask> criteria =
                 Criteria.builder()
                         .setTypes(Image.class, CategoryMask.class)
-                        .optModelUrls(url)
+                        .optModelUrls("djl://ai.djl.pytorch/deeplabv3")
                         .optTranslatorFactory(new SemanticSegmentationTranslatorFactory())
                         .optEngine("PyTorch")
                         .optProgress(new ProgressBar())
diff --git a/examples/src/main/java/ai/djl/examples/inference/StyleTransfer.java b/examples/src/main/java/ai/djl/examples/inference/cv/StyleTransfer.java
similarity index 85%
rename from examples/src/main/java/ai/djl/examples/inference/StyleTransfer.java
rename to examples/src/main/java/ai/djl/examples/inference/cv/StyleTransfer.java
index d6eaf9fab..34fef8a89 100644
--- a/examples/src/main/java/ai/djl/examples/inference/StyleTransfer.java
+++ b/examples/src/main/java/ai/djl/examples/inference/cv/StyleTransfer.java
@@ -10,7 +10,7 @@
  * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.
  */
-package ai.djl.examples.inference;
+package ai.djl.examples.inference.cv;
 
 import ai.djl.Application;
 import ai.djl.MalformedModelException;
@@ -32,6 +32,7 @@ import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.Locale;
 
 public final class StyleTransfer {
 
@@ -48,6 +49,14 @@ public final class StyleTransfer {
 
     public static void main(String[] args) throws IOException, ModelException, TranslateException {
         Artist artist = Artist.MONET;
+        if (args.length == 1) {
+            for (Artist a : Artist.values()) {
+                if (a.name().equalsIgnoreCase(args[0])) {
+                    artist = a;
+                    break;
+                }
+            }
+        }
         String imagePath = "src/test/resources/mountains.png";
         Image input = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
         Image output = transfer(input, artist);
@@ -61,11 +70,11 @@ public final class StyleTransfer {
     public static Image transfer(Image image, Artist artist)
             throws IOException, ModelNotFoundException, MalformedModelException,
                     TranslateException {
-        String modelName = "style_" + artist.toString().toLowerCase() + ".zip";
-        String modelUrl =
-                "https://mlrepo.djl.ai/model/cv/image_generation/ai/djl/pytorch/cyclegan/0.0.1/"
-                        + modelName;
+        // Use DJL PyTorch model zoo model, model can be found:
+        // https://mlrepo.djl.ai/model/cv/image_generation/ai/djl/pytorch/cyclegan/0.0.1/style_xxxx.zip
 
+        String modelName = "style_" + artist.toString().toLowerCase(Locale.ROOT);
+        String modelUrl = "djl://ai.djl.pytorch/cyclegan/0.0.1/" + modelName;
         Criteria<Image, Image> criteria =
                 Criteria.builder()
                         .optApplication(Application.CV.IMAGE_GENERATION)
diff --git a/examples/src/main/java/ai/djl/examples/inference/Yolov8Detection.java b/examples/src/main/java/ai/djl/examples/inference/cv/Yolov8Detection.java
similarity index 93%
rename from examples/src/main/java/ai/djl/examples/inference/Yolov8Detection.java
rename to examples/src/main/java/ai/djl/examples/inference/cv/Yolov8Detection.java
index 3d2cfb264..5aee1246c 100644
--- a/examples/src/main/java/ai/djl/examples/inference/Yolov8Detection.java
+++ b/examples/src/main/java/ai/djl/examples/inference/cv/Yolov8Detection.java
@@ -10,7 +10,7 @@
  * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.
  */
-package ai.djl.examples.inference;
+package ai.djl.examples.inference.cv;
 
 import ai.djl.ModelException;
 import ai.djl.inference.Predictor;
@@ -40,7 +40,7 @@ public final class Yolov8Detection {
     private Yolov8Detection() {}
 
     public static void main(String[] args) throws IOException, ModelException, TranslateException {
-        DetectedObjects detection = Yolov8Detection.predict();
+        DetectedObjects detection = predict();
         logger.info("{}", detection);
     }
 
@@ -48,6 +48,8 @@ public final class Yolov8Detection {
         Path imgPath = Paths.get("src/test/resources/yolov8_test.jpg");
         Image img = ImageFactory.getInstance().fromFile(imgPath);
 
+        // Use DJL OnnxRuntime model zoo model, model can be found:
+        // https://mlrepo.djl.ai/model/cv/object_detection/ai/djl/onnxruntime/yolov8n/0.0.1/yolov8n.zip
         Criteria<Image, DetectedObjects> criteria =
                 Criteria.builder()
                         .setTypes(Image.class, DetectedObjects.class)
diff --git a/examples/src/main/java/ai/djl/examples/inference/cv/package-info.java b/examples/src/main/java/ai/djl/examples/inference/cv/package-info.java
new file mode 100644
index 000000000..715e3d6fe
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/inference/cv/package-info.java
@@ -0,0 +1,14 @@
+/*
+ * Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+/** Contains examples of CV. */
+package ai.djl.examples.inference.cv;
diff --git a/examples/src/main/java/ai/djl/examples/inference/nlp/RollingBatch.java b/examples/src/main/java/ai/djl/examples/inference/nlp/RollingBatch.java
index 6aac42d6f..a0aada163 100644
--- a/examples/src/main/java/ai/djl/examples/inference/nlp/RollingBatch.java
+++ b/examples/src/main/java/ai/djl/examples/inference/nlp/RollingBatch.java
@@ -13,7 +13,6 @@
 package ai.djl.examples.inference.nlp;
 
 import ai.djl.MalformedModelException;
-import ai.djl.examples.inference.ImageClassification;
 import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
 import ai.djl.inference.Predictor;
 import ai.djl.modality.nlp.generate.CausalLMOutput;
@@ -36,7 +35,7 @@ import java.util.Map;
 
 public final class RollingBatch {
 
-    private static final Logger logger = LoggerFactory.getLogger(ImageClassification.class);
+    private static final Logger logger = LoggerFactory.getLogger(RollingBatch.class);
 
     private RollingBatch() {}
 
diff --git a/examples/src/main/java/ai/djl/examples/inference/nlp/TextGeneration.java b/examples/src/main/java/ai/djl/examples/inference/nlp/TextGeneration.java
index 59cba679b..676188412 100644
--- a/examples/src/main/java/ai/djl/examples/inference/nlp/TextGeneration.java
+++ b/examples/src/main/java/ai/djl/examples/inference/nlp/TextGeneration.java
@@ -13,7 +13,6 @@
 package ai.djl.examples.inference.nlp;
 
 import ai.djl.MalformedModelException;
-import ai.djl.examples.inference.ImageClassification;
 import ai.djl.huggingface.tokenizers.Encoding;
 import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
 import ai.djl.inference.Predictor;
@@ -38,7 +37,7 @@ import java.io.IOException;
 
 public final class TextGeneration {
 
-    private static final Logger logger = LoggerFactory.getLogger(ImageClassification.class);
+    private static final Logger logger = LoggerFactory.getLogger(TextGeneration.class);
 
     private TextGeneration() {}
 


Test patch:
diff --git a/api/src/test/java/ai/djl/nn/BlockFactoryTest.java b/api/src/test/java/ai/djl/nn/BlockFactoryTest.java
index 230631486..54502d171 100644
--- a/api/src/test/java/ai/djl/nn/BlockFactoryTest.java
+++ b/api/src/test/java/ai/djl/nn/BlockFactoryTest.java
@@ -13,12 +13,19 @@
 package ai.djl.nn;
 
 import ai.djl.Model;
+import ai.djl.ModelException;
+import ai.djl.ndarray.NDList;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.ZooModel;
 
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
+import java.io.IOException;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 @SuppressWarnings("PMD.TestClassWithoutTestCases")
 public class BlockFactoryTest {
@@ -32,4 +39,35 @@ public class BlockFactoryTest {
             Assert.assertEquals(((LambdaBlock) block).getName(), "identity");
         }
     }
+
+    @Test
+    public void testOnesBlockFactory() throws ModelException, IOException {
+        OnesBlockFactory factory = new OnesBlockFactory();
+        Path path = Paths.get("build");
+        Criteria<NDList, NDList> criteria =
+                Criteria.builder()
+                        .setTypes(NDList.class, NDList.class)
+                        .optModelPath(path)
+                        .optArgument("blockFactory", "ai.djl.nn.OnesBlockFactory")
+                        .optArgument("block_shapes", "(1)s,(1)d,(1)u,(1)b,(1)i,(1)l,(1)B,(1)f,(1)")
+                        .optArgument("block_names", "1,2,3,4,5,6,7,8,9")
+                        .optOption("hasParameter", "false")
+                        .build();
+
+        try (ZooModel<NDList, NDList> model = criteria.loadModel()) {
+            Block block = model.getBlock();
+            Assert.assertTrue(block instanceof LambdaBlock);
+
+            Map<String, String> args = new ConcurrentHashMap<>();
+            args.put("block_shapes", "1,2");
+            block = factory.newBlock(model, path, args);
+            Assert.assertTrue(block instanceof LambdaBlock);
+
+            args.put("block_shapes", "(1)a");
+            Assert.assertThrows(
+                    () -> {
+                        factory.newBlock(model, path, args);
+                    });
+        }
+    }
 }
diff --git a/api/src/test/java/ai/djl/translate/ServingTranslatorFactoryTest.java b/api/src/test/java/ai/djl/translate/ServingTranslatorFactoryTest.java
new file mode 100644
index 000000000..61faab1d8
--- /dev/null
+++ b/api/src/test/java/ai/djl/translate/ServingTranslatorFactoryTest.java
@@ -0,0 +1,136 @@
+/*
+ * Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.translate;
+
+import ai.djl.ModelException;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.nn.Blocks;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.ZooModel;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+
+public class ServingTranslatorFactoryTest {
+
+    @Test
+    public void test() throws IOException, TranslateException, ModelException {
+        Path path = Paths.get("build/model");
+        Files.createDirectories(path);
+        Criteria<Input, Output> criteria =
+                Criteria.builder()
+                        .setTypes(Input.class, Output.class)
+                        .optModelPath(path)
+                        .optModelName("identity")
+                        .optArgument("application", "image_classification")
+                        .optOption("hasParameter", "false")
+                        .optBlock(Blocks.identityBlock())
+                        .build();
+
+        try (ZooModel<Input, Output> model = criteria.loadModel()) {
+            String className = model.getTranslator().getClass().getSimpleName();
+            Assert.assertEquals(className, "ImageServingTranslator");
+        }
+
+        criteria =
+                Criteria.builder()
+                        .setTypes(Input.class, Output.class)
+                        .optModelPath(path)
+                        .optModelName("identity")
+                        // HF style task name
+                        .optArgument("task", "fill-mask")
+                        .optOption("hasParameter", "false")
+                        .optBlock(Blocks.identityBlock())
+                        .build();
+
+        try (ZooModel<Input, Output> model = criteria.loadModel()) {
+            String className = model.getTranslator().getClass().getSimpleName();
+            // Tokenizer is not in class path fallback to NoopServingTranslator
+            Assert.assertEquals(className, "NoopServingTranslator");
+        }
+
+        criteria =
+                Criteria.builder()
+                        .setTypes(Input.class, Output.class)
+                        .optModelPath(path)
+                        .optModelName("identity")
+                        // HF style task name
+                        .optArgument("task", "question-answering")
+                        .optOption("hasParameter", "false")
+                        .optBlock(Blocks.identityBlock())
+                        .build();
+
+        try (ZooModel<Input, Output> model = criteria.loadModel()) {
+            String className = model.getTranslator().getClass().getSimpleName();
+            // Tokenizer is not in class path fallback to NoopServingTranslator
+            Assert.assertEquals(className, "NoopServingTranslator");
+        }
+
+        criteria =
+                Criteria.builder()
+                        .setTypes(Input.class, Output.class)
+                        .optModelPath(path)
+                        .optModelName("identity")
+                        // HF style task name
+                        .optArgument("task", "sentence-similarity")
+                        .optOption("hasParameter", "false")
+                        .optBlock(Blocks.identityBlock())
+                        .build();
+
+        try (ZooModel<Input, Output> model = criteria.loadModel()) {
+            String className = model.getTranslator().getClass().getSimpleName();
+            // Tokenizer is not in class path fallback to NoopServingTranslator
+            Assert.assertEquals(className, "NoopServingTranslator");
+        }
+
+        criteria =
+                Criteria.builder()
+                        .setTypes(Input.class, Output.class)
+                        .optModelPath(path)
+                        .optModelName("identity")
+                        // HF style task name
+                        .optArgument("task", "text-classification")
+                        .optOption("hasParameter", "false")
+                        .optBlock(Blocks.identityBlock())
+                        .build();
+
+        try (ZooModel<Input, Output> model = criteria.loadModel()) {
+            String className = model.getTranslator().getClass().getSimpleName();
+            // Tokenizer is not in class path fallback to NoopServingTranslator
+            Assert.assertEquals(className, "NoopServingTranslator");
+        }
+
+        criteria =
+                Criteria.builder()
+                        .setTypes(Input.class, Output.class)
+                        .optModelPath(path)
+                        .optModelName("identity")
+                        // HF style task name
+                        .optArgument("task", "token-classification")
+                        .optOption("hasParameter", "false")
+                        .optBlock(Blocks.identityBlock())
+                        .build();
+
+        try (ZooModel<Input, Output> model = criteria.loadModel()) {
+            String className = model.getTranslator().getClass().getSimpleName();
+            // Tokenizer is not in class path fallback to NoopServingTranslator
+            Assert.assertEquals(className, "NoopServingTranslator");
+        }
+    }
+}
diff --git a/examples/src/test/java/ai/djl/examples/inference/ActionRecognitionTest.java b/examples/src/test/java/ai/djl/examples/inference/ActionRecognitionTest.java
index f2b450737..acf83fbb1 100644
--- a/examples/src/test/java/ai/djl/examples/inference/ActionRecognitionTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/ActionRecognitionTest.java
@@ -13,6 +13,7 @@
 package ai.djl.examples.inference;
 
 import ai.djl.ModelException;
+import ai.djl.examples.inference.cv.ActionRecognition;
 import ai.djl.modality.Classifications;
 import ai.djl.testing.TestRequirements;
 import ai.djl.translate.TranslateException;
diff --git a/examples/src/test/java/ai/djl/examples/inference/BigGANTest.java b/examples/src/test/java/ai/djl/examples/inference/BigGANTest.java
index 00ad1df9a..307a6e915 100644
--- a/examples/src/test/java/ai/djl/examples/inference/BigGANTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/BigGANTest.java
@@ -13,6 +13,7 @@
 package ai.djl.examples.inference;
 
 import ai.djl.ModelException;
+import ai.djl.examples.inference.cv.BigGAN;
 import ai.djl.modality.cv.Image;
 import ai.djl.testing.TestRequirements;
 import ai.djl.translate.TranslateException;
diff --git a/examples/src/test/java/ai/djl/examples/inference/InstanceSegmentationTest.java b/examples/src/test/java/ai/djl/examples/inference/InstanceSegmentationTest.java
index 33812b409..d17d6dcae 100644
--- a/examples/src/test/java/ai/djl/examples/inference/InstanceSegmentationTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/InstanceSegmentationTest.java
@@ -13,6 +13,7 @@
 package ai.djl.examples.inference;
 
 import ai.djl.ModelException;
+import ai.djl.examples.inference.cv.InstanceSegmentation;
 import ai.djl.modality.Classifications;
 import ai.djl.modality.cv.output.DetectedObjects;
 import ai.djl.testing.TestRequirements;
diff --git a/examples/src/test/java/ai/djl/examples/inference/MaskDetectionTest.java b/examples/src/test/java/ai/djl/examples/inference/MaskDetectionTest.java
index db8f4fb74..818be1f65 100644
--- a/examples/src/test/java/ai/djl/examples/inference/MaskDetectionTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/MaskDetectionTest.java
@@ -13,6 +13,7 @@
 package ai.djl.examples.inference;
 
 import ai.djl.ModelException;
+import ai.djl.examples.inference.cv.MaskDetection;
 import ai.djl.modality.Classifications;
 import ai.djl.modality.cv.output.DetectedObjects;
 import ai.djl.testing.TestRequirements;
diff --git a/examples/src/test/java/ai/djl/examples/inference/ObjectDetectionTest.java b/examples/src/test/java/ai/djl/examples/inference/ObjectDetectionTest.java
index d798dd354..5f612ec87 100644
--- a/examples/src/test/java/ai/djl/examples/inference/ObjectDetectionTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/ObjectDetectionTest.java
@@ -13,6 +13,7 @@
 package ai.djl.examples.inference;
 
 import ai.djl.ModelException;
+import ai.djl.examples.inference.cv.ObjectDetection;
 import ai.djl.modality.Classifications;
 import ai.djl.modality.cv.output.DetectedObjects;
 import ai.djl.testing.TestRequirements;
diff --git a/examples/src/test/java/ai/djl/examples/inference/ObjectDetectionWithTensorflowSavedModelTest.java b/examples/src/test/java/ai/djl/examples/inference/ObjectDetectionWithTensorflowSavedModelTest.java
index f208d034f..c3849598a 100644
--- a/examples/src/test/java/ai/djl/examples/inference/ObjectDetectionWithTensorflowSavedModelTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/ObjectDetectionWithTensorflowSavedModelTest.java
@@ -14,6 +14,7 @@
 package ai.djl.examples.inference;
 
 import ai.djl.ModelException;
+import ai.djl.examples.inference.cv.ObjectDetectionWithTensorflowSavedModel;
 import ai.djl.modality.Classifications;
 import ai.djl.modality.cv.output.DetectedObjects;
 import ai.djl.testing.TestRequirements;
diff --git a/examples/src/test/java/ai/djl/examples/inference/PoseEstimationTest.java b/examples/src/test/java/ai/djl/examples/inference/PoseEstimationTest.java
index f963eb956..da9277206 100644
--- a/examples/src/test/java/ai/djl/examples/inference/PoseEstimationTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/PoseEstimationTest.java
@@ -13,6 +13,7 @@
 package ai.djl.examples.inference;
 
 import ai.djl.ModelException;
+import ai.djl.examples.inference.cv.PoseEstimation;
 import ai.djl.modality.cv.output.Joints;
 import ai.djl.testing.TestRequirements;
 import ai.djl.translate.TranslateException;
diff --git a/examples/src/test/java/ai/djl/examples/inference/StyleTransferTest.java b/examples/src/test/java/ai/djl/examples/inference/StyleTransferTest.java
index c6dcec599..8b4d4510a 100644
--- a/examples/src/test/java/ai/djl/examples/inference/StyleTransferTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/StyleTransferTest.java
@@ -13,6 +13,7 @@
 package ai.djl.examples.inference;
 
 import ai.djl.MalformedModelException;
+import ai.djl.examples.inference.cv.StyleTransfer;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.ImageFactory;
 import ai.djl.repository.zoo.ModelNotFoundException;
diff --git a/examples/src/test/java/ai/djl/examples/inference/Yolov8DetectionTest.java b/examples/src/test/java/ai/djl/examples/inference/Yolov8DetectionTest.java
index 35e3fc434..780d90dc9 100644
--- a/examples/src/test/java/ai/djl/examples/inference/Yolov8DetectionTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/Yolov8DetectionTest.java
@@ -13,6 +13,7 @@
 package ai.djl.examples.inference;
 
 import ai.djl.ModelException;
+import ai.djl.examples.inference.cv.Yolov8Detection;
 import ai.djl.modality.Classifications;
 import ai.djl.modality.cv.output.DetectedObjects;
 import ai.djl.testing.TestRequirements;
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java b/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java
index 5949ad406..e72d3cf54 100644
--- a/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java
+++ b/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java
@@ -13,7 +13,7 @@
 package ai.djl.examples.training;
 
 import ai.djl.ModelException;
-import ai.djl.examples.inference.ImageClassification;
+import ai.djl.examples.inference.cv.ImageClassification;
 import ai.djl.modality.Classifications;
 import ai.djl.testing.TestRequirements;
 import ai.djl.training.TrainingResult;
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java
index 476eed24c..2712dc424 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java
@@ -21,14 +21,18 @@ import ai.djl.modality.Output;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.Block;
+import ai.djl.nn.Blocks;
 import ai.djl.nn.LambdaBlock;
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ZooModel;
 import ai.djl.testing.Assertions;
 import ai.djl.translate.TranslateException;
 import ai.djl.util.JsonUtils;
+import ai.djl.util.PairList;
+import ai.djl.util.Utils;
 
 import org.testng.Assert;
 import org.testng.annotations.Test;
@@ -50,15 +54,6 @@ public class TextEmbeddingTranslatorTest {
             throws ModelException, IOException, TranslateException {
         String text = "This is an example sentence";
 
-        Block block =
-                new LambdaBlock(
-                        a -> {
-                            NDManager manager = a.getManager();
-                            NDArray arr = manager.ones(new Shape(1, 7, 384));
-                            arr.setName("last_hidden_state");
-                            return new NDList(arr);
-                        },
-                        "model");
         Path modelDir = Paths.get("build/model");
         Files.createDirectories(modelDir);
         try (NDManager manager = NDManager.newBaseManager("Rust")) {
@@ -84,8 +79,10 @@ public class TextEmbeddingTranslatorTest {
                 Criteria.builder()
                         .setTypes(String.class, float[].class)
                         .optModelPath(modelDir)
-                        .optBlock(block)
                         .optEngine("PyTorch")
+                        .optArgument("blockFactory", "ai.djl.nn.OnesBlockFactory")
+                        .optArgument("block_shapes", "(1,7,384)")
+                        .optArgument("block_names", "last_hidden_state")
                         .optArgument("tokenizer", "bert-base-uncased")
                         .optOption("hasParameter", "false")
                         .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
@@ -103,7 +100,9 @@ public class TextEmbeddingTranslatorTest {
                 Criteria.builder()
                         .setTypes(String.class, float[].class)
                         .optModelPath(modelDir)
-                        .optBlock(block)
+                        .optArgument("blockFactory", "ai.djl.nn.OnesBlockFactory")
+                        .optArgument("block_shapes", "(1,7,384)")
+                        .optArgument("block_names", "last_hidden_state")
                         .optEngine("PyTorch")
                         .optArgument("tokenizer", "bert-base-uncased")
                         .optArgument("pooling", "max")
@@ -123,7 +122,9 @@ public class TextEmbeddingTranslatorTest {
                 Criteria.builder()
                         .setTypes(String.class, float[].class)
                         .optModelPath(modelDir)
-                        .optBlock(block)
+                        .optArgument("blockFactory", "ai.djl.nn.OnesBlockFactory")
+                        .optArgument("block_shapes", "(1,7,384)")
+                        .optArgument("block_names", "last_hidden_state")
                         .optEngine("PyTorch")
                         .optArgument("tokenizer", "bert-base-uncased")
                         .optArgument("pooling", "mean_sqrt_len")
@@ -143,7 +144,9 @@ public class TextEmbeddingTranslatorTest {
                 Criteria.builder()
                         .setTypes(String.class, float[].class)
                         .optModelPath(modelDir)
-                        .optBlock(block)
+                        .optArgument("blockFactory", "ai.djl.nn.OnesBlockFactory")
+                        .optArgument("block_shapes", "(1,7,384)")
+                        .optArgument("block_names", "last_hidden_state")
                         .optEngine("PyTorch")
                         .optArgument("tokenizer", "bert-base-uncased")
                         .optArgument("pooling", "weightedmean")
@@ -163,7 +166,9 @@ public class TextEmbeddingTranslatorTest {
                 Criteria.builder()
                         .setTypes(String.class, float[].class)
                         .optModelPath(modelDir)
-                        .optBlock(block)
+                        .optArgument("blockFactory", "ai.djl.nn.OnesBlockFactory")
+                        .optArgument("block_shapes", "(1,7,384)")
+                        .optArgument("block_names", "last_hidden_state")
                         .optEngine("PyTorch")
                         .optArgument("tokenizer", "bert-base-uncased")
                         .optArgument("dense", "linear.safetensors")
@@ -183,7 +188,9 @@ public class TextEmbeddingTranslatorTest {
                 Criteria.builder()
                         .setTypes(Input.class, Output.class)
                         .optModelPath(modelDir)
-                        .optBlock(block)
+                        .optArgument("blockFactory", "ai.djl.nn.OnesBlockFactory")
+                        .optArgument("block_shapes", "(1,7,384)")
+                        .optArgument("block_names", "last_hidden_state")
                         .optEngine("PyTorch")
                         .optArgument("tokenizer", "bert-base-uncased")
                         .optArgument("pooling", "cls")
@@ -212,7 +219,9 @@ public class TextEmbeddingTranslatorTest {
         }
 
         try (Model model = Model.newInstance("test")) {
-            model.setBlock(block);
+            PairList<DataType, Shape> pairs = new PairList<>();
+            pairs.add(DataType.FLOAT32, new Shape(1, 7, 384));
+            model.setBlock(Blocks.onesBlock(pairs, Utils.EMPTY_ARRAY));
             Map<String, String> options = new HashMap<>();
             options.put("hasParameter", "false");
             model.load(modelDir, "test", options);