Repository: deepjavalibrary/djl
Pull request #3400 (2024-08-08 22:24:38)

[api] Support TEI input format to reranking 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/api/src/main/java/ai/djl/modality/nlp/TextPrompt.java b/api/src/main/java/ai/djl/modality/nlp/TextPrompt.java
index 2793fea25..202e15e59 100644
--- a/api/src/main/java/ai/djl/modality/nlp/TextPrompt.java
+++ b/api/src/main/java/ai/djl/modality/nlp/TextPrompt.java
@@ -71,6 +71,12 @@ public final class TextPrompt {
      */
     public static TextPrompt parseInput(Input input) throws TranslateException {
         String contentType = input.getProperty("Content-Type", null);
+        if (contentType != null) {
+            int pos = contentType.indexOf(';');
+            if (pos > 0) {
+                contentType = contentType.substring(0, pos);
+            }
+        }
         String text = input.getData().getAsString();
         if (!"application/json".equals(contentType)) {
             return new TextPrompt(text);
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 75338df45..6c92a3c39 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
@@ -25,18 +25,17 @@ import ai.djl.util.JsonUtils;
 import ai.djl.util.PairList;
 import ai.djl.util.StringPair;
 
+import com.google.gson.JsonArray;
 import com.google.gson.JsonElement;
+import com.google.gson.JsonObject;
 import com.google.gson.JsonParseException;
-import com.google.gson.reflect.TypeToken;
 
-import java.lang.reflect.Type;
+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> {
 
-    private static final Type LIST_TYPE = new TypeToken<List<StringPair>>() {}.getType();
-
     private Translator<StringPair, float[]> translator;
 
     /**
@@ -63,31 +62,65 @@ public class CrossEncoderServingTranslator implements NoBatchifyTranslator<Input
         }
 
         String contentType = input.getProperty("Content-Type", null);
-        StringPair pair;
+        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);
-                    List<StringPair> inputs = JsonUtils.GSON.fromJson(json, LIST_TYPE);
+                    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);
-                }
-
-                pair = JsonUtils.GSON.fromJson(json, StringPair.class);
-                if (pair.getKey() == null || pair.getValue() == null) {
-                    throw new TranslateException("Missing key or value in json.");
+                } 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) {
-                throw new TranslateException("Missing key or value in input.");
+            if (key != null && value != null) {
+                pair = new StringPair(key, value);
             }
-            pair = new StringPair(key, value);
+        }
+
+        if (pair == null) {
+            throw new TranslateException("Missing key or value in input.");
         }
 
         NDList ret = translator.processInput(ctx, pair);
@@ -115,4 +148,18 @@ 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());
+        }
+        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.");
+    }
 }
diff --git a/api/src/main/java/ai/djl/modality/nlp/translator/QaServingTranslator.java b/api/src/main/java/ai/djl/modality/nlp/translator/QaServingTranslator.java
index 4abc0a1b5..45e3f59af 100644
--- a/api/src/main/java/ai/djl/modality/nlp/translator/QaServingTranslator.java
+++ b/api/src/main/java/ai/djl/modality/nlp/translator/QaServingTranslator.java
@@ -65,6 +65,12 @@ public class QaServingTranslator implements NoBatchifyTranslator<Input, Output>
         }
 
         String contentType = input.getProperty("Content-Type", null);
+        if (contentType != null) {
+            int pos = contentType.indexOf(';');
+            if (pos > 0) {
+                contentType = contentType.substring(0, pos);
+            }
+        }
         QAInput qa;
         if ("application/json".equals(contentType)) {
             String json = input.getData().getAsString();


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 14336182d..b59e840e7 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
@@ -38,11 +38,13 @@ import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 public class CrossEncoderTranslatorTest {
 
     @Test
+    @SuppressWarnings("unchecked")
     public void testCrossEncoderTranslator()
             throws ModelException, IOException, TranslateException {
         String text1 = "Sentence 1";
@@ -119,6 +121,48 @@ public class CrossEncoderTranslatorTest {
             float[] buf = (float[]) res.getData().getAsObject();
             Assert.assertEquals(buf[0], 0.32455865, 0.0001);
 
+            input = new Input();
+            input.add("text", text1);
+            input.add("text_pair", text2);
+            res = predictor.predict(input);
+            buf = (float[]) res.getData().getAsObject();
+            Assert.assertEquals(buf[0], 0.32455865, 0.0001);
+
+            input = new Input();
+            input.addProperty("Content-Type", "application/json; charset=utf-8");
+            input.add("data", "{\"text\": \"" + text1 + "\", \"text_pair\": \"" + text2 + "\"}");
+            res = predictor.predict(input);
+            buf = (float[]) res.getData().getAsObject();
+            Assert.assertEquals(buf[0], 0.32455865, 0.0001);
+
+            input = new Input();
+            input.addProperty("Content-Type", "application/json; charset=utf-8");
+            input.add("data", "{\"key\": \"" + text1 + "\", \"value\": \"" + text2 + "\"}");
+            res = predictor.predict(input);
+            buf = (float[]) res.getData().getAsObject();
+            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);
+            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);
+            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);
+            Assert.assertEquals(buf[0], 0.32455865, 0.0001);
+
             Assert.assertThrows(TranslateException.class, () -> predictor.predict(new Input()));
 
             Assert.assertThrows(
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/QuestionAnsweringTranslatorTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/QuestionAnsweringTranslatorTest.java
index 14e030e76..acb283734 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/QuestionAnsweringTranslatorTest.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/QuestionAnsweringTranslatorTest.java
@@ -123,7 +123,7 @@ public class QuestionAnsweringTranslatorTest {
                     TranslateException.class,
                     () -> {
                         Input req = new Input();
-                        req.addProperty("Content-Type", "application/json");
+                        req.addProperty("Content-Type", "application/json; charset=utf-8");
                         req.add("Invalid json");
                         predictor.predict(req);
                     });
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java
index 3dddf301f..15a670860 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java
@@ -213,7 +213,7 @@ public class TextEmbeddingTranslatorTest {
             Map<String, String> map = new HashMap<>();
             map.put("inputs", text);
             input.add(JsonUtils.GSON.toJson(map));
-            input.addProperty("Content-Type", "application/json");
+            input.addProperty("Content-Type", "application/json; charset=utf-8");
             out = predictor.predict(input);
             res = (float[]) out.getData().getAsObject();
             Assert.assertEquals(res.length, 384);