app/ml_models/classifier.py

from transformers import pipeline
import logging


from app.ml_models.classifier_loader import ClassifierLoader


logger = logging.getLogger(__name__)


class Classifier:
    def __init__(self, model_name: str = "toxic-bert") -> None:
        self.model = None
        self.tokenizer = None
        self.model_name = model_name
        self.classifier = None

    def initialize_classifier(self) -> None:
        loader = ClassifierLoader(self.model_name)
        self.model = loader.load_model()
        self.tokenizer = loader.load_tokenizer()
        self.classifier = pipeline(
            "text-classification",
            model=self.model,
            tokenizer=self.tokenizer,
            device=-1,
            top_k=None,
        )

    def predict_nsfw(self, content: str) -> dict:
        if self.classifier is None:
            raise RuntimeError(
                "Model not initialized. Please call `initialize_classifier()` first."
            )
        results = self.classifier(content)

        prediction = {}
        for result in results[0]:
            prediction[result["label"]] = result["score"]
        
        return prediction