diff --git a/rationai/mlkit/data/datasets/__init__.py b/rationai/mlkit/data/datasets/__init__.py
index bf56e5d..a8d52a8 100644
--- a/rationai/mlkit/data/datasets/__init__.py
+++ b/rationai/mlkit/data/datasets/__init__.py
@@ -1,6 +1,6 @@
+from rationai.mlkit.data.datasets.load_dataset import load_dataset
 from rationai.mlkit.data.datasets.meta_tiled_slides import MetaTiledSlides
 from rationai.mlkit.data.datasets.openslide_tiles_dataset import OpenSlideTilesDataset
-from rationai.mlkit.data.datasets.slides_tiles_loader import SlidesTilesLoader
 
 
-__all__ = ["MetaTiledSlides", "OpenSlideTilesDataset", "SlidesTilesLoader"]
+__all__ = ["MetaTiledSlides", "OpenSlideTilesDataset", "load_dataset"]
diff --git a/rationai/mlkit/data/datasets/load_dataset.py b/rationai/mlkit/data/datasets/load_dataset.py
new file mode 100644
index 0000000..a5d004d
--- /dev/null
+++ b/rationai/mlkit/data/datasets/load_dataset.py
@@ -0,0 +1,130 @@
+from collections.abc import Iterable
+from concurrent.futures import ThreadPoolExecutor
+from itertools import chain
+from pathlib import Path
+from typing import Any
+
+from datasets import Dataset as HFDataset
+from datasets import DatasetDict as HFDatasetDict
+from datasets import concatenate_datasets
+from datasets import load_dataset as hf_load_dataset
+from mlflow.artifacts import download_artifacts
+
+
+def load_dataset[T: (HFDataset, HFDatasetDict)](
+    *,
+    paths: Iterable[Path | str] | None = None,
+    uris: Iterable[str] | None = None,
+    dataset: T | None = None,
+    entities: Iterable[str] | None = None,
+    hf_kwargs: dict[str, Any] | None = None,
+) -> T:
+    """Load parquet-backed HF datasets from local paths and/or MLflow artifact URIs.
+
+    The result is always a ``DatasetDict`` keyed by split name. If the source
+    data has no split structure, everything falls back to ``"train"``.
+    Multiple sources are concatenated split-by-split.
+
+    Args:
+        paths: Local filesystem paths to search for parquet files.
+        uris: MLflow artifact URIs to download and load.
+        dataset: An already-loaded ``DatasetDict`` to merge with the above.
+        entities: If given, only these entity names are loaded. For URIs this
+            also limits what gets downloaded. ``None`` loads all entities.
+        hf_kwargs: Extra kwargs forwarded to ``datasets.load_dataset``.
+            Defaults to ``{"path": "parquet"}``. Any ``"split"`` key is
+            ignored — use the ``splits`` parameter instead.
+    """
+    if not (paths or uris or dataset):
+        raise ValueError("At least one of paths, uris or dataset must be provided.")
+
+    if hf_kwargs is None:
+        hf_kwargs = {"path": "parquet"}
+
+    _entities = frozenset(entities) if entities is not None else None
+    datasets: list[HFDatasetDict] = [dataset] if dataset is not None else []
+
+    if paths or uris:
+        datasets.extend(_load_sources(paths or [], uris or [], hf_kwargs, _entities))
+
+    types = {type(ds) for ds in datasets}
+    if len(types) > 1:
+        raise TypeError(
+            f"Cannot merge datasets of mixed types: {', '.join(t.__name__ for t in types)}"
+        )
+
+    if all(isinstance(ds, HFDataset) for ds in datasets):
+        return concatenate_datasets(datasets)
+
+    return _merge(datasets)
+
+
+def _merge(datasets: Iterable[HFDatasetDict]) -> HFDatasetDict:
+    """Merge a sequence of DatasetDicts split-by-split.
+
+    Splits present in only a subset of the dicts are still included —
+    they are concatenated from whichever dicts contain them.
+    """
+    datasets = list(datasets)
+    all_splits = set().union(*[dd.keys() for dd in datasets])
+    return HFDatasetDict({
+        split: concatenate_datasets([dd[split] for dd in datasets if split in dd])
+        for split in all_splits
+    })
+
+
+def _download_uri(uri: str, entities: frozenset[str] | None) -> list[Path]:
+    """Download an MLflow artifact URI, returning local paths.
+
+    When ``entities`` is specified only the per-entity subdirectories are
+    downloaded (``{uri}/{entity}``), avoiding pulling unused data.
+    Returns a list because a single URI may expand to multiple local paths
+    (one per entity).
+    """
+    if entities is None:
+        return [Path(download_artifacts(artifact_uri=uri))]
+
+    with ThreadPoolExecutor() as executor:
+        return [
+            Path(p)
+            for p in executor.map(
+                lambda entity: download_artifacts(artifact_uri=str(Path(uri, entity))),
+                entities,
+            )
+        ]
+
+
+def _load_sources(
+    paths: Iterable[str | Path],
+    uris: Iterable[str],
+    hf_kwargs: dict[str, Any],
+    entities: frozenset[str] | None,
+) -> list[HFDataset | HFDatasetDict]:
+    """Load parquet files from local paths and MLflow URIs into DatasetDicts.
+
+    Each source is loaded independently.
+    """
+    with ThreadPoolExecutor() as executor:
+        artifacts_paths = list(
+            chain.from_iterable(
+                executor.map(lambda uri: _download_uri(uri, entities), uris)
+            )
+        )
+
+    resolved = [Path(p) for p in (*paths, *artifacts_paths)]
+
+    if not resolved:
+        return []
+
+    try:
+        return [
+            hf_load_dataset(
+                **hf_kwargs,
+                **({"data_dir": str(p)} if p.is_dir() else {"data_files": str(p)}),
+            )
+            for p in resolved
+        ]
+        
+    except Exception as e:
+        msg = "Failed to load Parquet files."
+        raise RuntimeError(msg) from e
diff --git a/rationai/mlkit/data/datasets/meta_tiled_slides.py b/rationai/mlkit/data/datasets/meta_tiled_slides.py
index ff61c79..b323b99 100644
--- a/rationai/mlkit/data/datasets/meta_tiled_slides.py
+++ b/rationai/mlkit/data/datasets/meta_tiled_slides.py
@@ -1,13 +1,12 @@
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
-from pathlib import Path
-from typing import Any, TypeVar
+from typing import TypeVar
 
+import numpy as np
+import pyarrow as pa
 from datasets import Dataset as HFDataset
 from torch.utils.data import ConcatDataset, Dataset
 
-from rationai.mlkit.data.datasets.slides_tiles_loader import SlidesTilesLoader
-
 
 T = TypeVar("T", covariant=True)
 
@@ -25,52 +24,41 @@ class MetaTiledSlides(ConcatDataset[T], ABC):
 
     def __init__(
         self,
-        *,
-        paths: Iterable[Path | str] | None = None,
-        uris: Iterable[str] | None = None,
-        slides_and_tiles: tuple[HFDataset, HFDataset] | None = None,
-        hf_kwargs: dict[str, Any] | None = None,
+        slides: HFDataset,
+        tiles: HFDataset
     ) -> None:
         """Load slides and tiles from MLFlow artifacts.
 
         Args:
-            paths: List of directories to load slides and tiles from. Each
-                directory must include either single files (`slides.parquet`
-                and `tiles.parquet`) or subdirectories (`slides/` and `tiles/`)
-                containing chunked Parquet files.
-            uris: List of MLFlow artifact URIs pointing to folders containing
-                either single files (`slides.parquet` and `tiles.parquet`) or
-                subdirectories (`slides/` and `tiles/`) containing chunked
-                Parquet files.
-            slides_and_tiles: Tuple containing the slides and tiles Datasets.
-            hf_kwargs: Additional keyword arguments to pass to HuggingFace's
-                `load_dataset` function. Defaults to `{"path": "parquet", "split": "train"}`.
+            slides: Dataset containing slide metadata.
+            tiles: Dataset containing tile metadata.
         """
-        self._meta = SlidesTilesLoader(
-            paths=paths,
-            uris=uris,
-            slides_and_tiles=slides_and_tiles,
-            hf_kwargs=hf_kwargs,
-        )
-        self.slides = self._meta.slides
-        self.tiles = self._meta.tiles
-        super().__init__(self.generate_datasets())
+        self.slides = slides
+        self.tiles = tiles
 
-    def filter_tiles_by_slide(self, slide_id: str | bytes) -> HFDataset:
-        """Returns a view of the dataset using a slice or indices.
+        self._slide_id_to_indices = self._build_tile_index(self.tiles)
 
-        This function creates a view of the `self.tiles` dataset that contains only
-        the tiles belonging to the specified slide. It uses the precomputed
-        `_slide_id_to_indices` mapping to efficiently retrieve the relevant tiles
-        without copying data.
+        super().__init__(self.generate_datasets())
 
-        Args:
-            slide_id: The ID of the slide to filter tiles.
+    def filter_tiles_by_slide(self, slide_id: str | bytes) -> HFDataset:
+            """Returns a view of the dataset using a slice or indices.
+    
+            This function creates a view of the `self.tiles` dataset that contains only
+            the tiles belonging to the specified slide. It uses the precomputed
+            `_slide_id_to_indices` mapping to efficiently retrieve the relevant tiles
+            without copying data.
+    
+            Args:
+                slide_id: The ID of the slide to filter tiles.
+    
+            Returns:
+                A view of the tiles dataset containing only the tiles for the specified slide.
+            """
+            tile_indices = self._slide_id_to_indices.get(
+                slide_id, pa.scalar([], type=pa.list_(pa.int64()))
+            )
+            return self.tiles.select(tile_indices.values.to_numpy())
 
-        Returns:
-            A view of the tiles dataset containing only the tiles for the specified slide.
-        """
-        return self._meta.filter_tiles_by_slide(slide_id)
 
     @abstractmethod
     def generate_datasets(self) -> Iterable[Dataset[T]]:
@@ -90,3 +78,43 @@ def generate_datasets(self) -> Iterable[Dataset[T]]:
             )
             ```
         """
+
+    @staticmethod
+    def _build_tile_index(tiles: HFDataset) -> dict[str | bytes, pa.ListScalar]:
+        """Creates a fast lookup table for slide indices.
+
+        This function builds a mapping from `slide_id` to the list of indices in the
+        `tiles` dataset that correspond to that slide.
+
+        Args:
+            tiles: A dataset containing a `slide_id` column.
+
+        Returns:
+            A dictionary mapping each `slide_id` to a list of indices in the `tiles` dataset.
+        """
+        if len(tiles) == 0:
+            return {}
+
+        slide_ids = tiles.data.column("slide_id")
+        num_rows = len(slide_ids)
+
+        # group_by requires the "large" variants for string/binary columns
+        current_type = slide_ids.type
+        if pa.types.is_string(current_type):
+            slide_ids = slide_ids.cast(pa.large_string())
+        elif pa.types.is_binary(current_type):
+            slide_ids = slide_ids.cast(pa.large_binary())
+
+        # np.arange is used here because PyArrow can wrap it with zero-copy overhead
+        row_indices = pa.array(np.arange(num_rows, dtype=np.int64))
+        table = pa.Table.from_arrays(
+            [slide_ids, row_indices], names=["slide_id", "idx"]
+        )
+
+        # "list" aggregates all indices for a given slide_id into a single Arrow List scalar
+        grouped = table.group_by("slide_id").aggregate([("idx", "list")])
+
+        # Keep values as PyArrow ListScalars to avoid materializing them in Python
+        keys = grouped.column("slide_id").to_numpy()
+        values_array = grouped.column("idx_list")
+        return {key: values_array[i] for i, key in enumerate(keys)}
\ No newline at end of file
diff --git a/rationai/mlkit/data/datasets/slides_tiles_loader.py b/rationai/mlkit/data/datasets/slides_tiles_loader.py
deleted file mode 100644
index 7359965..0000000
--- a/rationai/mlkit/data/datasets/slides_tiles_loader.py
+++ /dev/null
@@ -1,192 +0,0 @@
-from collections.abc import Iterable
-from concurrent.futures import ThreadPoolExecutor
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-import pyarrow as pa
-from datasets import Dataset as HFDataset
-from datasets import concatenate_datasets, load_dataset
-from mlflow.artifacts import download_artifacts
-
-
-class SlidesTilesLoader:
-    """Loads and concatenates slides/tiles metadata."""
-
-    def __init__(
-        self,
-        *,
-        paths: Iterable[Path | str] | None = None,
-        uris: Iterable[str] | None = None,
-        slides_and_tiles: tuple[HFDataset, HFDataset] | None = None,
-        hf_kwargs: dict[str, Any] | None = None,
-    ) -> None:
-        """Load slides and tiles from local paths, MLFlow URIs, or preloaded datasets.
-
-        Args:
-            paths: List of directories to load slides and tiles from. Each
-                directory must include either single files (`slides.parquet`
-                and `tiles.parquet`) or subdirectories (`slides/` and `tiles/`)
-                containing chunked Parquet files.
-            uris: List of MLFlow artifact URIs pointing to folders containing
-                either single files (`slides.parquet` and `tiles.parquet`) or
-                subdirectories (`slides/` and `tiles/`) containing chunked
-                Parquet files.
-            slides_and_tiles: Tuple containing the slides and tiles Datasets.
-            hf_kwargs: Additional keyword arguments to pass to HuggingFace's
-                `load_dataset` function. Defaults to `{"path": "parquet", "split": "train"}`.
-        """
-        if not (paths or uris or slides_and_tiles):
-            raise ValueError(
-                "At least one of paths, uris or slides_and_tiles must be provided."
-            )
-
-        if hf_kwargs is None:
-            hf_kwargs = {"path": "parquet", "split": "train"}
-
-        slides = []
-        tiles = []
-
-        if paths or uris:
-            s, t = self.load_slides_and_tiles(paths or [], uris or [], hf_kwargs)
-            slides.append(s)
-            tiles.append(t)
-
-        if slides_and_tiles is not None:
-            slides.append(slides_and_tiles[0])
-            tiles.append(slides_and_tiles[1])
-
-        self.slides = concatenate_datasets(slides) if len(slides) > 1 else slides[0]
-        self.tiles = concatenate_datasets(tiles) if len(tiles) > 1 else tiles[0]
-        self._slide_id_to_indices = self._build_tile_index(self.tiles)
-
-    @staticmethod
-    def _build_tile_index(tiles: HFDataset) -> dict[str | bytes, pa.ListScalar]:
-        """Creates a fast lookup table for slide indices.
-
-        This function builds a mapping from `slide_id` to the list of indices in the
-        `tiles` dataset that correspond to that slide.
-
-        Args:
-            tiles: A dataset containing a `slide_id` column.
-
-        Returns:
-            A dictionary mapping each `slide_id` to a list of indices in the `tiles` dataset.
-        """
-        if len(tiles) == 0:
-            return {}
-
-        # 1. Grab the column directly from the underlying PyArrow Table
-        slide_ids = tiles.data.column("slide_id")
-        num_rows = len(slide_ids)
-
-        # 2. Handle the "Large" type conversion
-        current_type = slide_ids.type
-        if pa.types.is_string(current_type):
-            slide_ids = slide_ids.cast(pa.large_string())
-        elif pa.types.is_binary(current_type):
-            slide_ids = slide_ids.cast(pa.large_binary())
-
-        # 3. Generate sequential row indices
-        # np.arange is used here because PyArrow can wrap it instantly with zero-copy overhead
-        row_indices = pa.array(np.arange(num_rows, dtype=np.int64))
-
-        # 4. Combine them into a lightweight PyArrow Table
-        table = pa.Table.from_arrays(
-            [slide_ids, row_indices], names=["slide_id", "idx"]
-        )
-
-        # 5. Perform the native Arrow groupby and aggregate
-        # The "list" function aggregates all indices for a given slide_id into a single Arrow List scalar
-        grouped = table.group_by("slide_id").aggregate([("idx", "list")])
-
-        # 6. Extract keys to Python, but KEEP values as PyArrow ListScalars
-        keys = grouped.column("slide_id").to_numpy()
-        values_array = grouped.column("idx_list")
-
-        # Map the string key to the PyArrow ListScalar
-        return {key: values_array[i] for i, key in enumerate(keys)}
-
-    def filter_tiles_by_slide(self, slide_id: str | bytes) -> HFDataset:
-        """Returns a view of the dataset using a slice or indices.
-
-        This function creates a view of the `self.tiles` dataset that contains only
-        the tiles belonging to the specified slide. It uses the precomputed
-        `_slide_id_to_indices` mapping to efficiently retrieve the relevant tiles
-        without copying data.
-
-        Args:
-            slide_id: The ID of the slide to filter tiles.
-
-        Returns:
-            A view of the tiles dataset containing only the tiles for the specified slide.
-        """
-        tile_indices = self._slide_id_to_indices.get(
-            slide_id, pa.scalar([], type=pa.list_(pa.int64()))
-        )
-        return self.tiles.select(tile_indices.values.to_numpy())
-
-    @staticmethod
-    def load_slides_and_tiles(
-        paths: Iterable[str | Path], uris: Iterable[str], hf_kwargs: dict[str, Any]
-    ) -> tuple[HFDataset, HFDataset]:
-        """Load slides and tiles parquets from local storage and MLFlow artifacts.
-
-        Args:
-            paths: List of directories to load slides and tiles from. Each
-                directory must include either single files (`slides.parquet`
-                and `tiles.parquet`) or subdirectories (`slides/` and `tiles/`)
-                containing chunked Parquet files.
-            uris: List of MLFlow artifact URIs pointing to folders containing
-                either single files (`slides.parquet` and `tiles.parquet`) or
-                subdirectories (`slides/` and `tiles/`) containing chunked
-                Parquet files.
-            hf_kwargs: Additional keyword arguments to pass to HuggingFace's
-                `load_dataset` function.
-
-        Raises:
-            RuntimeError: If the data cannot be loaded from the specified URIs.
-
-        Returns:
-            A tuple containing the slides and tiles Datasets.
-        """
-        # Parallelize MLFlow downloads (I/O Bound)
-        with ThreadPoolExecutor() as executor:
-            artifacts_paths = list(
-                executor.map(lambda uri: download_artifacts(artifact_uri=uri), uris)
-            )
-
-        search_dirs = [Path(p) for p in (*paths, *artifacts_paths)]
-
-        # Handle empty datasets
-        if not len(search_dirs):
-            return HFDataset.from_dict({}), HFDataset.from_dict({})
-
-        def resolve_search_path(partition: str) -> list[dict[str, str]]:
-            return [
-                {"data_dir": str(path / partition)}
-                if (path / partition).is_dir()
-                else {"data_files": str(path / f"{partition}.parquet")}
-                for path in search_dirs
-            ]
-
-        try:
-            slides_ds = concatenate_datasets(
-                [
-                    load_dataset(**hf_kwargs, **datasource)
-                    for datasource in resolve_search_path("slides")
-                ]
-            )
-
-            tiles_ds = concatenate_datasets(
-                [
-                    load_dataset(**hf_kwargs, **datasource)
-                    for datasource in resolve_search_path("tiles")
-                ]
-            )
-
-            return slides_ds, tiles_ds
-
-        except Exception as e:
-            msg = "Failed to load Parquet files."
-            raise RuntimeError(msg) from e
diff --git a/rationai/mlkit/data/samplers/stratified_batch_sampler.py b/rationai/mlkit/data/samplers/stratified_batch_sampler.py
index 9db18c0..c2603de 100644
--- a/rationai/mlkit/data/samplers/stratified_batch_sampler.py
+++ b/rationai/mlkit/data/samplers/stratified_batch_sampler.py
@@ -67,9 +67,9 @@ def __iter__(self) -> Iterator[list[int]]:
                 if not len(indices[group_idx]):
                     indices.pop(group_idx)
 
-            batch_indices = list(batch_indices)
-            random.shuffle(batch_indices)
-            yield batch_indices
+            batch_list = list(batch_indices)
+            random.shuffle(batch_list)
+            yield batch_list
 
     def __len__(self) -> int:
         return self._indices_size(self.data_indices).sum() // self.batch_size
@@ -89,16 +89,14 @@ class PDMStratifiedBatchSampler(StratifiedBatchSampler):
 
     This sampler is designed to create balanced batches from a DataFrame by
     stratifying samples based on a specified column.
-
-
     """
 
     def __init__(
         self,
         data: pd.DataFrame,
-        stratify_by: None,
+        stratify_by: str | list[str],
         batch_size: int,
-        **kwargs: dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         """Initializes the PDMStratifiedBatchSampler with DataFrame and batch size.
 
diff --git a/rationai/mlkit/data/shard_parquet.py b/rationai/mlkit/data/shard_parquet.py
index 95d2654..a0a4271 100644
--- a/rationai/mlkit/data/shard_parquet.py
+++ b/rationai/mlkit/data/shard_parquet.py
@@ -31,7 +31,6 @@ def shard_parquet(
         AssertionError: If `rows_per_shard` or `row_group_size` are not strictly positive,
             or if `rows_per_shard` is not perfectly divisible by `row_group_size`.
     """
-    # --- Input Validation ---
     assert rows_per_shard > 0, "rows_per_shard must be greater than 0"
     assert row_group_size > 0, "row_group_size must be greater than 0"
 
@@ -40,31 +39,25 @@ def shard_parquet(
         "rows_per_shard must be divisible by row_group_size"
     )
 
-    # --- Setup Output Directory ---
     output_dir = Path(output_dir)
     output_dir.mkdir(parents=True, exist_ok=True)
 
-    # --- Read and Shard Process ---
     with pq.ParquetFile(input_file) as parquet_file:
         _logger.info(f"Total rows in source: {parquet_file.metadata.num_rows}")
 
-        # Initialize tracking variables
         shard_idx = 0
         current_shard_rows = 0
         writer = None
 
         try:
-            # Iterate through the source file in memory-efficient chunks (batches)
             for batch in parquet_file.iter_batches(batch_size=row_group_size):
                 if writer is None:
                     out_path = output_dir / f"shard_{shard_idx:05d}.parquet"
                     writer = pq.ParquetWriter(out_path, batch.schema)
 
-                # Write the current batch
                 writer.write_batch(batch)
                 current_shard_rows += batch.num_rows
 
-                # Check if the current shard has reached its maximum capacity
                 if current_shard_rows >= rows_per_shard:
                     writer.close()
                     writer = None
diff --git a/rationai/mlkit/lightning/loggers/mlflow.py b/rationai/mlkit/lightning/loggers/mlflow.py
index 0f56026..a65a888 100644
--- a/rationai/mlkit/lightning/loggers/mlflow.py
+++ b/rationai/mlkit/lightning/loggers/mlflow.py
@@ -150,9 +150,8 @@ def _log_checkpoint(self, key: str, path: str) -> None:
                 self.run_id, tmpdir, f"{MLFLOW_CHECKPOINT_PATH}/{key}"
             )
 
-    # Ensures that MLFlow logged checkpoints are in sync with those saved by the trainer.
     def _scan_and_log_checkpoints(self, checkpoint_callback: ModelCheckpoint) -> None:
-        """Scan checkpoints and log them to MLFlow if not already logged."""
+        """Scan checkpoints and log them to MLFlow, keeping them in sync with the trainer."""
         checkpoints = self._scan_checkpoints(checkpoint_callback)
 
         logged_checkpoints = {
diff --git a/rationai/mlkit/metrics/aggregated_metric_collection.py b/rationai/mlkit/metrics/aggregated_metric_collection.py
index 8da1dce..3487928 100644
--- a/rationai/mlkit/metrics/aggregated_metric_collection.py
+++ b/rationai/mlkit/metrics/aggregated_metric_collection.py
@@ -1,6 +1,6 @@
 from collections import defaultdict
-from collections.abc import Sequence
-from typing import Any
+from collections.abc import Callable, Sequence
+from typing import Any, cast
 
 from torch import Tensor
 from torchmetrics import Metric, MetricCollection
@@ -79,11 +79,19 @@ def __init__(
         aggregator: Aggregator,
         prefix: str | None = None,
     ) -> None:
-        super().__init__(metrics, prefix=prefix)
-
-        self.aggregators: dict[str, Aggregator] = defaultdict(aggregator.clone)
-
-    def update(  # pylint: disable=arguments-differ
+        super().__init__(
+            cast(
+                "Metric | MetricCollection | Sequence[Metric | MetricCollection] | dict[str, Metric | MetricCollection]",
+                metrics,
+            ),
+            prefix=prefix,
+        )
+
+        self.aggregators: dict[str, Aggregator] = defaultdict(
+            cast("Callable[[], Aggregator]", aggregator.clone)
+        )
+
+    def update(  # type: ignore[override]
         self, preds: Tensor, targets: Tensor, keys: list[str], **kwargs: Any
     ) -> None:
         kwargs_t = ({k: v[i] for k, v in kwargs.items()} for i in range(len(preds)))
diff --git a/rationai/mlkit/metrics/aggregators.py b/rationai/mlkit/metrics/aggregators.py
index 47057ab..a3ffe29 100644
--- a/rationai/mlkit/metrics/aggregators.py
+++ b/rationai/mlkit/metrics/aggregators.py
@@ -42,6 +42,10 @@ def compute(self) -> tuple[Tensor, Tensor]:
 class MeanAggregator(Aggregator):
     """Aggregator to compute the mean of predictions and targets."""
 
+    preds: Tensor
+    targets: Tensor
+    count: Tensor
+
     def __init__(self) -> None:
         super().__init__()
         self.add_state("preds", default=torch.tensor(0.0), dist_reduce_fx="sum")
@@ -58,11 +62,6 @@ def compute(self) -> tuple[Tensor, Tensor]:
 
 
 class HeatmapAggregator(Aggregator):
-    preds: list[Tensor]
-    targets: list[Tensor]
-    xs: list[Tensor]
-    ys: list[Tensor]
-
     """Abstract aggregator covering the prediction heatmap generation.
 
     Arguments:
@@ -70,6 +69,11 @@ class HeatmapAggregator(Aggregator):
         stride_tile (int): Tile stride.
     """
 
+    preds: list[Tensor]
+    targets: list[Tensor]
+    xs: list[Tensor]
+    ys: list[Tensor]
+
     def __init__(
         self,
         extent_tile: int,
diff --git a/rationai/mlkit/metrics/lazy_metric_dict.py b/rationai/mlkit/metrics/lazy_metric_dict.py
index a3781e5..d110788 100644
--- a/rationai/mlkit/metrics/lazy_metric_dict.py
+++ b/rationai/mlkit/metrics/lazy_metric_dict.py
@@ -1,7 +1,7 @@
 from copy import deepcopy
-from typing import Any
+from typing import Any, cast
 
-from deprecated import deprecated
+from deprecated import deprecated  # type: ignore[import-untyped]
 from torch.nn import ModuleDict
 from torchmetrics import Metric, MetricCollection
 
@@ -19,11 +19,15 @@ def update(self, *args: Any, key: str, **kwargs: Any) -> None:  # type: ignore[o
         if key not in self:
             self.add_module(key, deepcopy(self.metric))
 
-        self[key].update(*args, **kwargs)
+        cast("Metric | MetricCollection", self[key]).update(*args, **kwargs)
 
     def compute(self) -> dict[str, Any]:
-        return {k: v.compute() for k, v in self.items() if k != "metric"}
+        return {
+            k: cast("Metric | MetricCollection", v).compute()
+            for k, v in self.items()
+            if k != "metric"
+        }
 
     def reset(self) -> None:
         for metric in self.values():
-            metric.reset()
+            cast("Metric | MetricCollection", metric).reset()
diff --git a/rationai/mlkit/metrics/nested_metric_collection.py b/rationai/mlkit/metrics/nested_metric_collection.py
index 41e3656..1db822f 100644
--- a/rationai/mlkit/metrics/nested_metric_collection.py
+++ b/rationai/mlkit/metrics/nested_metric_collection.py
@@ -35,7 +35,7 @@ class NestedMetricCollection(MetricCollection):
         >>> # Create the NestedMetricCollection, setting 'slide' as the unique identifier for grouping. The class names are provided for multi-class metrics.
         >>> nested_metrics = NestedMetricCollection(
         ...     metrics,
-        ...     key_name="slide"
+        ...     key_name="slide",
         ...     class_names=["A", "B", "C"],
         ... )
 
@@ -84,7 +84,7 @@ def __init__(
         self.class_names = class_names
         self.sep = sep
 
-    def update(  # pylint: disable=arguments-differ
+    def update(  # type: ignore[override]
         self, preds: Tensor, targets: Tensor, keys: list[str]
     ) -> None:
         for pred, target, key in zip(preds, targets, keys, strict=True):
@@ -101,7 +101,7 @@ def update(  # pylint: disable=arguments-differ
                 self[new_name].update(pred.unsqueeze(0), target.unsqueeze(0))
 
     def compute(self) -> dict[str, Any]:
-        divided_metrics = defaultdict(dict)
+        divided_metrics: defaultdict[str, dict[str, Any]] = defaultdict(dict)
         for name, value in super().compute().items():
             key, subkey = name.split(self.sep, maxsplit=1)
 
@@ -112,7 +112,7 @@ def compute(self) -> dict[str, Any]:
                 # handle multi-class metrics without averaging
                 assert len(value.shape) == 1
                 if self.class_names is None:
-                    self.class_names = list(range(len(value)))
+                    self.class_names = [str(i) for i in range(len(value))]
 
                 if len(value) != len(self.class_names):
                     raise ValueError(
diff --git a/rationai/mlkit/mlflow/__init__.py b/rationai/mlkit/mlflow/__init__.py
new file mode 100644
index 0000000..46edf7e
--- /dev/null
+++ b/rationai/mlkit/mlflow/__init__.py
@@ -0,0 +1,4 @@
+from rationai.mlkit.mlflow.parquet_dataset import ParquetDataset, from_parquet
+
+
+__all__ = ["ParquetDataset", "from_parquet"]
diff --git a/rationai/mlkit/mlflow/parquet_dataset.py b/rationai/mlkit/mlflow/parquet_dataset.py
new file mode 100644
index 0000000..51b4bb3
--- /dev/null
+++ b/rationai/mlkit/mlflow/parquet_dataset.py
@@ -0,0 +1,210 @@
+import hashlib
+import json
+import logging
+from functools import cached_property
+from typing import Any
+
+import pyarrow.dataset as ds
+from mlflow.data.dataset import Dataset
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.types.schema import Schema
+from mlflow.types.utils import _infer_schema
+
+
+_logger = logging.getLogger(__name__)
+
+
+class ParquetDataset(Dataset):
+    """Lazy-loaded Parquet dataset (single file or sharded directory) with MLflow tracking."""
+
+    def __init__(
+        self,
+        path: str,
+        source: DatasetSource,
+        target_col: str | None = None,
+        name: str | None = None,
+        digest: str | None = None,
+    ) -> None:
+        """Initializes the ParquetDataset.
+
+        Args:
+            path: Local path or URI to the Parquet file or directory.
+            source: The source of the parquet dataset.
+            target_col: The name of the column representing the target variable. Optional.
+            name: The name of the dataset. If unspecified, a name is automatically generated.
+            digest: The digest (hash) of the dataset. If unspecified, a fast metadata-based
+                digest is automatically computed to avoid hashing massive files.
+        """
+        self._path = path
+        self._target_col = target_col
+        self._ds = ds.dataset(self._path, format="parquet")
+        super().__init__(source=source, name=name, digest=digest)
+
+    # ── MLflow Dataset interface ──────────────────────────────────────────────
+
+    @property
+    def data_type(self) -> str:
+        return "parquet"
+
+    @property
+    def source(self) -> DatasetSource:
+        return self._source
+
+    @property
+    def target_col(self) -> str | None:
+        return self._target_col
+
+    @property
+    def dataset(self) -> ds.Dataset:
+        return self._ds
+
+    # ── digest ────────────────────────────────────────────────────────────────
+
+    def _compute_digest(self) -> str:
+        """Fast metadata-based digest — hashes schema + sorted file paths, never reads data."""
+        hasher = hashlib.md5()
+        hasher.update(str(self._ds.schema).encode())
+        for f in sorted(self._ds.files):
+            hasher.update(f.encode())
+        return hasher.hexdigest()
+
+    # ── profile ───────────────────────────────────────────────────────────────
+
+    @property
+    def profile(self) -> dict[str, Any]:
+        """Row counts and structural metadata read from Parquet footers (no data blocks loaded)."""
+        total_rows = 0
+        for fragment in self._ds.get_fragments():
+            if hasattr(fragment, "metadata") and fragment.metadata is not None:
+                total_rows += fragment.metadata.num_rows
+            else:
+                total_rows += fragment.count_rows()
+        return {
+            "num_files": len(self._ds.files),
+            "total_rows": total_rows,
+            "num_columns": len(self._ds.schema.names),
+            "backend_format": "parquet",
+        }
+
+    # ── schema ────────────────────────────────────────────────────────────────
+
+    @cached_property
+    def schema(self) -> Schema | None:
+        try:
+            import pyarrow as pa
+            from mlflow.types.schema import Array, ColSpec, DataType, TensorSpec
+
+            pa_schema = self._ds.schema
+
+            def _is_scalar(t: pa.DataType) -> bool:
+                return (
+                    pa.types.is_integer(t) or pa.types.is_floating(t)
+                    or pa.types.is_boolean(t) or pa.types.is_string(t)
+                    or pa.types.is_large_string(t) or pa.types.is_binary(t)
+                    or pa.types.is_date(t) or pa.types.is_timestamp(t)
+                )
+
+            def _leaf_dtype(t: pa.DataType) -> DataType | None:
+                if pa.types.is_boolean(t):
+                    return DataType.boolean
+                if pa.types.is_integer(t):
+                    return DataType.long
+                if pa.types.is_floating(t):
+                    return DataType.double
+                return None
+
+            def _array_colspec(field: pa.Field) -> ColSpec | None:
+                """Build a nested Array ColSpec for fixed-shape tensor/list columns.
+
+                MLflow's Schema requires all-ColSpec or all-TensorSpec — never mixed —
+                so array/tensor columns are represented as ColSpec(Array(...)) rather
+                than TensorSpec, to stay homogeneous with the scalar columns below.
+                """
+                t = field.type
+                # Tensor extension types (Ray's ArrowTensorType, PyArrow's native
+                # fixed_shape_tensor) expose .shape plus a leaf element type — Ray uses
+                # .scalar_type, PyArrow uses .value_type. Their .storage_type is a
+                # *flattened* list (e.g. large_list<uint8>), so it can't be used to
+                # recover per-dimension shape and must not be unwrapped for ndims.
+                if isinstance(t, pa.ExtensionType):
+                    shape = getattr(t, "shape", None)
+                    leaf_type = getattr(t, "scalar_type", None) or getattr(t, "value_type", None)
+                    if shape is not None and leaf_type is not None:
+                        leaf = _leaf_dtype(leaf_type)
+                        if leaf is not None:
+                            arr: DataType | Array = leaf
+                            for _ in range(len(shape)):
+                                arr = Array(arr)
+                            return ColSpec(arr, name=field.name)
+                    t = t.storage_type  # unknown extension: fall through as a plain list
+                dims = 0
+                while pa.types.is_fixed_size_list(t) or pa.types.is_list(t) or pa.types.is_large_list(t):
+                    dims += 1
+                    t = t.value_type
+                if dims == 0:
+                    return None
+                leaf = _leaf_dtype(t)
+                if leaf is None:
+                    return None
+                arr = leaf
+                for _ in range(dims):
+                    arr = Array(arr)
+                return ColSpec(arr, name=field.name)
+
+            scalar_fields = [f for f in pa_schema if _is_scalar(f.type)]
+            array_specs = [
+                spec for f in pa_schema
+                if not _is_scalar(f.type)
+                if (spec := _array_colspec(f)) is not None
+            ]
+
+            if not scalar_fields and not array_specs:
+                return None
+
+            specs: list[ColSpec | TensorSpec] = list(array_specs)
+            if scalar_fields:
+                empty_table = pa.table({f.name: pa.array([], type=f.type) for f in scalar_fields})
+                scalar_schema = _infer_schema(empty_table.to_pandas())
+                specs = list(scalar_schema.inputs) + specs
+
+            return Schema(specs)
+        except Exception as exc:
+            _logger.warning("Failed to infer schema for Parquet dataset: %s", exc)
+            return None
+
+    # ── serialisation ─────────────────────────────────────────────────────────
+
+    def to_dict(self) -> dict[str, str]:
+        config = super().to_dict()
+        if self.schema is not None:
+            config["schema"] = json.dumps({"mlflow_colspec": self.schema.to_dict()})
+        config["profile"] = json.dumps(self.profile)
+        return config
+
+
+# ── factory ───────────────────────────────────────────────────────────────────
+
+def from_parquet(
+    path: str,
+    source: str | DatasetSource | None = None,
+    target_col: str | None = None,
+    name: str | None = None,
+    digest: str | None = None,
+) -> ParquetDataset:
+    """Construct a ParquetDataset from a single file or a directory of shards.
+
+    Example::
+
+        dataset = from_parquet("/path/to/tiles/", target_col="tumor")
+        mlflow.log_input(dataset, context="tiles")
+    """
+    from mlflow.data.code_dataset_source import CodeDatasetSource
+    from mlflow.data.dataset_source_registry import resolve_dataset_source
+    from mlflow.tracking.context import registry
+
+    if source is not None:
+        resolved_source = source if isinstance(source, DatasetSource) else resolve_dataset_source(source)
+    else:
+        resolved_source = CodeDatasetSource(tags=registry.resolve_tags())
+
+    return ParquetDataset(path=path, source=resolved_source, target_col=target_col, name=name, digest=digest)
diff --git a/rationai/mlkit/with_cli_args.py b/rationai/mlkit/with_cli_args.py
index ebddc0e..ea309ac 100644
--- a/rationai/mlkit/with_cli_args.py
+++ b/rationai/mlkit/with_cli_args.py
@@ -35,21 +35,13 @@ def with_cli_args(
     def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
         @wraps(func)
         def wrapper(*args: Any, **kwargs: Any) -> Any:
-            # 1. Save original state
             original_argv = sys.argv[:]
-
-            # 2. Deconstruct existing argv
-            # sys.argv[0] is the script name
-            script_name = [sys.argv[0]]
-            user_provided_args = sys.argv[1:]
-
-            # 3. Reconstruct: [Script] + [Start] + [User] + [End]
+            script_name, user_provided_args = sys.argv[:1], sys.argv[1:]
             sys.argv = script_name + prepend + user_provided_args + append
 
             try:
                 return func(*args, **kwargs)
             finally:
-                # 4. Restore original state guarantees safety
                 sys.argv = original_argv
 
         return wrapper
diff --git a/uv.lock b/uv.lock
index 6c5d097..cb69bbf 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2287,7 +2287,7 @@ dependencies = [
 
 [[package]]
 name = "rationai-mlkit"
-version = "0.4.0"
+version = "0.4.1"
 source = { virtual = "." }
 dependencies = [
     { name = "datasets" },