Repository: deepjavalibrary/djl
Pull request #1584 (2022-04-19 09:37:08)

HF Tokenizer: get charspans 
## Description ##

This PR extends the hugging face Java/Rust interface. I added a method to return the original char spans for each token result, together with the encoding object. 

This function is useful in Named Entity Recognition tasks (NER), specifically to reconstruct entities together, since the tokens retrieved from the tokenization lose whitespace information.


Patch:
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java
index 0deff350d..105ec60af 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java
@@ -15,6 +15,7 @@ package ai.djl.pytorch.jni;
 import ai.djl.util.ClassLoaderUtils;
 import ai.djl.util.Platform;
 import ai.djl.util.Utils;
+import ai.djl.util.cuda.CudaUtils;
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
@@ -102,6 +103,7 @@ public final class LibUtils {
             // PyTorch 1.8.1 libtorch_cpu.dylib cannot be loaded individually
             return;
         }
+        boolean isCuda = libTorch.flavor.contains("cu");
         List<String> deferred =
                 Arrays.asList(
                         System.mapLibraryName("fbgemm"),
@@ -118,6 +120,12 @@ public final class LibUtils {
             paths.filter(
                             path -> {
                                 String name = path.getFileName().toString();
+                                if (!isCuda
+                                        && name.contains("nvrtc")
+                                        && name.contains("cudart")
+                                        && name.contains("nvTools")) {
+                                    return false;
+                                }
                                 return !loadLater.contains(name)
                                         && Files.isRegularFile(path)
                                         && !name.endsWith(JNI_LIB_NAME)
@@ -139,6 +147,14 @@ public final class LibUtils {
                 loadNativeLibrary(libDir.resolve("cudnn64_7.dll").toString());
             }
 
+            if (!isCuda) {
+                deferred =
+                        Arrays.asList(
+                                System.mapLibraryName("fbgemm"),
+                                System.mapLibraryName("torch_cpu"),
+                                System.mapLibraryName("torch"));
+            }
+
             for (String dep : deferred) {
                 Path path = libDir.resolve(dep);
                 if (Files.exists(path)) {
@@ -289,6 +305,7 @@ public final class LibUtils {
             if (Files.exists(path)) {
                 return new LibTorch(dir.toAbsolutePath(), platform, flavor);
             }
+            Utils.deleteQuietly(dir);
 
             Matcher m = VERSION_PATTERN.matcher(version);
             if (!m.matches()) {
@@ -498,7 +515,11 @@ public final class LibUtils {
             if (flavor == null) {
                 flavor = System.getProperty("PYTORCH_FLAVOR");
                 if (flavor == null) {
-                    flavor = "cpu-precxx11";
+                    if (CudaUtils.getGpuCount() > 0) {
+                        flavor = "cu" + CudaUtils.getCudaVersionString() + "-precxx11";
+                    } else {
+                        flavor = "cpu-precxx11";
+                    }
                 }
             }
         }
diff --git a/extensions/tokenizers/rust/src/lib.rs b/extensions/tokenizers/rust/src/lib.rs
index 34e532de8..2a826c3f5 100644
--- a/extensions/tokenizers/rust/src/lib.rs
+++ b/extensions/tokenizers/rust/src/lib.rs
@@ -16,11 +16,13 @@ extern crate tokenizers as tk;
 
 use std::str::FromStr;
 use tk::tokenizer::{EncodeInput, Encoding};
-use tk::FromPretrainedParameters;
 use tk::Tokenizer;
+use tk::{FromPretrainedParameters, Offsets};
 
-use jni::objects::{JObject, JString};
-use jni::sys::{jboolean, jlong, jlongArray, jobjectArray, jsize, JNI_TRUE};
+use jni::objects::{JClass, JMethodID, JObject, JString, JValue};
+use jni::sys::{
+    jboolean, jint, jlong, jlongArray, jobjectArray, jsize, JNI_TRUE,
+};
 use jni::JNIEnv;
 
 #[no_mangle]
@@ -256,7 +258,8 @@ pub extern "system" fn Java_ai_djl_huggingface_tokenizers_jni_TokenizersLibrary_
         .unwrap();
     for (i, token) in tokens.iter().enumerate() {
         let item: JString = env.new_string(&token).unwrap();
-        env.set_object_array_element(array, i as jsize, item).unwrap();
+        env.set_object_array_element(array, i as jsize, item)
+            .unwrap();
     }
     array
 }
@@ -299,6 +302,48 @@ 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_getTokenCharSpans(
+    env: JNIEnv,
+    _: JObject,
+    handle: jlong,
+) -> jobjectArray {
+    let encoding = cast_handle::<Encoding>(handle);
+    let tokens = encoding.get_tokens();
+    let len = tokens.len() as jsize;
+
+    let array: jobjectArray = env
+        .new_object_array(
+            len,
+            "ai/djl/huggingface/tokenizers/jni/CharSpan",
+            JObject::null(),
+        )
+        .unwrap();
+    for (i, _) in tokens.iter().enumerate() {
+        let opt_offsets: Option<(usize, Offsets)> = encoding.token_to_chars(i);
+        match &opt_offsets {
+            Some((_, offsets)) => {
+                let class_id = "ai/djl/huggingface/tokenizers/jni/CharSpan";
+                let method_id = "<init>";
+                let params = "(II)V";
+                let cls: JClass = env.find_class(class_id).unwrap();
+                let constructor: JMethodID = env.get_method_id(cls, method_id, params).unwrap();
+                let offsets_vec: Vec<JValue> = vec![
+                    JValue::Int((*offsets).0 as jint),
+                    JValue::Int((*offsets).1 as jint),
+                ];
+                let obj = env
+                    .new_object_unchecked(cls, constructor, &offsets_vec[..])
+                    .unwrap();
+                env.set_object_array_element(array, i as jsize, obj)
+                    .unwrap();
+            }
+            None => {}
+        }
+    }
+    array
+}
+
 fn to_handle<T: 'static>(val: T) -> jlong {
     let handle = Box::into_raw(Box::new(val)) as jlong;
     handle
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 9d894add6..dfbee95de 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
@@ -12,6 +12,8 @@
  */
 package ai.djl.huggingface.tokenizers;
 
+import ai.djl.huggingface.tokenizers.jni.CharSpan;
+
 /** A class holds token encoding information. */
 public class Encoding {
 
@@ -21,6 +23,7 @@ public class Encoding {
     private long[] wordIds;
     private long[] attentionMask;
     private long[] specialTokenMask;
+    private CharSpan[] charTokenSpans;
 
     Encoding(
             long[] ids,
@@ -28,13 +31,15 @@ public class Encoding {
             String[] tokens,
             long[] wordIds,
             long[] attentionMask,
-            long[] specialTokenMask) {
+            long[] specialTokenMask,
+            CharSpan[] charTokenSpans) {
         this.ids = ids;
         this.typeIds = typeIds;
         this.tokens = tokens;
         this.wordIds = wordIds;
         this.attentionMask = attentionMask;
         this.specialTokenMask = specialTokenMask;
+        this.charTokenSpans = charTokenSpans;
     }
 
     /**
@@ -90,4 +95,13 @@ public class Encoding {
     public long[] getSpecialTokenMask() {
         return specialTokenMask;
     }
+
+    /**
+     * Returns char token spans.
+     *
+     * @return char token spans
+     */
+    public CharSpan[] getCharTokenSpans() {
+        return charTokenSpans;
+    }
 }
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 3bf130d50..1d540faaf 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
@@ -12,6 +12,7 @@
  */
 package ai.djl.huggingface.tokenizers;
 
+import ai.djl.huggingface.tokenizers.jni.CharSpan;
 import ai.djl.huggingface.tokenizers.jni.LibUtils;
 import ai.djl.huggingface.tokenizers.jni.TokenizersLibrary;
 import ai.djl.modality.nlp.preprocess.Tokenizer;
@@ -200,8 +201,10 @@ public final class HuggingFaceTokenizer extends NativeResource<Long> implements
         long[] wordIds = TokenizersLibrary.LIB.getWordIds(encoding);
         long[] attentionMask = TokenizersLibrary.LIB.getAttentionMask(encoding);
         long[] specialTokenMask = TokenizersLibrary.LIB.getSpecialTokenMask(encoding);
+        CharSpan[] charSpans = TokenizersLibrary.LIB.getTokenCharSpans(encoding);
 
         TokenizersLibrary.LIB.deleteEncoding(encoding);
-        return new Encoding(ids, typeIds, tokens, wordIds, attentionMask, specialTokenMask);
+        return new Encoding(
+                ids, typeIds, tokens, wordIds, 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
new file mode 100644
index 000000000..7478d0e8b
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/jni/CharSpan.java
@@ -0,0 +1,51 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * 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. */
+public class CharSpan {
+
+    private final int start;
+    private final int end;
+
+    /**
+     * Constructs a new {@code CharSpan} instance.
+     *
+     * @param start the start position
+     * @param end the end position
+     */
+    public CharSpan(int start, int end) {
+        this.start = start;
+        this.end = end;
+    }
+
+    /**
+     * Returns the start position.
+     *
+     * @return the start position
+     */
+    public int getStart() {
+        return start;
+    }
+
+    /**
+     * Returns the end position.
+     *
+     * @return the end position
+     */
+    public int getEnd() {
+        return end;
+    }
+}
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 cb8b68308..7ca5bd284 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
@@ -45,4 +45,6 @@ public final class TokenizersLibrary {
     public native long[] getAttentionMask(long encoding);
 
     public native long[] getSpecialTokenMask(long encoding);
+
+    public native CharSpan[] getTokenCharSpans(long encoding);
 }


Test patch:
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java b/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java
index 5c29059be..d4ee92bb6 100644
--- a/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java
+++ b/examples/src/test/java/ai/djl/examples/training/TrainMnistTest.java
@@ -50,7 +50,9 @@ public class TrainMnistTest {
 
         Classifications classifications = ImageClassification.predict();
         Classifications.Classification best = classifications.best();
-        Assert.assertEquals(best.getClassName(), "0");
+        if (Boolean.getBoolean("nightly")) {
+            Assert.assertEquals(best.getClassName(), "0");
+        }
         double probability = best.getProbability();
         Assert.assertTrue(probability > expectedProb && probability <= 1);
     }
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 4f559e430..b4cd211bb 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
@@ -13,6 +13,7 @@
 
 package ai.djl.huggingface.tokenizers;
 
+import ai.djl.huggingface.tokenizers.jni.CharSpan;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
@@ -49,6 +50,34 @@ public class HuggingFaceTokenizerTest {
             Assert.assertEquals(attentionMask, encoding.getAttentionMask());
             Assert.assertEquals(specialTokenMask, encoding.getSpecialTokenMask());
 
+            CharSpan[] charSpansExpected = {
+                null,
+                new CharSpan(0, 5),
+                new CharSpan(5, 6),
+                new CharSpan(7, 8),
+                new CharSpan(8, 9),
+                new CharSpan(9, 12),
+                new CharSpan(12, 13),
+                new CharSpan(14, 17),
+                new CharSpan(18, 21),
+                new CharSpan(22, 25),
+                new CharSpan(26, 30),
+                new CharSpan(31, 32),
+                null
+            };
+            int expectedLength = charSpansExpected.length;
+            CharSpan[] charSpansResult = encoding.getCharTokenSpans();
+
+            Assert.assertEquals(expectedLength, charSpansResult.length);
+            Assert.assertEquals(charSpansExpected[0], charSpansResult[0]);
+            Assert.assertEquals(
+                    charSpansExpected[expectedLength - 1], charSpansResult[expectedLength - 1]);
+
+            for (int i = 1; i < expectedLength - 1; i++) {
+                Assert.assertEquals(charSpansExpected[i].getStart(), charSpansResult[i].getStart());
+                Assert.assertEquals(charSpansExpected[i].getEnd(), charSpansResult[i].getEnd());
+            }
+
             Encoding[] encodings = tokenizer.batchEncode(Arrays.asList(inputs));
             Assert.assertEquals(encodings.length, 2);
         }