Repository: deepjavalibrary/djl
Pull request #2817 (2023-10-22 21:44:21)

[huggingface] Adds CrossEncoderTranslator
## 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://app.codecov.io/gh/deepjavalibrary/djl/pull/2817?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Attention: `1376 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`bb5073f`)](https://app.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) 72.08% compared to head [(`3f24d78`)](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2817?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.24%.
> Report is 898 commits behind head on master.

: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    #2817      +/-   ##
============================================
+ Coverage     72.08%   72.24%   +0.15%     
- Complexity     5126     7141    +2015     
============================================
  Files           473      707     +234     
  Lines         21970    31813    +9843     
  Branches       2351     3300     +949     
============================================
+ Hits          15838    22982    +7144     
- Misses         4925     7261    +2336     
- Partials       1207     1570     +363     
```


| [Files](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2817?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [...ava/ai/djl/inference/streaming/StreamingBlock.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2817?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL2luZmVyZW5jZS9zdHJlYW1pbmcvU3RyZWFtaW5nQmxvY2suamF2YQ==) | `100.00% <100.00%> (ø)` | |
| [api/src/main/java/ai/djl/metric/Dimension.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2817?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21ldHJpYy9EaW1lbnNpb24uamF2YQ==) | `100.00% <100.00%> (ø)` | |
| [api/src/main/java/ai/djl/metric/Unit.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2817?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21ldHJpYy9Vbml0LmphdmE=) | `100.00% <100.00%> (ø)` | |
| [api/src/main/java/ai/djl/modality/audio/Audio.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2817?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2F1ZGlvL0F1ZGlvLmphdmE=) | `100.00% <100.00%> (ø)` | |
| [api/src/main/java/ai/djl/modality/cv/Image.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2817?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://app.codecov.io/gh/deepjavalibrary/djl/pull/2817?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L011bHRpQm94UHJpb3IuamF2YQ==) | `76.00% <ø> (ø)` | |
| [...ava/ai/djl/modality/cv/output/DetectedObjects.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2817?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9EZXRlY3RlZE9iamVjdHMuamF2YQ==) | `96.29% <100.00%> (+1.29%)` | :arrow_up: |
| [...rc/main/java/ai/djl/modality/cv/output/Joints.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2817?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9Kb2ludHMuamF2YQ==) | `71.42% <100.00%> (ø)` | |
| [.../main/java/ai/djl/modality/cv/output/Landmark.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2817?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9MYW5kbWFyay5qYXZh) | `100.00% <ø> (ø)` | |
| [...i/djl/modality/cv/transform/RandomResizedCrop.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2817?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zZm9ybS9SYW5kb21SZXNpemVkQ3JvcC5qYXZh) | `94.11% <100.00%> (+5.22%)` | :arrow_up: |
| ... and [225 more](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2817?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

... and [378 files with indirect coverage changes](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2817/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2817?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).


Patch:
diff --git a/api/src/main/java/ai/djl/modality/nlp/translator/CrossEncoderServingTranslator.java b/api/src/main/java/ai/djl/modality/nlp/translator/CrossEncoderServingTranslator.java
new file mode 100644
index 000000000..e62167a34
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/translator/CrossEncoderServingTranslator.java
@@ -0,0 +1,115 @@
+/*
+ * Copyright 2023 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.ndarray.BytesSupplier;
+import ai.djl.ndarray.NDList;
+import ai.djl.translate.Batchifier;
+import ai.djl.translate.NoBatchifyTranslator;
+import ai.djl.translate.TranslateException;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorContext;
+import ai.djl.util.JsonUtils;
+import ai.djl.util.PairList;
+import ai.djl.util.StringPair;
+
+import com.google.gson.JsonElement;
+import com.google.gson.JsonParseException;
+
+/** A {@link Translator} that can handle generic cross encoder {@link Input} and {@link Output}. */
+public class CrossEncoderServingTranslator implements NoBatchifyTranslator<Input, Output> {
+
+    private Translator<StringPair, float[]> translator;
+    private Translator<StringPair[], float[][]> batchTranslator;
+
+    /**
+     * Constructs a {@code CrossEncoderServingTranslator} instance.
+     *
+     * @param translator a {@code Translator} processes question answering input
+     */
+    public CrossEncoderServingTranslator(Translator<StringPair, float[]> translator) {
+        this.translator = translator;
+        this.batchTranslator = translator.toBatchTranslator();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void prepare(TranslatorContext ctx) throws Exception {
+        translator.prepare(ctx);
+        batchTranslator.prepare(ctx);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, Input input) throws Exception {
+        PairList<String, BytesSupplier> content = input.getContent();
+        if (content.isEmpty()) {
+            throw new TranslateException("Input data is empty.");
+        }
+
+        String contentType = input.getProperty("Content-Type", null);
+        StringPair pair;
+        if ("application/json".equals(contentType)) {
+            String json = input.getData().getAsString();
+            try {
+                JsonElement element = JsonUtils.GSON.fromJson(json, JsonElement.class);
+                if (element.isJsonArray()) {
+                    ctx.setAttachment("batch", Boolean.TRUE);
+                    StringPair[] inputs = JsonUtils.GSON.fromJson(json, StringPair[].class);
+                    return batchTranslator.processInput(ctx, inputs);
+                }
+
+                pair = JsonUtils.GSON.fromJson(json, StringPair.class);
+                if (pair.getKey() == null || pair.getValue() == null) {
+                    throw new TranslateException("Missing key or value in json.");
+                }
+            } catch (JsonParseException e) {
+                throw new TranslateException("Input is not a valid json.", e);
+            }
+        } else {
+            String key = input.getAsString("key");
+            String value = input.getAsString("value");
+            if (key == null || value == null) {
+                throw new TranslateException("Missing key or value in input.");
+            }
+            pair = new StringPair(key, value);
+        }
+
+        NDList ret = translator.processInput(ctx, pair);
+        Batchifier batchifier = translator.getBatchifier();
+        if (batchifier != null) {
+            NDList[] batch = {ret};
+            return batchifier.batchify(batch);
+        }
+        return ret;
+    }
+
+    /** {@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(batchTranslator.processOutput(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;
+    }
+}
diff --git a/api/src/main/java/ai/djl/util/StringPair.java b/api/src/main/java/ai/djl/util/StringPair.java
new file mode 100644
index 000000000..a42e73961
--- /dev/null
+++ b/api/src/main/java/ai/djl/util/StringPair.java
@@ -0,0 +1,27 @@
+/*
+ * Copyright 2023 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.util;
+
+/** A class containing the string key-value pair. */
+public class StringPair extends Pair<String, String> {
+
+    /**
+     * Constructs a {@code Pair} instance with key and value.
+     *
+     * @param key the key
+     * @param value the value
+     */
+    public StringPair(String key, String value) {
+        super(key, value);
+    }
+}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderBatchTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderBatchTranslator.java
new file mode 100644
index 000000000..6f43c7cb4
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderBatchTranslator.java
@@ -0,0 +1,69 @@
+/*
+ * Copyright 2023 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.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.translate.Batchifier;
+import ai.djl.translate.NoBatchifyTranslator;
+import ai.djl.translate.TranslateException;
+import ai.djl.translate.TranslatorContext;
+import ai.djl.util.PairList;
+import ai.djl.util.StringPair;
+
+import java.util.Arrays;
+
+/** The translator for Huggingface cross encoder model. */
+public class CrossEncoderBatchTranslator implements NoBatchifyTranslator<StringPair[], float[][]> {
+
+    private HuggingFaceTokenizer tokenizer;
+    private boolean includeTokenTypes;
+    private Batchifier batchifier;
+
+    CrossEncoderBatchTranslator(
+            HuggingFaceTokenizer tokenizer, boolean includeTokenTypes, Batchifier batchifier) {
+        this.tokenizer = tokenizer;
+        this.includeTokenTypes = includeTokenTypes;
+        this.batchifier = batchifier;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, StringPair[] inputs)
+            throws TranslateException {
+        NDManager manager = ctx.getNDManager();
+        PairList<String, String> list = new PairList<>(Arrays.asList(inputs));
+        Encoding[] encodings = tokenizer.batchEncode(list);
+        NDList[] batch = new NDList[encodings.length];
+        for (int i = 0; i < encodings.length; ++i) {
+            batch[i] = encodings[i].toNDList(manager, includeTokenTypes);
+        }
+        return batchifier.batchify(batch);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public float[][] processOutput(TranslatorContext ctx, NDList list) {
+        NDList[] batch = batchifier.unbatchify(list);
+        float[][] ret = new float[batch.length][];
+        for (int i = 0; i < batch.length; ++i) {
+            NDArray logits = list.get(0);
+            NDArray result = logits.getNDArrayInternal().sigmoid();
+            ret[i] = result.toFloatArray();
+        }
+        return ret;
+    }
+}
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
new file mode 100644
index 000000000..b88347bc6
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslator.java
@@ -0,0 +1,149 @@
+/*
+ * Copyright 2023 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.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.translate.ArgumentsUtil;
+import ai.djl.translate.Batchifier;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorContext;
+import ai.djl.util.StringPair;
+
+import java.io.IOException;
+import java.util.Map;
+
+/** The translator for Huggingface cross encoder model. */
+public class CrossEncoderTranslator implements Translator<StringPair, float[]> {
+
+    private HuggingFaceTokenizer tokenizer;
+    private boolean includeTokenTypes;
+    private Batchifier batchifier;
+
+    CrossEncoderTranslator(
+            HuggingFaceTokenizer tokenizer, boolean includeTokenTypes, Batchifier batchifier) {
+        this.tokenizer = tokenizer;
+        this.includeTokenTypes = includeTokenTypes;
+        this.batchifier = batchifier;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Batchifier getBatchifier() {
+        return batchifier;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, StringPair input) {
+        Encoding encoding = tokenizer.encode(input.getKey(), input.getValue());
+        ctx.setAttachment("encoding", encoding);
+        return encoding.toNDList(ctx.getNDManager(), includeTokenTypes);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public float[] processOutput(TranslatorContext ctx, NDList list) {
+        NDArray logits = list.get(0);
+        NDArray result = logits.getNDArrayInternal().sigmoid();
+        return result.toFloatArray();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public CrossEncoderBatchTranslator toBatchTranslator(Batchifier batchifier) {
+        tokenizer.enableBatch();
+        return new CrossEncoderBatchTranslator(tokenizer, includeTokenTypes, batchifier);
+    }
+
+    /**
+     * Creates a builder to build a {@code CrossEncoderTranslator}.
+     *
+     * @param tokenizer the tokenizer
+     * @return a new builder
+     */
+    public static Builder builder(HuggingFaceTokenizer tokenizer) {
+        return new Builder(tokenizer);
+    }
+
+    /**
+     * Creates a builder to build a {@code CrossEncoderTranslator}.
+     *
+     * @param tokenizer the tokenizer
+     * @param arguments the models' arguments
+     * @return a new builder
+     */
+    public static Builder builder(HuggingFaceTokenizer tokenizer, Map<String, ?> arguments) {
+        Builder builder = builder(tokenizer);
+        builder.configure(arguments);
+
+        return builder;
+    }
+
+    /** The builder for question answering translator. */
+    public static final class Builder {
+
+        private HuggingFaceTokenizer tokenizer;
+        private boolean includeTokenTypes;
+        private Batchifier batchifier = Batchifier.STACK;
+
+        Builder(HuggingFaceTokenizer tokenizer) {
+            this.tokenizer = tokenizer;
+        }
+
+        /**
+         * Sets if include token types for the {@link Translator}.
+         *
+         * @param includeTokenTypes true to include token types
+         * @return this builder
+         */
+        public Builder optIncludeTokenTypes(boolean includeTokenTypes) {
+            this.includeTokenTypes = includeTokenTypes;
+            return this;
+        }
+
+        /**
+         * Sets the {@link Batchifier} for the {@link Translator}.
+         *
+         * @param batchifier true to include token types
+         * @return this builder
+         */
+        public Builder optBatchifier(Batchifier batchifier) {
+            this.batchifier = batchifier;
+            return this;
+        }
+
+        /**
+         * Configures the builder with the model arguments.
+         *
+         * @param arguments the model arguments
+         */
+        public void configure(Map<String, ?> arguments) {
+            optIncludeTokenTypes(ArgumentsUtil.booleanValue(arguments, "includeTokenTypes"));
+            String batchifierStr = ArgumentsUtil.stringValue(arguments, "batchifier", "stack");
+            optBatchifier(Batchifier.fromString(batchifierStr));
+        }
+
+        /**
+         * Builds the translator.
+         *
+         * @return the new translator
+         * @throws IOException if I/O error occurs
+         */
+        public CrossEncoderTranslator build() throws IOException {
+            return new CrossEncoderTranslator(tokenizer, includeTokenTypes, batchifier);
+        }
+    }
+}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslatorFactory.java
new file mode 100644
index 000000000..f4f9af02c
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslatorFactory.java
@@ -0,0 +1,80 @@
+/*
+ * Copyright 2023 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.Model;
+import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.nlp.translator.CrossEncoderServingTranslator;
+import ai.djl.translate.TranslateException;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorFactory;
+import ai.djl.util.Pair;
+import ai.djl.util.StringPair;
+
+import java.io.IOException;
+import java.io.Serializable;
+import java.lang.reflect.Type;
+import java.nio.file.Path;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+/** A {@link TranslatorFactory} that creates a {@link CrossEncoderTranslatorFactory} instance. */
+public class CrossEncoderTranslatorFactory implements TranslatorFactory, Serializable {
+
+    private static final long serialVersionUID = 1L;
+
+    private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>();
+
+    static {
+        SUPPORTED_TYPES.add(new Pair<>(StringPair.class, float[].class));
+        SUPPORTED_TYPES.add(new Pair<>(StringPair[].class, float[][].class));
+        SUPPORTED_TYPES.add(new Pair<>(Input.class, Output.class));
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Set<Pair<Type, Type>> getSupportedTypes() {
+        return SUPPORTED_TYPES;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments)
+            throws TranslateException {
+        Path modelPath = model.getModelPath();
+        try {
+            HuggingFaceTokenizer tokenizer =
+                    HuggingFaceTokenizer.builder(arguments)
+                            .optTokenizerPath(modelPath)
+                            .optManager(model.getNDManager())
+                            .build();
+            CrossEncoderTranslator translator =
+                    CrossEncoderTranslator.builder(tokenizer, arguments).build();
+            if (input == StringPair.class && output == float[].class) {
+                return (Translator<I, O>) translator;
+            } else if (input == StringPair[].class && output == float[][].class) {
+                return (Translator<I, O>) translator.toBatchTranslator();
+            } else if (input == Input.class && output == Output.class) {
+                return (Translator<I, O>) new CrossEncoderServingTranslator(translator);
+            }
+            throw new IllegalArgumentException("Unsupported input/output types.");
+        } catch (IOException e) {
+            throw new TranslateException("Failed to load tokenizer.", e);
+        }
+    }
+}


Test patch:
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/CrossEncoderTranslatorTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/CrossEncoderTranslatorTest.java
new file mode 100644
index 000000000..f3ee102e3
--- /dev/null
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/CrossEncoderTranslatorTest.java
@@ -0,0 +1,204 @@
+/*
+ * Copyright 2023 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.Model;
+import ai.djl.ModelException;
+import ai.djl.huggingface.translator.CrossEncoderTranslatorFactory;
+import ai.djl.inference.Predictor;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.nn.Block;
+import ai.djl.nn.LambdaBlock;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.ZooModel;
+import ai.djl.translate.TranslateException;
+import ai.djl.util.JsonUtils;
+import ai.djl.util.StringPair;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.HashMap;
+import java.util.Map;
+
+public class CrossEncoderTranslatorTest {
+
+    @Test
+    public void testCrossEncoderTranslator()
+            throws ModelException, IOException, TranslateException {
+        String text1 = "Sentence 1";
+        String text2 = "Sentence 2";
+        Block block =
+                new LambdaBlock(
+                        a -> {
+                            NDManager manager = a.getManager();
+                            NDArray array = manager.create(new float[] {-0.7329f});
+                            return new NDList(array);
+                        },
+                        "model");
+        Path modelDir = Paths.get("build/model");
+        Files.createDirectories(modelDir);
+
+        Criteria<StringPair, float[]> criteria =
+                Criteria.builder()
+                        .setTypes(StringPair.class, float[].class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-cased")
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new CrossEncoderTranslatorFactory())
+                        .build();
+
+        try (ZooModel<StringPair, float[]> model = criteria.loadModel();
+                Predictor<StringPair, float[]> predictor = model.newPredictor()) {
+            StringPair input = new StringPair(text1, text2);
+            float[] res = predictor.predict(input);
+            Assert.assertEquals(res[0], 0.32456556f, 0.0001);
+        }
+
+        Criteria<Input, Output> criteria2 =
+                Criteria.builder()
+                        .setTypes(Input.class, Output.class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-cased")
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new CrossEncoderTranslatorFactory())
+                        .build();
+
+        try (ZooModel<Input, Output> model = criteria2.loadModel();
+                Predictor<Input, Output> predictor = model.newPredictor()) {
+            Input input = new Input();
+            input.add("key", text1);
+            input.add("value", text2);
+            Output res = predictor.predict(input);
+            float[] buf = (float[]) res.getData().getAsObject();
+            Assert.assertEquals(buf[0], 0.32455865, 0.0001);
+
+            Assert.assertThrows(TranslateException.class, () -> predictor.predict(new Input()));
+
+            Assert.assertThrows(
+                    TranslateException.class,
+                    () -> {
+                        Input req = new Input();
+                        req.add("something", "false");
+                        predictor.predict(req);
+                    });
+
+            Assert.assertThrows(
+                    TranslateException.class,
+                    () -> {
+                        Input req = new Input();
+                        req.addProperty("Content-Type", "application/json");
+                        req.add("Invalid json");
+                        predictor.predict(req);
+                    });
+
+            Assert.assertThrows(
+                    TranslateException.class,
+                    () -> {
+                        Input req = new Input();
+                        req.addProperty("Content-Type", "application/json");
+                        req.add(JsonUtils.GSON.toJson(new StringPair(text1, null)));
+                        predictor.predict(req);
+                    });
+        }
+
+        try (Model model = Model.newInstance("test")) {
+            model.setBlock(block);
+            Map<String, String> options = new HashMap<>();
+            options.put("hasParameter", "false");
+            model.load(modelDir, "test", options);
+
+            CrossEncoderTranslatorFactory factory = new CrossEncoderTranslatorFactory();
+            Map<String, String> arguments = new HashMap<>();
+
+            Assert.assertThrows(
+                    TranslateException.class,
+                    () -> factory.newInstance(String.class, Integer.class, model, arguments));
+
+            arguments.put("tokenizer", "bert-base-cased");
+
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> factory.newInstance(String.class, Integer.class, model, arguments));
+        }
+    }
+
+    @Test
+    public void testCrossEncoderBatchTranslator()
+            throws ModelException, IOException, TranslateException {
+        StringPair pair1 = new StringPair("Sentence 1", "Sentence 2");
+        StringPair pair2 = new StringPair("Sentence 3", "Sentence 4");
+
+        Block block =
+                new LambdaBlock(
+                        a -> {
+                            NDManager manager = a.getManager();
+                            NDArray array = manager.create(new float[][] {{-0.7329f}, {-0.7329f}});
+                            return new NDList(array);
+                        },
+                        "model");
+        Path modelDir = Paths.get("build/model");
+        Files.createDirectories(modelDir);
+
+        Criteria<StringPair[], float[][]> criteria =
+                Criteria.builder()
+                        .setTypes(StringPair[].class, float[][].class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-cased")
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new CrossEncoderTranslatorFactory())
+                        .build();
+
+        try (ZooModel<StringPair[], float[][]> model = criteria.loadModel();
+                Predictor<StringPair[], float[][]> predictor = model.newPredictor()) {
+            StringPair[] inputs = {pair1, pair2};
+            float[][] res = predictor.predict(inputs);
+            Assert.assertEquals(res[1][0], 0.32455865, 0.0001);
+        }
+
+        Criteria<Input, Output> criteria2 =
+                Criteria.builder()
+                        .setTypes(Input.class, Output.class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-cased")
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new CrossEncoderTranslatorFactory())
+                        .build();
+
+        try (ZooModel<Input, Output> model = criteria2.loadModel();
+                Predictor<Input, Output> predictor = model.newPredictor()) {
+            Input input = new Input();
+            input.add(JsonUtils.GSON.toJson(new StringPair[] {pair1, pair2}));
+            input.addProperty("Content-Type", "application/json");
+            Output out = predictor.predict(input);
+            float[][] buf = (float[][]) out.getData().getAsObject();
+            Assert.assertEquals(buf[0][0], 0.32455865, 0.0001);
+        }
+    }
+}