classify-intent / components /model_loader.py
aman-s-affinsys's picture
feat: added logging
0c746c9
raw
history blame contribute delete
818 Bytes
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