extract an abstract model for sparse and dense sentence transformer t…

…ranslator Signed-off-by: xinyual <[email protected]>
opensearch-project · Sep 26, 2023 · a837080 · a837080
1 parent 4e3dc78
commit a837080
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 73 deletions.
diff --git a/...thms/src/main/java/org/opensearch/ml/engine/algorithms/SentenceTransformerTranslator.java b/...thms/src/main/java/org/opensearch/ml/engine/algorithms/SentenceTransformerTranslator.java
@@ -0,0 +1,62 @@
+package org.opensearch.ml.engine.algorithms;
+
+import ai.djl.huggingface.tokenizers.Encoding;
+import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+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.Batchifier;
+import ai.djl.translate.ServingTranslator;
+import ai.djl.translate.TranslatorContext;
+import org.opensearch.ml.common.output.model.MLResultDataType;
+import org.opensearch.ml.common.output.model.ModelTensor;
+import org.opensearch.ml.common.output.model.ModelTensors;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+
+abstract public class SentenceTransformerTranslator implements ServingTranslator {
+    protected HuggingFaceTokenizer tokenizer;
+
+    @Override
+    public Batchifier getBatchifier() {
+        return Batchifier.STACK;
+    }
+    @Override
+    public void prepare(TranslatorContext ctx) throws IOException {
+        Path path = ctx.getModel().getModelPath();
+        tokenizer = HuggingFaceTokenizer.builder().optPadding(true).optTokenizerPath(path.resolve("tokenizer.json")).build();
+    }
+
+    @Override
+    public NDList processInput(TranslatorContext ctx, Input input) {
+        String sentence = input.getAsString(0);
+        NDManager manager = ctx.getNDManager();
+        NDList ndList = new NDList();
+        Encoding encodings = tokenizer.encode(sentence);
+        long[] indices = encodings.getIds();
+        long[] attentionMask = encodings.getAttentionMask();
+
+        NDArray indicesArray = manager.create(indices);
+        indicesArray.setName("input1.input_ids");
+
+        NDArray attentionMaskArray = manager.create(attentionMask);
+        attentionMaskArray.setName("input1.attention_mask");
+
+        ndList.add(indicesArray);
+        ndList.add(attentionMaskArray);
+        return ndList;
+    }
+
+    @Override
+    public void setArguments(Map<String, ?> arguments) {
+    }
+}
diff --git a/...in/java/org/opensearch/ml/engine/algorithms/sparse_encoding/SparseEncodingTranslator.java b/...in/java/org/opensearch/ml/engine/algorithms/sparse_encoding/SparseEncodingTranslator.java
@@ -19,6 +19,7 @@
 import org.opensearch.ml.common.output.model.MLResultDataType;
 import org.opensearch.ml.common.output.model.ModelTensor;
 import org.opensearch.ml.common.output.model.ModelTensors;
+import org.opensearch.ml.engine.algorithms.SentenceTransformerTranslator;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -27,38 +28,7 @@
 
 import static org.opensearch.ml.common.CommonValue.ML_MAP_RESPONSE_KEY;
 
-public class SparseEncodingTranslator implements ServingTranslator {
-    private HuggingFaceTokenizer tokenizer;
-
-    @Override
-    public Batchifier getBatchifier() {
-        return Batchifier.STACK;
-    }
-    @Override
-    public void prepare(TranslatorContext ctx) throws IOException {
-        Path path = ctx.getModel().getModelPath();
-        tokenizer = HuggingFaceTokenizer.builder().optPadding(true).optTokenizerPath(path.resolve("tokenizer.json")).build();
-    }
-
-    @Override
-    public NDList processInput(TranslatorContext ctx, Input input) {
-        String sentence = input.getAsString(0);
-        NDManager manager = ctx.getNDManager();
-        NDList ndList = new NDList();
-        Encoding encodings = tokenizer.encode(sentence);
-        long[] indices = encodings.getIds();
-        long[] attentionMask = encodings.getAttentionMask();
-
-        NDArray indicesArray = manager.create(indices);
-        indicesArray.setName("input1.input_ids");
-
-        NDArray attentionMaskArray = manager.create(attentionMask);
-        attentionMaskArray.setName("input1.attention_mask");
-
-        ndList.add(indicesArray);
-        ndList.add(attentionMaskArray);
-        return ndList;
-    }
+public class SparseEncodingTranslator extends SentenceTransformerTranslator {
     private Map<String, Float>  convertOutput(NDArray array)
     {
         Map<String, Float> map = new HashMap<>();
@@ -94,8 +64,4 @@ public Output processOutput(TranslatorContext ctx, NDList list) {
         output.add(modelTensorOutput.toBytes());
         return output;
     }
-
-    @Override
-    public void setArguments(Map<String, ?> arguments) {
-    }
 }
diff --git a/...earch/ml/engine/algorithms/text_embedding/SentenceTransformerTextEmbeddingTranslator.java b/...earch/ml/engine/algorithms/text_embedding/SentenceTransformerTextEmbeddingTranslator.java
@@ -19,6 +19,7 @@
 import org.opensearch.ml.common.output.model.MLResultDataType;
 import org.opensearch.ml.common.output.model.ModelTensor;
 import org.opensearch.ml.common.output.model.ModelTensors;
+import org.opensearch.ml.engine.algorithms.SentenceTransformerTranslator;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -28,39 +29,7 @@
 import java.util.List;
 import java.util.Map;
 
-public class SentenceTransformerTextEmbeddingTranslator implements ServingTranslator {
-    private HuggingFaceTokenizer tokenizer;
-
-    @Override
-    public Batchifier getBatchifier() {
-        return Batchifier.STACK;
-    }
-    @Override
-    public void prepare(TranslatorContext ctx) throws IOException {
-        Path path = ctx.getModel().getModelPath();
-        tokenizer = HuggingFaceTokenizer.builder().optPadding(true).optTokenizerPath(path.resolve("tokenizer.json")).build();
-    }
-
-    @Override
-    public NDList processInput(TranslatorContext ctx, Input input) {
-        String sentence = input.getAsString(0);
-        NDManager manager = ctx.getNDManager();
-        NDList ndList = new NDList();
-        Encoding encodings = tokenizer.encode(sentence);
-        long[] indices = encodings.getIds();
-        long[] attentionMask = encodings.getAttentionMask();
-
-        NDArray indicesArray = manager.create(indices);
-        indicesArray.setName("input1.input_ids");
-
-        NDArray attentionMaskArray = manager.create(attentionMask);
-        attentionMaskArray.setName("input1.attention_mask");
-
-        ndList.add(indicesArray);
-        ndList.add(attentionMaskArray);
-        return ndList;
-    }
-
+public class SentenceTransformerTextEmbeddingTranslator extends SentenceTransformerTranslator {
     @Override
     public Output processOutput(TranslatorContext ctx, NDList list) {
         Output output = new Output(200, "OK");
@@ -89,8 +58,4 @@ public Output processOutput(TranslatorContext ctx, NDList list) {
         output.add(modelTensorOutput.toBytes());
         return output;
     }
-
-    @Override
-    public void setArguments(Map<String, ?> arguments) {
-    }
 }