Skip to content

GitLab

Unverified Commit 6deac5c8 authored by Thomas's avatar Thomas Committed by GitHub
Browse files

Adding type hints for TFXLnet (#19344) 

* Added type hints for TF: XLNet

* Added type hints for TF: XLNet

* Added type hints for TF: XLNet

* Added type hints for TF: XLNet

* Added type hints for TF: XLNet

* Added type hints for TF: XLNet

* Add type hints for XLnet (TF)
* Added type hints for XLnet (TF)

* Update src/transformers/models/xlnet/modeling_tf_xlnet.py
parent 7036c956
Hide whitespace changes
Inline Side-by-side
Showing with 43 additions and 43 deletions
src/transformers/models/xlnet/modeling_tf_xlnet.py
View file @ 6deac5c8
...@@ -196,11 +196,11 @@ class TFXLNetRelativeAttention(tf.keras.layers.Layer): ...@@ -196,11 +196,11 @@ class TFXLNetRelativeAttention(tf.keras.layers.Layer):
attn_mask_g, attn_mask_g,
r, r,
seg_mat, seg_mat,
mems, mems: Optional[Union[np.ndarray, tf.Tensor]] = None,
target_mapping, target_mapping: Optional[Union[np.ndarray, tf.Tensor]] = None,
head_mask, head_mask: Optional[Union[np.ndarray, tf.Tensor]] = None,
output_attentions, output_attentions: Optional[bool] = False,
training=False, training: bool = False,
): ):
if g is not None: if g is not None:
# Two-stream attention with relative positional encoding. # Two-stream attention with relative positional encoding.
...@@ -370,11 +370,11 @@ class TFXLNetLayer(tf.keras.layers.Layer): ...@@ -370,11 +370,11 @@ class TFXLNetLayer(tf.keras.layers.Layer):
attn_mask, attn_mask,
pos_emb, pos_emb,
seg_mat, seg_mat,
mems, mems: Optional[Union[np.ndarray, tf.Tensor]] = None,
target_mapping, target_mapping: Optional[Union[np.ndarray, tf.Tensor]] = None,
head_mask, head_mask: Optional[Union[np.ndarray, tf.Tensor]] = None,
output_attentions, output_attentions: Optional[bool] = False,
training=False, training: bool = False,
): ):
outputs = self.rel_attn( outputs = self.rel_attn(
output_h, output_h,
...@@ -583,20 +583,20 @@ class TFXLNetMainLayer(tf.keras.layers.Layer): ...@@ -583,20 +583,20 @@ class TFXLNetMainLayer(tf.keras.layers.Layer):
@unpack_inputs @unpack_inputs
def call( def call(
self, self,
input_ids=None, input_ids: Optional[TFModelInputType] = None,
attention_mask=None, attention_mask: Optional[Union[np.ndarray, tf.Tensor]] = None,
mems=None, mems: Optional[Union[np.ndarray, tf.Tensor]] = None,
perm_mask=None, perm_mask: Optional[Union[np.ndarray, tf.Tensor]] = None,
target_mapping=None, target_mapping: Optional[Union[np.ndarray, tf.Tensor]] = None,
token_type_ids=None, token_type_ids: Optional[Union[np.ndarray, tf.Tensor]] = None,
input_mask=None, input_mask: Optional[Union[np.ndarray, tf.Tensor]] = None,
head_mask=None, head_mask: Optional[Union[np.ndarray, tf.Tensor]] = None,
inputs_embeds=None, inputs_embeds: Optional[Union[np.ndarray, tf.Tensor]] = None,
use_mems=None, use_mems: Optional[bool] = None,
output_attentions=None, output_attentions: Optional[bool] = None,
output_hidden_states=None, output_hidden_states: Optional[bool] = None,
return_dict=None, return_dict: Optional[bool] = None,
training=False, training: bool = False,
): ):
if training and use_mems is None: if training and use_mems is None:
...@@ -1152,20 +1152,20 @@ class TFXLNetModel(TFXLNetPreTrainedModel): ...@@ -1152,20 +1152,20 @@ class TFXLNetModel(TFXLNetPreTrainedModel):
) )
def call( def call(
self, self,
input_ids=None, input_ids: Optional[TFModelInputType] = None,
attention_mask=None, attention_mask: Optional[Union[np.ndarray, tf.Tensor]] = None,
mems=None, mems: Optional[Union[np.ndarray, tf.Tensor]] = None,
perm_mask=None, perm_mask: Optional[Union[np.ndarray, tf.Tensor]] = None,
target_mapping=None, target_mapping: Optional[Union[np.ndarray, tf.Tensor]] = None,
token_type_ids=None, token_type_ids: Optional[Union[np.ndarray, tf.Tensor]] = None,
input_mask=None, input_mask: Optional[Union[np.ndarray, tf.Tensor]] = None,
head_mask=None, head_mask: Optional[Union[np.ndarray, tf.Tensor]] = None,
inputs_embeds=None, inputs_embeds: Optional[Union[np.ndarray, tf.Tensor]] = None,
use_mems=None, use_mems: Optional[bool] = None,
output_attentions=None, output_attentions: Optional[bool] = None,
output_hidden_states=None, output_hidden_states: Optional[bool] = None,
return_dict=None, return_dict: Optional[bool] = None,
training=False, training: bool = False,
): ):
outputs = self.transformer( outputs = self.transformer(
input_ids=input_ids, input_ids=input_ids,
...@@ -1275,7 +1275,7 @@ class TFXLNetLMHeadModel(TFXLNetPreTrainedModel, TFCausalLanguageModelingLoss): ...@@ -1275,7 +1275,7 @@ class TFXLNetLMHeadModel(TFXLNetPreTrainedModel, TFCausalLanguageModelingLoss):
output_hidden_states: Optional[bool] = None, output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None, return_dict: Optional[bool] = None,
labels: Optional[Union[np.ndarray, tf.Tensor]] = None, labels: Optional[Union[np.ndarray, tf.Tensor]] = None,
training: Optional[bool] = False, training: bool = False,
) -> Union[TFXLNetLMHeadModelOutput, Tuple[tf.Tensor]]: ) -> Union[TFXLNetLMHeadModelOutput, Tuple[tf.Tensor]]:
r""" r"""
labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*): labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
...@@ -1406,7 +1406,7 @@ class TFXLNetForSequenceClassification(TFXLNetPreTrainedModel, TFSequenceClassif ...@@ -1406,7 +1406,7 @@ class TFXLNetForSequenceClassification(TFXLNetPreTrainedModel, TFSequenceClassif
output_hidden_states: Optional[bool] = None, output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None, return_dict: Optional[bool] = None,
labels: Optional[Union[np.ndarray, tf.Tensor]] = None, labels: Optional[Union[np.ndarray, tf.Tensor]] = None,
training: Optional[bool] = False, training: bool = False,
) -> Union[TFXLNetForSequenceClassificationOutput, Tuple[tf.Tensor]]: ) -> Union[TFXLNetForSequenceClassificationOutput, Tuple[tf.Tensor]]:
r""" r"""
labels (`tf.Tensor` of shape `(batch_size,)`, *optional*): labels (`tf.Tensor` of shape `(batch_size,)`, *optional*):
...@@ -1512,7 +1512,7 @@ class TFXLNetForMultipleChoice(TFXLNetPreTrainedModel, TFMultipleChoiceLoss): ...@@ -1512,7 +1512,7 @@ class TFXLNetForMultipleChoice(TFXLNetPreTrainedModel, TFMultipleChoiceLoss):
output_hidden_states: Optional[bool] = None, output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None, return_dict: Optional[bool] = None,
labels: Optional[Union[np.ndarray, tf.Tensor]] = None, labels: Optional[Union[np.ndarray, tf.Tensor]] = None,
training: Optional[bool] = False, training: bool = False,
) -> Union[TFXLNetForMultipleChoiceOutput, Tuple[tf.Tensor]]: ) -> Union[TFXLNetForMultipleChoiceOutput, Tuple[tf.Tensor]]:
r""" r"""
labels (`tf.Tensor` of shape `(batch_size,)`, *optional*): labels (`tf.Tensor` of shape `(batch_size,)`, *optional*):
...@@ -1633,7 +1633,7 @@ class TFXLNetForTokenClassification(TFXLNetPreTrainedModel, TFTokenClassificatio ...@@ -1633,7 +1633,7 @@ class TFXLNetForTokenClassification(TFXLNetPreTrainedModel, TFTokenClassificatio
output_hidden_states: Optional[bool] = None, output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None, return_dict: Optional[bool] = None,
labels: Optional[Union[np.ndarray, tf.Tensor]] = None, labels: Optional[Union[np.ndarray, tf.Tensor]] = None,
training: Optional[bool] = False, training: bool = False,
) -> Union[TFXLNetForTokenClassificationOutput, Tuple[tf.Tensor]]: ) -> Union[TFXLNetForTokenClassificationOutput, Tuple[tf.Tensor]]:
r""" r"""
labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*): labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
...@@ -1720,7 +1720,7 @@ class TFXLNetForQuestionAnsweringSimple(TFXLNetPreTrainedModel, TFQuestionAnswer ...@@ -1720,7 +1720,7 @@ class TFXLNetForQuestionAnsweringSimple(TFXLNetPreTrainedModel, TFQuestionAnswer
return_dict: Optional[bool] = None, return_dict: Optional[bool] = None,
start_positions: Optional[Union[np.ndarray, tf.Tensor]] = None, start_positions: Optional[Union[np.ndarray, tf.Tensor]] = None,
end_positions: Optional[Union[np.ndarray, tf.Tensor]] = None, end_positions: Optional[Union[np.ndarray, tf.Tensor]] = None,
training: Optional[bool] = False, training: bool = False,
) -> Union[TFXLNetForQuestionAnsweringSimpleOutput, Tuple[tf.Tensor]]: ) -> Union[TFXLNetForQuestionAnsweringSimpleOutput, Tuple[tf.Tensor]]:
r""" r"""
start_positions (`tf.Tensor` of shape `(batch_size,)`, *optional*): start_positions (`tf.Tensor` of shape `(batch_size,)`, *optional*):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment