Repository: deepjavalibrary/djl
Pull request #1828 (2022-07-24 22:46:22)

Added Huggingface QuestionAnsweringTranslator
Adds question answering Translator that leverage Huggingface tokenizer.

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

```diff
@@             Coverage Diff              @@
##             master    #1828      +/-   ##
============================================
- Coverage     72.08%   69.71%   -2.38%     
- Complexity     5126     5521     +395     
============================================
  Files           473      526      +53     
  Lines         21970    24421    +2451     
  Branches       2351     2657     +306     
============================================
+ Hits          15838    17026    +1188     
- Misses         4925     6083    +1158     
- Partials       1207     1312     +105     
```


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

------

[Continue to review full report at Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1828?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1828?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Last update [587763a...a584e8e](https://codecov.io/gh/deepjavalibrary/djl/pull/1828?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/extensions/tokenizers/build.gradle b/extensions/tokenizers/build.gradle
index 017228b86..8aeceb893 100644
--- a/extensions/tokenizers/build.gradle
+++ b/extensions/tokenizers/build.gradle
@@ -6,8 +6,8 @@ dependencies {
     testImplementation("org.testng:testng:${testng_version}") {
         exclude group: "junit", module: "junit"
     }
-
-    testImplementation(project(":testing"))
+    testImplementation project(":engines:pytorch:pytorch-engine")
+    testImplementation project(":testing")
     testImplementation "org.slf4j:slf4j-simple:${slf4j_version}"
 }
 
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslator.java
new file mode 100644
index 000000000..ebc3798bd
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslator.java
@@ -0,0 +1,207 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.huggingface.translator;
+
+import ai.djl.huggingface.tokenizers.Encoding;
+import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
+import ai.djl.modality.nlp.qa.QAInput;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.translate.ArgumentsUtil;
+import ai.djl.translate.Batchifier;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorContext;
+
+import java.io.IOException;
+import java.nio.file.Path;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+/** The translator for Huggingface question answering model. */
+public class QuestionAnsweringTranslator implements Translator<QAInput, String> {
+
+    private HuggingFaceTokenizer tokenizer;
+    private boolean includeTokenTypes;
+    private Batchifier batchifier;
+
+    QuestionAnsweringTranslator(
+            HuggingFaceTokenizer tokenizer, boolean includeTokenTypes, Batchifier batchifier) {
+        this.tokenizer = tokenizer;
+        this.includeTokenTypes = includeTokenTypes;
+        this.batchifier = batchifier;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Batchifier getBatchifier() {
+        return batchifier;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, QAInput input) {
+        NDManager manager = ctx.getNDManager();
+        Encoding encoding = tokenizer.encode(input.getQuestion(), input.getParagraph());
+        String[] tokens = encoding.getTokens();
+        ctx.setAttachment("tokens", Arrays.asList(tokens));
+        long[] indices = encoding.getIds();
+        long[] attentionMask = encoding.getAttentionMask();
+        NDList ndList = new NDList(3);
+        ndList.add(manager.create(indices));
+        ndList.add(manager.create(attentionMask));
+        if (includeTokenTypes) {
+            long[] typeIds = encoding.getTypeIds();
+            ndList.add(manager.create(typeIds));
+        }
+        return ndList;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    @SuppressWarnings("unchecked")
+    public String processOutput(TranslatorContext ctx, NDList list) {
+        NDArray startLogits = list.get(0);
+        NDArray endLogits = list.get(1);
+        int startIdx = (int) startLogits.argMax().getLong();
+        int endIdx = (int) endLogits.argMax().getLong();
+        if (startIdx > endIdx) {
+            int tmp = startIdx;
+            startIdx = endIdx;
+            endIdx = tmp;
+        }
+        List<String> tokens = (List<String>) ctx.getAttachment("tokens");
+        return tokenizer.buildSentence(tokens.subList(startIdx, endIdx + 1));
+    }
+
+    /**
+     * Creates a builder to build a {@code QuestionAnsweringTranslator}.
+     *
+     * @return a new builder
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /**
+     * Creates a builder to build a {@code QuestionAnsweringTranslator}.
+     *
+     * @param arguments the models' arguments
+     * @return a new builder
+     */
+    public static Builder builder(Map<String, ?> arguments) {
+        Builder builder = builder();
+        builder.configure(arguments);
+
+        return builder;
+    }
+
+    /** The builder for question answering translator. */
+    public static final class Builder {
+
+        private String tokenizerName;
+        private Path tokenizerPath;
+        private boolean includeTokenTypes;
+        private boolean addSpecialTokens = true;
+        private Batchifier batchifier = Batchifier.STACK;
+
+        /**
+         * Sets the name of the tokenizer for the {@link Translator}.
+         *
+         * @param tokenizerName the name of the tokenizer
+         * @return this builder
+         */
+        public Builder optTokenizerName(String tokenizerName) {
+            this.tokenizerName = tokenizerName;
+            return this;
+        }
+
+        /**
+         * Sets the file path of the tokenizer for the {@link Translator}.
+         *
+         * @param tokenizerPath the name of the tokenizer
+         * @return this builder
+         */
+        public Builder optTokenizerPath(Path tokenizerPath) {
+            this.tokenizerPath = tokenizerPath;
+            return this;
+        }
+
+        /**
+         * Sets if add special tokens for the {@link Translator}.
+         *
+         * @param addSpecialTokens true to add special tokens
+         * @return this builder
+         */
+        public Builder optAddSpecialTokens(boolean addSpecialTokens) {
+            this.addSpecialTokens = addSpecialTokens;
+            return this;
+        }
+
+        /**
+         * Sets if include token types for the {@link Translator}.
+         *
+         * @param includeTokenTypes true to include token types
+         * @return this builder
+         */
+        public Builder optIncludeTokenTypes(boolean includeTokenTypes) {
+            this.includeTokenTypes = includeTokenTypes;
+            return this;
+        }
+
+        /**
+         * Sets the {@link Batchifier} for the {@link Translator}.
+         *
+         * @param batchifier true to include token types
+         * @return this builder
+         */
+        public Builder optBatchifer(Batchifier batchifier) {
+            this.batchifier = batchifier;
+            return this;
+        }
+
+        /**
+         * Configures the builder with the model arguments.
+         *
+         * @param arguments the model arguments
+         */
+        public void configure(Map<String, ?> arguments) {
+            optTokenizerName(ArgumentsUtil.stringValue(arguments, "tokenizer"));
+            optIncludeTokenTypes(ArgumentsUtil.booleanValue(arguments, "includeTokenTypes"));
+            optAddSpecialTokens(ArgumentsUtil.booleanValue(arguments, "addSpecialTokens", true));
+            optBatchifer(
+                    Batchifier.fromString(
+                            ArgumentsUtil.stringValue(arguments, "batchifier", "stack")));
+        }
+
+        /**
+         * Builds the translator.
+         *
+         * @return the new translator
+         * @throws IOException if I/O error occurs
+         */
+        public QuestionAnsweringTranslator build() throws IOException {
+            HuggingFaceTokenizer tokenizer;
+            Map<String, String> options = new ConcurrentHashMap<>();
+            options.put("addSpecialTokens", String.valueOf(addSpecialTokens));
+            if (tokenizerName != null) {
+                tokenizer = HuggingFaceTokenizer.newInstance(tokenizerName, options);
+            } else {
+                tokenizer = HuggingFaceTokenizer.newInstance(tokenizerPath, options);
+            }
+            return new QuestionAnsweringTranslator(tokenizer, includeTokenTypes, batchifier);
+        }
+    }
+}
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
new file mode 100644
index 000000000..44f6325a1
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslatorFactory.java
@@ -0,0 +1,69 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.huggingface.translator;
+
+import ai.djl.Model;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.nlp.qa.QAInput;
+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;
+
+import java.io.IOException;
+import java.lang.reflect.Type;
+import java.nio.file.Path;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+/** A {@link TranslatorFactory} that creates a {@link QuestionAnsweringTranslator} instance. */
+public class QuestionAnsweringTranslatorFactory implements TranslatorFactory {
+
+    private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>();
+
+    static {
+        SUPPORTED_TYPES.add(new Pair<>(QAInput.class, String.class));
+        SUPPORTED_TYPES.add(new Pair<>(Input.class, Output.class));
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Set<Pair<Type, Type>> getSupportedTypes() {
+        return SUPPORTED_TYPES;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Translator<?, ?> newInstance(
+            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
+            throws TranslateException {
+        Path modelPath = model.getModelPath();
+        try {
+            QuestionAnsweringTranslator translator =
+                    QuestionAnsweringTranslator.builder(arguments)
+                            .optTokenizerPath(modelPath)
+                            .build();
+            if (input == QAInput.class && output == String.class) {
+                return translator;
+            } else if (input == Input.class && output == Output.class) {
+                return new QaServingTranslator(translator);
+            }
+            throw new IllegalArgumentException("Unsupported input/output types.");
+        } catch (IOException e) {
+            throw new TranslateException("Failed to load tokenizer.", e);
+        }
+    }
+}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/package-info.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/package-info.java
new file mode 100644
index 000000000..b1a0e422d
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/package-info.java
@@ -0,0 +1,15 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+
+/** Contains {@link ai.djl.translate.Translator} classes that leverage Huggingface tokenizers. */
+package ai.djl.huggingface.translator;


Test patch:
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TranslatorTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TranslatorTest.java
new file mode 100644
index 000000000..46ea82cfb
--- /dev/null
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TranslatorTest.java
@@ -0,0 +1,97 @@
+/*
+ * 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.tokenizers;
+
+import ai.djl.Model;
+import ai.djl.ModelException;
+import ai.djl.huggingface.translator.QuestionAnsweringTranslatorFactory;
+import ai.djl.inference.Predictor;
+import ai.djl.modality.nlp.qa.QAInput;
+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.LambdaBlock;
+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.Utils;
+
+import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+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 TranslatorTest {
+
+    @AfterClass
+    public void tierDown() {
+        Utils.deleteQuietly(Paths.get("build/model"));
+    }
+
+    @Test
+    public void testQATranslator() throws ModelException, IOException, TranslateException {
+        TestRequirements.notArm();
+
+        String question = "When did BBC Japan start broadcasting?";
+        String paragraph =
+                "BBC Japan was a general entertainment Channel. "
+                        + "Which operated between December 2004 and April 2006. "
+                        + "It ceased operations after its Japanese distributor folded.";
+        QAInput input = new QAInput(question, paragraph);
+
+        Block block =
+                new LambdaBlock(
+                        a -> {
+                            NDManager manager = a.getManager();
+                            long[][] start = new long[1][36];
+                            long[][] end = new long[1][36];
+                            start[0][20] = 1;
+                            end[0][21] = 1;
+                            NDArray arr1 = manager.create(start);
+                            NDArray arr2 = manager.create(end);
+                            return new NDList(arr1, arr2);
+                        },
+                        "model");
+        Path modelDir = Paths.get("build/model");
+        Files.createDirectories(modelDir);
+        Model testModel = Model.newInstance("model");
+        testModel.setBlock(block);
+        block.initialize(testModel.getNDManager(), DataType.FLOAT32, new Shape(1));
+        testModel.save(modelDir, "model");
+        testModel.close();
+
+        Criteria<QAInput, String> criteria =
+                Criteria.builder()
+                        .setTypes(QAInput.class, String.class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-cased")
+                        .optTranslatorFactory(new QuestionAnsweringTranslatorFactory())
+                        .build();
+
+        try (ZooModel<QAInput, String> model = criteria.loadModel();
+                Predictor<QAInput, String> predictor = model.newPredictor()) {
+            String res = predictor.predict(input);
+            Assert.assertEquals(res, "December 2004");
+        }
+    }
+}