[tokenizers] Make encoding types to int

deepjavalibrary · Sep 12, 2024 · 2c14fde · 2c14fde
1 parent 8247e57
commit 2c14fde
Showing 1 changed file with 8 additions and 3 deletions.
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
@@ -16,6 +16,8 @@
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
 
+import java.util.Arrays;
+
 /** A class holds token encoding information. */
 public class Encoding {
 
@@ -59,10 +61,13 @@ protected Encoding(
      */
     public NDList toNDList(NDManager manager, boolean withTokenType) {
         NDList list = new NDList(withTokenType ? 3 : 2);
-        list.add(manager.create(ids));
-        list.add(manager.create(attentionMask));
+        int[] ids2 = Arrays.stream(ids).mapToInt(i -> (int) i).toArray();
+        int[] attentionMask2 = Arrays.stream(attentionMask).mapToInt(i -> (int) i).toArray();
+        list.add(manager.create(ids2));
+        list.add(manager.create(attentionMask2));
         if (withTokenType) {
-            list.add(manager.create(typeIds));
+            int[] typeIds2 = Arrays.stream(typeIds).mapToInt(i -> (int) i).toArray();
+            list.add(manager.create(typeIds2));
         }
         return list;
     }