Return distiset from Pipeline.run

[plaguss]

Description

This PR creates a Distiset class which is a wrapper around a dictionary containing the internal datasets.Dataset generated during the Pipeline.run method. Each key corresponds to the leaf_step name in internal DAG, and each value is a datasets.Dataset. It has two methods:

• push_to_hub: to push the Distiset to the hub, where each configuration corresponds to one of the subsets`
• train_test_split: which transforms each one of the internal datasets.Dataset to a datasets.DatasetDict (all the subsets with the same train/test sizes.

The Pipeline.run method after finishing will locate in the cache folder the (parquet) files written via _WriteBuffer and generate the Distiset.

Dummy example:

with Pipeline() as pipeline:
    load_hub_dataset = LoadHubDataset(name="load_dataset", batch_size=8)
    rename_columns = RenameColumns(name="rename_columns", input_batch_size=12)
    generate_response = GenerateResponse(
        name="generate_response", input_batch_size=16
    )

    load_hub_dataset.connect(rename_columns)
    rename_columns.connect(generate_response)

    ds = pipeline.run(
        parameters={
            "load_dataset": {
                "repo_id": "plaguss/test",
                "split": "train",
            },
            "rename_columns": {
                "rename_mappings": {
                    "prompt": "instruction",
                },
            },
        }
    )
# >>> ds["generate_response"]
# Distiset({"leaf_step_1": Dataset(...), "leaf_step_2": Dataset(...)})

Closes #373

[alvarobartt]

LGTM so far, but still need to test it further!

[alvarobartt]

Unrelated to this PR, but maybe it's time to create argilla-internal-testing in Hugging Face Hub? 😆

[plaguss]

yes... the current workflow is weird at best 😆

[gabrielmbmb]

LGTM! We can rename some files, but looks good :)