@@ -0,0 +1,1343 @@
+import fire%0Aimport pandas as pd%0Afrom neuralforecast.data.datasets.m3 import M3, M3Info%0Afrom neuralforecast.data.datasets.m4 import M4, M4Info%0Afrom neuralforecast.data.datasets.tourism import Tourism, TourismInfo%0A%0A%0Adict_datasets = %7B%0A    'Tourism': (Tourism, TourismInfo),%0A    'M3': (M3, M3Info),%0A    'M4': (M4, M4Info)%0A%7D%0A%0Adef get_data(directory: str, dataset: str, group: str, train: bool = True):%0A    if dataset not in dict_datasets.keys():%0A        raise Exception(f'dataset %7Bdataset%7D not found')%0A%0A    dataclass, datainfo = dict_datasets%5Bdataset%5D%0A    if group not in datainfo.groups:%0A        raise Exception(f'group %7Bgroup%7D not found for %7Bdataset%7D')%0A%0A    Y_df, *_ = dataclass.load(directory, group)%0A%0A    horizon = datainfo%5Bgroup%5D.horizon%0A    freq = datainfo%5Bgroup%5D.freq%0A    seasonality = datainfo%5Bgroup%5D.seasonality%0A    Y_df_test = Y_df.groupby('unique_id').tail(horizon)%0A    Y_df = Y_df.drop(Y_df_test.index)%0A%0A    if train:%0A        return Y_df, horizon, freq, seasonality%0A%0A    return Y_df_test, horizon, freq, seasonality%0A%0Adef save_data(dataset: str, group: str, train: bool = True):%0A    df, *_ = get_data('data', dataset, group, train)%0A    if train:%0A        df.to_csv(f'data/%7Bdataset%7D-%7Bgroup%7D.csv', index=False)%0A    else:%0A        df.to_csv(f'data/%7Bdataset%7D-%7Bgroup%7D-test.csv', index=False)%0A%0A%0Aif __name__==%22__main__%22:%0A    fire.Fire(save_data)%0A

add data neuralprophet exp