Repository: deepjavalibrary/djl
Pull request #3642 (2025-03-18 00:41:53)

[tokenizers] Adds SparseRetrievalTranslator
## 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: :warning: Please install the !['codecov app svg image'](https://github.com/codecov/engineering-team/assets/152432831/e90313f4-9d3a-4b63-8b54-cfe14e7ec20d) to ensure uploads and comments are reliably processed by Codecov.

## [Codecov](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3642?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Attention: Patch coverage is `81.76471%` with `31 lines` in your changes missing coverage. Please review.
> Project coverage is 59.74%. Comparing base [(`8aae174`)](https://app.codecov.io/gh/deepjavalibrary/djl/commit/8aae1746b394a8f3a94ee787ad69b8b3d4cc65de?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) to head [(`9ead1b4`)](https://app.codecov.io/gh/deepjavalibrary/djl/commit/9ead1b457e8d7cc1b8bbda7b244e7b8b12ed650e?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> Report is 11 commits behind head on master.

| [Files with missing lines](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3642?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Patch % | Lines |
|---|---|---|
| [...p/translator/SparseRetrievalServingTranslator.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3642?src=pr&el=tree&filepath=api%2Fsrc%2Fmain%2Fjava%2Fai%2Fdjl%2Fmodality%2Fnlp%2Ftranslator%2FSparseRetrievalServingTranslator.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC90cmFuc2xhdG9yL1NwYXJzZVJldHJpZXZhbFNlcnZpbmdUcmFuc2xhdG9yLmphdmE=) | 61.11% | [21 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3642?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...gingface/translator/SparseRetrievalTranslator.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3642?src=pr&el=tree&filepath=extensions%2Ftokenizers%2Fsrc%2Fmain%2Fjava%2Fai%2Fdjl%2Fhuggingface%2Ftranslator%2FSparseRetrievalTranslator.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZXh0ZW5zaW9ucy90b2tlbml6ZXJzL3NyYy9tYWluL2phdmEvYWkvZGpsL2h1Z2dpbmdmYWNlL3RyYW5zbGF0b3IvU3BhcnNlUmV0cmlldmFsVHJhbnNsYXRvci5qYXZh) | 91.56% | [3 Missing and 4 partials :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3642?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...ace/translator/TextEmbeddingTranslatorFactory.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3642?src=pr&el=tree&filepath=extensions%2Ftokenizers%2Fsrc%2Fmain%2Fjava%2Fai%2Fdjl%2Fhuggingface%2Ftranslator%2FTextEmbeddingTranslatorFactory.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZXh0ZW5zaW9ucy90b2tlbml6ZXJzL3NyYy9tYWluL2phdmEvYWkvZGpsL2h1Z2dpbmdmYWNlL3RyYW5zbGF0b3IvVGV4dEVtYmVkZGluZ1RyYW5zbGF0b3JGYWN0b3J5LmphdmE=) | 66.66% | [1 Missing and 2 partials :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3642?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |

:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

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


```diff
@@             Coverage Diff              @@
##             master    #3642      +/-   ##
============================================
- Coverage     62.05%   59.74%   -2.32%     
+ Complexity     6235     6024     -211     
============================================
  Files           685      694       +9     
  Lines         33508    33963     +455     
  Branches       3572     3627      +55     
============================================
- Hits          20795    20292     -503     
- Misses        11129    12135    +1006     
+ Partials       1584     1536      -48     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3642?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>

Patch:
diff --git a/api/src/main/java/ai/djl/modality/nlp/EmbeddingOutput.java b/api/src/main/java/ai/djl/modality/nlp/EmbeddingOutput.java
new file mode 100644
index 000000000..532549a1b
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/EmbeddingOutput.java
@@ -0,0 +1,68 @@
+/*
+ * Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.nlp;
+
+import com.google.gson.annotations.SerializedName;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+/** A class represents text embedding output. */
+public class EmbeddingOutput {
+
+    @SerializedName("dense_vecs")
+    private float[] denseEmbedding;
+
+    @SerializedName("lexical_weights")
+    private LinkedHashMap<String, Float> lexicalWeights; // NOPMD
+
+    /**
+     * Returns the dense embedding.
+     *
+     * @return the dense embedding
+     */
+    public float[] getDenseEmbedding() {
+        return denseEmbedding;
+    }
+
+    /**
+     * Sets the dense embedding.
+     *
+     * @param denseEmbedding the dense embedding
+     */
+    public void setDenseEmbedding(float[] denseEmbedding) {
+        this.denseEmbedding = denseEmbedding;
+    }
+
+    /**
+     * Returns the token weights map.
+     *
+     * @return the token weights map
+     */
+    public Map<String, Float> getLexicalWeights() {
+        return lexicalWeights;
+    }
+
+    /**
+     * Adds the token weights to the output.
+     *
+     * @param tokenId the token id
+     * @param weights the token weights
+     */
+    public void addTokenWeights(String tokenId, float weights) {
+        if (lexicalWeights == null) {
+            lexicalWeights = new LinkedHashMap<>(); // NOPMD
+        }
+        lexicalWeights.put(tokenId, weights);
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/nlp/translator/SparseRetrievalServingTranslator.java b/api/src/main/java/ai/djl/modality/nlp/translator/SparseRetrievalServingTranslator.java
new file mode 100644
index 000000000..12531ff53
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/translator/SparseRetrievalServingTranslator.java
@@ -0,0 +1,135 @@
+/*
+ * Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.nlp.translator;
+
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.nlp.EmbeddingOutput;
+import ai.djl.modality.nlp.TextPrompt;
+import ai.djl.ndarray.BytesSupplier;
+import ai.djl.ndarray.NDList;
+import ai.djl.translate.Batchifier;
+import ai.djl.translate.TranslateException;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorContext;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/** A {@link Translator} that can handle generic text embedding {@link Input} and {@link Output}. */
+public class SparseRetrievalServingTranslator implements Translator<Input, Output> {
+
+    private Translator<String, EmbeddingOutput> translator;
+
+    /**
+     * Constructs a {@code SparseRetrievalServingTranslator} instance.
+     *
+     * @param translator a {@code Translator} processes text embedding input
+     */
+    public SparseRetrievalServingTranslator(Translator<String, EmbeddingOutput> translator) {
+        this.translator = translator;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void prepare(TranslatorContext ctx) throws Exception {
+        translator.prepare(ctx);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, Input input) throws Exception {
+        if (input.getContent().isEmpty()) {
+            throw new TranslateException("Input data is empty.");
+        }
+
+        TextPrompt prompt = TextPrompt.parseInput(input);
+        if (prompt.isBatch()) {
+            ctx.setAttachment("batch", Boolean.TRUE);
+            return translator.batchProcessInput(ctx, prompt.getBatch());
+        }
+
+        NDList ret = translator.processInput(ctx, prompt.getText());
+        Batchifier batchifier = translator.getBatchifier();
+        if (batchifier != null) {
+            NDList[] batch = {ret};
+            return batchifier.batchify(batch);
+        }
+        return ret;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    @SuppressWarnings("PMD.SignatureDeclareThrowsException")
+    public NDList batchProcessInput(TranslatorContext ctx, List<Input> inputs) throws Exception {
+        int[] mapping = new int[inputs.size()];
+        List<String> prompts = new ArrayList<>(mapping.length);
+        for (int i = 0; i < mapping.length; ++i) {
+            TextPrompt prompt = TextPrompt.parseInput(inputs.get(i));
+            if (prompt.isBatch()) {
+                List<String> batch = prompt.getBatch();
+                mapping[i] = batch.size();
+                prompts.addAll(batch);
+            } else {
+                mapping[i] = -1;
+                prompts.add(prompt.getText());
+            }
+        }
+        ctx.setAttachment("mapping", mapping);
+        return translator.batchProcessInput(ctx, prompts);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Output processOutput(TranslatorContext ctx, NDList list) throws Exception {
+        Output output = new Output();
+        output.addProperty("Content-Type", "application/json");
+        if (ctx.getAttachment("batch") != null) {
+            output.add(BytesSupplier.wrapAsJson(translator.batchProcessOutput(ctx, list)));
+        } else {
+            Batchifier batchifier = translator.getBatchifier();
+            if (batchifier != null) {
+                list = batchifier.unbatchify(list)[0];
+            }
+            output.add(BytesSupplier.wrapAsJson(translator.processOutput(ctx, list)));
+        }
+        return output;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    @SuppressWarnings("PMD.SignatureDeclareThrowsException")
+    public List<Output> batchProcessOutput(TranslatorContext ctx, NDList list) throws Exception {
+        List<EmbeddingOutput> outputs = translator.batchProcessOutput(ctx, list);
+        int[] mapping = (int[]) ctx.getAttachment("mapping");
+        List<Output> ret = new ArrayList<>(mapping.length);
+        int index = 0;
+        for (int size : mapping) {
+            Output output = new Output();
+            output.addProperty("Content-Type", "application/json");
+            if (size == -1) {
+                // non-batching
+                output.add(BytesSupplier.wrapAsJson(outputs.get(index++)));
+            } else {
+                // client side batching
+                List<EmbeddingOutput> embeddings = new ArrayList<>(size);
+                for (int j = 0; j < size; ++j) {
+                    embeddings.add(outputs.get(index++));
+                }
+                output.add(BytesSupplier.wrapAsJson(embeddings));
+            }
+            ret.add(output);
+        }
+        return ret;
+    }
+}
diff --git a/api/src/main/java/ai/djl/util/JsonSerializable.java b/api/src/main/java/ai/djl/util/JsonSerializable.java
index e321ec21e..603aaf1c4 100644
--- a/api/src/main/java/ai/djl/util/JsonSerializable.java
+++ b/api/src/main/java/ai/djl/util/JsonSerializable.java
@@ -52,7 +52,7 @@ public interface JsonSerializable extends Serializable, BytesSupplier {
     /** {@inheritDoc} */
     JsonElement serialize();
 
-    /** A customized Gson serializer to serialize the {@code Segmentation} object. */
+    /** A customized Gson serializer to serialize the {@code JsonSerializable} object. */
     final class Serializer implements JsonSerializer<JsonSerializable> {
 
         /** {@inheritDoc} */
diff --git a/api/src/main/java/ai/djl/util/PairList.java b/api/src/main/java/ai/djl/util/PairList.java
index e7a2eed74..36dd34ac4 100644
--- a/api/src/main/java/ai/djl/util/PairList.java
+++ b/api/src/main/java/ai/djl/util/PairList.java
@@ -14,11 +14,11 @@ package ai.djl.util;
 
 import java.util.ArrayList;
 import java.util.Iterator;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Objects;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 
@@ -369,7 +369,7 @@ public class PairList<K, V> implements Iterable<Pair<K, V>> {
      */
     public Map<K, V> toMap(boolean checkDuplicate) {
         int size = keys.size();
-        Map<K, V> map = new ConcurrentHashMap<>(size * 3 / 2);
+        Map<K, V> map = new LinkedHashMap<>(size * 3 / 2); // NOPMD
         for (int i = 0; i < size; ++i) {
             if (map.put(keys.get(i), values.get(i)) != null && checkDuplicate) {
                 throw new IllegalStateException("Duplicate keys: " + keys.get(i));
diff --git a/examples/src/main/java/ai/djl/examples/inference/nlp/SparseRetrieval.java b/examples/src/main/java/ai/djl/examples/inference/nlp/SparseRetrieval.java
new file mode 100644
index 000000000..3feff5907
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/inference/nlp/SparseRetrieval.java
@@ -0,0 +1,65 @@
+/*
+ * Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.examples.inference.nlp;
+
+import ai.djl.ModelException;
+import ai.djl.huggingface.translator.TextEmbeddingTranslatorFactory;
+import ai.djl.inference.Predictor;
+import ai.djl.modality.nlp.EmbeddingOutput;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.ZooModel;
+import ai.djl.training.util.ProgressBar;
+import ai.djl.translate.TranslateException;
+import ai.djl.util.JsonUtils;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.List;
+
+public final class SparseRetrieval {
+
+    private static final Logger logger = LoggerFactory.getLogger(SparseRetrieval.class);
+
+    private SparseRetrieval() {}
+
+    public static void main(String[] args) throws ModelException, IOException, TranslateException {
+        List<EmbeddingOutput> ret = predict();
+        for (EmbeddingOutput result : ret) {
+            logger.info("text embedding: {}", JsonUtils.GSON_PRETTY.toJson(result));
+        }
+    }
+
+    public static List<EmbeddingOutput> predict()
+            throws ModelException, IOException, TranslateException {
+        Criteria<String, EmbeddingOutput> criteria =
+                Criteria.builder()
+                        .setTypes(String.class, EmbeddingOutput.class)
+                        .optModelUrls("djl://ai.djl.huggingface.pytorch/BAAI/bge-m3")
+                        .optEngine("PyTorch")
+                        .optArgument("sparse", true) // use SparseRetrievalTranslator
+                        .optArgument("returnDenseEmbedding", false) // only returns sparse embedding
+                        .optArgument("sparseLinear", "sparse_linear.safetensors")
+                        .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
+                        .optProgress(new ProgressBar())
+                        .build();
+        List<String> inputs =
+                Arrays.asList("This is an example sentence", "What is sparse retrieval?");
+        try (ZooModel<String, EmbeddingOutput> model = criteria.loadModel();
+                Predictor<String, EmbeddingOutput> predictor = model.newPredictor()) {
+            return predictor.batchPredict(inputs);
+        }
+    }
+}
diff --git a/examples/src/main/java/ai/djl/examples/inference/nlp/ZeroShotClassification.java b/examples/src/main/java/ai/djl/examples/inference/nlp/ZeroShotClassification.java
index 0fc85a3fb..581ec809c 100644
--- a/examples/src/main/java/ai/djl/examples/inference/nlp/ZeroShotClassification.java
+++ b/examples/src/main/java/ai/djl/examples/inference/nlp/ZeroShotClassification.java
@@ -19,6 +19,7 @@ import ai.djl.modality.nlp.translator.ZeroShotClassificationInput;
 import ai.djl.modality.nlp.translator.ZeroShotClassificationOutput;
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ZooModel;
+import ai.djl.training.util.ProgressBar;
 import ai.djl.translate.TranslateException;
 import ai.djl.util.JsonUtils;
 
@@ -51,6 +52,7 @@ public final class ZeroShotClassification {
                         .optModelUrls("djl://ai.djl.huggingface.pytorch/facebook/bart-large-mnli")
                         .optEngine("PyTorch")
                         .optTranslatorFactory(new ZeroShotClassificationTranslatorFactory())
+                        .optProgress(new ProgressBar())
                         .build();
         String prompt = "one day I will see the world";
         String[] candidates = {"travel", "cooking", "dancing", "exploration"};
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslator.java
index 843bb5716..5cf275c0f 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslator.java
@@ -143,7 +143,7 @@ public class CrossEncoderTranslator implements Translator<StringPair, float[]> {
         return builder;
     }
 
-    /** The builder for question answering translator. */
+    /** The builder for cross encoder translator. */
     public static final class Builder {
 
         private HuggingFaceTokenizer tokenizer;
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslator.java
index 7fe20e0fe..b37cf1535 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslator.java
@@ -24,7 +24,6 @@ import ai.djl.translate.TranslateException;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorContext;
 
-import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -250,9 +249,8 @@ public class FillMaskTranslator implements Translator<String, Classifications> {
          * Builds the translator.
          *
          * @return the new translator
-         * @throws IOException if I/O error occurs
          */
-        public FillMaskTranslator build() throws IOException {
+        public FillMaskTranslator build() {
             return new FillMaskTranslator(
                     tokenizer, maskedToken, topK, includeTokenTypes, int32, batchifier);
         }
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslator.java
index 92cba3726..e3bab067c 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslator.java
@@ -26,7 +26,6 @@ import ai.djl.translate.TranslatorContext;
 import ai.djl.util.JsonUtils;
 import ai.djl.util.PairList;
 
-import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -262,9 +261,8 @@ public class QuestionAnsweringTranslator implements Translator<QAInput, String>
          * Builds the translator.
          *
          * @return the new translator
-         * @throws IOException if I/O error occurs
          */
-        public QuestionAnsweringTranslator build() throws IOException {
+        public QuestionAnsweringTranslator build() {
             return new QuestionAnsweringTranslator(
                     tokenizer, includeTokenTypes, int32, batchifier, detail);
         }
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/SparseRetrievalTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/SparseRetrievalTranslator.java
new file mode 100644
index 000000000..e664eb752
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/SparseRetrievalTranslator.java
@@ -0,0 +1,231 @@
+/*
+ * Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.huggingface.translator;
+
+import ai.djl.huggingface.tokenizers.Encoding;
+import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
+import ai.djl.modality.nlp.EmbeddingOutput;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.nn.Activation;
+import ai.djl.translate.ArgumentsUtil;
+import ai.djl.translate.TranslateException;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorContext;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.FloatBuffer;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/** The translator handles sparse retrieval for Huggingface text embedding model. */
+public class SparseRetrievalTranslator implements Translator<String, EmbeddingOutput> {
+
+    private static final String[] SPECIAL_TOKENS = {
+        "cls_token", "eos_token", "pad_token", "unk_token"
+    };
+
+    private HuggingFaceTokenizer tokenizer;
+    private TextEmbeddingTranslator translator;
+    private boolean includeTokenTypes;
+    private boolean int32;
+    private boolean returnDenseEmbedding;
+    private Set<Long> unusedTokens;
+    private String sparseLinear;
+    private NDList sparseLinearModel;
+
+    SparseRetrievalTranslator(Builder builder) {
+        this.tokenizer = builder.tokenizer;
+        this.translator = builder.baseBuilder.build();
+        this.includeTokenTypes = builder.baseBuilder.includeTokenTypes;
+        this.int32 = builder.baseBuilder.int32;
+        this.returnDenseEmbedding = builder.returnDenseEmbedding;
+        this.sparseLinear = builder.sparseLinear;
+        Encoding encoding = tokenizer.encode(SPECIAL_TOKENS);
+        unusedTokens = Arrays.stream(encoding.getIds()).boxed().collect(Collectors.toSet());
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void prepare(TranslatorContext ctx) throws Exception {
+        NDManager manager = ctx.getPredictorManager().newSubManager();
+        if (returnDenseEmbedding) {
+            translator.prepare(ctx);
+        }
+        if (sparseLinear != null) {
+            Path file = Paths.get(sparseLinear);
+            if (!file.isAbsolute()) {
+                file = ctx.getModel().getModelPath().resolve(file);
+            }
+            if (Files.notExists(file)) {
+                throw new TranslateException("sparseLinear file does not exist: " + sparseLinear);
+            }
+            try (InputStream is = Files.newInputStream(file)) {
+                sparseLinearModel = NDList.decode(manager, is);
+            }
+        }
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, String input) {
+        return batchProcessInput(ctx, Collections.singletonList(input));
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList batchProcessInput(TranslatorContext ctx, List<String> inputs) {
+        NDManager manager = ctx.getNDManager();
+        Encoding[] encodings = tokenizer.batchEncode(inputs);
+        NDList list = Encoding.toNDList(encodings, manager, includeTokenTypes, int32);
+        ctx.setAttachment("encodings", encodings);
+        ctx.setAttachment("attentionMask", list.get(1));
+        return list;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public EmbeddingOutput processOutput(TranslatorContext ctx, NDList list) {
+        return Objects.requireNonNull(batchProcessOutput(ctx, list)).get(0);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public List<EmbeddingOutput> batchProcessOutput(TranslatorContext ctx, NDList list) {
+        Encoding[] encodings = (Encoding[]) ctx.getAttachment("encodings");
+        int batchSize = encodings.length;
+
+        List<EmbeddingOutput> embeddings = new ArrayList<>();
+        NDArray lastHiddenState = list.get("last_hidden_state");
+        if (lastHiddenState == null) {
+            lastHiddenState = list.get(0); // only pytorch returns "last_hidden_state" name
+        }
+        NDArray weight =
+                sparseLinearModel.get("weight").toType(lastHiddenState.getDataType(), false);
+        NDArray bias = sparseLinearModel.get("bias").toType(lastHiddenState.getDataType(), false);
+        NDArray array =
+                lastHiddenState.getNDArrayInternal().linear(lastHiddenState, weight, bias).get(0);
+        array = Activation.relu(array);
+        NDArray sparseVecs = array.squeeze(-1);
+        float[] data = sparseVecs.toFloatArray();
+        int index = 0;
+        for (Encoding encoding : encodings) {
+            long[] tokenIds = encoding.getIds();
+            EmbeddingOutput embedding = new EmbeddingOutput();
+            embeddings.add(embedding);
+            for (long idx : tokenIds) {
+                float w = data[index++];
+                if (!unusedTokens.contains(idx) && w > 0) {
+                    embedding.addTokenWeights(String.valueOf(idx), w);
+                }
+            }
+        }
+
+        if (returnDenseEmbedding) {
+            NDArray attentionMask = (NDArray) ctx.getAttachment("attentionMask");
+            NDArray output = translator.processEmbedding(list, attentionMask);
+            FloatBuffer fb = output.toByteBuffer().asFloatBuffer();
+            int denseEmbeddingSize = fb.remaining() / batchSize;
+            for (EmbeddingOutput embedding : embeddings) {
+                float[] buf = new float[denseEmbeddingSize];
+                fb.get(buf);
+                embedding.setDenseEmbedding(buf);
+            }
+        }
+        return embeddings;
+    }
+
+    /**
+     * Creates a builder to build a {@code SparseRetrievalTranslator}.
+     *
+     * @param tokenizer the tokenizer
+     * @param arguments the models' arguments
+     * @return a new builder
+     */
+    public static Builder builder(HuggingFaceTokenizer tokenizer, Map<String, ?> arguments) {
+        Builder builder = new Builder(tokenizer);
+        builder.configure(arguments);
+
+        return builder;
+    }
+
+    /** The builder for question answering translator. */
+    public static final class Builder {
+
+        HuggingFaceTokenizer tokenizer;
+        TextEmbeddingTranslator.Builder baseBuilder;
+        boolean returnDenseEmbedding;
+        String sparseLinear;
+
+        Builder(HuggingFaceTokenizer tokenizer) {
+            this.tokenizer = tokenizer;
+            baseBuilder = TextEmbeddingTranslator.builder(tokenizer);
+            sparseLinear = "sparse_linear.safetensors";
+        }
+
+        /**
+         * Sets if apply sigmoid for the {@link Translator}.
+         *
+         * @param returnDenseEmbedding true to output dense embedding
+         * @return this builder
+         */
+        public Builder optReturnDenseEmbedding(boolean returnDenseEmbedding) {
+            this.returnDenseEmbedding = returnDenseEmbedding;
+            return this;
+        }
+
+        /**
+         * Sets the sparse linear layer model file for the {@link Translator}.
+         *
+         * @param sparseLinear path to sparse linear layer model file
+         * @return this builder
+         */
+        public Builder optSparseLinear(String sparseLinear) {
+            this.sparseLinear = sparseLinear;
+            return this;
+        }
+
+        /**
+         * Configures the builder with the model arguments.
+         *
+         * @param arguments the model arguments
+         */
+        public void configure(Map<String, ?> arguments) {
+            baseBuilder.configure(arguments);
+            optReturnDenseEmbedding(
+                    ArgumentsUtil.booleanValue(arguments, "returnDenseEmbedding", false));
+            optSparseLinear(ArgumentsUtil.stringValue(arguments, "sparseLinear", sparseLinear));
+        }
+
+        /**
+         * Builds the translator.
+         *
+         * @return the new translator
+         * @throws IOException if I/O error occurs
+         */
+        public SparseRetrievalTranslator build() throws IOException {
+            return new SparseRetrievalTranslator(this);
+        }
+    }
+}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslator.java
index 82c9112d3..6f2959257 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslator.java
@@ -209,9 +209,8 @@ public class TextClassificationTranslator implements Translator<String, Classifi
          * Builds the translator.
          *
          * @return the new translator
-         * @throws IOException if I/O error occurs
          */
-        public TextClassificationTranslator build() throws IOException {
+        public TextClassificationTranslator build() {
             return new TextClassificationTranslator(
                     tokenizer, includeTokenTypes, int32, batchifier);
         }
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 1ef7f433b..cbbc81b4d 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
@@ -24,7 +24,6 @@ import ai.djl.translate.Batchifier;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorContext;
 
-import java.io.IOException;
 import java.io.InputStream;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -51,25 +50,16 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
     private NDList denseModel;
     private NDList layerNormModel;
 
-    TextEmbeddingTranslator(
-            HuggingFaceTokenizer tokenizer,
-            Batchifier batchifier,
-            String pooling,
-            boolean normalize,
-            boolean includeTokenTypes,
-            boolean int32,
-            String dense,
-            String denseActivation,
-            String layerNorm) {
-        this.tokenizer = tokenizer;
-        this.batchifier = batchifier;
-        this.pooling = pooling;
-        this.normalize = normalize;
-        this.includeTokenTypes = includeTokenTypes;
-        this.int32 = int32;
-        this.dense = dense;
-        this.denseActivation = denseActivation;
-        this.layerNorm = layerNorm;
+    TextEmbeddingTranslator(Builder builder) {
+        this.tokenizer = builder.tokenizer;
+        this.batchifier = builder.batchifier;
+        this.pooling = builder.pooling;
+        this.normalize = builder.normalize;
+        this.includeTokenTypes = builder.includeTokenTypes;
+        this.int32 = builder.int32;
+        this.dense = builder.dense;
+        this.denseActivation = builder.denseActivation;
+        this.layerNorm = builder.layerNorm;
     }
 
     /** {@inheritDoc} */
@@ -155,7 +145,7 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
         return ret;
     }
 
-    private NDArray processEmbedding(NDList list, NDArray attentionMask) {
+    NDArray processEmbedding(NDList list, NDArray attentionMask) {
         NDArray embedding = list.get("last_hidden_state");
         if (embedding == null) {
             // For Onnx model, NDArray name is not present
@@ -283,15 +273,15 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
     /** The builder for token classification translator. */
     public static final class Builder {
 
-        private HuggingFaceTokenizer tokenizer;
-        private Batchifier batchifier = Batchifier.STACK;
-        private boolean normalize = true;
-        private String pooling = "mean";
-        private boolean includeTokenTypes;
-        private boolean int32;
-        private String dense;
-        private String denseActivation;
-        private String layerNorm;
+        HuggingFaceTokenizer tokenizer;
+        Batchifier batchifier = Batchifier.STACK;
+        boolean normalize = true;
+        String pooling = "mean";
+        boolean includeTokenTypes;
+        boolean int32;
+        String dense;
+        String denseActivation;
+        String layerNorm;
 
         Builder(HuggingFaceTokenizer tokenizer) {
             this.tokenizer = tokenizer;
@@ -416,19 +406,9 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
          * Builds the translator.
          *
          * @return the new translator
-         * @throws IOException if I/O error occurs
          */
-        public TextEmbeddingTranslator build() throws IOException {
-            return new TextEmbeddingTranslator(
-                    tokenizer,
-                    batchifier,
-                    pooling,
-                    normalize,
-                    includeTokenTypes,
-                    int32,
-                    dense,
-                    denseActivation,
-                    layerNorm);
+        public TextEmbeddingTranslator build() {
+            return new TextEmbeddingTranslator(this);
         }
     }
 }
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java
index d91c35cf2..658b02c84 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java
@@ -16,7 +16,9 @@ import ai.djl.Model;
 import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
 import ai.djl.modality.Input;
 import ai.djl.modality.Output;
+import ai.djl.modality.nlp.EmbeddingOutput;
 import ai.djl.modality.nlp.translator.CrossEncoderServingTranslator;
+import ai.djl.modality.nlp.translator.SparseRetrievalServingTranslator;
 import ai.djl.modality.nlp.translator.TextEmbeddingServingTranslator;
 import ai.djl.translate.ArgumentsUtil;
 import ai.djl.translate.TranslateException;
@@ -42,6 +44,7 @@ public class TextEmbeddingTranslatorFactory implements TranslatorFactory, Serial
 
     static {
         SUPPORTED_TYPES.add(new Pair<>(String.class, float[].class));
+        SUPPORTED_TYPES.add(new Pair<>(String.class, EmbeddingOutput.class));
         SUPPORTED_TYPES.add(new Pair<>(StringPair.class, float[].class));
         SUPPORTED_TYPES.add(new Pair<>(Input.class, Output.class));
     }
@@ -74,6 +77,15 @@ public class TextEmbeddingTranslatorFactory implements TranslatorFactory, Serial
                     return (Translator<I, O>) new CrossEncoderServingTranslator(translator);
                 }
                 throw new IllegalArgumentException("Unsupported input/output types.");
+            } else if (ArgumentsUtil.booleanValue(arguments, "sparse")) {
+                SparseRetrievalTranslator translator =
+                        SparseRetrievalTranslator.builder(tokenizer, arguments).build();
+                if (input == String.class && output == EmbeddingOutput.class) {
+                    return (Translator<I, O>) translator;
+                } else if (input == Input.class && output == Output.class) {
+                    return (Translator<I, O>) new SparseRetrievalServingTranslator(translator);
+                }
+                throw new IllegalArgumentException("Unsupported input/output types.");
             }
 
             TextEmbeddingTranslator translator =
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslator.java
index 43aed0db2..3099e8e12 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslator.java
@@ -242,9 +242,8 @@ public class TokenClassificationTranslator implements Translator<String, NamedEn
          * Builds the translator.
          *
          * @return the new translator
-         * @throws IOException if I/O error occurs
          */
-        public TokenClassificationTranslator build() throws IOException {
+        public TokenClassificationTranslator build() {
             return new TokenClassificationTranslator(
                     tokenizer, includeTokenTypes, int32, softmax, batchifier);
         }
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/ZeroShotClassificationTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/ZeroShotClassificationTranslator.java
index 6159cfd17..7fb801bd2 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/ZeroShotClassificationTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/ZeroShotClassificationTranslator.java
@@ -177,9 +177,8 @@ public class ZeroShotClassificationTranslator
          * Builds the translator.
          *
          * @return the new translator
-         * @throws IOException if I/O error occurs
          */
-        public ZeroShotClassificationTranslator build() throws IOException {
+        public ZeroShotClassificationTranslator build() {
             return new ZeroShotClassificationTranslator(tokenizer, int32);
         }
     }
diff --git a/extensions/tokenizers/src/main/python/djl_converter/sentence_similarity_converter.py b/extensions/tokenizers/src/main/python/djl_converter/sentence_similarity_converter.py
index f46925aea..c508b37ce 100644
--- a/extensions/tokenizers/src/main/python/djl_converter/sentence_similarity_converter.py
+++ b/extensions/tokenizers/src/main/python/djl_converter/sentence_similarity_converter.py
@@ -20,6 +20,7 @@ import torch
 from transformers import AutoTokenizer, AutoModel, AutoConfig
 
 from djl_converter.huggingface_converter import HuggingfaceConverter, PipelineHolder
+from djl_converter.safetensors_convert import convert_file
 
 
 class SentenceSimilarityConverter(HuggingfaceConverter):
@@ -162,6 +163,21 @@ class SentenceSimilarityConverter(HuggingfaceConverter):
         if not normalize:
             args["normalize"] = "false"
 
+        try:
+            file = self.get_file(model_id, "sparse_linear.pt")
+            if os.path.exists(file):
+                target = os.path.join(temp_dir, "sparse_linear.safetensors")
+                convert_file(file, target)
+        except requests.exceptions.HTTPError:
+            pass
+
+        try:
+            file = self.get_file(model_id, "sparse_linear.safetensors")
+            if os.path.exists(file):
+                shutil.copyfile(file, os.path.join(temp_dir, "sparse_linear.safetensors"))
+        except requests.exceptions.HTTPError:
+            pass
+
         return args
 
     @staticmethod
diff --git a/integration/src/main/resources/log4j2.xml b/integration/src/main/resources/log4j2.xml
index cb6d3c6fb..a4fad32c7 100644
--- a/integration/src/main/resources/log4j2.xml
+++ b/integration/src/main/resources/log4j2.xml
@@ -1,5 +1,5 @@
 <?xml version="1.0" encoding="UTF-8"?>
-<Configuration status="INFO">
+<Configuration status="WARN">
   <Appenders>
     <Console name="console" target="SYSTEM_OUT">
       <PatternLayout


Test patch:
diff --git a/extensions/aws-ai/src/test/resources/log4j2.xml b/extensions/aws-ai/src/test/resources/log4j2.xml
index 358b5a2a0..9a7074ef2 100644
--- a/extensions/aws-ai/src/test/resources/log4j2.xml
+++ b/extensions/aws-ai/src/test/resources/log4j2.xml
@@ -1,5 +1,5 @@
 <?xml version="1.0" encoding="UTF-8"?>
-<Configuration status="INFO">
+<Configuration status="WARN">
     <Appenders>
         <Console name="console" target="SYSTEM_OUT">
             <PatternLayout
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/SparseRetrievalTranslatorTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/SparseRetrievalTranslatorTest.java
new file mode 100644
index 000000000..bc34c8792
--- /dev/null
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/SparseRetrievalTranslatorTest.java
@@ -0,0 +1,145 @@
+/*
+ * Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.huggingface.tokenizers;
+
+import ai.djl.ModelException;
+import ai.djl.huggingface.translator.TextEmbeddingTranslatorFactory;
+import ai.djl.inference.Predictor;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.nlp.EmbeddingOutput;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.ZooModel;
+import ai.djl.testing.Assertions;
+import ai.djl.translate.TranslateException;
+import ai.djl.util.JsonUtils;
+
+import com.google.gson.JsonArray;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class SparseRetrievalTranslatorTest {
+
+    @Test
+    public void testSparseRetrievalTranslator()
+            throws ModelException, IOException, TranslateException {
+        String[] text = {"This is an example sentence", "What is sparse retrieval?"};
+
+        Path modelDir = Paths.get("build/model");
+        Path sparseLinear = modelDir.resolve("sparse_linear.safetensors");
+        Files.createDirectories(modelDir);
+        try (NDManager manager = NDManager.newBaseManager("PyTorch")) {
+            NDArray weight = manager.ones(new Shape(1, 1024));
+            weight.setName("weight");
+            NDArray bias = manager.ones(new Shape(1));
+            bias.setName("bias");
+            NDList linear = new NDList(weight, bias);
+            try (OutputStream os = Files.newOutputStream(sparseLinear)) {
+                linear.encode(os, NDList.Encoding.SAFETENSORS);
+            }
+        }
+
+        Criteria<String, EmbeddingOutput> criteria =
+                Criteria.builder()
+                        .setTypes(String.class, EmbeddingOutput.class)
+                        .optModelPath(modelDir)
+                        .optEngine("PyTorch")
+                        .optArgument("blockFactory", "ai.djl.nn.OnesBlockFactory")
+                        .optArgument("block_shapes", "(2,9,1024)")
+                        .optArgument("block_names", "last_hidden_state")
+                        .optArgument("tokenizer", "bert-base-uncased")
+                        .optArgument("sparse", true)
+                        .optArgument("returnDenseEmbedding", true)
+                        .optArgument("pooling", "cls")
+                        .optArgument("sparseLinear", sparseLinear.toAbsolutePath().toString())
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
+                        .build();
+
+        try (ZooModel<String, EmbeddingOutput> model = criteria.loadModel();
+                Predictor<String, EmbeddingOutput> predictor = model.newPredictor()) {
+            EmbeddingOutput res = predictor.predict(text[0]);
+            Map<String, Float> tokenWeights = res.getLexicalWeights();
+            Assert.assertNotNull(res.getDenseEmbedding());
+            Assert.assertEquals(tokenWeights.size(), 5);
+            Assertions.assertAlmostEquals(tokenWeights.get("2023"), 1025.0);
+        }
+
+        Criteria<Input, Output> criteria2 =
+                Criteria.builder()
+                        .setTypes(Input.class, Output.class)
+                        .optModelPath(modelDir)
+                        .optArgument("blockFactory", "ai.djl.nn.OnesBlockFactory")
+                        .optArgument("block_shapes", "(2,9,1024)")
+                        .optArgument("block_names", "last_hidden_state")
+                        .optArgument("tokenizer", "bert-base-uncased")
+                        .optArgument("sparse", true)
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
+                        .build();
+
+        try (ZooModel<Input, Output> model = criteria2.loadModel();
+                Predictor<Input, Output> predictor = model.newPredictor()) {
+            Input input = new Input();
+            input.add(text[0]);
+            Output out = predictor.predict(input);
+            EmbeddingOutput res =
+                    JsonUtils.GSON.fromJson(out.getAsString(0), EmbeddingOutput.class);
+            Map<String, Float> tokenWeights = res.getLexicalWeights();
+            Assert.assertNull(res.getDenseEmbedding());
+            Assert.assertEquals(tokenWeights.size(), 5);
+            Assertions.assertAlmostEquals(tokenWeights.get("2023"), 1025.0);
+
+            // batch serving predict
+            Input input2 = new Input();
+            input2.add(text[1]);
+            List<Input> batchInput = Arrays.asList(input, input2);
+            List<Output> batchOutput = predictor.batchPredict(batchInput);
+            out = batchOutput.get(0);
+            res = JsonUtils.GSON.fromJson(out.getAsString(0), EmbeddingOutput.class);
+            tokenWeights = res.getLexicalWeights();
+            Assert.assertNull(res.getDenseEmbedding());
+            Assert.assertEquals(tokenWeights.size(), 5);
+            Assertions.assertAlmostEquals(tokenWeights.get("2023"), 1025.0);
+
+            // client side batch
+            input = new Input();
+            Map<String, String[]> map = new HashMap<>();
+            map.put("inputs", text);
+            input.add(JsonUtils.GSON.toJson(map));
+            input.addProperty("Content-Type", "application/json; charset=utf-8");
+            out = predictor.predict(input);
+            JsonArray batch = JsonUtils.GSON.fromJson(out.getAsString(0), JsonArray.class);
+            res = JsonUtils.GSON.fromJson(batch.get(0), EmbeddingOutput.class);
+            tokenWeights = res.getLexicalWeights();
+            Assert.assertNull(res.getDenseEmbedding());
+            Assert.assertEquals(tokenWeights.size(), 5);
+            Assertions.assertAlmostEquals(tokenWeights.get("2023"), 1025.0);
+        }
+    }
+}