Repository: deepjavalibrary/djl
Pull request #2154 (2022-11-12 01:09:39)

[tokenizers] Add batch decoding methods for tokenizers
## Description ##

Add batch decoding functionality from tokenizers


codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2154?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Base: **72.08**% // Head: **71.46**% // Decreases project coverage by **`-0.62%`** :warning:
> Coverage data is based on head [(`31bfdf1`)](https://codecov.io/gh/deepjavalibrary/djl/pull/2154?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) compared to base [(`bb5073f`)](https://codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> Patch coverage: 72.03% of modified lines in pull request are covered.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #2154      +/-   ##
============================================
- Coverage     72.08%   71.46%   -0.63%     
- Complexity     5126     6333    +1207     
============================================
  Files           473      626     +153     
  Lines         21970    27962    +5992     
  Branches       2351     3014     +663     
============================================
+ Hits          15838    19983    +4145     
- Misses         4925     6511    +1586     
- Partials       1207     1468     +261     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/2154?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [api/src/main/java/ai/djl/modality/cv/Image.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2154/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0ltYWdlLmphdmE=) | `69.23% <ø> (-4.11%)` | :arrow_down: |
| [...rc/main/java/ai/djl/modality/cv/MultiBoxPrior.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2154/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L011bHRpQm94UHJpb3IuamF2YQ==) | `76.00% <ø> (ø)` | |
| [...rc/main/java/ai/djl/modality/cv/output/Joints.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2154/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9Kb2ludHMuamF2YQ==) | `71.42% <ø> (ø)` | |
| [.../main/java/ai/djl/modality/cv/output/Landmark.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2154/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9MYW5kbWFyay5qYXZh) | `100.00% <ø> (ø)` | |
| [...main/java/ai/djl/modality/cv/output/Rectangle.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2154/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9SZWN0YW5nbGUuamF2YQ==) | `72.41% <0.00%> (ø)` | |
| [...i/djl/modality/cv/translator/BigGANTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2154/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvci5qYXZh) | `21.42% <0.00%> (-5.24%)` | :arrow_down: |
| [.../modality/cv/translator/ImageFeatureExtractor.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2154/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW1hZ2VGZWF0dXJlRXh0cmFjdG9yLmphdmE=) | `0.00% <0.00%> (ø)` | |
| [.../ai/djl/modality/cv/translator/YoloTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2154/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1RyYW5zbGF0b3IuamF2YQ==) | `27.77% <0.00%> (+18.95%)` | :arrow_up: |
| [...modality/cv/translator/wrapper/FileTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2154/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3Ivd3JhcHBlci9GaWxlVHJhbnNsYXRvci5qYXZh) | `44.44% <ø> (ø)` | |
| [...y/cv/translator/wrapper/InputStreamTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2154/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3Ivd3JhcHBlci9JbnB1dFN0cmVhbVRyYW5zbGF0b3IuamF2YQ==) | `44.44% <ø> (ø)` | |
| ... and [561 more](https://codecov.io/gh/deepjavalibrary/djl/pull/2154/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2154?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/extensions/tokenizers/rust/src/lib.rs b/extensions/tokenizers/rust/src/lib.rs
index 92038b222..0182327b8 100644
--- a/extensions/tokenizers/rust/src/lib.rs
+++ b/extensions/tokenizers/rust/src/lib.rs
@@ -428,6 +428,47 @@ pub extern "system" fn Java_ai_djl_huggingface_tokenizers_jni_TokenizersLibrary_
     ret
 }
 
+#[no_mangle]
+pub extern "system" fn Java_ai_djl_huggingface_tokenizers_jni_TokenizersLibrary_batchDecode(
+    env: JNIEnv,
+    _: JObject,
+    handle: jlong,
+    batch_ids: jobjectArray,
+    skip_special_tokens: jboolean,
+) -> jobjectArray {
+    let tokenizer = cast_handle::<Tokenizer>(handle);
+    let batch_len = env.get_array_length(batch_ids).unwrap();
+    let mut batch_decode_input: Vec<Vec<u32>> = Vec::new();
+    for i in 0..batch_len {
+        let item = env.get_object_array_element(batch_ids, i).unwrap();
+        let sequence_ids = env
+            .get_long_array_elements(*item, ReleaseMode::NoCopyBack)
+            .unwrap();
+        let sequence_ids_ptr = sequence_ids.as_ptr();
+        let sequence_len = sequence_ids.size().unwrap() as usize;
+        let mut decode_ids: Vec<u32> = Vec::new();
+        for i in 0..sequence_len {
+            unsafe {
+                let val = sequence_ids_ptr.add(i);
+                decode_ids.push(*val as u32);
+            }
+        }
+        batch_decode_input.push(decode_ids);
+    }
+    let decoding: Vec<String> = tokenizer
+        .decode_batch(batch_decode_input, skip_special_tokens == JNI_TRUE)
+        .unwrap();
+    let ret: jobjectArray = env
+        .new_object_array(batch_len, "java/lang/String", JObject::null())
+        .unwrap();
+    for (i, decode) in decoding.iter().enumerate() {
+        let item: JString = env.new_string(&decode).unwrap();
+        env.set_object_array_element(ret, i as jsize, item)
+            .unwrap();
+    }
+    ret
+}
+
 #[no_mangle]
 pub extern "system" fn Java_ai_djl_huggingface_tokenizers_jni_TokenizersLibrary_getTruncationStrategy(
     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 fd05caea6..8881292b8 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
@@ -339,6 +339,27 @@ public final class HuggingFaceTokenizer extends NativeResource<Long> implements
         return decode(ids, !addSpecialTokens);
     }
 
+    /**
+     * Returns the decoded Strings from the input batch ids.
+     *
+     * @param batchIds the batch of id sequences to decode
+     * @param skipSpecialTokens whether to remove special tokens in the decoding
+     * @return the decoded Strings from the input batch ids
+     */
+    public String[] batchDecode(long[][] batchIds, boolean skipSpecialTokens) {
+        return TokenizersLibrary.LIB.batchDecode(getHandle(), batchIds, skipSpecialTokens);
+    }
+
+    /**
+     * Returns the decoded Strings from the input batch ids.
+     *
+     * @param batchIds the batch of id sequences to decode
+     * @return the decoded Strings from the input batch ids
+     */
+    public String[] batchDecode(long[][] batchIds) {
+        return batchDecode(batchIds, !addSpecialTokens);
+    }
+
     /**
      * Creates a builder to build a {@code HuggingFaceTokenizer}.
      *
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 0e775d66c..b37947d78 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
@@ -35,6 +35,8 @@ public final class TokenizersLibrary {
 
     public native long[] batchEncode(long tokenizer, String[] inputs, boolean addSpecialTokens);
 
+    public native String[] batchDecode(long tokenizer, long[][] batchIds, boolean addSpecialTokens);
+
     public native void deleteEncoding(long encoding);
 
     public native long[] getTokenIds(long encoding);


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 2413e24e1..033e4afa4 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
@@ -414,4 +414,48 @@ public class HuggingFaceTokenizerTest {
             }
         }
     }
+
+    @Test
+    public void testBatchProcessing() throws IOException {
+        String[] inputs =
+                new String[] {
+                    "Hello there friend", "How are you today", "Good weather I'd say", "I am Happy!"
+                };
+        String[] outputsWithSpecialTokens =
+                new String[] {
+                    "[CLS] Hello there friend [SEP]",
+                    "[CLS] How are you today [SEP]",
+                    "[CLS] Good weather I ' d say [SEP]",
+                    "[CLS] I am Happy! [SEP]"
+                };
+        String[] outputsWithoutSpecialTokens =
+                new String[] {
+                    "Hello there friend",
+                    "How are you today",
+                    "Good weather I ' d say",
+                    "I am Happy!"
+                };
+        try (HuggingFaceTokenizer tokenizer =
+                HuggingFaceTokenizer.builder().optTokenizerName("bert-base-cased").build()) {
+
+            // default tokenizer with special tokens included
+            Encoding[] encodings = tokenizer.batchEncode(inputs);
+            long[][] batchIds =
+                    Arrays.stream(encodings).map(Encoding::getIds).toArray(long[][]::new);
+            String[] outputs = tokenizer.batchDecode(batchIds);
+            Assert.assertEquals(outputs, outputsWithSpecialTokens);
+
+            // encode with special tokens, decode with special tokens
+            encodings = tokenizer.batchEncode(inputs, true);
+            batchIds = Arrays.stream(encodings).map(Encoding::getIds).toArray(long[][]::new);
+            outputs = tokenizer.batchDecode(batchIds, false);
+            Assert.assertEquals(outputs, outputsWithSpecialTokens);
+
+            // encode without special tokens, decode without special tokens
+            encodings = tokenizer.batchEncode(inputs, false);
+            batchIds = Arrays.stream(encodings).map(Encoding::getIds).toArray(long[][]::new);
+            outputs = tokenizer.batchDecode(batchIds, true);
+            Assert.assertEquals(outputs, outputsWithoutSpecialTokens);
+        }
+    }
 }