Repository: deepjavalibrary/djl
Pull request #1940 (2022-08-23 16:54:47)

[tokenizer] Use builder for HuggingfaceTokenizer
1. Add builder for HuggingfaceTokenizer
2. Make HuggingfaceTokenizer immutable

Change-Id: I84724b83051cb2079046b49893d96485d8fdcd3b

## 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/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizer.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizer.java
index 8ff7edd2b..33d8190cc 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizer.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizer.java
@@ -30,6 +30,7 @@ import java.nio.file.Path;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * {@code HuggingFaceTokenizer} is a Huggingface tokenizer implementation of the {@link Tokenizer}
@@ -338,81 +339,25 @@ public final class HuggingFaceTokenizer extends NativeResource<Long> implements
     }
 
     /**
-     * Enables or Disables default truncation behavior for the tokenizer.
+     * Creates a builder to build a {@code HuggingFaceTokenizer}.
      *
-     * @param enabled whether to enable default truncation behavior
+     * @return a new builder
      */
-    public void setTruncation(boolean enabled) {
-        TruncationStrategy strategy =
-                enabled ? TruncationStrategy.LONGEST_FIRST : TruncationStrategy.DO_NOT_TRUNCATE;
-        if (strategy != truncation) {
-            truncation = strategy;
-            updateTruncationAndPadding();
-        }
-    }
-
-    /** Enables truncation to only truncate the first item. */
-    public void setTruncateFirstOnly() {
-        if (truncation != TruncationStrategy.ONLY_FIRST) {
-            truncation = TruncationStrategy.ONLY_FIRST;
-            updateTruncationAndPadding();
-        }
-    }
-
-    /** Enables truncation to only truncate the second item. */
-    public void setTruncateSecondOnly() {
-        if (truncation != TruncationStrategy.ONLY_SECOND) {
-            truncation = TruncationStrategy.ONLY_SECOND;
-            updateTruncationAndPadding();
-        }
+    public static Builder builder() {
+        return new Builder();
     }
 
     /**
-     * Enables or Disables default padding behavior for the tokenizer.
+     * Creates a builder to build a {@code HuggingFaceTokenizer}.
      *
-     * @param enabled whether to enable default padding behavior
+     * @param arguments the models' arguments
+     * @return a new builder
      */
-    public void setPadding(boolean enabled) {
-        PaddingStrategy p = enabled ? PaddingStrategy.LONGEST : PaddingStrategy.DO_NOT_PAD;
-        if (padding != p) {
-            padding = p;
-            updateTruncationAndPadding();
-        }
-    }
+    public static Builder builder(Map<String, ?> arguments) {
+        Builder builder = builder();
+        builder.configure(arguments);
 
-    /**
-     * Enables padding to pad sequences to previously specified maxLength, or modelMaxLength if not
-     * specified.
-     */
-    public void setPaddingMaxLength() {
-        if (padding != PaddingStrategy.MAX_LENGTH) {
-            padding = PaddingStrategy.MAX_LENGTH;
-            updateTruncationAndPadding();
-        }
-    }
-
-    /**
-     * Sets maxLength for padding and truncation.
-     *
-     * @param maxLength the length to truncate and/or pad sequences to
-     */
-    public void setMaxLength(int maxLength) {
-        if (this.maxLength != maxLength) {
-            this.maxLength = maxLength;
-            updateTruncationAndPadding();
-        }
-    }
-
-    /**
-     * Sets padToMultipleOf for padding.
-     *
-     * @param padToMultipleOf the multiple of sequences should be padded to
-     */
-    public void setPadToMultipleOf(int padToMultipleOf) {
-        if (this.padToMultipleOf != padToMultipleOf) {
-            this.padToMultipleOf = padToMultipleOf;
-            updateTruncationAndPadding();
-        }
+        return builder;
     }
 
     /*
@@ -535,4 +480,152 @@ public final class HuggingFaceTokenizer extends NativeResource<Long> implements
             throw new IllegalArgumentException("Invalid PaddingStrategy: " + value);
         }
     }
+
+    /** The builder for creating huggingface tokenizer. */
+    public static final class Builder {
+
+        private Path tokenizerPath;
+        private Map<String, String> options;
+
+        Builder() {
+            options = new ConcurrentHashMap<>();
+            options.put("addSpecialTokens", "true");
+        }
+
+        /**
+         * Sets the name of the tokenizer.
+         *
+         * @param tokenizerName the name of the tokenizer
+         * @return this builder
+         */
+        public Builder optTokenizerName(String tokenizerName) {
+            options.put("tokenizer", tokenizerName);
+            return this;
+        }
+
+        /**
+         * Sets the file path of the tokenizer.
+         *
+         * @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.
+         *
+         * @param addSpecialTokens true to add special tokens
+         * @return this builder
+         */
+        public Builder optAddSpecialTokens(boolean addSpecialTokens) {
+            options.put("addSpecialTokens", String.valueOf(addSpecialTokens));
+            return this;
+        }
+
+        /**
+         * Enables or Disables default truncation behavior for the tokenizer.
+         *
+         * @param enabled whether to enable default truncation behavior
+         * @return this builder
+         */
+        public Builder optTruncation(boolean enabled) {
+            options.put("truncation", String.valueOf(enabled));
+            return this;
+        }
+
+        /**
+         * Enables truncation to only truncate the first item.
+         *
+         * @return this builder
+         */
+        public Builder optTruncateFirstOnly() {
+            options.put("truncation", TruncationStrategy.ONLY_FIRST.name());
+            return this;
+        }
+
+        /**
+         * Enables truncation to only truncate the second item.
+         *
+         * @return this builder
+         */
+        public Builder optTruncateSecondOnly() {
+            options.put("truncation", TruncationStrategy.ONLY_SECOND.name());
+            return this;
+        }
+
+        /**
+         * Enables or Disables default padding behavior for the tokenizer.
+         *
+         * @param enabled whether to enable default padding behavior
+         * @return this builder
+         */
+        public Builder optPadding(boolean enabled) {
+            options.put("padding", String.valueOf(enabled));
+            return this;
+        }
+
+        /**
+         * Enables padding to pad sequences to previously specified maxLength, or modelMaxLength if
+         * not specified.
+         *
+         * @return this builder
+         */
+        public Builder optPadToMaxLength() {
+            options.put("padding", PaddingStrategy.MAX_LENGTH.name());
+            return this;
+        }
+
+        /**
+         * Sets maxLength for padding and truncation.
+         *
+         * @param maxLength the length to truncate and/or pad sequences to
+         * @return this builder
+         */
+        public Builder optMaxLength(int maxLength) {
+            options.put("maxLength", String.valueOf(maxLength));
+            return this;
+        }
+
+        /**
+         * Sets padToMultipleOf for padding.
+         *
+         * @param padToMultipleOf the multiple of sequences should be padded to
+         * @return this builder
+         */
+        public Builder optPadToMultipleOf(int padToMultipleOf) {
+            options.put("padToMultipleOf", String.valueOf(padToMultipleOf));
+            return this;
+        }
+
+        /**
+         * Configures the builder with the arguments.
+         *
+         * @param arguments the arguments
+         */
+        public void configure(Map<String, ?> arguments) {
+            for (Map.Entry<String, ?> entry : arguments.entrySet()) {
+                options.put(entry.getKey(), entry.getValue().toString());
+            }
+        }
+
+        /**
+         * Builds the translator.
+         *
+         * @return the new translator
+         * @throws IOException when IO operation fails in loading a resource
+         */
+        public HuggingFaceTokenizer build() throws IOException {
+            String tokenizerName = options.get("tokenizer");
+            if (tokenizerName != null) {
+                return HuggingFaceTokenizer.newInstance(tokenizerName, options);
+            }
+            if (tokenizerPath == null) {
+                throw new IllegalArgumentException("Missing tokenizer path.");
+            }
+            return HuggingFaceTokenizer.newInstance(tokenizerPath, options);
+        }
+    }
 }
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
index 92586788c..b9e96723a 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslator.java
@@ -25,12 +25,10 @@ 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 fill mask model. */
 public class FillMaskTranslator implements Translator<String, Classifications> {
@@ -102,20 +100,22 @@ public class FillMaskTranslator implements Translator<String, Classifications> {
     /**
      * Creates a builder to build a {@code FillMaskTranslator}.
      *
+     * @param tokenizer the tokenizer
      * @return a new builder
      */
-    public static Builder builder() {
-        return new Builder();
+    public static Builder builder(HuggingFaceTokenizer tokenizer) {
+        return new Builder(tokenizer);
     }
 
     /**
      * Creates a builder to build a {@code FillMaskTranslator}.
      *
+     * @param tokenizer the tokenizer
      * @param arguments the models' arguments
      * @return a new builder
      */
-    public static Builder builder(Map<String, ?> arguments) {
-        Builder builder = builder();
+    public static Builder builder(HuggingFaceTokenizer tokenizer, Map<String, ?> arguments) {
+        Builder builder = builder(tokenizer);
         builder.configure(arguments);
 
         return builder;
@@ -124,13 +124,15 @@ public class FillMaskTranslator implements Translator<String, Classifications> {
     /** The builder for fill mask translator. */
     public static final class Builder {
 
+        private HuggingFaceTokenizer tokenizer;
         private String maskedToken = "[MASK]";
         private int topK = 5;
-        private String tokenizerName;
-        private Path tokenizerPath;
-        private boolean addSpecialTokens = true;
         private Batchifier batchifier = Batchifier.STACK;
 
+        Builder(HuggingFaceTokenizer tokenizer) {
+            this.tokenizer = tokenizer;
+        }
+
         /**
          * Sets the id of the mask {@link Translator}.
          *
@@ -153,39 +155,6 @@ public class FillMaskTranslator implements Translator<String, Classifications> {
             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}.
          *
@@ -205,8 +174,6 @@ public class FillMaskTranslator implements Translator<String, Classifications> {
         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));
         }
@@ -218,14 +185,6 @@ public class FillMaskTranslator implements Translator<String, Classifications> {
          * @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
index d3db2dab1..515abe53d 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslatorFactory.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslatorFactory.java
@@ -13,6 +13,7 @@
 package ai.djl.huggingface.translator;
 
 import ai.djl.Model;
+import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
 import ai.djl.modality.Classifications;
 import ai.djl.modality.Input;
 import ai.djl.modality.Output;
@@ -52,8 +53,10 @@ public class FillMaskTranslatorFactory implements TranslatorFactory {
             throws TranslateException {
         Path modelPath = model.getModelPath();
         try {
+            HuggingFaceTokenizer tokenizer =
+                    HuggingFaceTokenizer.builder(arguments).optTokenizerPath(modelPath).build();
             FillMaskTranslator translator =
-                    FillMaskTranslator.builder(arguments).optTokenizerPath(modelPath).build();
+                    FillMaskTranslator.builder(tokenizer, arguments).build();
             if (input == String.class && output == Classifications.class) {
                 return translator;
             } else if (input == Input.class && output == Output.class) {
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 5dcc7ccab..e0fe288ab 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
@@ -25,9 +25,7 @@ import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorContext;
 
 import java.io.IOException;
-import java.nio.file.Path;
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
 
 /** The translator for Huggingface question answering model. */
 public class QuestionAnsweringTranslator implements Translator<QAInput, String> {
@@ -95,20 +93,22 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
     /**
      * Creates a builder to build a {@code QuestionAnsweringTranslator}.
      *
+     * @param tokenizer the tokenizer
      * @return a new builder
      */
-    public static Builder builder() {
-        return new Builder();
+    public static Builder builder(HuggingFaceTokenizer tokenizer) {
+        return new Builder(tokenizer);
     }
 
     /**
      * Creates a builder to build a {@code QuestionAnsweringTranslator}.
      *
+     * @param tokenizer the tokenizer
      * @param arguments the models' arguments
      * @return a new builder
      */
-    public static Builder builder(Map<String, ?> arguments) {
-        Builder builder = builder();
+    public static Builder builder(HuggingFaceTokenizer tokenizer, Map<String, ?> arguments) {
+        Builder builder = builder(tokenizer);
         builder.configure(arguments);
 
         return builder;
@@ -117,43 +117,12 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
     /** The builder for question answering translator. */
     public static final class Builder {
 
-        private String tokenizerName;
-        private Path tokenizerPath;
+        private HuggingFaceTokenizer tokenizer;
         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 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;
+        Builder(HuggingFaceTokenizer tokenizer) {
+            this.tokenizer = tokenizer;
         }
 
         /**
@@ -184,9 +153,7 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
          * @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));
             String batchifierStr = ArgumentsUtil.stringValue(arguments, "batchifier", "stack");
             optBatchifier(Batchifier.fromString(batchifierStr));
         }
@@ -198,14 +165,6 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
          * @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
index 44f6325a1..4de590d70 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslatorFactory.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslatorFactory.java
@@ -13,6 +13,7 @@
 package ai.djl.huggingface.translator;
 
 import ai.djl.Model;
+import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
 import ai.djl.modality.Input;
 import ai.djl.modality.Output;
 import ai.djl.modality.nlp.qa.QAInput;
@@ -52,10 +53,10 @@ public class QuestionAnsweringTranslatorFactory implements TranslatorFactory {
             throws TranslateException {
         Path modelPath = model.getModelPath();
         try {
+            HuggingFaceTokenizer tokenizer =
+                    HuggingFaceTokenizer.builder(arguments).optTokenizerPath(modelPath).build();
             QuestionAnsweringTranslator translator =
-                    QuestionAnsweringTranslator.builder(arguments)
-                            .optTokenizerPath(modelPath)
-                            .build();
+                    QuestionAnsweringTranslator.builder(tokenizer, arguments).build();
             if (input == QAInput.class && output == String.class) {
                 return translator;
             } else if (input == Input.class && output == Output.class) {
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
index 7eceb2eff..2a7093632 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslator.java
@@ -32,7 +32,6 @@ 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[]> {
@@ -111,20 +110,22 @@ public class TokenClassificationTranslator implements Translator<String, NamedEn
     /**
      * Creates a builder to build a {@code TokenClassificationTranslator}.
      *
+     * @param tokenizer the tokenizer
      * @return a new builder
      */
-    public static Builder builder() {
-        return new Builder();
+    public static Builder builder(HuggingFaceTokenizer tokenizer) {
+        return new Builder(tokenizer);
     }
 
     /**
      * Creates a builder to build a {@code TokenClassificationTranslator}.
      *
+     * @param tokenizer the tokenizer
      * @param arguments the models' arguments
      * @return a new builder
      */
-    public static Builder builder(Map<String, ?> arguments) {
-        Builder builder = builder();
+    public static Builder builder(HuggingFaceTokenizer tokenizer, Map<String, ?> arguments) {
+        Builder builder = builder(tokenizer);
         builder.configure(arguments);
 
         return builder;
@@ -133,42 +134,11 @@ public class TokenClassificationTranslator implements Translator<String, NamedEn
     /** The builder for token classification translator. */
     public static final class Builder {
 
-        private String tokenizerName;
-        private Path tokenizerPath;
-        private boolean addSpecialTokens = true;
+        private HuggingFaceTokenizer tokenizer;
         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;
+        Builder(HuggingFaceTokenizer tokenizer) {
+            this.tokenizer = tokenizer;
         }
 
         /**
@@ -188,8 +158,6 @@ public class TokenClassificationTranslator implements Translator<String, NamedEn
          * @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));
         }
@@ -201,14 +169,6 @@ public class TokenClassificationTranslator implements Translator<String, NamedEn
          * @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);
         }
     }
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslatorFactory.java
index 556b7e4eb..e03cbc82a 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslatorFactory.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslatorFactory.java
@@ -13,6 +13,7 @@
 package ai.djl.huggingface.translator;
 
 import ai.djl.Model;
+import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
 import ai.djl.modality.Input;
 import ai.djl.modality.Output;
 import ai.djl.modality.nlp.translator.NamedEntity;
@@ -52,10 +53,10 @@ public class TokenClassificationTranslatorFactory implements TranslatorFactory {
             throws TranslateException {
         Path modelPath = model.getModelPath();
         try {
+            HuggingFaceTokenizer tokenizer =
+                    HuggingFaceTokenizer.builder(arguments).optTokenizerPath(modelPath).build();
             TokenClassificationTranslator translator =
-                    TokenClassificationTranslator.builder(arguments)
-                            .optTokenizerPath(modelPath)
-                            .build();
+                    TokenClassificationTranslator.builder(tokenizer, arguments).build();
             if (input == String.class && output == NamedEntity[].class) {
                 return translator;
             } else if (input == Input.class && output == Output.class) {


Test patch:
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizerTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizerTest.java
index bf5ffe74b..75c4d4492 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizerTest.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizerTest.java
@@ -93,7 +93,10 @@ public class HuggingFaceTokenizerTest {
                 HuggingFaceTokenizer.newInstance("bert-base-cased", options)) {
             String[] exp = {"Hello", ",", "y", "'", "all", "!", "How", "are", "you", "[UNK]", "?"};
 
-            Encoding encoding = tokenizer.encode(Arrays.asList(inputs));
+            Encoding encoding = tokenizer.encode(inputs);
+            Assert.assertEquals(encoding.getTokens(), exp);
+
+            encoding = tokenizer.encode(Arrays.asList(inputs));
             Assert.assertEquals(encoding.getTokens(), exp);
         }
 
@@ -109,105 +112,80 @@ public class HuggingFaceTokenizerTest {
     }
 
     @Test
-    public void testTokenizerDecoding() {
-        List<long[]> testIds =
-                Arrays.asList(
-                        new long[] {
-                            101, 8667, 117, 194, 112, 1155, 106, 1731, 1132, 1128, 100, 136, 102
-                        },
-                        new long[] {
-                            101, 3570, 1110, 170, 21162, 1285, 119, 2750, 4250, 146, 112, 173, 1474,
-                            102
-                        });
-        List<String> expectedDecodingsNoSpecialTokens =
-                Arrays.asList(
-                        "Hello, y ' all! How are you?",
-                        "Today is a sunny day. Good weather I ' d say");
-        List<String> expectedDecodingsWithSpecialTokens =
-                Arrays.asList(
-                        "[CLS] Hello, y ' all! How are you [UNK]? [SEP]",
-                        "[CLS] Today is a sunny day. Good weather I ' d say [SEP]");
+    public void testTokenizerDecoding() throws IOException {
+        long[][] testIds = {
+            {101, 8667, 117, 194, 112, 1155, 106, 1731, 1132, 1128, 100, 136, 102},
+            {101, 3570, 1110, 170, 21162, 1285, 119, 2750, 4250, 146, 112, 173, 1474, 102}
+        };
+        String[] expectedDecodedNoSpecialTokens = {
+            "Hello, y ' all! How are you?", "Today is a sunny day. Good weather I ' d say"
+        };
+        String[] expectedDecodedWithSpecialTokens = {
+            "[CLS] Hello, y ' all! How are you [UNK]? [SEP]",
+            "[CLS] Today is a sunny day. Good weather I ' d say [SEP]"
+        };
         try (HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.newInstance("bert-base-cased")) {
-            for (int i = 0; i < testIds.size(); ++i) {
-                Assert.assertEquals(
-                        tokenizer.decode(testIds.get(i)),
-                        expectedDecodingsWithSpecialTokens.get(i));
-                Assert.assertEquals(
-                        tokenizer.decode(testIds.get(i), true),
-                        expectedDecodingsNoSpecialTokens.get(i));
+            for (int i = 0; i < testIds.length; ++i) {
+                String ret = tokenizer.decode(testIds[i]);
+                Assert.assertEquals(ret, expectedDecodedWithSpecialTokens[i]);
+                ret = tokenizer.decode(testIds[i], true);
+                Assert.assertEquals(ret, expectedDecodedNoSpecialTokens[i]);
             }
         }
 
-        Map<String, String> options = new ConcurrentHashMap<>();
-        options.put("addSpecialTokens", "false");
         try (HuggingFaceTokenizer tokenizer =
-                HuggingFaceTokenizer.newInstance("bert-base-cased", options)) {
-            for (int i = 0; i < testIds.size(); ++i) {
-                Assert.assertEquals(
-                        tokenizer.decode(testIds.get(i)), expectedDecodingsNoSpecialTokens.get(i));
-                Assert.assertEquals(
-                        tokenizer.decode(testIds.get(i), false),
-                        expectedDecodingsWithSpecialTokens.get(i));
+                HuggingFaceTokenizer.builder()
+                        .optTokenizerName("bert-base-cased")
+                        .optAddSpecialTokens(false)
+                        .build()) {
+            for (int i = 0; i < testIds.length; ++i) {
+                String ret = tokenizer.decode(testIds[i]);
+                Assert.assertEquals(ret, expectedDecodedNoSpecialTokens[i]);
+                ret = tokenizer.decode(testIds[i], false);
+                Assert.assertEquals(ret, expectedDecodedWithSpecialTokens[i]);
             }
         }
     }
 
     @Test
-    public void testMaxLengthTruncationAndAllPaddings() {
-        List<String> inputs =
-                Arrays.asList(
-                        "Hello, y'all! How are you?",
-                        "Today is a sunny day. Good weather I'd say",
-                        "I am happy");
+    public void testMaxLengthTruncationAndAllPaddings() throws IOException {
+        String[] inputs = {
+            "Hello, y'all! How are you?", "Today is a sunny day. Good weather I'd say", "I am happy"
+        };
 
         Map<String, String> options = new ConcurrentHashMap<>();
-        options.put("truncation", "longest_first");
-        options.put("padding", "longest");
+        options.put("tokenizer", "bert-base-cased");
+        options.put("truncation", "longest_first"); // true
+        options.put("padding", "longest"); // true
         options.put("maxLength", "10");
 
-        try (HuggingFaceTokenizer tokenizer =
-                HuggingFaceTokenizer.newInstance("bert-base-cased", options)) {
-            List<long[]> expectedIdsWithTruncationAndPadding =
-                    Arrays.asList(
-                            new long[] {101, 8667, 117, 194, 112, 1155, 106, 1731, 1132, 102},
-                            new long[] {101, 3570, 1110, 170, 21162, 1285, 119, 2750, 4250, 102},
-                            new long[] {101, 146, 1821, 2816, 102, 0, 0, 0, 0, 0});
-            Encoding[] encodings = tokenizer.batchEncode(inputs);
-            for (int i = 0; i < encodings.length; ++i) {
-                Assert.assertEquals(
-                        encodings[i].getIds(), expectedIdsWithTruncationAndPadding.get(i));
+        try (HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.builder(options).build()) {
+            int expectedSize = 10;
+            for (Encoding encoding : tokenizer.batchEncode(inputs)) {
+                Assert.assertEquals(encoding.getIds().length, expectedSize);
             }
+        }
 
-            List<long[]> expectedIdsWithTruncationAndPaddingToMultiple =
-                    Arrays.asList(
-                            new long[] {101, 8667, 117, 194, 112, 1155, 106, 1731, 1132, 102, 0, 0},
-                            new long[] {
-                                101, 3570, 1110, 170, 21162, 1285, 119, 2750, 4250, 102, 0, 0
-                            },
-                            new long[] {101, 146, 1821, 2816, 102, 0, 0, 0, 0, 0, 0, 0});
-            tokenizer.setPadToMultipleOf(3);
-            encodings = tokenizer.batchEncode(inputs);
-            for (int i = 0; i < encodings.length; ++i) {
-                Assert.assertEquals(
-                        encodings[i].getIds(),
-                        expectedIdsWithTruncationAndPaddingToMultiple.get(i));
+        options.put("padToMultipleOf", "3");
+        try (HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.builder(options).build()) {
+            int expectedSize = 12;
+            for (Encoding encoding : tokenizer.batchEncode(inputs)) {
+                Assert.assertEquals(encoding.getIds().length, expectedSize);
             }
+        }
 
-            List<long[]> expectedIdsNoPadding =
-                    Arrays.asList(
-                            new long[] {101, 8667, 117, 194, 112, 1155, 106, 1731, 1132, 102},
-                            new long[] {101, 3570, 1110, 170, 21162, 1285, 119, 2750, 4250, 102},
-                            new long[] {101, 146, 1821, 2816, 102});
-            tokenizer.setPadding(false);
-            encodings = tokenizer.batchEncode(inputs);
+        options.put("padding", "false");
+        try (HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.builder(options).build()) {
+            int[] expectedIdsNoPadding = {10, 10, 5};
+            Encoding[] encodings = tokenizer.batchEncode(inputs);
             for (int i = 0; i < encodings.length; ++i) {
-                Assert.assertEquals(encodings[i].getIds(), expectedIdsNoPadding.get(i));
+                Assert.assertEquals(encodings[i].getIds().length, expectedIdsNoPadding[i]);
             }
         }
     }
 
     @Test
-    public void testMaxModelLengthTruncationAndAllPaddings() {
+    public void testMaxModelLengthTruncationAndAllPaddings() throws IOException {
         String repeat = "hi,";
         int numRepeats = 513;
         int capacity = numRepeats * 3;
@@ -223,18 +201,22 @@ public class HuggingFaceTokenizerTest {
                 Assert.assertEquals(
                         encodings[i].getIds().length, expectedNumberOfIdsNoTruncationNoPadding[i]);
             }
+        }
 
-            tokenizer.setTruncation(true);
-            int[] expectedNumberOfIdsDefaultTruncationNoPadding = new int[] {512, 7};
-            encodings = tokenizer.batchEncode(inputs);
+        Map<String, String> options = new ConcurrentHashMap<>();
+        options.put("tokenizer", "bert-base-cased");
+        options.put("truncation", "true");
+        try (HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.builder(options).build()) {
+            int[] expectedSize = new int[] {512, 7};
+            Encoding[] encodings = tokenizer.batchEncode(inputs);
             for (int i = 0; i < encodings.length; ++i) {
-                Assert.assertEquals(
-                        encodings[i].getIds().length,
-                        expectedNumberOfIdsDefaultTruncationNoPadding[i]);
+                Assert.assertEquals(encodings[i].getIds().length, expectedSize[i]);
             }
+        }
 
-            tokenizer.setPadding(true);
-            encodings = tokenizer.batchEncode(inputs);
+        options.put("padding", "true");
+        try (HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.builder(options).build()) {
+            Encoding[] encodings = tokenizer.batchEncode(inputs);
             for (Encoding encoding : encodings) {
                 Assert.assertEquals(encoding.getIds().length, 512);
             }
@@ -242,139 +224,123 @@ public class HuggingFaceTokenizerTest {
     }
 
     @Test
-    public void testNoTruncationAndAllPaddings() {
-        List<String> inputs =
-                Arrays.asList(
-                        "Hello, y'all! How are you?",
-                        "Today is a sunny day. Good weather I'd say",
-                        "I am happy");
+    public void testNoTruncationAndAllPaddings() throws IOException {
+        String[] inputs = {
+            "Hello, y'all! How are you?", "Today is a sunny day. Good weather I'd say", "I am happy"
+        };
         try (HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.newInstance("bert-base-cased")) {
-            List<long[]> expectedIdsNoTruncationOrPadding =
-                    Arrays.asList(
-                            new long[] {
-                                101, 8667, 117, 194, 112, 1155, 106, 1731, 1132, 1128, 136, 102
-                            },
-                            new long[] {
-                                101, 3570, 1110, 170, 21162, 1285, 119, 2750, 4250, 146, 112, 173,
-                                1474, 102
-                            },
-                            new long[] {101, 146, 1821, 2816, 102});
-
+            int[] expectedSize = {12, 14, 5};
             Encoding[] encodings = tokenizer.batchEncode(inputs);
             for (int i = 0; i < encodings.length; ++i) {
-                Assert.assertEquals(encodings[i].getIds(), expectedIdsNoTruncationOrPadding.get(i));
+                Assert.assertEquals(encodings[i].getIds().length, expectedSize[i]);
             }
-            tokenizer.setPadding(true);
-            List<long[]> expectedIdsNoTruncationDefaultPadding =
-                    Arrays.asList(
-                            new long[] {
-                                101, 8667, 117, 194, 112, 1155, 106, 1731, 1132, 1128, 136, 102, 0,
-                                0
-                            },
-                            new long[] {
-                                101, 3570, 1110, 170, 21162, 1285, 119, 2750, 4250, 146, 112, 173,
-                                1474, 102
-                            },
-                            new long[] {101, 146, 1821, 2816, 102, 0, 0, 0, 0, 0, 0, 0, 0, 0});
-            encodings = tokenizer.batchEncode(inputs);
-            for (int i = 0; i < encodings.length; ++i) {
-                Assert.assertEquals(
-                        encodings[i].getIds(), expectedIdsNoTruncationDefaultPadding.get(i));
+        }
+
+        Map<String, String> options = new ConcurrentHashMap<>();
+        options.put("tokenizer", "bert-base-cased");
+        options.put("padding", "true");
+        try (HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.builder(options).build()) {
+            int expectedSize = 14;
+            for (Encoding encoding : tokenizer.batchEncode(inputs)) {
+                Assert.assertEquals(encoding.getIds().length, expectedSize);
             }
+        }
 
-            tokenizer.setPaddingMaxLength();
-            encodings = tokenizer.batchEncode(inputs);
-            for (int i = 0; i < encodings.length; ++i) {
-                Assert.assertEquals(encodings[i].getIds().length, 512);
+        try (HuggingFaceTokenizer tokenizer =
+                HuggingFaceTokenizer.builder()
+                        .optTokenizerName("bert-base-cased")
+                        .optPadToMaxLength()
+                        .build()) {
+            for (Encoding encoding : tokenizer.batchEncode(inputs)) {
+                Assert.assertEquals(encoding.getIds().length, 512);
             }
+        }
 
-            tokenizer.setMaxLength(15);
-            List<long[]> expectedIdsNoTruncationMaxLengthPadding =
-                    Arrays.asList(
-                            new long[] {
-                                101, 8667, 117, 194, 112, 1155, 106, 1731, 1132, 1128, 136, 102, 0,
-                                0, 0
-                            },
-                            new long[] {
-                                101, 3570, 1110, 170, 21162, 1285, 119, 2750, 4250, 146, 112, 173,
-                                1474, 102, 0
-                            },
-                            new long[] {101, 146, 1821, 2816, 102, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0});
-            encodings = tokenizer.batchEncode(inputs);
-            for (int i = 0; i < encodings.length; ++i) {
-                Assert.assertEquals(
-                        encodings[i].getIds(), expectedIdsNoTruncationMaxLengthPadding.get(i));
+        try (HuggingFaceTokenizer tokenizer =
+                HuggingFaceTokenizer.builder()
+                        .optTokenizerName("bert-base-cased")
+                        .optPadToMaxLength()
+                        .optMaxLength(15)
+                        .build()) {
+            int expectedSize = 15;
+            for (Encoding encoding : tokenizer.batchEncode(inputs)) {
+                Assert.assertEquals(encoding.getIds().length, expectedSize);
             }
+        }
 
-            tokenizer.setPadToMultipleOf(8);
-            List<long[]> expectedIdsNoTruncationMaxLengthPadToMultiplePadding =
-                    Arrays.asList(
-                            new long[] {
-                                101, 8667, 117, 194, 112, 1155, 106, 1731, 1132, 1128, 136, 102, 0,
-                                0, 0, 0
-                            },
-                            new long[] {
-                                101, 3570, 1110, 170, 21162, 1285, 119, 2750, 4250, 146, 112, 173,
-                                1474, 102, 0, 0
-                            },
-                            new long[] {
-                                101, 146, 1821, 2816, 102, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
-                            });
-            encodings = tokenizer.batchEncode(inputs);
-            for (int i = 0; i < encodings.length; ++i) {
-                Assert.assertEquals(
-                        encodings[i].getIds(),
-                        expectedIdsNoTruncationMaxLengthPadToMultiplePadding.get(i));
+        try (HuggingFaceTokenizer tokenizer =
+                HuggingFaceTokenizer.builder()
+                        .optTokenizerName("bert-base-cased")
+                        .optPadToMaxLength()
+                        .optMaxLength(15)
+                        .optPadToMultipleOf(8)
+                        .build()) {
+            int expectedSize = 16;
+            for (Encoding encoding : tokenizer.batchEncode(inputs)) {
+                Assert.assertEquals(encoding.getIds().length, expectedSize);
             }
         }
     }
 
     @Test
-    public void testTruncationAndPaddingForPairInputs() {
-        Map<String, String> options = new ConcurrentHashMap<>();
-        options.put("truncation", "only_first");
-        options.put("padding", "longest");
-        options.put("maxLength", "8");
+    public void testTruncationAndPaddingForPairInputs() throws IOException {
         String text = "Hello there my friend";
         String textPair = "How are you";
 
         try (HuggingFaceTokenizer tokenizer =
-                HuggingFaceTokenizer.newInstance("bert-base-cased", options)) {
+                HuggingFaceTokenizer.builder()
+                        .optTokenizerName("bert-base-cased")
+                        .optTruncateFirstOnly()
+                        .optPadding(true)
+                        .optMaxLength(8)
+                        .build()) {
+            Encoding encoding = tokenizer.encode(text, textPair);
+            Assert.assertEquals(encoding.getIds().length, 8);
+        }
+
+        try (HuggingFaceTokenizer tokenizer =
+                HuggingFaceTokenizer.builder()
+                        .optTokenizerName("bert-base-cased")
+                        .optTruncateFirstOnly()
+                        .optPadding(true)
+                        .optMaxLength(8)
+                        .optPadToMultipleOf(3)
+                        .build()) {
             Encoding encoding = tokenizer.encode(text, textPair);
-            long[] expectedId = new long[] {101, 8667, 1175, 102, 1731, 1132, 1128, 102};
-            Assert.assertEquals(encoding.getIds(), expectedId);
-
-            tokenizer.setPadToMultipleOf(3);
-            encoding = tokenizer.encode(text, textPair);
-            long[] expectedIdPadToMultiple =
-                    new long[] {101, 8667, 1175, 102, 1731, 1132, 1128, 102, 0};
-            Assert.assertEquals(encoding.getIds(), expectedIdPadToMultiple);
+            Assert.assertEquals(encoding.getIds().length, 9);
         }
 
-        options.put("truncation", "only_second");
-        options.remove("padding");
-        options.put("maxLength", "8");
+        try (HuggingFaceTokenizer tokenizer =
+                HuggingFaceTokenizer.builder()
+                        .optTokenizerName("bert-base-cased")
+                        .optTruncateSecondOnly()
+                        .optMaxLength(8)
+                        .build()) {
+            Encoding encoding = tokenizer.encode(text, textPair);
+            Assert.assertEquals(encoding.getIds().length, 8);
+        }
 
         try (HuggingFaceTokenizer tokenizer =
-                HuggingFaceTokenizer.newInstance("bert-base-cased", options)) {
+                HuggingFaceTokenizer.builder()
+                        .optTokenizerName("bert-base-cased")
+                        .optTruncateSecondOnly()
+                        .optMaxLength(8)
+                        .optPadToMultipleOf(3)
+                        .optPadding(true)
+                        .build()) {
             Encoding encoding = tokenizer.encode(text, textPair);
-            long[] expectedId = new long[] {101, 8667, 1175, 1139, 1910, 102, 1731, 102};
-            Assert.assertEquals(encoding.getIds(), expectedId);
-
-            tokenizer.setPadToMultipleOf(3);
-            tokenizer.setPadding(true);
-            encoding = tokenizer.encode(text, textPair);
-            long[] expectedIdPadToMultiple =
-                    new long[] {101, 8667, 1175, 1139, 1910, 102, 1731, 102, 0};
-            Assert.assertEquals(encoding.getIds(), expectedIdPadToMultiple);
+            Assert.assertEquals(encoding.getIds().length, 9);
         }
 
-        options.put("truncation", "longest_first");
         try (HuggingFaceTokenizer tokenizer =
-                HuggingFaceTokenizer.newInstance("bert-base-cased", options)) {
+                HuggingFaceTokenizer.builder()
+                        .optTokenizerName("bert-base-cased")
+                        .optTruncation(true)
+                        .optPadding(true)
+                        .optMaxLength(8)
+                        .build()) {
             Encoding encoding = tokenizer.encode(text, textPair);
-            long[] expectedId = new long[] {101, 8667, 1175, 1139, 102, 1731, 1132, 102};
-            Assert.assertEquals(encoding.getIds(), expectedId);
+            Assert.assertEquals(encoding.getIds().length, 8);
         }
     }
 }