Repository: deepjavalibrary/djl
Pull request #3712 (2025-06-10 21:00:43)

Align ZeroShotClassificationTranslator with Hugging Face behavior
This refactor improves the ZeroShotClassificationTranslator implementation to better reflect the behavior of Hugging Face’s zero-shot-classification pipeline in Python. It replaces several hardcoded assumptions with dynamic logic that adapts based on model configuration.

### Motivation

The translator was optimized to work with the "djl://ai.djl.huggingface.pytorch/facebook/bart-large-mnli" model. When a different model was used, it would either output wrong predictions or throw errors altogether. 

Some models on Hugging Face define their label mapping in a non-standard order. Without reading and honoring label2id, the results returned by the translator may be incorrect or misleading. The previous implementation also failed on models that require token type IDs. These issues are addressed in this refactor.

The new implementation closely mirrors the logic used in the Hugging Face Transformers ZeroShotClassificationPipeline, improving correctness and interoperability with a wider range of models.

⸻

I created a repo to compare the outputs of the new refactored translator with the outputs of Python's Transformers library: https://github.com/raphaeldelio/deep-java-library-zero-shot-classification/

I compared the outputs of four different models and they were all consistent:

MoritzLaurer/DeBERTa-v3-large-mnli-fever-anli-ling-wanli
Facebook/bart-large-mnli    
MoritzLaurer/bge-m3-zeroshot-v2.0   
tasksource/ModernBERT-base-nli

### Summary of changes

1. Dynamic entailment and contradiction label IDs
	•	Previous implementation assumed fixed label indices (entailment = 2, contradiction = 0).
	•	This change parses the label2id map from config.json, if available, to dynamically determine the correct label positions.
	•	Fallback values are retained for backward compatibility.
	•	The builder now allows manually overriding these IDs if needed.

2. Support for token_type_ids
	•	Added inference of withTokenType based on type_vocab_size and model_type in the model config.
	•	Common transformer families like BERT, ALBERT, and DeBERTa use token type embeddings and require this flag.
	•	The builder supports overriding withTokenType explicitly.

3. Updated multi-label and single-label scoring logic
	•	Multi-label classification now applies a softmax over the [contradiction, entailment] logits for each candidate and returns the entailment probability. This matches Hugging Face’s logic.
	•	Single-label classification uses only the entailment logits, then applies softmax across all candidates to normalize scores. This avoids including irrelevant logits from the contradiction class.
	•	Sorting and label-score pairing were updated to correctly reflect probabilities and preserve candidate order.

4. Config-aware translator preparation
	•	During prepare(), the model’s config.json is parsed to extract relevant metadata (label mappings, type vocab size, model type).
	•	Failures in config parsing are logged but do not break inference.

5. Extended builder API

	•	The builder now supports:
	•	withTokenType(boolean)
	•	int32(boolean)
	•	entailmentId(int)
	•	contradictionId(int)
	•	This allows downstream users to override behavior explicitly when needed.

### Impact

	•	Existing behavior is preserved for models that follow standard label conventions.
	•	Compatibility is improved for models with custom label IDs and token type requirements.
	•	The new translator is more robust and predictable

raphaeldelio: > Thank you for your contribution.
> 
> Please follow our development guide line: https://docs.djl.ai/master/docs/development/development_guideline.html#development-guideline sepcifically to fix coding style and pass the build checks:
> 
> 1. format the java code with: `./gradlew fJ`
> 2. pass all checks: `./gradlew check`

Thank you for referencing the guidelines. I'd missed it. 

I ran the formatting commands, fixed the javadocs, and applied your suggestions.

Patch:
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 581ec809c..276bce38e 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
@@ -49,7 +49,8 @@ public final class ZeroShotClassification {
                         .setTypes(
                                 ZeroShotClassificationInput.class,
                                 ZeroShotClassificationOutput.class)
-                        .optModelUrls("djl://ai.djl.huggingface.pytorch/facebook/bart-large-mnli")
+                        .optModelUrls(
+                                "djl://ai.djl.huggingface.pytorch/MoritzLaurer/ModernBERT-large-zeroshot-v2.0")
                         .optEngine("PyTorch")
                         .optTranslatorFactory(new ZeroShotClassificationTranslatorFactory())
                         .optProgress(new ProgressBar())
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 7fb801bd2..5bfbc8243 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
@@ -29,10 +29,25 @@ 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 ai.djl.util.JsonUtils;
+import ai.djl.util.Pair;
+
+import com.google.gson.JsonElement;
+import com.google.gson.JsonObject;
+import com.google.gson.JsonParseException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
+import java.io.Reader;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.UUID;
 
@@ -40,21 +55,85 @@ import java.util.UUID;
 public class ZeroShotClassificationTranslator
         implements NoBatchifyTranslator<ZeroShotClassificationInput, ZeroShotClassificationOutput> {
 
+    private static final Logger logger =
+            LoggerFactory.getLogger(ZeroShotClassificationTranslator.class);
+
     private HuggingFaceTokenizer tokenizer;
+    private int entailmentId;
+    private int contradictionId;
+    private boolean tokenTypeId;
     private boolean int32;
     private Predictor<NDList, NDList> predictor;
 
-    ZeroShotClassificationTranslator(HuggingFaceTokenizer tokenizer, boolean int32) {
+    ZeroShotClassificationTranslator(
+            HuggingFaceTokenizer tokenizer, boolean tokenTypeId, boolean int32) {
         this.tokenizer = tokenizer;
+        this.tokenTypeId = tokenTypeId;
         this.int32 = int32;
     }
 
+    ZeroShotClassificationTranslator(
+            HuggingFaceTokenizer tokenizer,
+            boolean tokenTypeId,
+            boolean int32,
+            int entailmentId,
+            int contradictionId) {
+        this(tokenizer, tokenTypeId, int32);
+        this.entailmentId = entailmentId;
+        this.contradictionId = contradictionId;
+    }
+
     /** {@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);
+
+        Path configFile = model.getModelPath().resolve("config.json");
+        if (!Files.isRegularFile(configFile)) {
+            return;
+        }
+        try (Reader reader = Files.newBufferedReader(configFile)) {
+            JsonObject config = JsonUtils.GSON.fromJson(reader, JsonObject.class);
+            if (config.has("label2id")) {
+                JsonObject label2Id = config.getAsJsonObject("label2id");
+                for (Map.Entry<String, JsonElement> entry : label2Id.entrySet()) {
+                    String key = entry.getKey().toLowerCase(Locale.ROOT);
+                    int value = entry.getValue().getAsInt();
+                    if (key.startsWith("entail")) {
+                        entailmentId = value;
+                    } else if (key.startsWith("contra")) {
+                        contradictionId = value;
+                    }
+                }
+            }
+
+            boolean inferredWithTokenType = false; // Default assumption
+            if (config.has("type_vocab_size")) {
+                JsonElement typeVocabSizeObj = config.get("type_vocab_size");
+                if (typeVocabSizeObj.isJsonPrimitive()) {
+                    int typeVocabSize = typeVocabSizeObj.getAsInt();
+                    if (typeVocabSize > 1) {
+                        inferredWithTokenType = true;
+                    }
+                }
+            }
+
+            if (!inferredWithTokenType && config.has("model_type")) {
+                String modelType = config.get("model_type").getAsString().toLowerCase(Locale.ROOT);
+                if ("bert".equals(modelType)
+                        || "albert".equals(modelType)
+                        || "xlnet".equals(modelType)
+                        || modelType.startsWith("deberta")) {
+                    inferredWithTokenType = true;
+                }
+            }
+
+            tokenTypeId = inferredWithTokenType;
+        } catch (IOException | JsonParseException e) {
+            logger.error("Failed to read or parse config.json for label2id", e);
+        }
     }
 
     /** {@inheritDoc} */
@@ -82,32 +161,68 @@ public class ZeroShotClassificationTranslator
         for (String candidate : candidates) {
             String hypothesis = applyTemplate(template, candidate);
             Encoding encoding = tokenizer.encode(input.getText(), hypothesis);
-            NDList in = encoding.toNDList(manager, false, int32);
+            NDList in = encoding.toNDList(manager, tokenTypeId, 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];
+        NDArray combinedLogits = NDArrays.concat(output);
+
+        String[] finalLabels;
+        double[] finalScores;
+
+        if (input.isMultiLabel() || candidates.length == 1) {
+            NDArray entailmentScores;
+            if (combinedLogits.getShape().get(1) == 2) {
+                // Binary classification: [not entail, entail]
+                NDArray probs = combinedLogits.softmax(1);
+                entailmentScores = probs.get(":, " + entailmentId);
+            } else {
+                // 3-class NLI output (e.g., entailment, neutral, contradiction)
+                NDArray entailContrLogits =
+                        combinedLogits.get(
+                                new NDIndex(
+                                        ":, {}",
+                                        manager.create(new int[] {contradictionId, entailmentId})));
+                NDArray scoresProbs = entailContrLogits.softmax(1);
+                entailmentScores = scoresProbs.get(":, 1");
+            }
+
+            float[] floatScores = entailmentScores.toFloatArray();
+
+            List<Pair<Double, String>> pairs = new ArrayList<>();
+            for (int i = 0; i < floatScores.length; i++) {
+                Pair<Double, String> pair = new Pair<>((double) floatScores[i], candidates[i]);
+                pairs.add(pair);
+            }
+            pairs.sort(
+                    Comparator.comparingDouble((Pair<Double, String> e) -> e.getKey()).reversed());
+
+            finalLabels = new String[candidates.length];
+            finalScores = new double[candidates.length];
+            for (int i = 0; i < candidates.length; i++) {
+                finalLabels[i] = pairs.get(i).getValue();
+                finalScores[i] = pairs.get(i).getKey();
+            }
+        } else { // Single-label classification (len(candidate_labels) > 1 and not multi_label)
+            NDArray entailLogits = combinedLogits.get(":, " + entailmentId);
+            NDArray exp = entailLogits.exp();
+            NDArray sum = exp.sum();
+            NDArray normalizedScores = exp.div(sum); // Probabilities sum to 1 across candidates
+
+            long[] indices = normalizedScores.argSort(-1, false).toLongArray();
+            float[] probabilities = normalizedScores.toFloatArray();
+
+            finalLabels = new String[candidates.length];
+            finalScores = new double[candidates.length];
+            for (int i = 0; i < finalLabels.length; ++i) {
+                int index = (int) indices[i];
+                finalLabels[i] = candidates[index];
+                finalScores[i] = probabilities[index];
+            }
         }
 
-        return new ZeroShotClassificationOutput(input.getText(), labels, scores);
+        return new ZeroShotClassificationOutput(input.getText(), finalLabels, finalScores);
     }
 
     private String applyTemplate(String template, String arg) {
@@ -147,29 +262,68 @@ public class ZeroShotClassificationTranslator
     public static final class Builder {
 
         private HuggingFaceTokenizer tokenizer;
+        private boolean tokenTypeId;
         private boolean int32;
+        private int entailmentId = 2;
+        private int contradictionId;
 
         Builder(HuggingFaceTokenizer tokenizer) {
             this.tokenizer = tokenizer;
         }
 
         /**
-         * Sets if use int32 datatype for the {@link Translator}.
+         * Specifies whether to include token type IDs in the input tensors.
          *
-         * @param int32 true to include token types
-         * @return this builder
+         * @param tokenTypeId {@code true} to include token type IDs, {@code false} to omit them
+         * @return this builder instance for method chaining
+         */
+        public Builder optTokenTypeId(boolean tokenTypeId) {
+            this.tokenTypeId = tokenTypeId;
+            return this;
+        }
+
+        /**
+         * Specifies whether to use int32 as the data type for input token tensors.
+         *
+         * @param int32 {@code true} to use int32 inputs, {@code false} to use the default type
+         * @return this builder instance for method chaining
          */
         public Builder optInt32(boolean int32) {
             this.int32 = int32;
             return this;
         }
 
+        /**
+         * Optional: Set custom entailment ID if different from default (2). This value usually
+         * comes from the model's `config.json` `label2id` mapping.
+         *
+         * @param entailmentId The index for the 'entailment' label.
+         * @return this builder
+         */
+        public Builder optEntailmentId(int entailmentId) {
+            this.entailmentId = entailmentId;
+            return this;
+        }
+
+        /**
+         * Optional: Set custom contradiction ID if different from default (0). This value usually
+         * comes from the model's `config.json` `label2id` mapping.
+         *
+         * @param contradictionId The index for the 'contradiction' label.
+         * @return this builder
+         */
+        public Builder optContradictionId(int contradictionId) {
+            this.contradictionId = contradictionId;
+            return this;
+        }
+
         /**
          * Configures the builder with the model arguments.
          *
          * @param arguments the model arguments
          */
         public void configure(Map<String, ?> arguments) {
+            optTokenTypeId(ArgumentsUtil.booleanValue(arguments, "tokenTypeId"));
             optInt32(ArgumentsUtil.booleanValue(arguments, "int32"));
         }
 
@@ -177,9 +331,11 @@ public class ZeroShotClassificationTranslator
          * Builds the translator.
          *
          * @return the new translator
+         * @throws IOException if I/O error occurs
          */
-        public ZeroShotClassificationTranslator build() {
-            return new ZeroShotClassificationTranslator(tokenizer, int32);
+        public ZeroShotClassificationTranslator build() throws IOException {
+            return new ZeroShotClassificationTranslator(
+                    tokenizer, tokenTypeId, int32, entailmentId, contradictionId);
         }
     }
 }


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
index 548a9d309..edfa86ad9 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/ZeroShotClassificationTranslatorTest.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/ZeroShotClassificationTranslatorTest.java
@@ -89,7 +89,7 @@ public class ZeroShotClassificationTranslatorTest {
                         model.newPredictor()) {
             ZeroShotClassificationInput input = new ZeroShotClassificationInput(text, candidates);
             ZeroShotClassificationOutput res = predictor.predict(input);
-            Assertions.assertAlmostEquals(res.getScores()[1], 0.940441);
+            Assertions.assertAlmostEquals(res.getScores()[1], 0.251434);
             Assert.assertEquals(res.getLabels()[1], "exploration");
         }
 
@@ -113,7 +113,7 @@ public class ZeroShotClassificationTranslatorTest {
             Output out = predictor.predict(input);
             ZeroShotClassificationOutput res =
                     (ZeroShotClassificationOutput) out.getData().getAsObject();
-            Assertions.assertAlmostEquals(res.getScores()[1], 0.251434);
+            Assertions.assertAlmostEquals(res.getScores()[1], 0.940441);
             Assert.assertEquals(res.getLabels()[0], "travel");
         }