File size: 259 Bytes
f7a8d72
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
from transformers import AutoModelForSequenceClassification
from src.config import Config


def get_model():
    model = AutoModelForSequenceClassification.from_pretrained(
        Config.MODEL_NAME,
        num_labels=Config.NUM_LABELS
    )
    return model