From a0300f07bda5687299db58819188b3aba197d7d4 Mon Sep 17 00:00:00 2001 From: James Briggs <james.briggs@hotmail.com> Date: Sat, 1 Jun 2024 17:59:17 +0800 Subject: [PATCH] fix: force download of ViT encoder --- tests/unit/encoders/test_vit.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/unit/encoders/test_vit.py b/tests/unit/encoders/test_vit.py index 72d01b34..9970e310 100644 --- a/tests/unit/encoders/test_vit.py +++ b/tests/unit/encoders/test_vit.py @@ -4,8 +4,14 @@ import torch from PIL import Image from semantic_router.encoders import VitEncoder +from transformers import AutoTokenizer, AutoModel test_model_name = "hf-internal-testing/tiny-random-vit" + +# force the model download +tokenizer = AutoTokenizer.from_pretrained(test_model_name, force_download=True) +model = AutoModel.from_pretrained(test_model_name, force_download=True) + vit_encoder = VitEncoder(name=test_model_name) embed_dim = 32 -- GitLab