def partition_dataset_by_category(
    dataset_uri: str,
    destination_path: str,
    hf_loading_vars: dict,
    hf_parsing_map: dict,
) -> dict[str, str]:
    """If has_categories is enabled, we partition the dataset into a separate.

    dataset for each category value in the data and write each partition to a
    local file.

    Args:
        dataset_uri (str): Location of dataset.
        destination_path (str): Base destination path, we will write a separate partition off this URI for each category.
        hf_loading_vars (Dict): A dictionary containing keyword arguments to be passed into `load_dataset` if dataset is being pulled from HF.
        hf_parsing_map (Dict): A dictionary containing a mapping from HF columns to ICL dataset keys. The dictionary should be formatted {icl_key:[hf_key1, hf_key1]}.
            Column contents will be concatenated with ' ' separating them. If not included, will load the columns already present in the HF dataset.


    Raises:
        MissingConditionalImportError: If datasets not installed raise exception.
        Exception: If 'category' key missing from dataset, raise exception.

    Returns:
        Dict[str, str]: Mapping of category names to partitioned dataset local files names.
    """
    if dataset_uri.startswith('hf://'):
        dataset_uri = dataset_uri.replace('hf://', '')
        dataset = load_dataset(dataset_uri, **hf_loading_vars)
        assert isinstance(dataset,
                          HFDataset) or isinstance(dataset, IterableDataset)
        if hf_parsing_map:
            dataset_parsing_func = lambda example: {
                k: ' '.join([str(example[col]) for col in v])
                for k, v in hf_parsing_map.items()
            }
            assert hasattr(dataset, 'column_names')
            dataset = dataset.map(
                dataset_parsing_func,
                remove_columns=dataset.column_names,
            )
    else:
        with dist.local_rank_zero_download_and_wait(destination_path):
            if dist.get_local_rank() == 0:
                get_file(dataset_uri, destination_path, overwrite=True)
        dataset = load_dataset(
            'json',
            data_files=destination_path,
            split='train',
            streaming=False,
        )
    assert isinstance(dataset,
                      HFDataset) or isinstance(dataset, IterableDataset)
    assert hasattr(dataset, 'features')
    assert dataset.features is not None
    if 'category' not in dataset.features.keys():
        raise Exception(
            f"""Attempted to partition dataset by `category` \
            but it doesn't have a `category` key. \
            Got keys: {str(list(dataset.features.keys()))}""",
        )
    categories = sorted(
        set(
            dataset['category'],
        ),
    )  # pyright: ignore[reportIndexIssue, reportGeneralTypeIssues]
    output_files = {}
    for cat in categories:
        path = destination_path.split('/')
        cat_dest = '/'.join(path[:-1]) + f'/{cat}_{path[-1]}'
        tmp_path_to_broadcast = str(os.path.abspath(cat_dest))
        gathered_paths = dist.all_gather_object(tmp_path_to_broadcast)
        if dist.get_local_rank() == 0:
            subset = [
                l for l in dataset if
                l['category'] == cat  # pyright: ignore[reportGeneralTypeIssues]
            ]  # pyright: ignore[reportArgumentType, reportCallIssue]
            with open(gathered_paths[0], 'w', encoding='utf8') as f:
                for l in subset:
                    f.write(json.dumps(l, ensure_ascii=False) + '\n')
        output_files[cat] = cat_dest
    return output_files
If has_categories is enabled, we partition the dataset into a separate.

    dataset for each category value in the data and write each partition to a
    local file.

    Args:
        dataset_uri (str): Location of dataset.
        destination_path (str): Base destination path, we will write a separate partition off this URI for each category.
        hf_loading_vars (Dict): A dictionary containing keyword arguments to be passed into `load_dataset` if dataset is being pulled from HF.
        hf_parsing_map (Dict): A dictionary containing a mapping from HF columns to ICL dataset keys. The dictionary should be formatted {icl_key:[hf_key1, hf_key1]}.
            Column contents will be concatenated with ' ' separating them. If not included, will load the columns already present in the HF dataset.


    Raises:
        MissingConditionalImportError: If datasets not installed raise exception.
        Exception: If 'category' key missing from dataset, raise exception.

    Returns:
        Dict[str, str]: Mapping of category names to partitioned dataset local files names.
    
llmfoundry/eval/datasets/in_context_learning_evaluation.py