Repository: deepjavalibrary/djl
Pull request #3579 (2025-01-14 05:23:54)

[api] Adds ZeroShotClassification support
## 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


lordofthejars: That's great to see this. Thank you very much

Patch:
diff --git a/api/src/main/java/ai/djl/Application.java b/api/src/main/java/ai/djl/Application.java
index 0f3d49d5d..0cc7ef409 100644
--- a/api/src/main/java/ai/djl/Application.java
+++ b/api/src/main/java/ai/djl/Application.java
@@ -97,6 +97,9 @@ public class Application {
             case "nlp/token_classification":
             case "token_classification":
                 return NLP.TOKEN_CLASSIFICATION;
+            case "nlp/zero_shot_classification":
+            case "zero_shot_classification":
+                return NLP.ZERO_SHOT_CLASSIFICATION;
             case "nlp/word_embedding":
             case "word_embedding":
                 return NLP.WORD_EMBEDDING;
@@ -280,6 +283,12 @@ public class Application {
          */
         Application SENTIMENT_ANALYSIS = new Application("nlp/sentiment_analysis");
 
+        /**
+         * An application that classifies text into arbitrary label, a specific case of {@link
+         * #TEXT_CLASSIFICATION}.
+         */
+        Application ZERO_SHOT_CLASSIFICATION = new Application("nlp/zero_shot_classification");
+
         /**
          * A natural language understanding application that assigns a label to some tokens in a
          * text.
diff --git a/api/src/main/java/ai/djl/inference/Predictor.java b/api/src/main/java/ai/djl/inference/Predictor.java
index 5f6edb424..c5e850895 100644
--- a/api/src/main/java/ai/djl/inference/Predictor.java
+++ b/api/src/main/java/ai/djl/inference/Predictor.java
@@ -144,6 +144,9 @@ public class Predictor<I, O> implements AutoCloseable {
     protected NDList predictInternal(TranslatorContext ctx, NDList ndList)
             throws TranslateException {
         logger.trace("Predictor input data: {}", ndList);
+        if (ndList.isEmpty()) {
+            return new NDList();
+        }
         return block.forward(parameterStore, ndList, false);
     }
 
diff --git a/api/src/main/java/ai/djl/modality/nlp/translator/ZeroShotClassificationInput.java b/api/src/main/java/ai/djl/modality/nlp/translator/ZeroShotClassificationInput.java
new file mode 100644
index 000000000..c2345fc13
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/translator/ZeroShotClassificationInput.java
@@ -0,0 +1,124 @@
+/*
+ * 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.translate.TranslateException;
+import ai.djl.util.JsonUtils;
+
+import com.google.gson.JsonParseException;
+import com.google.gson.annotations.SerializedName;
+
+/** A class that represents a {@code ZeroShotClassificationInput} object. */
+public class ZeroShotClassificationInput {
+
+    private String text;
+
+    @SerializedName("candidate_labels")
+    private String[] candidates;
+
+    @SerializedName("multi_label")
+    private boolean multiLabel;
+
+    @SerializedName("hypothesis_template")
+    private String hypothesisTemplate;
+
+    /**
+     * Constructs a new {@code ZeroShotClassificationInput} instance.
+     *
+     * @param text the text to classify
+     * @param candidates the candidate labels
+     */
+    public ZeroShotClassificationInput(String text, String[] candidates) {
+        this(text, candidates, false);
+    }
+
+    /**
+     * Constructs a new {@code ZeroShotClassificationInput} instance.
+     *
+     * @param text the text to classify
+     * @param candidates the candidate labels
+     * @param multiLabel true to classify multiple labels
+     */
+    public ZeroShotClassificationInput(String text, String[] candidates, boolean multiLabel) {
+        this(text, candidates, multiLabel, null);
+    }
+
+    /**
+     * Constructs a new {@code ZeroShotClassificationInput} instance.
+     *
+     * @param text the text to classify
+     * @param candidates the candidate labels
+     * @param multiLabel true to classify multiple labels
+     * @param hypothesisTemplate the custom template
+     */
+    public ZeroShotClassificationInput(
+            String text, String[] candidates, boolean multiLabel, String hypothesisTemplate) {
+        this.text = text;
+        this.candidates = candidates;
+        this.multiLabel = multiLabel;
+        this.hypothesisTemplate = hypothesisTemplate;
+    }
+
+    /**
+     * Returns the {@code ZeroShotClassificationInput} from the {@link Input}.
+     *
+     * @param input the input object
+     * @return the {@code ZeroShotClassificationInput} from the {@link Input}
+     * @throws TranslateException if the input is invalid
+     */
+    public static ZeroShotClassificationInput parseInput(Input input) throws TranslateException {
+        String text = input.getData().getAsString();
+        try {
+            return JsonUtils.GSON.fromJson(text, ZeroShotClassificationInput.class);
+        } catch (JsonParseException e) {
+            throw new TranslateException("Input is not a valid json.", e);
+        }
+    }
+
+    /**
+     * Returns the text.
+     *
+     * @return the text to be classified
+     */
+    public String getText() {
+        return text;
+    }
+
+    /**
+     * Returns the candidate labels.
+     *
+     * @return the candidate labels
+     */
+    public String[] getCandidates() {
+        return candidates;
+    }
+
+    /**
+     * Returns {@code true} if to classify multiple labels.
+     *
+     * @return {@code true} if to classify multiple labels
+     */
+    public boolean isMultiLabel() {
+        return multiLabel;
+    }
+
+    /**
+     * Returns the custom template.
+     *
+     * @return the custom template
+     */
+    public String getHypothesisTemplate() {
+        return hypothesisTemplate == null ? "This example is {}." : hypothesisTemplate;
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/nlp/translator/ZeroShotClassificationOutput.java b/api/src/main/java/ai/djl/modality/nlp/translator/ZeroShotClassificationOutput.java
new file mode 100644
index 000000000..975f9e103
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/translator/ZeroShotClassificationOutput.java
@@ -0,0 +1,61 @@
+/*
+ * 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;
+
+/** A class that represents a {@code ZeroShotClassificationOutput} object. */
+public class ZeroShotClassificationOutput {
+
+    private String sequence;
+    private String[] labels;
+    private double[] scores;
+
+    /**
+     * Constructs a new {@code ZeroShotClassificationOutput} instance.
+     *
+     * @param sequence the input text
+     * @param labels the labels
+     * @param scores the scores of the labels
+     */
+    public ZeroShotClassificationOutput(String sequence, String[] labels, double[] scores) {
+        this.sequence = sequence;
+        this.labels = labels;
+        this.scores = scores;
+    }
+
+    /**
+     * Returns the input text.
+     *
+     * @return the input text
+     */
+    public String getSequence() {
+        return sequence;
+    }
+
+    /**
+     * Returns the labels in sorted order.
+     *
+     * @return the labels in sorted order
+     */
+    public String[] getLabels() {
+        return labels;
+    }
+
+    /**
+     * Returns the scores of the labels.
+     *
+     * @return the scores of the labels
+     */
+    public double[] getScores() {
+        return scores;
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/nlp/translator/ZeroShotClassificationServingTranslator.java b/api/src/main/java/ai/djl/modality/nlp/translator/ZeroShotClassificationServingTranslator.java
new file mode 100644
index 000000000..15d503fce
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/translator/ZeroShotClassificationServingTranslator.java
@@ -0,0 +1,79 @@
+/*
+ * 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.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;
+
+/**
+ * A {@link Translator} that can handle generic zero-shot-classification {@link Input} and {@link
+ * Output}.
+ */
+public class ZeroShotClassificationServingTranslator
+        implements NoBatchifyTranslator<Input, Output> {
+
+    private Translator<ZeroShotClassificationInput, ZeroShotClassificationOutput> translator;
+
+    /**
+     * Constructs a {@code TokenClassificationServingTranslator} instance.
+     *
+     * @param translator a {@code Translator} processes token classification input
+     */
+    public ZeroShotClassificationServingTranslator(
+            Translator<ZeroShotClassificationInput, ZeroShotClassificationOutput> 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.");
+        }
+
+        ZeroShotClassificationInput prompt = ZeroShotClassificationInput.parseInput(input);
+        NDList ret = translator.processInput(ctx, prompt);
+        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");
+        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/examples/src/main/java/ai/djl/examples/inference/nlp/ZeroShotClassification.java b/examples/src/main/java/ai/djl/examples/inference/nlp/ZeroShotClassification.java
new file mode 100644
index 000000000..0fc85a3fb
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/inference/nlp/ZeroShotClassification.java
@@ -0,0 +1,67 @@
+/*
+ * 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.ZeroShotClassificationTranslatorFactory;
+import ai.djl.inference.Predictor;
+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.translate.TranslateException;
+import ai.djl.util.JsonUtils;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+
+public final class ZeroShotClassification {
+
+    private static final Logger logger = LoggerFactory.getLogger(ZeroShotClassification.class);
+
+    private ZeroShotClassification() {}
+
+    public static void main(String[] args) throws ModelException, IOException, TranslateException {
+        ZeroShotClassificationOutput ret = predict(false);
+        logger.info("{}", JsonUtils.GSON_PRETTY.toJson(ret));
+
+        ret = predict(true);
+        logger.info("{}", JsonUtils.GSON_PRETTY.toJson(ret));
+    }
+
+    public static ZeroShotClassificationOutput predict(boolean multiLabels)
+            throws ModelException, IOException, TranslateException {
+        Criteria<ZeroShotClassificationInput, ZeroShotClassificationOutput> criteria =
+                Criteria.builder()
+                        .setTypes(
+                                ZeroShotClassificationInput.class,
+                                ZeroShotClassificationOutput.class)
+                        .optModelUrls("djl://ai.djl.huggingface.pytorch/facebook/bart-large-mnli")
+                        .optEngine("PyTorch")
+                        .optTranslatorFactory(new ZeroShotClassificationTranslatorFactory())
+                        .build();
+        String prompt = "one day I will see the world";
+        String[] candidates = {"travel", "cooking", "dancing", "exploration"};
+
+        try (ZooModel<ZeroShotClassificationInput, ZeroShotClassificationOutput> model =
+                        criteria.loadModel();
+                Predictor<ZeroShotClassificationInput, ZeroShotClassificationOutput> predictor =
+                        model.newPredictor()) {
+            ZeroShotClassificationInput input =
+                    new ZeroShotClassificationInput(prompt, candidates, multiLabels);
+            return predictor.predict(input);
+        }
+    }
+}
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
new file mode 100644
index 000000000..a8d651848
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/ZeroShotClassificationTranslator.java
@@ -0,0 +1,186 @@
+/*
+ * 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.Model;
+import ai.djl.ModelException;
+import ai.djl.huggingface.tokenizers.Encoding;
+import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
+import ai.djl.inference.Predictor;
+import ai.djl.modality.nlp.translator.ZeroShotClassificationInput;
+import ai.djl.modality.nlp.translator.ZeroShotClassificationOutput;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDArrays;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.index.NDIndex;
+import ai.djl.translate.ArgumentsUtil;
+import ai.djl.translate.Batchifier;
+import ai.djl.translate.NoBatchifyTranslator;
+import ai.djl.translate.NoopTranslator;
+import ai.djl.translate.TranslateException;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorContext;
+
+import java.io.IOException;
+import java.util.Map;
+import java.util.UUID;
+
+/** The translator for Huggingface zero-shot-classification model. */
+public class ZeroShotClassificationTranslator
+        implements NoBatchifyTranslator<ZeroShotClassificationInput, ZeroShotClassificationOutput> {
+
+    private HuggingFaceTokenizer tokenizer;
+    private boolean int32;
+    private Predictor<NDList, NDList> predictor;
+
+    ZeroShotClassificationTranslator(HuggingFaceTokenizer tokenizer, boolean int32) {
+        this.tokenizer = tokenizer;
+        this.int32 = int32;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void prepare(TranslatorContext ctx) throws IOException, ModelException {
+        Model model = ctx.getModel();
+        predictor = model.newPredictor(new NoopTranslator(null));
+        ctx.getPredictorManager().attachInternal(UUID.randomUUID().toString(), predictor);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, ZeroShotClassificationInput input) {
+        ctx.setAttachment("input", input);
+        return new NDList();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public ZeroShotClassificationOutput processOutput(TranslatorContext ctx, NDList list)
+            throws TranslateException {
+        ZeroShotClassificationInput input =
+                (ZeroShotClassificationInput) ctx.getAttachment("input");
+
+        String template = input.getHypothesisTemplate();
+        String[] candidates = input.getCandidates();
+        if (candidates == null || candidates.length == 0) {
+            throw new TranslateException("Missing candidates in input");
+        }
+
+        NDManager manager = ctx.getNDManager();
+        NDList output = new NDList(candidates.length);
+        for (String candidate : candidates) {
+            String hypothesis = applyTemplate(template, candidate);
+            Encoding encoding = tokenizer.encode(input.getText(), hypothesis);
+            NDList in = encoding.toNDList(manager, false, int32);
+            NDList batch = Batchifier.STACK.batchify(new NDList[] {in});
+            output.add(predictor.predict(batch).get(0));
+        }
+
+        NDArray logits = NDArrays.concat(output);
+        if (input.isMultiLabel()) {
+            logits = logits.get(":, -1");
+            logits = logits.softmax(-1);
+        } else {
+            logits = logits.get(new NDIndex(":, {}", manager.create(new int[] {0, 2})));
+            logits = logits.softmax(1);
+            logits = logits.get(":, -1");
+        }
+        long[] indices = logits.argSort(-1, false).toLongArray();
+        float[] probabilities = logits.toFloatArray();
+
+        String[] labels = new String[candidates.length];
+        double[] scores = new double[candidates.length];
+        for (int i = 0; i < labels.length; ++i) {
+            int index = (int) indices[i];
+            labels[i] = candidates[index];
+            scores[i] = probabilities[index];
+        }
+
+        return new ZeroShotClassificationOutput(input.getText(), labels, scores);
+    }
+
+    private String applyTemplate(String template, String arg) {
+        int pos = template.indexOf("{}");
+        if (pos == -1) {
+            return template + arg;
+        }
+        int len = template.length();
+        return template.substring(0, pos) + arg + template.substring(pos + 1, len);
+    }
+
+    /**
+     * Creates a builder to build a {@code ZeroShotClassificationTranslator}.
+     *
+     * @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 ZeroShotClassificationTranslator}.
+     *
+     * @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 zero-shot classification translator. */
+    public static final class Builder {
+
+        private HuggingFaceTokenizer tokenizer;
+        private boolean int32;
+
+        Builder(HuggingFaceTokenizer tokenizer) {
+            this.tokenizer = tokenizer;
+        }
+
+        /**
+         * Sets if use int32 datatype for the {@link Translator}.
+         *
+         * @param int32 true to include token types
+         * @return this builder
+         */
+        public Builder optInt32(boolean int32) {
+            this.int32 = int32;
+            return this;
+        }
+
+        /**
+         * Configures the builder with the model arguments.
+         *
+         * @param arguments the model arguments
+         */
+        public void configure(Map<String, ?> arguments) {
+            optInt32(ArgumentsUtil.booleanValue(arguments, "int32"));
+        }
+
+        /**
+         * Builds the translator.
+         *
+         * @return the new translator
+         * @throws IOException if I/O error occurs
+         */
+        public ZeroShotClassificationTranslator build() throws IOException {
+            return new ZeroShotClassificationTranslator(tokenizer, int32);
+        }
+    }
+}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/ZeroShotClassificationTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/ZeroShotClassificationTranslatorFactory.java
new file mode 100644
index 000000000..d314008ce
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/ZeroShotClassificationTranslatorFactory.java
@@ -0,0 +1,80 @@
+/*
+ * 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.Model;
+import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.nlp.translator.ZeroShotClassificationInput;
+import ai.djl.modality.nlp.translator.ZeroShotClassificationOutput;
+import ai.djl.modality.nlp.translator.ZeroShotClassificationServingTranslator;
+import ai.djl.translate.TranslateException;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorFactory;
+import ai.djl.util.Pair;
+
+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 ZeroShotClassificationTranslator} instance. */
+public class ZeroShotClassificationTranslatorFactory 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<>(ZeroShotClassificationInput.class, ZeroShotClassificationOutput.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();
+            ZeroShotClassificationTranslator translator =
+                    ZeroShotClassificationTranslator.builder(tokenizer, arguments).build();
+            if (input == ZeroShotClassificationInput.class
+                    && output == ZeroShotClassificationOutput.class) {
+                return (Translator<I, O>) translator;
+            } else if (input == Input.class && output == Output.class) {
+                return (Translator<I, O>) new ZeroShotClassificationServingTranslator(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/ZeroShotClassificationTranslatorTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/ZeroShotClassificationTranslatorTest.java
new file mode 100644
index 000000000..655733012
--- /dev/null
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/ZeroShotClassificationTranslatorTest.java
@@ -0,0 +1,141 @@
+/*
+ * 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.Model;
+import ai.djl.ModelException;
+import ai.djl.huggingface.translator.ZeroShotClassificationTranslatorFactory;
+import ai.djl.inference.Predictor;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.nlp.translator.ZeroShotClassificationInput;
+import ai.djl.modality.nlp.translator.ZeroShotClassificationOutput;
+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.testing.Assertions;
+import ai.djl.translate.TranslateException;
+import ai.djl.util.JsonUtils;
+
+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 ZeroShotClassificationTranslatorTest {
+
+    private static final float[][][] ARRAYS = {
+        {{-2.6497f, 0.8422f, 2.0454f}},
+        {{2.9208f, 0.1431f, -3.3346f}},
+        {{1.6723f, 1.0156f, -3.1728f}},
+        {{-1.795f, 0.983f, 0.9644f}}
+    };
+
+    private static int index;
+
+    @Test
+    public void testZeroShotClassificationTranslator()
+            throws ModelException, IOException, TranslateException {
+        String text = "one day I will see the world";
+        String[] candidates = {"travel", "cooking", "dancing", "exploration"};
+
+        Block block =
+                new LambdaBlock(
+                        a -> {
+                            NDManager manager = a.getManager();
+                            NDArray arr = manager.create(ARRAYS[index++ % 4]);
+                            return new NDList(arr);
+                        },
+                        "model");
+
+        Path modelDir = Paths.get("build/model");
+        Files.createDirectories(modelDir);
+
+        Criteria<ZeroShotClassificationInput, ZeroShotClassificationOutput> criteria =
+                Criteria.builder()
+                        .setTypes(
+                                ZeroShotClassificationInput.class,
+                                ZeroShotClassificationOutput.class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-uncased")
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new ZeroShotClassificationTranslatorFactory())
+                        .build();
+
+        try (ZooModel<ZeroShotClassificationInput, ZeroShotClassificationOutput> model =
+                        criteria.loadModel();
+                Predictor<ZeroShotClassificationInput, ZeroShotClassificationOutput> predictor =
+                        model.newPredictor()) {
+            ZeroShotClassificationInput input = new ZeroShotClassificationInput(text, candidates);
+            ZeroShotClassificationOutput res = predictor.predict(input);
+            Assertions.assertAlmostEquals(res.getScores()[1], 0.940441);
+            Assert.assertEquals(res.getLabels()[1], "exploration");
+        }
+
+        Criteria<Input, Output> criteria2 =
+                Criteria.builder()
+                        .setTypes(Input.class, Output.class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-uncased")
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new ZeroShotClassificationTranslatorFactory())
+                        .build();
+
+        try (ZooModel<Input, Output> model = criteria2.loadModel();
+                Predictor<Input, Output> predictor = model.newPredictor()) {
+            Input input = new Input();
+            ZeroShotClassificationInput zsc =
+                    new ZeroShotClassificationInput(text, candidates, true);
+            input.add(JsonUtils.toJson(zsc));
+            Output out = predictor.predict(input);
+            ZeroShotClassificationOutput res =
+                    (ZeroShotClassificationOutput) out.getData().getAsObject();
+            Assertions.assertAlmostEquals(res.getScores()[1], 0.251434);
+            Assert.assertEquals(res.getLabels()[0], "travel");
+        }
+
+        try (Model model = Model.newInstance("test")) {
+            model.setBlock(block);
+            Map<String, String> options = new HashMap<>();
+            options.put("hasParameter", "false");
+            model.load(modelDir, "test", options);
+
+            ZeroShotClassificationTranslatorFactory factory =
+                    new ZeroShotClassificationTranslatorFactory();
+            Map<String, String> arguments = new HashMap<>();
+
+            Assert.assertThrows(
+                    TranslateException.class,
+                    () -> factory.newInstance(String.class, Integer.class, model, arguments));
+
+            arguments.put("tokenizer", "bert-base-uncased");
+
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> factory.newInstance(String.class, Integer.class, model, arguments));
+        }
+    }
+}