text: Title: How to set_epoch with interleave_datasets?
  Body: 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')
  ```