Repository: deepjavalibrary/djl
Pull request #3459 (2024-09-08 20:07:25)

[api] Optimized text embedding post processing performance
## 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/api/src/main/java/ai/djl/modality/nlp/translator/CrossEncoderServingTranslator.java b/api/src/main/java/ai/djl/modality/nlp/translator/CrossEncoderServingTranslator.java
index 6c92a3c39..394d0173c 100644
--- a/api/src/main/java/ai/djl/modality/nlp/translator/CrossEncoderServingTranslator.java
+++ b/api/src/main/java/ai/djl/modality/nlp/translator/CrossEncoderServingTranslator.java
@@ -17,7 +17,6 @@ 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;
@@ -34,7 +33,7 @@ import java.util.ArrayList;
 import java.util.List;
 
 /** A {@link Translator} that can handle generic cross encoder {@link Input} and {@link Output}. */
-public class CrossEncoderServingTranslator implements NoBatchifyTranslator<Input, Output> {
+public class CrossEncoderServingTranslator implements Translator<Input, Output> {
 
     private Translator<StringPair, float[]> translator;
 
@@ -56,74 +55,13 @@ public class CrossEncoderServingTranslator implements NoBatchifyTranslator<Input
     /** {@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.");
+        ReRankingInput in = ReRankingInput.parseInput(input);
+        if (in.batch != null) {
+            ctx.setAttachment("batch", Boolean.TRUE);
+            return translator.batchProcessInput(ctx, in.batch);
         }
 
-        String contentType = input.getProperty("Content-Type", null);
-        if (contentType != null) {
-            int pos = contentType.indexOf(';');
-            if (pos > 0) {
-                contentType = contentType.substring(0, pos);
-            }
-        }
-        StringPair pair = null;
-        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);
-                    JsonArray array = element.getAsJsonArray();
-                    int size = array.size();
-                    List<StringPair> inputs = new ArrayList<>(size);
-                    for (int i = 0; i < size; ++i) {
-                        JsonObject obj = array.get(i).getAsJsonObject();
-                        inputs.add(parseStringPair(obj));
-                    }
-                    return translator.batchProcessInput(ctx, inputs);
-                } else if (element.isJsonObject()) {
-                    JsonObject obj = element.getAsJsonObject();
-                    JsonElement query = obj.get("query");
-                    if (query != null) {
-                        String key = query.getAsString();
-                        JsonArray texts = obj.get("texts").getAsJsonArray();
-                        int size = texts.size();
-                        List<StringPair> inputs = new ArrayList<>(size);
-                        for (int i = 0; i < size; ++i) {
-                            String value = texts.get(i).getAsString();
-                            inputs.add(new StringPair(key, value));
-                        }
-                        ctx.setAttachment("batch", Boolean.TRUE);
-                        return translator.batchProcessInput(ctx, inputs);
-                    } else {
-                        pair = parseStringPair(obj);
-                    }
-                } else {
-                    throw new TranslateException("Unexpected json type");
-                }
-            } catch (JsonParseException e) {
-                throw new TranslateException("Input is not a valid json.", e);
-            }
-        } else {
-            String text = input.getAsString("text");
-            String textPair = input.getAsString("text_pair");
-            if (text != null && textPair != null) {
-                pair = new StringPair(text, textPair);
-            }
-            String key = input.getAsString("key");
-            String value = input.getAsString("value");
-            if (key != null && value != null) {
-                pair = new StringPair(key, value);
-            }
-        }
-
-        if (pair == null) {
-            throw new TranslateException("Missing key or value in input.");
-        }
-
-        NDList ret = translator.processInput(ctx, pair);
+        NDList ret = translator.processInput(ctx, in.pair);
         Batchifier batchifier = translator.getBatchifier();
         if (batchifier != null) {
             NDList[] batch = {ret};
@@ -132,6 +70,27 @@ public class CrossEncoderServingTranslator implements NoBatchifyTranslator<Input
         return ret;
     }
 
+    /** {@inheritDoc} */
+    @Override
+    @SuppressWarnings("PMD.SignatureDeclareThrowsException")
+    public NDList batchProcessInput(TranslatorContext ctx, List<Input> inputs) throws Exception {
+        int[] mapping = new int[inputs.size()];
+        List<StringPair> prompts = new ArrayList<>(mapping.length);
+        for (int i = 0; i < mapping.length; ++i) {
+            ReRankingInput in = ReRankingInput.parseInput(inputs.get(i));
+            if (in.batch != null) {
+                List<StringPair> batch = in.batch;
+                mapping[i] = batch.size();
+                prompts.addAll(batch);
+            } else {
+                mapping[i] = -1;
+                prompts.add(in.pair);
+            }
+        }
+        ctx.setAttachment("mapping", mapping);
+        return translator.batchProcessInput(ctx, prompts);
+    }
+
     /** {@inheritDoc} */
     @Override
     public Output processOutput(TranslatorContext ctx, NDList list) throws Exception {
@@ -149,17 +108,126 @@ public class CrossEncoderServingTranslator implements NoBatchifyTranslator<Input
         return output;
     }
 
-    private StringPair parseStringPair(JsonObject json) throws TranslateException {
-        JsonElement text = json.get("text");
-        JsonElement textPair = json.get("text_pair");
-        if (text != null && textPair != null) {
-            return new StringPair(text.getAsString(), textPair.getAsString());
+    /** {@inheritDoc} */
+    @Override
+    @SuppressWarnings("PMD.SignatureDeclareThrowsException")
+    public List<Output> batchProcessOutput(TranslatorContext ctx, NDList list) throws Exception {
+        List<float[]> outputs = translator.batchProcessOutput(ctx, list);
+        int[] mapping = (int[]) ctx.getAttachment("mapping");
+        List<Output> ret = new ArrayList<>(mapping.length);
+        int index = 0;
+        for (int size : mapping) {
+            Output output = new Output();
+            output.addProperty("Content-Type", "application/json");
+            if (size == -1) {
+                // non-batching
+                output.add(BytesSupplier.wrapAsJson(outputs.get(index++)));
+            } else {
+                // client side batching
+                float[][] embeddings = new float[size][];
+                for (int j = 0; j < size; ++j) {
+                    embeddings[j] = outputs.get(index++);
+                }
+                output.add(BytesSupplier.wrapAsJson(embeddings));
+            }
+            ret.add(output);
+        }
+        return ret;
+    }
+
+    private static final class ReRankingInput {
+
+        private StringPair pair;
+        private List<StringPair> batch;
+
+        ReRankingInput(StringPair pair) {
+            this.pair = pair;
+        }
+
+        ReRankingInput(List<StringPair> batch) {
+            this.batch = batch;
         }
-        JsonElement key = json.get("key");
-        JsonElement value = json.get("value");
-        if (key != null && value != null) {
-            return new StringPair(key.getAsString(), value.getAsString());
+
+        static ReRankingInput parseInput(Input input) throws TranslateException {
+            PairList<String, BytesSupplier> content = input.getContent();
+            if (content.isEmpty()) {
+                throw new TranslateException("Input data is empty.");
+            }
+
+            String contentType = input.getProperty("Content-Type", null);
+            if (contentType != null) {
+                int pos = contentType.indexOf(';');
+                if (pos > 0) {
+                    contentType = contentType.substring(0, pos);
+                }
+            }
+            StringPair pair = null;
+            if ("application/json".equals(contentType)) {
+                String json = input.getData().getAsString();
+                try {
+                    JsonElement element = JsonUtils.GSON.fromJson(json, JsonElement.class);
+                    if (element.isJsonArray()) {
+                        JsonArray array = element.getAsJsonArray();
+                        int size = array.size();
+                        List<StringPair> batch = new ArrayList<>(size);
+                        for (int i = 0; i < size; ++i) {
+                            JsonObject obj = array.get(i).getAsJsonObject();
+                            batch.add(parseStringPair(obj));
+                        }
+                        return new ReRankingInput(batch);
+                    } else if (element.isJsonObject()) {
+                        JsonObject obj = element.getAsJsonObject();
+                        JsonElement query = obj.get("query");
+                        if (query != null) {
+                            String key = query.getAsString();
+                            JsonArray texts = obj.get("texts").getAsJsonArray();
+                            int size = texts.size();
+                            List<StringPair> batch = new ArrayList<>(size);
+                            for (int i = 0; i < size; ++i) {
+                                String value = texts.get(i).getAsString();
+                                batch.add(new StringPair(key, value));
+                            }
+                            return new ReRankingInput(batch);
+                        } else {
+                            pair = parseStringPair(obj);
+                        }
+                    } else {
+                        throw new TranslateException("Unexpected json type");
+                    }
+                } catch (JsonParseException e) {
+                    throw new TranslateException("Input is not a valid json.", e);
+                }
+            } else {
+                String text = input.getAsString("text");
+                String textPair = input.getAsString("text_pair");
+                if (text != null && textPair != null) {
+                    pair = new StringPair(text, textPair);
+                }
+                String key = input.getAsString("key");
+                String value = input.getAsString("value");
+                if (key != null && value != null) {
+                    pair = new StringPair(key, value);
+                }
+            }
+
+            if (pair == null) {
+                throw new TranslateException("Missing key or value in input.");
+            }
+            return new ReRankingInput(pair);
+        }
+
+        private static StringPair parseStringPair(JsonObject json) throws TranslateException {
+            JsonElement text = json.get("text");
+            JsonElement textPair = json.get("text_pair");
+            if (text != null && textPair != null) {
+                return new StringPair(text.getAsString(), textPair.getAsString());
+            }
+            JsonElement key = json.get("key");
+            JsonElement value = json.get("value");
+            if (key != null && value != null) {
+                return new StringPair(key.getAsString(), value.getAsString());
+            }
+            throw new TranslateException("Missing text or text_pair in json.");
         }
-        throw new TranslateException("Missing text or text_pair in json.");
     }
 }
diff --git a/api/src/main/java/ai/djl/modality/nlp/translator/TextEmbeddingServingTranslator.java b/api/src/main/java/ai/djl/modality/nlp/translator/TextEmbeddingServingTranslator.java
index 52dcc593c..fd0c723ad 100644
--- a/api/src/main/java/ai/djl/modality/nlp/translator/TextEmbeddingServingTranslator.java
+++ b/api/src/main/java/ai/djl/modality/nlp/translator/TextEmbeddingServingTranslator.java
@@ -107,7 +107,7 @@ public class TextEmbeddingServingTranslator implements Translator<Input, Output>
 
     /** {@inheritDoc} */
     @Override
-    @SuppressWarnings({"PMD.SignatureDeclareThrowsException", "unchecked"})
+    @SuppressWarnings("PMD.SignatureDeclareThrowsException")
     public List<Output> batchProcessOutput(TranslatorContext ctx, NDList list) throws Exception {
         List<float[]> outputs = translator.batchProcessOutput(ctx, list);
         int[] mapping = (int[]) ctx.getAttachment("mapping");
diff --git a/extensions/tokenizers/build.gradle.kts b/extensions/tokenizers/build.gradle.kts
index 30fd0e4dc..c4437da53 100644
--- a/extensions/tokenizers/build.gradle.kts
+++ b/extensions/tokenizers/build.gradle.kts
@@ -89,7 +89,7 @@ tasks {
                     downloadPath gzipInto file
                 }
 
-                if ("text_embedding" != task)
+                if (task !in arrayOf("text_embedding", "text_classification"))
                     continue
 
                 file = prefix / task / "ai.djl.huggingface.rust.json"
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
index 61663519a..ac5e4ed55 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslator.java
@@ -26,6 +26,7 @@ import ai.djl.util.StringPair;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
@@ -88,14 +89,29 @@ public class CrossEncoderTranslator implements Translator<StringPair, float[]> {
     /** {@inheritDoc} */
     @Override
     public List<float[]> batchProcessOutput(TranslatorContext ctx, NDList list) {
-        NDList[] batches = batchifier.unbatchify(list);
-        List<float[]> ret = new ArrayList<>(batches.length);
-        for (NDList batch : batches) {
-            NDArray result = batch.get(0);
-            if (sigmoid) {
+        if (sigmoid) {
+            NDList[] batches = batchifier.unbatchify(list);
+            List<float[]> ret = new ArrayList<>(batches.length);
+            for (NDList batch : batches) {
+                NDArray result = batch.get(0);
                 result = result.getNDArrayInternal().sigmoid();
+                ret.add(result.toFloatArray());
             }
-            ret.add(result.toFloatArray());
+            return ret;
+        }
+        NDArray array = list.get(0);
+        int batchSize = Math.toIntExact(array.size(0));
+        float[] buf = list.get(0).toFloatArray();
+        if (batchSize == 1) {
+            return Collections.singletonList(buf);
+        }
+
+        int length = buf.length / batchSize;
+        List<float[]> ret = new ArrayList<>(batchSize);
+        for (int i = 0; i < batchSize; ++i) {
+            float[] f = new float[length];
+            System.arraycopy(buf, i * length, f, 0, length);
+            ret.add(f);
         }
         return ret;
     }
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
index 8a9520561..0d91734d5 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
@@ -30,6 +30,7 @@ import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
@@ -155,14 +156,20 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
     /** {@inheritDoc} */
     @Override
     public List<float[]> batchProcessOutput(TranslatorContext ctx, NDList list) {
-        int batchSize = Math.toIntExact(list.head().size(0));
         NDArray attentionMask = (NDArray) ctx.getAttachment("attentionMask");
         NDArray output = processEmbedding(list, attentionMask);
+        int batchSize = Math.toIntExact(output.size(0));
+        float[] buf = output.toFloatArray();
+        if (batchSize == 1) {
+            return Collections.singletonList(buf);
+        }
+
+        int length = buf.length / batchSize;
         List<float[]> ret = new ArrayList<>(batchSize);
-        NDList splitList = output.split(batchSize);
-        for (int i = 0; i < batchSize; i++) {
-            NDArray array = splitList.get(i);
-            ret.add(array.toFloatArray());
+        for (int i = 0; i < batchSize; ++i) {
+            float[] f = new float[length];
+            System.arraycopy(buf, i * length, f, 0, length);
+            ret.add(f);
         }
         return ret;
     }
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index 752dd8d6a..0ceae7b4e 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -16,32 +16,32 @@ fasttext = "0.9.2"
 sentencepiece = "0.2.0"
 tokenizers = "0.20.0"
 
-slf4j = "2.0.13"
+slf4j = "2.0.16"
 log4j_slf4j = "2.23.1"
 gson = "2.11.0"
 jna = "5.14.0"
-commonCli = "1.8.0"
-commonsCompress = "1.26.2"
+commonCli = "1.9.0"
+commonsCompress = "1.27.1"
 commonsCsv = "1.11.0"
-commonsLogging = "1.3.3"
+commonsLogging = "1.3.4"
 
-awssdk = "2.26.19"
+awssdk = "2.27.18"
 ffmpeg = "6.1.1-1.5.10"
 hadoop = "3.4.0"
 javacpp = "1.5.10"
 javacv = "1.5.10"
 openpnpOpencv = "4.9.0-0"
-spark = "3.5.1"
+spark = "3.5.2"
 tablesaw = "0.43.1"
 rapis = "22.12.0"
-protobuf = "3.25.3"
+protobuf = "3.25.4"
 jackson = "2.17.2"
 jtransforms = "3.1"
 
 antlr = "4.11.1"
 testng = "7.10.2"
 junit = "4.13.2"
-mockito = "5.12.0"
+mockito = "5.13.0"
 
 [libraries]
 slf4j-api = { module = "org.slf4j:slf4j-api", version.ref = "slf4j" }


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 b59e840e7..e874e4380 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
@@ -37,6 +37,7 @@ import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -44,7 +45,6 @@ import java.util.Map;
 public class CrossEncoderTranslatorTest {
 
     @Test
-    @SuppressWarnings("unchecked")
     public void testCrossEncoderTranslator()
             throws ModelException, IOException, TranslateException {
         String text1 = "Sentence 1";
@@ -146,21 +146,21 @@ public class CrossEncoderTranslatorTest {
             input.addProperty("Content-Type", "application/json");
             input.add("data", "{\"query\": \"" + text1 + "\", \"texts\": [\"" + text2 + "\"]}");
             res = predictor.predict(input);
-            buf = ((List<float[]>) res.getData().getAsObject()).get(0);
+            buf = ((float[][]) res.getData().getAsObject())[0];
             Assert.assertEquals(buf[0], 0.32455865, 0.0001);
 
             input = new Input();
             input.addProperty("Content-Type", "application/json");
             input.add("data", "{\"query\": \"" + text1 + "\", \"texts\": [\"" + text2 + "\"]}");
             res = predictor.predict(input);
-            buf = ((List<float[]>) res.getData().getAsObject()).get(0);
+            buf = ((float[][]) res.getData().getAsObject())[0];
             Assert.assertEquals(buf[0], 0.32455865, 0.0001);
 
             input = new Input();
             input.addProperty("Content-Type", "application/json");
             input.add("data", "[{\"text\": \"" + text1 + "\", \"text_pair\": \"" + text2 + "\"}]");
             res = predictor.predict(input);
-            buf = ((List<float[]>) res.getData().getAsObject()).get(0);
+            buf = ((float[][]) res.getData().getAsObject())[0];
             Assert.assertEquals(buf[0], 0.32455865, 0.0001);
 
             Assert.assertThrows(TranslateException.class, () -> predictor.predict(new Input()));
@@ -213,4 +213,56 @@ public class CrossEncoderTranslatorTest {
                     () -> factory.newInstance(String.class, Integer.class, model, arguments));
         }
     }
+
+    @Test
+    public void testCrossEncoderTranslatorServingBatch()
+            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}, {-0.7329f}, {-0.7329f}});
+                            return new NDList(array);
+                        },
+                        "model");
+        Path modelDir = Paths.get("build/model");
+        Files.createDirectories(modelDir);
+
+        Criteria<Input, Output> criteria =
+                Criteria.builder()
+                        .setTypes(Input.class, Output.class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-cased")
+                        .optArgument("reranking", true)
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
+                        .build();
+
+        try (ZooModel<Input, Output> model = criteria.loadModel();
+                Predictor<Input, Output> predictor = model.newPredictor()) {
+            Input input1 = new Input();
+            input1.add("text", text1);
+            input1.add("text_pair", text2);
+
+            Input input2 = new Input();
+            input2.addProperty("Content-Type", "application/json; charset=utf-8");
+            input2.add(
+                    "data",
+                    "{\"query\": \"query\", \"texts\": [\"" + text1 + "\", \"" + text2 + "\"]}");
+            List<Input> batchInput = Arrays.asList(input1, input2);
+
+            List<Output> batchOutput = predictor.batchPredict(batchInput);
+            Assert.assertEquals(batchOutput.size(), 2);
+            float[] ret1 = (float[]) batchOutput.get(0).getData().getAsObject();
+            float[][] ret2 = (float[][]) batchOutput.get(1).getData().getAsObject();
+            Assert.assertEquals(ret1[0], 0.32455865, 0.0001);
+            Assert.assertEquals(ret2[1][0], 0.32455865, 0.0001);
+        }
+    }
 }