Repository: deepjavalibrary/djl
Pull request #1876 (2022-08-07 06:00:10)

Adds huggingface FillMaskTranslator
Adds huggingface `FillMaskTranslator` 

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

```diff
@@             Coverage Diff              @@
##             master    #1876      +/-   ##
============================================
- Coverage     72.08%   69.81%   -2.28%     
- Complexity     5126     5561     +435     
============================================
  Files           473      529      +56     
  Lines         21970    24603    +2633     
  Branches       2351     2676     +325     
============================================
+ Hits          15838    17176    +1338     
- Misses         4925     6112    +1187     
- Partials       1207     1315     +108     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1876?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/1876/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/1876/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/1876/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/1876/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/1876/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/1876/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/1876/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/1876/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/1876/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/1876/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 [425 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1876/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
index e6c128728..01cb98991 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
@@ -1055,10 +1055,13 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
     /** {@inheritDoc} */
     @Override
     public PtNDArray argSort(int axis, boolean ascending) {
-        if (!ascending) {
-            throw new UnsupportedOperationException("Only support ascending!");
+        PtNDArray arr = JniUtils.argSort(this, axis, false);
+        if (ascending) {
+            return arr;
         }
-        return JniUtils.argSort(this, axis, false);
+        PtNDArray flip = JniUtils.flip(arr, new long[] {axis});
+        arr.close();
+        return flip;
     }
 
     /** {@inheritDoc} */
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslator.java
new file mode 100644
index 000000000..884460d8a
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslator.java
@@ -0,0 +1,232 @@
+/*
+ * 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.Classifications;
+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.TranslateException;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorContext;
+
+import java.io.IOException;
+import java.nio.file.Path;
+import java.util.ArrayList;
+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 FillMaskTranslator implements Translator<String, Classifications> {
+
+    private HuggingFaceTokenizer tokenizer;
+    private String maskToken;
+    private long maskTokenId;
+    private int topK;
+    private Batchifier batchifier;
+
+    FillMaskTranslator(
+            HuggingFaceTokenizer tokenizer, String maskToken, int topK, Batchifier batchifier) {
+        this.tokenizer = tokenizer;
+        this.maskToken = maskToken;
+        this.topK = topK;
+        this.batchifier = batchifier;
+        Encoding encoding = tokenizer.encode(maskToken, false);
+        maskTokenId = encoding.getIds()[0];
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Batchifier getBatchifier() {
+        return batchifier;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, String input) throws TranslateException {
+        NDManager manager = ctx.getNDManager();
+        Encoding encoding = tokenizer.encode(input);
+        long[] indices = encoding.getIds();
+        int maskIndex = -1;
+        for (int i = 0; i < indices.length; ++i) {
+            if (indices[i] == maskTokenId) {
+                if (maskIndex != -1) {
+                    throw new TranslateException("Only one mask supported.");
+                }
+                maskIndex = i;
+            }
+        }
+        if (maskIndex == -1) {
+            throw new TranslateException("Mask token " + maskToken + " not found.");
+        }
+        ctx.setAttachment("maskIndex", maskIndex);
+        long[] attentionMask = encoding.getAttentionMask();
+        NDList ndList = new NDList(2);
+        ndList.add(manager.create(indices));
+        ndList.add(manager.create(attentionMask));
+        return ndList;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Classifications processOutput(TranslatorContext ctx, NDList list) {
+        int maskIndex = (int) ctx.getAttachment("maskIndex");
+        NDArray prob = list.get(0).get(maskIndex).softmax(0);
+        NDArray array = prob.argSort(0, false);
+        long[] classIds = new long[topK];
+        List<Double> probabilities = new ArrayList<>(topK);
+        for (int i = 0; i < topK; ++i) {
+            classIds[i] = array.getLong(i);
+            probabilities.add((double) prob.getFloat(classIds[i]));
+        }
+        String[] classes = tokenizer.decode(classIds).trim().split(" ");
+        return new Classifications(Arrays.asList(classes), probabilities);
+    }
+
+    /**
+     * 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 maskedToken = "[MASK]";
+        private int topK = 5;
+        private String tokenizerName;
+        private Path tokenizerPath;
+        private boolean addSpecialTokens = true;
+        private Batchifier batchifier = Batchifier.STACK;
+
+        /**
+         * Sets the id of the mask {@link Translator}.
+         *
+         * @param maskedToken the id of the mask
+         * @return this builder
+         */
+        public Builder optMaskToken(String maskedToken) {
+            this.maskedToken = maskedToken;
+            return this;
+        }
+
+        /**
+         * Set the topK number of classes to be displayed.
+         *
+         * @param topK the number of top classes to return
+         * @return this builder
+         */
+        public Builder optTopK(int topK) {
+            this.topK = topK;
+            return this;
+        }
+
+        /**
+         * 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 path 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 the {@link Batchifier} for the {@link Translator}.
+         *
+         * @param batchifier true to include token types
+         * @return this builder
+         */
+        public Builder optBatchifier(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) {
+            optMaskToken(ArgumentsUtil.stringValue(arguments, "maskToken", "[MASK]"));
+            optTopK(ArgumentsUtil.intValue(arguments, "topK", 5));
+            optTokenizerName(ArgumentsUtil.stringValue(arguments, "tokenizer"));
+            optAddSpecialTokens(ArgumentsUtil.booleanValue(arguments, "addSpecialTokens", true));
+            String batchifierStr = ArgumentsUtil.stringValue(arguments, "batchifier", "stack");
+            optBatchifier(Batchifier.fromString(batchifierStr));
+        }
+
+        /**
+         * Builds the translator.
+         *
+         * @return the new translator
+         * @throws IOException if I/O error occurs
+         */
+        public FillMaskTranslator 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 FillMaskTranslator(tokenizer, maskedToken, topK, batchifier);
+        }
+    }
+}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslatorFactory.java
new file mode 100644
index 000000000..d3db2dab1
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslatorFactory.java
@@ -0,0 +1,67 @@
+/*
+ * 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.Classifications;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.nlp.translator.TextClassificationServingTranslator;
+import ai.djl.translate.TranslateException;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorFactory;
+import ai.djl.util.Pair;
+
+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 FillMaskTranslator} instance. */
+public class FillMaskTranslatorFactory implements TranslatorFactory {
+
+    private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>();
+
+    static {
+        SUPPORTED_TYPES.add(new Pair<>(String.class, Classifications.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 {
+            FillMaskTranslator translator =
+                    FillMaskTranslator.builder(arguments).optTokenizerPath(modelPath).build();
+            if (input == String.class && output == Classifications.class) {
+                return translator;
+            } else if (input == Input.class && output == Output.class) {
+                return new TextClassificationServingTranslator(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/QuestionAnsweringTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslator.java
index c8f7a2669..5dcc7ccab 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslator.java
@@ -137,7 +137,7 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
         /**
          * Sets the file path of the tokenizer for the {@link Translator}.
          *
-         * @param tokenizerPath the name of the tokenizer
+         * @param tokenizerPath the path of the tokenizer
          * @return this builder
          */
         public Builder optTokenizerPath(Path tokenizerPath) {
@@ -173,7 +173,7 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
          * @param batchifier true to include token types
          * @return this builder
          */
-        public Builder optBatchifer(Batchifier batchifier) {
+        public Builder optBatchifier(Batchifier batchifier) {
             this.batchifier = batchifier;
             return this;
         }
@@ -187,9 +187,8 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
             optTokenizerName(ArgumentsUtil.stringValue(arguments, "tokenizer"));
             optIncludeTokenTypes(ArgumentsUtil.booleanValue(arguments, "includeTokenTypes"));
             optAddSpecialTokens(ArgumentsUtil.booleanValue(arguments, "addSpecialTokens", true));
-            optBatchifer(
-                    Batchifier.fromString(
-                            ArgumentsUtil.stringValue(arguments, "batchifier", "stack")));
+            String batchifierStr = ArgumentsUtil.stringValue(arguments, "batchifier", "stack");
+            optBatchifier(Batchifier.fromString(batchifierStr));
         }
 
         /**
diff --git a/extensions/tokenizers/src/main/python/.gitignore b/extensions/tokenizers/src/main/python/.gitignore
index bee8a64b7..a63c8dced 100644
--- a/extensions/tokenizers/src/main/python/.gitignore
+++ b/extensions/tokenizers/src/main/python/.gitignore
@@ -1,1 +1,3 @@
 __pycache__
+model/
+processed_models.json
diff --git a/extensions/tokenizers/src/main/python/huggingface_converter.py b/extensions/tokenizers/src/main/python/huggingface_converter.py
index f81404c86..9023d8553 100644
--- a/extensions/tokenizers/src/main/python/huggingface_converter.py
+++ b/extensions/tokenizers/src/main/python/huggingface_converter.py
@@ -146,7 +146,8 @@ class HuggingfaceConverter:
         if not result:
             return False, reason, -1
 
-        size = self.save_model_zoo(task, model_id, temp_dir)
+        size = self.save_model_zoo(task, model_id, temp_dir,
+                                   hf_pipeline.tokenizer.mask_token)
 
         return True, None, size
 
@@ -183,7 +184,8 @@ class HuggingfaceConverter:
         script_module.save(model_file)
         return model_file
 
-    def save_model_zoo(self, task: str, model_id: str, temp_dir: str):
+    def save_model_zoo(self, task: str, model_id: str, temp_dir: str,
+                       mask_token: str):
         artifact_ids = model_id.split("/")
         model_name = artifact_ids[-1]
 
@@ -197,9 +199,11 @@ class HuggingfaceConverter:
         # Save serving.properties
         serving_file = os.path.join(temp_dir, "serving.properties")
         with open(serving_file, 'w') as f:
-            f.write(
-                f"engine=PyTorch\noption.modelName={model_name}\noption.mapLocation=true\n"
-                f"translatorFactory={translator}")
+            f.write(f"engine=PyTorch\n"
+                    f"option.modelName={model_name}\n"
+                    f"option.mapLocation=true\n"
+                    f"maskToken={mask_token}\n"
+                    f"translatorFactory={translator}")
 
         # Save model as .zip file
         zip_file = os.path.join(model_dir, f"{model_name}.zip")
@@ -257,11 +261,10 @@ class HuggingfaceConverter:
                         f"pipeline output differs from expected: {pipeline_output}"
                     )
         elif task == "fill-mask":
-            masked_token_id = tokenizer.mask_token_id
-            masked_index = torch.nonzero(
-                input_ids.squeeze(0) == masked_token_id,
-                as_tuple=False).squeeze(0)
-            logits = out['logits'][0, masked_index]
+            mask_token_id = tokenizer.mask_token_id
+            mask_index = torch.nonzero(input_ids.squeeze(0) == mask_token_id,
+                                       as_tuple=False).squeeze(0)
+            logits = out['logits'][0, mask_index]
             answer = torch.argmax(logits)
             prediction = tokenizer.decode(answer).strip()
 


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
index 5dd79f231..fa85d05e8 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TranslatorTest.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TranslatorTest.java
@@ -13,8 +13,10 @@
 package ai.djl.huggingface.tokenizers;
 
 import ai.djl.ModelException;
+import ai.djl.huggingface.translator.FillMaskTranslatorFactory;
 import ai.djl.huggingface.translator.QuestionAnsweringTranslatorFactory;
 import ai.djl.inference.Predictor;
+import ai.djl.modality.Classifications;
 import ai.djl.modality.nlp.qa.QAInput;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
@@ -88,4 +90,45 @@ public class TranslatorTest {
             Assert.assertEquals(res, "December 2004");
         }
     }
+
+    @Test
+    public void testFillMaskTranslator() throws ModelException, IOException, TranslateException {
+        TestRequirements.notArm();
+
+        String input = "Hello I'm a [MASK] model.";
+
+        Block block =
+                new LambdaBlock(
+                        a -> {
+                            NDManager manager = a.getManager();
+                            float[][] logits = new float[10][4828];
+                            logits[6][4827] = 5;
+                            logits[6][2535] = 4;
+                            logits[6][2047] = 3;
+                            logits[6][3565] = 2;
+                            logits[6][2986] = 1;
+                            NDArray arr = manager.create(logits);
+                            arr = arr.expandDims(0);
+                            return new NDList(arr);
+                        },
+                        "model");
+        Path modelDir = Paths.get("build/model");
+        Files.createDirectories(modelDir);
+        Criteria<String, Classifications> criteria =
+                Criteria.builder()
+                        .setTypes(String.class, Classifications.class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-uncased")
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new FillMaskTranslatorFactory())
+                        .build();
+
+        try (ZooModel<String, Classifications> model = criteria.loadModel();
+                Predictor<String, Classifications> predictor = model.newPredictor()) {
+            Classifications res = predictor.predict(input);
+            Assert.assertEquals(res.best().getClassName(), "fashion");
+        }
+    }
 }