Review fixes

qdrant · Jan 16, 2025 · 3a0b570 · 3a0b570
1 parent 5518a88
commit 3a0b570
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git 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
@@ -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)