speedy-utils 1.1.32__py3-none-any.whl → 1.1.34__py3-none-any.whl

llm_utils/lm/llm.py

@@ -63,9 +63,13 @@ class LLM(
         vllm_cmd: str | None = None,
         vllm_timeout: int = 1200,
         vllm_reuse: bool = True,
+        verbose=False,
         **model_kwargs,
     ):
         """Initialize LLMTask."""
+        if verbose:
+            available_models = LLM.list_models(client=client)
+            logger.info(f'Available models: {available_models}')
         self.instruction = instruction
         self.input_model = input_model
         self.output_model = output_model

llm_utils/lm/utils.py

@@ -282,7 +282,11 @@ def get_base_client(
             return MOpenAI(
                 base_url=f"http://localhost:{port}/v1", api_key=api_key, cache=cache
             )
-        raise ValueError("Either client or vllm_cmd must be provided.")
+        # Use default port 8000 when client is None
+        logger.info("No client specified, using default port 8000 at http://localhost:8000/v1")
+        return MOpenAI(
+            base_url="http://localhost:8000/v1", api_key=api_key, cache=cache
+        )
     if isinstance(client, int):
         return MOpenAI(
             base_url=f"http://localhost:{client}/v1", api_key=api_key, cache=cache

speedy_utils/__imports.py

@@ -34,7 +34,6 @@ import re
 import sys
 import textwrap
 import threading
-import time
 import traceback
 import types
 import uuid
@@ -77,6 +76,9 @@ from typing import (
 )
 
 import cachetools
+
+# Direct imports (previously lazy-loaded)
+import numpy as np
 import psutil
 from fastcore.parallel import parallel
 from json_repair import loads as jloads
@@ -84,52 +86,61 @@ from loguru import logger
 from tqdm import tqdm
 
 
-# Direct imports (previously lazy-loaded)
-import numpy as np
 tabulate = __import__('tabulate').tabulate
 import xxhash
 
+
 # Optional imports - lazy loaded for performance
 def _get_pandas():
     """Lazy import pandas."""
     try:
         import pandas as pd
+
         return pd
     except ImportError:
         return None
 
+
 def _get_ray():
     """Lazy import ray."""
     try:
         import ray
+
         return ray
     except ImportError:
         return None
 
+
 def _get_matplotlib():
     """Lazy import matplotlib."""
     try:
         import matplotlib
+
         return matplotlib
     except ImportError:
         return None
 
+
 def _get_matplotlib_pyplot():
     """Lazy import matplotlib.pyplot."""
     try:
         import matplotlib.pyplot as plt
+
         return plt
     except ImportError:
         return None
 
+
 def _get_ipython_core():
     """Lazy import IPython.core.getipython."""
     try:
         from IPython.core.getipython import get_ipython
-        return get_ipython
+
+        return get_ipython()
     except ImportError:
         return None
 
+
 # Cache for lazy imports
 _pandas_cache = None
 _ray_cache = None
@@ -137,9 +148,11 @@ _matplotlib_cache = None
 _plt_cache = None
 _get_ipython_cache = None
 
+
 # Lazy import classes for performance-critical modules
 class _LazyModule:
     """Lazy module loader that imports only when accessed."""
+
     def __init__(self, import_func, cache_var_name):
         self._import_func = import_func
         self._cache_var_name = cache_var_name
@@ -168,9 +181,10 @@ class _LazyModule:
 
     def __repr__(self):
         if self._module is None:
-            return f"<LazyModule: not loaded>"
+            return '<LazyModule: not loaded>'
         return repr(self._module)
 
+
 # Create lazy loaders for top slow imports (import only when accessed)
 pd = _LazyModule(_get_pandas, '_pandas_cache')
 ray = _LazyModule(_get_ray, '_ray_cache')
@@ -200,11 +214,12 @@ try:
 except ImportError:
     BaseModel = None
 if TYPE_CHECKING:
+    import matplotlib.pyplot as plt
     import numpy as np
     import pandas as pd
     import ray
     import torch
-    import matplotlib.pyplot as plt
+
     # xxhash
     import xxhash  # type: ignore
     from IPython.core.getipython import get_ipython  # type: ignore

speedy_utils/multi_worker/process.py

@@ -1,6 +1,3 @@
-# ray_multi_process.py
-
-
 from ..__imports import *
 
 
@@ -8,6 +5,12 @@ SPEEDY_RUNNING_PROCESSES: list[psutil.Process] = []
 _SPEEDY_PROCESSES_LOCK = threading.Lock()
 
 
+# /mnt/data/anhvth8/venvs/Megatron-Bridge-Host/lib/python3.12/site-packages/ray/_private/worker.py:2046: FutureWarning: Tip: In future versions of Ray, Ray will no longer override accelerator visible devices env var if num_gpus=0 or num_gpus=None (default). To enable this behavior and turn off this error message, set RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0
+# turn off future warning and verbose task logs
+os.environ["RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO"] = "0"
+os.environ["RAY_DEDUP_LOGS"] = "0"
+os.environ["RAY_LOG_TO_STDERR"] = "0"
+
 def _prune_dead_processes() -> None:
     """Remove dead processes from tracking list."""
     with _SPEEDY_PROCESSES_LOCK:
@@ -78,6 +81,7 @@ def _track_multiprocessing_processes() -> None:
 
 def _build_cache_dir(func: Callable, items: list[Any]) -> Path:
     """Build cache dir with function name + timestamp."""
+    import datetime
     func_name = getattr(func, '__name__', 'func')
     now = datetime.datetime.now()
     stamp = now.strftime('%m%d_%Hh%Mm%Ss')
@@ -85,9 +89,8 @@ def _build_cache_dir(func: Callable, items: list[Any]) -> Path:
     path = Path('.cache') / run_id
     path.mkdir(parents=True, exist_ok=True)
     return path
-
-
-def wrap_dump(func: Callable, cache_dir: Path | None):
+_DUMP_THREADS = []
+def wrap_dump(func: Callable, cache_dir: Path | None, dump_in_thread: bool = True):
     """Wrap a function so results are dumped to .pkl when cache_dir is set."""
     if cache_dir is None:
         return func
@@ -95,8 +98,24 @@ def wrap_dump(func: Callable, cache_dir: Path | None):
     def wrapped(x, *args, **kwargs):
         res = func(x, *args, **kwargs)
         p = cache_dir / f'{uuid.uuid4().hex}.pkl'
-        with open(p, 'wb') as fh:
-            pickle.dump(res, fh)
+
+        def save():
+            with open(p, 'wb') as fh:
+                pickle.dump(res, fh)
+            # Clean trash to avoid bloating memory
+            # print(f'Thread count: {threading.active_count()}')
+            # print(f'Saved result to {p}')
+
+        if dump_in_thread:
+            thread = threading.Thread(target=save)
+            _DUMP_THREADS.append(thread)
+            # count thread
+            # print(f'Thread count: {threading.active_count()}')
+            while threading.active_count() > 16:
+                time.sleep(0.1)
+            thread.start()
+        else:
+            save()
         return str(p)
 
     return wrapped
@@ -109,20 +128,28 @@ RAY_WORKER = None
 
 def ensure_ray(workers: int, pbar: tqdm | None = None):
     """Initialize or reinitialize Ray with a given worker count, log to bar postfix."""
+    import ray as _ray_module
+    import logging
+
     global RAY_WORKER
-    if not ray.is_initialized() or workers != RAY_WORKER:
-        if ray.is_initialized() and pbar:
+    # shutdown when worker count changes or if Ray not initialized
+    if not _ray_module.is_initialized() or workers != RAY_WORKER:
+        if _ray_module.is_initialized() and pbar:
             pbar.set_postfix_str(f'Restarting Ray {workers} workers')
-            ray.shutdown()
+            _ray_module.shutdown()
         t0 = time.time()
-        ray.init(num_cpus=workers, ignore_reinit_error=True)
+        _ray_module.init(
+            num_cpus=workers,
+            ignore_reinit_error=True,
+            logging_level=logging.ERROR,
+            log_to_driver=False,
+        )
         took = time.time() - t0
         _track_ray_processes()  # Track Ray worker processes
         if pbar:
             pbar.set_postfix_str(f'ray.init {workers} took {took:.2f}s')
         RAY_WORKER = workers
 
-
 def multi_process(
     func: Callable[[Any], Any],
     items: Iterable[Any] | None = None,
@@ -134,6 +161,8 @@ def multi_process(
     # backend: str = "ray",   # "seq", "ray", or "fastcore"
     backend: Literal['seq', 'ray', 'mp', 'threadpool', 'safe'] = 'mp',
     desc: str | None = None,
+    shared_kwargs: list[str] | None = None,
+    dump_in_thread: bool = True,
     **func_kwargs: Any,
 ) -> list[Any]:
     """
@@ -146,13 +175,55 @@ def multi_process(
         - "threadpool": run in parallel with thread pool
         - "safe": run in parallel with thread pool (explicitly safe for tests)
 
+    shared_kwargs:
+        - Optional list of kwarg names that should be shared via Ray's zero-copy object store
+        - Only works with Ray backend
+        - Useful for large objects (e.g., models, datasets) that should be shared across workers
+        - Example: shared_kwargs=['model', 'tokenizer'] for sharing large ML models
+
+    dump_in_thread:
+        - Whether to dump results to disk in a separate thread (default: True)
+        - If False, dumping is done synchronously, which may block but ensures data is saved before returning
+
     If lazy_output=True, every result is saved to .pkl and
     the returned list contains file paths.
     """
 
     # default backend selection
     if backend is None:
-        backend = 'ray' if _HAS_RAY else 'mp'
+        try:
+            import ray as _ray_module
+            backend = 'ray'
+        except ImportError:
+            backend = 'mp'
+
+    # Validate shared_kwargs
+    if shared_kwargs:
+        # Validate that all shared_kwargs are valid kwargs for the function
+        sig = inspect.signature(func)
+        valid_params = set(sig.parameters.keys())
+
+        for kw in shared_kwargs:
+            if kw not in func_kwargs:
+                raise ValueError(
+                    f"shared_kwargs key '{kw}' not found in provided func_kwargs"
+                )
+            # Check if parameter exists in function signature or if function accepts **kwargs
+            has_var_keyword = any(
+                p.kind == inspect.Parameter.VAR_KEYWORD
+                for p in sig.parameters.values()
+            )
+            if kw not in valid_params and not has_var_keyword:
+                raise ValueError(
+                    f"shared_kwargs key '{kw}' is not a valid parameter for function '{func.__name__}'. "
+                    f"Valid parameters: {valid_params}"
+                )
+
+        # Only allow shared_kwargs with Ray backend
+        if backend != 'ray':
+            raise ValueError(
+                f"shared_kwargs only supported with 'ray' backend, got '{backend}'"
+            )
 
     # unify items
     # unify items and coerce to concrete list so we can use len() and
@@ -169,7 +240,7 @@ def multi_process(
 
     # build cache dir + wrap func
     cache_dir = _build_cache_dir(func, items) if lazy_output else None
-    f_wrapped = wrap_dump(func, cache_dir)
+    f_wrapped = wrap_dump(func, cache_dir, dump_in_thread)
 
     total = len(items)
     if desc:
@@ -181,8 +252,6 @@ def multi_process(
     ) as pbar:
         # ---- sequential backend ----
         if backend == 'seq':
-            pbar.set_postfix_str('backend=seq')
-            results = []
             for x in items:
                 results.append(f_wrapped(x, **func_kwargs))
                 pbar.update(1)
@@ -190,19 +259,46 @@ def multi_process(
 
         # ---- ray backend ----
         if backend == 'ray':
-            pbar.set_postfix_str('backend=ray')
-            ensure_ray(workers, pbar)
+            import ray as _ray_module
 
-            @ray.remote
-            def _task(x):
-                return f_wrapped(x, **func_kwargs)
-
-            refs = [_task.remote(x) for x in items]
+            ensure_ray(workers, pbar)
+            shared_refs = {}
+            regular_kwargs = {}
+
+            if shared_kwargs:
+                for kw in shared_kwargs:
+                    # Put large objects in Ray's object store (zero-copy)
+                    shared_refs[kw] = _ray_module.put(func_kwargs[kw])
+                    pbar.set_postfix_str(f'ray: shared `{kw}` via object store')
+
+                # Remaining kwargs are regular
+                regular_kwargs = {
+                    k: v for k, v in func_kwargs.items()
+                    if k not in shared_kwargs
+                }
+            else:
+                regular_kwargs = func_kwargs
+
+            @_ray_module.remote
+            def _task(x, shared_refs_dict, regular_kwargs_dict):
+                # Dereference shared objects (zero-copy for numpy arrays)
+                import ray as _ray_in_task
+                dereferenced = {k: _ray_in_task.get(v) for k, v in shared_refs_dict.items()}
+                # Merge with regular kwargs
+                all_kwargs = {**dereferenced, **regular_kwargs_dict}
+                return f_wrapped(x, **all_kwargs)
+
+            refs = [
+                _task.remote(x, shared_refs, regular_kwargs) for x in items
+            ]
 
             results = []
+            t_start = time.time()
             for r in refs:
-                results.append(ray.get(r))
+                results.append(_ray_module.get(r))
                 pbar.update(1)
+            t_end = time.time()
+            print(f"Ray processing took {t_end - t_start:.2f}s for {total} items")
             return results
 
         # ---- fastcore backend ----

{speedy_utils-1.1.32.dist-info → speedy_utils-1.1.34.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: speedy-utils
-Version: 1.1.32
+Version: 1.1.34
 Summary: Fast and easy-to-use package for data science
 Project-URL: Homepage, https://github.com/anhvth/speedy
 Project-URL: Repository, https://github.com/anhvth/speedy

{speedy_utils-1.1.32.dist-info → speedy_utils-1.1.34.dist-info}/RECORD

@@ -6,13 +6,13 @@ llm_utils/chat_format/transform.py,sha256=PJ2g9KT1GSbWuAs7giEbTpTAffpU9QsIXyRlbf
 llm_utils/chat_format/utils.py,sha256=M2EctZ6NeHXqFYufh26Y3CpSphN0bdZm5xoNaEJj5vg,1251
 llm_utils/lm/__init__.py,sha256=lFE2DZRpj6eRMo11kx7oRLyYOP2FuDmz08mAcq-cYew,730
 llm_utils/lm/base_prompt_builder.py,sha256=_TzYMsWr-SsbA_JNXptUVN56lV5RfgWWTrFi-E8LMy4,12337
-llm_utils/lm/llm.py,sha256=Ryne4why4VgRCvTJW1SLJciv2pcYYROGP0AwkLVdGPg,16299
+llm_utils/lm/llm.py,sha256=C8Z8l6Ljs7uVX-zabLcDCdTf3fpGxfljaYRM0patHUQ,16469
 llm_utils/lm/llm_signature.py,sha256=vV8uZgLLd6ZKqWbq0OPywWvXAfl7hrJQnbtBF-VnZRU,1244
 llm_utils/lm/lm_base.py,sha256=Bk3q34KrcCK_bC4Ryxbc3KqkiPL39zuVZaBQ1i6wJqs,9437
 llm_utils/lm/mixins.py,sha256=on83g-JO2SpZ0digOpU8mooqFBX6w7Bc-DeGzVoVCX8,14536
 llm_utils/lm/openai_memoize.py,sha256=rYrSFPpgO7adsjK1lVdkJlhqqIw_13TCW7zU8eNwm3o,5185
 llm_utils/lm/signature.py,sha256=K1hvCAqoC5CmsQ0Y_ywnYy2fRb5JzmIK8OS-hjH-5To,9971
-llm_utils/lm/utils.py,sha256=t-RSR-ffOs2HT67JfgqRhbZX9wbxlJNtQMro4ZyuVVM,12461
+llm_utils/lm/utils.py,sha256=dEKFta8S6Mm4LjIctcpFlEGL9RnmLm5DHd2TA70UWuA,12649
 llm_utils/lm/async_lm/__init__.py,sha256=j0xK49ooZ0Dm5GstGGHbmPMrPjd3mOXoJ1H7eAL_Z4g,122
 llm_utils/lm/async_lm/_utils.py,sha256=mB-AueWJJatTx0PXqd_oWc6Kz36cfgDmDTKgiXafCJI,6106
 llm_utils/lm/async_lm/async_llm_task.py,sha256=2PWW4vPW2jYUiGmYFo4-DHrmX5Jm8Iw_1qo6EPL-ytE,18611
@@ -27,7 +27,7 @@ llm_utils/vector_cache/cli.py,sha256=MAvnmlZ7j7_0CvIcSyK4TvJlSRFWYkm4wE7zSq3KR8k
 llm_utils/vector_cache/core.py,sha256=VXuYJy1AX22NHKvIXRriETip5RrmQcNp73-g-ZT774o,30950
 llm_utils/vector_cache/types.py,sha256=CpMZanJSTeBVxQSqjBq6pBVWp7u2-JRcgY9t5jhykdQ,438
 llm_utils/vector_cache/utils.py,sha256=OsiRFydv8i8HiJtPL9hh40aUv8I5pYfg2zvmtDi4DME,1446
-speedy_utils/__imports.py,sha256=PhHqZWwVOKAbbXoWxZLVVyurGmZhui3boQ7Nji002cQ,7795
+speedy_utils/__imports.py,sha256=V0YzkDK4-QkK_IDXY1be6C6_STuNhXAKIp4_dM0coQs,7800
 speedy_utils/__init__.py,sha256=VkKqS4eHXd8YeDu2TAQ3Osqy70RSufUL1sECDoYzqvM,2685
 speedy_utils/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 speedy_utils/common/clock.py,sha256=raLtMGIgzrRej5kUt7hOUm2ZZw2THVPo-q8dMvdZOxw,7354
@@ -41,16 +41,16 @@ speedy_utils/common/utils_io.py,sha256=w9AxMD_8V3Wyo_0o9OtXjVQS8Z3KhxQiOkrl2p8Np
 speedy_utils/common/utils_misc.py,sha256=ZRJCS7OJxybpVm1sasoeCYRW2TaaGCXj4DySYlQeVR8,2227
 speedy_utils/common/utils_print.py,sha256=AGDB7mgJnO00QkJBH6kJb46738q3GzMUZPwtQ248vQw,4763
 speedy_utils/multi_worker/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-speedy_utils/multi_worker/process.py,sha256=O3BpGH7iL_2vh_ezwyHb28lvkVADABpTUnhKHbiEe8I,10542
+speedy_utils/multi_worker/process.py,sha256=jk2K3oNnul1jop4g2U7-6GAekJ4fCyXCbj39WWAwXWQ,14925
 speedy_utils/multi_worker/thread.py,sha256=k4Ff4R2W0Ehet1zJ5nHQOfcsvOjnJzU6A2I18qw7_6M,21320
 speedy_utils/scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 speedy_utils/scripts/mpython.py,sha256=aZvusJLKa3APVhabbFUAEo873VBm8Fym7HKGmVW4LyE,3843
 speedy_utils/scripts/openapi_client_codegen.py,sha256=GModmmhkvGnxljK4KczyixKDrk-VEcLaW5I0XT6tzWo,9657
 vision_utils/README.md,sha256=AIDZZj8jo_QNrEjFyHwd00iOO431s-js-M2dLtVTn3I,5740
-vision_utils/__init__.py,sha256=XsLxy1Fn33Zxu6hTFl3NEWfxGjuQQ-0Wmoh6lU9NZ_o,257
-vision_utils/io_utils.py,sha256=q41pffN632HbMmzcBzfg2Z7DvZZgoAQCdD9jHLqDgjc,26603
-vision_utils/plot.py,sha256=v73onfH8KbGHigw5KStUPqbLyJqIEOvvJaqtaoGKrls,12032
-speedy_utils-1.1.32.dist-info/METADATA,sha256=ElLAOdGyTiqq33ON2WgHX4grgtlFqWRncupD0DivCBk,8048
-speedy_utils-1.1.32.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-speedy_utils-1.1.32.dist-info/entry_points.txt,sha256=1rrFMfqvaMUE9hvwGiD6vnVh98kmgy0TARBj-v0Lfhs,244
-speedy_utils-1.1.32.dist-info/RECORD,,
+vision_utils/__init__.py,sha256=hF54sT6FAxby8kDVhOvruy4yot8O-Ateey5n96O1pQM,284
+vision_utils/io_utils.py,sha256=pI0Va6miesBysJcllK6NXCay8HpGZsaMWwlsKB2DMgA,26510
+vision_utils/plot.py,sha256=HkNj3osA3moPuupP1VguXfPPOW614dZO5tvC-EFKpKM,12028
+speedy_utils-1.1.34.dist-info/METADATA,sha256=diZ6MTVGRDDhsbxoK9eBydHrbW2I6rvYG8lXXzJnJEU,8048
+speedy_utils-1.1.34.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+speedy_utils-1.1.34.dist-info/entry_points.txt,sha256=1rrFMfqvaMUE9hvwGiD6vnVh98kmgy0TARBj-v0Lfhs,244
+speedy_utils-1.1.34.dist-info/RECORD,,

vision_utils/__init__.py

@@ -1,4 +1,11 @@
-from .io_utils import read_images, read_images_cpu, read_images_gpu, ImageMmap, ImageMmapDynamic
+from .io_utils import (
+    ImageMmap,
+    ImageMmapDynamic,
+    read_images,
+    read_images_cpu,
+    read_images_gpu,
+)
 from .plot import plot_images_notebook
 
-__all__ = ['plot_images_notebook', 'read_images_cpu', 'read_images_gpu', 'read_images', 'ImageMmap', 'ImageMmapDynamic']
+
+__all__ = ['plot_images_notebook', 'read_images_cpu', 'read_images_gpu', 'read_images', 'ImageMmap', 'ImageMmapDynamic']

vision_utils/io_utils.py

@@ -3,14 +3,16 @@ from __future__ import annotations
 # type: ignore
 import os
 import time
-from pathlib import Path
-from typing import Sequence, Tuple, TYPE_CHECKING
 from multiprocessing import cpu_count
+from pathlib import Path
+from typing import TYPE_CHECKING, Sequence, Tuple
 
 import numpy as np
 from PIL import Image
+
 from speedy_utils import identify
 
+
 try:
     from torch.utils.data import Dataset
 except ImportError:
@@ -438,12 +440,11 @@ class ImageMmap(Dataset):
                     if img is None:
                         if self.safe:
                             raise ValueError(f"Failed to load image: {path}")
-                        else:
-                            # Failed to load, write zeros
-                            print(f"Warning: Failed to load {path}, using zeros")
-                            mm[global_idx] = np.zeros(
-                                (self.H, self.W, self.C), dtype=self.dtype
-                            )
+                        # Failed to load, write zeros
+                        print(f"Warning: Failed to load {path}, using zeros")
+                        mm[global_idx] = np.zeros(
+                            (self.H, self.W, self.C), dtype=self.dtype
+                        )
                     else:
                         # Clip to valid range and ensure correct dtype
                         if self.dtype == np.uint8:
@@ -625,9 +626,10 @@ class ImageMmapDynamic(Dataset):
           - data file: concatenated flattened images in path order
           - meta: JSON with offsets, shapes, dtype, total_elems, paths, n
         """
-        from tqdm import tqdm
         import json
 
+        from tqdm import tqdm
+
         print(f"Building dynamic mmap cache for {self.n} images...")
         # We don't know total size up front -> write sequentially
         offsets = np.zeros(self.n, dtype=np.int64)
@@ -660,11 +662,10 @@ class ImageMmapDynamic(Dataset):
                     if img is None:
                         if self.safe:
                             raise ValueError(f"Failed to load image: {path}")
-                        else:
-                            print(
-                                f"Warning: Failed to load {path}, storing 1x1x3 zeros"
-                            )
-                            img = np.zeros((1, 1, 3), dtype=self.dtype)
+                        print(
+                            f"Warning: Failed to load {path}, storing 1x1x3 zeros"
+                        )
+                        img = np.zeros((1, 1, 3), dtype=self.dtype)
 
                     # Clip to valid range for uint8
                     if self.dtype == np.uint8:

vision_utils/plot.py

@@ -1,8 +1,8 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union
 
-import numpy as np
 import matplotlib.pyplot as plt
+import numpy as np
 
 
 if TYPE_CHECKING:
@@ -311,7 +311,7 @@ def visualize_tensor(img_tensor, mode='hwc', normalize=True, max_cols=8):
     mpl_available, plt = _check_matplotlib_available()
     if not mpl_available:
         raise ImportError("matplotlib is required for plotting. Install it with: pip install matplotlib")
-    
+
     if mode == 'chw':
         img_tensor = img_tensor.permute(1, 2, 0)
         imgs = [img_tensor]