Repository: deepjavalibrary/djl
Pull request #3338 (2024-07-14 08:01:22)

[tokenizer] Supports cross encoder for text classification model
## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here



Patch:
diff --git a/.github/workflows/native_jni_s3_pytorch.yml b/.github/workflows/native_jni_s3_pytorch.yml
index a84b2bc03..92840065d 100644
--- a/.github/workflows/native_jni_s3_pytorch.yml
+++ b/.github/workflows/native_jni_s3_pytorch.yml
@@ -272,7 +272,12 @@ jobs:
   build-pytorch-jni-aarch64:
     if: github.repository == 'deepjavalibrary/djl'
     runs-on: [ self-hosted, aarch64 ]
-    container: amazonlinux:2
+    container:
+      image: amazonlinux:2
+      env:
+        JAVA_HOME: /usr/lib/jvm/java-17-amazon-corretto.aarch64
+        AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+        AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
     timeout-minutes: 30
     needs: create-aarch64-runner
     steps:
@@ -288,22 +293,15 @@ jobs:
           ln -sf /usr/bin/gcc10-ar /usr/bin/ar
           ln -sf /usr/bin/cmake3 /usr/bin/cmake
           pip3 install awscli --upgrade
-      - uses: actions/checkout@v3
+      - uses: taiki-e/checkout-action@v1
       - name: Release JNI prep
         run: |
-          export JAVA_HOME=/usr/lib/jvm/java-17-amazon-corretto.aarch64
           export PATH=$PATH:$JAVA_HOME/bin
           PYTORCH_VERSION=${{ github.event.inputs.pt_version }}
           export PYTORCH_VERSION=${PYTORCH_VERSION:-$(awk -F '=' '/pytorch/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)}
           echo $PYTORCH_VERSION
           ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pprecxx11 -Ppt_version=$PYTORCH_VERSION
           ./gradlew -Pjni -Ppt_version=$PYTORCH_VERSION :integration:test "-Dai.djl.default_engine=PyTorch"
-      - name: Configure AWS Credentials
-        uses: aws-actions/configure-aws-credentials@v2
-        with:
-          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
-          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
-          aws-region: us-east-2
       - name: Copy files to S3 with the AWS CLI
         run: |
           PYTORCH_VERSION=${{ github.event.inputs.pt_version }}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslatorFactory.java
index 295f014c5..8d5b3f962 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslatorFactory.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslatorFactory.java
@@ -17,11 +17,14 @@ import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
 import ai.djl.modality.Classifications;
 import ai.djl.modality.Input;
 import ai.djl.modality.Output;
+import ai.djl.modality.nlp.translator.CrossEncoderServingTranslator;
 import ai.djl.modality.nlp.translator.TextClassificationServingTranslator;
+import ai.djl.translate.ArgumentsUtil;
 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;
@@ -40,6 +43,7 @@ public class TextClassificationTranslatorFactory implements TranslatorFactory, S
 
     static {
         SUPPORTED_TYPES.add(new Pair<>(String.class, Classifications.class));
+        SUPPORTED_TYPES.add(new Pair<>(StringPair.class, float[].class));
         SUPPORTED_TYPES.add(new Pair<>(Input.class, Output.class));
     }
 
@@ -62,6 +66,17 @@ public class TextClassificationTranslatorFactory implements TranslatorFactory, S
                             .optTokenizerPath(modelPath)
                             .optManager(model.getNDManager())
                             .build();
+            if (ArgumentsUtil.booleanValue(arguments, "reranking")) {
+                CrossEncoderTranslator translator =
+                        CrossEncoderTranslator.builder(tokenizer, arguments).build();
+                if (input == StringPair.class && output == float[].class) {
+                    return (Translator<I, O>) translator;
+                } else if (input == Input.class && output == Output.class) {
+                    return (Translator<I, O>) new CrossEncoderServingTranslator(translator);
+                }
+                throw new IllegalArgumentException("Unsupported input/output types.");
+            }
+
             TextClassificationTranslator translator =
                     TextClassificationTranslator.builder(tokenizer, arguments).build();
             if (input == String.class && output == Classifications.class) {


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
index 33cbd9bd5..14336182d 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/CrossEncoderTranslatorTest.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/CrossEncoderTranslatorTest.java
@@ -14,6 +14,7 @@ package ai.djl.huggingface.tokenizers;
 
 import ai.djl.Model;
 import ai.djl.ModelException;
+import ai.djl.huggingface.translator.TextClassificationTranslatorFactory;
 import ai.djl.huggingface.translator.TextEmbeddingTranslatorFactory;
 import ai.djl.inference.Predictor;
 import ai.djl.modality.Input;
@@ -77,6 +78,26 @@ public class CrossEncoderTranslatorTest {
             Assert.assertEquals(res[0], 0.32456556f, 0.0001);
         }
 
+        criteria =
+                Criteria.builder()
+                        .setTypes(StringPair.class, float[].class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-cased")
+                        .optArgument("tokenizerPath", modelDir)
+                        .optArgument("reranking", true)
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new TextClassificationTranslatorFactory())
+                        .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)