# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/05_models.roberta.deep_hierarchical_classifiers.ipynb.

# %% ../../../nbs/05_models.roberta.deep_hierarchical_classifiers.ipynb 3
from __future__ import annotations
import torch
from transformers.models.roberta.configuration_roberta import RobertaConfig
from transformers.models.roberta.modeling_roberta import RobertaModel
from transformers.modeling_outputs import SequenceClassifierOutput
from transformers.models.roberta.modeling_roberta import RobertaPreTrainedModel
from sklearn.preprocessing import MultiLabelBinarizer

# %% auto 0
__all__ = ['build_DHC_conditional_mask', 'loss_for_DHC', 'RobertaConcatHeadDHCRoot', 'RobertaSimpleHSCDHCSequenceClassification',
           'RobertaHSCDHCSequenceClassification']

# %% ../../../nbs/05_models.roberta.deep_hierarchical_classifiers.ipynb 5
def build_DHC_conditional_mask(df_labels,
                                  label1,label2):
    """
    This is really similar to `build_standard_condition_mask`, 
    but we don't concatenate l1 1-hot matrix and l2 1-hot matrix
    """
    df_labels = df_labels.drop_duplicates().sort_values([label1,label2])
    _d = df_labels.groupby([label1])[label2].apply(list).to_dict()
    
    mlb = MultiLabelBinarizer()
    results = mlb.fit_transform([v for k,v in sorted(_d.items())])
    return torch.from_numpy(results).float()

# %% ../../../nbs/05_models.roberta.deep_hierarchical_classifiers.ipynb 10
def _check_hierarchy(l1_pred,l2_pred,dhc_mask,size_l2):
    # Check if the predicted class at L2 is a children of the class predicted at L1 for the entire batch
    # if is a children return 0, else 1
    _l1_pred = dhc_mask[l1_pred] # bs,size_l2
    _l2_pred = torch.nn.functional.one_hot(l2_pred, num_classes=size_l2) # bs, size_l2
    return (~torch.mul(_l1_pred,_l2_pred).sum(axis=1).bool()).float()

# %% ../../../nbs/05_models.roberta.deep_hierarchical_classifiers.ipynb 11
def loss_for_DHC(l1_repr_logits, # Head 1's logit output
                 l2_repr_logits, # Head 2's logit output
                 labels_l1, # True label for head 1
                 labels_l2, # True label for head 2
                 dhc_mask, # A one-hot matrix between classes of head 1 and 2
                 lloss_weight=1.0, # Weight for Layer Loss (lloss)
                 dloss_weight=0.8 # Weight for Dependence Loss (dloss)
                ):
    "Reference: https://github.com/Ugenteraan/Deep_Hierarchical_Classification/blob/main/model/hierarchical_loss.py"
    size_l2 = dhc_mask.shape[1]
    
    # calculate lloss
    lloss = 0
    lloss += torch.nn.CrossEntropyLoss()(l1_repr_logits, labels_l1)
    lloss += torch.nn.CrossEntropyLoss()(l2_repr_logits, labels_l2)
    lloss = lloss_weight*lloss


    # calculate dloss
    l1_pred = torch.argmax(l1_repr_logits, dim=1)
    l2_pred = torch.argmax(l2_repr_logits, dim=1)
    D_l = _check_hierarchy(l1_pred,l2_pred,dhc_mask,size_l2)
    l1_not_match_true_label = (~(l1_pred == labels_l1)).float()
    l2_not_match_true_label = (~(l2_pred == labels_l2)).float()
    dloss_l1 = torch.nn.CrossEntropyLoss()(l1_repr_logits*D_l[:,None]*l1_not_match_true_label[:,None], labels_l1)
    dloss_l2 = torch.nn.CrossEntropyLoss()(l2_repr_logits*D_l[:,None]*l2_not_match_true_label[:,None], labels_l2)
    dloss = dloss_weight*(dloss_l1 + dloss_l2)  
    
    
    loss = lloss + dloss
    
    return loss

# %% ../../../nbs/05_models.roberta.deep_hierarchical_classifiers.ipynb 13
class RobertaConcatHeadDHCRoot(torch.nn.Module):
    """
    Concatenated head for Roberta DHC Classification Model. 
    """
    def __init__(self,
                 config, # HuggingFace model configuration
                 classifier_dropout=0.1, # Dropout ratio (for dropout layer right before the last nn.Linear)
                 last_hidden_size=768, # Last hidden size (before the last nn.Linear)
                 layer2concat=4, # number of hidden layer to concatenate (counting from top)
                 **kwargs
                ):
        super().__init__()
        self.classifier_dropout=classifier_dropout
        self.last_hidden_size=last_hidden_size
        self.dropout = torch.nn.Dropout(classifier_dropout)
        self.pre_classifier = torch.nn.Linear(layer2concat*config.hidden_size,
                                              last_hidden_size)    
    def forward(self, inp, **kwargs):
        x = inp
        x = self.dropout(x)
        x = self.pre_classifier(x)
        x = torch.tanh(x)
#         x = torch.relu(x)
        return x

# %% ../../../nbs/05_models.roberta.deep_hierarchical_classifiers.ipynb 15
class RobertaSimpleHSCDHCSequenceClassification(RobertaPreTrainedModel):
    """
    Roberta Simple-DHC Architecture with Hidden-State-Concatenation for Sequence Classification task
    """
    config_class = RobertaConfig

    def __init__(self, 
                 config, # HuggingFace model configuration
                 dhc_mask, # A one-hot matrix between classes of head 1 and 2
                 lloss_weight=1.0, # Weight for Layer Loss (lloss)
                 dloss_weight=0.8, # Weight for Dependence Loss (dloss)
                 layer2concat=4, # number of hidden layer to concatenate (counting from top)
                 device=None # CPU or GPU
                ):
        
        super().__init__(config)
        self.training_device = device if device is not None else torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.size_l1 = dhc_mask.shape[0]
        self.size_l2 = dhc_mask.shape[1]
        self.head_class_sizes=[self.size_l1,self.size_l2] # will be useful for metric calculation later

        self.dhc_mask = dhc_mask.to(self.training_device)
        self.layer2concat=layer2concat
    
        self.lloss_weight = lloss_weight
        self.dloss_weight = dloss_weight
        
        self.body_model = RobertaModel(config, add_pooling_layer=False)
        
        self.linear_L1_logit = torch.nn.Linear(layer2concat*config.hidden_size,self.size_l1)

        self.linear_L2 = torch.nn.Linear(layer2concat*config.hidden_size, abs(self.size_l2-self.size_l1))
        self.linear_L2_logit = torch.nn.Linear(self.size_l2,self.size_l2)
        
    
    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None,
                labels=None, 
                output_attentions= None,
                output_hidden_states= None,
                return_dict= None,
                **kwargs):
        # Use model body to get encoder representations
        # the only ones we need for now are input_ids and attention_mask
        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        outputs = self.body_model(input_ids, 
                                  attention_mask=attention_mask,
                                  token_type_ids=token_type_ids, 
                                  output_attentions=output_attentions,
                                  output_hidden_states=output_hidden_states,
                                  return_dict=return_dict,
                                  **kwargs)
        
        hidden_states = outputs['hidden_states'] # tuples with len 13 (number of layer/block)
        # each with shape: (bs,seq_len,hidden_size_len), e.g. for phobert: (bs,256, 768)
        # Note: hidden_size_len = embedding_size
        
        hidden_concat = torch.cat([hidden_states[i][:,0] for i in range(-1,-self.layer2concat-1,-1)],
                                  -1) # (bs,768*4)
        
        # root_representation
        root_repr = hidden_concat # (bs,768*4) 
        
        l1_repr_logits = self.linear_L1_logit(root_repr) # (bs,size_l1)
        
        l2_repr = self.linear_L2(root_repr) # (bs,size_l2-size_l1)
        l2_repr = torch.cat((l2_repr,l1_repr_logits),dim=1) # (bs, size_l2)
        l2_repr_logits = self.linear_L2_logit(l2_repr) # (bs, size_l2) 
        
        loss = None
        
        if labels is not None:
            # labels shape: (bs,2), first is L1, second is L2
            labels_l1 = labels[:,0].view(-1) #(bs,)
            labels_l2 = labels[:,1].view(-1) #(bs,)

            loss = loss_for_DHC(l1_repr_logits, l2_repr_logits, 
                                labels_l1, labels_l2, 
                                self.dhc_mask,
                                self.lloss_weight,self.dloss_weight)
            
        # Return model output object
        return SequenceClassifierOutput(loss=loss, logits=(l1_repr_logits,l2_repr_logits),
                                     hidden_states=None,
                                     attentions=outputs.attentions)

# %% ../../../nbs/05_models.roberta.deep_hierarchical_classifiers.ipynb 17
class RobertaHSCDHCSequenceClassification(RobertaPreTrainedModel):
    """
    Roberta DHC Architecture with Hidden-State-Concatenation for Sequence Classification task
    """    
    config_class = RobertaConfig

    def __init__(self, 
                 config, # HuggingFace model configuration
                 dhc_mask, # A one-hot matrix between classes of head 1 and 2
                 classifier_dropout=0.1, # Dropout ratio (for dropout layer right before the last nn.Linear)
                 last_hidden_size=768, # Last hidden size (before the last nn.Linear)
                 linear_l1_size=None, # last hidden size for head 1
                 linear_l2_size=None, # last hidden size for head 2
                 lloss_weight=1.0, # Weight for Layer Loss (lloss)
                 dloss_weight=0.8, # Weight for Dependence Loss (dloss)
                 layer2concat=4, # number of hidden layer to concatenate (counting from top)
                 device=None # CPU or GPU
                ):
        
        super().__init__(config)
        self.training_device = device if device is not None else torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.size_l1 = dhc_mask.shape[0]
        self.size_l2 = dhc_mask.shape[1]
        self.head_class_sizes=[self.size_l1,self.size_l2] # will be useful for metric calculation later

        self.dhc_mask = dhc_mask.to(self.training_device)
        self.layer2concat=layer2concat

        
        if linear_l1_size is None: linear_l1_size = (last_hidden_size+self.size_l1)//2 # 389
        if linear_l2_size is None: linear_l2_size = (last_hidden_size+self.size_l2)//2 # 417
            
        self.lloss_weight = lloss_weight
        self.dloss_weight = dloss_weight
        
        self.body_model = RobertaModel(config, add_pooling_layer=False)
        self.root_representation = RobertaConcatHeadDHCRoot(config=config,
                                                            classifier_dropout=classifier_dropout,
                                                            last_hidden_size=last_hidden_size)
        
        self.linear_L1 = torch.nn.Linear(last_hidden_size, linear_l1_size)
        self.linear_L2 = torch.nn.Linear(last_hidden_size, linear_l2_size)
        
        self.linear_L1_logit = torch.nn.Linear(linear_l1_size,self.size_l1)
        self.linear_L2_logit = torch.nn.Linear(linear_l1_size + linear_l2_size,self.size_l2)
        

    
    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None,
                labels=None, 
                output_attentions= None,
                output_hidden_states= None,
                return_dict= None,
                **kwargs):
        # Use model body to get encoder representations
        # the only ones we need for now are input_ids and attention_mask
        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        outputs = self.body_model(input_ids, 
                                  attention_mask=attention_mask,
                                  token_type_ids=token_type_ids, 
                                  output_attentions=output_attentions,
                                  output_hidden_states=output_hidden_states,
                                  return_dict=return_dict,
                                  **kwargs)
        
        hidden_states = outputs['hidden_states'] # tuples with len 13 (number of layer/block)
        # each with shape: (bs,seq_len,hidden_size_len), e.g. for phobert: (bs,256, 768)
        # Note: hidden_size_len = embedding_size
        
        hidden_concat = torch.cat([hidden_states[i][:,0] for i in range(-1,-self.layer2concat-1,-1)],
                                  -1) # (bs,768*4)
        
        # root_representation
        root_repr = self.root_representation(hidden_concat) # (bs,768) 
        
        l1_repr = self.linear_L1(root_repr) # (bs,linear_L1_size)
        l1_repr_logits = self.linear_L1_logit(l1_repr) # (bs,size_l1)
        l2_repr = self.linear_L2(root_repr) # (bs,linear_L2_size)
        l2_repr = torch.cat((l2_repr,l1_repr),dim=1) # (bs, linear_l1_size + linear_l2_size)
        l2_repr_logits = self.linear_L2_logit(l2_repr) # (bs, size_l2) 

        
        loss = None
        
        if labels is not None:
            # labels shape: (bs,2), first is L3, second is L4
            labels_l1 = labels[:,0].view(-1) #(bs,)
            labels_l2 = labels[:,1].view(-1) #(bs,)

            loss = loss_for_DHC(l1_repr_logits, l2_repr_logits, 
                                labels_l1, labels_l2, 
                                self.dhc_mask,
                                self.lloss_weight,self.dloss_weight)
            
        # Return model output object
        return SequenceClassifierOutput(loss=loss, logits=(l1_repr_logits,l2_repr_logits),
                                     hidden_states=None,
                                     attentions=outputs.attentions)
