Repository: deepjavalibrary/djl
Pull request #1827 (2022-07-24 18:16:20)

[pt-model-zoo] Make PtBertQATranslator compatible with huggingface model
Change-Id: I16ce250daea8cf2a07a39b064cff9f9bfa58399d

## 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/android/core/src/main/java/ai/djl/android/core/BitmapImageFactory.java b/android/core/src/main/java/ai/djl/android/core/BitmapImageFactory.java
index 573fef16a..f5e29a6db 100644
--- a/android/core/src/main/java/ai/djl/android/core/BitmapImageFactory.java
+++ b/android/core/src/main/java/ai/djl/android/core/BitmapImageFactory.java
@@ -26,7 +26,6 @@ import java.net.URL;
 import java.nio.ByteBuffer;
 import java.nio.file.Path;
 import java.util.List;
-import java.util.stream.IntStream;
 
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.ImageFactory;
@@ -88,27 +87,27 @@ public class BitmapImageFactory extends ImageFactory {
         Shape shape = array.getShape();
         if (shape.dimension() != 3) {
             throw new IllegalArgumentException("Shape should only have three dimension follow CHW");
-        }
-        if (shape.get(0) == 1) {
+        } else if (shape.get(0) == 1 || shape.get(2) == 1) {
             throw new UnsupportedOperationException("Grayscale image is not supported");
         }
         if (NDImageUtils.isCHW(shape)) {
-            array = array.transpose(2, 0, 1);
+            array = array.transpose(1, 2, 0);
             shape = array.getShape();
         }
-        int height = (int) shape.get(1);
-        int width = (int) shape.get(2);
-        int imageArea = width * height;
+
+        int height = (int) shape.get(0);
+        int width = (int) shape.get(1);
         Bitmap bitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888);
         int[] raw = array.toType(DataType.UINT8, false).toUint8Array();
-        IntStream.range(0, imageArea).parallel().forEach(ele -> {
-            int x = ele % width;
-            int y = ele / width;
-            int red = raw[ele] & 0xFF;
-            int green = raw[ele + imageArea] & 0xFF;
-            int blue = raw[ele + imageArea * 2] & 0xFF;
-            bitmap.setPixel(x, y, Color.argb(255, red, green, blue));
-        });
+        for (int h = 0; h < height; ++h) {
+            for (int w = 0; w < width; ++w) {
+                int pos = (h * width + w) * 3;
+                int red = raw[pos];
+                int green = raw[pos + 1];
+                int blue = raw[pos + 2];
+                bitmap.setPixel(w, h, Color.argb(255, red, green, blue));
+            }
+        }
         return new BitMapWrapper(bitmap);
     }
 
diff --git a/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java b/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
index 9a59979c8..c02c36464 100644
--- a/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
@@ -39,7 +39,6 @@ import java.io.OutputStream;
 import java.nio.ByteBuffer;
 import java.nio.file.Path;
 import java.util.List;
-import java.util.stream.IntStream;
 
 import javax.imageio.ImageIO;
 
@@ -91,27 +90,26 @@ public class BufferedImageFactory extends ImageFactory {
         } else if (shape.get(0) == 1 || shape.get(2) == 1) {
             throw new UnsupportedOperationException("Grayscale image is not supported");
         }
-        int height = (int) shape.get(0);
-        int width = (int) shape.get(1);
         if (NDImageUtils.isCHW(shape)) {
-            height = (int) shape.get(1);
-            width = (int) shape.get(2);
+            array = array.transpose(1, 2, 0);
+            shape = array.getShape();
         }
-        int imageArea = width * height;
+        int height = (int) shape.get(0);
+        int width = (int) shape.get(1);
         BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
         int[] raw = array.toType(DataType.UINT8, false).toUint8Array();
-        int finalWidth = width;
-        IntStream.range(0, imageArea)
-                .forEach(
-                        ele -> {
-                            int x = ele % finalWidth;
-                            int y = ele / finalWidth;
-                            int red = raw[ele] & 0xFF;
-                            int green = raw[ele + imageArea] & 0xFF;
-                            int blue = raw[ele + imageArea * 2] & 0xFF;
-                            int rgb = (red << 16) | (green << 8) | blue;
-                            image.setRGB(x, y, rgb);
-                        });
+        int[] pixels = new int[width * height];
+        for (int h = 0; h < height; ++h) {
+            for (int w = 0; w < width; ++w) {
+                int index = h * width + w;
+                int pos = index * 3;
+                int red = raw[pos];
+                int green = raw[pos + 1];
+                int blue = raw[pos + 2];
+                pixels[index] = (red << 16) | (green << 8) | blue;
+            }
+        }
+        image.setRGB(0, 0, width, height, pixels, 0, width);
         return new BufferedImageWrapper(image);
     }
 
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslator.java b/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslator.java
index 65d7ce624..b9e0b5726 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslator.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslator.java
@@ -91,23 +91,27 @@ public class SemanticSegmentationTranslator extends BaseImageTranslator<Image> {
             // change color of pixels in image array where objects have been detected
             for (int h = 0; h < height; h++) {
                 for (int w = 0; w < width; w++) {
+                    int index = h * width + w;
                     int maxi = 0;
                     double maxnum = -Double.MAX_VALUE;
                     for (int i = 0; i < CLASSNUM; i++) {
                         // get score for each i at the h,w pixel of the image
-                        float score = scores[i * (imageSize) + h * width + w];
+                        float score = scores[i * (imageSize) + index];
                         if (score > maxnum) {
                             maxnum = score;
                             maxi = i;
                         }
                     }
                     if (maxi > 0) {
-                        bb.put(h * width + w, colors[maxi][r]);
-                        bb.put(g * imageSize + h * width + w, colors[maxi][g]);
-                        bb.put(b * imageSize + h * width + w, colors[maxi][b]);
+                        bb.put(colors[maxi][r]);
+                        bb.put(colors[maxi][g]);
+                        bb.put(colors[maxi][b]);
+                    } else {
+                        bb.position(bb.position() + 3);
                     }
                 }
             }
+            bb.rewind();
             int originW = (int) ctx.getAttachment("originalWidth");
             int originH = (int) ctx.getAttachment("originalHeight");
             NDArray fullImage =
diff --git a/api/src/main/java/ai/djl/modality/nlp/bert/BertFullTokenizer.java b/api/src/main/java/ai/djl/modality/nlp/bert/BertFullTokenizer.java
index 0e7c28490..f19c2accc 100644
--- a/api/src/main/java/ai/djl/modality/nlp/bert/BertFullTokenizer.java
+++ b/api/src/main/java/ai/djl/modality/nlp/bert/BertFullTokenizer.java
@@ -76,7 +76,7 @@ public class BertFullTokenizer extends BertTokenizer {
 
     /** {@inheritDoc} */
     @Override
-    public String tokenToString(List<String> tokens) {
+    public String buildSentence(List<String> tokens) {
         return String.join(" ", tokens).replace(" ##", "").trim();
     }
 
diff --git a/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/nlp/qa/MxBertQATranslator.java b/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/nlp/qa/MxBertQATranslator.java
index b15e93b7d..6ab3c16fb 100644
--- a/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/nlp/qa/MxBertQATranslator.java
+++ b/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/nlp/qa/MxBertQATranslator.java
@@ -113,7 +113,7 @@ public class MxBertQATranslator extends QATranslator {
         NDArray endLogits = output.get(1).reshape(new Shape(1, -1));
         int startIdx = (int) startLogits.argMax(1).getLong();
         int endIdx = (int) endLogits.argMax(1).getLong();
-        return tokenizer.tokenToString(tokens.subList(startIdx, endIdx + 1));
+        return tokenizer.buildSentence(tokens.subList(startIdx, endIdx + 1));
     }
 
     /**
diff --git a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/qa/PtBertQATranslator.java b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/qa/PtBertQATranslator.java
index cf66d40ad..6a301dfd0 100644
--- a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/qa/PtBertQATranslator.java
+++ b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/qa/PtBertQATranslator.java
@@ -91,9 +91,11 @@ public class PtBertQATranslator extends QATranslator {
         int startIdx = (int) startLogits.argMax().getLong();
         int endIdx = (int) endLogits.argMax().getLong();
         if (startIdx >= endIdx) {
-            return "";
+            int tmp = startIdx;
+            startIdx = endIdx;
+            endIdx = tmp;
         }
-        return tokenizer.tokenToString(tokens.subList(startIdx, endIdx + 1));
+        return tokenizer.buildSentence(tokens.subList(startIdx, endIdx + 1));
     }
 
     /**
diff --git a/extensions/tokenizers/rust/src/lib.rs b/extensions/tokenizers/rust/src/lib.rs
index 2a826c3f5..a02c906de 100644
--- a/extensions/tokenizers/rust/src/lib.rs
+++ b/extensions/tokenizers/rust/src/lib.rs
@@ -105,6 +105,39 @@ pub extern "system" fn Java_ai_djl_huggingface_tokenizers_jni_TokenizersLibrary_
     }
 }
 
+#[no_mangle]
+pub extern "system" fn Java_ai_djl_huggingface_tokenizers_jni_TokenizersLibrary_encodeDual(
+    env: JNIEnv,
+    _: JObject,
+    handle: jlong,
+    text: JString,
+    text_pair: JString,
+    add_special_tokens: jboolean,
+) -> jlong {
+    let tokenizer = cast_handle::<Tokenizer>(handle);
+    let sequence1: String = env
+        .get_string(text)
+        .expect("Couldn't get text string!")
+        .into();
+    let sequence2: String = env
+        .get_string(text_pair)
+        .expect("Couldn't get text_pair string!")
+        .into();
+
+    let input_sequence1 = tk::InputSequence::from(sequence1);
+    let input_sequence2 = tk::InputSequence::from(sequence2);
+    let encoded_input = EncodeInput::Dual(input_sequence1, input_sequence2);
+    let encoding = tokenizer.encode(encoded_input, add_special_tokens == JNI_TRUE);
+
+    match encoding {
+        Ok(output) => to_handle(output),
+        Err(err) => {
+            env.throw(err.to_string()).unwrap();
+            0
+        }
+    }
+}
+
 #[no_mangle]
 pub extern "system" fn Java_ai_djl_huggingface_tokenizers_jni_TokenizersLibrary_encodeList(
     env: JNIEnv,
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 77bc17dba..fac3c1f39 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
@@ -35,9 +35,12 @@ public final class HuggingFaceTokenizer extends NativeResource<Long> implements
 
     private boolean addSpecialTokens;
 
-    private HuggingFaceTokenizer(long handle, boolean addSpecialTokens) {
+    private HuggingFaceTokenizer(long handle, Map<String, String> options) {
         super(handle);
-        this.addSpecialTokens = addSpecialTokens;
+        this.addSpecialTokens =
+                options == null
+                        || !options.containsKey("addSpecialTokens")
+                        || Boolean.parseBoolean(options.get("addSpecialTokens"));
     }
 
     /**
@@ -61,11 +64,7 @@ public final class HuggingFaceTokenizer extends NativeResource<Long> implements
         LibUtils.checkStatus();
 
         long handle = TokenizersLibrary.LIB.createTokenizer(identifier);
-        boolean addSpecialTokens =
-                options == null
-                        || !options.containsKey("addSpecialTokens")
-                        || Boolean.parseBoolean(options.get("addSpecialTokens"));
-        return new HuggingFaceTokenizer(handle, addSpecialTokens);
+        return new HuggingFaceTokenizer(handle, options);
     }
 
     /**
@@ -89,6 +88,9 @@ public final class HuggingFaceTokenizer extends NativeResource<Long> implements
      */
     public static HuggingFaceTokenizer newInstance(Path modelPath, Map<String, String> options)
             throws IOException {
+        if (Files.isDirectory(modelPath)) {
+            modelPath = modelPath.resolve("tokenizer.json");
+        }
         try (InputStream is = Files.newInputStream(modelPath)) {
             return newInstance(is, options);
         }
@@ -108,9 +110,7 @@ public final class HuggingFaceTokenizer extends NativeResource<Long> implements
         String json = Utils.toString(is);
 
         long handle = TokenizersLibrary.LIB.createTokenizerFromString(json);
-        boolean addSpecialTokens =
-                options != null && Boolean.parseBoolean(options.get("addSpecialTokens"));
-        return new HuggingFaceTokenizer(handle, addSpecialTokens);
+        return new HuggingFaceTokenizer(handle, options);
     }
 
     /** {@inheritDoc} */
@@ -124,7 +124,7 @@ public final class HuggingFaceTokenizer extends NativeResource<Long> implements
     @Override
     public String buildSentence(List<String> tokens) {
         // TODO:
-        return String.join(" ", tokens);
+        return String.join(" ", tokens).replace(" ##", "").trim();
     }
 
     /** {@inheritDoc} */
@@ -139,11 +139,24 @@ public final class HuggingFaceTokenizer extends NativeResource<Long> implements
     /**
      * Returns the {@code Encoding} of the input sentence.
      *
-     * @param input the input sentence
+     * @param text the input sentence
+     * @return the {@code Encoding} of the input sentence
+     */
+    public Encoding encode(String text) {
+        long encoding = TokenizersLibrary.LIB.encode(getHandle(), text, addSpecialTokens);
+        return toEncoding(encoding);
+    }
+
+    /**
+     * Returns the {@code Encoding} of the input sentence.
+     *
+     * @param text the input sentence
+     * @param textPair the second input sentence
      * @return the {@code Encoding} of the input sentence
      */
-    public Encoding encode(String input) {
-        long encoding = TokenizersLibrary.LIB.encode(getHandle(), input, addSpecialTokens);
+    public Encoding encode(String text, String textPair) {
+        long encoding =
+                TokenizersLibrary.LIB.encodeDual(getHandle(), text, textPair, addSpecialTokens);
         return toEncoding(encoding);
     }
 
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 7ca5bd284..b0471385e 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
@@ -26,7 +26,10 @@ public final class TokenizersLibrary {
 
     public native void deleteTokenizer(long handle);
 
-    public native long encode(long tokenizer, String input, boolean addSpecialTokens);
+    public native long encode(long tokenizer, String text, boolean addSpecialTokens);
+
+    public native long encodeDual(
+            long tokenizer, String text, String textPair, boolean addSpecialTokens);
 
     public native long encodeList(long tokenizer, String[] inputs, boolean addSpecialTokens);
 


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 909589c08..566bf6186 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
@@ -15,10 +15,14 @@ package ai.djl.huggingface.tokenizers;
 
 import ai.djl.huggingface.tokenizers.jni.CharSpan;
 import ai.djl.testing.TestRequirements;
+import ai.djl.training.util.DownloadUtils;
 
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
+import java.io.IOException;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
@@ -27,7 +31,7 @@ import java.util.concurrent.ConcurrentHashMap;
 public class HuggingFaceTokenizerTest {
 
     @Test
-    public void testTokenizer() {
+    public void testTokenizer() throws IOException {
         TestRequirements.notArm();
 
         String input = "Hello, y'all! How are you 😁 ?";
@@ -74,9 +78,8 @@ public class HuggingFaceTokenizerTest {
             CharSpan[] charSpansResult = encoding.getCharTokenSpans();
 
             Assert.assertEquals(expectedLength, charSpansResult.length);
-            Assert.assertEquals(charSpansExpected[0], charSpansResult[0]);
-            Assert.assertEquals(
-                    charSpansExpected[expectedLength - 1], charSpansResult[expectedLength - 1]);
+            Assert.assertNull(charSpansResult[0]);
+            Assert.assertNull(charSpansResult[expectedLength - 1]);
 
             for (int i = 1; i < expectedLength - 1; i++) {
                 Assert.assertEquals(charSpansExpected[i].getStart(), charSpansResult[i].getStart());
@@ -96,5 +99,15 @@ public class HuggingFaceTokenizerTest {
             Encoding encoding = tokenizer.encode(Arrays.asList(inputs));
             Assert.assertEquals(encoding.getTokens(), exp);
         }
+
+        DownloadUtils.download(
+                "https://huggingface.co/bert-base-cased/raw/main/tokenizer.json",
+                "build/tokenizer/tokenizer.json");
+        Path path = Paths.get("build/tokenizer");
+        try (HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.newInstance(path)) {
+            long[] ids = {101, 8667, 102, 1731, 1132, 1128, 102};
+            Encoding encoding = tokenizer.encode("Hello", "How are you");
+            Assert.assertEquals(encoding.getIds(), ids);
+        }
     }
 }