import pickle
from conversation import mkdir
from datasets import load_dataset

dataset_id = 'OpenAssistant/oasst2'

mkdir('../datasets')
pickle_filename = './datasets/' + dataset_id.replace('/', '_') + '.pickle'
dataset = load_dataset(dataset_id)
with open(pickle_filename, 'wb') as file:
    pickle.dump(dataset, file)
print('Saved as pickle to ' + pickle_filename)