How do I change the classification head of a model?

🤗Transformers
22

Thank you! This fixed my problem too!

It was weird that this didn’t work:
AutoModelForSequenceClassification.from_pretrained(“huggingface/CodeBERTa-language-id”, num_labels=15)

but this did:

config = AutoConfig.from_pretrained(“huggingface/CodeBERTa-language-id”)
config.num_labels = 15
model = AutoModelForSequenceClassification.from_config(config)

Answering to tolgayan, the point is that this gets trained, you’re fine tuning the model.

2 Likes