File size: 505 Bytes
a6bbf63
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
from transformers import RobertaTokenizerFast, RobertaForSequenceClassification
from utils.commons import MODEL_SAVE_PATH

class ModelLoader:
    def __init__(self):
        self.tokenizer = None
        self.model = None
        
    def load_model(self):
        """Load model from local path"""
        self.tokenizer = RobertaTokenizerFast.from_pretrained(MODEL_SAVE_PATH)
        self.model = RobertaForSequenceClassification.from_pretrained(MODEL_SAVE_PATH)
        return self.tokenizer, self.model