kaiko-eva 0.0.0.dev6__py3-none-any.whl

eva/__init__.py

@@ -0,0 +1,33 @@
+"""eva public API."""
+
+from eva.core import (
+    DataLoader,
+    DataloadersSchema,
+    DataModule,
+    DatasetsSchema,
+    HeadModule,
+    InferenceModule,
+    Interface,
+    Trainer,
+    callbacks,
+    data,
+    metrics,
+    models,
+)
+from eva.core.data import datasets
+
+__all__ = [
+    "DataLoader",
+    "DataloadersSchema",
+    "DataModule",
+    "DatasetsSchema",
+    "HeadModule",
+    "InferenceModule",
+    "Interface",
+    "Trainer",
+    "callbacks",
+    "data",
+    "metrics",
+    "models",
+    "datasets",
+]

eva/__main__.py

@@ -0,0 +1,18 @@
+"""Main entry-point module."""
+
+from eva.core.cli import cli
+
+
+def main() -> None:
+    """Main entry-point.
+
+    The CLI fetches the input arguments from `sys.argv`.
+
+    For usage information, execute:
+        $ eva --help
+    """
+    cli()
+
+
+if __name__ == "__main__":
+    main()

eva/__version__.py

@@ -0,0 +1,25 @@
+"""Fetches the version of the library."""
+
+from importlib import metadata
+
+
+def _fetch_version(package_name: str) -> str:
+    """Fetches the version of an installed package.
+
+    If it fails to do so, it returns a "*", indicating
+    that the package has been installed as editable.
+
+    Args:
+        package_name: The name of the package to fetch
+            the version of.
+
+    Returns:
+        A string representing the version of the library.
+    """
+    try:
+        return metadata.version(package_name)
+    except metadata.PackageNotFoundError:
+        return "*"
+
+
+__version__ = _fetch_version("eva")

eva/core/__init__.py

@@ -0,0 +1,19 @@
+"""eva core API."""
+
+from eva.core.cli import cli
+from eva.core.data import DataLoader, DataloadersSchema, DataModule, DatasetsSchema
+from eva.core.interface import Interface
+from eva.core.models import HeadModule, InferenceModule
+from eva.core.trainers import Trainer
+
+__all__ = [
+    "cli",
+    "DataLoader",
+    "DataloadersSchema",
+    "DataModule",
+    "DatasetsSchema",
+    "Interface",
+    "HeadModule",
+    "InferenceModule",
+    "Trainer",
+]

eva/core/callbacks/__init__.py

@@ -0,0 +1,5 @@
+"""Callbacks API."""
+
+from eva.core.callbacks.writers import EmbeddingsWriter
+
+__all__ = ["EmbeddingsWriter"]

eva/core/callbacks/writers/__init__.py

@@ -0,0 +1,5 @@
+"""Callbacks API."""
+
+from eva.core.callbacks.writers.embeddings import EmbeddingsWriter
+
+__all__ = ["EmbeddingsWriter"]

eva/core/callbacks/writers/embeddings.py

@@ -0,0 +1,169 @@
+"""Embeddings writer."""
+
+import csv
+import io
+import os
+from typing import Any, Dict, Sequence
+
+import lightning.pytorch as pl
+import torch
+from lightning.pytorch import callbacks
+from loguru import logger
+from torch import multiprocessing, nn
+from typing_extensions import override
+
+from eva.core.callbacks.writers.typings import QUEUE_ITEM
+from eva.core.models.modules.typings import INPUT_BATCH
+from eva.core.utils import multiprocessing as eva_multiprocessing
+
+
+class EmbeddingsWriter(callbacks.BasePredictionWriter):
+    """Callback for writing generated embeddings to disk."""
+
+    def __init__(
+        self,
+        output_dir: str,
+        backbone: nn.Module | None = None,
+        dataloader_idx_map: Dict[int, str] | None = None,
+        group_key: str | None = None,
+        overwrite: bool = True,
+    ) -> None:
+        """Initializes a new EmbeddingsWriter instance.
+
+        This callback writes the embedding files in a separate process to avoid blocking the
+        main process where the model forward pass is executed.
+
+        Args:
+            output_dir: The directory where the embeddings will be saved.
+            backbone: A model to be used as feature extractor. If `None`,
+                it will be expected that the input batch returns the features directly.
+            dataloader_idx_map: A dictionary mapping dataloader indices to their respective
+                names (e.g. train, val, test).
+            group_key: The metadata key to group the embeddings by. If specified, the
+                embedding files will be saved in subdirectories named after the group_key.
+                If specified, the key must be present in the metadata of the input batch.
+            overwrite: Whether to overwrite the output directory. Defaults to True.
+        """
+        super().__init__(write_interval="batch")
+
+        self._output_dir = output_dir
+        self._backbone = backbone
+        self._dataloader_idx_map = dataloader_idx_map or {}
+        self._group_key = group_key
+        self._overwrite = overwrite
+
+        self._write_queue: multiprocessing.Queue
+        self._write_process: eva_multiprocessing.Process
+
+    @override
+    def on_predict_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        os.makedirs(self._output_dir, exist_ok=self._overwrite)
+        self._initialize_write_process()
+        self._write_process.start()
+
+        if self._backbone is not None:
+            self._backbone = self._backbone.to(pl_module.device)
+            self._backbone.eval()
+
+    @override
+    def write_on_batch_end(
+        self,
+        trainer: pl.Trainer,
+        pl_module: pl.LightningModule,
+        prediction: Any,
+        batch_indices: Sequence[int],
+        batch: INPUT_BATCH,
+        batch_idx: int,
+        dataloader_idx: int,
+    ) -> None:
+        dataset = trainer.predict_dataloaders[dataloader_idx].dataset  # type: ignore
+        _, targets, metadata = INPUT_BATCH(*batch)
+        split = self._dataloader_idx_map.get(dataloader_idx)
+
+        embeddings = self._get_embeddings(prediction)
+        for local_idx, global_idx in enumerate(batch_indices[: len(embeddings)]):
+            input_name, save_name = self._construct_save_name(
+                dataset.filename(global_idx), metadata, local_idx
+            )
+            embeddings_buffer, target_buffer = io.BytesIO(), io.BytesIO()
+            torch.save(embeddings[local_idx].clone(), embeddings_buffer)
+            torch.save(targets[local_idx], target_buffer)  # type: ignore
+            item = QUEUE_ITEM(embeddings_buffer, target_buffer, input_name, save_name, split)
+            self._write_queue.put(item)
+
+        self._write_process.check_exceptions()
+
+    @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}")
+
+    def _initialize_write_process(self) -> None:
+        self._write_queue = multiprocessing.Queue()
+        self._write_process = eva_multiprocessing.Process(
+            target=_process_write_queue, args=(self._write_queue, self._output_dir, self._overwrite)
+        )
+
+    def _get_embeddings(self, prediction: torch.Tensor) -> torch.Tensor:
+        """Returns the embeddings from predictions."""
+        if self._backbone is None:
+            return prediction
+
+        with torch.no_grad():
+            return self._backbone(prediction)
+
+    def _construct_save_name(self, input_name, metadata, local_idx):
+        group_name = metadata[self._group_key][local_idx] if self._group_key else None
+        save_name = os.path.splitext(input_name)[0] + ".pt"
+        if group_name:
+            save_name = os.path.join(group_name, save_name)
+        return input_name, save_name
+
+
+def _process_write_queue(
+    write_queue: multiprocessing.Queue, output_dir: str, overwrite: bool = False
+) -> None:
+    manifest_file, manifest_writer = _init_manifest(output_dir, overwrite)
+    while True:
+        item = write_queue.get()
+        if item is None:
+            break
+
+        prediction_buffer, target_buffer, input_name, save_name, split = QUEUE_ITEM(*item)
+        _save_prediction(prediction_buffer, save_name, output_dir)
+        _update_manifest(target_buffer, input_name, save_name, split, manifest_writer)
+
+    manifest_file.close()
+
+
+def _save_prediction(prediction_buffer: io.BytesIO, save_name: str, output_dir: str) -> None:
+    save_path = os.path.join(output_dir, save_name)
+    prediction = torch.load(io.BytesIO(prediction_buffer.getbuffer()), map_location="cpu")
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    torch.save(prediction, save_path)
+
+
+def _init_manifest(output_dir: str, overwrite: bool = False) -> tuple[io.TextIOWrapper, Any]:
+    manifest_path = os.path.join(output_dir, "manifest.csv")
+    if os.path.exists(manifest_path) and not overwrite:
+        raise FileExistsError(
+            f"Manifest file already exists at {manifest_path}. This likely means that the "
+            "embeddings have been computed before. Consider using `eva fit` instead "
+            "of `eva predict_fit` or `eva predict`."
+        )
+    manifest_file = open(manifest_path, "w", newline="")
+    manifest_writer = csv.writer(manifest_file)
+    manifest_writer.writerow(["origin", "embeddings", "target", "split"])
+    return manifest_file, manifest_writer
+
+
+def _update_manifest(
+    target_buffer: io.BytesIO,
+    input_name: str,
+    save_name: str,
+    split: str | None,
+    manifest_writer,
+) -> None:
+    target = torch.load(io.BytesIO(target_buffer.getbuffer()), map_location="cpu")
+    manifest_writer.writerow([input_name, save_name, target.item(), split])

eva/core/callbacks/writers/typings.py

@@ -0,0 +1,23 @@
+"""Typing definitions for the writer callback functions."""
+
+import io
+from typing import NamedTuple
+
+
+class QUEUE_ITEM(NamedTuple):
+    """The default input batch data scheme."""
+
+    prediction_buffer: io.BytesIO
+    """IO buffer containing the prediction tensor"""
+
+    target_buffer: io.BytesIO
+    """IO buffer containing the target tensor"""
+
+    input_name: str
+    """Name of the original input file that was used to generate the embedding."""
+
+    save_name: str
+    """Name to store the generated embedding"""
+
+    split: str | None
+    """The dataset split the item belongs to (e.g. train, val, test)."""

eva/core/cli/__init__.py

@@ -0,0 +1,5 @@
+"""CLI API."""
+
+from eva.core.cli.cli import cli
+
+__all__ = ["cli"]

eva/core/cli/cli.py

@@ -0,0 +1,19 @@
+"""eva's main cli manager."""
+
+import jsonargparse
+
+from eva.__version__ import __version__
+from eva.core import interface
+from eva.core.cli import logo, setup
+
+
+def cli() -> object:
+    """Main CLI factory."""
+    logo.print_cli_logo()
+    setup.setup()
+    return jsonargparse.CLI(
+        interface.Interface,
+        parser_mode="omegaconf",
+        fail_untyped=False,
+        version=f"version {__version__}",
+    )

eva/core/cli/logo.py

@@ -0,0 +1,38 @@
+"""CLI logos."""
+
+_EVA_LOGO: str = r"""
+      _____   ____ _
+     / _ \ \ / / _` |
+    |  __/\ V / (_| |
+     \___| \_/ \__,_|
+
+         kaiko.ai
+"""
+
+
+ANSI_LOGO_COLOR = "\33[0;34m"
+"""ANSI main logo color."""
+
+ANSI_COLOR_RESET = "\33[0m"
+"""ANSI color reset code."""
+
+
+def _print_logo(
+    logo: str,
+    prefix: str = "",
+    suffix: str = "",
+) -> None:
+    r"""Prints an ASCII terminal art logo in terminal.
+
+    Args:
+        logo: The desired art logo to print.
+        prefix: Characters to add before the logo.
+        suffix: Characters to add after the logo.
+    """
+    colored_logo = f"{ANSI_LOGO_COLOR}{logo}{ANSI_COLOR_RESET}"
+    print(prefix + colored_logo + suffix)
+
+
+def print_cli_logo() -> None:
+    """Prints the CLI logo."""
+    _print_logo(_EVA_LOGO, suffix="\n")

eva/core/cli/setup.py

@@ -0,0 +1,89 @@
+"""Operations which are executed when the CLI is triggered."""
+
+import os
+import sys
+import warnings
+
+import jsonargparse
+import yaml
+from lightning_fabric.utilities import seed as pl_seed
+from loguru import logger
+
+from eva.core.utils import workers
+
+
+def _configure_random_seed(seed: int | None = None) -> None:
+    """Sets the global random seed.
+
+    Args:
+        seed: The seed number to use. If `None`, it will read the seed from
+            `EVA_GLOBAL_SEED` env variable. If `None` and the `EVA_GLOBAL_SEED`
+            env variable is not set, then the seed defaults to `42`. If `None`
+            and the `EVA_GLOBAL_SEED` is set to `False`, it will not set the seed.
+    """
+    effective_seed = seed or os.environ.get("EVA_GLOBAL_SEED", default=42)
+    if isinstance(effective_seed, str):
+        effective_seed = yaml.safe_load(effective_seed)
+        if not isinstance(effective_seed, (bool, int)):
+            raise ValueError(
+                f"Invalid 'EVA_GLOBAL_SEED' value '{effective_seed}'. "
+                "It should be an integer or a boolean value."
+            )
+
+    if isinstance(effective_seed, bool) and effective_seed is False:
+        return
+
+    pl_seed.seed_everything(seed=int(effective_seed), workers=True)
+
+
+def _configure_jsonargparse() -> None:
+    """Configures the `jsonargparse` library."""
+    jsonargparse.set_config_read_mode(
+        urls_enabled=True,
+        fsspec_enabled=True,
+    )
+
+
+def _initialize_logger() -> None:
+    """Initializes, manipulates and customizes the logger.
+
+    This customizable logger can be used by just importing `loguru`
+    from everywhere as follows:
+    >>> from loguru import logger
+    >>> logger.info(...)
+    """
+    logger.remove()
+    logger.add(
+        sys.stderr,
+        format="<blue>{time:HH:mm:ss}</blue>"
+        " :: <bold><level>{level}</level></bold>"
+        " :: {message}",
+        colorize=True,
+        level="INFO",
+    )
+
+
+def _suppress_warnings() -> None:
+    """Suppress all warnings from all subprocesses."""
+    if not sys.warnoptions:
+        warnings.simplefilter("ignore")
+        os.environ["PYTHONWARNINGS"] = "ignore"
+
+
+def _enable_mps_fallback() -> None:
+    """It enables the MPS fallback in torch.
+
+    Note that this action has to take place before importing torch.
+    """
+    if os.environ.get("PYTORCH_ENABLE_MPS_FALLBACK") is None:
+        os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
+
+@workers.main_worker_only
+def setup() -> None:
+    """Sets up the environment before the module is imported."""
+    _configure_random_seed()
+    _configure_jsonargparse()
+    _initialize_logger()
+    _suppress_warnings()
+    _enable_mps_fallback()

eva/core/data/__init__.py

@@ -0,0 +1,14 @@
+"""Data API."""
+
+from eva.core.data.dataloaders import DataLoader
+from eva.core.data.datamodules import DataloadersSchema, DataModule, DatasetsSchema
+from eva.core.data.datasets import Dataset, TorchDataset
+
+__all__ = [
+    "DataLoader",
+    "DataloadersSchema",
+    "DataModule",
+    "DatasetsSchema",
+    "Dataset",
+    "TorchDataset",
+]

eva/core/data/dataloaders/__init__.py

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

eva/core/data/dataloaders/dataloader.py

@@ -0,0 +1,80 @@
+"""Core Dataloader module."""
+
+import dataclasses
+import multiprocessing
+from typing import Callable
+
+from torch.utils.data import dataloader
+
+from eva.core.data import datasets, samplers
+
+
+@dataclasses.dataclass
+class DataLoader:
+    """The `DataLoader` combines a dataset and a sampler.
+
+    It provides an iterable over the given dataset.
+    """
+
+    batch_size: int | None = 1
+    """How many samples per batch to load.
+
+    Set to `None` for iterable dataset where dataset produces batches.
+    """
+
+    shuffle: bool = False
+    """Whether to shuffle the data at every epoch."""
+
+    sampler: samplers.Sampler | None = None
+    """Defines the strategy to draw samples from the dataset.
+
+    Can be any Iterable with `__len__` implemented. If specified, shuffle must
+    not be specified.
+    """
+
+    batch_sampler: samplers.Sampler | None = None
+    """Like `sampler`, but returns a batch of indices at a time.
+
+    Mutually exclusive with `batch_size`, `shuffle`, `sampler` and `drop_last`.
+    """
+
+    num_workers: int = multiprocessing.cpu_count()
+    """How many workers to use for loading the data.
+
+    By default, it will use the number of CPUs available.
+    """
+
+    collate_fn: Callable | None = None
+    """The batching process."""
+
+    pin_memory: bool = True
+    """Will copy Tensors into CUDA pinned memory before returning them."""
+
+    drop_last: bool = False
+    """Drops the last incomplete batch."""
+
+    persistent_workers: bool = True
+    """Will keep the worker processes after a dataset has been consumed once."""
+
+    prefetch_factor: int | None = 2
+    """Number of batches loaded in advance by each worker."""
+
+    def __call__(self, dataset: datasets.TorchDataset) -> dataloader.DataLoader:
+        """Returns the dataloader on the provided dataset.
+
+        Args:
+            dataset: dataset from which to load the data.
+        """
+        return dataloader.DataLoader(
+            dataset=dataset,
+            batch_size=self.batch_size,
+            shuffle=self.shuffle,
+            sampler=self.sampler,
+            batch_sampler=self.batch_sampler,
+            num_workers=self.num_workers,
+            collate_fn=self.collate_fn,
+            pin_memory=self.pin_memory,
+            drop_last=self.drop_last,
+            persistent_workers=self.persistent_workers,
+            prefetch_factor=self.prefetch_factor,
+        )

eva/core/data/datamodules/__init__.py

@@ -0,0 +1,6 @@
+"""Datamodules API."""
+
+from eva.core.data.datamodules.datamodule import DataModule
+from eva.core.data.datamodules.schemas import DataloadersSchema, DatasetsSchema
+
+__all__ = ["DataModule", "DataloadersSchema", "DatasetsSchema"]

eva/core/data/datamodules/call.py

@@ -0,0 +1,33 @@
+"""Helper dataset calling methods."""
+
+from typing import Any, Iterable
+
+from eva.core.data import datasets as datasets_lib
+
+
+def call_method_if_exists(objects: Iterable[Any], /, method: str) -> None:
+    """Calls a desired `method` from the datasets if exists.
+
+    Args:
+        objects: An iterable of objects.
+        method: The dataset method name to call if exists.
+    """
+    for _object in _recursive_iter(objects):
+        if hasattr(_object, method):
+            fn = getattr(_object, method)
+            fn()
+
+
+def _recursive_iter(objects: Iterable[Any], /) -> Iterable[datasets_lib.TorchDataset]:
+    """Iterates thought an iterable of objects and their respective iterable values.
+
+    Args:
+        objects: The objects to iterate from.
+
+    Yields:
+        The individual object class.
+    """
+    for _object in objects:
+        if not isinstance(_object, list):
+            _object = [_object]
+        yield from _object