Unverified Commit 7af80f66 authored by Sylvain Gugger's avatar Sylvain Gugger Committed by GitHub
Browse files

Convert docstrings of modeling files (#14850)

* Convert file_utils docstrings to Markdown

* Test on BERT

* Return block indent

* Temporarily disable doc styler

* Remove from quality checks as well

* Remove doc styler mess

* Remove check from circleCI

* Fix typo

* Convert file_utils docstrings to Markdown

* Test on BERT

* Return block indent

* Temporarily disable doc styler

* Remove from quality checks as well

* Remove doc styler mess

* Remove check from circleCI

* Fix typo

* Let's go on all other model files

* Add templates too

* Styling and quality
parent 2a337346
...@@ -848,7 +848,7 @@ jobs: ...@@ -848,7 +848,7 @@ jobs:
- run: isort --check-only examples tests src utils - run: isort --check-only examples tests src utils
- run: python utils/custom_init_isort.py --check_only - run: python utils/custom_init_isort.py --check_only
- run: flake8 examples tests src utils - run: flake8 examples tests src utils
- run: python utils/style_doc.py src/transformers docs/source --max_len 119 --check_only # - run: python utils/style_doc.py src/transformers docs/source --max_len 119 --check_only
check_repository_consistency: check_repository_consistency:
working_directory: ~/transformers working_directory: ~/transformers
......
...@@ -48,13 +48,13 @@ quality: ...@@ -48,13 +48,13 @@ quality:
isort --check-only $(check_dirs) isort --check-only $(check_dirs)
python utils/custom_init_isort.py --check_only python utils/custom_init_isort.py --check_only
flake8 $(check_dirs) flake8 $(check_dirs)
python utils/style_doc.py src/transformers docs/source --max_len 119 --check_only # python utils/style_doc.py src/transformers docs/source --max_len 119 --check_only
# Format source code automatically and check is there are any problems left that need manual fixing # Format source code automatically and check is there are any problems left that need manual fixing
extra_style_checks: extra_style_checks:
python utils/custom_init_isort.py python utils/custom_init_isort.py
python utils/style_doc.py src/transformers docs/source --max_len 119 # python utils/style_doc.py src/transformers docs/source --max_len 119
# this target runs checks on all files and potentially modifies some of them # this target runs checks on all files and potentially modifies some of them
......
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
...@@ -517,25 +517,24 @@ class AlbertPreTrainedModel(PreTrainedModel): ...@@ -517,25 +517,24 @@ class AlbertPreTrainedModel(PreTrainedModel):
@dataclass @dataclass
class AlbertForPreTrainingOutput(ModelOutput): class AlbertForPreTrainingOutput(ModelOutput):
""" """
Output type of :class:`~transformers.AlbertForPreTraining`. Output type of [`AlbertForPreTraining`].
Args: Args:
loss (`optional`, returned when ``labels`` is provided, ``torch.FloatTensor`` of shape :obj:`(1,)`): loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
Total loss as the sum of the masked language modeling loss and the next sequence prediction Total loss as the sum of the masked language modeling loss and the next sequence prediction
(classification) loss. (classification) loss.
prediction_logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`): prediction_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
sop_logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, 2)`): sop_logits (`torch.FloatTensor` of shape `(batch_size, 2)`):
Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation
before SoftMax). before SoftMax).
hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``): hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
of shape :obj:`(batch_size, sequence_length, hidden_size)`. of shape `(batch_size, sequence_length, hidden_size)`.
Hidden-states of the model at the output of each layer plus the initial embedding outputs. Hidden-states of the model at the output of each layer plus the initial embedding outputs.
attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True`` is passed or when ``config.output_attentions=True``): attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape :obj:`(batch_size, num_heads, Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`.
sequence_length, sequence_length)`.
Attentions weights after the attention softmax, used to compute the weighted average in the self-attention Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
heads. heads.
...@@ -550,69 +549,67 @@ class AlbertForPreTrainingOutput(ModelOutput): ...@@ -550,69 +549,67 @@ class AlbertForPreTrainingOutput(ModelOutput):
ALBERT_START_DOCSTRING = r""" ALBERT_START_DOCSTRING = r"""
This model inherits from :class:`~transformers.PreTrainedModel`. Check the superclass documentation for the generic This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic
methods the library implements for all its model (such as downloading or saving, resizing the input embeddings, methods the library implements for all its model (such as downloading or saving, resizing the input embeddings,
pruning heads etc.) pruning heads etc.)
This model is also a PyTorch `torch.nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`__ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module)
subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to
general usage and behavior. general usage and behavior.
Args: Args:
config (:class:`~transformers.AlbertConfig`): Model configuration class with all the parameters of the model. config ([`AlbertConfig`]): Model configuration class with all the parameters of the model.
Initializing with a config file does not load the weights associated with the model, only the Initializing with a config file does not load the weights associated with the model, only the
configuration. Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model
weights. weights.
""" """
ALBERT_INPUTS_DOCSTRING = r""" ALBERT_INPUTS_DOCSTRING = r"""
Args: Args:
input_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`): input_ids (`torch.LongTensor` of shape `({0})`):
Indices of input sequence tokens in the vocabulary. Indices of input sequence tokens in the vocabulary.
Indices can be obtained using :class:`~transformers.AlbertTokenizer`. See Indices can be obtained using [`AlbertTokenizer`]. See
:meth:`transformers.PreTrainedTokenizer.__call__` and :meth:`transformers.PreTrainedTokenizer.encode` for [`PreTrainedTokenizer.__call__`] and [`PreTrainedTokenizer.encode`] for
details. details.
`What are input IDs? <../glossary.html#input-ids>`__ [What are input IDs?](../glossary#input-ids)
attention_mask (:obj:`torch.FloatTensor` of shape :obj:`({0})`, `optional`): attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
- 1 for tokens that are **not masked**, - 1 for tokens that are **not masked**,
- 0 for tokens that are **masked**. - 0 for tokens that are **masked**.
`What are attention masks? <../glossary.html#attention-mask>`__ [What are attention masks?](../glossary#attention-mask)
token_type_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`): token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
Segment token indices to indicate first and second portions of the inputs. Indices are selected in ``[0, Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, 1]`:
1]``:
- 0 corresponds to a `sentence A` token, - 0 corresponds to a *sentence A* token,
- 1 corresponds to a `sentence B` token. - 1 corresponds to a *sentence B* token.
`What are token type IDs? <../glossary.html#token-type-ids>`_ [What are token type IDs?](../glossary#token-type-ids)
position_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`): position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
Indices of positions of each input sequence tokens in the position embeddings. Selected in the range ``[0, Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.max_position_embeddings - 1]`.
config.max_position_embeddings - 1]``.
`What are position IDs? <../glossary.html#position-ids>`_ [What are position IDs?](../glossary#position-ids)
head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`): head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
Mask to nullify selected heads of the self-attention modules. Mask values selected in ``[0, 1]``: Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
- 1 indicates the head is **not masked**, - 1 indicates the head is **not masked**,
- 0 indicates the head is **masked**. - 0 indicates the head is **masked**.
inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`({0}, hidden_size)`, `optional`): inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation. Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
This is useful if you want more control over how to convert :obj:`input_ids` indices into associated This is useful if you want more control over how to convert `input_ids` indices into associated
vectors than the model's internal embedding lookup matrix. vectors than the model's internal embedding lookup matrix.
output_attentions (:obj:`bool`, `optional`): output_attentions (`bool`, *optional*):
Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
tensors for more detail. tensors for more detail.
output_hidden_states (:obj:`bool`, `optional`): output_hidden_states (`bool`, *optional*):
Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
more detail. more detail.
return_dict (:obj:`bool`, `optional`): return_dict (`bool`, *optional*):
Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple. Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
""" """
...@@ -787,19 +784,19 @@ class AlbertForPreTraining(AlbertPreTrainedModel): ...@@ -787,19 +784,19 @@ class AlbertForPreTraining(AlbertPreTrainedModel):
return_dict=None, return_dict=None,
): ):
r""" r"""
labels (``torch.LongTensor`` of shape ``(batch_size, sequence_length)``, `optional`): labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
Labels for computing the masked language modeling loss. Indices should be in ``[-100, 0, ..., Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored
config.vocab_size]`` (see ``input_ids`` docstring) Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
(masked), the loss is only computed for the tokens with labels in ``[0, ..., config.vocab_size]`` sentence_order_label (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
sentence_order_label (``torch.LongTensor`` of shape ``(batch_size,)``, `optional`):
Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair
(see :obj:`input_ids` docstring) Indices should be in ``[0, 1]``. ``0`` indicates original order (sequence (see `input_ids` docstring) Indices should be in `[0, 1]`. `0` indicates original order (sequence
A, then sequence B), ``1`` indicates switched order (sequence B, then sequence A). A, then sequence B), `1` indicates switched order (sequence B, then sequence A).
Returns: Returns:
Example:: Example:
```python
>>> from transformers import AlbertTokenizer, AlbertForPreTraining >>> from transformers import AlbertTokenizer, AlbertForPreTraining
>>> import torch >>> import torch
...@@ -811,8 +808,7 @@ class AlbertForPreTraining(AlbertPreTrainedModel): ...@@ -811,8 +808,7 @@ class AlbertForPreTraining(AlbertPreTrainedModel):
>>> prediction_logits = outputs.prediction_logits >>> prediction_logits = outputs.prediction_logits
>>> sop_logits = outputs.sop_logits >>> sop_logits = outputs.sop_logits
```"""
"""
return_dict = return_dict if return_dict is not None else self.config.use_return_dict return_dict = return_dict if return_dict is not None else self.config.use_return_dict
outputs = self.albert( outputs = self.albert(
...@@ -938,10 +934,9 @@ class AlbertForMaskedLM(AlbertPreTrainedModel): ...@@ -938,10 +934,9 @@ class AlbertForMaskedLM(AlbertPreTrainedModel):
return_dict=None, return_dict=None,
): ):
r""" r"""
labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`): labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
Labels for computing the masked language modeling loss. Indices should be in ``[-100, 0, ..., Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored
config.vocab_size]`` (see ``input_ids`` docstring) Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
(masked), the loss is only computed for the tokens with labels in ``[0, ..., config.vocab_size]``
""" """
return_dict = return_dict if return_dict is not None else self.config.use_return_dict return_dict = return_dict if return_dict is not None else self.config.use_return_dict
...@@ -1018,10 +1013,9 @@ class AlbertForSequenceClassification(AlbertPreTrainedModel): ...@@ -1018,10 +1013,9 @@ class AlbertForSequenceClassification(AlbertPreTrainedModel):
return_dict=None, return_dict=None,
): ):
r""" r"""
labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`): labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
Labels for computing the sequence classification/regression loss. Indices should be in ``[0, ..., Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss),
config.num_labels - 1]``. If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss), If `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy).
""" """
return_dict = return_dict if return_dict is not None else self.config.use_return_dict return_dict = return_dict if return_dict is not None else self.config.use_return_dict
...@@ -1125,9 +1119,8 @@ class AlbertForTokenClassification(AlbertPreTrainedModel): ...@@ -1125,9 +1119,8 @@ class AlbertForTokenClassification(AlbertPreTrainedModel):
return_dict=None, return_dict=None,
): ):
r""" r"""
labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`): 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 - Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
1]``.
""" """
return_dict = return_dict if return_dict is not None else self.config.use_return_dict return_dict = return_dict if return_dict is not None else self.config.use_return_dict
...@@ -1217,13 +1210,13 @@ class AlbertForQuestionAnswering(AlbertPreTrainedModel): ...@@ -1217,13 +1210,13 @@ class AlbertForQuestionAnswering(AlbertPreTrainedModel):
return_dict=None, return_dict=None,
): ):
r""" r"""
start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`): 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. 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 (:obj:`sequence_length`). Position outside of the 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. sequence are not taken into account for computing the loss.
end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`): 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. 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 (:obj:`sequence_length`). Position outside of the 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. sequence are not taken into account for computing the loss.
""" """
return_dict = return_dict if return_dict is not None else self.config.use_return_dict return_dict = return_dict if return_dict is not None else self.config.use_return_dict
...@@ -1316,10 +1309,9 @@ class AlbertForMultipleChoice(AlbertPreTrainedModel): ...@@ -1316,10 +1309,9 @@ class AlbertForMultipleChoice(AlbertPreTrainedModel):
return_dict=None, return_dict=None,
): ):
r""" r"""
labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`): labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
Labels for computing the multiple choice classification loss. Indices should be in ``[0, ..., Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., num_choices-1]` where *num_choices* is the size of the second dimension of the input tensors. (see
num_choices-1]`` where `num_choices` is the size of the second dimension of the input tensors. (see *input_ids* above)
`input_ids` above)
""" """
return_dict = return_dict if return_dict is not None else self.config.use_return_dict return_dict = return_dict if return_dict is not None else self.config.use_return_dict
num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1] num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
......
...@@ -56,22 +56,21 @@ _TOKENIZER_FOR_DOC = "AlbertTokenizer" ...@@ -56,22 +56,21 @@ _TOKENIZER_FOR_DOC = "AlbertTokenizer"
@flax.struct.dataclass @flax.struct.dataclass
class FlaxAlbertForPreTrainingOutput(ModelOutput): class FlaxAlbertForPreTrainingOutput(ModelOutput):
""" """
Output type of :class:`~transformers.FlaxAlbertForPreTraining`. Output type of [`FlaxAlbertForPreTraining`].
Args: Args:
prediction_logits (:obj:`jnp.ndarray` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`): prediction_logits (`jnp.ndarray` of shape `(batch_size, sequence_length, config.vocab_size)`):
Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
sop_logits (:obj:`jnp.ndarray` of shape :obj:`(batch_size, 2)`): sop_logits (`jnp.ndarray` of shape `(batch_size, 2)`):
Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation
before SoftMax). before SoftMax).
hidden_states (:obj:`tuple(jnp.ndarray)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``): hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
Tuple of :obj:`jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of
shape :obj:`(batch_size, sequence_length, hidden_size)`. shape `(batch_size, sequence_length, hidden_size)`.
Hidden-states of the model at the output of each layer plus the initial embedding outputs. Hidden-states of the model at the output of each layer plus the initial embedding outputs.
attentions (:obj:`tuple(jnp.ndarray)`, `optional`, returned when ``output_attentions=True`` is passed or when ``config.output_attentions=True``): attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
Tuple of :obj:`jnp.ndarray` (one for each layer) of shape :obj:`(batch_size, num_heads, sequence_length, Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`.
sequence_length)`.
Attentions weights after the attention softmax, used to compute the weighted average in the self-attention Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
heads. heads.
...@@ -85,70 +84,67 @@ class FlaxAlbertForPreTrainingOutput(ModelOutput): ...@@ -85,70 +84,67 @@ class FlaxAlbertForPreTrainingOutput(ModelOutput):
ALBERT_START_DOCSTRING = r""" ALBERT_START_DOCSTRING = r"""
This model inherits from :class:`~transformers.FlaxPreTrainedModel`. Check the superclass documentation for the This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the
generic methods the library implements for all its model (such as downloading, saving and converting weights from generic methods the library implements for all its model (such as downloading, saving and converting weights from
PyTorch models) PyTorch models)
This model is also a Flax Linen `flax.linen.Module This model is also a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module) subclass. Use it as a regular Flax linen Module
<https://flax.readthedocs.io/en/latest/flax.linen.html#module>`__ subclass. Use it as a regular Flax linen Module
and refer to the Flax documentation for all matter related to general usage and behavior. and refer to the Flax documentation for all matter related to general usage and behavior.
Finally, this model supports inherent JAX features such as: Finally, this model supports inherent JAX features such as:
- `Just-In-Time (JIT) compilation <https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit>`__ - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)
- `Automatic Differentiation <https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation>`__ - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)
- `Vectorization <https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap>`__ - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)
- `Parallelization <https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap>`__ - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)
Parameters: Parameters:
config (:class:`~transformers.AlbertConfig`): Model configuration class with all the parameters of the model. config ([`AlbertConfig`]): Model configuration class with all the parameters of the model.
Initializing with a config file does not load the weights associated with the model, only the Initializing with a config file does not load the weights associated with the model, only the
configuration. Check out the :meth:`~transformers.FlaxPreTrainedModel.from_pretrained` method to load the configuration. Check out the [`~FlaxPreTrainedModel.from_pretrained`] method to load the
model weights. model weights.
dtype (:obj:`jax.numpy.dtype`, `optional`, defaults to :obj:`jax.numpy.float32`): dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):
The data type of the computation. Can be one of :obj:`jax.numpy.float32`, :obj:`jax.numpy.float16` (on The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on
GPUs) and :obj:`jax.numpy.bfloat16` (on TPUs). GPUs) and `jax.numpy.bfloat16` (on TPUs).
This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If
specified all the computation will be performed with the given ``dtype``. specified all the computation will be performed with the given `dtype`.
**Note that this only specifies the dtype of the computation and does not influence the dtype of model **Note that this only specifies the dtype of the computation and does not influence the dtype of model
parameters.** parameters.**
If you wish to change the dtype of the model parameters, see If you wish to change the dtype of the model parameters, see
:meth:`~transformers.FlaxPreTrainedModel.to_fp16` and :meth:`~transformers.FlaxPreTrainedModel.to_bf16`. [`~FlaxPreTrainedModel.to_fp16`] and [`~FlaxPreTrainedModel.to_bf16`].
""" """
ALBERT_INPUTS_DOCSTRING = r""" ALBERT_INPUTS_DOCSTRING = r"""
Args: Args:
input_ids (:obj:`numpy.ndarray` of shape :obj:`({0})`): input_ids (`numpy.ndarray` of shape `({0})`):
Indices of input sequence tokens in the vocabulary. Indices of input sequence tokens in the vocabulary.
Indices can be obtained using :class:`~transformers.AlbertTokenizer`. See Indices can be obtained using [`AlbertTokenizer`]. See
:meth:`transformers.PreTrainedTokenizer.encode` and :func:`transformers.PreTrainedTokenizer.__call__` for [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for
details. details.
`What are input IDs? <../glossary.html#input-ids>`__ [What are input IDs?](../glossary#input-ids)
attention_mask (:obj:`numpy.ndarray` of shape :obj:`({0})`, `optional`): attention_mask (`numpy.ndarray` of shape `({0})`, *optional*):
Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
- 1 for tokens that are **not masked**, - 1 for tokens that are **not masked**,
- 0 for tokens that are **masked**. - 0 for tokens that are **masked**.
`What are attention masks? <../glossary.html#attention-mask>`__ [What are attention masks?](../glossary#attention-mask)
token_type_ids (:obj:`numpy.ndarray` of shape :obj:`({0})`, `optional`): token_type_ids (`numpy.ndarray` of shape `({0})`, *optional*):
Segment token indices to indicate first and second portions of the inputs. Indices are selected in ``[0, Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, 1]`:
1]``:
- 0 corresponds to a `sentence A` token, - 0 corresponds to a *sentence A* token,
- 1 corresponds to a `sentence B` token. - 1 corresponds to a *sentence B* token.
`What are token type IDs? <../glossary.html#token-type-ids>`__ [What are token type IDs?](../glossary#token-type-ids)
position_ids (:obj:`numpy.ndarray` of shape :obj:`({0})`, `optional`): position_ids (`numpy.ndarray` of shape `({0})`, *optional*):
Indices of positions of each input sequence tokens in the position embeddings. Selected in the range ``[0, Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.max_position_embeddings - 1]`.
config.max_position_embeddings - 1]``. return_dict (`bool`, *optional*):
return_dict (:obj:`bool`, `optional`): Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple.
""" """
......
...@@ -163,7 +163,7 @@ class TFAlbertEmbeddings(tf.keras.layers.Layer): ...@@ -163,7 +163,7 @@ class TFAlbertEmbeddings(tf.keras.layers.Layer):
Applies embedding based on inputs tensor. Applies embedding based on inputs tensor.
Returns: Returns:
final_embeddings (:obj:`tf.Tensor`): output embedding tensor. final_embeddings (`tf.Tensor`): output embedding tensor.
""" """
if input_ids is None and inputs_embeds is None: if input_ids is None and inputs_embeds is None:
raise ValueError("Need to provide either `input_ids` or `input_embeds`.") raise ValueError("Need to provide either `input_ids` or `input_embeds`.")
...@@ -648,22 +648,21 @@ class TFAlbertMainLayer(tf.keras.layers.Layer): ...@@ -648,22 +648,21 @@ class TFAlbertMainLayer(tf.keras.layers.Layer):
@dataclass @dataclass
class TFAlbertForPreTrainingOutput(ModelOutput): class TFAlbertForPreTrainingOutput(ModelOutput):
""" """
Output type of :class:`~transformers.TFAlbertForPreTraining`. Output type of [`TFAlbertForPreTraining`].
Args: Args:
prediction_logits (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`): prediction_logits (`tf.Tensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
sop_logits (:obj:`tf.Tensor` of shape :obj:`(batch_size, 2)`): sop_logits (`tf.Tensor` of shape `(batch_size, 2)`):
Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation
before SoftMax). before SoftMax).
hidden_states (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``): hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
Tuple of :obj:`tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of
shape :obj:`(batch_size, sequence_length, hidden_size)`. shape `(batch_size, sequence_length, hidden_size)`.
Hidden-states of the model at the output of each layer plus the initial embedding outputs. Hidden-states of the model at the output of each layer plus the initial embedding outputs.
attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True`` is passed or when ``config.output_attentions=True``): attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
Tuple of :obj:`tf.Tensor` (one for each layer) of shape :obj:`(batch_size, num_heads, sequence_length, Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`.
sequence_length)`.
Attentions weights after the attention softmax, used to compute the weighted average in the self-attention Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
heads. heads.
...@@ -678,92 +677,92 @@ class TFAlbertForPreTrainingOutput(ModelOutput): ...@@ -678,92 +677,92 @@ class TFAlbertForPreTrainingOutput(ModelOutput):
ALBERT_START_DOCSTRING = r""" ALBERT_START_DOCSTRING = r"""
This model inherits from :class:`~transformers.TFPreTrainedModel`. Check the superclass documentation for the This model inherits from [`TFPreTrainedModel`]. Check the superclass documentation for the
generic methods the library implements for all its model (such as downloading or saving, resizing the input generic methods the library implements for all its model (such as downloading or saving, resizing the input
embeddings, pruning heads etc.) embeddings, pruning heads etc.)
This model is also a `tf.keras.Model <https://www.tensorflow.org/api_docs/python/tf/keras/Model>`__ subclass. Use This model is also a [tf.keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model) subclass. Use
it as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage it as a regular TF 2.0 Keras Model and refer to the TF 2.0 documentation for all matter related to general usage
and behavior. and behavior.
.. note:: <Tip>
TF 2.0 models accepts two formats as inputs: TF 2.0 models accepts two formats as inputs:
- having all inputs as keyword arguments (like PyTorch models), or - having all inputs as keyword arguments (like PyTorch models), or
- having all inputs as a list, tuple or dict in the first positional arguments. - having all inputs as a list, tuple or dict in the first positional arguments.
This second option is useful when using :meth:`tf.keras.Model.fit` method which currently requires having all This second option is useful when using [`tf.keras.Model.fit`] method which currently requires having all
the tensors in the first argument of the model call function: :obj:`model(inputs)`. the tensors in the first argument of the model call function: `model(inputs)`.
If you choose this second option, there are three possibilities you can use to gather all the input Tensors in If you choose this second option, there are three possibilities you can use to gather all the input Tensors in
the first positional argument : the first positional argument :
- a single Tensor with :obj:`input_ids` only and nothing else: :obj:`model(inputs_ids)` - a single Tensor with `input_ids` only and nothing else: `model(inputs_ids)`
- a list of varying length with one or several input Tensors IN THE ORDER given in the docstring: - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
:obj:`model([input_ids, attention_mask])` or :obj:`model([input_ids, attention_mask, token_type_ids])` `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])`
- a dictionary with one or several input Tensors associated to the input names given in the docstring: - a dictionary with one or several input Tensors associated to the input names given in the docstring:
:obj:`model({"input_ids": input_ids, "token_type_ids": token_type_ids})` `model({"input_ids": input_ids, "token_type_ids": token_type_ids})`
</Tip>
Args: Args:
config (:class:`~transformers.AlbertConfig`): Model configuration class with all the parameters of the model. config ([`AlbertConfig`]): Model configuration class with all the parameters of the model.
Initializing with a config file does not load the weights associated with the model, only the Initializing with a config file does not load the weights associated with the model, only the
configuration. Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model
weights. weights.
""" """
ALBERT_INPUTS_DOCSTRING = r""" ALBERT_INPUTS_DOCSTRING = r"""
Args: Args:
input_ids (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`({0})`): input_ids (`Numpy array` or `tf.Tensor` of shape `({0})`):
Indices of input sequence tokens in the vocabulary. Indices of input sequence tokens in the vocabulary.
Indices can be obtained using :class:`~transformers.AlbertTokenizer`. See Indices can be obtained using [`AlbertTokenizer`]. See
:func:`transformers.PreTrainedTokenizer.__call__` and :func:`transformers.PreTrainedTokenizer.encode` for [`PreTrainedTokenizer.__call__`] and [`PreTrainedTokenizer.encode`] for
details. details.
`What are input IDs? <../glossary.html#input-ids>`__ [What are input IDs?](../glossary#input-ids)
attention_mask (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`({0})`, `optional`): attention_mask (`Numpy array` or `tf.Tensor` of shape `({0})`, *optional*):
Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
- 1 for tokens that are **not masked**, - 1 for tokens that are **not masked**,
- 0 for tokens that are **masked**. - 0 for tokens that are **masked**.
`What are attention masks? <../glossary.html#attention-mask>`__ [What are attention masks?](../glossary#attention-mask)
token_type_ids (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`({0})`, `optional`): token_type_ids (`Numpy array` or `tf.Tensor` of shape `({0})`, *optional*):
Segment token indices to indicate first and second portions of the inputs. Indices are selected in ``[0, Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, 1]`:
1]``:
- 0 corresponds to a `sentence A` token, - 0 corresponds to a *sentence A* token,
- 1 corresponds to a `sentence B` token. - 1 corresponds to a *sentence B* token.
`What are token type IDs? <../glossary.html#token-type-ids>`_ [What are token type IDs?](../glossary#token-type-ids)
position_ids (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`({0})`, `optional`): position_ids (`Numpy array` or `tf.Tensor` of shape `({0})`, *optional*):
Indices of positions of each input sequence tokens in the position embeddings. Selected in the range ``[0, Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.max_position_embeddings - 1]`.
config.max_position_embeddings - 1]``.
`What are position IDs? <../glossary.html#position-ids>`_ [What are position IDs?](../glossary#position-ids)
head_mask (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`): head_mask (`Numpy array` or `tf.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
Mask to nullify selected heads of the self-attention modules. Mask values selected in ``[0, 1]``: Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
- 1 indicates the head is **not masked**, - 1 indicates the head is **not masked**,
- 0 indicates the head is **masked**. - 0 indicates the head is **masked**.
inputs_embeds (:obj:`tf.Tensor` of shape :obj:`({0}, hidden_size)`, `optional`): inputs_embeds (`tf.Tensor` of shape `({0}, hidden_size)`, *optional*):
Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation. Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
This is useful if you want more control over how to convert :obj:`input_ids` indices into associated This is useful if you want more control over how to convert `input_ids` indices into associated
vectors than the model's internal embedding lookup matrix. vectors than the model's internal embedding lookup matrix.
output_attentions (:obj:`bool`, `optional`): output_attentions (`bool`, *optional*):
Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
tensors for more detail. This argument can be used only in eager mode, in graph mode the value in the tensors for more detail. This argument can be used only in eager mode, in graph mode the value in the
config will be used instead. config will be used instead.
output_hidden_states (:obj:`bool`, `optional`): output_hidden_states (`bool`, *optional*):
Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
more detail. This argument can be used only in eager mode, in graph mode the value in the config will be more detail. This argument can be used only in eager mode, in graph mode the value in the config will be
used instead. used instead.
return_dict (:obj:`bool`, `optional`): return_dict (`bool`, *optional*):
Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple. This Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. This
argument can be used in eager mode, in graph mode the value will always be set to True. argument can be used in eager mode, in graph mode the value will always be set to True.
training (:obj:`bool`, `optional`, defaults to :obj:`False`): training (`bool`, *optional*, defaults to `False`):
Whether or not to use the model in training mode (some modules like dropout modules have different Whether or not to use the model in training mode (some modules like dropout modules have different
behaviors between training and evaluation). behaviors between training and evaluation).
""" """
...@@ -1019,10 +1018,9 @@ class TFAlbertForMaskedLM(TFAlbertPreTrainedModel, TFMaskedLanguageModelingLoss) ...@@ -1019,10 +1018,9 @@ class TFAlbertForMaskedLM(TFAlbertPreTrainedModel, TFMaskedLanguageModelingLoss)
**kwargs, **kwargs,
) -> Union[TFMaskedLMOutput, Tuple[tf.Tensor]]: ) -> Union[TFMaskedLMOutput, Tuple[tf.Tensor]]:
r""" r"""
labels (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length)`, `optional`): labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
Labels for computing the masked language modeling loss. Indices should be in ``[-100, 0, ..., Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored
config.vocab_size]`` (see ``input_ids`` docstring) Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
(masked), the loss is only computed for the tokens with labels in ``[0, ..., config.vocab_size]``
""" """
inputs = input_processing( inputs = input_processing(
func=self.call, func=self.call,
...@@ -1124,10 +1122,9 @@ class TFAlbertForSequenceClassification(TFAlbertPreTrainedModel, TFSequenceClass ...@@ -1124,10 +1122,9 @@ class TFAlbertForSequenceClassification(TFAlbertPreTrainedModel, TFSequenceClass
**kwargs, **kwargs,
) -> Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]: ) -> Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:
r""" r"""
labels (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`): labels (`tf.Tensor` of shape `(batch_size,)`, *optional*):
Labels for computing the sequence classification/regression loss. Indices should be in ``[0, ..., Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss),
config.num_labels - 1]``. If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss), If `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy).
""" """
inputs = input_processing( inputs = input_processing(
func=self.call, func=self.call,
...@@ -1233,9 +1230,8 @@ class TFAlbertForTokenClassification(TFAlbertPreTrainedModel, TFTokenClassificat ...@@ -1233,9 +1230,8 @@ class TFAlbertForTokenClassification(TFAlbertPreTrainedModel, TFTokenClassificat
**kwargs, **kwargs,
) -> Union[TFTokenClassifierOutput, Tuple[tf.Tensor]]: ) -> Union[TFTokenClassifierOutput, Tuple[tf.Tensor]]:
r""" r"""
labels (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length)`, `optional`): labels (`tf.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
Labels for computing the token classification loss. Indices should be in ``[0, ..., config.num_labels - Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
1]``.
""" """
inputs = input_processing( inputs = input_processing(
func=self.call, func=self.call,
...@@ -1335,13 +1331,13 @@ class TFAlbertForQuestionAnswering(TFAlbertPreTrainedModel, TFQuestionAnsweringL ...@@ -1335,13 +1331,13 @@ class TFAlbertForQuestionAnswering(TFAlbertPreTrainedModel, TFQuestionAnsweringL
**kwargs, **kwargs,
) -> Union[TFQuestionAnsweringModelOutput, Tuple[tf.Tensor]]: ) -> Union[TFQuestionAnsweringModelOutput, Tuple[tf.Tensor]]:
r""" r"""
start_positions (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`): start_positions (`tf.Tensor` of shape `(batch_size,)`, *optional*):
Labels for position (index) of the start of the labelled span for computing the token classification loss. 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 (:obj:`sequence_length`). Position outside of the 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. sequence are not taken into account for computing the loss.
end_positions (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`): end_positions (`tf.Tensor` of shape `(batch_size,)`, *optional*):
Labels for position (index) of the end of the labelled span for computing the token classification loss. 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 (:obj:`sequence_length`). Position outside of the 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. sequence are not taken into account for computing the loss.
""" """
inputs = input_processing( inputs = input_processing(
...@@ -1462,10 +1458,9 @@ class TFAlbertForMultipleChoice(TFAlbertPreTrainedModel, TFMultipleChoiceLoss): ...@@ -1462,10 +1458,9 @@ class TFAlbertForMultipleChoice(TFAlbertPreTrainedModel, TFMultipleChoiceLoss):
**kwargs, **kwargs,
) -> Union[TFMultipleChoiceModelOutput, Tuple[tf.Tensor]]: ) -> Union[TFMultipleChoiceModelOutput, Tuple[tf.Tensor]]:
r""" r"""
labels (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`): labels (`tf.Tensor` of shape `(batch_size,)`, *optional*):
Labels for computing the multiple choice classification loss. Indices should be in ``[0, ..., Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., num_choices]` where `num_choices` is the size of the second dimension of the input tensors. (See
num_choices]`` where :obj:`num_choices` is the size of the second dimension of the input tensors. (See `input_ids` above)
:obj:`input_ids` above)
""" """
inputs = input_processing( inputs = input_processing(
func=self.call, func=self.call,
......
...@@ -59,170 +59,165 @@ _TOKENIZER_FOR_DOC = "BartTokenizer" ...@@ -59,170 +59,165 @@ _TOKENIZER_FOR_DOC = "BartTokenizer"
BART_START_DOCSTRING = r""" BART_START_DOCSTRING = r"""
This model inherits from :class:`~transformers.FlaxPreTrainedModel`. Check the superclass documentation for the This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the
generic methods the library implements for all its model (such as downloading or saving, resizing the input generic methods the library implements for all its model (such as downloading or saving, resizing the input
embeddings, pruning heads etc.) embeddings, pruning heads etc.)
This model is also a Flax Linen `flax.nn.Module This model is also a Flax Linen [flax.nn.Module](https://flax.readthedocs.io/en/latest/_autosummary/flax.nn.module.html) subclass. Use it as a regular Flax
<https://flax.readthedocs.io/en/latest/_autosummary/flax.nn.module.html>`__ subclass. Use it as a regular Flax
Module and refer to the Flax documentation for all matter related to general usage and behavior. Module and refer to the Flax documentation for all matter related to general usage and behavior.
Finally, this model supports inherent JAX features such as: Finally, this model supports inherent JAX features such as:
- `Just-In-Time (JIT) compilation <https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit>`__ - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)
- `Automatic Differentiation <https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation>`__ - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)
- `Vectorization <https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap>`__ - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)
- `Parallelization <https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap>`__ - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)
Parameters: Parameters:
config (:class:`~transformers.BartConfig`): Model configuration class with all the parameters of the model. config ([`BartConfig`]): Model configuration class with all the parameters of the model.
Initializing with a config file does not load the weights associated with the model, only the Initializing with a config file does not load the weights associated with the model, only the
configuration. Check out the :meth:`~transformers.FlaxPreTrainedModel.from_pretrained` method to load the configuration. Check out the [`~FlaxPreTrainedModel.from_pretrained`] method to load the
model weights. model weights.
dtype (:obj:`jax.numpy.dtype`, `optional`, defaults to :obj:`jax.numpy.float32`): dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):
The data type of the computation. Can be one of :obj:`jax.numpy.float32`, :obj:`jax.numpy.float16` (on The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on
GPUs) and :obj:`jax.numpy.bfloat16` (on TPUs). GPUs) and `jax.numpy.bfloat16` (on TPUs).
This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If
specified all the computation will be performed with the given ``dtype``. specified all the computation will be performed with the given `dtype`.
**Note that this only specifies the dtype of the computation and does not influence the dtype of model **Note that this only specifies the dtype of the computation and does not influence the dtype of model
parameters.** parameters.**
If you wish to change the dtype of the model parameters, see If you wish to change the dtype of the model parameters, see
:meth:`~transformers.FlaxPreTrainedModel.to_fp16` and :meth:`~transformers.FlaxPreTrainedModel.to_bf16`. [`~FlaxPreTrainedModel.to_fp16`] and [`~FlaxPreTrainedModel.to_bf16`].
""" """
BART_INPUTS_DOCSTRING = r""" BART_INPUTS_DOCSTRING = r"""
Args: Args:
input_ids (:obj:`jnp.ndarray` of shape :obj:`(batch_size, sequence_length)`): input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):
Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
it. it.
Indices can be obtained using :class:`~transformers.BartTokenizer`. See Indices can be obtained using [`BartTokenizer`]. See
:meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for
details. details.
`What are input IDs? <../glossary.html#input-ids>`__ [What are input IDs?](../glossary#input-ids)
attention_mask (:obj:`jnp.ndarray` of shape :obj:`(batch_size, sequence_length)`, `optional`): attention_mask (`jnp.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
- 1 for tokens that are **not masked**, - 1 for tokens that are **not masked**,
- 0 for tokens that are **masked**. - 0 for tokens that are **masked**.
`What are attention masks? <../glossary.html#attention-mask>`__ [What are attention masks?](../glossary#attention-mask)
decoder_input_ids (:obj:`jnp.ndarray` of shape :obj:`(batch_size, target_sequence_length)`, `optional`): decoder_input_ids (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):
Indices of decoder input sequence tokens in the vocabulary. Indices of decoder input sequence tokens in the vocabulary.
Indices can be obtained using :class:`~transformers.BartTokenizer`. See Indices can be obtained using [`BartTokenizer`]. See
:meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for
details. details.
`What are decoder input IDs? <../glossary.html#decoder-input-ids>`__ [What are decoder input IDs?](../glossary#decoder-input-ids)
For translation and summarization training, :obj:`decoder_input_ids` should be provided. If no For translation and summarization training, `decoder_input_ids` should be provided. If no
:obj:`decoder_input_ids` is provided, the model will create this tensor by shifting the :obj:`input_ids` to `decoder_input_ids` is provided, the model will create this tensor by shifting the `input_ids` to
the right for denoising pre-training following the paper. the right for denoising pre-training following the paper.
decoder_attention_mask (:obj:`jnp.ndarray` of shape :obj:`(batch_size, target_sequence_length)`, `optional`): decoder_attention_mask (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):
Default behavior: generate a tensor that ignores pad tokens in :obj:`decoder_input_ids`. Causal mask will Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will
also be used by default. also be used by default.
If you want to change padding behavior, you should modify to your needs. See diagram 1 in `the paper If you want to change padding behavior, you should modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more information on the default strategy.
<https://arxiv.org/abs/1910.13461>`__ for more information on the default strategy. position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
position_ids (:obj:`numpy.ndarray` of shape :obj:`(batch_size, sequence_length)`, `optional`): Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.max_position_embeddings - 1]`.
Indices of positions of each input sequence tokens in the position embeddings. Selected in the range ``[0, decoder_position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
config.max_position_embeddings - 1]``.
decoder_position_ids (:obj:`numpy.ndarray` of shape :obj:`(batch_size, sequence_length)`, `optional`):
Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the
range ``[0, config.max_position_embeddings - 1]``. range `[0, config.max_position_embeddings - 1]`.
output_attentions (:obj:`bool`, `optional`): output_attentions (`bool`, *optional*):
Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
tensors for more detail. tensors for more detail.
output_hidden_states (:obj:`bool`, `optional`): output_hidden_states (`bool`, *optional*):
Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
more detail. more detail.
return_dict (:obj:`bool`, `optional`): return_dict (`bool`, *optional*):
Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple. Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
""" """
BART_ENCODE_INPUTS_DOCSTRING = r""" BART_ENCODE_INPUTS_DOCSTRING = r"""
Args: Args:
input_ids (:obj:`jnp.ndarray` of shape :obj:`(batch_size, sequence_length)`): input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):
Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
it. it.
Indices can be obtained using :class:`~transformers.BartTokenizer`. See Indices can be obtained using [`BartTokenizer`]. See
:meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for
details. details.
`What are input IDs? <../glossary.html#input-ids>`__ [What are input IDs?](../glossary#input-ids)
attention_mask (:obj:`jnp.ndarray` of shape :obj:`(batch_size, sequence_length)`, `optional`): attention_mask (`jnp.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
- 1 for tokens that are **not masked**, - 1 for tokens that are **not masked**,
- 0 for tokens that are **masked**. - 0 for tokens that are **masked**.
`What are attention masks? <../glossary.html#attention-mask>`__ [What are attention masks?](../glossary#attention-mask)
position_ids (:obj:`numpy.ndarray` of shape :obj:`(batch_size, sequence_length)`, `optional`): position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
Indices of positions of each input sequence tokens in the position embeddings. Selected in the range ``[0, Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.max_position_embeddings - 1]`.
config.max_position_embeddings - 1]``. output_attentions (`bool`, *optional*):
output_attentions (:obj:`bool`, `optional`): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned
tensors for more detail. tensors for more detail.
output_hidden_states (:obj:`bool`, `optional`): output_hidden_states (`bool`, *optional*):
Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
more detail. more detail.
return_dict (:obj:`bool`, `optional`): return_dict (`bool`, *optional*):
Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple. Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
""" """
BART_DECODE_INPUTS_DOCSTRING = r""" BART_DECODE_INPUTS_DOCSTRING = r"""
Args: Args:
decoder_input_ids (:obj:`jnp.ndarray` of shape :obj:`(batch_size, target_sequence_length)`): decoder_input_ids (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`):
Indices of decoder input sequence tokens in the vocabulary. Indices of decoder input sequence tokens in the vocabulary.
Indices can be obtained using :class:`~transformers.BartTokenizer`. See Indices can be obtained using [`BartTokenizer`]. See
:meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for
details. details.
`What are decoder input IDs? <../glossary.html#decoder-input-ids>`__ [What are decoder input IDs?](../glossary#decoder-input-ids)
For translation and summarization training, :obj:`decoder_input_ids` should be provided. If no For translation and summarization training, `decoder_input_ids` should be provided. If no
:obj:`decoder_input_ids` is provided, the model will create this tensor by shifting the :obj:`input_ids` to `decoder_input_ids` is provided, the model will create this tensor by shifting the `input_ids` to
the right for denoising pre-training following the paper. the right for denoising pre-training following the paper.
encoder_outputs (:obj:`tuple(tuple(jnp.ndarray)`): encoder_outputs (`tuple(tuple(jnp.ndarray)`):
Tuple consists of (:obj:`last_hidden_state`, `optional`: :obj:`hidden_states`, `optional`: Tuple consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*:
:obj:`attentions`) :obj:`last_hidden_state` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `attentions`) `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`,
`optional`) is a sequence of hidden-states at the output of the last layer of the encoder. Used in the *optional*) is a sequence of hidden-states at the output of the last layer of the encoder. Used in the
cross-attention of the decoder. cross-attention of the decoder.
encoder_attention_mask (:obj:`jnp.ndarray` of shape :obj:`(batch_size, sequence_length)`, `optional`): encoder_attention_mask (`jnp.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
- 1 for tokens that are **not masked**, - 1 for tokens that are **not masked**,
- 0 for tokens that are **masked**. - 0 for tokens that are **masked**.
`What are attention masks? <../glossary.html#attention-mask>`__ [What are attention masks?](../glossary#attention-mask)
decoder_attention_mask (:obj:`jnp.ndarray` of shape :obj:`(batch_size, target_sequence_length)`, `optional`): decoder_attention_mask (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):
Default behavior: generate a tensor that ignores pad tokens in :obj:`decoder_input_ids`. Causal mask will Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will
also be used by default. also be used by default.
If you want to change padding behavior, you should modify to your needs. See diagram 1 in `the paper If you want to change padding behavior, you should modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more information on the default strategy.
<https://arxiv.org/abs/1910.13461>`__ for more information on the default strategy. decoder_position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):
decoder_position_ids (:obj:`numpy.ndarray` of shape :obj:`(batch_size, sequence_length)`, `optional`):
Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the
range ``[0, config.max_position_embeddings - 1]``. range `[0, config.max_position_embeddings - 1]`.
past_key_values (:obj:`Dict[str, np.ndarray]`, `optional`, returned by ``init_cache`` or when passing previous ``past_key_values``): past_key_values (`Dict[str, np.ndarray]`, *optional*, returned by `init_cache` or when passing previous `past_key_values`):
Dictionary of pre-computed hidden-states (key and values in the attention blocks) that can be used for fast Dictionary of pre-computed hidden-states (key and values in the attention blocks) that can be used for fast
auto-regressive decoding. Pre-computed key and value hidden-states are of shape `[batch_size, max_length]`. auto-regressive decoding. Pre-computed key and value hidden-states are of shape *[batch_size, max_length]*.
output_attentions (:obj:`bool`, `optional`): output_attentions (`bool`, *optional*):
Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
tensors for more detail. tensors for more detail.
output_hidden_states (:obj:`bool`, `optional`): output_hidden_states (`bool`, *optional*):
Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
more detail. more detail.
return_dict (:obj:`bool`, `optional`): return_dict (`bool`, *optional*):
Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple. Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
""" """
...@@ -967,15 +962,14 @@ class FlaxBartPreTrainedModel(FlaxPreTrainedModel): ...@@ -967,15 +962,14 @@ class FlaxBartPreTrainedModel(FlaxPreTrainedModel):
def init_cache(self, batch_size, max_length, encoder_outputs): def init_cache(self, batch_size, max_length, encoder_outputs):
r""" r"""
Args: Args:
batch_size (:obj:`int`): batch_size (`int`):
batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache. batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.
max_length (:obj:`int`): max_length (`int`):
maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized
cache. cache.
encoder_outputs (:obj:`Union[FlaxBaseModelOutput, tuple(tuple(jnp.ndarray)]`): encoder_outputs (`Union[FlaxBaseModelOutput, tuple(tuple(jnp.ndarray)]`):
``encoder_outputs`` consists of (:obj:`last_hidden_state`, `optional`: :obj:`hidden_states`, `encoder_outputs` consists of (`last_hidden_state`, *optional*: `hidden_states`,
`optional`: :obj:`attentions`). :obj:`last_hidden_state` of shape :obj:`(batch_size, sequence_length, *optional*: `attentions`). `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*) is a sequence of hidden-states at the output of the last layer of the
hidden_size)`, `optional`) is a sequence of hidden-states at the output of the last layer of the
encoder. Used in the cross-attention of the decoder. encoder. Used in the cross-attention of the decoder.
""" """
# init input variables to retrieve cache # init input variables to retrieve cache
......
...@@ -46,23 +46,22 @@ BEIT_PRETRAINED_MODEL_ARCHIVE_LIST = [ ...@@ -46,23 +46,22 @@ BEIT_PRETRAINED_MODEL_ARCHIVE_LIST = [
@dataclass @dataclass
class BeitModelOutputWithPooling(BaseModelOutputWithPooling): class BeitModelOutputWithPooling(BaseModelOutputWithPooling):
""" """
Class for outputs of :class:`~transformers.BeitModel`. Class for outputs of [`BeitModel`].
Args: Args:
last_hidden_state (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`): last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
Sequence of hidden-states at the output of the last layer of the model. Sequence of hidden-states at the output of the last layer of the model.
pooler_output (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, hidden_size)`): pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`):
Average of the last layer hidden states of the patch tokens (excluding the `[CLS]` token) if Average of the last layer hidden states of the patch tokens (excluding the *[CLS]* token) if
`config.use_mean_pooling` is set to True. If set to False, then the final hidden state of the `[CLS]` token *config.use_mean_pooling* is set to True. If set to False, then the final hidden state of the *[CLS]* token
will be returned. will be returned.
hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``): hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
of shape :obj:`(batch_size, sequence_length, hidden_size)`. of shape `(batch_size, sequence_length, hidden_size)`.
Hidden-states of the model at the output of each layer plus the initial embedding outputs. Hidden-states of the model at the output of each layer plus the initial embedding outputs.
attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True`` is passed or when ``config.output_attentions=True``): attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape :obj:`(batch_size, num_heads, Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`.
sequence_length, sequence_length)`.
Attentions weights after the attention softmax, used to compute the weighted average in the self-attention Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
heads. heads.
...@@ -548,37 +547,37 @@ class BeitPreTrainedModel(PreTrainedModel): ...@@ -548,37 +547,37 @@ class BeitPreTrainedModel(PreTrainedModel):
BEIT_START_DOCSTRING = r""" BEIT_START_DOCSTRING = r"""
This model is a PyTorch `torch.nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`_ subclass. Use This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use
it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
behavior. behavior.
Parameters: Parameters:
config (:class:`~transformers.BeitConfig`): Model configuration class with all the parameters of the model. config ([`BeitConfig`]): Model configuration class with all the parameters of the model.
Initializing with a config file does not load the weights associated with the model, only the Initializing with a config file does not load the weights associated with the model, only the
configuration. Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model
weights. weights.
""" """
BEIT_INPUTS_DOCSTRING = r""" BEIT_INPUTS_DOCSTRING = r"""
Args: Args:
pixel_values (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_channels, height, width)`): pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
Pixel values. Pixel values can be obtained using :class:`~transformers.BeitFeatureExtractor`. See Pixel values. Pixel values can be obtained using [`BeitFeatureExtractor`]. See
:meth:`transformers.BeitFeatureExtractor.__call__` for details. [`BeitFeatureExtractor.__call__`] for details.
head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`): head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
Mask to nullify selected heads of the self-attention modules. Mask values selected in ``[0, 1]``: Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
- 1 indicates the head is **not masked**, - 1 indicates the head is **not masked**,
- 0 indicates the head is **masked**. - 0 indicates the head is **masked**.
output_attentions (:obj:`bool`, `optional`): output_attentions (`bool`, *optional*):
Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
tensors for more detail. tensors for more detail.
output_hidden_states (:obj:`bool`, `optional`): output_hidden_states (`bool`, *optional*):
Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
more detail. more detail.
return_dict (:obj:`bool`, `optional`): return_dict (`bool`, *optional*):
Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple. Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
""" """
...@@ -733,18 +732,18 @@ class BeitForMaskedImageModeling(BeitPreTrainedModel): ...@@ -733,18 +732,18 @@ class BeitForMaskedImageModeling(BeitPreTrainedModel):
return_dict=None, return_dict=None,
): ):
r""" r"""
bool_masked_pos (:obj:`torch.BoolTensor` of shape :obj:`(batch_size, num_patches)`): bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`):
Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`): labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
Labels for computing the image classification/regression loss. Indices should be in :obj:`[0, ..., Labels for computing the image classification/regression loss. Indices should be in `[0, ..., config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss),
config.num_labels - 1]`. If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss), If `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
Returns: Returns:
Examples:: Examples:
```python
>>> from transformers import BeitFeatureExtractor, BeitForMaskedImageModeling >>> from transformers import BeitFeatureExtractor, BeitForMaskedImageModeling
>>> from PIL import Image >>> from PIL import Image
>>> import requests >>> import requests
...@@ -758,7 +757,7 @@ class BeitForMaskedImageModeling(BeitPreTrainedModel): ...@@ -758,7 +757,7 @@ class BeitForMaskedImageModeling(BeitPreTrainedModel):
>>> inputs = feature_extractor(images=image, return_tensors="pt") >>> inputs = feature_extractor(images=image, return_tensors="pt")
>>> outputs = model(**inputs) >>> outputs = model(**inputs)
>>> logits = outputs.logits >>> logits = outputs.logits
""" ```"""
return_dict = return_dict if return_dict is not None else self.config.use_return_dict return_dict = return_dict if return_dict is not None else self.config.use_return_dict
outputs = self.beit( outputs = self.beit(
...@@ -823,15 +822,15 @@ class BeitForImageClassification(BeitPreTrainedModel): ...@@ -823,15 +822,15 @@ class BeitForImageClassification(BeitPreTrainedModel):
return_dict=None, return_dict=None,
): ):
r""" r"""
labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`): labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
Labels for computing the image classification/regression loss. Indices should be in :obj:`[0, ..., Labels for computing the image classification/regression loss. Indices should be in `[0, ..., config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss),
config.num_labels - 1]`. If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss), If `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
Returns: Returns:
Examples:: Examples:
```python
>>> from transformers import BeitFeatureExtractor, BeitForImageClassification >>> from transformers import BeitFeatureExtractor, BeitForImageClassification
>>> from PIL import Image >>> from PIL import Image
>>> import requests >>> import requests
...@@ -848,7 +847,7 @@ class BeitForImageClassification(BeitPreTrainedModel): ...@@ -848,7 +847,7 @@ class BeitForImageClassification(BeitPreTrainedModel):
>>> # model predicts one of the 1000 ImageNet classes >>> # model predicts one of the 1000 ImageNet classes
>>> predicted_class_idx = logits.argmax(-1).item() >>> predicted_class_idx = logits.argmax(-1).item()
>>> print("Predicted class:", model.config.id2label[predicted_class_idx]) >>> print("Predicted class:", model.config.id2label[predicted_class_idx])
""" ```"""
return_dict = return_dict if return_dict is not None else self.config.use_return_dict return_dict = return_dict if return_dict is not None else self.config.use_return_dict
outputs = self.beit( outputs = self.beit(
...@@ -1157,15 +1156,15 @@ class BeitForSemanticSegmentation(BeitPreTrainedModel): ...@@ -1157,15 +1156,15 @@ class BeitForSemanticSegmentation(BeitPreTrainedModel):
return_dict=None, return_dict=None,
): ):
r""" r"""
labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, height, width)`, `optional`): labels (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*):
Ground truth semantic segmentation maps for computing the loss. Indices should be in :obj:`[0, ..., Ground truth semantic segmentation maps for computing the loss. Indices should be in `[0, ..., config.num_labels - 1]`. If `config.num_labels > 1`, a classification loss is computed
config.num_labels - 1]`. If :obj:`config.num_labels > 1`, a classification loss is computed
(Cross-Entropy). (Cross-Entropy).
Returns: Returns:
Examples:: Examples:
```python
>>> from transformers import BeitFeatureExtractor, BeitForSemanticSegmentation >>> from transformers import BeitFeatureExtractor, BeitForSemanticSegmentation
>>> from PIL import Image >>> from PIL import Image
>>> import requests >>> import requests
...@@ -1180,7 +1179,7 @@ class BeitForSemanticSegmentation(BeitPreTrainedModel): ...@@ -1180,7 +1179,7 @@ class BeitForSemanticSegmentation(BeitPreTrainedModel):
>>> outputs = model(**inputs) >>> outputs = model(**inputs)
>>> # logits are of shape (batch_size, num_labels, height/4, width/4) >>> # logits are of shape (batch_size, num_labels, height/4, width/4)
>>> logits = outputs.logits >>> logits = outputs.logits
""" ```"""
return_dict = return_dict if return_dict is not None else self.config.use_return_dict return_dict = return_dict if return_dict is not None else self.config.use_return_dict
output_hidden_states = ( output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
......
...@@ -44,76 +44,74 @@ from .configuration_beit import BeitConfig ...@@ -44,76 +44,74 @@ from .configuration_beit import BeitConfig
@flax.struct.dataclass @flax.struct.dataclass
class FlaxBeitModelOutputWithPooling(FlaxBaseModelOutputWithPooling): class FlaxBeitModelOutputWithPooling(FlaxBaseModelOutputWithPooling):
""" """
Class for outputs of :class:`~transformers.FlaxBeitModel`. Class for outputs of [`FlaxBeitModel`].
Args: Args:
last_hidden_state (:obj:`jnp.ndarray` of shape :obj:`(batch_size, sequence_length, hidden_size)`): last_hidden_state (`jnp.ndarray` of shape `(batch_size, sequence_length, hidden_size)`):
Sequence of hidden-states at the output of the last layer of the model. Sequence of hidden-states at the output of the last layer of the model.
pooler_output (:obj:`jnp.ndarray` of shape :obj:`(batch_size, hidden_size)`): pooler_output (`jnp.ndarray` of shape `(batch_size, hidden_size)`):
Average of the last layer hidden states of the patch tokens (excluding the `[CLS]` token) if Average of the last layer hidden states of the patch tokens (excluding the *[CLS]* token) if
`config.use_mean_pooling` is set to True. If set to False, then the final hidden state of the `[CLS]` token *config.use_mean_pooling* is set to True. If set to False, then the final hidden state of the *[CLS]* token
will be returned. will be returned.
hidden_states (:obj:`tuple(jnp.ndarray)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``): hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
Tuple of :obj:`jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of
shape :obj:`(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each
layer plus the initial embedding outputs. layer plus the initial embedding outputs.
attentions (:obj:`tuple(jnp.ndarray)`, `optional`, returned when ``output_attentions=True`` is passed or when ``config.output_attentions=True``): attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
Tuple of :obj:`jnp.ndarray` (one for each layer) of shape :obj:`(batch_size, num_heads, sequence_length, Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
the self-attention heads. the self-attention heads.
""" """
BEIT_START_DOCSTRING = r""" BEIT_START_DOCSTRING = r"""
This model inherits from :class:`~transformers.FlaxPreTrainedModel`. Check the superclass documentation for the This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the
generic methods the library implements for all its model (such as downloading, saving and converting weights from generic methods the library implements for all its model (such as downloading, saving and converting weights from
PyTorch models) PyTorch models)
This model is also a Flax Linen `flax.linen.Module This model is also a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module) subclass. Use it as a regular Flax linen Module
<https://flax.readthedocs.io/en/latest/flax.linen.html#module>`__ subclass. Use it as a regular Flax linen Module
and refer to the Flax documentation for all matter related to general usage and behavior. and refer to the Flax documentation for all matter related to general usage and behavior.
Finally, this model supports inherent JAX features such as: Finally, this model supports inherent JAX features such as:
- `Just-In-Time (JIT) compilation <https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit>`__ - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)
- `Automatic Differentiation <https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation>`__ - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)
- `Vectorization <https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap>`__ - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)
- `Parallelization <https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap>`__ - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)
Parameters: Parameters:
config (:class:`~transformers.BeitConfig`): Model configuration class with all the parameters of the model. config ([`BeitConfig`]): Model configuration class with all the parameters of the model.
Initializing with a config file does not load the weights associated with the model, only the Initializing with a config file does not load the weights associated with the model, only the
configuration. Check out the :meth:`~transformers.FlaxPreTrainedModel.from_pretrained` method to load the configuration. Check out the [`~FlaxPreTrainedModel.from_pretrained`] method to load the
model weights. model weights.
dtype (:obj:`jax.numpy.dtype`, `optional`, defaults to :obj:`jax.numpy.float32`): dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):
The data type of the computation. Can be one of :obj:`jax.numpy.float32`, :obj:`jax.numpy.float16` (on The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on
GPUs) and :obj:`jax.numpy.bfloat16` (on TPUs). GPUs) and `jax.numpy.bfloat16` (on TPUs).
This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If
specified all the computation will be performed with the given ``dtype``. specified all the computation will be performed with the given `dtype`.
**Note that this only specifies the dtype of the computation and does not influence the dtype of model **Note that this only specifies the dtype of the computation and does not influence the dtype of model
parameters.** parameters.**
If you wish to change the dtype of the model parameters, see If you wish to change the dtype of the model parameters, see
:meth:`~transformers.FlaxPreTrainedModel.to_fp16` and :meth:`~transformers.FlaxPreTrainedModel.to_bf16`. [`~FlaxPreTrainedModel.to_fp16`] and [`~FlaxPreTrainedModel.to_bf16`].
""" """
BEIT_INPUTS_DOCSTRING = r""" BEIT_INPUTS_DOCSTRING = r"""
Args: Args:
pixel_values (:obj:`numpy.ndarray` of shape :obj:`(batch_size, num_channels, height, width)`): pixel_values (`numpy.ndarray` of shape `(batch_size, num_channels, height, width)`):
Pixel values. Pixel values can be obtained using :class:`~transformers.BeitFeatureExtractor`. See Pixel values. Pixel values can be obtained using [`BeitFeatureExtractor`]. See
:meth:`transformers.BeitFeatureExtractor.__call__` for details. [`BeitFeatureExtractor.__call__`] for details.
output_attentions (:obj:`bool`, `optional`): output_attentions (`bool`, *optional*):
Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
tensors for more detail. tensors for more detail.
output_hidden_states (:obj:`bool`, `optional`): output_hidden_states (`bool`, *optional*):
Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
more detail. more detail.
return_dict (:obj:`bool`, `optional`): return_dict (`bool`, *optional*):
Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple. Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
""" """
...@@ -810,13 +808,14 @@ class FlaxBeitForMaskedImageModeling(FlaxBeitPreTrainedModel): ...@@ -810,13 +808,14 @@ class FlaxBeitForMaskedImageModeling(FlaxBeitPreTrainedModel):
FLAX_BEIT_MLM_DOCSTRING = """ FLAX_BEIT_MLM_DOCSTRING = """
bool_masked_pos (:obj:`numpy.ndarray` of shape :obj:`(batch_size, num_patches)`): bool_masked_pos (`numpy.ndarray` of shape `(batch_size, num_patches)`):
Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
Returns: Returns:
Examples:: Examples:
```python
>>> from transformers import BeitFeatureExtractor, BeitForMaskedImageModeling >>> from transformers import BeitFeatureExtractor, BeitForMaskedImageModeling
>>> from PIL import Image >>> from PIL import Image
>>> import requests >>> import requests
...@@ -830,6 +829,7 @@ FLAX_BEIT_MLM_DOCSTRING = """ ...@@ -830,6 +829,7 @@ FLAX_BEIT_MLM_DOCSTRING = """
>>> inputs = feature_extractor(images=image, return_tensors="np") >>> inputs = feature_extractor(images=image, return_tensors="np")
>>> outputs = model(**inputs) >>> outputs = model(**inputs)
>>> logits = outputs.logits >>> logits = outputs.logits
```
""" """
overwrite_call_docstring(FlaxBeitForMaskedImageModeling, FLAX_BEIT_MLM_DOCSTRING) overwrite_call_docstring(FlaxBeitForMaskedImageModeling, FLAX_BEIT_MLM_DOCSTRING)
......
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