Update `NewTokenClassifierOutput `

modeling.py

@@ -16,6 +16,7 @@
 """PyTorch NEW model."""
 
 import math
+from dataclasses import dataclass
 from typing import List, Optional, Tuple, Union
 
 import torch
@@ -30,7 +31,7 @@ from transformers.modeling_outputs import (
     MultipleChoiceModelOutput,
     QuestionAnsweringModelOutput,
     SequenceClassifierOutput,
-    TokenClassifierOutput,
+    ModelOutput,
 )
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import logging
@@ -1249,6 +1250,15 @@ class NewForMultipleChoice(NewPreTrainedModel):
         )
 
 
+@dataclass
+class NewTokenClassifierOutput(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    last_hidden_state: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
 class NewForTokenClassification(NewPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
@@ -1277,7 +1287,7 @@ class NewForTokenClassification(NewPreTrainedModel):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         unpad_inputs: Optional[bool] = None,
-    ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]:
+    ) -> Union[Tuple[torch.Tensor], NewTokenClassifierOutput]:
         r"""
         labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
             Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
@@ -1311,9 +1321,10 @@ class NewForTokenClassification(NewPreTrainedModel):
             output = (logits,) + outputs[2:]
             return ((loss,) + output) if loss is not None else output
 
-        return TokenClassifierOutput(
+        return NewTokenClassifierOutput(
             loss=loss,
             logits=logits,
+            last_hidden_state=sequence_output,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
         )