lhoestq
commented
4 months ago Cool ! Thanks for diving into it :)
Your implementation is great and indeed supports shuffling and batching, you just need to additionally account for state_dict (for dataset checkpointing+resuming)
That being said, I believe the implementation can be made simpler by relying on IterableDataset.map() which already implements all this. Maybe something like
def batch(self, batch_size: int, drop_last_batch: bool = False) -> "IterableDataset":
def batch(unbatched: dict[str, list]) -> dict[str, list]:
return {k: [v] for k, v in unbatched}
return self.map(batch, batched=True, batch_size=batch_size, drop_last_batch=drop_last_batch)And this way no need to reimplement everything !
(my only small concern is that it's not an Arrow-optimized function so it requires the examples to be manipulated as python objects even if the original data is in Arrow format (e.g. when streaming Parquet files) but it's not a big deal and we can see later if we need to optimize this)
lappemic
HuggingFaceDocBuilderDev
github-actions[bot]
I've taken a try at implementing a batched
IterableDatasetas requested in issue #6279. This PR adds a newBatchedExamplesIterableclass and a.batch()method to theIterableDatasetclass.The main changes are:
BatchedExamplesIterablethat groups examples into batches..batch()method forIterableDatasetto easily create batched versions.I'm not sure if this is exactly what you had in mind and also have not fully tested it atm, so I'd really appreciate your feedback. Does this seem like it's heading in the right direction? I'm happy to make any changes or explore different approaches if needed.
Pinging @lhoestq