add test for inference

huggingface · Jan 16, 2024 · 2b83992 · 2b83992
1 parent 65988b8
commit 2b83992
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git 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
@@ -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",
         ]