From 4d9e45f3ef624cab41f605d7439862ce23ca806a Mon Sep 17 00:00:00 2001 From: David Reguera <33068707+nablabits@users.noreply.github.com> Date: Fri, 25 Aug 2023 16:12:54 +0200 Subject: [PATCH] Add type hints for several pytorch models (batch-3) (#25705) * Add missing type hints for ErnieM family * Add missing type hints for EsmForProteinFolding model * Add missing type hints for Graphormer model * Add type hints for InstructBlipQFormer model * Add missing type hints for LayoutLMForMaskedLM model * Add missing type hints for LukeForEntitySpanClassification model --- .../models/ernie_m/modeling_ernie_m.py | 12 ++++----- .../models/esm/modeling_esmfold.py | 4 +-- .../models/graphormer/modeling_graphormer.py | 4 +-- .../instructblip/modeling_instructblip.py | 26 +++++++++---------- .../models/layoutlm/modeling_layoutlm.py | 4 +-- src/transformers/models/luke/modeling_luke.py | 2 +- 6 files changed, 26 insertions(+), 26 deletions(-) diff --git a/src/transformers/models/ernie_m/modeling_ernie_m.py b/src/transformers/models/ernie_m/modeling_ernie_m.py index 82e4023949130b..9c53ddd73c8540 100755 --- a/src/transformers/models/ernie_m/modeling_ernie_m.py +++ b/src/transformers/models/ernie_m/modeling_ernie_m.py @@ -538,7 +538,7 @@ def forward( output_hidden_states: Optional[bool] = None, output_attentions: Optional[bool] = None, return_dict: Optional[bool] = None, - ): + ) -> Union[Tuple[torch.FloatTensor], BaseModelOutputWithPoolingAndCrossAttentions]: if input_ids is not None and inputs_embeds is not None: raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time.") @@ -647,7 +647,7 @@ def forward( output_attentions: Optional[bool] = None, return_dict: Optional[bool] = True, labels: Optional[torch.Tensor] = None, - ): + ) -> Union[Tuple[torch.FloatTensor], SequenceClassifierOutput]: r""" labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., @@ -744,7 +744,7 @@ def forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = True, - ): + ) -> Union[Tuple[torch.FloatTensor], MultipleChoiceModelOutput]: r""" labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., @@ -837,7 +837,7 @@ def forward( output_attentions: Optional[bool] = None, return_dict: Optional[bool] = True, labels: Optional[torch.Tensor] = None, - ): + ) -> Union[Tuple[torch.FloatTensor], TokenClassifierOutput]: 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]`. @@ -914,7 +914,7 @@ def forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = True, - ): + ) -> Union[Tuple[torch.FloatTensor], 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. @@ -1003,7 +1003,7 @@ def forward( output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = True, - ): + ) -> Union[Tuple[torch.FloatTensor], QuestionAnsweringModelOutput]: r""" start_positions (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for position (index) for computing the start_positions loss. Position outside of the sequence are diff --git a/src/transformers/models/esm/modeling_esmfold.py b/src/transformers/models/esm/modeling_esmfold.py index 6e6c9fcb4aea08..ee99a76e4961cb 100644 --- a/src/transformers/models/esm/modeling_esmfold.py +++ b/src/transformers/models/esm/modeling_esmfold.py @@ -2086,11 +2086,11 @@ def _af2_to_esm_from_vocab_list(vocab_list: List[str]) -> torch.Tensor: def forward( self, input_ids: torch.Tensor, - attention_mask: torch.Tensor = None, + attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, masking_pattern: Optional[torch.Tensor] = None, num_recycles: Optional[int] = None, - ): + ) -> EsmForProteinFoldingOutput: r""" Returns: diff --git a/src/transformers/models/graphormer/modeling_graphormer.py b/src/transformers/models/graphormer/modeling_graphormer.py index 82ffd4b1637d54..7257415656a20a 100755 --- a/src/transformers/models/graphormer/modeling_graphormer.py +++ b/src/transformers/models/graphormer/modeling_graphormer.py @@ -816,8 +816,8 @@ def forward( out_degree: torch.LongTensor, spatial_pos: torch.LongTensor, attn_edge_type: torch.LongTensor, - perturb=None, - masked_tokens=None, + perturb: Optional[torch.FloatTensor] = None, + masked_tokens: None = None, return_dict: Optional[bool] = None, **unused, ) -> Union[Tuple[torch.LongTensor], BaseModelOutputWithNoAttention]: diff --git a/src/transformers/models/instructblip/modeling_instructblip.py b/src/transformers/models/instructblip/modeling_instructblip.py index ea736c28400cb0..082900a6652f80 100644 --- a/src/transformers/models/instructblip/modeling_instructblip.py +++ b/src/transformers/models/instructblip/modeling_instructblip.py @@ -1124,19 +1124,19 @@ def get_extended_attention_mask( def forward( self, - input_ids, - attention_mask=None, - position_ids=None, - query_embeds=None, - head_mask=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - past_key_values=None, - use_cache=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - ): + input_ids: torch.LongTensor, + attention_mask: Optional[torch.FloatTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + query_embeds: Optional[torch.Tensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.FloatTensor], BaseModelOutputWithPoolingAndCrossAttentions]: r""" encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if diff --git a/src/transformers/models/layoutlm/modeling_layoutlm.py b/src/transformers/models/layoutlm/modeling_layoutlm.py index b1f453f1dfd081..884a2799728b47 100644 --- a/src/transformers/models/layoutlm/modeling_layoutlm.py +++ b/src/transformers/models/layoutlm/modeling_layoutlm.py @@ -891,8 +891,8 @@ def forward( head_mask: Optional[torch.FloatTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, - encoder_hidden_states=None, - encoder_attention_mask=None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, diff --git a/src/transformers/models/luke/modeling_luke.py b/src/transformers/models/luke/modeling_luke.py index 52b947e8097793..6913ede09d1c7b 100644 --- a/src/transformers/models/luke/modeling_luke.py +++ b/src/transformers/models/luke/modeling_luke.py @@ -1664,7 +1664,7 @@ def __init__(self, config): def forward( self, input_ids: Optional[torch.LongTensor] = None, - attention_mask=None, + attention_mask: Optional[torch.FloatTensor] = None, token_type_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, entity_ids: Optional[torch.LongTensor] = None,