Comment by rao-v

You could probably speed this up a lot by getting the token log probs and finding your category based on the highest log prob token that is in the category (you may need multiple steps if your categories share token prefixes)