Spaces:
Sleeping
Sleeping
File size: 560 Bytes
402d8e8 07cffcd 402d8e8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
from datasets import DatasetDict, Dataset
import pandas as pd
def convert_df_to_dataset_dict(df_train:pd.DataFrame, df_test:pd.DataFrame)-> DatasetDict:
mapper = {"#2_tweet": "tweet", "#3_country_label": "label"}
df_train = df_train.rename(columns=mapper)
df_test = df_test.rename(columns=mapper)
columns_to_keep = ["tweet", "label"]
train_dataset = Dataset.from_pandas(df_train[columns_to_keep])
test_dataset = Dataset.from_pandas(df_test[columns_to_keep])
return DatasetDict({'train': train_dataset, 'test': test_dataset}) |