from datasets import load_dataset
squad = load_dataset("squad", split="train[:5000]")

Split the dataset's train split into a train and test set with the [~datasets.Dataset.train_test_split] method:

squad = squad.train_test_split(test_size=0.2)

Then take a look at an example: