File size: 818 Bytes
a6bbf63
0c746c9
 
 
a6bbf63
 
 
 
 
 
 
 
0c746c9
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
from transformers import RobertaTokenizerFast, RobertaForSequenceClassification
from utils.commons import MODEL_SAVE_PATH, setup_logging

logger = setup_logging("model_loader.log")

class ModelLoader:
    def __init__(self):
        self.tokenizer = None
        self.model = None
        
    def load_model(self):
        """Load model from local path"""
        try:
            logger.info("Loading tokenizer and model...")
            self.tokenizer = RobertaTokenizerFast.from_pretrained(MODEL_SAVE_PATH)
            self.model = RobertaForSequenceClassification.from_pretrained(MODEL_SAVE_PATH)
            logger.info("Tokenizer and model loaded successfully.")
            return self.tokenizer, self.model
        except Exception as e:
            logger.error(f"Error loading model: {e}")
            raise