File size: 294 Bytes
9a1ddeb
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
from transformers import AutoModelForSequenceClassification

def load_model(model_name, num_labels, id2label, label2id):
    return AutoModelForSequenceClassification.from_pretrained(
        model_name,
        num_labels=num_labels,
        id2label=id2label,
        label2id=label2id
    )