model_sentiment.py

import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data
from transformers import BartTokenizer, BartForConditionalGeneration, BartConfig, BertForSequenceClassification
import math
import numpy as np
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

"""MLP baseline using sentence level eeg"""
# using sent level EEG, MLP baseline for sentiment
class BaselineMLPSentence(nn.Module):
    def __init__(self, input_dim = 840, hidden_dim = 128, output_dim = 3):
        super(BaselineMLPSentence, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim) 
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.relu2 = nn.ReLU()
        self.fc3 = nn.Linear(hidden_dim, output_dim) # positive, negative, neutral  
        self.dropout = nn.Dropout(0.25)

    def forward(self, x):
        out = self.fc1(x)
        out = self.relu1(out)
        out = self.fc2(out)
        out = self.relu2(out)
        out = self.dropout(out)
        out = self.fc3(out)
        return out


"""bidirectional LSTM baseline using word level eeg"""
class BaselineLSTM(nn.Module):
    def __init__(self, input_dim = 840, hidden_dim = 256, output_dim = 3, num_layers = 1):
        super(BaselineLSTM, self).__init__()
        
        self.hidden_dim = hidden_dim

        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers = 1, batch_first = True, bidirectional = True)

        self.hidden2sentiment = nn.Linear(hidden_dim*2, output_dim)

    def forward(self, x_packed):
        # input: (N,seq_len,input_dim)
        # print(x_packed.data.size())
        lstm_out, _ = self.lstm(x_packed)
        last_hidden_state = pad_packed_sequence(lstm_out, batch_first = True)[0][:,-1,:]
        # print(last_hidden_state.size())
        out = self.hidden2sentiment(last_hidden_state)
        return out

""" Bert Baseline: Finetuning from a pretrained language model Bert"""
class NaiveFineTunePretrainedBert(nn.Module):
    def __init__(self, input_dim = 840, hidden_dim = 768, output_dim = 3, pretrained_checkpoint = None):
        super(NaiveFineTunePretrainedBert, self).__init__()
        # mapping hidden states dimensioin
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.pretrained_Bert = BertForSequenceClassification.from_pretrained('bert-base-cased',num_labels=3)
        
        if pretrained_checkpoint is not None:
            self.pretrained_Bert.load_state_dict(torch.load(pretrained_checkpoint))

    def forward(self, input_embeddings_batch, input_masks_batch, labels):
        embedding = F.relu(self.fc1(input_embeddings_batch))
        out = self.pretrained_Bert(inputs_embeds = embedding, attention_mask = input_masks_batch, labels = labels, return_dict = True)
        return out

""" Finetuning from a pretrained language model BART, two step training"""
class FineTunePretrainedTwoStep(nn.Module):
    def __init__(self, pretrained_layers, in_feature = 840, d_model = 1024, additional_encoder_nhead=8, additional_encoder_dim_feedforward = 2048):
        super(FineTunePretrainedTwoStep, self).__init__()
        
        self.pretrained_layers = pretrained_layers
        # additional transformer encoder, following BART paper about 
        self.additional_encoder_layer = nn.TransformerEncoderLayer(d_model=in_feature, nhead=additional_encoder_nhead,  dim_feedforward = additional_encoder_dim_feedforward, batch_first=True)
        self.additional_encoder = nn.TransformerEncoder(self.additional_encoder_layer, num_layers=6)
        
        # NOTE: add positional embedding?
        # print('[INFO]adding positional embedding')
        # self.positional_embedding = PositionalEncoding(in_feature)

        self.fc1 = nn.Linear(in_feature, d_model)

    def forward(self, input_embeddings_batch, input_masks_batch, input_masks_invert, labels):
        """input_embeddings_batch: batch_size*Seq_len*840"""
        """input_mask: 1 is not masked, 0 is masked"""
        """input_masks_invert: 1 is masked, 0 is not masked"""
        """labels: sentitment labels 0,1,2"""
        
        # NOTE: add positional embedding?
        # input_embeddings_batch = self.positional_embedding(input_embeddings_batch) 

        # use src_key_padding_masks
        encoded_embedding = self.additional_encoder(input_embeddings_batch, src_key_padding_mask = input_masks_invert) 
        # encoded_embedding = self.additional_encoder(input_embeddings_batch) 
        
        encoded_embedding = F.relu(self.fc1(encoded_embedding))
        out = self.pretrained_layers(inputs_embeds = encoded_embedding, attention_mask = input_masks_batch, return_dict = True, labels = labels)                    
        
        return out

""" Zero-shot sentiment discovery using a finetuned generation model and a sentiment model pretrained on text """
class ZeroShotSentimentDiscovery(nn.Module):
    def __init__(self, brain2text_translator, sentiment_classifier, translation_tokenizer, sentiment_tokenizer, device = 'cpu'):
        # only for inference
        super(ZeroShotSentimentDiscovery, self).__init__()
        
        self.brain2text_translator = brain2text_translator
        self.sentiment_classifier = sentiment_classifier
        self.translation_tokenizer = translation_tokenizer
        self.sentiment_tokenizer = sentiment_tokenizer
        self.device = device
    

    def forward(self, input_embeddings_batch, input_masks_batch, input_masks_invert, target_ids_batch_converted, sentiment_labels):
        """input_embeddings_batch: batch_size*Seq_len*840"""
        """input_mask: 1 is not masked, 0 is masked"""
        """input_masks_invert: 1 is masked, 0 is not masked"""
        """labels: sentitment labels 0,1,2"""
        
        def logits2PredString(logits):
            probs = logits[0].softmax(dim = 1)
            # print('probs size:', probs.size())
            values, predictions = probs.topk(1)
            # print('predictions before squeeze:',predictions.size())
            predictions = torch.squeeze(predictions)
            predict_string = self.translation_tokenizer.decode(predictions)
            return predict_string

        # only works on batch is one
        assert input_embeddings_batch.size()[0] == 1

        seq2seqLMoutput = self.brain2text_translator(input_embeddings_batch, input_masks_batch, input_masks_invert, target_ids_batch_converted)
        predict_string = logits2PredString(seq2seqLMoutput.logits)
        predict_string = predict_string.split('</s></s>')[0]
        predict_string = predict_string.replace('<s>','')
        print('predict string:', predict_string)
        re_tokenized = self.sentiment_tokenizer(predict_string, return_tensors='pt', return_attention_mask = True)
        input_ids = re_tokenized['input_ids'].to(self.device) # batch = 1
        attn_mask = re_tokenized['attention_mask'].to(self.device) # batch = 1

        out = self.sentiment_classifier(input_ids = input_ids, attention_mask = attn_mask, return_dict = True, labels = sentiment_labels)

        return out


""" Miscellaneous: jointly learn generation and classification (not working well) """
class BartClassificationHead(nn.Module):
    # from transformers: https://huggingface.co/transformers/_modules/transformers/models/bart/modeling_bart.html
    """Head for sentence-level classification tasks."""
    def __init__(
        self,
        input_dim: int,
        inner_dim: int,
        num_classes: int,
        pooler_dropout: float,
    ):
        super().__init__()
        self.dense = nn.Linear(input_dim, inner_dim)
        self.dropout = nn.Dropout(p=pooler_dropout)
        self.out_proj = nn.Linear(inner_dim, num_classes)

    def forward(self, hidden_states: torch.Tensor):
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.dense(hidden_states)
        hidden_states = torch.tanh(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.out_proj(hidden_states)
        return hidden_states

class JointBrainTranslatorSentimentClassifier(nn.Module):
    def __init__(self, pretrained_layers, in_feature = 840, d_model = 1024, additional_encoder_nhead=8, additional_encoder_dim_feedforward = 2048, num_labels = 3):
        super(JointBrainTranslatorSentimentClassifier, self).__init__()
        
        self.pretrained_generator = pretrained_layers
        # additional transformer encoder, following BART paper about 
        self.additional_encoder_layer = nn.TransformerEncoderLayer(d_model=in_feature, nhead=additional_encoder_nhead,  dim_feedforward = additional_encoder_dim_feedforward, batch_first=True)
        self.additional_encoder = nn.TransformerEncoder(self.additional_encoder_layer, num_layers=6)
        self.fc1 = nn.Linear(in_feature, d_model)
        self.num_labels = num_labels

        self.pooler = Pooler(d_model)
        self.classifier = BartClassificationHead(input_dim = d_model, inner_dim = d_model, num_classes = num_labels, pooler_dropout = pretrained_layers.config.classifier_dropout)

    def forward(self, input_embeddings_batch, input_masks_batch, input_masks_invert, target_ids_batch_converted, sentiment_labels):
        """input_embeddings_batch: batch_size*Seq_len*840"""
        """input_mask: 1 is not masked, 0 is masked"""
        """input_masks_invert: 1 is masked, 0 is not masked"""
        
        # NOTE: add positional embedding?
        # input_embeddings_batch = self.positional_embedding(input_embeddings_batch) 

        # use src_key_padding_masks
        encoded_embedding = self.additional_encoder(input_embeddings_batch, src_key_padding_mask = input_masks_invert) 
        
        # encoded_embedding = self.additional_encoder(input_embeddings_batch) 
        encoded_embedding = F.relu(self.fc1(encoded_embedding))
        LMoutput = self.pretrained_generator(inputs_embeds = encoded_embedding, attention_mask = input_masks_batch, return_dict = True, labels = target_ids_batch_converted, output_hidden_states = True)                    
        hidden_states = LMoutput.decoder_hidden_states # N, seq_len, hidden_dim
        # print('hidden states len:', len(hidden_states))
        last_hidden_states = hidden_states[-1]
        # print('last hidden states size:', last_hidden_states.size())
        sentence_representation = self.pooler(last_hidden_states)
 
        classification_logits = self.classifier(sentence_representation) 
        loss_fct = nn.CrossEntropyLoss()
        classification_loss = loss_fct(classification_logits.view(-1, self.num_labels), sentiment_labels.view(-1))
        classification_output = {'loss':classification_loss,'logits':classification_logits}
        # print('successful one forward!!!!')
        return LMoutput, classification_output


""" helper modules """
# modified from BertPooler
class Pooler(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.dense = nn.Linear(hidden_size, hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states):
        # We "pool" the model by simply taking the hidden state corresponding
        # to the first token.
        first_token_tensor = hidden_states[:, 0]
        pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(pooled_output)
        return pooled_output

# from https://pytorch.org/tutorials/beginner/transformer_tutorial.html
class PositionalEncoding(nn.Module):

    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        # print('[DEBUG] input size:', x.size())
        # print('[DEBUG] positional embedding size:', self.pe.size())
        x = x + self.pe[:x.size(0), :]
        # print('[DEBUG] output x with pe size:', x.size())
        return self.dropout(x)