from datasets import load_dataset, Dataset, load_from_disk


def main():
    promptbook = load_dataset('NYUSHPRP/ModelCofferPromptBook', split='train')
    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()
    # load()
    test()