models.py

import torch
import torch.nn as nn
from torch.nn import CrossEntropyLoss, MSELoss
import math
import torch.nn.functional as F

from transformers import GPT2Model
from transformers import AutoModelWithLMHead, AutoTokenizer

from from_epitome.roberta import RobertaForTokenClassification, RobertaModel
from from_epitome.modeling_bert import BertEmbeddings, BertLayerNorm, BertModel, BertPreTrainedModel, gelu
from from_epitome.configuration_roberta import RobertaConfig
from from_epitome.file_utils import add_start_docstrings, add_start_docstrings_to_callable


ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP = {
    "roberta-base": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-base-pytorch_model.bin",
    "roberta-large": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-large-pytorch_model.bin",
    "roberta-large-mnli": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-large-mnli-pytorch_model.bin",
    "distilroberta-base": "https://s3.amazonaws.com/models.huggingface.co/bert/distilroberta-base-pytorch_model.bin",
    "roberta-base-openai-detector": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-base-openai-detector-pytorch_model.bin",
    "roberta-large-openai-detector": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-large-openai-detector-pytorch_model.bin",
}


class Norm(nn.Module):
    def __init__(self, d_model, eps = 1e-6):
        super().__init__()

        self.size = d_model
        # create two learnable parameters to calibrate normalisation
        self.alpha = nn.Parameter(torch.ones(self.size))
        self.bias = nn.Parameter(torch.zeros(self.size))
        self.eps = eps
    def forward(self, x):
        norm = self.alpha * (x - x.mean(dim=-1, keepdim=True)) \
        / (x.std(dim=-1, keepdim=True) + self.eps) + self.bias
        return norm


class MultiHeadAttention(nn.Module):
    def __init__(self, heads, d_model, dropout = 0.1):
        super().__init__()

        self.d_model = d_model
        self.d_k = d_model // heads
        self.h = heads

        self.q_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)
        self.out = nn.Linear(d_model, d_model)

    def forward(self, q, k, v, mask=None):

        bs = q.size(0)

        k = self.k_linear(k).view(bs, -1, self.h, self.d_k)
        q = self.q_linear(q).view(bs, -1, self.h, self.d_k)
        v = self.v_linear(v).view(bs, -1, self.h, self.d_k)

        k = k.transpose(1,2)
        q = q.transpose(1,2)
        v = v.transpose(1,2)

        scores = self.attention(q, k, v, self.d_k, mask, self.dropout)

        concat = scores.transpose(1,2).contiguous()\
        .view(bs, -1, self.d_model)

        output = self.out(concat)

        return output

    def attention(self, q, k, v, d_k, mask=None, dropout=None):

        scores = torch.matmul(q, k.transpose(-2, -1)) /  math.sqrt(d_k)

        if mask is not None:
            mask = mask.unsqueeze(1)
            scores = scores.masked_fill(mask == 0, -1e9)

        scores = F.softmax(scores, dim=-1)

        if dropout is not None:
            scores = dropout(scores)

        output = torch.matmul(scores, v)

        return output


class SeekerEncoder(BertPreTrainedModel):
    config_class = RobertaConfig
    pretrained_model_archive_map = ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP
    base_model_prefix = "roberta"

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.roberta = RobertaModel(config)
        self.init_weights()

    def get_input_embeddings(self):
        return self.roberta.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.roberta.embeddings.word_embeddings = value


class ResponderEncoder(BertPreTrainedModel):
    config_class = RobertaConfig
    pretrained_model_archive_map = ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP
    base_model_prefix = "roberta"

    def __init__(self, config):
        super().__init__(config)
        self.roberta = RobertaModel(config)
        self.init_weights()

    def get_input_embeddings(self):
        return self.roberta.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.roberta.embeddings.word_embeddings = value

class BiEncoderAttentionWithRationaleClassification(nn.Module):

    def __init__(self, hidden_dropout_prob=0.2, rationale_num_labels=2, empathy_num_labels=3, hidden_size=768, attn_heads = 1):
        super().__init__()

        self.dropout = nn.Dropout(hidden_dropout_prob)
        self.rationale_classifier = nn.Linear(hidden_size, rationale_num_labels)
        self.attn = MultiHeadAttention(attn_heads, hidden_size)
        self.norm = Norm(hidden_size)
        self.rationale_num_labels = rationale_num_labels
        self.empathy_num_labels = empathy_num_labels
        self.empathy_classifier = RobertaClassificationHead(hidden_size = 768)

        self.apply(self._init_weights)

        self.seeker_encoder = SeekerEncoder.from_pretrained(
                                "roberta-base", # Use the 12-layer BERT model, with an uncased vocab.
                                output_attentions = False, # Whether the model returns attentions weights.
                                output_hidden_states = False)

        self.responder_encoder = ResponderEncoder.from_pretrained(
                                "roberta-base", # Use the 12-layer BERT model, with an uncased vocab.
                                output_attentions = False, # Whether the model returns attentions weights.
                                output_hidden_states = False)


    def _init_weights(self, module):
        """ Initialize the weights """
        if isinstance(module, (nn.Linear, nn.Embedding)):
            # Slightly different from the TF version which uses truncated_normal for initialization
            # cf https://github.com/pytorch/pytorch/pull/5617
            initializer_range=0.02
            module.weight.data.normal_(mean=0.0, std=initializer_range)
        elif isinstance(module, BertLayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
        if isinstance(module, nn.Linear) and module.bias is not None:
            module.bias.data.zero_()


    # @add_start_docstrings_to_callable(ROBERTA_INPUTS_DOCSTRING)
    def forward(
        self,
        input_ids_SP=None,
        input_ids_RP=None,
        attention_mask_SP=None,
        attention_mask_RP=None,
        token_type_ids_SP=None,
        token_type_ids_RP=None,
        position_ids_SP=None,
        position_ids_RP=None,
        head_mask_SP=None,
        head_mask_RP=None,
        inputs_embeds_SP=None,
        inputs_embeds_RP=None,
        empathy_labels=None,
        rationale_labels=None,
        lambda_EI=1,
        lambda_RE=0.1
    ):
        outputs_SP = self.seeker_encoder.roberta(
            input_ids_SP,
            attention_mask=attention_mask_SP,
            token_type_ids=token_type_ids_SP,
            position_ids=position_ids_SP,
            head_mask=head_mask_SP,
            inputs_embeds=inputs_embeds_SP,
        )


        outputs_RP = self.responder_encoder.roberta(
            input_ids_RP,
            attention_mask=attention_mask_RP,
            token_type_ids=token_type_ids_RP,
            position_ids=position_ids_RP,
            head_mask=head_mask_RP,
            inputs_embeds=inputs_embeds_RP,
        )

        sequence_output_SP = outputs_SP[0]
        sequence_output_RP = outputs_RP[0]

        sequence_output_RP = sequence_output_RP + self.dropout(self.attn(sequence_output_RP, sequence_output_SP, sequence_output_SP))

        logits_empathy = self.empathy_classifier(sequence_output_RP[:, 0, :]) # (sequence_output_RP[:, 0, :]) #(torch.tanh(concat_tensor))

        sequence_output = self.dropout(sequence_output_RP)
        logits_rationales = self.rationale_classifier(sequence_output)
        outputs = (logits_empathy,logits_rationales) + outputs_RP[2:]

        loss_rationales = 0.0
        loss_empathy = 0.0

        if rationale_labels is not None:
            loss_fct = CrossEntropyLoss()
            # Only keep active parts of the loss
            if attention_mask_RP is not None:
                active_loss = attention_mask_RP.view(-1) == 1
                active_logits = logits_rationales.view(-1, self.rationale_num_labels)
                active_labels = torch.where(
                    active_loss, rationale_labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(rationale_labels)
                )
                loss_rationales = loss_fct(active_logits, active_labels)
            else:
                loss_rationales = loss_fct(logits_rationales.view(-1, self.rationale_num_labels), rationale_labels.view(-1))

        if empathy_labels is not None:
            loss_fct = CrossEntropyLoss()
            loss_empathy = loss_fct(logits_empathy.view(-1, self.empathy_num_labels), empathy_labels.view(-1))

            loss = lambda_EI * loss_empathy + lambda_RE * loss_rationales

            outputs = (loss, loss_empathy, loss_rationales) + outputs

        return outputs  # (loss), (scores_empathy, scores_rationales), (hidden_states), (attentions)



class RobertaClassificationHead(nn.Module):
    """Head for sentence-level classification tasks."""

    def __init__(self, hidden_dropout_prob=0.1, hidden_size=768, empathy_num_labels=3):
        super().__init__()

        self.dense = nn.Linear(hidden_size, hidden_size)
        self.dropout = nn.Dropout(hidden_dropout_prob)
        self.out_proj = nn.Linear(hidden_size, empathy_num_labels)

    def forward(self, features, **kwargs):
        x = features[:, :]  # take <s> token (equiv. to [CLS])
        x = self.dropout(x)
        x = self.dense(x)
        x = torch.relu(x)
        x = self.dropout(x)
        x = self.out_proj(x)
        return x