def load_dataset_task(
    args,
    batch_size=1,
    limit_size=None,
    ref_dataset=None,
):
    """Loads dataset based on args by creating a task"""
    if not args.data or not args.subset or not args.task:
        raise Exception(
            "Please provide necessary arguments to load the dataset - data, subset and task"
        )
    task = tasks.setup_task(args)

    task.load_dataset(args.subset)
    if not limit_size:
        limit_size = len(task.dataset(args.subset))

    iter = task.get_batch_iterator(
        dataset=task.dataset(args.subset), max_sentences=batch_size
    ).next_epoch_itr(shuffle=False)
    dataset = []
    for i, sample in enumerate(iter):
        sample = {
            "id": task.datasets[args.subset].ids[sample["id"].item()],
            "net_input": {
                "src_tokens": sample["net_input"]["src_tokens"],
                "src_lengths": sample["net_input"]["src_lengths"],
            },
        }
        dataset.append(sample)
        if i == limit_size - 1:
            break

    if ref_dataset:
        try:
            ids = get_ids_from_dataset(ref_dataset)
        except Exception as e:
            raise Exception(f"{e} - Cannot extract ids from reference dataset")

        filtered_dataset = []
        for sample in dataset:
            if (
                sample["id"] in ids
                or sample["id"][5:] in ids
                or f"dev_{sample['id']}" in ids
            ):
                filtered_dataset.append(sample)
        dataset = filtered_dataset

    max_len, min_len, avg_len = get_dataset_stats(dataset)
    print(
        f"{args.subset} dataset stats : num_samples={len(dataset)} max_len = {max_len} min_len = {min_len} avg_len = {avg_len}"
    )

    return dataset
Loads dataset based on args by creating a task
understand_icl/fairseq/examples/speech_to_speech/benchmarking/data_utils.py