add function name annotation

Signed-off-by: HenryL27 <[email protected]>
HenryL27 · Nov 10, 2023 · c40b374 · c40b374
1 parent d855243
commit c40b374
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/.../org/opensearch/ml/engine/algorithms/text_similarity/TextSimilarityCrossEncoderModel.java b/.../org/opensearch/ml/engine/algorithms/text_similarity/TextSimilarityCrossEncoderModel.java
@@ -21,20 +21,23 @@
 import java.util.List;
 
 import org.apache.commons.lang3.tuple.Pair;
+import org.opensearch.ml.common.FunctionName;
 import org.opensearch.ml.common.dataset.MLInputDataset;
 import org.opensearch.ml.common.dataset.TextSimilarityInputDataSet;
 import org.opensearch.ml.common.input.MLInput;
 import org.opensearch.ml.common.model.MLModelConfig;
 import org.opensearch.ml.common.output.model.ModelTensorOutput;
 import org.opensearch.ml.common.output.model.ModelTensors;
 import org.opensearch.ml.engine.algorithms.DLModel;
+import org.opensearch.ml.engine.annotation.Function;
 
 import ai.djl.modality.Input;
 import ai.djl.modality.Output;
 import ai.djl.translate.TranslateException;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 
+@Function(FunctionName.TEXT_SIMILARITY)
 public class TextSimilarityCrossEncoderModel extends DLModel {
 
     @Override