Repository: deepjavalibrary/djl
Pull request #1983 (2022-09-02 07:00:22)

Adds TextClassificationTranslator


Patch:
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/PretrainedConfig.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/PretrainedConfig.java
new file mode 100644
index 000000000..c873dcd4c
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/PretrainedConfig.java
@@ -0,0 +1,26 @@
+/*
+ * Copyright 2022 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 com.google.gson.annotations.SerializedName;
+
+import java.util.Map;
+
+/** A class representing HuggingFace config.json file. */
+public class PretrainedConfig {
+
+    @SerializedName("problem_type")
+    String problemType;
+
+    Map<String, String> id2label;
+}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslator.java
new file mode 100644
index 000000000..a8cffd350
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslator.java
@@ -0,0 +1,165 @@
+/*
+ * Copyright 2022 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.huggingface.tokenizers.Encoding;
+import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
+import ai.djl.modality.Classifications;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.translate.ArgumentsUtil;
+import ai.djl.translate.Batchifier;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorContext;
+import ai.djl.util.JsonUtils;
+
+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.List;
+import java.util.Map;
+
+/** The translator for Huggingface text classification model. */
+public class TextClassificationTranslator implements Translator<String, Classifications> {
+
+    private HuggingFaceTokenizer tokenizer;
+    private Batchifier batchifier;
+    private PretrainedConfig config;
+
+    TextClassificationTranslator(HuggingFaceTokenizer tokenizer, Batchifier batchifier) {
+        this.tokenizer = tokenizer;
+        this.batchifier = batchifier;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Batchifier getBatchifier() {
+        return batchifier;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void prepare(TranslatorContext ctx) throws IOException {
+        Path path = ctx.getModel().getModelPath();
+        Path file = path.resolve("config.json");
+        try (Reader reader = Files.newBufferedReader(file)) {
+            config = JsonUtils.GSON.fromJson(reader, PretrainedConfig.class);
+        }
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, String input) {
+        NDManager manager = ctx.getNDManager();
+        Encoding encoding = tokenizer.encode(input);
+        ctx.setAttachment("encoding", encoding);
+        long[] indices = encoding.getIds();
+        long[] attentionMask = encoding.getAttentionMask();
+        NDList ndList = new NDList(2);
+        ndList.add(manager.create(indices));
+        ndList.add(manager.create(attentionMask));
+        return ndList;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Classifications processOutput(TranslatorContext ctx, NDList list) {
+        NDArray logits = list.get(0);
+        int size = config.id2label.size();
+        if ("multi_label_classification".equals(config.problemType) || size == 1) {
+            logits = logits.getNDArrayInternal().sigmoid();
+        } else if ("single_label_classification".equals(config.problemType) || size > 1) {
+            logits = logits.softmax(0);
+        }
+        long[] indices = logits.argSort(-1, false).toLongArray();
+        float[] buf = logits.toFloatArray();
+        List<String> classes = new ArrayList<>(size);
+        List<Double> probabilities = new ArrayList<>(size);
+        for (long l : indices) {
+            int index = Math.toIntExact(l);
+            classes.add(config.id2label.get(String.valueOf(index)));
+            probabilities.add((double) buf[index]);
+        }
+
+        return new Classifications(classes, probabilities);
+    }
+
+    /**
+     * Creates a builder to build a {@code TextClassificationTranslator}.
+     *
+     * @param tokenizer the tokenizer
+     * @return a new builder
+     */
+    public static Builder builder(HuggingFaceTokenizer tokenizer) {
+        return new Builder(tokenizer);
+    }
+
+    /**
+     * Creates a builder to build a {@code TextClassificationTranslator}.
+     *
+     * @param tokenizer the tokenizer
+     * @param arguments the models' arguments
+     * @return a new builder
+     */
+    public static Builder builder(HuggingFaceTokenizer tokenizer, Map<String, ?> arguments) {
+        Builder builder = builder(tokenizer);
+        builder.configure(arguments);
+
+        return builder;
+    }
+
+    /** The builder for token classification translator. */
+    public static final class Builder {
+
+        private HuggingFaceTokenizer tokenizer;
+        private Batchifier batchifier = Batchifier.STACK;
+
+        Builder(HuggingFaceTokenizer tokenizer) {
+            this.tokenizer = tokenizer;
+        }
+
+        /**
+         * Sets the {@link Batchifier} for the {@link Translator}.
+         *
+         * @param batchifier true to include token types
+         * @return this builder
+         */
+        public Builder optBatchifier(Batchifier batchifier) {
+            this.batchifier = batchifier;
+            return this;
+        }
+
+        /**
+         * Configures the builder with the model arguments.
+         *
+         * @param arguments the model arguments
+         */
+        public void configure(Map<String, ?> arguments) {
+            String batchifierStr = ArgumentsUtil.stringValue(arguments, "batchifier", "stack");
+            optBatchifier(Batchifier.fromString(batchifierStr));
+        }
+
+        /**
+         * Builds the translator.
+         *
+         * @return the new translator
+         * @throws IOException if I/O error occurs
+         */
+        public TextClassificationTranslator build() throws IOException {
+            return new TextClassificationTranslator(tokenizer, batchifier);
+        }
+    }
+}
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
new file mode 100644
index 000000000..a20031cac
--- /dev/null
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextClassificationTranslatorFactory.java
@@ -0,0 +1,73 @@
+/*
+ * Copyright 2022 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.Classifications;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.nlp.translator.TextClassificationServingTranslator;
+import ai.djl.translate.TranslateException;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorFactory;
+import ai.djl.util.Pair;
+
+import java.io.IOException;
+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 TextClassificationTranslator} instance. */
+public class TextClassificationTranslatorFactory implements TranslatorFactory {
+
+    private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>();
+
+    static {
+        SUPPORTED_TYPES.add(new Pair<>(String.class, Classifications.class));
+        SUPPORTED_TYPES.add(new Pair<>(Input.class, Output.class));
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Set<Pair<Type, Type>> getSupportedTypes() {
+        return SUPPORTED_TYPES;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Translator<?, ?> newInstance(
+            Class<?> input, Class<?> output, Model model, Map<String, ?> arguments)
+            throws TranslateException {
+        Path modelPath = model.getModelPath();
+        try {
+            HuggingFaceTokenizer tokenizer =
+                    HuggingFaceTokenizer.builder(arguments)
+                            .optTokenizerPath(modelPath)
+                            .optManager(model.getNDManager())
+                            .build();
+            TextClassificationTranslator translator =
+                    TextClassificationTranslator.builder(tokenizer, arguments).build();
+            if (input == String.class && output == Classifications.class) {
+                return translator;
+            } else if (input == Input.class && output == Output.class) {
+                return new TextClassificationServingTranslator(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/TextEmbeddingTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
index 0187ae3fa..0c505e3ad 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
@@ -119,7 +119,7 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
          * @param batchifier true to include token types
          * @return this builder
          */
-        public TextEmbeddingTranslator.Builder optBatchifier(Batchifier batchifier) {
+        public Builder optBatchifier(Batchifier batchifier) {
             this.batchifier = batchifier;
             return this;
         }
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslator.java
index 2a7093632..6486f22a6 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslator.java
@@ -38,7 +38,7 @@ public class TokenClassificationTranslator implements Translator<String, NamedEn
 
     private HuggingFaceTokenizer tokenizer;
     private Batchifier batchifier;
-    private Config config;
+    private PretrainedConfig config;
 
     TokenClassificationTranslator(HuggingFaceTokenizer tokenizer, Batchifier batchifier) {
         this.tokenizer = tokenizer;
@@ -57,7 +57,7 @@ public class TokenClassificationTranslator implements Translator<String, NamedEn
         Path path = ctx.getModel().getModelPath();
         Path file = path.resolve("config.json");
         try (Reader reader = Files.newBufferedReader(file)) {
-            config = JsonUtils.GSON.fromJson(reader, Config.class);
+            config = JsonUtils.GSON.fromJson(reader, PretrainedConfig.class);
         }
     }
 
@@ -172,9 +172,4 @@ public class TokenClassificationTranslator implements Translator<String, NamedEn
             return new TokenClassificationTranslator(tokenizer, batchifier);
         }
     }
-
-    private static final class Config {
-
-        Map<String, String> id2label;
-    }
 }


Test patch:
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TranslatorTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TranslatorTest.java
index c37743565..67f40885e 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TranslatorTest.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TranslatorTest.java
@@ -16,6 +16,7 @@ import ai.djl.Model;
 import ai.djl.ModelException;
 import ai.djl.huggingface.translator.FillMaskTranslatorFactory;
 import ai.djl.huggingface.translator.QuestionAnsweringTranslatorFactory;
+import ai.djl.huggingface.translator.TextClassificationTranslatorFactory;
 import ai.djl.huggingface.translator.TextEmbeddingTranslatorFactory;
 import ai.djl.huggingface.translator.TokenClassificationTranslatorFactory;
 import ai.djl.inference.Predictor;
@@ -397,4 +398,91 @@ public class TranslatorTest {
                     () -> factory.newInstance(String.class, Integer.class, model, arguments));
         }
     }
+
+    @Test
+    public void testTextClassificationTranslator()
+            throws ModelException, IOException, TranslateException {
+        String text = "DJL is the best.";
+
+        Block block =
+                new LambdaBlock(
+                        a -> {
+                            NDManager manager = a.getManager();
+                            float[] logits = new float[] {0.02f, 0.2f, 0.97f};
+                            NDArray arr = manager.create(logits, new Shape(1, 3));
+                            arr.setName("logits");
+                            return new NDList(arr);
+                        },
+                        "model");
+        Path modelDir = Paths.get("build/model");
+        Files.createDirectories(modelDir);
+
+        Path path = modelDir.resolve("config.json");
+        Map<String, Map<String, String>> map = new HashMap<>();
+        Map<String, String> id2label = new HashMap<>();
+        id2label.put("0", "LABEL_0");
+        id2label.put("1", "LABEL_1");
+        id2label.put("2", "LABEL_2");
+        map.put("id2label", id2label);
+        try (Writer writer = Files.newBufferedWriter(path)) {
+            writer.write(JsonUtils.GSON.toJson(map));
+        }
+
+        Criteria<String, Classifications> criteria =
+                Criteria.builder()
+                        .setTypes(String.class, Classifications.class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-uncased")
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new TextClassificationTranslatorFactory())
+                        .build();
+
+        try (ZooModel<String, Classifications> model = criteria.loadModel();
+                Predictor<String, Classifications> predictor = model.newPredictor()) {
+            Classifications res = predictor.predict(text);
+            Assert.assertEquals(res.best().getClassName(), "LABEL_2");
+        }
+
+        Criteria<Input, Output> criteria2 =
+                Criteria.builder()
+                        .setTypes(Input.class, Output.class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optEngine("PyTorch")
+                        .optArgument("tokenizer", "bert-base-uncased")
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new TextClassificationTranslatorFactory())
+                        .build();
+
+        try (ZooModel<Input, Output> model = criteria2.loadModel();
+                Predictor<Input, Output> predictor = model.newPredictor()) {
+            Input input = new Input();
+            input.add(text);
+            Output out = predictor.predict(input);
+            Classifications res = (Classifications) out.getData();
+            Assert.assertEquals(res.best().getClassName(), "LABEL_2");
+        }
+
+        try (Model model = Model.newInstance("test")) {
+            model.setBlock(block);
+            Map<String, String> options = new HashMap<>();
+            options.put("hasParameter", "false");
+            model.load(modelDir, "test", options);
+
+            TextClassificationTranslatorFactory factory = new TextClassificationTranslatorFactory();
+            Map<String, String> arguments = new HashMap<>();
+
+            Assert.assertThrows(
+                    TranslateException.class,
+                    () -> factory.newInstance(String.class, Integer.class, model, arguments));
+
+            arguments.put("tokenizer", "bert-base-uncased");
+
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> factory.newInstance(String.class, Integer.class, model, arguments));
+        }
+    }
 }