Repository: deepjavalibrary/djl
Pull request #3559 (2024-12-14 18:07:57)

[api] Allows to configure custom batchifer
## 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/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index 3d93c9878..aaefd2f88 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -1016,15 +1016,13 @@ public interface NDArray extends NDResource, BytesSupplier {
         return getManager().create(getShape());
     }
 
-    ////////////////////////////////////////
-    ////////////////////////////////////////
+    /*
     // Operations
-    ////////////////////////////////////////
-    ////////////////////////////////////////
+    */
 
-    ////////////////////////////////////////
+    /*
     // Operations: Element Comparison
-    ////////////////////////////////////////
+    */
 
     /**
      * Returns {@code true} if all elements in this {@code NDArray} are equal to the {@link Number}.
@@ -1375,9 +1373,9 @@ public interface NDArray extends NDResource, BytesSupplier {
      */
     NDArray lte(NDArray other);
 
-    ////////////////////////////////////////
+    /*
     // Operations: Element Arithmetic
-    ////////////////////////////////////////
+    */
 
     /**
      * Adds a number to this {@code NDArray} element-wise.
@@ -2006,9 +2004,9 @@ public interface NDArray extends NDResource, BytesSupplier {
      */
     NDArray minimum(NDArray other);
 
-    ////////////////////////////////////////
+    /*
     // Operations: Basic Numeric
-    ////////////////////////////////////////
+    */
 
     /**
      * Returns the numerical negative {@code NDArray} element-wise.
@@ -2516,9 +2514,9 @@ public interface NDArray extends NDResource, BytesSupplier {
      */
     NDArray toRadians();
 
-    ////////////////////////////////////////
+    /*
     // Operations: Reduction
-    ////////////////////////////////////////
+    */
 
     /**
      * Returns the maximum of this {@code NDArray}.
@@ -3161,9 +3159,9 @@ public interface NDArray extends NDResource, BytesSupplier {
      */
     NDArray trace(int offset, int axis1, int axis2);
 
-    ////////////////////////////////////////
+    /*
     // Operations: Shapes and Arrays Manipulation
-    ////////////////////////////////////////
+    */
 
     /**
      * Splits this {@code NDArray} into multiple sub{@code NDArray}s given sections along first
@@ -3885,9 +3883,9 @@ public interface NDArray extends NDResource, BytesSupplier {
         return getNDArrayInternal().concat(new NDList(array), axis);
     }
 
-    ////////////////////////////////////////
+    /*
     // Operations: Logical Op
-    ////////////////////////////////////////
+    */
 
     /**
      * Returns the truth value of this {@code NDArray} AND the other {@code NDArray} element-wise.
@@ -4010,9 +4008,9 @@ public interface NDArray extends NDResource, BytesSupplier {
      */
     NDArray logicalNot();
 
-    ////////////////////////////////////////
+    /*
     // Operations: Other
-    ////////////////////////////////////////
+    */
 
     /**
      * Returns the indices that would sort this {@code NDArray}.
diff --git a/api/src/main/java/ai/djl/ndarray/NDArrays.java b/api/src/main/java/ai/djl/ndarray/NDArrays.java
index 0e1c0922a..6391b4076 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArrays.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArrays.java
@@ -32,9 +32,9 @@ public final class NDArrays {
         }
     }
 
-    ////////////////////////////////////////
+    /*
     // Operations: Element Comparison
-    ////////////////////////////////////////
+    */
 
     /**
      * Returns {@code true} if all elements in {@link NDArray} a are equal to {@link NDArray} b.
@@ -788,9 +788,9 @@ public final class NDArrays {
         return data.sequenceMask(sequenceLength);
     }
 
-    ////////////////////////////////////////
+    /*
     // Operations: Element Arithmetic
-    ////////////////////////////////////////
+    */
 
     /**
      * Adds a number to the {@link NDArray} element-wise.
diff --git a/api/src/main/java/ai/djl/ndarray/internal/NDArrayEx.java b/api/src/main/java/ai/djl/ndarray/internal/NDArrayEx.java
index b25852ad5..c4b874220 100644
--- a/api/src/main/java/ai/djl/ndarray/internal/NDArrayEx.java
+++ b/api/src/main/java/ai/djl/ndarray/internal/NDArrayEx.java
@@ -28,9 +28,9 @@ import java.util.List;
 @SuppressWarnings("MissingJavadocMethod")
 public interface NDArrayEx {
 
-    ////////////////////////////////////////
+    /*
     // NDArrays
-    ////////////////////////////////////////
+    */
 
     /**
      * Applies reverse division with a scalar - i.e., (n / thisArrayValues).
@@ -144,9 +144,9 @@ public interface NDArrayEx {
      */
     NDArray rpowi(Number n);
 
-    ////////////////////////////////////////
+    /*
     // Activations
-    ////////////////////////////////////////
+    */
 
     /**
      * Computes rectified linear activation.
@@ -179,9 +179,9 @@ public interface NDArrayEx {
         return getArray().exp().add(1).log2().tanh().mul(getArray());
     }
 
-    ////////////////////////////////////////
+    /*
     // Pooling Operations
-    ////////////////////////////////////////
+    */
 
     NDArray maxPool(Shape kernelShape, Shape stride, Shape padding, boolean ceilMode);
 
@@ -201,9 +201,9 @@ public interface NDArrayEx {
 
     NDArray globalLpPool(float normType);
 
-    ////////////////////////////////////////
+    /*
     // Optimizer
-    ////////////////////////////////////////
+    */
 
     void adadeltaUpdate(
             NDList inputs,
@@ -268,9 +268,9 @@ public interface NDArrayEx {
             float momentum,
             boolean lazyUpdate);
 
-    ////////////////////////////////////////
+    /*
     // Neural network
-    ////////////////////////////////////////
+    */
 
     NDList convolution(
             NDArray input,
@@ -398,9 +398,9 @@ public interface NDArrayEx {
             boolean bidirectional,
             boolean batchFirst);
 
-    ////////////////////////////////////////
+    /*
     // Image and CV
-    ////////////////////////////////////////
+    */
 
     /**
      * Normalizes a NDArray of shape CHW or NCHW with mean and standard deviation.
@@ -489,9 +489,9 @@ public interface NDArrayEx {
     // TODO: add TorchVision support
     NDArray randomColorJitter(float brightness, float contrast, float saturation, float hue);
 
-    ////////////////////////////////////////
+    /*
     // Miscellaneous
-    ////////////////////////////////////////
+    */
 
     /**
      * Returns an {@link NDArrayIndexer}.
diff --git a/api/src/main/java/ai/djl/translate/Batchifier.java b/api/src/main/java/ai/djl/translate/Batchifier.java
index 32fa23f94..9ffcb781d 100644
--- a/api/src/main/java/ai/djl/translate/Batchifier.java
+++ b/api/src/main/java/ai/djl/translate/Batchifier.java
@@ -13,6 +13,7 @@
 package ai.djl.translate;
 
 import ai.djl.ndarray.NDList;
+import ai.djl.util.ClassLoaderUtils;
 
 import java.io.Serializable;
 import java.util.Arrays;
@@ -47,7 +48,12 @@ public interface Batchifier extends Serializable {
             case "none":
                 return null;
             default:
-                throw new IllegalArgumentException("Invalid batchifier name");
+                ClassLoader cl = ClassLoaderUtils.getContextClassLoader();
+                Batchifier b = ClassLoaderUtils.initClass(cl, Batchifier.class, name);
+                if (b == null) {
+                    throw new IllegalArgumentException("Invalid batchifier name: " + name);
+                }
+                return b;
         }
     }
 
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArrayEx.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArrayEx.java
index 00a7a0f25..4ea8d07fa 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArrayEx.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArrayEx.java
@@ -31,6 +31,7 @@ import java.util.Arrays;
 import java.util.List;
 
 /** {@code MxNDArrayEx} is the MXNet implementation of the {@link NDArrayEx}. */
+@SuppressWarnings("dangling-doc-comments")
 class MxNDArrayEx implements NDArrayEx {
 
     private MxNDArray array;
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/jna/JnaUtils.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/jna/JnaUtils.java
index 3ec9aab02..cb31f3f35 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/jna/JnaUtils.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/jna/JnaUtils.java
@@ -51,7 +51,7 @@ import java.util.concurrent.ConcurrentHashMap;
 /**
  * A class containing utilities to interact with the MXNet Engine's Java Native Access (JNA) layer.
  */
-@SuppressWarnings("MissingJavadocMethod")
+@SuppressWarnings({"MissingJavadocMethod", "dangling-doc-comments"})
 public final class JnaUtils {
 
     public static final ObjectPool<PointerByReference> REFS =
diff --git a/examples/src/main/java/ai/djl/examples/inference/nlp/BertQaInference.java b/examples/src/main/java/ai/djl/examples/inference/nlp/BertQaInference.java
index d02da5437..507bf2589 100644
--- a/examples/src/main/java/ai/djl/examples/inference/nlp/BertQaInference.java
+++ b/examples/src/main/java/ai/djl/examples/inference/nlp/BertQaInference.java
@@ -48,7 +48,7 @@ public final class BertQaInference {
 
     public static void main(String[] args) throws IOException, TranslateException, ModelException {
         String answer = BertQaInference.predict();
-        logger.info("Answer: {}", answer);
+        logger.info("Output: {}", answer);
     }
 
     public static String predict() throws IOException, TranslateException, ModelException {
@@ -69,6 +69,7 @@ public final class BertQaInference {
                                 "djl://ai.djl.huggingface.pytorch/deepset/minilm-uncased-squad2")
                         .optEngine("PyTorch")
                         .optTranslatorFactory(new QuestionAnsweringTranslatorFactory())
+                        .optArgument("detail", true)
                         .optProgress(new ProgressBar())
                         .build();
 
diff --git a/extensions/tokenizers/rust/src/lib.rs b/extensions/tokenizers/rust/src/lib.rs
index 3d4cd1264..a7d1072c9 100644
--- a/extensions/tokenizers/rust/src/lib.rs
+++ b/extensions/tokenizers/rust/src/lib.rs
@@ -39,8 +39,8 @@ use tk::models::bpe::BPE;
 use tk::tokenizer::{EncodeInput, Encoding};
 use tk::utils::padding::{PaddingParams, PaddingStrategy};
 use tk::utils::truncation::{TruncationParams, TruncationStrategy};
-use tk::Tokenizer;
 use tk::Offsets;
+use tk::Tokenizer;
 
 #[cfg(not(target_os = "android"))]
 use tk::FromPretrainedParameters;
@@ -407,6 +407,31 @@ pub extern "system" fn Java_ai_djl_huggingface_tokenizers_jni_TokenizersLibrary_
     array
 }
 
+#[no_mangle]
+pub extern "system" fn Java_ai_djl_huggingface_tokenizers_jni_TokenizersLibrary_getSequenceIds<
+    'local,
+>(
+    env: JNIEnv<'local>,
+    _: JObject,
+    handle: jlong,
+) -> JLongArray<'local> {
+    let encoding = cast_handle::<Encoding>(handle);
+    let sequence_ids = encoding.get_sequence_ids();
+    let len = sequence_ids.len() as jsize;
+    let mut long_ids: Vec<jlong> = Vec::new();
+    for i in sequence_ids {
+        if let Some(sequence_id) = i {
+            long_ids.push(sequence_id as jlong)
+        } else {
+            long_ids.push(-1)
+        }
+    }
+
+    let array = env.new_long_array(len).unwrap();
+    env.set_long_array_region(&array, 0, &long_ids).unwrap();
+    array
+}
+
 #[no_mangle]
 pub extern "system" fn Java_ai_djl_huggingface_tokenizers_jni_TokenizersLibrary_getTokens<
     'local,
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 301490a2b..4877f53aa 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
@@ -25,6 +25,7 @@ public class Encoding {
     private long[] typeIds;
     private String[] tokens;
     private long[] wordIds;
+    private long[] sequenceIds;
     private long[] attentionMask;
     private long[] specialTokenMask;
     private CharSpan[] charTokenSpans;
@@ -36,6 +37,7 @@ public class Encoding {
             long[] typeIds,
             String[] tokens,
             long[] wordIds,
+            long[] sequenceIds,
             long[] attentionMask,
             long[] specialTokenMask,
             CharSpan[] charTokenSpans,
@@ -45,6 +47,7 @@ public class Encoding {
         this.typeIds = typeIds;
         this.tokens = tokens;
         this.wordIds = wordIds;
+        this.sequenceIds = sequenceIds;
         this.attentionMask = attentionMask;
         this.specialTokenMask = specialTokenMask;
         this.charTokenSpans = charTokenSpans;
@@ -109,6 +112,15 @@ public class Encoding {
         return wordIds;
     }
 
+    /**
+     * Returns the sequence ids.
+     *
+     * @return the sequence ids
+     */
+    public long[] getSequenceIds() {
+        return sequenceIds;
+    }
+
     /**
      * Returns the attention masks.
      *
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 cd6eaf289..c4c04c4df 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
@@ -623,6 +623,7 @@ public final class HuggingFaceTokenizer extends NativeResource<Long> implements
         long[] typeIds = TokenizersLibrary.LIB.getTypeIds(encoding);
         String[] tokens = TokenizersLibrary.LIB.getTokens(encoding);
         long[] wordIds = TokenizersLibrary.LIB.getWordIds(encoding);
+        long[] sequenceIds = TokenizersLibrary.LIB.getSequenceIds(encoding);
         long[] attentionMask = TokenizersLibrary.LIB.getAttentionMask(encoding);
         long[] specialTokenMask = TokenizersLibrary.LIB.getSpecialTokenMask(encoding);
         CharSpan[] charSpans = TokenizersLibrary.LIB.getTokenCharSpans(encoding);
@@ -646,6 +647,7 @@ public final class HuggingFaceTokenizer extends NativeResource<Long> implements
                 typeIds,
                 tokens,
                 wordIds,
+                sequenceIds,
                 attentionMask,
                 specialTokenMask,
                 charSpans,
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/jni/CharSpan.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/jni/CharSpan.java
index 7478d0e8b..9a15980b4 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/jni/CharSpan.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/jni/CharSpan.java
@@ -10,8 +10,6 @@
  * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
  * and limitations under the License.
  */
-
-/** A class holds character span information. */
 package ai.djl.huggingface.tokenizers.jni;
 
 /** A class holds character span information. */
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/jni/TokenizersLibrary.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/jni/TokenizersLibrary.java
index 7ef895091..994694fe2 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/jni/TokenizersLibrary.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/jni/TokenizersLibrary.java
@@ -50,6 +50,8 @@ public final class TokenizersLibrary {
 
     public native long[] getWordIds(long encoding);
 
+    public native long[] getSequenceIds(long encoding);
+
     public native String[] getTokens(long encoding);
 
     public native long[] getAttentionMask(long encoding);
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 f0b0fac02..1b3b4b098 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
@@ -23,12 +23,14 @@ 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 ai.djl.util.PairList;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /** The translator for Huggingface question answering model. */
 public class QuestionAnsweringTranslator implements Translator<QAInput, String> {
@@ -36,12 +38,17 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
     private HuggingFaceTokenizer tokenizer;
     private boolean includeTokenTypes;
     private Batchifier batchifier;
+    private boolean detail;
 
     QuestionAnsweringTranslator(
-            HuggingFaceTokenizer tokenizer, boolean includeTokenTypes, Batchifier batchifier) {
+            HuggingFaceTokenizer tokenizer,
+            boolean includeTokenTypes,
+            Batchifier batchifier,
+            boolean detail) {
         this.tokenizer = tokenizer;
         this.includeTokenTypes = includeTokenTypes;
         this.batchifier = batchifier;
+        this.detail = detail;
     }
 
     /** {@inheritDoc} */
@@ -102,6 +109,27 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
             startLogits = startLogits.duplicate();
             endLogits = endLogits.duplicate();
         }
+        if (detail) {
+            // exclude undesired sequences
+            long[] sequenceIds = encoding.getSequenceIds();
+            List<Integer> undesired = new ArrayList<>();
+            for (int i = 0; i < sequenceIds.length; ++i) {
+                if (sequenceIds[i] == 0) {
+                    undesired.add(i);
+                }
+            }
+            int[] idx = undesired.stream().mapToInt(Integer::intValue).toArray();
+            NDIndex ndIndex = new NDIndex("{}", list.getManager().create(idx));
+            startLogits.set(ndIndex, -100000f);
+            endLogits.set(ndIndex, -100000f);
+
+            // normalize
+            startLogits = startLogits.sub(startLogits.max()).exp();
+            startLogits = startLogits.div(startLogits.sum());
+            endLogits = endLogits.sub(endLogits.max()).exp();
+            endLogits = endLogits.div(endLogits.sum());
+        }
+
         // exclude <CLS>, TODO: exclude impossible ids properly and handle max answer length
         startLogits.set(new NDIndex(0), -100000);
         endLogits.set(new NDIndex(0), -100000);
@@ -111,12 +139,26 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
             int tmp = startIdx;
             startIdx = endIdx;
             endIdx = tmp;
+            NDArray tmpArray = startLogits;
+            startLogits = endLogits;
+            endLogits = tmpArray;
         }
         long[] indices = encoding.getIds();
         int len = endIdx - startIdx + 1;
         long[] ids = new long[len];
         System.arraycopy(indices, startIdx, ids, 0, len);
-        return tokenizer.decode(ids).trim();
+        String answer = tokenizer.decode(ids).trim();
+        if (detail) {
+            float score = startLogits.getFloat(startIdx) * endLogits.getFloat(endIdx);
+
+            Map<String, Object> dict = new ConcurrentHashMap<>();
+            dict.put("score", score);
+            dict.put("start", startIdx);
+            dict.put("end", endIdx);
+            dict.put("answer", answer);
+            return JsonUtils.toJson(dict);
+        }
+        return answer;
     }
 
     /**
@@ -149,6 +191,7 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
         private HuggingFaceTokenizer tokenizer;
         private boolean includeTokenTypes;
         private Batchifier batchifier = Batchifier.STACK;
+        private boolean detail;
 
         Builder(HuggingFaceTokenizer tokenizer) {
             this.tokenizer = tokenizer;
@@ -176,6 +219,17 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
             return this;
         }
 
+        /**
+         * Sets if output detail for the {@link Translator}.
+         *
+         * @param detail true to output detail
+         * @return this builder
+         */
+        public Builder optDetail(boolean detail) {
+            this.detail = detail;
+            return this;
+        }
+
         /**
          * Configures the builder with the model arguments.
          *
@@ -184,6 +238,7 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
         public void configure(Map<String, ?> arguments) {
             optIncludeTokenTypes(ArgumentsUtil.booleanValue(arguments, "includeTokenTypes"));
             String batchifierStr = ArgumentsUtil.stringValue(arguments, "batchifier", "stack");
+            optDetail(ArgumentsUtil.booleanValue(arguments, "detail"));
             optBatchifier(Batchifier.fromString(batchifierStr));
         }
 
@@ -194,7 +249,8 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
          * @throws IOException if I/O error occurs
          */
         public QuestionAnsweringTranslator build() throws IOException {
-            return new QuestionAnsweringTranslator(tokenizer, includeTokenTypes, batchifier);
+            return new QuestionAnsweringTranslator(
+                    tokenizer, includeTokenTypes, batchifier, detail);
         }
     }
 }


Test patch:
diff --git a/api/src/test/java/ai/djl/translate/BatchifierTest.java b/api/src/test/java/ai/djl/translate/BatchifierTest.java
index c7ef460e9..92afd9ea4 100644
--- a/api/src/test/java/ai/djl/translate/BatchifierTest.java
+++ b/api/src/test/java/ai/djl/translate/BatchifierTest.java
@@ -22,6 +22,17 @@ public class BatchifierTest {
         Batchifier batchifier = Batchifier.fromString("stack");
         Assert.assertEquals(batchifier, Batchifier.STACK);
 
+        batchifier = Batchifier.fromString("none");
+        Assert.assertNull(batchifier);
+
+        batchifier = Batchifier.fromString("padding");
+        Assert.assertNotNull(batchifier);
+        Assert.assertEquals(batchifier.getClass(), SimplePaddingStackBatchifier.class);
+
+        batchifier = Batchifier.fromString("ai.djl.translate.SimplePaddingStackBatchifier");
+        Assert.assertNotNull(batchifier);
+        Assert.assertEquals(batchifier.getClass(), SimplePaddingStackBatchifier.class);
+
         Assert.assertThrows(() -> Batchifier.fromString("invalid"));
     }
 }
diff --git a/basicdataset/src/test/java/ai/djl/basicdataset/MnistTest.java b/basicdataset/src/test/java/ai/djl/basicdataset/MnistTest.java
index 8d0e609cc..792ff77b7 100644
--- a/basicdataset/src/test/java/ai/djl/basicdataset/MnistTest.java
+++ b/basicdataset/src/test/java/ai/djl/basicdataset/MnistTest.java
@@ -157,7 +157,7 @@ public class MnistTest {
                             .optUsage(Dataset.Usage.TEST)
                             .optRepository(repository)
                             .setSampling(32, false)
-                            .optLabelBatchifier(new StackBatchifier() {})
+                            .optLabelBatchifier(new StackBatchifier())
                             .build();
 
             try (Trainer trainer = model.newTrainer(config)) {
diff --git a/examples/src/test/java/ai/djl/examples/inference/nlp/BertQaTest.java b/examples/src/test/java/ai/djl/examples/inference/nlp/BertQaTest.java
index 21a70ff67..1a464d8f7 100644
--- a/examples/src/test/java/ai/djl/examples/inference/nlp/BertQaTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/nlp/BertQaTest.java
@@ -15,6 +15,9 @@ package ai.djl.examples.inference.nlp;
 import ai.djl.ModelException;
 import ai.djl.testing.TestRequirements;
 import ai.djl.translate.TranslateException;
+import ai.djl.util.JsonUtils;
+
+import com.google.gson.JsonObject;
 
 import org.testng.Assert;
 import org.testng.annotations.Test;
@@ -28,6 +31,8 @@ public class BertQaTest {
         TestRequirements.linux();
 
         String result = BertQaInference.predict();
-        Assert.assertEquals(result, "december 2004");
+        JsonObject json = JsonUtils.GSON.fromJson(result, JsonObject.class);
+        String answer = json.get("answer").getAsString();
+        Assert.assertEquals(answer, "december 2004");
     }
 }
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 4e11fee74..44103f3a8 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
@@ -67,6 +67,7 @@ public class HuggingFaceTokenizerTest {
             long[] ids = {101, 8667, 117, 194, 112, 1155, 106, 1731, 1132, 1128, 100, 136, 102};
             long[] typeIds = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
             long[] wordIds = {-1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, -1};
+            long[] sequenceIds = {-1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -1};
             long[] attentionMask = {1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1};
             long[] specialTokenMask = {1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1};
 
@@ -74,6 +75,7 @@ public class HuggingFaceTokenizerTest {
             Assert.assertEquals(ids, encoding.getIds());
             Assert.assertEquals(typeIds, encoding.getTypeIds());
             Assert.assertEquals(wordIds, encoding.getWordIds());
+            Assert.assertEquals(sequenceIds, encoding.getSequenceIds());
             Assert.assertEquals(attentionMask, encoding.getAttentionMask());
             Assert.assertEquals(specialTokenMask, encoding.getSpecialTokenMask());
 
@@ -104,6 +106,10 @@ public class HuggingFaceTokenizerTest {
                 Assert.assertEquals(charSpansExpected[i].getEnd(), charSpansResult[i].getEnd());
             }
 
+            encoding = tokenizer.encode(inputs[0], inputs[1]);
+            sequenceIds = new long[] {-1, 0, 0, 0, 0, 0, 0, -1, 1, 1, 1, 1, 1, -1};
+            Assert.assertEquals(encoding.getSequenceIds(), sequenceIds);
+
             Assert.assertThrows(() -> tokenizer.encode((String) null));
             Assert.assertThrows(() -> tokenizer.encode(new String[] {null}));
             Assert.assertThrows(() -> tokenizer.encode(null, null));