from transformers import PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput
from .configuration_moderation import FelaModeratorConfig
from .modeling import FELAModerationV2, ModerationConfig

_FIELDS = set(ModerationConfig.__dataclass_fields__.keys())


class FelaModeratorForSequenceClassification(PreTrainedModel):
    config_class = FelaModeratorConfig
    base_model_prefix = "model"

    def __init__(self, config):
        super().__init__(config)
        cfg = ModerationConfig(
            **{k: getattr(config, k) for k in _FIELDS if hasattr(config, k)}
        )
        self.model = FELAModerationV2(cfg, n_tax=config.n_tax)
        self.post_init()

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        labels=None,
        task="taxonomy",
        **kwargs,
    ):
        out = self.model(input_ids, attention_mask, task=task)
        logits = out[task] if isinstance(out, dict) else out
        return SequenceClassifierOutput(logits=logits)