update pipeline unit test model to a tiny one

Striveworks · Mar 7, 2024 · 483a2d6 · 483a2d6
1 parent bef23e6
commit 483a2d6
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/runtimes/huggingface/tests/test_pipeline.py b/runtimes/huggingface/tests/test_pipeline.py
@@ -7,7 +7,8 @@
 
 
 def test_sentence_transformers_pipeline():
-    pretrained_model = "sentence-transformers/all-distilroberta-v1"
+    # A tiny sentence transformers model for testing, which only takes 17.6MB of memory
+    pretrained_model = "sentence-transformers-testing/stsb-bert-tiny-safetensors"
     st_embeder = SentenceTransformer(pretrained_model)
 
     hf_settings = HuggingFaceSettings(