[huggingface] Adds CrossEncoderTranslator

api/src/main/java/ai/djl/modality/nlp/translator/CrossEncoderServingTranslator.java

@@ -0,0 +1,115 @@
+/*
+ * 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.modality.nlp.translator;
+
+import ai.djl.modality.Input;
+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;
+import ai.djl.util.JsonUtils;
+import ai.djl.util.PairList;
+import ai.djl.util.StringPair;
+
+import com.google.gson.JsonElement;
+import com.google.gson.JsonParseException;
+
+/** A {@link Translator} that can handle generic cross encoder {@link Input} and {@link Output}. */
+public class CrossEncoderServingTranslator implements NoBatchifyTranslator<Input, Output> {
+
+ private Translator<StringPair, float[]> translator;
+ private Translator<StringPair[], float[][]> batchTranslator;
+
+ /**
+ * Constructs a {@code CrossEncoderServingTranslator} instance.
+ *
+ * @param translator a {@code Translator} processes question answering input
+ */
+ public CrossEncoderServingTranslator(Translator<StringPair, float[]> translator) {
+ this.translator = translator;
+ this.batchTranslator = translator.toBatchTranslator();
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public void prepare(TranslatorContext ctx) throws Exception {
+ translator.prepare(ctx);
+ batchTranslator.prepare(ctx);
+ }
+
+ /** {@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.");
+ }
+
+ String contentType = input.getProperty("Content-Type", null);
+ StringPair pair;
+ 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);
+ StringPair[] inputs = JsonUtils.GSON.fromJson(json, StringPair[].class);
+ return batchTranslator.processInput(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.");
+ }
+ } catch (JsonParseException e) {
+ throw new TranslateException("Input is not a valid json.", e);
+ }
+ } else {
+ String key = input.getAsString("key");
+ String value = input.getAsString("value");
+ if (key == null || value == null) {
+ throw new TranslateException("Missing key or value in input.");
+ }
+ pair = new StringPair(key, value);
+ }
+
+ NDList ret = translator.processInput(ctx, pair);
+ Batchifier batchifier = translator.getBatchifier();
+ if (batchifier != null) {
+ NDList[] batch = {ret};
+ return batchifier.batchify(batch);
+ }
+ return ret;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public Output processOutput(TranslatorContext ctx, NDList list) throws Exception {
+ Output output = new Output();
+ output.addProperty("Content-Type", "application/json");
+ if (ctx.getAttachment("batch") != null) {
+ output.add(BytesSupplier.wrapAsJson(batchTranslator.processOutput(ctx, list)));
+ } else {
+ Batchifier batchifier = translator.getBatchifier();
+ if (batchifier != null) {
+ list = batchifier.unbatchify(list)[0];
+ }
+ output.add(BytesSupplier.wrapAsJson(translator.processOutput(ctx, list)));
+ }
+ return output;
+ }
+}

api/src/main/java/ai/djl/util/StringPair.java

@@ -0,0 +1,27 @@
+/*
+ * 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.util;
+
+/** A class containing the string key-value pair. */
+public class StringPair extends Pair<String, String> {
+
+ /**
+ * Constructs a {@code Pair} instance with key and value.
+ *
+ * @param key the key
+ * @param value the value
+ */
+ public StringPair(String key, String value) {
+ super(key, value);
+ }
+}

extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderBatchTranslator.java

@@ -0,0 +1,69 @@
+/*
+ * 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.huggingface.tokenizers.Encoding;
+import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.translate.Batchifier;
+import ai.djl.translate.NoBatchifyTranslator;
+import ai.djl.translate.TranslateException;
+import ai.djl.translate.TranslatorContext;
+import ai.djl.util.PairList;
+import ai.djl.util.StringPair;
+
+import java.util.Arrays;
+
+/** The translator for Huggingface cross encoder model. */
+public class CrossEncoderBatchTranslator implements NoBatchifyTranslator<StringPair[], float[][]> {
+
+ private HuggingFaceTokenizer tokenizer;
+ private boolean includeTokenTypes;
+ private Batchifier batchifier;
+
+ CrossEncoderBatchTranslator(
+ HuggingFaceTokenizer tokenizer, boolean includeTokenTypes, Batchifier batchifier) {
+ this.tokenizer = tokenizer;
+ this.includeTokenTypes = includeTokenTypes;
+ this.batchifier = batchifier;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public NDList processInput(TranslatorContext ctx, StringPair[] inputs)
+ throws TranslateException {
+ NDManager manager = ctx.getNDManager();
+ PairList<String, String> list = new PairList<>(Arrays.asList(inputs));
+ Encoding[] encodings = tokenizer.batchEncode(list);
+ NDList[] batch = new NDList[encodings.length];
+ for (int i = 0; i < encodings.length; ++i) {
+ batch[i] = encodings[i].toNDList(manager, includeTokenTypes);
+ }
+ return batchifier.batchify(batch);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public float[][] processOutput(TranslatorContext ctx, NDList list) {
+ NDList[] batch = batchifier.unbatchify(list);
+ float[][] ret = new float[batch.length][];
+ for (int i = 0; i < batch.length; ++i) {
+ NDArray logits = list.get(0);
+ NDArray result = logits.getNDArrayInternal().sigmoid();
+ ret[i] = result.toFloatArray();
+ }
+ return ret;
+ }
+}

extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslator.java

@@ -0,0 +1,149 @@
+/*
+ * 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.huggingface.tokenizers.Encoding;
+import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.translate.ArgumentsUtil;
+import ai.djl.translate.Batchifier;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorContext;
+import ai.djl.util.StringPair;
+
+import java.io.IOException;
+import java.util.Map;
+
+/** The translator for Huggingface cross encoder model. */
+public class CrossEncoderTranslator implements Translator<StringPair, float[]> {
+
+ private HuggingFaceTokenizer tokenizer;
+ private boolean includeTokenTypes;
+ private Batchifier batchifier;
+
+ CrossEncoderTranslator(
+ HuggingFaceTokenizer tokenizer, boolean includeTokenTypes, Batchifier batchifier) {
+ this.tokenizer = tokenizer;
+ this.includeTokenTypes = includeTokenTypes;
+ this.batchifier = batchifier;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public Batchifier getBatchifier() {
+ return batchifier;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public NDList processInput(TranslatorContext ctx, StringPair input) {
+ Encoding encoding = tokenizer.encode(input.getKey(), input.getValue());
+ ctx.setAttachment("encoding", encoding);
+ return encoding.toNDList(ctx.getNDManager(), includeTokenTypes);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public float[] processOutput(TranslatorContext ctx, NDList list) {
+ NDArray logits = list.get(0);
+ NDArray result = logits.getNDArrayInternal().sigmoid();
+ return result.toFloatArray();
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public CrossEncoderBatchTranslator toBatchTranslator(Batchifier batchifier) {
+ tokenizer.enableBatch();
+ return new CrossEncoderBatchTranslator(tokenizer, includeTokenTypes, batchifier);
+ }
+
+ /**
+ * Creates a builder to build a {@code CrossEncoderTranslator}.
+ *
+ * @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 CrossEncoderTranslator}.
+ *
+ * @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 question answering translator. */
+ public static final class Builder {
+
+ private HuggingFaceTokenizer tokenizer;
+ private boolean includeTokenTypes;
+ private Batchifier batchifier = Batchifier.STACK;
+
+ Builder(HuggingFaceTokenizer tokenizer) {
+ this.tokenizer = tokenizer;
+ }
+
+ /**
+ * Sets if include token types for the {@link Translator}.
+ *
+ * @param includeTokenTypes true to include token types
+ * @return this builder
+ */
+ public Builder optIncludeTokenTypes(boolean includeTokenTypes) {
+ this.includeTokenTypes = includeTokenTypes;
+ return this;
+ }
+
+ /**
+ * 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) {
+ optIncludeTokenTypes(ArgumentsUtil.booleanValue(arguments, "includeTokenTypes"));
+ 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 CrossEncoderTranslator build() throws IOException {
+ return new CrossEncoderTranslator(tokenizer, includeTokenTypes, batchifier);
+ }
+ }
+}