Add resplitting functionality to Flower Datasets

[adam-narozniak]

Issue

The datasets downloaded from Hugging Face come with certain splits. Yet, users might want to use different divisions of the whole dataset, which is currently not possible.

Description

fds = FederatedDataset(dataset="mnist", partitioners={"train": 10})
# then I can
fds.load_partition(10, "train")

But what if the dataset had three splits "train", "valid", "test" (not 2 like "mnist").
In that case, you might want to have a single dataset from which 10 partitions are created.

And the reverse might hold true for just a single-split dataset. Sometimes datasets have just train split and need to create a centralized dataset. This is currently impossible.

Proposal

Enable the missing functionality described above (and add tests to make sure they are met).

• Introduce the resplitter keyword in the FederatedDataset
• Enable two major ways of creation of the new dataset (with different splits)
• The first one Callable[[DatasetDict], DatasetDict]] You might perform as sophisticated a change as you wish - just provide this as resplitter to the FederatedDataset (Note: all the checks to use that correctly are on the user side)
• The second option is Dict[Tuple[str, ...], str] (called for convenience ResplitStrategy). Here is an example {("train", "valid"): "bigger_train"}. We create a "bigger_train" split from the "train" and "valid" splits. From this object a Resplitter (newly introduced class) is created. That is essentially Callable[[DatasetDict], DatasetDict]] with additional check if the splits are used correctly (you can use only the existing splits, and you cannot create a new dataset that has two splits with the same name)
  This works as follows:
  First option

  def resplit(dataset: DatasetDict) -> DatasetDict:
      return DatasetDict(
          {
              "bigger_train": concatenate_datasets(
                  [dataset["train"], dataset["valid]]
              ),
              "test": dataset["test"]
          }
      )

  fds = FederatedDataset(
      dataset=self.dataset_name, resplitter=resplit, partitioners={"bigger_train": 100}
  )
  bigger_train = fds.load_full("bigger_train")
  # Or just load a partition
  partition_from_bigger = fds.load_partition("bigger_train": 100)

The second option (that does the same thing) ResplitterStrategy specification => internally Resplitter creation

fds = FederatedDataset(
            dataset=self.dataset_name, resplitter={("train", "valid"): "bigger_train"}, partitioners={"bigger_train": 100}
     )
# if I made a mistake using here 
# e.g. resplitter = {("train", "split-that-does- not-exist"): "bigger_train"}
# I'll get a meaningful error
# similarly if I did sth like that 
# resplitter = {("train",): "new", ("valid", ): "new"}
# there can't be two splits named "new" so I'll get an error
bigger_train = fds.load_full("bigger_train")
# Or just load a partition
partition_from_bigger = fds.load_partition("bigger_train": 100)

[adam-narozniak]

Maybe the MergeResplitter should be simply Merger or SplitMerger? WDYT?

[adam-narozniak]

Also, one more thing. I don't think the Dict[Tuple[str, ...], str] (the Tuple as the key) is that common. It matches more naturally the flow of operation: FROM (some keywords that already exist) -----create---> TO (new keywords). WDYT? We might change that.

[jafermarq]

Very cool functionality. I just left a single comment.

[jafermarq]

Also, one more thing. I don't think the Dict[Tuple[str, ...], str] (the Tuple as the key) is that common. It matches more naturally the flow of operation: FROM (some keywords that already exist) -----create---> TO (new keywords). WDYT? We might change that.

This felt fine when I wast testing the splitter functionality.

[adam-narozniak]

I've switched the keys and values. I also fixed the resplit_dataset_if_needed (I violated the single responsibility principle there and it had complex initialization and then resplit functionality, inti now moved to utils; that also created circular import of Resplitter so I created a common.typing which of which the name can be changed but typing alone didn't work so if common.typing is not ok, then maybe types.py?)