Update modeling_tunbert.py

modeling_tunbert.py

@@ -2,15 +2,17 @@ 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)
+      
+    def forward(self,tensor):
+        out1 = self.layer0(tensor)
+        return self.layer1(out1)
 
 
 class TunBERT(PreTrainedModel):
@@ -22,15 +24,20 @@ class TunBERT(PreTrainedModel):
         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 :
+        outputs = self.BertModel(input_ids,token_type_ids,attention_mask)
+        sequence_output = self.dropout(outputs.last_hidden_state)
+        logits = self.classifier(sequence_output)
+        # every sentence is surrounded by [cls] in the beginning and [sep] in the end
+        # the [cls] token is used in bert to identify the class of the sentence
+        # meaning that we need only the first token of each sentence 
+        # and the model representation of the rest of the sentence does not concern us
+        logits = logits[:,0,:] # [bs, seq, class] 
+        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)
-    
+        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")