support includeTokenTypes in TextEmbeddingBatchTranslator (#3032)

* support includeTokenTypes in TextEmbeddingBatchTranslator Co-authored-by: Frank Liu <frankfliu2000@gmail.com>
deepjavalibrary · Apr 26, 2024 · 8055c01 · 8055c01
1 parent b65ba10
commit 8055c01
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/.../tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingBatchTranslator.java b/.../tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingBatchTranslator.java
@@ -28,16 +28,19 @@ public class TextEmbeddingBatchTranslator implements NoBatchifyTranslator<String
     private Batchifier batchifier;
     private boolean normalize;
     private String pooling;
+    private boolean includeTokenTypes;
 
     TextEmbeddingBatchTranslator(
             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} */
@@ -48,7 +51,7 @@ public NDList processInput(TranslatorContext ctx, String[] input) {
         ctx.setAttachment("encodings", encodings);
         NDList[] batch = new NDList[encodings.length];
         for (int i = 0; i < encodings.length; ++i) {
-            batch[i] = encodings[i].toNDList(manager, false);
+            batch[i] = encodings[i].toNDList(manager, includeTokenTypes);
         }
         return batchifier.batchify(batch);
     }

diff --git a/...sions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java b/...sions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
@@ -81,7 +81,8 @@ public float[] processOutput(TranslatorContext ctx, NDList list) {
     @Override
     public TextEmbeddingBatchTranslator toBatchTranslator(Batchifier batchifier) {
         tokenizer.enableBatch();
-        return new TextEmbeddingBatchTranslator(tokenizer, batchifier, pooling, normalize);
+        return new TextEmbeddingBatchTranslator(
+                tokenizer, batchifier, pooling, normalize, includeTokenTypes);
     }
 
     static NDArray processEmbedding(