import torch.nn as nn
from transformers import  PreTrainedModel, BertModel
from transformers.modeling_outputs import SequenceClassifierOutput
from .config_tunbert import TunBertConfig
class classifier(nn.Module):
  def __init__(self,config):
    super().__init__()

    self.layer0 = nn.Linear(in_features=config.hidden_size, out_features=config.hidden_size, bias=True)
    self.layer1 = nn.Linear(in_features=config.hidden_size, out_features=config.type_vocab_size, bias=True)
  def forward(self,tensor):
    out1 = self.layer0(tensor)
    return self.layer1(out1)


class TunBERT(PreTrainedModel):
    config_class = TunBertConfig
    def __init__(self, config):
        super().__init__(config)
        self.BertModel = BertModel(config)
        self.dropout = nn.Dropout(p=0.1, inplace=False)
        self.classifier = classifier(config)

    def forward(self,input_ids=None,token_type_ids=None,attention_mask=None,labels=None) :
      outputs = self.BertModel(input_ids,token_type_ids,attention_mask)
      sequence_output = self.dropout(outputs.last_hidden_state)
      logits = self.classifier(sequence_output)
      loss =None
      if labels is not None :
        loss_func = nn.CrossEntropyLoss()
        loss = loss_func(logits.view(-1,self.config.type_vocab_size),labels.view(-1))
      return SequenceClassifierOutput(loss = loss, logits= logits, hidden_states=outputs.last_hidden_state,attentions=outputs.attentions)
    

TunBertConfig.register_for_auto_class()
TunBERT.register_for_auto_class("AutoModelForSequenceClassification")