add error message for non-torchscript cross encoders

Signed-off-by: HenryL27 <[email protected]>
HenryL27 · Nov 28, 2023 · cd4a583 · cd4a583
1 parent 0731463
commit cd4a583
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 3 deletions.
diff --git a/.../org/opensearch/ml/engine/algorithms/text_similarity/TextSimilarityCrossEncoderModel.java b/.../org/opensearch/ml/engine/algorithms/text_similarity/TextSimilarityCrossEncoderModel.java
@@ -17,6 +17,8 @@
  */
 package org.opensearch.ml.engine.algorithms.text_similarity;
 
+import static org.opensearch.ml.engine.ModelHelper.PYTORCH_ENGINE;
+
 import java.util.ArrayList;
 import java.util.List;
 
@@ -58,8 +60,15 @@ public ModelTensorOutput predict(String modelId, MLInput mlInput) throws Transla
     }
 
     @Override
-    public Translator<Input, Output> getTranslator(String engine, MLModelConfig modelConfig) {
-        return new TextSimilarityTranslator();
+    public Translator<Input, Output> getTranslator(String engine, MLModelConfig modelConfig) throws IllegalArgumentException {
+        if(PYTORCH_ENGINE.equals(engine)) {
+            return new TextSimilarityTranslator();
+        } else {
+            throw new IllegalArgumentException("Wrong deep learning engine [" 
+                + engine 
+                + "]. Only TORCH_SCRIPT is supported for function name " 
+                + FunctionName.TEXT_SIMILARITY.name());
+        }
     }
 
     @Override

diff --git a/.../opensearch/ml/engine/algorithms/text_similarity/TextSimilarityCrossEncoderModelTest.java b/.../opensearch/ml/engine/algorithms/text_similarity/TextSimilarityCrossEncoderModelTest.java
@@ -190,6 +190,22 @@ public void initModel_predict_TorchScript_CrossEncoder() throws URISyntaxExcepti
         textSimilarityCrossEncoderModel.close();
     }
 
+    @Test
+    public void initModel_predict_ONNX_CrossEncoder_ThenFail() throws URISyntaxException {
+        model = MLModel.builder()
+                .modelFormat(MLModelFormat.ONNX)
+                .name("test_model_name")
+                .modelId("test_model_id")
+                .algorithm(FunctionName.TEXT_SIMILARITY)
+                .version("1.0.0")
+                .modelState(MLModelState.TRAINED)
+                .build();
+        assertThrows("Wrong deep learning engine [OnnxRuntime]. Only TORCH_SCRIPT is supported for function name TEXT_SIMILARITY",
+            MLException.class,
+            () -> textSimilarityCrossEncoderModel.initModel(model, params, encryptor)
+        );
+    }
+
     @Test
     public void initModel_NullModelHelper() throws URISyntaxException {
         Map<String, Object> params = new HashMap<>();
@@ -265,7 +281,7 @@ public void predict_AfterModelClosed() {
         log.info(e.getMessage());
         assert (e.getMessage().startsWith("Failed to inference TEXT_SIMILARITY"));
     }
-
+    
     @After
     public void tearDown() {
         FileUtils.deleteFileQuietly(mlCachePath);