Oh I just noticed that the json loader features=
parameter doesn’t do class label encoding so it fails (see issue here).
As a workaround you can do
dataset = load_dataset("json", data_files=data_files)
dataset = dataset.map(features.encode_example, features=features)