Thanks for your help! It was way too late when I wrote this 
It turned out that the easiest way to solve my initial question is this:
config = AutoConfig.from_pretrained(model_checkpoint)
config.num_labels = 2
model = AutoModelForTokenClassification.from_config(config)