load.py

from datasets import load_dataset

# Load dataset
dataset = load_dataset("imdb")

# Save to disk
dataset.save_to_disk("./imdb_dataset")

# Later: reload
# from datasets import load_from_disk
# dataset = load_from_disk("./imdb_dataset")