Skip to content

Commit df2babc

Browse files
robsmith155elusenji
robsmith155
authored and
elusenji
committed
Add type hints for BigBirdPegasus and Data2VecText PyTorch models (huggingface#17123)
* Add type hints for remaining BigBirdPegasus models Here I added type hints to the BigBirdPegasusForCausalLM class. * Add missing type hints for Data2VecText models Added type hints to the Data2VecTextForCausalLM, Data2VecTextForMaskedLM, Data2VecTextForMultipleChoice, Data2VecTextForQuestionAnswering, Data2VecTextForSequenceClassification, and Data2VecTextForTokenClassification classes.
1 parent 64f5141 commit df2babc

File tree

2 files changed

+87
-87
lines changed

2 files changed

+87
-87
lines changed

‎src/transformers/models/bigbird_pegasus/modeling_bigbird_pegasus.py

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2931,20 +2931,20 @@ def get_decoder(self):
29312931
@replace_return_docstrings(output_type=CausalLMOutputWithCrossAttentions, config_class=_CONFIG_FOR_DOC)
29322932
def forward(
29332933
self,
2934-
input_ids=None,
2935-
attention_mask=None,
2936-
encoder_hidden_states=None,
2937-
encoder_attention_mask=None,
2938-
head_mask=None,
2939-
cross_attn_head_mask=None,
2940-
past_key_values=None,
2941-
inputs_embeds=None,
2942-
labels=None,
2943-
use_cache=None,
2944-
output_attentions=None,
2945-
output_hidden_states=None,
2946-
return_dict=None,
2947-
):
2934+
input_ids: torch.LongTensor = None,
2935+
attention_mask: Optional[torch.Tensor] = None,
2936+
encoder_hidden_states: Optional[torch.FloatTensor] = None,
2937+
encoder_attention_mask: Optional[torch.FloatTensor] = None,
2938+
head_mask: Optional[torch.Tensor] = None,
2939+
cross_attn_head_mask: Optional[torch.Tensor] = None,
2940+
past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
2941+
inputs_embeds: Optional[torch.FloatTensor] = None,
2942+
labels: Optional[torch.LongTensor] = None,
2943+
use_cache: Optional[bool] = None,
2944+
output_attentions: Optional[bool] = None,
2945+
output_hidden_states: Optional[bool] = None,
2946+
return_dict: Optional[bool] = None,
2947+
) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
29482948
r"""
29492949
Args:
29502950
input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):

‎src/transformers/models/data2vec/modeling_data2vec_text.py

Lines changed: 73 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -908,21 +908,21 @@ def set_output_embeddings(self, new_embeddings):
908908
@replace_return_docstrings(output_type=CausalLMOutputWithCrossAttentions, config_class=_CONFIG_FOR_DOC)
909909
def forward(
910910
self,
911-
input_ids=None,
912-
attention_mask=None,
913-
token_type_ids=None,
914-
position_ids=None,
915-
head_mask=None,
916-
inputs_embeds=None,
917-
encoder_hidden_states=None,
918-
encoder_attention_mask=None,
919-
labels=None,
920-
past_key_values=None,
921-
use_cache=None,
922-
output_attentions=None,
923-
output_hidden_states=None,
924-
return_dict=None,
925-
):
911+
input_ids: Optional[torch.LongTensor] = None,
912+
attention_mask: Optional[torch.FloatTensor] = None,
913+
token_type_ids: Optional[torch.LongTensor] = None,
914+
position_ids: Optional[torch.LongTensor] = None,
915+
head_mask: Optional[torch.FloatTensor] = None,
916+
inputs_embeds: Optional[torch.FloatTensor] = None,
917+
encoder_hidden_states: Optional[torch.FloatTensor] = None,
918+
encoder_attention_mask: Optional[torch.FloatTensor] = None,
919+
labels: Optional[torch.LongTensor] = None,
920+
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
921+
use_cache: Optional[bool] = None,
922+
output_attentions: Optional[bool] = None,
923+
output_hidden_states: Optional[bool] = None,
924+
return_dict: Optional[bool] = None,
925+
) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
926926
r"""
927927
encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
928928
Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
@@ -1069,19 +1069,19 @@ def set_output_embeddings(self, new_embeddings):
10691069
)
10701070
def forward(
10711071
self,
1072-
input_ids=None,
1073-
attention_mask=None,
1074-
token_type_ids=None,
1075-
position_ids=None,
1076-
head_mask=None,
1077-
inputs_embeds=None,
1078-
encoder_hidden_states=None,
1079-
encoder_attention_mask=None,
1080-
labels=None,
1081-
output_attentions=None,
1082-
output_hidden_states=None,
1083-
return_dict=None,
1084-
):
1072+
input_ids: Optional[torch.LongTensor] = None,
1073+
attention_mask: Optional[torch.FloatTensor] = None,
1074+
token_type_ids: Optional[torch.LongTensor] = None,
1075+
position_ids: Optional[torch.LongTensor] = None,
1076+
head_mask: Optional[torch.FloatTensor] = None,
1077+
inputs_embeds: Optional[torch.FloatTensor] = None,
1078+
encoder_hidden_states: Optional[torch.FloatTensor] = None,
1079+
encoder_attention_mask: Optional[torch.FloatTensor] = None,
1080+
labels: Optional[torch.LongTensor] = None,
1081+
output_attentions: Optional[bool] = None,
1082+
output_hidden_states: Optional[bool] = None,
1083+
return_dict: Optional[bool] = None,
1084+
) -> Union[Tuple, MaskedLMOutput]:
10851085
r"""
10861086
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
10871087
Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
@@ -1183,17 +1183,17 @@ def __init__(self, config):
11831183
)
11841184
def forward(
11851185
self,
1186-
input_ids=None,
1187-
attention_mask=None,
1188-
token_type_ids=None,
1189-
position_ids=None,
1190-
head_mask=None,
1191-
inputs_embeds=None,
1192-
labels=None,
1193-
output_attentions=None,
1194-
output_hidden_states=None,
1195-
return_dict=None,
1196-
):
1186+
input_ids: Optional[torch.LongTensor] = None,
1187+
attention_mask: Optional[torch.FloatTensor] = None,
1188+
token_type_ids: Optional[torch.LongTensor] = None,
1189+
position_ids: Optional[torch.LongTensor] = None,
1190+
head_mask: Optional[torch.FloatTensor] = None,
1191+
inputs_embeds: Optional[torch.FloatTensor] = None,
1192+
labels: Optional[torch.LongTensor] = None,
1193+
output_attentions: Optional[bool] = None,
1194+
output_hidden_states: Optional[bool] = None,
1195+
return_dict: Optional[bool] = None,
1196+
) -> Union[Tuple, SequenceClassifierOutput]:
11971197
r"""
11981198
labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
11991199
Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
@@ -1282,17 +1282,17 @@ def __init__(self, config):
12821282
)
12831283
def forward(
12841284
self,
1285-
input_ids=None,
1286-
token_type_ids=None,
1287-
attention_mask=None,
1288-
labels=None,
1289-
position_ids=None,
1290-
head_mask=None,
1291-
inputs_embeds=None,
1292-
output_attentions=None,
1293-
output_hidden_states=None,
1294-
return_dict=None,
1295-
):
1285+
input_ids: Optional[torch.LongTensor] = None,
1286+
token_type_ids: Optional[torch.LongTensor] = None,
1287+
attention_mask: Optional[torch.FloatTensor] = None,
1288+
labels: Optional[torch.LongTensor] = None,
1289+
position_ids: Optional[torch.LongTensor] = None,
1290+
head_mask: Optional[torch.FloatTensor] = None,
1291+
inputs_embeds: Optional[torch.FloatTensor] = None,
1292+
output_attentions: Optional[bool] = None,
1293+
output_hidden_states: Optional[bool] = None,
1294+
return_dict: Optional[bool] = None,
1295+
) -> Union[Tuple, MultipleChoiceModelOutput]:
12961296
r"""
12971297
labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
12981298
Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
@@ -1380,17 +1380,17 @@ def __init__(self, config):
13801380
)
13811381
def forward(
13821382
self,
1383-
input_ids=None,
1384-
attention_mask=None,
1385-
token_type_ids=None,
1386-
position_ids=None,
1387-
head_mask=None,
1388-
inputs_embeds=None,
1389-
labels=None,
1390-
output_attentions=None,
1391-
output_hidden_states=None,
1392-
return_dict=None,
1393-
):
1383+
input_ids: Optional[torch.LongTensor] = None,
1384+
attention_mask: Optional[torch.FloatTensor] = None,
1385+
token_type_ids: Optional[torch.LongTensor] = None,
1386+
position_ids: Optional[torch.LongTensor] = None,
1387+
head_mask: Optional[torch.FloatTensor] = None,
1388+
inputs_embeds: Optional[torch.FloatTensor] = None,
1389+
labels: Optional[torch.LongTensor] = None,
1390+
output_attentions: Optional[bool] = None,
1391+
output_hidden_states: Optional[bool] = None,
1392+
return_dict: Optional[bool] = None,
1393+
) -> Union[Tuple, TokenClassifierOutput]:
13941394
r"""
13951395
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
13961396
Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
@@ -1484,18 +1484,18 @@ def __init__(self, config):
14841484
)
14851485
def forward(
14861486
self,
1487-
input_ids=None,
1488-
attention_mask=None,
1489-
token_type_ids=None,
1490-
position_ids=None,
1491-
head_mask=None,
1492-
inputs_embeds=None,
1493-
start_positions=None,
1494-
end_positions=None,
1495-
output_attentions=None,
1496-
output_hidden_states=None,
1497-
return_dict=None,
1498-
):
1487+
input_ids: Optional[torch.LongTensor] = None,
1488+
attention_mask: Optional[torch.FloatTensor] = None,
1489+
token_type_ids: Optional[torch.LongTensor] = None,
1490+
position_ids: Optional[torch.LongTensor] = None,
1491+
head_mask: Optional[torch.FloatTensor] = None,
1492+
inputs_embeds: Optional[torch.FloatTensor] = None,
1493+
start_positions: Optional[torch.LongTensor] = None,
1494+
end_positions: Optional[torch.LongTensor] = None,
1495+
output_attentions: Optional[bool] = None,
1496+
output_hidden_states: Optional[bool] = None,
1497+
return_dict: Optional[bool] = None,
1498+
) -> Union[Tuple, QuestionAnsweringModelOutput]:
14991499
r"""
15001500
start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
15011501
Labels for position (index) of the start of the labelled span for computing the token classification loss.

0 commit comments

Comments
 (0)