[text-embedding] Remove CrossEncoderTranslatorFactory in favor of Tex…

…tEmbeddingTranslatorFactory (#3239)
deepjavalibrary · Jun 5, 2024 · bd7f4ca · bd7f4ca
1 parent f2a1c60
commit bd7f4ca
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 81 deletions.
diff --git a/...tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslatorFactory.java b/...tokenizers/src/main/java/ai/djl/huggingface/translator/CrossEncoderTranslatorFactory.java
diff --git a/...okenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java b/...okenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslatorFactory.java
@@ -16,11 +16,14 @@
 import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
 import ai.djl.modality.Input;
 import ai.djl.modality.Output;
+import ai.djl.modality.nlp.translator.CrossEncoderServingTranslator;
 import ai.djl.modality.nlp.translator.TextEmbeddingServingTranslator;
+import ai.djl.translate.ArgumentsUtil;
 import ai.djl.translate.TranslateException;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 import ai.djl.util.Pair;
+import ai.djl.util.StringPair;
 
 import java.io.IOException;
 import java.io.Serializable;
@@ -39,6 +42,7 @@ public class TextEmbeddingTranslatorFactory implements TranslatorFactory, Serial
 
  static {
  SUPPORTED_TYPES.add(new Pair<>(String.class, float[].class));
+ SUPPORTED_TYPES.add(new Pair<>(StringPair.class, float[].class));
  SUPPORTED_TYPES.add(new Pair<>(Input.class, Output.class));
  }
 
@@ -61,6 +65,17 @@ public <I, O> Translator<I, O> newInstance(
  .optTokenizerPath(modelPath)
  .optManager(model.getNDManager())
  .build();
+ if (ArgumentsUtil.booleanValue(arguments, "reranking")) {
+ CrossEncoderTranslator translator =
+ CrossEncoderTranslator.builder(tokenizer, arguments).build();
+ if (input == StringPair.class && output == float[].class) {
+ return (Translator<I, O>) translator;
+ } else if (input == Input.class && output == Output.class) {
+ return (Translator<I, O>) new CrossEncoderServingTranslator(translator);
+ }
+ throw new IllegalArgumentException("Unsupported input/output types.");
+ }
+
  TextEmbeddingTranslator translator =
  TextEmbeddingTranslator.builder(tokenizer, arguments).build();
  if (input == String.class && output == float[].class) {

diff --git a/...ns/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/CrossEncoderTranslatorTest.java b/...ns/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/CrossEncoderTranslatorTest.java
@@ -14,7 +14,7 @@
 
 import ai.djl.Model;
 import ai.djl.ModelException;
-import ai.djl.huggingface.translator.CrossEncoderTranslatorFactory;
+import ai.djl.huggingface.translator.TextEmbeddingTranslatorFactory;
 import ai.djl.inference.Predictor;
 import ai.djl.modality.Input;
 import ai.djl.modality.Output;
@@ -65,8 +65,9 @@ public void testCrossEncoderTranslator()
  .optEngine("PyTorch")
  .optArgument("tokenizer", "bert-base-cased")
  .optArgument("tokenizerPath", modelDir)
+ .optArgument("reranking", true)
  .optOption("hasParameter", "false")
- .optTranslatorFactory(new CrossEncoderTranslatorFactory())
+ .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
  .build();
 
  try (ZooModel<StringPair, float[]> model = criteria.loadModel();
@@ -83,8 +84,9 @@ public void testCrossEncoderTranslator()
  .optBlock(block)
  .optEngine("PyTorch")
  .optArgument("tokenizer", "bert-base-cased")
+ .optArgument("reranking", true)
  .optOption("hasParameter", "false")
- .optTranslatorFactory(new CrossEncoderTranslatorFactory())
+ .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
  .build();
 
  try (ZooModel<Input, Output> model = criteria2.loadModel();
@@ -131,14 +133,15 @@ public void testCrossEncoderTranslator()
  options.put("hasParameter", "false");
  model.load(modelDir, "test", options);
 
- CrossEncoderTranslatorFactory factory = new CrossEncoderTranslatorFactory();
+ TextEmbeddingTranslatorFactory factory = new TextEmbeddingTranslatorFactory();
  Map<String, String> arguments = new HashMap<>();
 
  Assert.assertThrows(
  TranslateException.class,
  () -> factory.newInstance(String.class, Integer.class, model, arguments));
 
  arguments.put("tokenizer", "bert-base-cased");
+ arguments.put("reranking", "true");
 
  Assert.assertThrows(
  IllegalArgumentException.class,