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
)
|