diff --git a/fastembed/late_interaction_multimodal/onnx_multimodal_model.py b/fastembed/late_interaction_multimodal/onnx_multimodal_model.py index c9a12b22..7f084fb8 100644 --- a/fastembed/late_interaction_multimodal/onnx_multimodal_model.py +++ b/fastembed/late_interaction_multimodal/onnx_multimodal_model.py @@ -2,7 +2,7 @@ import os from multiprocessing import get_all_start_methods from pathlib import Path -from typing import Any, Iterable, Optional, Sequence, Type, Union +from typing import Any, Iterable, Optional, Sequence, Type, Union, get_args import numpy as np from PIL import Image @@ -186,7 +186,7 @@ def _embed_images( ) -> Iterable[T]: is_small = False - if isinstance(images, (str, Path, Image.Image)): + if isinstance(images, get_args(ImageInput)): images = [images] is_small = True diff --git a/tests/test_late_interaction_multimodal.py b/tests/test_late_interaction_multimodal.py index 9d7c0b90..dc75d9df 100644 --- a/tests/test_late_interaction_multimodal.py +++ b/tests/test_late_interaction_multimodal.py @@ -58,7 +58,6 @@ def test_batch_embedding(): for value in result: batch_size, token_num, abridged_dim = expected_result.shape assert np.allclose(value[:token_num, :abridged_dim], expected_result, atol=1e-3) - break if is_ci: delete_model_cache(model.model._model_dir)