How to set_epoch with interleave_datasets?
Let's say I have dataset A which has 100k examples, and dataset B which has 100m examples.

I want to train on an interleaved dataset of A+B, with stopping_strategy='all_exhausted' so dataset B doesn't repeat any examples. But every time A is exhausted I want it to be reshuffled (eg. calling set_epoch)

Of course I want to interleave as IterableDatasets / streaming mode so B doesn't have to get tokenized completely at the start.

How could I achieve this? I was thinking something like, if I wrap dataset A in some new IterableDataset with from_generator() and manually call set_epoch before interleaving it? But I'm not sure how to keep the number of shards in that dataset...

Something like

```
dataset_a = load_dataset(...)
dataset_b = load_dataset(...)

def epoch_shuffled_dataset(ds):
  # How to make this maintain the number of shards in ds??
  for epoch in itertools.count():
    ds.set_epoch(epoch)
    yield from iter(ds)

shuffled_dataset_a = IterableDataset.from_generator(epoch_shuffled_dataset, gen_kwargs={'ds': dataset_a})
interleaved = interleave_datasets([shuffled_dataset_a, dataset_b], probs, stopping_strategy='all_exhausted')
```