Repository: deepjavalibrary/djl (java)

Issue:
Can user create custom NLP Dataset by extending from ai.djl.basicdataset.nlp.TextDataset?
I'm trying to create custom dataset class from ai.djl.basicdataset.nlp.TextDataset, but the Builder inside TextDataset has a package-level constructor which I can't extends from other package. Is that a bug? Should the constructor be public?
I' m using dependency below:
```xml
<dependency>
    <groupId>ai.djl</groupId>
    <artifactId>basicdataset</artifactId>
    <version>0.20.0</version>
</dependency>
```

Pull request description:
[dataset] Makes TestDataset constructor protected
Fixes #2269

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


Patch:
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/nlp/TextDataset.java b/basicdataset/src/main/java/ai/djl/basicdataset/nlp/TextDataset.java
index 50ac4e80e..de2fb916e 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/nlp/TextDataset.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/nlp/TextDataset.java
@@ -197,7 +197,7 @@ public abstract class TextDataset extends RandomAccessDataset {
         protected Usage usage;
 
         /** Constructs a new builder. */
-        Builder() {
+        protected Builder() {
             repository = BasicDatasets.REPOSITORY;
             groupId = BasicDatasets.GROUP_ID;
             usage = Usage.TRAIN;
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/aws-ai/src/test/java/ai/djl/aws/sagemaker/SageMakerTest.java b/extensions/aws-ai/src/test/java/ai/djl/aws/sagemaker/SageMakerTest.java
index 6d9d56366..e3cec9a9c 100644
--- a/extensions/aws-ai/src/test/java/ai/djl/aws/sagemaker/SageMakerTest.java
+++ b/extensions/aws-ai/src/test/java/ai/djl/aws/sagemaker/SageMakerTest.java
@@ -28,6 +28,7 @@ import org.testng.annotations.Test;
 
 import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
 import software.amazon.awssdk.core.exception.SdkClientException;
+import software.amazon.awssdk.services.sagemaker.model.SageMakerException;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -82,7 +83,7 @@ public class SageMakerTest {
                 sageMaker.deleteEndpointConfig(true);
                 sageMaker.deleteSageMakerModel(true);
             }
-        } catch (SdkClientException e) {
+        } catch (SdkClientException | SageMakerException e) {
             throw new SkipException("Skip tests that requires permission.", e);
         }
     }
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)