Skip to content

Commit

Permalink
Fix
Browse files Browse the repository at this point in the history
  • Loading branch information
albertvillanova committed Dec 6, 2023
1 parent 86a3559 commit 5dc2a42
Showing 1 changed file with 5 additions and 8 deletions.
13 changes: 5 additions & 8 deletions src/datasets/dataset_dict.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,19 +56,16 @@ def _check_values_features(self):
f"All datasets in `DatasetDict` should have the same features but features for '{item_a[0]}' and '{item_b[0]}' don't match: {item_a[1].features} != {item_b[1].features}"
)

def __del__(self):
for dataset in self.values():
if hasattr(dataset, "_data"):
del dataset._data
if hasattr(dataset, "_indices"):
del dataset._indices

def __enter__(self):
return self

def __exit__(self, exc_type, exc_val, exc_tb):
# Here `del` is used to del the pyarrow tables. This properly closes the files used for memory mapped tables
self.__del__()
for dataset in self.values():
if hasattr(dataset, "_data"):
del dataset._data
if hasattr(dataset, "_indices"):
del dataset._indices

def __getitem__(self, k) -> Dataset:
if isinstance(k, (str, NamedSplit)) or len(self) == 0:
Expand Down

0 comments on commit 5dc2a42

Please sign in to comment.