import torch
import torch.nn as nn
from transformers import RobertaModel, RobertaPreTrainedModel


class RobertaMultiTask(RobertaPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.num_labels      = config.num_labels
        self.roberta         = RobertaModel(config)
        self.dropout         = nn.Dropout(config.hidden_dropout_prob)
        self.classifier      = nn.Linear(config.hidden_size, config.num_labels)
        self.span_classifier = nn.Linear(config.hidden_size, 2)
        self.post_init()

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        labels=None,
        span_labels=None
    ):
        outputs = self.roberta(
            input_ids,
            attention_mask=attention_mask
        )
        sequence_output = self.dropout(outputs.last_hidden_state)
        pooled_output   = self.dropout(outputs.pooler_output)

        logits      = self.classifier(pooled_output)
        span_logits = self.span_classifier(sequence_output)

        loss = None
        if labels is not None and span_labels is not None:
            cls_loss = nn.CrossEntropyLoss()(
                logits.view(-1, self.num_labels),
                labels.view(-1)
            )
            span_loss = nn.CrossEntropyLoss(ignore_index=-100)(
                span_logits.view(-1, 2),
                span_labels.view(-1)
            )
            loss = cls_loss + 0.3 * span_loss

        return {
            "loss":        loss,
            "logits":      logits,
            "span_logits": span_logits
        }