from datasets import load_dataset, Dataset, load_from_disk


def main():
    promptbook = load_dataset('NYUSHPRP/ModelCofferPromptBook', split='train')
    print(promptbook)
    promptbook.save_to_disk('./promptbook')
    #
    # roster = load_dataset('NYUSHPRP/ModelCofferRoster', split='train')
    # roster.save_to_disk('./roster')


def load():
    roster = load_from_disk('./roster')
    print(roster)


def test():
    promptbook = load_from_disk('./promptbook')
    print(promptbook[0]['image'])


if __name__ == '__main__':
    main()