Repository: deepjavalibrary/djl
Pull request #3571 (2024-12-30 20:16:39)

[tokenizers] Add int32 option to encoding
## 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


frankfliu: Fixes: #3566

Patch:
diff --git a/README.md b/README.md
index cda6533df..47d0c20be 100644
--- a/README.md
+++ b/README.md
@@ -3,7 +3,7 @@
 
 ![Continuous](https://github.com/deepjavalibrary/djl/workflows/Continuous/badge.svg)
 [![Docs](https://github.com/deepjavalibrary/djl/workflows/Docs/badge.svg)](https://docs.djl.ai/master/index.html)
-![Nightly Publish](https://github.com/deepjavalibrary/djl/workflows/Nightly%20Publish/badge.svg)
+[![Nightly Publish](https://github.com/deepjavalibrary/djl/workflows/Nightly%20Publish/badge.svg)](https://github.com/deepjavalibrary/djl/actions/workflows/nightly_publish.yml)
 
 # Deep Java Library (DJL)
 
diff --git a/android/gradle.properties b/android/gradle.properties
index 78b3fa76c..56211c900 100644
--- a/android/gradle.properties
+++ b/android/gradle.properties
@@ -17,6 +17,6 @@ org.gradle.jvmargs=-Xmx1536m
 android.useAndroidX=true
 # Automatically convert third-party libraries to use AndroidX
 android.enableJetifier=true
-djl_version=0.31.0
+djl_version=0.32.0
 pytorch_version=2.5.1
-tokenizers_version=0.20.3
+tokenizers_version=0.21.0
diff --git a/api/src/main/java/ai/djl/modality/cv/transform/Resize.java b/api/src/main/java/ai/djl/modality/cv/transform/Resize.java
index f2bddc8f5..1947b6973 100644
--- a/api/src/main/java/ai/djl/modality/cv/transform/Resize.java
+++ b/api/src/main/java/ai/djl/modality/cv/transform/Resize.java
@@ -57,6 +57,24 @@ public class Resize implements Transform {
         this.interpolation = interpolation;
     }
 
+    /**
+     * Returns the width.
+     *
+     * @return the width
+     */
+    public int getWidth() {
+        return width;
+    }
+
+    /**
+     * Returns the height.
+     *
+     * @return the height
+     */
+    public int getHeight() {
+        return height;
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray transform(NDArray array) {
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/BaseImageTranslator.java b/api/src/main/java/ai/djl/modality/cv/translator/BaseImageTranslator.java
index 3a66df732..675c4fb2f 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/BaseImageTranslator.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/BaseImageTranslator.java
@@ -116,6 +116,26 @@ public abstract class BaseImageTranslator<T> implements Translator<Image, T> {
          */
         public T setPipeline(Pipeline pipeline) {
             this.pipeline = pipeline;
+            for (Transform transform : pipeline.getTransforms()) {
+                if (transform instanceof Resize) {
+                    Resize resize = (Resize) transform;
+                    width = resize.getWidth();
+                    height = resize.getHeight();
+                }
+            }
+            return self();
+        }
+
+        /**
+         * Sets the image size.
+         *
+         * @param width the image width
+         * @param height the image height
+         * @return this builder
+         */
+        public T setImageSize(int width, int height) {
+            this.width = width;
+            this.height = height;
             return self();
         }
 
@@ -129,6 +149,12 @@ public abstract class BaseImageTranslator<T> implements Translator<Image, T> {
             if (pipeline == null) {
                 pipeline = new Pipeline();
             }
+            if (transform instanceof Resize) {
+                Resize resize = (Resize) transform;
+                width = resize.getWidth();
+                height = resize.getHeight();
+            }
+
             pipeline.add(transform);
             return self();
         }
@@ -167,13 +193,14 @@ public abstract class BaseImageTranslator<T> implements Translator<Image, T> {
             } else if (!"false".equals(resize)) {
                 String[] tokens = resize.split("\\s*,\\s*");
                 if (tokens.length > 1) {
-                    addTransform(
-                            new Resize(
-                                    (int) Double.parseDouble(tokens[0]),
-                                    (int) Double.parseDouble(tokens[1])));
+                    setImageSize(
+                            (int) Double.parseDouble(tokens[0]),
+                            (int) Double.parseDouble(tokens[1]));
                 } else {
-                    addTransform(new Resize((int) Double.parseDouble(tokens[0])));
+                    int size = (int) Double.parseDouble(tokens[0]);
+                    setImageSize(size, size);
                 }
+                addTransform(new Resize(width, height));
             }
             if (ArgumentsUtil.booleanValue(arguments, "centerCrop", false)) {
                 addTransform(new CenterCrop(width, height));
diff --git a/api/src/main/java/ai/djl/translate/Pipeline.java b/api/src/main/java/ai/djl/translate/Pipeline.java
index d31fcf315..3de467fdd 100644
--- a/api/src/main/java/ai/djl/translate/Pipeline.java
+++ b/api/src/main/java/ai/djl/translate/Pipeline.java
@@ -17,6 +17,7 @@ import ai.djl.ndarray.NDList;
 import ai.djl.util.Pair;
 import ai.djl.util.PairList;
 
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -176,6 +177,15 @@ public class Pipeline {
         return new NDList(arrays);
     }
 
+    /**
+     * Returns the list of transforms.
+     *
+     * @return the list of transforms
+     */
+    public List<Transform> getTransforms() {
+        return transforms.values();
+    }
+
     private static final class IndexKey {
         private String key;
         private int index;
diff --git a/engines/pytorch/pytorch-engine/README.md b/engines/pytorch/pytorch-engine/README.md
index 4623fd032..973e0bff9 100644
--- a/engines/pytorch/pytorch-engine/README.md
+++ b/engines/pytorch/pytorch-engine/README.md
@@ -129,7 +129,7 @@ export PYTORCH_FLAVOR=cpu
 **Note:**
 
 - x86_64 macOS is no longer supported since 0.30.0
-- PyTorch 1.13+ doesn't support mac 11 any more, you must use DJL 0.19.0 ane lower version.
+- PyTorch 1.13+ doesn't support mac 11 any more, you must use DJL 0.19.0 or a lower version.
 
 ### macOS M1
 
diff --git a/extensions/tokenizers/rust/Cargo.toml b/extensions/tokenizers/rust/Cargo.toml
index ebfa6a93f..f33a58bdc 100644
--- a/extensions/tokenizers/rust/Cargo.toml
+++ b/extensions/tokenizers/rust/Cargo.toml
@@ -13,7 +13,7 @@ candle-flash-attn = { version = "*", optional = true }
 candle-cublaslt = { git = "https://github.com/huggingface/candle-cublaslt", rev = "cf789b7dd6d4abb19b03b9556442f94f0588b4a0", optional = true }
 candle-layer-norm = { git = "https://github.com/xyang16/candle-layer-norm", rev = "e574de6a7f88bafbede8edf9ee43170c6a8ce51a", optional = true }
 candle-rotary = { git = "https://github.com/huggingface/candle-rotary", rev = "0a718a0856569a92f3112e64f10d07e4447822e8", optional = true }
-tokenizers = { version = "0.20.3" }
+tokenizers = { version = "0.21.0" }
 half = "2.4.0"
 tracing = "0.1.40"
 safetensors = "0.4.3"
@@ -34,7 +34,7 @@ openssl-src = { git = "https://github.com/alexcrichton/openssl-src-rs", rev = "c
 openssl = { version = "0.10", features = ["vendored"] }
 
 [target.'cfg(not(target_os = "android"))'.dependencies]
-tokenizers = { version = "0.20.0", features = ["http"] }
+tokenizers = { version = "0.21.0", features = ["http"] }
 
 [lib]
 crate-type = ["cdylib"]
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/Encoding.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/Encoding.java
index 4877f53aa..8670a0163 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/Encoding.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/Encoding.java
@@ -13,6 +13,7 @@
 package ai.djl.huggingface.tokenizers;
 
 import ai.djl.huggingface.tokenizers.jni.CharSpan;
+import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
 
@@ -55,23 +56,85 @@ public class Encoding {
         this.overflowing = overflowing;
     }
 
+    /**
+     * Returns the {@link NDList} representation of the encodings.
+     *
+     * @param encodings the {@code Encoding} batch
+     * @param manager the {@link NDManager} to create the NDList
+     * @param withTokenType true to include the token type id
+     * @param int32 true to use int32 datatype
+     * @return the {@link NDList}
+     */
+    public static NDList toNDList(
+            Encoding[] encodings, NDManager manager, boolean withTokenType, boolean int32) {
+        NDList list = new NDList();
+        if (!int32) {
+            long[][] ids = new long[encodings.length][];
+            long[][] attentionMask = new long[encodings.length][];
+            long[][] typeIds = new long[encodings.length][];
+            for (int i = 0; i < encodings.length; i++) {
+                ids[i] = encodings[i].getIds();
+                attentionMask[i] = encodings[i].getAttentionMask();
+                if (withTokenType) {
+                    typeIds[i] = encodings[i].getTypeIds();
+                }
+            }
+            list.add(manager.create(ids));
+            NDArray inputAttentionMask = manager.create(attentionMask);
+            list.add(inputAttentionMask);
+            if (withTokenType) {
+                list.add(manager.create(typeIds));
+            }
+            return list;
+        }
+
+        int[][] ids = new int[encodings.length][];
+        int[][] attentionMask = new int[encodings.length][];
+        int[][] typeIds = new int[encodings.length][];
+        for (int i = 0; i < encodings.length; i++) {
+            ids[i] = Arrays.stream(encodings[i].getIds()).mapToInt(l -> (int) l).toArray();
+            attentionMask[i] =
+                    Arrays.stream(encodings[i].getAttentionMask()).mapToInt(l -> (int) l).toArray();
+            if (withTokenType) {
+                typeIds[i] =
+                        Arrays.stream(encodings[i].getTypeIds()).mapToInt(l -> (int) l).toArray();
+            }
+        }
+        list.add(manager.create(ids));
+        NDArray inputAttentionMask = manager.create(attentionMask);
+        list.add(inputAttentionMask);
+        if (withTokenType) {
+            list.add(manager.create(typeIds));
+        }
+        return list;
+    }
+
     /**
      * Returns the {@link NDList} representation of the encoding.
      *
      * @param manager the {@link NDManager} to create the NDList
      * @param withTokenType true to include the token type id
+     * @param int32 true to use int32 datatype
      * @return the {@link NDList}
      */
-    public NDList toNDList(NDManager manager, boolean withTokenType) {
+    public NDList toNDList(NDManager manager, boolean withTokenType, boolean int32) {
         // Converting encoding to int32 NDList because candle can't convert int64 to fp16 in cuda
         NDList list = new NDList(withTokenType ? 3 : 2);
-        int[] intIds = Arrays.stream(ids).mapToInt(i -> (int) i).toArray();
-        int[] intAttentionMask = Arrays.stream(attentionMask).mapToInt(i -> (int) i).toArray();
-        list.add(manager.create(intIds));
-        list.add(manager.create(intAttentionMask));
-        if (withTokenType) {
-            int[] intTypeIds = Arrays.stream(typeIds).mapToInt(i -> (int) i).toArray();
-            list.add(manager.create(intTypeIds));
+        if (int32) {
+            int[] intIds = Arrays.stream(ids).mapToInt(i -> (int) i).toArray();
+            int[] intAttentionMask = Arrays.stream(attentionMask).mapToInt(i -> (int) i).toArray();
+            list.add(manager.create(intIds));
+            list.add(manager.create(intAttentionMask));
+            if (withTokenType) {
+                int[] intTypeIds = Arrays.stream(typeIds).mapToInt(i -> (int) i).toArray();
+                list.add(manager.create(intTypeIds));
+            }
+        } else {
+            list.add(manager.create(ids));
+            list.add(manager.create(attentionMask));
+            if (withTokenType) {
+                list.add(manager.create(typeIds));
+            }
         }
         return list;
     }
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslator.java
index ac5e4ed55..843bb5716 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslator.java
@@ -35,16 +35,19 @@ public class CrossEncoderTranslator implements Translator<StringPair, float[]> {
 
     private HuggingFaceTokenizer tokenizer;
     private boolean includeTokenTypes;
+    private boolean int32;
     private boolean sigmoid;
     private Batchifier batchifier;
 
     CrossEncoderTranslator(
             HuggingFaceTokenizer tokenizer,
             boolean includeTokenTypes,
+            boolean int32,
             boolean sigmoid,
             Batchifier batchifier) {
         this.tokenizer = tokenizer;
         this.includeTokenTypes = includeTokenTypes;
+        this.int32 = int32;
         this.sigmoid = sigmoid;
         this.batchifier = batchifier;
     }
@@ -60,7 +63,7 @@ public class CrossEncoderTranslator implements Translator<StringPair, float[]> {
     public NDList processInput(TranslatorContext ctx, StringPair input) {
         Encoding encoding = tokenizer.encode(input.getKey(), input.getValue());
         ctx.setAttachment("encoding", encoding);
-        return encoding.toNDList(ctx.getNDManager(), includeTokenTypes);
+        return encoding.toNDList(ctx.getNDManager(), includeTokenTypes, int32);
     }
 
     /** {@inheritDoc} */
@@ -71,7 +74,7 @@ public class CrossEncoderTranslator implements Translator<StringPair, float[]> {
         Encoding[] encodings = tokenizer.batchEncode(list);
         NDList[] batch = new NDList[encodings.length];
         for (int i = 0; i < encodings.length; ++i) {
-            batch[i] = encodings[i].toNDList(manager, includeTokenTypes);
+            batch[i] = encodings[i].toNDList(manager, includeTokenTypes, int32);
         }
         return batchifier.batchify(batch);
     }
@@ -145,6 +148,7 @@ public class CrossEncoderTranslator implements Translator<StringPair, float[]> {
 
         private HuggingFaceTokenizer tokenizer;
         private boolean includeTokenTypes;
+        private boolean int32;
         private boolean sigmoid = true;
         private Batchifier batchifier = Batchifier.STACK;
 
@@ -163,6 +167,17 @@ public class CrossEncoderTranslator implements Translator<StringPair, float[]> {
             return this;
         }
 
+        /**
+         * Sets if use int32 datatype for the {@link Translator}.
+         *
+         * @param int32 true to include token types
+         * @return this builder
+         */
+        public Builder optInt32(boolean int32) {
+            this.int32 = int32;
+            return this;
+        }
+
         /**
          * Sets if apply sigmoid for the {@link Translator}.
          *
@@ -192,6 +207,7 @@ public class CrossEncoderTranslator implements Translator<StringPair, float[]> {
          */
         public void configure(Map<String, ?> arguments) {
             optIncludeTokenTypes(ArgumentsUtil.booleanValue(arguments, "includeTokenTypes"));
+            optInt32(ArgumentsUtil.booleanValue(arguments, "int32"));
             optSigmoid(ArgumentsUtil.booleanValue(arguments, "sigmoid", true));
             String batchifierStr = ArgumentsUtil.stringValue(arguments, "batchifier", "stack");
             optBatchifier(Batchifier.fromString(batchifierStr));
@@ -204,7 +220,8 @@ public class CrossEncoderTranslator implements Translator<StringPair, float[]> {
          * @throws IOException if I/O error occurs
          */
         public CrossEncoderTranslator build() throws IOException {
-            return new CrossEncoderTranslator(tokenizer, includeTokenTypes, sigmoid, batchifier);
+            return new CrossEncoderTranslator(
+                    tokenizer, includeTokenTypes, int32, sigmoid, batchifier);
         }
     }
 }
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 6eaf75597..7fe20e0fe 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
@@ -38,6 +38,7 @@ public class FillMaskTranslator implements Translator<String, Classifications> {
     private long maskTokenId;
     private int topK;
     private boolean includeTokenTypes;
+    private boolean int32;
     private Batchifier batchifier;
 
     FillMaskTranslator(
@@ -45,11 +46,13 @@ public class FillMaskTranslator implements Translator<String, Classifications> {
             String maskToken,
             int topK,
             boolean includeTokenTypes,
+            boolean int32,
             Batchifier batchifier) {
         this.tokenizer = tokenizer;
         this.maskToken = maskToken;
         this.topK = topK;
         this.includeTokenTypes = includeTokenTypes;
+        this.int32 = int32;
         this.batchifier = batchifier;
         Encoding encoding = tokenizer.encode(maskToken, false, false);
         maskTokenId = encoding.getIds()[0];
@@ -68,7 +71,7 @@ public class FillMaskTranslator implements Translator<String, Classifications> {
         long[] indices = encoding.getIds();
         int maskIndex = getMaskIndex(indices);
         ctx.setAttachment("maskIndex", maskIndex);
-        return encoding.toNDList(ctx.getNDManager(), includeTokenTypes);
+        return encoding.toNDList(ctx.getNDManager(), includeTokenTypes, int32);
     }
 
     /** {@inheritDoc} */
@@ -83,7 +86,7 @@ public class FillMaskTranslator implements Translator<String, Classifications> {
         for (int i = 0; i < batch.length; ++i) {
             long[] indices = encodings[i].getIds();
             maskIndices[i] = getMaskIndex(indices);
-            batch[i] = encodings[i].toNDList(manager, includeTokenTypes);
+            batch[i] = encodings[i].toNDList(manager, includeTokenTypes, int32);
         }
         return batchifier.batchify(batch);
     }
@@ -167,6 +170,7 @@ public class FillMaskTranslator implements Translator<String, Classifications> {
         private String maskedToken = "[MASK]";
         private int topK = 5;
         private boolean includeTokenTypes;
+        private boolean int32;
         private Batchifier batchifier = Batchifier.STACK;
 
         Builder(HuggingFaceTokenizer tokenizer) {
@@ -206,6 +210,17 @@ public class FillMaskTranslator implements Translator<String, Classifications> {
             return this;
         }
 
+        /**
+         * Sets if use int32 datatype for the {@link Translator}.
+         *
+         * @param int32 true to include token types
+         * @return this builder
+         */
+        public Builder optInt32(boolean int32) {
+            this.int32 = int32;
+            return this;
+        }
+
         /**
          * Sets the {@link Batchifier} for the {@link Translator}.
          *
@@ -224,6 +239,7 @@ public class FillMaskTranslator implements Translator<String, Classifications> {
          */
         public void configure(Map<String, ?> arguments) {
             optMaskToken(ArgumentsUtil.stringValue(arguments, "maskToken", "[MASK]"));
+            optInt32(ArgumentsUtil.booleanValue(arguments, "int32"));
             optTopK(ArgumentsUtil.intValue(arguments, "topK", 5));
             optIncludeTokenTypes(ArgumentsUtil.booleanValue(arguments, "includeTokenTypes"));
             String batchifierStr = ArgumentsUtil.stringValue(arguments, "batchifier", "stack");
@@ -238,7 +254,7 @@ public class FillMaskTranslator implements Translator<String, Classifications> {
          */
         public FillMaskTranslator build() throws IOException {
             return new FillMaskTranslator(
-                    tokenizer, maskedToken, topK, includeTokenTypes, batchifier);
+                    tokenizer, maskedToken, topK, includeTokenTypes, int32, batchifier);
         }
     }
 }
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 1b3b4b098..92cba3726 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
@@ -37,16 +37,19 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
 
     private HuggingFaceTokenizer tokenizer;
     private boolean includeTokenTypes;
+    private boolean int32;
     private Batchifier batchifier;
     private boolean detail;
 
     QuestionAnsweringTranslator(
             HuggingFaceTokenizer tokenizer,
             boolean includeTokenTypes,
+            boolean int32,
             Batchifier batchifier,
             boolean detail) {
         this.tokenizer = tokenizer;
         this.includeTokenTypes = includeTokenTypes;
+        this.int32 = int32;
         this.batchifier = batchifier;
         this.detail = detail;
     }
@@ -62,7 +65,7 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
     public NDList processInput(TranslatorContext ctx, QAInput input) {
         Encoding encoding = tokenizer.encode(input.getQuestion(), input.getParagraph());
         ctx.setAttachment("encoding", encoding);
-        return encoding.toNDList(ctx.getNDManager(), includeTokenTypes);
+        return encoding.toNDList(ctx.getNDManager(), includeTokenTypes, int32);
     }
 
     /** {@inheritDoc} */
@@ -77,7 +80,7 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
         ctx.setAttachment("encodings", encodings);
         NDList[] batch = new NDList[encodings.length];
         for (int i = 0; i < encodings.length; ++i) {
-            batch[i] = encodings[i].toNDList(manager, includeTokenTypes);
+            batch[i] = encodings[i].toNDList(manager, includeTokenTypes, int32);
         }
         return batchifier.batchify(batch);
     }
@@ -190,6 +193,7 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
 
         private HuggingFaceTokenizer tokenizer;
         private boolean includeTokenTypes;
+        private boolean int32;
         private Batchifier batchifier = Batchifier.STACK;
         private boolean detail;
 
@@ -208,6 +212,17 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
             return this;
         }
 
+        /**
+         * Sets if use int32 datatype for the {@link Translator}.
+         *
+         * @param int32 true to include token types
+         * @return this builder
+         */
+        public Builder optInt32(boolean int32) {
+            this.int32 = int32;
+            return this;
+        }
+
         /**
          * Sets the {@link Batchifier} for the {@link Translator}.
          *
@@ -237,6 +252,7 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
          */
         public void configure(Map<String, ?> arguments) {
             optIncludeTokenTypes(ArgumentsUtil.booleanValue(arguments, "includeTokenTypes"));
+            optInt32(ArgumentsUtil.booleanValue(arguments, "int32"));
             String batchifierStr = ArgumentsUtil.stringValue(arguments, "batchifier", "stack");
             optDetail(ArgumentsUtil.booleanValue(arguments, "detail"));
             optBatchifier(Batchifier.fromString(batchifierStr));
@@ -250,7 +266,7 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
          */
         public QuestionAnsweringTranslator build() throws IOException {
             return new QuestionAnsweringTranslator(
-                    tokenizer, includeTokenTypes, batchifier, detail);
+                    tokenizer, includeTokenTypes, int32, batchifier, detail);
         }
     }
 }
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslator.java
index 25224d207..82c9112d3 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslator.java
@@ -37,13 +37,18 @@ public class TextClassificationTranslator implements Translator<String, Classifi
 
     private HuggingFaceTokenizer tokenizer;
     private boolean includeTokenTypes;
+    private boolean int32;
     private Batchifier batchifier;
     private PretrainedConfig config;
 
     TextClassificationTranslator(
-            HuggingFaceTokenizer tokenizer, boolean includeTokenTypes, Batchifier batchifier) {
+            HuggingFaceTokenizer tokenizer,
+            boolean includeTokenTypes,
+            boolean int32,
+            Batchifier batchifier) {
         this.tokenizer = tokenizer;
         this.includeTokenTypes = includeTokenTypes;
+        this.int32 = int32;
         this.batchifier = batchifier;
     }
 
@@ -67,7 +72,7 @@ public class TextClassificationTranslator implements Translator<String, Classifi
     @Override
     public NDList processInput(TranslatorContext ctx, String input) {
         Encoding encoding = tokenizer.encode(input);
-        return encoding.toNDList(ctx.getNDManager(), includeTokenTypes);
+        return encoding.toNDList(ctx.getNDManager(), includeTokenTypes, int32);
     }
 
     /** {@inheritDoc} */
@@ -77,7 +82,7 @@ public class TextClassificationTranslator implements Translator<String, Classifi
         Encoding[] encodings = tokenizer.batchEncode(inputs);
         NDList[] batch = new NDList[encodings.length];
         for (int i = 0; i < encodings.length; ++i) {
-            batch[i] = encodings[i].toNDList(manager, includeTokenTypes);
+            batch[i] = encodings[i].toNDList(manager, includeTokenTypes, int32);
         }
         return batchifier.batchify(batch);
     }
@@ -148,6 +153,7 @@ public class TextClassificationTranslator implements Translator<String, Classifi
 
         private HuggingFaceTokenizer tokenizer;
         private boolean includeTokenTypes;
+        private boolean int32;
         private Batchifier batchifier = Batchifier.STACK;
 
         Builder(HuggingFaceTokenizer tokenizer) {
@@ -165,6 +171,17 @@ public class TextClassificationTranslator implements Translator<String, Classifi
             return this;
         }
 
+        /**
+         * Sets if use int32 datatype for the {@link Translator}.
+         *
+         * @param int32 true to include token types
+         * @return this builder
+         */
+        public Builder optInt32(boolean int32) {
+            this.int32 = int32;
+            return this;
+        }
+
         /**
          * Sets the {@link Batchifier} for the {@link Translator}.
          *
@@ -183,6 +200,7 @@ public class TextClassificationTranslator implements Translator<String, Classifi
          */
         public void configure(Map<String, ?> arguments) {
             optIncludeTokenTypes(ArgumentsUtil.booleanValue(arguments, "includeTokenTypes"));
+            optInt32(ArgumentsUtil.booleanValue(arguments, "int32"));
             String batchifierStr = ArgumentsUtil.stringValue(arguments, "batchifier", "stack");
             optBatchifier(Batchifier.fromString(batchifierStr));
         }
@@ -194,7 +212,8 @@ public class TextClassificationTranslator implements Translator<String, Classifi
          * @throws IOException if I/O error occurs
          */
         public TextClassificationTranslator build() throws IOException {
-            return new TextClassificationTranslator(tokenizer, includeTokenTypes, batchifier);
+            return new TextClassificationTranslator(
+                    tokenizer, includeTokenTypes, int32, batchifier);
         }
     }
 }
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
index 0d91734d5..81e69b25e 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
@@ -44,6 +44,7 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
     private boolean normalize;
     private String pooling;
     private boolean includeTokenTypes;
+    private boolean int32;
     private String dense;
     private String denseActivation;
     private String layerNorm;
@@ -56,6 +57,7 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
             String pooling,
             boolean normalize,
             boolean includeTokenTypes,
+            boolean int32,
             String dense,
             String denseActivation,
             String layerNorm) {
@@ -64,6 +66,7 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
         this.pooling = pooling;
         this.normalize = normalize;
         this.includeTokenTypes = includeTokenTypes;
+        this.int32 = int32;
         this.dense = dense;
         this.denseActivation = denseActivation;
         this.layerNorm = layerNorm;
@@ -108,14 +111,8 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
     public NDList processInput(TranslatorContext ctx, String input) {
         NDManager manager = ctx.getNDManager();
         Encoding encoding = tokenizer.encode(input);
-        NDList list = new NDList();
-        list.add(manager.create(encoding.getIds()));
-        NDArray inputAttentionMask = manager.create(encoding.getAttentionMask());
-        list.add(inputAttentionMask);
-        ctx.setAttachment("attentionMask", inputAttentionMask);
-        if (includeTokenTypes) {
-            list.add(manager.create(encoding.getTypeIds()));
-        }
+        NDList list = encoding.toNDList(manager, includeTokenTypes, int32);
+        ctx.setAttachment("attentionMask", list.get(1));
         return list;
     }
 
@@ -124,24 +121,8 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
     public NDList batchProcessInput(TranslatorContext ctx, List<String> inputs) {
         NDManager manager = ctx.getNDManager();
         Encoding[] encodings = tokenizer.batchEncode(inputs);
-        long[][] ids = new long[encodings.length][];
-        long[][] attentionMask = new long[encodings.length][];
-        long[][] typeIds = new long[encodings.length][];
-        for (int i = 0; i < encodings.length; i++) {
-            ids[i] = encodings[i].getIds();
-            attentionMask[i] = encodings[i].getAttentionMask();
-            if (includeTokenTypes) {
-                typeIds[i] = encodings[i].getTypeIds();
-            }
-        }
-        NDList list = new NDList();
-        list.add(manager.create(ids));
-        NDArray inputAttentionMask = manager.create(attentionMask);
-        list.add(inputAttentionMask);
-        ctx.setAttachment("attentionMask", inputAttentionMask);
-        if (includeTokenTypes) {
-            list.add(manager.create(typeIds));
-        }
+        NDList list = Encoding.toNDList(encodings, manager, includeTokenTypes, int32);
+        ctx.setAttachment("attentionMask", list.get(1));
         return list;
     }
 
@@ -290,6 +271,7 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
         private boolean normalize = true;
         private String pooling = "mean";
         private boolean includeTokenTypes;
+        private boolean int32;
         private String dense;
         private String denseActivation;
         private String layerNorm;
@@ -351,6 +333,17 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
             return this;
         }
 
+        /**
+         * Sets if use int32 datatype for the {@link Translator}.
+         *
+         * @param int32 true to include token types
+         * @return this builder
+         */
+        public Builder optInt32(boolean int32) {
+            this.int32 = int32;
+            return this;
+        }
+
         /**
          * Sets the dense layer model file for the {@link Translator}.
          *
@@ -395,6 +388,7 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
             optNormalize(ArgumentsUtil.booleanValue(arguments, "normalize", true));
             optPoolingMode(ArgumentsUtil.stringValue(arguments, "pooling", "mean"));
             optIncludeTokenTypes(ArgumentsUtil.booleanValue(arguments, "includeTokenTypes"));
+            optInt32(ArgumentsUtil.booleanValue(arguments, "int32"));
             optDense(ArgumentsUtil.stringValue(arguments, "dense"));
             optDenseActivation(ArgumentsUtil.stringValue(arguments, "denseActivation"));
             optLayerNorm(ArgumentsUtil.stringValue(arguments, "layerNorm"));
@@ -413,6 +407,7 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
                     pooling,
                     normalize,
                     includeTokenTypes,
+                    int32,
                     dense,
                     denseActivation,
                     layerNorm);
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 c22b4011a..43aed0db2 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
@@ -38,6 +38,7 @@ public class TokenClassificationTranslator implements Translator<String, NamedEn
 
     private HuggingFaceTokenizer tokenizer;
     private boolean includeTokenTypes;
+    private boolean int32;
     private boolean softmax;
     private Batchifier batchifier;
     private PretrainedConfig config;
@@ -45,10 +46,12 @@ public class TokenClassificationTranslator implements Translator<String, NamedEn
     TokenClassificationTranslator(
             HuggingFaceTokenizer tokenizer,
             boolean includeTokenTypes,
+            boolean int32,
             boolean softmax,
             Batchifier batchifier) {
         this.tokenizer = tokenizer;
         this.includeTokenTypes = includeTokenTypes;
+        this.int32 = int32;
         this.softmax = softmax;
         this.batchifier = batchifier;
     }
@@ -74,7 +77,7 @@ public class TokenClassificationTranslator implements Translator<String, NamedEn
     public NDList processInput(TranslatorContext ctx, String input) {
         Encoding encoding = tokenizer.encode(input);
         ctx.setAttachment("encoding", encoding);
-        return encoding.toNDList(ctx.getNDManager(), includeTokenTypes);
+        return encoding.toNDList(ctx.getNDManager(), includeTokenTypes, int32);
     }
 
     /** {@inheritDoc} */
@@ -85,7 +88,7 @@ public class TokenClassificationTranslator implements Translator<String, NamedEn
         ctx.setAttachment("encodings", encodings);
         NDList[] batch = new NDList[encodings.length];
         for (int i = 0; i < encodings.length; ++i) {
-            batch[i] = encodings[i].toNDList(manager, includeTokenTypes);
+            batch[i] = encodings[i].toNDList(manager, includeTokenTypes, int32);
         }
         return batchifier.batchify(batch);
     }
@@ -170,6 +173,7 @@ public class TokenClassificationTranslator implements Translator<String, NamedEn
 
         private HuggingFaceTokenizer tokenizer;
         private boolean includeTokenTypes;
+        private boolean int32;
         private boolean softmax = true;
         private Batchifier batchifier = Batchifier.STACK;
 
@@ -188,6 +192,17 @@ public class TokenClassificationTranslator implements Translator<String, NamedEn
             return this;
         }
 
+        /**
+         * Sets if use int32 datatype for the {@link Translator}.
+         *
+         * @param int32 true to include token types
+         * @return this builder
+         */
+        public Builder optInt32(boolean int32) {
+            this.int32 = int32;
+            return this;
+        }
+
         /**
          * Sets if implement softmax operation for the {@link Translator}.
          *
@@ -217,6 +232,7 @@ public class TokenClassificationTranslator implements Translator<String, NamedEn
          */
         public void configure(Map<String, ?> arguments) {
             optIncludeTokenTypes(ArgumentsUtil.booleanValue(arguments, "includeTokenTypes"));
+            optInt32(ArgumentsUtil.booleanValue(arguments, "int32"));
             optSoftmax(ArgumentsUtil.booleanValue(arguments, "softmax", true));
             String batchifierStr = ArgumentsUtil.stringValue(arguments, "batchifier", "stack");
             optBatchifier(Batchifier.fromString(batchifierStr));
@@ -230,7 +246,7 @@ public class TokenClassificationTranslator implements Translator<String, NamedEn
          */
         public TokenClassificationTranslator build() throws IOException {
             return new TokenClassificationTranslator(
-                    tokenizer, includeTokenTypes, softmax, batchifier);
+                    tokenizer, includeTokenTypes, int32, softmax, batchifier);
         }
     }
 }
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index 9a0221074..ddd35ca4e 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -14,7 +14,7 @@ lightgbm = "3.2.110"
 tensorrt = "8.4.1"
 fasttext = "0.9.2"
 sentencepiece = "0.2.0"
-tokenizers = "0.20.3"
+tokenizers = "0.21.0"
 
 slf4j = "2.0.16"
 log4j_slf4j = "2.24.1"


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 44103f3a8..d44e426f7 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
@@ -39,7 +39,7 @@ public class HuggingFaceTokenizerTest {
     public void testVersion() {
         try (HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.newInstance("bert-base-cased")) {
             String djlVersion = Engine.getDjlVersion();
-            Assert.assertEquals(tokenizer.getVersion(), "0.20.3-" + djlVersion);
+            Assert.assertEquals(tokenizer.getVersion(), "0.21.0-" + djlVersion);
         }
     }