fix: max_size in DataDomain and ShimmerDataset (#180)

add max_size to DataDomain and changed max_size default value to None
instead -1 in ShimmerDataset

Author: leopoldmt

shimmer/data/dataset.py

@@ -55,7 +55,7 @@ def __init__(
         dataset_path: str | Path,
         split: str,
         domain_classes: Mapping[DomainDesc, type[DataDomain]],
-        max_size: int = -1,
+        max_size: int | None = None,
         transforms: Mapping[str, Callable[[Any], Any]] | None = None,
         domain_args: Mapping[str, Any] | None = None,
     ):
@@ -87,14 +87,15 @@ def __init__(
             self.domains[domain.kind] = domain_cls(
                 dataset_path,
                 split,
+                max_size,
                 transform,
                 self.domain_args.get(domain.kind, None),
             )
 
         lengths = {len(domain) for domain in self.domains.values()}
         assert len(lengths) == 1, "Domains have different lengths"
         self.dataset_size = next(iter(lengths))
-        if self.max_size != -1:
+        if self.max_size is not None:
             assert (
                 self.max_size <= self.dataset_size
             ), "Max sizes can only be lower than actual size."

shimmer/data/domain.py

@@ -23,6 +23,7 @@ def __init__(
         self,
         dataset_path: str | Path,
         split: str,
+        max_size: int | None = None,
         transform: Callable[[Any], _T] | None = None,
         additional_args: dict[str, Any] | None = None,
     ) -> None: