Finetuning llama for classification

I have some domain specific datasets, so I want to fine-tune llama or other LLMs.
Other than putting classification head on the top, I want to train the model in the generative way, letting models output text that can be thought of labels.

But the issues is that, how can I restrict model generated text to the labels that I want? I’m worried if the model would generate texts that are invalid (text that are not label)

It would be a great help!

You could look at the probabilities on the lm_head logit associated with the first token in each of your answers to turn your lm head into a classifier. It works well if you can formulate your problem such that the answers have unique token ids try yes no answers or multiple choice a), b), c)