Add IterableDataset __repr__ (#6480)

add iterabledataset repr
huggingface · Dec 8, 2023 · 79e94fc · 79e94fc
1 parent dd9044c
commit 79e94fc
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/src/datasets/dataset_dict.py b/src/datasets/dataset_dict.py
@@ -1815,6 +1815,11 @@ def push_to_hub(
 
 
 class IterableDatasetDict(dict):
+    def __repr__(self):
+        repr = "\n".join([f"{k}: {v}" for k, v in self.items()])
+        repr = re.sub(r"^", " " * 4, repr, 0, re.M)
+        return f"IterableDatasetDict({{\n{repr}\n}})"
+
     def with_format(
         self,
         type: Optional[str] = None,

diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py
@@ -1220,6 +1220,9 @@ def __init__(
         self._token_per_repo_id: Dict[str, Union[str, bool, None]] = token_per_repo_id or {}
         _maybe_add_torch_iterable_dataset_parent_class(self.__class__)
 
+    def __repr__(self):
+        return f"IterableDataset({{\n    features: {list(self._info.features.keys()) if self._info.features is not None else 'Unknown'},\n    n_shards: {self.n_shards}\n}})"
+
     def __getstate__(self):
         return self.__dict__