Repository: deepjavalibrary/djl
Pull request #2314 (2023-01-10 06:11:58)

[tokenizer] Update TextEmbedding pooling model name
## 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/2314?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: **74.22**% // Increases project coverage by **`+2.13%`** :tada:
> Coverage data is based on head [(`ff986da`)](https://codecov.io/gh/deepjavalibrary/djl/pull/2314?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: 74.96% of modified lines in pull request are covered.

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


```diff
@@             Coverage Diff              @@
##             master    #2314      +/-   ##
============================================
+ Coverage     72.08%   74.22%   +2.13%     
- Complexity     5126     6674    +1548     
============================================
  Files           473      656     +183     
  Lines         21970    29045    +7075     
  Branches       2351     3002     +651     
============================================
+ Hits          15838    21559    +5721     
- Misses         4925     6022    +1097     
- Partials       1207     1464     +257     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/2314?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/2314?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/2314?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/2314?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/2314?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/2314?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/2314?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/2314?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/2314?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/2314?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/2314?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 [598 more](https://codecov.io/gh/deepjavalibrary/djl/pull/2314?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/2314?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 544a13d70..4047927d7 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
@@ -92,18 +92,18 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
         long[] attentionMask = encoding.getAttentionMask();
         NDArray inputAttentionMask = manager.create(attentionMask).toType(DataType.FLOAT32, true);
         switch (pooling) {
-            case "mean_tokens":
+            case "mean":
                 return meanPool(embeddings, inputAttentionMask, false);
-            case "mean_sqrt_len_tokens":
+            case "mean_sqrt_len":
                 return meanPool(embeddings, inputAttentionMask, true);
-            case "max_tokens":
+            case "max":
                 return maxPool(embeddings, inputAttentionMask);
-            case "weightedmean_tokens":
+            case "weightedmean":
                 return weightedMeanPool(embeddings, inputAttentionMask);
-            case "cls_token":
+            case "cls":
                 return embeddings.get(0);
             default:
-                throw new AssertionError("Unexpected pooling model: " + pooling);
+                throw new AssertionError("Unexpected pooling mode: " + pooling);
         }
     }
 
@@ -171,7 +171,7 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
         private HuggingFaceTokenizer tokenizer;
         private Batchifier batchifier = Batchifier.STACK;
         private boolean normalize = true;
-        private String pooling = "mean_tokens";
+        private String pooling = "mean";
 
         Builder(HuggingFaceTokenizer tokenizer) {
             this.tokenizer = tokenizer;
@@ -189,7 +189,7 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
         }
 
         /**
-         * Sets the normalize for the {@link Translator}.
+         * Sets the {@code normalize} for the {@link Translator}.
          *
          * @param normalize true to normalize the embeddings
          * @return this builder
@@ -206,14 +206,14 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
          * @return this builder
          */
         public Builder optPoolingMode(String poolingMode) {
-            if (!"mean_tokens".equals(poolingMode)
-                    && !"max_tokens".equals(poolingMode)
-                    && !"cls_token".equals(poolingMode)
-                    && !"mean_sqrt_len_tokens".equals(poolingMode)
-                    && !"weightedmean_tokens".equals(poolingMode)) {
+            if (!"mean".equals(poolingMode)
+                    && !"max".equals(poolingMode)
+                    && !"cls".equals(poolingMode)
+                    && !"mean_sqrt_len".equals(poolingMode)
+                    && !"weightedmean".equals(poolingMode)) {
                 throw new IllegalArgumentException(
-                        "Invalid pooling model, must be one of [mean_tokens, max_tokens,"
-                                + " cls_token, mean_sqrt_len_tokens, weightedmean_tokens].");
+                        "Invalid pooling model, must be one of [mean, max, cls, mean_sqrt_len,"
+                                + " weightedmean].");
             }
             this.pooling = poolingMode;
             return this;
@@ -228,7 +228,7 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
             String batchifierStr = ArgumentsUtil.stringValue(arguments, "batchifier", "stack");
             optBatchifier(Batchifier.fromString(batchifierStr));
             optNormalize(ArgumentsUtil.booleanValue(arguments, "normalize", true));
-            optPoolingMode(ArgumentsUtil.stringValue(arguments, "pooling", "mean_tokens"));
+            optPoolingMode(ArgumentsUtil.stringValue(arguments, "pooling", "mean"));
         }
 
         /**
diff --git a/extensions/tokenizers/src/main/python/sentence_similarity_converter.py b/extensions/tokenizers/src/main/python/sentence_similarity_converter.py
index ec9d446e4..c96975ab5 100644
--- a/extensions/tokenizers/src/main/python/sentence_similarity_converter.py
+++ b/extensions/tokenizers/src/main/python/sentence_similarity_converter.py
@@ -66,13 +66,13 @@ class SentenceSimilarityConverter(HuggingfaceConverter):
                 with open(file, "r") as f:
                     pooling = json.load(f)
                     if pooling.get("pooling_mode_cls_token"):
-                        args["pooling"] = "cls_token"
+                        args["pooling"] = "cls"
                     elif pooling.get("pooling_mode_max_tokens"):
-                        args["pooling"] = "max_tokens"
+                        args["pooling"] = "max"
                     elif pooling.get("pooling_mode_mean_sqrt_len_tokens"):
-                        args["pooling"] = "mean_sqrt_len_tokens"
+                        args["pooling"] = "mean_sqrt_len"
                     elif pooling.get("pooling_mode_weightedmean_tokens"):
-                        args["pooling"] = "weightedmean_tokens"
+                        args["pooling"] = "weightedmean"
                     elif pooling.get("pooling_mode_lasttoken"):
                         args["pooling"] = "lasttoken"
         except requests.exceptions.HTTPError:


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 821afb117..84b3d3d5d 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
@@ -365,7 +365,7 @@ public class TranslatorTest {
                         .optBlock(block)
                         .optEngine("PyTorch")
                         .optArgument("tokenizer", "bert-base-uncased")
-                        .optArgument("pooling", "max_tokens")
+                        .optArgument("pooling", "max")
                         .optOption("hasParameter", "false")
                         .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
                         .build();
@@ -385,7 +385,7 @@ public class TranslatorTest {
                         .optBlock(block)
                         .optEngine("PyTorch")
                         .optArgument("tokenizer", "bert-base-uncased")
-                        .optArgument("pooling", "mean_sqrt_len_tokens")
+                        .optArgument("pooling", "mean_sqrt_len")
                         .optOption("hasParameter", "false")
                         .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
                         .build();
@@ -405,7 +405,7 @@ public class TranslatorTest {
                         .optBlock(block)
                         .optEngine("PyTorch")
                         .optArgument("tokenizer", "bert-base-uncased")
-                        .optArgument("pooling", "weightedmean_tokens")
+                        .optArgument("pooling", "weightedmean")
                         .optOption("hasParameter", "false")
                         .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
                         .build();
@@ -424,7 +424,7 @@ public class TranslatorTest {
                         .optBlock(block)
                         .optEngine("PyTorch")
                         .optArgument("tokenizer", "bert-base-uncased")
-                        .optArgument("pooling", "cls_token")
+                        .optArgument("pooling", "cls")
                         .optOption("hasParameter", "false")
                         .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
                         .build();