The type is guaranteed because you only read the probability from your set of choice tokens

You don't actually use the "next token" that the model chooses

That's what the author is doing in this part

  token_ids = [model.tokenize(text=label.encode(), add_bos=False)[0] for label in labels]
  choice_logits = numpy.asarray([logits[token_id] for token_id in token_ids])
  logprobs = choice_logits - numpy.logaddexp.reduce(choice_logits)
  probabilities = numpy.exp(logprobs)

This works because the model is always producing probabilities for all tokens