Repository: deepjavalibrary/djl
Pull request #3239 (2024-06-05 18:12:06)

[text-embedding] Remove CrossEncoderTranslatorFactory in favor of TextEmbeddingTranslatorFactory
## 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_s3_pytorch_android.yml b/.github/workflows/native_s3_pytorch_android.yml
index 8887704a9..b21600813 100644
--- a/.github/workflows/native_s3_pytorch_android.yml
+++ b/.github/workflows/native_s3_pytorch_android.yml
@@ -10,7 +10,6 @@ jobs:
       matrix:
         format: ["armeabi-v7a", "arm64-v8a", "x86" ,"x86_64"]
     env:
-      PYTORCH_VERSION: "2.1.2"
       NDK_VERSION: "21.1.6352462"
     steps:
       - uses: actions/checkout@v4
@@ -25,6 +24,7 @@ jobs:
           python-version: '3.x'
       - name: clone pytorch
         run: |
+          PYTORCH_VERSION=$(awk -F '=' '/pytorch/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)
           git clone https://github.com/pytorch/pytorch android_pytorch_tmp -b "v${PYTORCH_VERSION}" --recursive
       - name: install Python Dependencies
         run: pip install numpy ninja pyyaml mkl mkl-include setuptools cmake cffi typing_extensions
@@ -46,5 +46,6 @@ jobs:
           aws-region: us-east-2
       - name: Copy files to S3 with the AWS CLI
         run: |
+          PYTORCH_VERSION=$(awk -F '=' '/pytorch/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)
           aws s3 cp android_pytorch_tmp/build_android/${{ matrix.format }}_native.zip s3://djl-ai/publish/pytorch/${PYTORCH_VERSION}/android_native/
           aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/pytorch/${PYTORCH_VERSION}/android_native*"
diff --git a/.github/workflows/nightly_android.yml b/.github/workflows/nightly_android.yml
index b46528bc2..773eb8dea 100644
--- a/.github/workflows/nightly_android.yml
+++ b/.github/workflows/nightly_android.yml
@@ -10,7 +10,7 @@ on:
 jobs:
   build:
     if: github.repository == 'deepjavalibrary/djl'
-    runs-on: macos-latest
+    runs-on: ubuntu-latest
     strategy:
       matrix:
         api-level: [ 26 ]
diff --git a/android/gradle.properties b/android/gradle.properties
index 4c93af453..3d06462ba 100644
--- a/android/gradle.properties
+++ b/android/gradle.properties
@@ -17,5 +17,5 @@ org.gradle.jvmargs=-Xmx1536m
 android.useAndroidX=true
 # Automatically convert third-party libraries to use AndroidX
 android.enableJetifier=true
-djl_version=0.27.0
-pytorch_version=2.1.1
+djl_version=0.28.0
+pytorch_version=2.2.2
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
deleted file mode 100644
index b3b765bbf..000000000
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslatorFactory.java
+++ /dev/null
@@ -1,77 +0,0 @@
-/*
- * 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<>(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 == 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);
-        }
-    }
-}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java
index 1f8dd5a51..d91c35cf2 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java
@@ -16,11 +16,14 @@ 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.modality.nlp.translator.TextEmbeddingServingTranslator;
+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;
@@ -39,6 +42,7 @@ public class TextEmbeddingTranslatorFactory implements TranslatorFactory, Serial
 
     static {
         SUPPORTED_TYPES.add(new Pair<>(String.class, float[].class));
+        SUPPORTED_TYPES.add(new Pair<>(StringPair.class, float[].class));
         SUPPORTED_TYPES.add(new Pair<>(Input.class, Output.class));
     }
 
@@ -61,6 +65,17 @@ public class TextEmbeddingTranslatorFactory implements TranslatorFactory, Serial
                             .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.");
+            }
+
             TextEmbeddingTranslator translator =
                     TextEmbeddingTranslator.builder(tokenizer, arguments).build();
             if (input == String.class && output == float[].class) {
diff --git a/tools/gradle/android-formatter.gradle b/tools/gradle/android-formatter.gradle
index 6ade6d3e1..7ffe40638 100644
--- a/tools/gradle/android-formatter.gradle
+++ b/tools/gradle/android-formatter.gradle
@@ -20,11 +20,7 @@ import com.google.googlejavaformat.java.RemoveUnusedImports
 
 class JavaFormatterPlugin implements Plugin<Project> {
     void apply(Project project) {
-        tasks.register('formatJava') {
-            def resultFilePath = "build/formatJava-result.txt"
-            inputs.files(project.sourceSets*.allSource)
-            inputs.files(project.fileTree('generated-src'))
-            outputs.file(project.file(resultFilePath))
+        project.task('formatJava') {
             doLast {
                 if (project.getRootProject() == project) {
                     return
@@ -39,7 +35,6 @@ class JavaFormatterPlugin implements Plugin<Project> {
                     }
                 }
             }
-            project.file(resultFilePath).write('Success')
         }
     }
 }


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 2a98f63db..33cbd9bd5 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,7 +14,7 @@ package ai.djl.huggingface.tokenizers;
 
 import ai.djl.Model;
 import ai.djl.ModelException;
-import ai.djl.huggingface.translator.CrossEncoderTranslatorFactory;
+import ai.djl.huggingface.translator.TextEmbeddingTranslatorFactory;
 import ai.djl.inference.Predictor;
 import ai.djl.modality.Input;
 import ai.djl.modality.Output;
@@ -65,8 +65,9 @@ public class CrossEncoderTranslatorTest {
                         .optEngine("PyTorch")
                         .optArgument("tokenizer", "bert-base-cased")
                         .optArgument("tokenizerPath", modelDir)
+                        .optArgument("reranking", true)
                         .optOption("hasParameter", "false")
-                        .optTranslatorFactory(new CrossEncoderTranslatorFactory())
+                        .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
                         .build();
 
         try (ZooModel<StringPair, float[]> model = criteria.loadModel();
@@ -83,8 +84,9 @@ public class CrossEncoderTranslatorTest {
                         .optBlock(block)
                         .optEngine("PyTorch")
                         .optArgument("tokenizer", "bert-base-cased")
+                        .optArgument("reranking", true)
                         .optOption("hasParameter", "false")
-                        .optTranslatorFactory(new CrossEncoderTranslatorFactory())
+                        .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
                         .build();
 
         try (ZooModel<Input, Output> model = criteria2.loadModel();
@@ -131,7 +133,7 @@ public class CrossEncoderTranslatorTest {
             options.put("hasParameter", "false");
             model.load(modelDir, "test", options);
 
-            CrossEncoderTranslatorFactory factory = new CrossEncoderTranslatorFactory();
+            TextEmbeddingTranslatorFactory factory = new TextEmbeddingTranslatorFactory();
             Map<String, String> arguments = new HashMap<>();
 
             Assert.assertThrows(
@@ -139,6 +141,7 @@ public class CrossEncoderTranslatorTest {
                     () -> factory.newInstance(String.class, Integer.class, model, arguments));
 
             arguments.put("tokenizer", "bert-base-cased");
+            arguments.put("reranking", "true");
 
             Assert.assertThrows(
                     IllegalArgumentException.class,