Repository: deepjavalibrary/djl
Pull request #2261 (2022-12-25 02:15:48)

[tokenizer] Adds more pooling mode form TextEmbedding
## 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/2261?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.56**% // Decreases project coverage by **`-0.52%`** :warning:
> Coverage data is based on head [(`025adb8`)](https://codecov.io/gh/deepjavalibrary/djl/pull/2261?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    #2261      +/-   ##
============================================
- Coverage     72.08%   71.56%   -0.53%     
- Complexity     5126     6364    +1238     
============================================
  Files           473      631     +158     
  Lines         21970    28147    +6177     
  Branches       2351     2994     +643     
============================================
+ Hits          15838    20142    +4304     
- Misses         4925     6539    +1614     
- Partials       1207     1466     +259     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/2261?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/2261/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/2261/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/2261/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/2261/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/2261/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/2261/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/2261/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/2261/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/2261/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/2261/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/2261/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/2261?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/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_creation.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_creation.cc
index ebf015014..6386f3e9e 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_creation.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_creation.cc
@@ -160,13 +160,12 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchSparseCoo(
   API_END_RETURN();
 }
 
-
 JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchHannWindow(
-        JNIEnv* env, jobject jthis, jlong jn_fft, jboolean jperoidic, jintArray jdevice) {
-    API_BEGIN()
-        const auto device = utils::GetDeviceFromJDevice(env, jdevice);
-        auto options = torch::TensorOptions().device(device);
-        const auto* result_ptr = new torch::Tensor(torch::hann_window(jn_fft, jperoidic, options));
-        return reinterpret_cast<uintptr_t>(result_ptr);
-    API_END_RETURN()
+    JNIEnv* env, jobject jthis, jlong jn_fft, jboolean jperoidic, jintArray jdevice) {
+  API_BEGIN()
+  const auto device = utils::GetDeviceFromJDevice(env, jdevice);
+  auto options = torch::TensorOptions().device(device);
+  const auto* result_ptr = new torch::Tensor(torch::hann_window(jn_fft, jperoidic, options));
+  return reinterpret_cast<uintptr_t>(result_ptr);
+  API_END_RETURN()
 }
diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_other.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_other.cc
index 8dda3cdc3..586f82143 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_other.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_other.cc
@@ -26,44 +26,45 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchFlatten(
 }
 
 JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchFft(
-        JNIEnv* env, jobject jthis, jlong jhandle, jlong jn, jlong jaxis) {
-    API_BEGIN()
-        const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
-        const auto* result_ptr = new torch::Tensor(torch::fft_fft(*tensor_ptr, jn, jaxis));
-        return reinterpret_cast<uintptr_t>(result_ptr);
-    API_END_RETURN()
+    JNIEnv* env, jobject jthis, jlong jhandle, jlong jn, jlong jaxis) {
+  API_BEGIN()
+  const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+  const auto* result_ptr = new torch::Tensor(torch::fft_fft(*tensor_ptr, jn, jaxis));
+  return reinterpret_cast<uintptr_t>(result_ptr);
+  API_END_RETURN()
 }
 
-
-JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchStft(
-        JNIEnv* env, jobject jthis, jlong jhandle, jlong jn_fft, jlong jhop_length, jlong jwindow, jboolean jcenter, jboolean jnormalize, jboolean jreturn_complex) {
-    API_BEGIN()
-        const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
-        if (jwindow == -1L) {
-            const auto* result_ptr = new torch::Tensor(tensor_ptr->stft(jn_fft, jhop_length, c10::nullopt, c10::nullopt, jcenter, "reflect", jnormalize, c10::nullopt, jreturn_complex));
-            return reinterpret_cast<uintptr_t>(result_ptr);
-        } else {
-            const auto* window_ptr = reinterpret_cast<torch::Tensor*>(jwindow);
-            const auto* result_ptr = new torch::Tensor(tensor_ptr->stft(jn_fft, jhop_length, c10::nullopt, *window_ptr, jcenter, "reflect", jnormalize, c10::nullopt, jreturn_complex));
-            return reinterpret_cast<uintptr_t>(result_ptr);
-        }
-    API_END_RETURN()
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchStft(JNIEnv* env, jobject jthis, jlong jhandle,
+    jlong jn_fft, jlong jhop_length, jlong jwindow, jboolean jcenter, jboolean jnormalize, jboolean jreturn_complex) {
+  API_BEGIN()
+  const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+  if (jwindow == -1L) {
+    const auto* result_ptr = new torch::Tensor(tensor_ptr->stft(jn_fft, jhop_length, c10::nullopt, c10::nullopt,
+        jcenter, "reflect", jnormalize, c10::nullopt, jreturn_complex));
+    return reinterpret_cast<uintptr_t>(result_ptr);
+  } else {
+    const auto* window_ptr = reinterpret_cast<torch::Tensor*>(jwindow);
+    const auto* result_ptr = new torch::Tensor(tensor_ptr->stft(
+        jn_fft, jhop_length, c10::nullopt, *window_ptr, jcenter, "reflect", jnormalize, c10::nullopt, jreturn_complex));
+    return reinterpret_cast<uintptr_t>(result_ptr);
+  }
+  API_END_RETURN()
 }
 
 JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchViewAsReal(
-        JNIEnv* env, jobject jthis, jlong jhandle) {
-    API_BEGIN()
-        const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
-        const auto* result_ptr = new torch::Tensor(torch::view_as_real_copy(*tensor_ptr));
-        return reinterpret_cast<uintptr_t>(result_ptr);
-    API_END_RETURN()
+    JNIEnv* env, jobject jthis, jlong jhandle) {
+  API_BEGIN()
+  const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+  const auto* result_ptr = new torch::Tensor(torch::view_as_real_copy(*tensor_ptr));
+  return reinterpret_cast<uintptr_t>(result_ptr);
+  API_END_RETURN()
 }
 
 JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchViewAsComplex(
-        JNIEnv* env, jobject jthis, jlong jhandle) {
-    API_BEGIN()
-        const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
-        const auto* result_ptr = new torch::Tensor(torch::view_as_complex_copy(*tensor_ptr));
-        return reinterpret_cast<uintptr_t>(result_ptr);
-    API_END_RETURN()
+    JNIEnv* env, jobject jthis, jlong jhandle) {
+  API_BEGIN()
+  const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+  const auto* result_ptr = new torch::Tensor(torch::view_as_complex_copy(*tensor_ptr));
+  return reinterpret_cast<uintptr_t>(result_ptr);
+  API_END_RETURN()
 }
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 282d8d247..8940b4ba9 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
@@ -33,10 +33,18 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
 
     private HuggingFaceTokenizer tokenizer;
     private Batchifier batchifier;
-
-    TextEmbeddingTranslator(HuggingFaceTokenizer tokenizer, Batchifier batchifier) {
+    private boolean normalize;
+    private String pooling;
+
+    TextEmbeddingTranslator(
+            HuggingFaceTokenizer tokenizer,
+            Batchifier batchifier,
+            String pooling,
+            boolean normalize) {
         this.tokenizer = tokenizer;
         this.batchifier = batchifier;
+        this.pooling = pooling;
+        this.normalize = normalize;
     }
 
     /** {@inheritDoc} */
@@ -67,15 +75,44 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
         long[] attentionMask = encoding.getAttentionMask();
         NDManager manager = ctx.getNDManager();
         NDArray inputAttentionMask = manager.create(attentionMask).toType(DataType.FLOAT32, true);
+        switch (pooling) {
+            case "mean_tokens":
+                embeddings = meanPool(embeddings, inputAttentionMask);
+                break;
+            case "max_tokens":
+                embeddings = maxPool(embeddings, inputAttentionMask);
+                break;
+            case "cls_token":
+                embeddings = embeddings.get(0);
+                break;
+            default:
+                throw new AssertionError("Unexpected pooling model: " + pooling);
+        }
+        if (normalize) {
+            embeddings = embeddings.normalize(2, 0);
+        }
+
+        return embeddings.toFloatArray();
+    }
+
+    private NDArray meanPool(NDArray embeddings, NDArray attentionMask) {
         long[] shape = embeddings.getShape().getShape();
-        inputAttentionMask = inputAttentionMask.expandDims(-1).broadcast(shape);
-        NDArray inputAttentionMaskSum = inputAttentionMask.sum(AXIS);
+        attentionMask = attentionMask.expandDims(-1).broadcast(shape);
+        NDArray inputAttentionMaskSum = attentionMask.sum(AXIS);
         NDArray clamp = inputAttentionMaskSum.clip(1e-9, 1e12);
-        NDArray prod = embeddings.mul(inputAttentionMask);
+        NDArray prod = embeddings.mul(attentionMask);
         NDArray sum = prod.sum(AXIS);
-        embeddings = sum.div(clamp).normalize(2, 0);
+        return sum.div(clamp);
+    }
 
-        return embeddings.toFloatArray();
+    private NDArray maxPool(NDArray embeddings, NDArray inputAttentionMask) {
+        long[] shape = embeddings.getShape().getShape();
+        inputAttentionMask = inputAttentionMask.expandDims(-1).broadcast(shape);
+        inputAttentionMask = inputAttentionMask.eq(0);
+        embeddings = embeddings.duplicate();
+        embeddings.set(inputAttentionMask, -1e9); // Set padding tokens to large negative value
+
+        return embeddings.max(AXIS, true);
     }
 
     /**
@@ -107,6 +144,8 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
 
         private HuggingFaceTokenizer tokenizer;
         private Batchifier batchifier = Batchifier.STACK;
+        private boolean normalize = true;
+        private String pooling = "mean_tokens";
 
         Builder(HuggingFaceTokenizer tokenizer) {
             this.tokenizer = tokenizer;
@@ -123,6 +162,41 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
             return this;
         }
 
+        /**
+         * Sets the normalize for the {@link Translator}.
+         *
+         * @param normalize true to normalize the embeddings
+         * @return this builder
+         */
+        public Builder optNormalize(boolean normalize) {
+            this.normalize = normalize;
+            return this;
+        }
+
+        /**
+         * Sets the pooling for the {@link Translator}.
+         *
+         * @param poolingMode the pooling model, one of mean_pool, max_pool and cls
+         * @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)) {
+                throw new IllegalArgumentException(
+                        "Invalid pooling model, must be one of [mean_tokens, max_tokens,"
+                                + " cls_token].");
+            }
+            this.pooling = poolingMode;
+            return this;
+        }
+
         /**
          * Configures the builder with the model arguments.
          *
@@ -131,6 +205,8 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
         public void configure(Map<String, ?> arguments) {
             String batchifierStr = ArgumentsUtil.stringValue(arguments, "batchifier", "stack");
             optBatchifier(Batchifier.fromString(batchifierStr));
+            optNormalize(ArgumentsUtil.booleanValue(arguments, "normalize", true));
+            optPoolingMode(ArgumentsUtil.stringValue(arguments, "pooling", "mean_tokens"));
         }
 
         /**
@@ -140,7 +216,7 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
          * @throws IOException if I/O error occurs
          */
         public TextEmbeddingTranslator build() throws IOException {
-            return new TextEmbeddingTranslator(tokenizer, batchifier);
+            return new TextEmbeddingTranslator(tokenizer, batchifier, pooling, normalize);
         }
     }
 }
diff --git a/extensions/tokenizers/src/main/python/fill_mask_converter.py b/extensions/tokenizers/src/main/python/fill_mask_converter.py
index 6d9e10deb..ab0b5f444 100644
--- a/extensions/tokenizers/src/main/python/fill_mask_converter.py
+++ b/extensions/tokenizers/src/main/python/fill_mask_converter.py
@@ -59,5 +59,5 @@ class FillMaskConverter(HuggingfaceConverter):
         text = self.inputs.replace("[MASK]", tokenizer.mask_token)
         return tokenizer.encode_plus(text, return_tensors='pt')
 
-    def get_extra_arguments(self, hf_pipeline) -> dict:
+    def get_extra_arguments(self, hf_pipeline, model_id: str) -> dict:
         return {"maskToken": hf_pipeline.tokenizer.mask_token}
diff --git a/extensions/tokenizers/src/main/python/huggingface_converter.py b/extensions/tokenizers/src/main/python/huggingface_converter.py
index ec1aeebd4..5b4749d28 100644
--- a/extensions/tokenizers/src/main/python/huggingface_converter.py
+++ b/extensions/tokenizers/src/main/python/huggingface_converter.py
@@ -131,7 +131,7 @@ class HuggingfaceConverter:
 
         # Save serving.properties
         serving_file = os.path.join(temp_dir, "serving.properties")
-        arguments = self.get_extra_arguments(hf_pipeline)
+        arguments = self.get_extra_arguments(hf_pipeline, model_id)
         with open(serving_file, 'w') as f:
             f.write(f"engine=PyTorch\n"
                     f"option.modelName={model_name}\n"
@@ -194,7 +194,7 @@ class HuggingfaceConverter:
 
         return self.verify_jit_output(hf_pipeline, encoding, out)
 
-    def get_extra_arguments(self, hf_pipeline) -> dict:
+    def get_extra_arguments(self, hf_pipeline, model_id: str) -> dict:
         return {}
 
     def verify_jit_output(self, hf_pipeline, encoding, out):
diff --git a/extensions/tokenizers/src/main/python/sentence_similarity_converter.py b/extensions/tokenizers/src/main/python/sentence_similarity_converter.py
index 6e372b706..ec9d446e4 100644
--- a/extensions/tokenizers/src/main/python/sentence_similarity_converter.py
+++ b/extensions/tokenizers/src/main/python/sentence_similarity_converter.py
@@ -10,12 +10,16 @@
 # or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS"
 # BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for
 # the specific language governing permissions and limitations under the License.
+import json
 import logging
+import os
 
+import requests
 import torch
 from transformers import AutoTokenizer, AutoModel
 
 from huggingface_converter import HuggingfaceConverter
+from huggingface_hub import hf_hub_download
 
 
 class PipelineHolder(object):
@@ -53,5 +57,25 @@ class SentenceSimilarityConverter(HuggingfaceConverter):
 
         return True, None
 
-    def get_extra_arguments(self, hf_pipeline) -> dict:
-        return {"padding": "true"}
+    def get_extra_arguments(self, hf_pipeline, model_id: str) -> dict:
+        args = {"padding": "true"}
+        try:
+            file = hf_hub_download(repo_id=model_id,
+                                   filename="1_Pooling/config.json")
+            if os.path.exists(file):
+                with open(file, "r") as f:
+                    pooling = json.load(f)
+                    if pooling.get("pooling_mode_cls_token"):
+                        args["pooling"] = "cls_token"
+                    elif pooling.get("pooling_mode_max_tokens"):
+                        args["pooling"] = "max_tokens"
+                    elif pooling.get("pooling_mode_mean_sqrt_len_tokens"):
+                        args["pooling"] = "mean_sqrt_len_tokens"
+                    elif pooling.get("pooling_mode_weightedmean_tokens"):
+                        args["pooling"] = "weightedmean_tokens"
+                    elif pooling.get("pooling_mode_lasttoken"):
+                        args["pooling"] = "lasttoken"
+        except requests.exceptions.HTTPError:
+            logging.warning(f"{model_id}: 1_Pooling/config.json not found.")
+
+        return args


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 67f40885e..d012157b8 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,6 +357,26 @@ public class TranslatorTest {
             Assertions.assertAlmostEquals(res[0], 0.05103);
         }
 
+        // pooling_mode_max_token
+        criteria =
+                Criteria.builder()
+                        .setTypes(String.class, float[].class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-uncased")
+                        .optArgument("pooling", "max_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], 1.0);
+        }
+
         Criteria<Input, Output> criteria2 =
                 Criteria.builder()
                         .setTypes(Input.class, Output.class)
@@ -364,6 +384,7 @@ public class TranslatorTest {
                         .optBlock(block)
                         .optEngine("PyTorch")
                         .optArgument("tokenizer", "bert-base-uncased")
+                        .optArgument("pooling", "cls_token")
                         .optOption("hasParameter", "false")
                         .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
                         .build();