From 2b83992f34b2146d98c9bf54fce02e98b82ae5e7 Mon Sep 17 00:00:00 2001
From: JingyaHuang <jingya@huggingface.co>
Date: Tue, 16 Jan 2024 13:43:56 +0000
Subject: [PATCH] add test for inference

---
 tests/inference/inference_utils.py | 1 +
 tests/inference/test_modeling.py   | 1 +
 2 files changed, 2 insertions(+)

diff --git a/tests/inference/inference_utils.py b/tests/inference/inference_utils.py
index 8d4331ce4..46e64bb7b 100644
--- a/tests/inference/inference_utils.py
+++ b/tests/inference/inference_utils.py
@@ -44,6 +44,7 @@
     "mpnet": "hf-internal-testing/tiny-random-MPNetModel",
     "roberta": "hf-internal-testing/tiny-random-RobertaModel",
     "roformer": "hf-internal-testing/tiny-random-RoFormerModel",
+    "sentence-transformers-transformer": "BAAI/bge-small-en-v1.5",
     "stable-diffusion": "hf-internal-testing/tiny-stable-diffusion-torch",
     "stable-diffusion-xl": "echarlaix/tiny-random-stable-diffusion-xl",
     "xlm": "hf-internal-testing/tiny-random-XLMModel",
diff --git a/tests/inference/test_modeling.py b/tests/inference/test_modeling.py
index 3c12b0bbb..cc5225f03 100644
--- a/tests/inference/test_modeling.py
+++ b/tests/inference/test_modeling.py
@@ -172,6 +172,7 @@ class NeuronModelForFeatureExtractionIntegrationTest(NeuronModelTestMixin):
             "mobilebert",
             "roberta",
             "roformer",
+            "sentence-transformers-transformer",
             # "xlm",  # accuracy off compared to pytorch (not due to the padding)
             "xlm-roberta",
         ]