Skip to content

Instantly share code, notes, and snippets.

@kasperjunge
Created July 29, 2022 08:33
Show Gist options
  • Save kasperjunge/e0edba806cb0f3a1480c0bd418c2c1a7 to your computer and use it in GitHub Desktop.
Save kasperjunge/e0edba806cb0f3a1480c0bd418c2c1a7 to your computer and use it in GitHub Desktop.
Convert at dataset stored as a pandas dataframe to a Hugging Face DatasetDict.
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment