-
Notifications
You must be signed in to change notification settings - Fork 653
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Adds TextEmbeddingTranslator (#1953)
- Loading branch information
1 parent
e21b1ac
commit 99c6de0
Showing
7 changed files
with
378 additions
and
4 deletions.
There are no files selected for viewing
64 changes: 64 additions & 0 deletions
64
api/src/main/java/ai/djl/modality/nlp/translator/TextEmbeddingServingTranslator.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,64 @@ | ||
/* | ||
* Copyright 2021 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.NDList; | ||
import ai.djl.translate.Batchifier; | ||
import ai.djl.translate.Translator; | ||
import ai.djl.translate.TranslatorContext; | ||
import ai.djl.util.JsonUtils; | ||
|
||
/** A {@link Translator} that can handle generic text embedding {@link Input} and {@link Output}. */ | ||
public class TextEmbeddingServingTranslator implements Translator<Input, Output> { | ||
|
||
private Translator<String, float[]> translator; | ||
|
||
/** | ||
* Constructs a {@code TextEmbeddingServingTranslator} instance. | ||
* | ||
* @param translator a {@code Translator} processes text embedding input | ||
*/ | ||
public TextEmbeddingServingTranslator(Translator<String, float[]> translator) { | ||
this.translator = translator; | ||
} | ||
|
||
/** {@inheritDoc} */ | ||
@Override | ||
public Batchifier getBatchifier() { | ||
return translator.getBatchifier(); | ||
} | ||
|
||
/** {@inheritDoc} */ | ||
@Override | ||
public void prepare(TranslatorContext ctx) throws Exception { | ||
translator.prepare(ctx); | ||
} | ||
|
||
/** {@inheritDoc} */ | ||
@Override | ||
public NDList processInput(TranslatorContext ctx, Input input) throws Exception { | ||
String text = input.getData().getAsString(); | ||
return translator.processInput(ctx, text); | ||
} | ||
|
||
/** {@inheritDoc} */ | ||
@Override | ||
public Output processOutput(TranslatorContext ctx, NDList list) throws Exception { | ||
float[] ret = translator.processOutput(ctx, list); | ||
Output output = new Output(); | ||
output.add(JsonUtils.GSON_PRETTY.toJson(ret)); | ||
return output; | ||
} | ||
} |
147 changes: 147 additions & 0 deletions
147
...sions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,147 @@ | ||
/* | ||
* 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.ndarray.NDArray; | ||
import ai.djl.ndarray.NDList; | ||
import ai.djl.ndarray.NDManager; | ||
import ai.djl.ndarray.types.DataType; | ||
import ai.djl.translate.ArgumentsUtil; | ||
import ai.djl.translate.Batchifier; | ||
import ai.djl.translate.Translator; | ||
import ai.djl.translate.TranslatorContext; | ||
|
||
import java.io.IOException; | ||
import java.util.Map; | ||
|
||
/** The translator for Huggingface text embedding model. */ | ||
public class TextEmbeddingTranslator implements Translator<String, float[]> { | ||
|
||
private static final int[] AXIS = {0}; | ||
|
||
private HuggingFaceTokenizer tokenizer; | ||
private Batchifier batchifier; | ||
|
||
TextEmbeddingTranslator(HuggingFaceTokenizer tokenizer, Batchifier batchifier) { | ||
this.tokenizer = tokenizer; | ||
this.batchifier = batchifier; | ||
} | ||
|
||
/** {@inheritDoc} */ | ||
@Override | ||
public Batchifier getBatchifier() { | ||
return batchifier; | ||
} | ||
|
||
/** {@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 float[] processOutput(TranslatorContext ctx, NDList list) { | ||
NDArray embeddings = list.get("last_hidden_state"); | ||
Encoding encoding = (Encoding) ctx.getAttachment("encoding"); | ||
long[] attentionMask = encoding.getAttentionMask(); | ||
NDManager manager = ctx.getNDManager(); | ||
NDArray inputAttentionMask = manager.create(attentionMask).toType(DataType.FLOAT32, true); | ||
long[] shape = embeddings.getShape().getShape(); | ||
inputAttentionMask = inputAttentionMask.tile(shape[shape.length - 1]); | ||
inputAttentionMask = inputAttentionMask.reshape(embeddings.getShape()); | ||
NDArray inputAttentionMaskSum = inputAttentionMask.sum(AXIS); | ||
NDArray clamp = inputAttentionMaskSum.clip(1e-9, 1e12); | ||
NDArray prod = embeddings.mul(inputAttentionMask); | ||
NDArray sum = prod.sum(AXIS); | ||
embeddings = sum.div(clamp).normalize(2, 0); | ||
|
||
return embeddings.toFloatArray(); | ||
} | ||
|
||
/** | ||
* Creates a builder to build a {@code TextEmbeddingTranslator}. | ||
* | ||
* @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 TextEmbeddingTranslator}. | ||
* | ||
* @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 TextEmbeddingTranslator.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 TextEmbeddingTranslator build() throws IOException { | ||
return new TextEmbeddingTranslator(tokenizer, batchifier); | ||
} | ||
} | ||
} |
72 changes: 72 additions & 0 deletions
72
...okenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,72 @@ | ||
/* | ||
* 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.Input; | ||
import ai.djl.modality.Output; | ||
import ai.djl.modality.nlp.translator.TextEmbeddingServingTranslator; | ||
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 TextEmbeddingTranslator} instance. */ | ||
public class TextEmbeddingTranslatorFactory implements TranslatorFactory { | ||
|
||
private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>(); | ||
|
||
static { | ||
SUPPORTED_TYPES.add(new Pair<>(String.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 | ||
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(); | ||
TextEmbeddingTranslator translator = | ||
TextEmbeddingTranslator.builder(tokenizer, arguments).build(); | ||
if (input == String.class && output == float[].class) { | ||
return translator; | ||
} else if (input == Input.class && output == Output.class) { | ||
return new TextEmbeddingServingTranslator(translator); | ||
} | ||
throw new IllegalArgumentException("Unsupported input/output types."); | ||
} catch (IOException e) { | ||
throw new TranslateException("Failed to load tokenizer.", e); | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.