Convert at dataset stored as a pandas dataframe to a Hugging Face DatasetDict.

dataframe_to_huggingface_dataset.py

import pandas as pd
from datasets import Dataset, DatasetDict

def dataframe_to_huggingface_dataset(df: pd.DataFrame) -> DatasetDict:
    """Convert at dataframe with split column into Hugging Face DatasetDict.

    Args:
        df (pd.DataFrame): Dataset stored in a pandas dataframe with a 
        column named "split" that tells what split a datapoint belongs to.
        
        Example:
                                                        text  label  split
        0  Drama i X Factor: Teknik ødelægger optræden – ...  flash  train
        1  Hjemvendt Superliga-profil vil lægge stilen om...  sport  train
        2  Voldsom lejlighedsbrand: Tre opgange evakueret...    112  val
        3  Storbritannien: Forældre får livstid for æresd...    112  train
        4  Slut med Payet og Martial: EU-exit kan smadre ...  sport  test

    Returns:
        DatasetDict: Dataset wrapped in a Hugging Face dataset dictionary.
    """

    assert isinstance(df, pd.DataFrame), "df is required to be a pandas dataframe :("
    assert "split" in df, "Dataframe is required to have a column named 'split'."
    
    ds = DatasetDict()
    splits = df["split"].unique().tolist()
    for split in splits:
        ds[split] = Dataset.from_pandas(df.loc[df["split"] == split], split=split, preserve_index=False)
        
    return ds