Eventually got this working - appears that the error was in step 2, where I was combining the target column with the tokenized labels before creating the dataset.
I also needed to turn the tokenized labels into a dict.
# This version works
train_dataset = tf.data.Dataset.from_tensor_slices((
    dict(tokens_w_labels),
    df_raw['target'].values
))