Repository: deepjavalibrary/djl
Pull request #2272 (2022-12-30 17:51:26)

[tokenizer] Adds mean_sqrt_len and weightedmean pooling for TextEmbedding
Fixes: #2237

## 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


codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2272?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.55**% // Decreases project coverage by **`-0.53%`** :warning:
> Coverage data is based on head [(`abac407`)](https://codecov.io/gh/deepjavalibrary/djl/pull/2272?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: 73.77% of modified lines in pull request are covered.

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


```diff
@@             Coverage Diff              @@
##             master    #2272      +/-   ##
============================================
- Coverage     72.08%   71.55%   -0.54%     
- Complexity     5126     6365    +1239     
============================================
  Files           473      631     +158     
  Lines         21970    28164    +6194     
  Branches       2351     2997     +646     
============================================
+ Hits          15838    20152    +4314     
- Misses         4925     6547    +1622     
- Partials       1207     1465     +258     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/2272?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/2272/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/2272/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/2272/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/2272/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/2272/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/2272/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/2272/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/2272/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: |
| [...ain/java/ai/djl/modality/cv/util/NDImageUtils.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2272/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3V0aWwvTkRJbWFnZVV0aWxzLmphdmE=) | `59.21% <0.00%> (ø)` | |
| [api/src/main/java/ai/djl/modality/nlp/Decoder.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2272/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9EZWNvZGVyLmphdmE=) | `63.63% <ø> (ø)` | |
| ... and [564 more](https://codecov.io/gh/deepjavalibrary/djl/pull/2272/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/2272?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/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
index 8940b4ba9..b611a90d0 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
@@ -77,11 +77,17 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
         NDArray inputAttentionMask = manager.create(attentionMask).toType(DataType.FLOAT32, true);
         switch (pooling) {
             case "mean_tokens":
-                embeddings = meanPool(embeddings, inputAttentionMask);
+                embeddings = meanPool(embeddings, inputAttentionMask, false);
+                break;
+            case "mean_sqrt_len_tokens":
+                embeddings = meanPool(embeddings, inputAttentionMask, true);
                 break;
             case "max_tokens":
                 embeddings = maxPool(embeddings, inputAttentionMask);
                 break;
+            case "weightedmean_tokens":
+                embeddings = weightedMeanPool(embeddings, inputAttentionMask);
+                break;
             case "cls_token":
                 embeddings = embeddings.get(0);
                 break;
@@ -95,13 +101,16 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
         return embeddings.toFloatArray();
     }
 
-    private NDArray meanPool(NDArray embeddings, NDArray attentionMask) {
+    private NDArray meanPool(NDArray embeddings, NDArray attentionMask, boolean sqrt) {
         long[] shape = embeddings.getShape().getShape();
         attentionMask = attentionMask.expandDims(-1).broadcast(shape);
         NDArray inputAttentionMaskSum = attentionMask.sum(AXIS);
         NDArray clamp = inputAttentionMaskSum.clip(1e-9, 1e12);
         NDArray prod = embeddings.mul(attentionMask);
         NDArray sum = prod.sum(AXIS);
+        if (sqrt) {
+            return sum.div(clamp.sqrt());
+        }
         return sum.div(clamp);
     }
 
@@ -115,6 +124,17 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
         return embeddings.max(AXIS, true);
     }
 
+    private NDArray weightedMeanPool(NDArray embeddings, NDArray attentionMask) {
+        long[] shape = embeddings.getShape().getShape();
+        NDArray weight = embeddings.getManager().arange(1, shape[0] + 1);
+        weight = weight.expandDims(-1).broadcast(shape);
+
+        attentionMask = attentionMask.expandDims(-1).broadcast(shape).mul(weight);
+        NDArray maskSum = attentionMask.sum(AXIS);
+        NDArray embeddingSum = embeddings.mul(attentionMask).sum(AXIS);
+        return embeddingSum.div(maskSum);
+    }
+
     /**
      * Creates a builder to build a {@code TextEmbeddingTranslator}.
      *
@@ -180,18 +200,14 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
          * @return this builder
          */
         public Builder optPoolingMode(String poolingMode) {
-            if ("weightedmean_tokens".equals(poolingMode)
-                    || "lasttoken".equals(poolingMode)
-                    || "mean_sqrt_len_tokens".equals(poolingMode)) {
-                throw new UnsupportedOperationException(
-                        "Unsupported pooling model: " + poolingMode);
-            }
             if (!"mean_tokens".equals(poolingMode)
                     && !"max_tokens".equals(poolingMode)
-                    && !"cls_token".equals(poolingMode)) {
+                    && !"cls_token".equals(poolingMode)
+                    && !"mean_sqrt_len_tokens".equals(poolingMode)
+                    && !"weightedmean_tokens".equals(poolingMode)) {
                 throw new IllegalArgumentException(
                         "Invalid pooling model, must be one of [mean_tokens, max_tokens,"
-                                + " cls_token].");
+                                + " cls_token, mean_sqrt_len_tokens, weightedmean_tokens].");
             }
             this.pooling = poolingMode;
             return this;


Test patch:
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TranslatorTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TranslatorTest.java
index d012157b8..e8abc5d03 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TranslatorTest.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TranslatorTest.java
@@ -357,7 +357,7 @@ public class TranslatorTest {
             Assertions.assertAlmostEquals(res[0], 0.05103);
         }
 
-        // pooling_mode_max_token
+        // pooling_mode_max_tokens
         criteria =
                 Criteria.builder()
                         .setTypes(String.class, float[].class)
@@ -377,6 +377,46 @@ public class TranslatorTest {
             Assertions.assertAlmostEquals(res[0], 1.0);
         }
 
+        // pooling_mean_sqrt_len_tokens
+        criteria =
+                Criteria.builder()
+                        .setTypes(String.class, float[].class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-uncased")
+                        .optArgument("pooling", "mean_sqrt_len_tokens")
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
+                        .build();
+
+        try (ZooModel<String, float[]> model = criteria.loadModel();
+                Predictor<String, float[]> predictor = model.newPredictor()) {
+            float[] res = predictor.predict(text);
+            Assert.assertEquals(res.length, 384);
+            Assertions.assertAlmostEquals(res[0], 0.05103104);
+        }
+
+        // pooling_weightedmean_tokens
+        criteria =
+                Criteria.builder()
+                        .setTypes(String.class, float[].class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-uncased")
+                        .optArgument("pooling", "weightedmean_tokens")
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
+                        .build();
+
+        try (ZooModel<String, float[]> model = criteria.loadModel();
+                Predictor<String, float[]> predictor = model.newPredictor()) {
+            float[] res = predictor.predict(text);
+            Assert.assertEquals(res.length, 384);
+            Assertions.assertAlmostEquals(res[0], 0.05103104);
+        }
+
         Criteria<Input, Output> criteria2 =
                 Criteria.builder()
                         .setTypes(Input.class, Output.class)