Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fixes #2744, support onnx model for TextEmbeddingTranslator #2749

Merged
merged 1 commit into from
Aug 14, 2023
Merged

Fixes #2744, support onnx model for TextEmbeddingTranslator #2749

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 24 additions & 3 deletions ...sions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,16 +35,19 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
private Batchifier batchifier;
private boolean normalize;
private String pooling;
private boolean includeTokenTypes;

TextEmbeddingTranslator(
HuggingFaceTokenizer tokenizer,
Batchifier batchifier,
String pooling,
boolean normalize) {
boolean normalize,
boolean includeTokenTypes) {
this.tokenizer = tokenizer;
this.batchifier = batchifier;
this.pooling = pooling;
this.normalize = normalize;
this.includeTokenTypes = includeTokenTypes;
}

/** {@inheritDoc} */
Expand All @@ -58,7 +61,7 @@ public Batchifier getBatchifier() {
public NDList processInput(TranslatorContext ctx, String input) {
Encoding encoding = tokenizer.encode(input);
ctx.setAttachment("encoding", encoding);
return encoding.toNDList(ctx.getNDManager(), false);
return encoding.toNDList(ctx.getNDManager(), includeTokenTypes);
}

/** {@inheritDoc} */
Expand All @@ -84,6 +87,10 @@ public TextEmbeddingBatchTranslator toBatchTranslator(Batchifier batchifier) {
static NDArray processEmbedding(
NDManager manager, NDList list, Encoding encoding, String pooling) {
NDArray embedding = list.get("last_hidden_state");
if (embedding == null) {
// For Onnx model, NDArray name is not present
embedding = list.head();
}
long[] attentionMask = encoding.getAttentionMask();
NDArray inputAttentionMask = manager.create(attentionMask).toType(DataType.FLOAT32, true);
switch (pooling) {
Expand Down Expand Up @@ -167,6 +174,7 @@ public static final class Builder {
private Batchifier batchifier = Batchifier.STACK;
private boolean normalize = true;
private String pooling = "mean";
private boolean includeTokenTypes;

Builder(HuggingFaceTokenizer tokenizer) {
this.tokenizer = tokenizer;
Expand Down Expand Up @@ -214,6 +222,17 @@ public Builder optPoolingMode(String poolingMode) {
return this;
}

/**
* 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;
}

/**
* Configures the builder with the model arguments.
*
Expand All @@ -224,6 +243,7 @@ public void configure(Map<String, ?> arguments) {
optBatchifier(Batchifier.fromString(batchifierStr));
optNormalize(ArgumentsUtil.booleanValue(arguments, "normalize", true));
optPoolingMode(ArgumentsUtil.stringValue(arguments, "pooling", "mean"));
optIncludeTokenTypes(ArgumentsUtil.booleanValue(arguments, "includeTokenTypes"));
}

/**
Expand All @@ -233,7 +253,8 @@ public void configure(Map<String, ?> arguments) {
* @throws IOException if I/O error occurs
*/
public TextEmbeddingTranslator build() throws IOException {
return new TextEmbeddingTranslator(tokenizer, batchifier, pooling, normalize);
return new TextEmbeddingTranslator(
tokenizer, batchifier, pooling, normalize, includeTokenTypes);
}
}
}
Loading