kaiko-eva 0.3.3__py3-none-any.whl → 0.4.1__py3-none-any.whl

eva/core/callbacks/config.py

@@ -9,11 +9,13 @@ from typing import Any, Dict, List
 import lightning.pytorch as pl
 import yaml
 from lightning_fabric.utilities import cloud_io
+from loguru import logger
 from loguru import logger as cli_logger
 from omegaconf import OmegaConf
 from typing_extensions import TypeGuard, override
 
 from eva.core import loggers
+from eva.core.utils import distributed as dist_utils
 
 
 class ConfigurationLogger(pl.Callback):
@@ -39,8 +41,14 @@ class ConfigurationLogger(pl.Callback):
         pl_module: pl.LightningModule,
         stage: str | None = None,
     ) -> None:
-        log_dir = trainer.log_dir
-        if not _logdir_exists(log_dir):
+        if dist_utils.is_distributed():
+            logger.info("ConfigurationLogger skipped as not supported in distributed mode.")
+            # TODO: Enabling leads to deadlocks in DDP mode, but I could not yet figure out why.
+            return
+
+        if not trainer.is_global_zero or not _logdir_exists(
+            log_dir := trainer.log_dir, self._verbose
+        ):
             return
 
         configuration = _load_submitted_config()
@@ -51,6 +59,10 @@ class ConfigurationLogger(pl.Callback):
 
         save_as = os.path.join(log_dir, self._save_as)
         fs = cloud_io.get_filesystem(log_dir)
+
+        if not fs.exists(log_dir):
+            fs.makedirs(log_dir)
+
         with fs.open(save_as, "w") as output_file:
             yaml.dump(configuration, output_file, sort_keys=False)
 
@@ -126,7 +138,7 @@ def _type_resolver(mapping: Dict[str, Any]) -> Dict[str, Any]:
     for key, value in mapping.items():
         if isinstance(value, dict):
             formatted_value = _type_resolver(value)
-        elif isinstance(value, list) and isinstance(value[0], dict):
+        elif isinstance(value, list) and value and isinstance(value[0], dict):
             formatted_value = [_type_resolver(subvalue) for subvalue in value]
         else:
             try:
@@ -134,10 +146,7 @@ def _type_resolver(mapping: Dict[str, Any]) -> Dict[str, Any]:
                 formatted_value = (
                     value if isinstance(parsed_value, BuiltinFunctionType) else parsed_value
                 )
-
             except Exception:
                 formatted_value = value
-
         mapping[key] = formatted_value
-
     return mapping

eva/core/callbacks/writers/embeddings/base.py

@@ -7,6 +7,7 @@ from typing import Any, Dict, List, Sequence
 
 import lightning.pytorch as pl
 import torch
+import torch.distributed as dist
 from lightning.pytorch import callbacks
 from loguru import logger
 from torch import multiprocessing, nn
@@ -15,6 +16,7 @@ from typing_extensions import override
 from eva.core import utils
 from eva.core.callbacks.writers.embeddings.typings import QUEUE_ITEM
 from eva.core.models.modules.typings import INPUT_BATCH
+from eva.core.utils import distributed as dist_utils
 from eva.core.utils import multiprocessing as eva_multiprocessing
 
 
@@ -58,8 +60,9 @@ class EmbeddingsWriter(callbacks.BasePredictionWriter, abc.ABC):
         self._save_every_n = save_every_n
         self._metadata_keys = metadata_keys or []
 
-        self._write_queue: multiprocessing.Queue
-        self._write_process: eva_multiprocessing.Process
+        self._write_queue: multiprocessing.Queue | None = None
+        self._write_process: eva_multiprocessing.Process | None = None
+        self._is_rank_zero: bool = False
 
     @staticmethod
     @abc.abstractmethod
@@ -78,9 +81,13 @@ class EmbeddingsWriter(callbacks.BasePredictionWriter, abc.ABC):
 
     @override
     def on_predict_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
-        self._check_if_exists()
-        self._initialize_write_process()
-        self._write_process.start()
+        self._is_rank_zero = trainer.is_global_zero
+        if self._is_rank_zero:
+            self._check_if_exists()
+            self._initialize_write_process()
+            if self._write_process is None or self._write_queue is None:
+                raise RuntimeError("Failed to initialize embedding writer process.")
+            self._write_process.start()
 
         if self._backbone is not None:
             self._backbone = self._backbone.to(pl_module.device)
@@ -106,6 +113,7 @@ class EmbeddingsWriter(callbacks.BasePredictionWriter, abc.ABC):
         with torch.no_grad():
             embeddings = self._get_embeddings(prediction)
 
+        queue_items: List[QUEUE_ITEM] = []
         for local_idx, global_idx in enumerate(batch_indices[: len(embeddings)]):
             data_name = dataset.filename(global_idx)
             save_name = os.path.splitext(data_name)[0] + ".pt"
@@ -121,15 +129,41 @@ class EmbeddingsWriter(callbacks.BasePredictionWriter, abc.ABC):
                 split=split,
                 metadata=item_metadata,
             )
-            self._write_queue.put(item)
+            queue_items.append(item)
 
-        self._write_process.check_exceptions()
+        gathered_items = self._gather_queue_items(queue_items)
+        if self._is_rank_zero:
+            for item in gathered_items:
+                self._write_queue.put(item)  # type: ignore
+                self._write_process.check_exceptions()  # type: ignore
 
     @override
     def on_predict_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
-        self._write_queue.put(None)
-        self._write_process.join()
-        logger.info(f"Predictions and manifest saved to {self._output_dir}")
+        if dist_utils.is_distributed():
+            dist.barrier()
+
+        if self._is_rank_zero and self._write_queue is not None:
+            self._write_queue.put(None)
+            if self._write_process is not None:
+                self._write_process.join()
+            logger.info(f"Predictions and manifest saved to {self._output_dir}")
+
+    def _gather_queue_items(self, items: List[QUEUE_ITEM]) -> List[QUEUE_ITEM]:
+        """Gather queue items across distributed ranks, returning only on rank zero."""
+        if not dist_utils.is_distributed():
+            return items
+
+        world_size = dist.get_world_size()
+        object_list: List[List[QUEUE_ITEM]] = [[] for _ in range(world_size)]
+        dist.all_gather_object(object_list, items)
+
+        if self._is_rank_zero:
+            gathered: List[QUEUE_ITEM] = []
+            for rank_items in object_list:
+                gathered.extend(rank_items)
+            return gathered
+
+        return []
 
     def _initialize_write_process(self) -> None:
         self._write_queue = multiprocessing.Queue()

eva/core/cli/setup.py

@@ -59,7 +59,7 @@ def _initialize_logger() -> None:
         " :: <bold><level>{level}</level></bold>"
         " :: {message}",
         colorize=True,
-        level="INFO",
+        level=os.getenv("LOGURU_LEVEL", "INFO"),
     )
 
 

eva/core/data/dataloaders/__init__.py

@@ -1,6 +1,5 @@
 """Dataloaders API."""
 
-from eva.core.data.dataloaders.collate_fn import text_collate_fn
 from eva.core.data.dataloaders.dataloader import DataLoader
 
-__all__ = ["text_collate_fn", "DataLoader"]
+__all__ = ["DataLoader"]

eva/core/data/samplers/classification/balanced.py

@@ -1,9 +1,10 @@
 """Random class sampler for data loading."""
 
 from collections import defaultdict
-from typing import Dict, Iterator, List
+from typing import Dict, Iterator, List, Union
 
 import numpy as np
+import torch
 from loguru import logger
 from typing_extensions import override
 
@@ -32,7 +33,7 @@ class BalancedSampler(SamplerWithDataSource[int]):
         """
         self._num_samples = num_samples
         self._replacement = replacement
-        self._class_indices: Dict[int, List[int]] = defaultdict(list)
+        self._class_indices: Dict[Union[int, str], List[int]] = defaultdict(list)
         self._random_generator = np.random.default_rng(seed)
         self._indices: List[int] = []
 
@@ -62,20 +63,31 @@ class BalancedSampler(SamplerWithDataSource[int]):
         super().set_dataset(data_source)
         self._make_indices()
 
+    def _get_class_idx(self, idx):
+        """Load and validate the class index for a given sample index."""
+        if hasattr(self.data_source, "load_target"):
+            target = self.data_source.load_target(idx)  # type: ignore
+        else:
+            _, target, _ = DataSample(*self.data_source[idx])
+
+        if target is None:
+            raise ValueError("The dataset must return non-empty targets.")
+
+        if isinstance(target, str):
+            return target
+
+        if isinstance(target, torch.Tensor):
+            if target.numel() != 1:
+                raise ValueError("The dataset must return a single & scalar target.")
+            return int(target.item())
+
+        raise ValueError("Unsupported target type. Expected str or tensor-like object.")
+
     def _make_indices(self):
         """Samples the indices for each class in the dataset."""
         self._class_indices.clear()
         for idx in tqdm(range(len(self.data_source)), desc="Fetching class indices for sampler"):
-            if hasattr(self.data_source, "load_target"):
-                target = self.data_source.load_target(idx)  # type: ignore
-            else:
-                _, target, _ = DataSample(*self.data_source[idx])
-            if target is None:
-                raise ValueError("The dataset must return non-empty targets.")
-            if target.numel() != 1:
-                raise ValueError("The dataset must return a single & scalar target.")
-
-            class_idx = int(target.item())
+            class_idx = self._get_class_idx(idx)
             self._class_indices[class_idx].append(idx)
 
         if not self._replacement:

eva/core/data/samplers/random.py

@@ -2,6 +2,7 @@
 
 from typing import Optional
 
+import torch
 from torch.utils import data
 from typing_extensions import override
 
@@ -10,30 +11,36 @@ from eva.core.data.samplers.sampler import SamplerWithDataSource
 
 
 class RandomSampler(data.RandomSampler, SamplerWithDataSource[int]):
-    """Samples elements randomly."""
+    """Samples elements randomly from a MapDataset."""
 
     data_source: datasets.MapDataset  # type: ignore
 
     def __init__(
-        self, replacement: bool = False, num_samples: Optional[int] = None, generator=None
+        self,
+        replacement: bool = False,
+        num_samples: Optional[int] = None,
+        seed: Optional[int] = None,
     ) -> None:
-        """Initializes the random sampler.
+        """Initialize the random sampler.
 
         Args:
-            data_source: dataset to sample from
-            replacement: samples are drawn on-demand with replacement if ``True``, default=``False``
-            num_samples: number of samples to draw, default=`len(dataset)`.
-            generator: Generator used in sampling.
+            replacement: Samples are drawn on-demand with replacement if ``True``, default=``False``
+            num_samples: Number of samples to draw, default=``len(dataset)``.
+            seed: Optional seed for the random number generator.
         """
         self.replacement = replacement
         self._num_samples = num_samples
-        self.generator = generator
+        self._generator = None
+
+        if seed is not None:
+            self._generator = torch.Generator()
+            self._generator.manual_seed(seed)
 
     @override
     def set_dataset(self, data_source: datasets.MapDataset) -> None:
         super().__init__(
             data_source,
             replacement=self.replacement,
-            num_samples=self.num_samples,
-            generator=self.generator,
+            num_samples=self._num_samples,
+            generator=self._generator,
         )

eva/core/interface/interface.py

@@ -132,3 +132,24 @@ class Interface:
             n_runs=trainer.n_runs,
             verbose=trainer.n_runs > 1,
         )
+
+    def validate_test(
+        self,
+        trainer: eva_trainer.Trainer,
+        model: modules.ModelModule,
+        data: datamodules.DataModule,
+    ) -> None:
+        """Runs validation & test stages."""
+        if getattr(data.datasets, "val", None) is None:
+            raise ValueError("The provided data module does not contain a validation dataset.")
+        if getattr(data.datasets, "test", None) is None:
+            raise ValueError("The provided data module does not contain a test dataset.")
+
+        eva_trainer.run_evaluation_session(
+            base_trainer=trainer,
+            base_model=model,
+            datamodule=data,
+            stages=["validate", "test"],
+            n_runs=trainer.n_runs,
+            verbose=trainer.n_runs > 1,
+        )

eva/core/loggers/utils/wandb.py

@@ -5,6 +5,8 @@ from typing import Any, Dict
 
 from loguru import logger
 
+from eva.core.utils import requirements
+
 
 def rename_active_run(name: str) -> None:
     """Renames the current run."""
@@ -12,7 +14,8 @@ def rename_active_run(name: str) -> None:
 
     if wandb.run:
         wandb.run.name = name
-        wandb.run.save()
+        if requirements.below("wandb", "0.21.0"):
+            wandb.run.save()
     else:
         logger.warning("No active wandb run found that could be renamed.")
 

eva/core/models/modules/module.py

@@ -33,8 +33,8 @@ class ModelModule(pl.LightningModule):
         super().__init__()
 
         self._metrics = metrics or self.default_metrics
-        self._postprocess = postprocess or self.default_postprocess
 
+        self.postprocess = postprocess or self.default_postprocess
         self.metrics = metrics_lib.MetricModule.from_schema(self._metrics)
 
     @property
@@ -133,7 +133,7 @@ class ModelModule(pl.LightningModule):
         Returns:
             The updated outputs.
         """
-        self._postprocess(outputs)
+        self.postprocess(outputs)
         return memory.recursive_detach(outputs, to_cpu=self.metrics_device.type == "cpu")
 
     def _forward_and_log_metrics(

eva/core/models/wrappers/base.py

@@ -25,7 +25,7 @@ class BaseModel(nn.Module, Generic[InputType, OutputType]):
 
         self._output_transforms = transforms
 
-        self._model: Callable[..., OutputType] | nn.Module
+        self.model: Callable[..., OutputType] | nn.Module
 
     @override
     def forward(self, tensor: InputType) -> OutputType:
@@ -43,7 +43,7 @@ class BaseModel(nn.Module, Generic[InputType, OutputType]):
         Args:
             tensor: The input tensor to the model.
         """
-        return self._model(tensor)
+        return self.model(tensor)
 
     def _apply_transforms(self, tensor: OutputType) -> OutputType:
         if self._output_transforms is not None:

eva/core/models/wrappers/from_function.py

@@ -41,12 +41,12 @@ class ModelFromFunction(base.BaseModel[torch.Tensor, torch.Tensor]):
         self._arguments = arguments
         self._checkpoint_path = checkpoint_path
 
-        self.load_model()
+        self.model = self.load_model()
 
     @override
-    def load_model(self) -> None:
+    def load_model(self) -> nn.Module:
         class_path = jsonargparse.class_from_function(self._path, func_return=nn.Module)
         model = class_path(**self._arguments or {})
         if self._checkpoint_path is not None:
             _utils.load_model_weights(model, self._checkpoint_path)
-        self._model = model
+        return model

eva/core/models/wrappers/from_torchhub.py

@@ -52,12 +52,12 @@ class TorchHubModel(base.BaseModel[torch.Tensor, torch.Tensor]):
         self._trust_repo = trust_repo
         self._model_kwargs = model_kwargs or {}
 
-        self.load_model()
+        self.model = self.load_model()
 
     @override
-    def load_model(self) -> None:
+    def load_model(self) -> nn.Module:
         """Builds and loads the torch.hub model."""
-        self._model: nn.Module = torch.hub.load(
+        model: nn.Module = torch.hub.load(
             repo_or_dir=self._repo_or_dir,
             model=self._model_name,
             trust_repo=self._trust_repo,
@@ -66,21 +66,23 @@ class TorchHubModel(base.BaseModel[torch.Tensor, torch.Tensor]):
         )  # type: ignore
 
         if self._checkpoint_path:
-            _utils.load_model_weights(self._model, self._checkpoint_path)
+            _utils.load_model_weights(model, self._checkpoint_path)
 
         TorchHubModel.__name__ = self._model_name
 
+        return model
+
     @override
     def model_forward(self, tensor: torch.Tensor) -> torch.Tensor | List[torch.Tensor]:
         if self._out_indices is not None:
-            if not hasattr(self._model, "get_intermediate_layers"):
+            if not hasattr(self.model, "get_intermediate_layers"):
                 raise ValueError(
                     "Only models with `get_intermediate_layers` are supported "
                     "when using `out_indices`."
                 )
 
             return list(
-                self._model.get_intermediate_layers(
+                self.model.get_intermediate_layers(  # type: ignore
                     tensor,
                     self._out_indices,
                     reshape=True,
@@ -89,4 +91,4 @@ class TorchHubModel(base.BaseModel[torch.Tensor, torch.Tensor]):
                 )
             )
 
-        return self._model(tensor)
+        return self.model(tensor)

eva/core/models/wrappers/huggingface.py

@@ -4,6 +4,7 @@ from typing import Any, Callable, Dict
 
 import torch
 import transformers
+from torch import nn
 from typing_extensions import override
 
 from eva.core.models.wrappers import base
@@ -33,12 +34,10 @@ class HuggingFaceModel(base.BaseModel[torch.Tensor, torch.Tensor]):
         self._model_name_or_path = model_name_or_path
         self._model_kwargs = model_kwargs or {}
 
-        self.load_model()
+        self.model = self.load_model()
 
     @override
-    def load_model(self) -> None:
+    def load_model(self) -> nn.Module:
         # Use safetensors to avoid torch.load security vulnerability
         model_kwargs = {"use_safetensors": True, **self._model_kwargs}
-        self._model = transformers.AutoModel.from_pretrained(
-            self._model_name_or_path, **model_kwargs
-        )
+        return transformers.AutoModel.from_pretrained(self._model_name_or_path, **model_kwargs)

eva/core/models/wrappers/onnx.py

@@ -30,21 +30,21 @@ class ONNXModel(base.BaseModel[torch.Tensor, torch.Tensor]):
         self._path = path
         self._device = device
 
-        self.load_model()
+        self.model = self.load_model()
 
     @override
     def load_model(self) -> Any:
         if self._device == "cuda" and not torch.cuda.is_available():
             raise ValueError("Device is set to 'cuda', but CUDA is not available.")
         provider = "CUDAExecutionProvider" if self._device == "cuda" else "CPUExecutionProvider"
-        self._model = ort.InferenceSession(self._path, providers=[provider])  # type: ignore
+        return ort.InferenceSession(self._path, providers=[provider])  # type: ignore
 
     @override
     def model_forward(self, tensor: torch.Tensor) -> torch.Tensor:
         # TODO: Use IO binding to avoid copying the tensor to CPU.
         # https://onnxruntime.ai/docs/api/python/api_summary.html#data-on-device
-        if not isinstance(self._model, ort.InferenceSession):
+        if not isinstance(self.model, ort.InferenceSession):
             raise ValueError("Model is not loaded.")
-        inputs = {self._model.get_inputs()[0].name: tensor.detach().cpu().numpy()}
-        outputs = self._model.run(None, inputs)[0]
+        inputs = {self.model.get_inputs()[0].name: tensor.detach().cpu().numpy()}
+        outputs = self.model.run(None, inputs)[0]
         return torch.from_numpy(outputs).float().to(tensor.device)

eva/core/trainers/trainer.py

@@ -8,6 +8,7 @@ from lightning.pytorch import loggers as pl_loggers
 from lightning.pytorch import trainer as pl_trainer
 from lightning.pytorch.utilities import argparse
 from lightning_fabric.utilities import cloud_io
+from lightning_utilities.core.rank_zero import rank_zero_only
 from typing_extensions import override
 
 from eva.core import loggers as eva_loggers
@@ -30,6 +31,8 @@ class Trainer(pl_trainer.Trainer):
         default_root_dir: str = "logs",
         n_runs: int = 1,
         checkpoint_type: Literal["best", "last"] = "best",
+        accelerator: str = "auto",
+        devices: int = 1,
         **kwargs: Any,
     ) -> None:
         """Initializes the trainer.
@@ -44,9 +47,17 @@ class Trainer(pl_trainer.Trainer):
             n_runs: The amount of runs (fit and evaluate) to perform in an evaluation session.
             checkpoint_type: Wether to load the "best" or "last" checkpoint saved by the checkpoint
                 callback for evaluations on validation & test sets.
+            accelerator: The accelerator to use for training (e.g. "cpu", "gpu").
+            devices: The number of devices (GPUs) to use for training.
             kwargs: Kew-word arguments of ::class::`lightning.pytorch.Trainer`.
         """
-        super().__init__(*args, default_root_dir=default_root_dir, **kwargs)
+        super().__init__(
+            *args,
+            default_root_dir=default_root_dir,
+            accelerator=accelerator,
+            devices=devices,
+            **kwargs,
+        )
 
         self.checkpoint_type = checkpoint_type
         self.n_runs = n_runs
@@ -66,6 +77,7 @@ class Trainer(pl_trainer.Trainer):
     def log_dir(self) -> str | None:
         return self.strategy.broadcast(self._log_dir)
 
+    @rank_zero_only
     def init_logger_run(self, run_id: int | None) -> None:
         """Setup the loggers & log directories when starting a new run.
 

eva/core/utils/__init__.py

@@ -3,5 +3,6 @@
 from eva.core.utils.clone import clone
 from eva.core.utils.memory import to_cpu
 from eva.core.utils.operations import numeric_sort
+from eva.core.utils.paths import home_dir
 
-__all__ = ["clone", "to_cpu", "numeric_sort"]
+__all__ = ["clone", "to_cpu", "numeric_sort", "home_dir"]

eva/core/utils/distributed.py

@@ -0,0 +1,12 @@
+"""Utility functions for distributed training."""
+
+import torch.distributed as dist
+
+
+def is_distributed() -> bool:
+    """Check if current environment is distributed.
+
+    Returns:
+        bool: True if distributed environment (e.g. multiple gpu processes).
+    """
+    return dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1

eva/core/utils/paths.py

@@ -0,0 +1,14 @@
+"""Utility functions for handling paths."""
+
+import os
+
+
+def home_dir():
+    """Get eva's home directory for caching."""
+    torch_home = os.path.expanduser(
+        os.getenv(
+            "EVA_HOME",
+            os.path.join("~/.cache", "eva"),
+        )
+    )
+    return torch_home

eva/core/utils/requirements.py

@@ -3,10 +3,58 @@
 import importlib
 from typing import Dict
 
-from packaging import version
+import packaging.version
 
 
-def check_dependencies(requirements: Dict[str, str]) -> None:
+def fetch_version(name: str) -> str | None:
+    """Fetch the installed version of a package.
+
+    Args:
+        name: The name of the package.
+
+    Returns:
+        A string representing the installed version of the package, or None if not found.
+    """
+    try:
+        module = importlib.import_module(name)
+        return getattr(module, "__version__", None)
+    except ImportError:
+        return None
+
+
+def below(name: str, version: str) -> bool:
+    """Check if the installed version of a package is below a certain version.
+
+    Args:
+        name: The name of the package.
+        version: The version to compare against.
+
+    Returns:
+        True if the installed version is below the specified version, False otherwise.
+    """
+    actual = fetch_version(name)
+    if actual:
+        return packaging.version.parse(actual) < packaging.version.parse(version)
+    return False
+
+
+def above_or_equal(name: str, version: str) -> bool:
+    """Check if the installed version of a package is above a certain version.
+
+    Args:
+        name: The name of the package.
+        version: The version to compare against.
+
+    Returns:
+        True if the installed version is above the specified version, False otherwise.
+    """
+    actual = fetch_version(name)
+    if actual:
+        return packaging.version.parse(actual) >= packaging.version.parse(version)
+    return False
+
+
+def check_min_versions(requirements: Dict[str, str]) -> None:
     """Check installed package versions against requirements dict.
 
     Args:
@@ -17,10 +65,8 @@ def check_dependencies(requirements: Dict[str, str]) -> None:
         ImportError: If any package does not meet the minimum required version.
     """
     for package, min_version in requirements.items():
-        module = importlib.import_module(package)
-        actual = getattr(module, "__version__", None)
-        if actual and not (version.parse(actual) >= version.parse(min_version)):
+        if below(package, min_version):
             raise ImportError(
-                f"Package '{package}' version {actual} does not meet "
+                f"Package '{package}' version {fetch_version(package)} does not meet "
                 f"the minimum required version {min_version}."
             )

eva/language/__init__.py

@@ -1,6 +1,7 @@
 """eva language API."""
 
 try:
+    from eva.language import models
     from eva.language.data import datasets
 except ImportError as e:
     msg = (
@@ -10,4 +11,4 @@ except ImportError as e:
     )
     raise ImportError(str(e) + "\n\n" + msg) from e
 
-__all__ = ["datasets"]
+__all__ = ["models", "datasets"]