Repository: deepjavalibrary/djl
Pull request #1906 (2022-08-15 00:59:53)

Adds huggingface TokenClassificationTranslator

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

```diff
@@             Coverage Diff              @@
##             master    #1906      +/-   ##
============================================
- Coverage     72.08%   69.86%   -2.23%     
- Complexity     5126     5677     +551     
============================================
  Files           473      543      +70     
  Lines         21970    25069    +3099     
  Branches       2351     2712     +361     
============================================
+ Hits          15838    17515    +1677     
- Misses         4925     6216    +1291     
- Partials       1207     1338     +131     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1906?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/1906/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/1906/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/1906/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/1906/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/1906/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/1906/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/1906/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/1906/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/1906/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/1906/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 [443 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1906/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). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)


Patch:
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/NamedEntity.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/NamedEntity.java
new file mode 100644
index 000000000..19325d5e5
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/NamedEntity.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.translator;
+
+/** A class that represents a {@code NamedEntity} object. */
+public class NamedEntity {
+
+    private String entity;
+    private float score;
+    private int index;
+    private String word;
+    private int start;
+    private int end;
+
+    /**
+     * Constructs a new instance of {@code NamedEntity}.
+     *
+     * @param entity the class of the entity
+     * @param score the score of the entity
+     * @param index the position of the entity in the original sentence
+     * @param word the token of the entity
+     * @param start the start index of the word in the sentence
+     * @param end the end index of the word in the sentence
+     */
+    public NamedEntity(String entity, float score, int index, String word, int start, int end) {
+        this.entity = entity;
+        this.score = score;
+        this.index = index;
+        this.word = word;
+        this.start = start;
+        this.end = end;
+    }
+
+    /**
+     * Returns the class of the entity.
+     *
+     * @return the class of the entity
+     */
+    public String getEntity() {
+        return entity;
+    }
+
+    /**
+     * Returns the score of the entity.
+     *
+     * @return the score of the entity
+     */
+    public float getScore() {
+        return score;
+    }
+
+    /**
+     * Returns the position of the entity in the original sentence.
+     *
+     * @return the position of the entity in the original sentence
+     */
+    public int getIndex() {
+        return index;
+    }
+
+    /**
+     * Returns the token of the entity.
+     *
+     * @return the token of the entity
+     */
+    public String getWord() {
+        return word;
+    }
+
+    /**
+     * Returns the start index of the word in the sentence.
+     *
+     * @return the start index of the word in the sentence
+     */
+    public int getStart() {
+        return start;
+    }
+
+    /**
+     * Returns the end index of the word in the sentence.
+     *
+     * @return the end index of the word in the sentence
+     */
+    public int getEnd() {
+        return end;
+    }
+}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationServingTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationServingTranslator.java
new file mode 100644
index 000000000..172b836b4
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationServingTranslator.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.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.ndarray.BytesSupplier;
+import ai.djl.ndarray.NDList;
+import ai.djl.translate.Batchifier;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorContext;
+
+/**
+ * A {@link Translator} that can handle generic token classification {@link Input} and {@link
+ * Output}.
+ */
+public class TokenClassificationServingTranslator implements Translator<Input, Output> {
+
+    private Translator<String, NamedEntity[]> translator;
+
+    /**
+     * Constructs a {@code TokenClassificationServingTranslator} instance.
+     *
+     * @param translator a {@code Translator} processes token classification input
+     */
+    public TokenClassificationServingTranslator(Translator<String, NamedEntity[]> translator) {
+        this.translator = translator;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Batchifier getBatchifier() {
+        return translator.getBatchifier();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void prepare(TranslatorContext ctx) throws Exception {
+        translator.prepare(ctx);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, Input input) throws Exception {
+        String text = input.getData().getAsString();
+        return translator.processInput(ctx, text);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Output processOutput(TranslatorContext ctx, NDList list) throws Exception {
+        NamedEntity[] ret = translator.processOutput(ctx, list);
+        Output output = new Output();
+        output.add(BytesSupplier.wrapAsJson(ret));
+        return output;
+    }
+}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslator.java
new file mode 100644
index 000000000..5249f358b
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslator.java
@@ -0,0 +1,219 @@
+/*
+ * 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.huggingface.tokenizers.jni.CharSpan;
+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 ai.djl.util.JsonUtils;
+
+import java.io.IOException;
+import java.io.Reader;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+/** The translator for Huggingface token classification model. */
+public class TokenClassificationTranslator implements Translator<String, NamedEntity[]> {
+
+    private HuggingFaceTokenizer tokenizer;
+    private Batchifier batchifier;
+    private Config config;
+
+    TokenClassificationTranslator(HuggingFaceTokenizer tokenizer, Batchifier batchifier) {
+        this.tokenizer = tokenizer;
+        this.batchifier = batchifier;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Batchifier getBatchifier() {
+        return batchifier;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void prepare(TranslatorContext ctx) throws IOException {
+        Path path = ctx.getModel().getModelPath();
+        Path file = path.resolve("config.json");
+        try (Reader reader = Files.newBufferedReader(file)) {
+            config = JsonUtils.GSON.fromJson(reader, Config.class);
+        }
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, String input) {
+        NDManager manager = ctx.getNDManager();
+        Encoding encoding = tokenizer.encode(input);
+        ctx.setAttachment("encoding", encoding);
+        long[] indices = encoding.getIds();
+        long[] attentionMask = encoding.getAttentionMask();
+        NDList ndList = new NDList(2);
+        ndList.add(manager.create(indices));
+        ndList.add(manager.create(attentionMask));
+        return ndList;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NamedEntity[] processOutput(TranslatorContext ctx, NDList list) {
+        NDArray logits = list.get(0);
+        Encoding encoding = (Encoding) ctx.getAttachment("encoding");
+        long[] inputIds = encoding.getIds();
+        CharSpan[] offsetMapping = encoding.getCharTokenSpans();
+        long[] specialTokenMasks = encoding.getSpecialTokenMask();
+        NDArray probabilities = logits.softmax(1);
+        List<NamedEntity> entities = new ArrayList<>();
+
+        for (int i = 0; i < inputIds.length; ++i) {
+            if (specialTokenMasks[i] != 0) {
+                continue;
+            }
+
+            int entityIdx = (int) probabilities.get(i).argMax().getLong();
+            String entity = config.id2label.get(String.valueOf(entityIdx));
+
+            if (!"O".equals(entity)) {
+                float score = probabilities.get(i).getFloat(entityIdx);
+                String word = encoding.getTokens()[i];
+                int start = offsetMapping[i].getStart();
+                int end = offsetMapping[i].getEnd();
+
+                NamedEntity item = new NamedEntity(entity, score, i, word, start, end);
+                entities.add(item);
+            }
+        }
+        return entities.toArray(new NamedEntity[0]);
+    }
+
+    /**
+     * Creates a builder to build a {@code TokenClassificationTranslator}.
+     *
+     * @return a new builder
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /**
+     * Creates a builder to build a {@code TokenClassificationTranslator}.
+     *
+     * @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 token classification translator. */
+    public static final class Builder {
+
+        private String tokenizerName;
+        private Path tokenizerPath;
+        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 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) {
+            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 TokenClassificationTranslator 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 TokenClassificationTranslator(tokenizer, batchifier);
+        }
+    }
+
+    private static final class Config {
+
+        Map<String, String> id2label;
+    }
+}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslatorFactory.java
new file mode 100644
index 000000000..d5bf25518
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslatorFactory.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.Input;
+import ai.djl.modality.Output;
+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 TokenClassificationTranslator} instance. */
+public class TokenClassificationTranslatorFactory implements TranslatorFactory {
+
+    private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>();
+
+    static {
+        SUPPORTED_TYPES.add(new Pair<>(String.class, NamedEntity[].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 {
+            TokenClassificationTranslator translator =
+                    TokenClassificationTranslator.builder(arguments)
+                            .optTokenizerPath(modelPath)
+                            .build();
+            if (input == String.class && output == NamedEntity[].class) {
+                return translator;
+            } else if (input == Input.class && output == Output.class) {
+                return new TokenClassificationServingTranslator(translator);
+            }
+            throw new IllegalArgumentException("Unsupported input/output types.");
+        } catch (IOException e) {
+            throw new TranslateException("Failed to load tokenizer.", e);
+        }
+    }
+}


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 c4f05d446..72decd31b 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
@@ -15,7 +15,9 @@ package ai.djl.huggingface.tokenizers;
 import ai.djl.Model;
 import ai.djl.ModelException;
 import ai.djl.huggingface.translator.FillMaskTranslatorFactory;
+import ai.djl.huggingface.translator.NamedEntity;
 import ai.djl.huggingface.translator.QuestionAnsweringTranslatorFactory;
+import ai.djl.huggingface.translator.TokenClassificationTranslatorFactory;
 import ai.djl.inference.Predictor;
 import ai.djl.modality.Classifications;
 import ai.djl.modality.Input;
@@ -28,8 +30,10 @@ 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.Assertions;
 import ai.djl.testing.TestRequirements;
 import ai.djl.translate.TranslateException;
+import ai.djl.util.JsonUtils;
 import ai.djl.util.Utils;
 
 import org.testng.Assert;
@@ -37,6 +41,7 @@ import org.testng.annotations.AfterClass;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
+import java.io.Writer;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -217,7 +222,103 @@ public class TranslatorTest {
                     TranslateException.class,
                     () -> factory.newInstance(String.class, Integer.class, model, arguments));
 
-            arguments.put("tokenizer", "bert-base-cased");
+            arguments.put("tokenizer", "bert-base-uncased");
+
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> factory.newInstance(String.class, Integer.class, model, arguments));
+        }
+    }
+
+    @Test
+    public void testTokenClassificationTranslator()
+            throws ModelException, IOException, TranslateException {
+        TestRequirements.notArm();
+
+        String text = "My name is Wolfgang and I live in Berlin.";
+
+        Block block =
+                new LambdaBlock(
+                        a -> {
+                            NDManager manager = a.getManager();
+                            float[][] logits = new float[12][9];
+                            logits[4][3] = 1;
+                            logits[9][7] = 1;
+                            NDArray arr = manager.create(logits);
+                            arr = arr.expandDims(0);
+                            return new NDList(arr);
+                        },
+                        "model");
+        Path modelDir = Paths.get("build/model");
+        Files.createDirectories(modelDir);
+        Path path = modelDir.resolve("config.json");
+        Map<String, Map<String, String>> map = new HashMap<>();
+        Map<String, String> id2label = new HashMap<>();
+        id2label.put("0", "O");
+        id2label.put("3", "B-PER");
+        id2label.put("7", "B-LOC");
+        map.put("id2label", id2label);
+        try (Writer writer = Files.newBufferedWriter(path)) {
+            writer.write(JsonUtils.GSON.toJson(map));
+        }
+
+        Criteria<String, NamedEntity[]> criteria =
+                Criteria.builder()
+                        .setTypes(String.class, NamedEntity[].class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-uncased")
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new TokenClassificationTranslatorFactory())
+                        .build();
+
+        try (ZooModel<String, NamedEntity[]> model = criteria.loadModel();
+                Predictor<String, NamedEntity[]> predictor = model.newPredictor()) {
+            NamedEntity[] res = predictor.predict(text);
+            Assert.assertEquals(res[0].getEntity(), "B-PER");
+            Assertions.assertAlmostEquals(res[0].getScore(), 0.2536117);
+            Assert.assertEquals(res[0].getIndex(), 4);
+            Assert.assertEquals(res[0].getWord(), "wolfgang");
+            Assert.assertEquals(res[0].getStart(), 11);
+            Assert.assertEquals(res[0].getEnd(), 19);
+        }
+
+        Criteria<Input, Output> criteria2 =
+                Criteria.builder()
+                        .setTypes(Input.class, Output.class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-uncased")
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new TokenClassificationTranslatorFactory())
+                        .build();
+
+        try (ZooModel<Input, Output> model = criteria2.loadModel();
+                Predictor<Input, Output> predictor = model.newPredictor()) {
+            Input input = new Input();
+            input.add(text);
+            Output out = predictor.predict(input);
+            NamedEntity[] res = (NamedEntity[]) out.getData().getAsObject();
+            Assert.assertEquals(res[0].getEntity(), "B-PER");
+        }
+
+        try (Model model = Model.newInstance("test")) {
+            model.setBlock(block);
+            Map<String, String> options = new HashMap<>();
+            options.put("hasParameter", "false");
+            model.load(modelDir, "test", options);
+
+            TokenClassificationTranslatorFactory factory =
+                    new TokenClassificationTranslatorFactory();
+            Map<String, String> arguments = new HashMap<>();
+
+            Assert.assertThrows(
+                    TranslateException.class,
+                    () -> factory.newInstance(String.class, Integer.class, model, arguments));
+
+            arguments.put("tokenizer", "bert-base-uncased");
 
             Assert.assertThrows(
                     IllegalArgumentException.class,