Update custom_heads_flash_t5.py

custom_heads_flash_t5.py

@@ -12,7 +12,7 @@ from transformers.modeling_outputs import (
     SequenceClassifierOutput
 )
 
-from .modeling_flash_t5 import FlashT5PreTrainedModel, FlashT5Stack, FlashT5Model, FlashT5EncoderModel
+from .modeling_flash_t5 import FlashT5PreTrainedModel, FlashT5Stack, FlashT5Model
 from .configuration_flash_t5 import FlashT5Config
 
 
@@ -225,15 +225,20 @@ class FlashT5ForQuestionAnswering(FlashT5PreTrainedModel):
 
     def __init__(self, config: FlashT5Config):
         super().__init__(config)
-        self.transformer = FlashT5EncoderModel(config)
+        self.shared = nn.Embedding(config.vocab_size, config.d_model)
 
-        self.num_labels = config.num_labels
+        encoder_config = copy.deepcopy(config)
+        encoder_config.is_decoder = False
+        encoder_config.is_encoder_decoder = False
+        self.encoder = FlashT5Stack(encoder_config, self.shared)
         self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
 
         # Initialize weights and apply final processing
         self.post_init()
 
-        # Model parallel
+        self.qa_outputs.weight.data.normal_(mean=0.0, std=config.initializer_factor * 1.0)
+        self.qa_outputs.bias.data.zero_()
+
         self.model_parallel = False
 
     def forward(
@@ -242,37 +247,37 @@ class FlashT5ForQuestionAnswering(FlashT5PreTrainedModel):
         attention_mask: Optional[torch.FloatTensor] = None,
         head_mask: Optional[torch.FloatTensor] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
-        start_positions: Optional[torch.Tensor] = None,
-        end_positions: Optional[torch.Tensor] = None,
+        start_positions: Optional[torch.LongTensor] = None,
+        end_positions: Optional[torch.LongTensor] = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-    ) -> Union[Tuple[torch.FloatTensor], QuestionAnsweringModelOutput]:
+    ) -> Union[Tuple, QuestionAnsweringModelOutput]:
         r"""
-        start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
-            Labels for position (index) of the start of the labelled span for computing the token classification loss.
-            Positions are clamped to the length of the sequence (*sequence_length*). Position outside of the sequence
-            are not taken into account for computing the loss.
-        end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
-            Labels for position (index) of the end of the labelled span for computing the token classification loss.
-            Positions are clamped to the length of the sequence (*sequence_length*). Position outside of the sequence
-            are not taken into account for computing the loss.
-
         Returns:
-        """
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, MTxEncoderForQuestionAnswering
+
+        >>> tokenizer = AutoTokenizer.from_pretrained("MTx-small")
+        >>> model = MTxEncoderForQuestionAnswering.from_pretrained("MTx-small")
+        >>> input_ids = tokenizer(
+        ...     "Studies have been shown that owning a dog is good for you", return_tensors="pt"
+        ... ).input_ids  # Batch size 1
+        >>> outputs = model(input_ids=input_ids)
+        >>> start_logits = outputs.start_logits
+        >>> end_logits = outputs.end_logits
+        ```"""
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        encoder_outputs = self.transformer(
-            input_ids=input_ids,
+        outputs = self.encoder(
+            input_ids,
             attention_mask=attention_mask,
             inputs_embeds=inputs_embeds,
-            head_mask=head_mask,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
         )
-
-        sequence_output = encoder_outputs[0]
+        sequence_output = outputs[0]
 
         logits = self.qa_outputs(sequence_output)
         start_logits, end_logits = logits.split(1, dim=-1)
@@ -297,13 +302,13 @@ class FlashT5ForQuestionAnswering(FlashT5PreTrainedModel):
             total_loss = (start_loss + end_loss) / 2
 
         if not return_dict:
-            output = (start_logits, end_logits) + encoder_outputs[1:]
+            output = (start_logits, end_logits) + outputs[1:]
             return ((total_loss,) + output) if total_loss is not None else output
 
         return QuestionAnsweringModelOutput(
             loss=total_loss,
             start_logits=start_logits,
             end_logits=end_logits,
-            hidden_states=encoder_outputs.hidden_states,
-            attentions=encoder_outputs.attentions,
-        )
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )