Spaces:
Sleeping
Sleeping
File size: 1,551 Bytes
d68e65a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from app.ml_models.classifier_path_loader import ClassifierPathLoader
import logging
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification, pipeline
logger = logging.getLogger(__name__)
class ClassifierLoader:
def __init__(self, model_name: str):
self.model_name = model_name
self.model = None
self.tokenizer = None
path_loader = ClassifierPathLoader()
path_loader.set_model(self.model_name)
self.model_path = path_loader.get_model_path()
# If model doesn't exist, download it
if not self.model_path.exists():
model_name = "unitary/toxic-bert"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer.save_pretrained(self.model_path)
model.save_pretrained(self.model_path)
def load_model(self):
if self.model is None:
self.model = AutoModelForSequenceClassification.from_pretrained(
self.model_path
)
logger.info("[✅] Model loaded successfully.")
return self.model
def load_tokenizer(self):
if self.tokenizer is None:
self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
logger.info("[✅] Tokenizer loaded successfully.")
return self.tokenizer
|