FlaxBartModel

 of jnp.ndarray (one for each layer) of shape (batch_size, num_heads, sequence_length, sequence_length).
Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the
self-attention heads.
  The FlaxBartPreTrainedModel forward method, overrides the __call__ special method. Although the recipe for forward pass needs to be defined within this function, one should call the Module
instance afterwards instead of this since the former takes care of running the pre and post processing steps while
the latter silently ignores them.  Example:   Copied >>> from transformers import AutoTokenizer, FlaxBartModel
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
>>> model = FlaxBartModel.from_pretrained("facebook/bart-base")
>>> inputs = tokenizer("Hello, my dog is cute", return_tensors="jax")
>>> outputs = model(**inputs)
>>> last_hidden_states = outputs.last_hidden_state  encode  < source > ( input_ids: Array a


FlaxBartModel

ttention_mask: typing.Optional[jax.Array] = None position_ids: typing.Optional[jax.Array] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None train: bool = False params: dict = None dropout_rng: PRNGKey = None  ) → transformers.modeling_flax_outputs.FlaxBaseModelOutput or tuple(torch.FloatTensor)  Parameters   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
it.
Indices can be obtained using AutoTokenizer. See PreTrainedTokenizer.encode() and
PreTrainedTokenizer.call() for details.
What are input IDs?   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]:
1 for tokens that are not masked,
0 for tokens that are masked.
What are at


FlaxBartModel

tention masks?   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, config.max_position_embeddings - 1].   output_attentions (bool, optional) —
Whether or not to return the attentions tensors of all attention layers. See attentions under returned
tensors for more detail.   output_hidden_states (bool, optional) —
Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for
more detail.   return_dict (bool, optional) —
Whether or not to return a ModelOutput instead of a plain tuple.   Returns
transformers.modeling_flax_outputs.FlaxBaseModelOutput or tuple(torch.FloatTensor)

A transformers.modeling_flax_outputs.FlaxBaseModelOutput or a tuple of
torch.FloatTensor (if return_dict=False is passed or when config.return_dict=False) comprising various
elements depending on the configuration (<


FlaxBartModel

class 'transformers.models.bart.configuration_bart.BartConfig'>) and inputs.
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.
hidden_states (tuple(jnp.ndarray), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple of jnp.ndarray (one for the output of the embeddings + one for the output of each layer) of shape
(batch_size, sequence_length, hidden_size).
Hidden-states of the model at the output of each layer plus the initial embedding outputs.
attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — 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 the self-attention
heads.
   Examp


FlaxBartModel

le:   Copied >>> from transformers import AutoTokenizer, FlaxBartForConditionalGeneration
>>> model = FlaxBartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
>>> text = "My friends are cool but they eat too many carbs."
>>> inputs = tokenizer(text, max_length=1024, return_tensors="jax")
>>> encoder_outputs = model.encode(**inputs)  decode  < source > ( decoder_input_ids encoder_outputs encoder_attention_mask: typing.Optional[jax.Array] = None decoder_attention_mask: typing.Optional[jax.Array] = None decoder_position_ids: typing.Optional[jax.Array] = None past_key_values: dict = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None train: bool = False params: dict = None dropout_rng: PRNGKey = None  ) → transformers.modeling_flax_outputs.FlaxBaseModelOutputWithPastAndCrossAttentio


FlaxBartModel

ns or tuple(torch.FloatTensor)  Parameters   decoder_input_ids (jnp.ndarray of shape (batch_size, target_sequence_length)) —
Indices of decoder input sequence tokens in the vocabulary.
Indices can be obtained using AutoTokenizer. See PreTrainedTokenizer.encode() and
PreTrainedTokenizer.call() for details.
What are decoder input IDs?
For translation and summarization training, decoder_input_ids should be provided. If no
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.   encoder_outputs (tuple(tuple(jnp.ndarray)) —
Tuple consists of (last_hidden_state, optional: hidden_states, 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 encoder. Used in the cross-attention of the decoder.   encoder_attention_mask (jnp.ndarray of shape (batch_size, sequ


FlaxBartModel

ence_length), optional) —
Mask to avoid performing attention on padding token indices. Mask values selected in [0, 1]:
1 for tokens that are not masked,
0 for tokens that are masked.
What are attention masks?   decoder_attention_mask (jnp.ndarray of shape (batch_size, target_sequence_length), optional) —
Default behavior: generate a tensor that ignores pad tokens in decoder_input_ids. Causal mask will also
be used by default.
If you want to change padding behavior, you should modify to your needs. See diagram 1 in the
paper for more information on the default strategy.   decoder_position_ids (numpy.ndarray of shape (batch_size, sequence_length), optional) —
Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the
range [0, config.max_position_embeddings - 1].   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


FlaxBartModel

 (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].   output_attentions (bool, optional) —
Whether or not to return the attentions tensors of all attention layers. See attentions under returned
tensors for more detail.   output_hidden_states (bool, optional) —
Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for
more detail.   return_dict (bool, optional) —
Whether or not to return a ModelOutput instead of a plain tuple.   Returns
transformers.modeling_flax_outputs.FlaxBaseModelOutputWithPastAndCrossAttentions or tuple(torch.FloatTensor)

A transformers.modeling_flax_outputs.FlaxBaseModelOutputWithPastAndCrossAttentions or a tuple of
torch.FloatTensor (if return_dict=False is passed or when config.return_dict=False) comprising various
elements depending on the configuration (<class 'transf


FlaxBartModel

ormers.models.bart.configuration_bart.BartConfig'>) and inputs.
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.
If past_key_values is used only the last hidden-state of the sequences of shape (batch_size, 1, hidden_size) is output.
past_key_values (tuple(tuple(jnp.ndarray)), optional, returned when use_cache=True is passed or when config.use_cache=True) — Tuple of tuple(jnp.ndarray) of length config.n_layers, with each tuple having 2 tensors of shape
(batch_size, num_heads, sequence_length, embed_size_per_head)) and optionally if
config.is_encoder_decoder=True 2 additional tensors of shape (batch_size, num_heads, encoder_sequence_length, embed_size_per_head).
Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if
config.is_encoder_decoder=True in the cross-attention blocks) that can be used (see past_key_value


FlaxBartModel

s
input) to speed up sequential decoding.
hidden_states (tuple(jnp.ndarray), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple of jnp.ndarray (one for the output of the embeddings + one for the output of each layer) of shape
(batch_size, sequence_length, hidden_size).
Hidden-states of the model at the output of each layer plus the initial embedding outputs.
attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — 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 the self-attention
heads.
cross_attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True and config.add_cross_attention=True is passed or when config.output_attentions=True) — Tuple of jnp.ndarray (one for


FlaxBartModel

 each layer) of shape (batch_size, num_heads, sequence_length, sequence_length).
Attentions weights of the decoder’s cross-attention layer, after the attention softmax, used to compute the
weighted average in the cross-attention heads.
   Example:   Copied >>> import jax.numpy as jnp
>>> from transformers import AutoTokenizer, FlaxBartForConditionalGeneration
>>> model = FlaxBartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
>>> text = "My friends are cool but they eat too many carbs."
>>> inputs = tokenizer(text, max_length=1024, return_tensors="jax")
>>> encoder_outputs = model.encode(**inputs)
>>> decoder_start_token_id = model.config.decoder_start_token_id
>>> decoder_input_ids = jnp.ones((inputs.input_ids.shape[0], 1), dtype="i4") * decoder_start_token_id
>>> outputs = model.decode(decoder_input_ids, encoder_outputs)
>>> last_decoder_hidden_states = outputs.last_h


FlaxBartForConditionalGeneration

class transformers.FlaxBartForConditionalGeneration  < source > ( config: BartConfig input_shape: typing.Tuple[int] = (1, 1) seed: int = 0 dtype: dtype = <class 'jax.numpy.float32'> _do_init: bool = True **kwargs  )   Parameters   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
configuration. Check out the from_pretrained() method to load the model weights.   dtype (jax.numpy.dtype, optional, defaults to jax.numpy.float32) —
The data type of the computation. Can be one of jax.numpy.float32, jax.numpy.float16 (on 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
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
parameters.
If you wish to change the d


FlaxBartForConditionalGeneration

type of the model parameters, see to_fp16() and
to_bf16().    The BART Model with a language modeling head. Can be used for summarization.
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 embeddings, pruning heads
etc.) This model is also a Flax Linen
flax.nn.Module subclass. Use it as a
regular Flax 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: Just-In-Time (JIT) compilation Automatic Differentiation Vectorization Parallelization  __call__  < source > ( input_ids: Array attention_mask: typing.Optional[jax.Array] = None decoder_input_ids: typing.Optional[jax.Array] = None decoder_attention_mask: typing.Optional[jax.Array] = None position_ids: typing.Optional[jax.Array] = None decoder_position_ids: typing.Optional[jax.Array] = None


FlaxBartForConditionalGeneration

output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None train: bool = False params: dict = None dropout_rng: PRNGKey = None  ) → transformers.modeling_flax_outputs.FlaxSeq2SeqLMOutput or tuple(torch.FloatTensor)  Parameters   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
it.
Indices can be obtained using AutoTokenizer. See PreTrainedTokenizer.encode() and
PreTrainedTokenizer.call() for details.
What are input IDs?   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]:
1 for tokens that are not masked,
0 for tokens that are masked.
What are attention masks?   decoder_input_ids (jnp.ndarray of shape (batch_size, target_sequence_length), optional) —
Indices of deco


FlaxBartForConditionalGeneration

der input sequence tokens in the vocabulary.
Indices can be obtained using AutoTokenizer. See PreTrainedTokenizer.encode() and
PreTrainedTokenizer.call() for details.
What are decoder input IDs?
For translation and summarization training, decoder_input_ids should be provided. If no
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.   decoder_attention_mask (jnp.ndarray of shape (batch_size, target_sequence_length), optional) —
Default behavior: generate a tensor that ignores pad tokens in decoder_input_ids. Causal mask will also
be used by default.
If you want to change padding behavior, you should modify to your needs. See diagram 1 in the
paper for more information on the default strategy.   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, config.max_


FlaxBartForConditionalGeneration

position_embeddings - 1].   decoder_position_ids (numpy.ndarray of shape (batch_size, sequence_length), optional) —
Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the
range [0, config.max_position_embeddings - 1].   output_attentions (bool, optional) —
Whether or not to return the attentions tensors of all attention layers. See attentions under returned
tensors for more detail.   output_hidden_states (bool, optional) —
Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for
more detail.   return_dict (bool, optional) —
Whether or not to return a ModelOutput instead of a plain tuple.   Returns
transformers.modeling_flax_outputs.FlaxSeq2SeqLMOutput or tuple(torch.FloatTensor)

A transformers.modeling_flax_outputs.FlaxSeq2SeqLMOutput or a tuple of
torch.FloatTensor (if return_dict=False is passed or when config.return_dict=False) comprising various
elements depending on the configuration


FlaxBartForConditionalGeneration

(BartConfig) and inputs.
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).
past_key_values (tuple(tuple(jnp.ndarray)), optional, returned when use_cache=True is passed or when config.use_cache=True) — Tuple of tuple(jnp.ndarray) of length config.n_layers, with each tuple having 2 tensors of shape
(batch_size, num_heads, sequence_length, embed_size_per_head)) and 2 additional tensors of shape
(batch_size, num_heads, encoder_sequence_length, embed_size_per_head).
Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
blocks) that can be used (see past_key_values input) to speed up sequential decoding.
decoder_hidden_states (tuple(jnp.ndarray), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple of jnp.ndarray (one for the output of the embeddings + on


FlaxBartForConditionalGeneration

e for the output of each layer) of shape
(batch_size, sequence_length, hidden_size).
Hidden-states of the decoder at the output of each layer plus the initial embedding outputs.
decoder_attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — Tuple of jnp.ndarray (one for each layer) of shape (batch_size, num_heads, sequence_length, sequence_length).
Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the
self-attention heads.
cross_attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — Tuple of jnp.ndarray (one for each layer) of shape (batch_size, num_heads, sequence_length, sequence_length).
Attentions weights of the decoder’s cross-attention layer, after the attention softmax, used to compute the
weighted average in the cross-attention heads.
encoder_last_hidden_state (jnp.ndarray of s


FlaxBartForConditionalGeneration

hape (batch_size, sequence_length, hidden_size), optional) — Sequence of hidden-states at the output of the last layer of the encoder of the model.
encoder_hidden_states (tuple(jnp.ndarray), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple of jnp.ndarray (one for the output of the embeddings + one for the output of each layer) of shape
(batch_size, sequence_length, hidden_size).
Hidden-states of the encoder at the output of each layer plus the initial embedding outputs.
encoder_attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — Tuple of jnp.ndarray (one for each layer) of shape (batch_size, num_heads, sequence_length, sequence_length).
Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the
self-attention heads.
  The FlaxBartPreTrainedModel forward method, overrides the __call__ special method. A


FlaxBartForConditionalGeneration

lthough the recipe for forward pass needs to be defined within this function, one should call the Module
instance afterwards instead of this since the former takes care of running the pre and post processing steps while
the latter silently ignores them.  Summarization example:   Copied >>> from transformers import AutoTokenizer, FlaxBartForConditionalGeneration
>>> model = FlaxBartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
>>> ARTICLE_TO_SUMMARIZE = "My friends are cool but they eat too many carbs."
>>> inputs = tokenizer([ARTICLE_TO_SUMMARIZE], max_length=1024, return_tensors="np")
>>> # Generate Summary
>>> summary_ids = model.generate(inputs["input_ids"]).sequences
>>> print(tokenizer.batch_decode(summary_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False))  Mask filling example:   Copied >>> import jax
>>> from transformers import AutoTokenizer, FlaxBartForConditionalGen


FlaxBartForConditionalGeneration

eration
>>> model = FlaxBartForConditionalGeneration.from_pretrained("facebook/bart-large")
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large")
>>> TXT = "My friends are <mask> but they eat too many carbs."
>>> input_ids = tokenizer([TXT], return_tensors="jax")["input_ids"]
>>> logits = model(input_ids).logits
>>> masked_index = (input_ids[0] == tokenizer.mask_token_id).nonzero()[0].item()
>>> probs = jax.nn.softmax(logits[0, masked_index], axis=0)
>>> values, predictions = jax.lax.top_k(probs, k=1)
>>> tokenizer.decode(predictions).split()  encode  < source > ( input_ids: Array attention_mask: typing.Optional[jax.Array] = None position_ids: typing.Optional[jax.Array] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None train: bool = False params: dict = None dropout_rng: PRNGKey = None  ) → transformers.modeling_flax_outputs.FlaxBaseModelOutput or tuple(torch.FloatTensor)  P


FlaxBartForConditionalGeneration

arameters   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
it.
Indices can be obtained using AutoTokenizer. See PreTrainedTokenizer.encode() and
PreTrainedTokenizer.call() for details.
What are input IDs?   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]:
1 for tokens that are not masked,
0 for tokens that are masked.
What are attention masks?   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, config.max_position_embeddings - 1].   output_attentions (bool, optional) —
Whether or not to return the attentions tensors of all attention layers. See attentions under returned
tensors for more detail.   output_hidden


FlaxBartForConditionalGeneration

_states (bool, optional) —
Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for
more detail.   return_dict (bool, optional) —
Whether or not to return a ModelOutput instead of a plain tuple.   Returns
transformers.modeling_flax_outputs.FlaxBaseModelOutput or tuple(torch.FloatTensor)

A transformers.modeling_flax_outputs.FlaxBaseModelOutput or a tuple of
torch.FloatTensor (if return_dict=False is passed or when config.return_dict=False) comprising various
elements depending on the configuration (<class 'transformers.models.bart.configuration_bart.BartConfig'>) and inputs.
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.
hidden_states (tuple(jnp.ndarray), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple of jnp.ndarray (one for the output of the embeddings + one for the


FlaxBartForConditionalGeneration

 output of each layer) of shape
(batch_size, sequence_length, hidden_size).
Hidden-states of the model at the output of each layer plus the initial embedding outputs.
attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — 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 the self-attention
heads.
   Example:   Copied >>> from transformers import AutoTokenizer, FlaxBartForConditionalGeneration
>>> model = FlaxBartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
>>> text = "My friends are cool but they eat too many carbs."
>>> inputs = tokenizer(text, max_length=1024, return_tensors="jax")
>>> encoder_outputs = model.encode(**inputs)  decode  < source > ( decoder_input_ids encod


FlaxBartForConditionalGeneration

er_outputs encoder_attention_mask: typing.Optional[jax.Array] = None decoder_attention_mask: typing.Optional[jax.Array] = None decoder_position_ids: typing.Optional[jax.Array] = None past_key_values: dict = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None train: bool = False params: dict = None dropout_rng: PRNGKey = None  ) → transformers.modeling_flax_outputs.FlaxCausalLMOutputWithCrossAttentions or tuple(torch.FloatTensor)  Parameters   decoder_input_ids (jnp.ndarray of shape (batch_size, target_sequence_length)) —
Indices of decoder input sequence tokens in the vocabulary.
Indices can be obtained using AutoTokenizer. See PreTrainedTokenizer.encode() and
PreTrainedTokenizer.call() for details.
What are decoder input IDs?
For translation and summarization training, decoder_input_ids should be provided. If no
decoder_input_ids is provided, the model will create this tensor by shifting the


FlaxBartForConditionalGeneration

input_ids to the right
for denoising pre-training following the paper.   encoder_outputs (tuple(tuple(jnp.ndarray)) —
Tuple consists of (last_hidden_state, optional: hidden_states, 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 encoder. Used in the cross-attention of the decoder.   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]:
1 for tokens that are not masked,
0 for tokens that are masked.
What are attention masks?   decoder_attention_mask (jnp.ndarray of shape (batch_size, target_sequence_length), optional) —
Default behavior: generate a tensor that ignores pad tokens in decoder_input_ids. Causal mask will also
be used by default.
If you want to change padding behavior, you should modify to your needs. See diagram 1 in the
paper


FlaxBartForConditionalGeneration

 for more information on the default strategy.   decoder_position_ids (numpy.ndarray of shape (batch_size, sequence_length), optional) —
Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the
range [0, config.max_position_embeddings - 1].   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
auto-regressive decoding. Pre-computed key and value hidden-states are of shape [batch_size, max_length].   output_attentions (bool, optional) —
Whether or not to return the attentions tensors of all attention layers. See attentions under returned
tensors for more detail.   output_hidden_states (bool, optional) —
Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for
more detail.   return_dict (bool, optional) —
Whether or not to retu


FlaxBartForConditionalGeneration

rn a ModelOutput instead of a plain tuple.   Returns
transformers.modeling_flax_outputs.FlaxCausalLMOutputWithCrossAttentions or tuple(torch.FloatTensor)

A transformers.modeling_flax_outputs.FlaxCausalLMOutputWithCrossAttentions or a tuple of
torch.FloatTensor (if return_dict=False is passed or when config.return_dict=False) comprising various
elements depending on the configuration (<class 'transformers.models.bart.configuration_bart.BartConfig'>) and inputs.
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).
hidden_states (tuple(jnp.ndarray), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple of jnp.ndarray (one for the output of the embeddings + one for the output of each layer) of shape
(batch_size, sequence_length, hidden_size).
Hidden-states of the model at the output of each layer plus the


FlaxBartForConditionalGeneration

initial embedding outputs.
attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — 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 the self-attention
heads.
cross_attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — Tuple of jnp.ndarray (one for each layer) of shape (batch_size, num_heads, sequence_length, sequence_length).
Cross attentions weights after the attention softmax, used to compute the weighted average in the
cross-attention heads.
past_key_values (tuple(tuple(jnp.ndarray)), optional, returned when use_cache=True is passed or when config.use_cache=True) — Tuple of jnp.ndarray tuples of length config.n_layers, with each tuple containing the cached key, value
states of the self-att


FlaxBartForConditionalGeneration

ention and the cross-attention layers if model is used in encoder-decoder setting.
Only relevant if config.is_decoder = True.
Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see
past_key_values input) to speed up sequential decoding.
   Example:   Copied >>> import jax.numpy as jnp
>>> from transformers import AutoTokenizer, FlaxBartForConditionalGeneration
>>> model = FlaxBartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
>>> text = "My friends are cool but they eat too many carbs."
>>> inputs = tokenizer(text, max_length=1024, return_tensors="jax")
>>> encoder_outputs = model.encode(**inputs)
>>> decoder_start_token_id = model.config.decoder_start_token_id
>>> decoder_input_ids = jnp.ones((inputs.input_ids.shape[0], 1), dtype="i4") * decoder_start_token_id
>>> outputs = model.decode(decoder_input_ids, encoder_outputs)
>>> logits = outputs.logits


FlaxBartForSequenceClassification

class transformers.FlaxBartForSequenceClassification  < source > ( config: BartConfig input_shape: typing.Tuple[int] = (1, 1) seed: int = 0 dtype: dtype = <class 'jax.numpy.float32'> _do_init: bool = True **kwargs  )   Parameters   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
configuration. Check out the from_pretrained() method to load the model weights.   dtype (jax.numpy.dtype, optional, defaults to jax.numpy.float32) —
The data type of the computation. Can be one of jax.numpy.float32, jax.numpy.float16 (on 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
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
parameters.


FlaxBartForSequenceClassification

If you wish to change the dtype of the model parameters, see to_fp16() and
to_bf16().    Bart model with a sequence classification/head on top (a linear layer on top of the pooled output) e.g. for GLUE
tasks. 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 embeddings, pruning heads
etc.) This model is also a Flax Linen
flax.nn.Module subclass. Use it as a
regular Flax 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: Just-In-Time (JIT) compilation Automatic Differentiation Vectorization Parallelization  __call__  < source > ( input_ids: Array attention_mask: typing.Optional[jax.Array] = None decoder_input_ids: typing.Optional[jax.Array] = None decoder_attention_mask: typing.Optional[jax.Array] = None positi


FlaxBartForSequenceClassification

on_ids: typing.Optional[jax.Array] = None decoder_position_ids: typing.Optional[jax.Array] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None train: bool = False params: dict = None dropout_rng: PRNGKey = None  ) → transformers.modeling_flax_outputs.FlaxSeq2SeqSequenceClassifierOutput or tuple(torch.FloatTensor)  Parameters   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
it.
Indices can be obtained using AutoTokenizer. See PreTrainedTokenizer.encode() and
PreTrainedTokenizer.call() for details.
What are input IDs?   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]:
1 for tokens that are not masked,
0 for tokens that are


FlaxBartForSequenceClassification

masked.
What are attention masks?   decoder_input_ids (jnp.ndarray of shape (batch_size, target_sequence_length), optional) —
Indices of decoder input sequence tokens in the vocabulary.
Indices can be obtained using AutoTokenizer. See PreTrainedTokenizer.encode() and
PreTrainedTokenizer.call() for details.
What are decoder input IDs?
For translation and summarization training, decoder_input_ids should be provided. If no
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.   decoder_attention_mask (jnp.ndarray of shape (batch_size, target_sequence_length), optional) —
Default behavior: generate a tensor that ignores pad tokens in decoder_input_ids. Causal mask will also
be used by default.
If you want to change padding behavior, you should modify to your needs. See diagram 1 in the
paper for more information on the default strategy.   position_ids (numpy.ndarray


FlaxBartForSequenceClassification

of shape (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].   decoder_position_ids (numpy.ndarray of shape (batch_size, sequence_length), optional) —
Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the
range [0, config.max_position_embeddings - 1].   output_attentions (bool, optional) —
Whether or not to return the attentions tensors of all attention layers. See attentions under returned
tensors for more detail.   output_hidden_states (bool, optional) —
Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for
more detail.   return_dict (bool, optional) —
Whether or not to return a ModelOutput instead of a plain tuple.   Returns
transformers.modeling_flax_outputs.FlaxSeq2SeqSequenceClassifierOutput or tuple(torch.FloatTensor)

A t


FlaxBartForSequenceClassification

ransformers.modeling_flax_outputs.FlaxSeq2SeqSequenceClassifierOutput or a tuple of
torch.FloatTensor (if return_dict=False is passed or when config.return_dict=False) comprising various
elements depending on the configuration (BartConfig) and inputs.
logits (jnp.ndarray of shape (batch_size, config.num_labels)) — Classification (or regression if config.num_labels==1) scores (before SoftMax).
past_key_values (tuple(tuple(jnp.ndarray)), optional, returned when use_cache=True is passed or when config.use_cache=True) — Tuple of tuple(jnp.ndarray) of length config.n_layers, with each tuple having 2 tensors of shape
(batch_size, num_heads, sequence_length, embed_size_per_head)) and 2 additional tensors of shape
(batch_size, num_heads, encoder_sequence_length, embed_size_per_head).
Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
blocks) that can be used (see past_key_values input) to speed up sequential de


FlaxBartForSequenceClassification

coding.
decoder_hidden_states (tuple(jnp.ndarray), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple of jnp.ndarray (one for the output of the embeddings + one for the output of each layer) of shape
(batch_size, sequence_length, hidden_size).
Hidden-states of the decoder at the output of each layer plus the initial embedding outputs.
decoder_attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — Tuple of jnp.ndarray (one for each layer) of shape (batch_size, num_heads, sequence_length, sequence_length).
Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the
self-attention heads.
cross_attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — Tuple of jnp.ndarray (one for each layer) of shape (batch_size,


FlaxBartForSequenceClassification

 num_heads, sequence_length, sequence_length).
Attentions weights of the decoder’s cross-attention layer, after the attention softmax, used to compute the
weighted average in the cross-attention heads.
encoder_last_hidden_state (jnp.ndarray of shape (batch_size, sequence_length, hidden_size), optional) — Sequence of hidden-states at the output of the last layer of the encoder of the model.
encoder_hidden_states (tuple(jnp.ndarray), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple of jnp.ndarray (one for the output of the embeddings + one for the output of each layer) of shape
(batch_size, sequence_length, hidden_size).
Hidden-states of the encoder at the output of each layer plus the initial embedding outputs.
encoder_attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — Tuple of jnp.ndarray (one for each layer) of shape (ba


FlaxBartForSequenceClassification

tch_size, num_heads, sequence_length, sequence_length).
Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the
self-attention heads.
  The FlaxBartPreTrainedModel forward method, overrides the __call__ special method. Although the recipe for forward pass needs to be defined within this function, one should call the Module
instance afterwards instead of this since the former takes care of running the pre and post processing steps while
the latter silently ignores them.  Example:   Copied >>> from transformers import AutoTokenizer, FlaxBartForSequenceClassification
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
>>> model = FlaxBartForSequenceClassification.from_pretrained("facebook/bart-base")
>>> inputs = tokenizer("Hello, my dog is cute", return_tensors="jax")
>>> outputs = model(**inputs)
>>> logits = outputs.logits  encode  < source > ( input_ids: Array attention_mask: typing.Optional


FlaxBartForSequenceClassification

[jax.Array] = None position_ids: typing.Optional[jax.Array] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None train: bool = False params: dict = None dropout_rng: PRNGKey = None  ) → transformers.modeling_flax_outputs.FlaxBaseModelOutput or tuple(torch.FloatTensor)  Parameters   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
it.
Indices can be obtained using AutoTokenizer. See PreTrainedTokenizer.encode() and
PreTrainedTokenizer.call() for details.
What are input IDs?   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]:
1 for tokens that are not masked,
0 for tokens that are masked.
What are attention masks?   position_id


FlaxBartForSequenceClassification

s (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, config.max_position_embeddings - 1].   output_attentions (bool, optional) —
Whether or not to return the attentions tensors of all attention layers. See attentions under returned
tensors for more detail.   output_hidden_states (bool, optional) —
Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for
more detail.   return_dict (bool, optional) —
Whether or not to return a ModelOutput instead of a plain tuple.   Returns
transformers.modeling_flax_outputs.FlaxBaseModelOutput or tuple(torch.FloatTensor)

A transformers.modeling_flax_outputs.FlaxBaseModelOutput or a tuple of
torch.FloatTensor (if return_dict=False is passed or when config.return_dict=False) comprising various
elements depending on the configuration (<class 'transformers.models


FlaxBartForSequenceClassification

.bart.configuration_bart.BartConfig'>) and inputs.
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.
hidden_states (tuple(jnp.ndarray), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple of jnp.ndarray (one for the output of the embeddings + one for the output of each layer) of shape
(batch_size, sequence_length, hidden_size).
Hidden-states of the model at the output of each layer plus the initial embedding outputs.
attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — 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 the self-attention
heads.
   Example:   Copied >>> from tr


FlaxBartForSequenceClassification

ansformers import AutoTokenizer, FlaxBartForConditionalGeneration
>>> model = FlaxBartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
>>> text = "My friends are cool but they eat too many carbs."
>>> inputs = tokenizer(text, max_length=1024, return_tensors="jax")
>>> encoder_outputs = model.encode(**inputs)  decode  < source > ( decoder_input_ids encoder_outputs encoder_attention_mask: typing.Optional[jax.Array] = None decoder_attention_mask: typing.Optional[jax.Array] = None decoder_position_ids: typing.Optional[jax.Array] = None past_key_values: dict = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None train: bool = False params: dict = None dropout_rng: PRNGKey = None  ) → transformers.modeling_flax_outputs.FlaxBaseModelOutputWithPastAndCrossAttentions or tuple(torch.Floa


FlaxBartForSequenceClassification

tTensor)  Parameters   decoder_input_ids (jnp.ndarray of shape (batch_size, target_sequence_length)) —
Indices of decoder input sequence tokens in the vocabulary.
Indices can be obtained using AutoTokenizer. See PreTrainedTokenizer.encode() and
PreTrainedTokenizer.call() for details.
What are decoder input IDs?
For translation and summarization training, decoder_input_ids should be provided. If no
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.   encoder_outputs (tuple(tuple(jnp.ndarray)) —
Tuple consists of (last_hidden_state, optional: hidden_states, 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 encoder. Used in the cross-attention of the decoder.   encoder_attention_mask (jnp.ndarray of shape (batch_size, sequence_length), option


FlaxBartForSequenceClassification

al) —
Mask to avoid performing attention on padding token indices. Mask values selected in [0, 1]:
1 for tokens that are not masked,
0 for tokens that are masked.
What are attention masks?   decoder_attention_mask (jnp.ndarray of shape (batch_size, target_sequence_length), optional) —
Default behavior: generate a tensor that ignores pad tokens in decoder_input_ids. Causal mask will also
be used by default.
If you want to change padding behavior, you should modify to your needs. See diagram 1 in the
paper for more information on the default strategy.   decoder_position_ids (numpy.ndarray of shape (batch_size, sequence_length), optional) —
Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the
range [0, config.max_position_embeddings - 1].   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 i


FlaxBartForSequenceClassification

n 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].   output_attentions (bool, optional) —
Whether or not to return the attentions tensors of all attention layers. See attentions under returned
tensors for more detail.   output_hidden_states (bool, optional) —
Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for
more detail.   return_dict (bool, optional) —
Whether or not to return a ModelOutput instead of a plain tuple.   Returns
transformers.modeling_flax_outputs.FlaxBaseModelOutputWithPastAndCrossAttentions or tuple(torch.FloatTensor)

A transformers.modeling_flax_outputs.FlaxBaseModelOutputWithPastAndCrossAttentions or a tuple of
torch.FloatTensor (if return_dict=False is passed or when config.return_dict=False) comprising various
elements depending on the configuration (<class 'transformers.models.ba


FlaxBartForSequenceClassification

rt.configuration_bart.BartConfig'>) and inputs.
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.
If past_key_values is used only the last hidden-state of the sequences of shape (batch_size, 1, hidden_size) is output.
past_key_values (tuple(tuple(jnp.ndarray)), optional, returned when use_cache=True is passed or when config.use_cache=True) — Tuple of tuple(jnp.ndarray) of length config.n_layers, with each tuple having 2 tensors of shape
(batch_size, num_heads, sequence_length, embed_size_per_head)) and optionally if
config.is_encoder_decoder=True 2 additional tensors of shape (batch_size, num_heads, encoder_sequence_length, embed_size_per_head).
Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if
config.is_encoder_decoder=True in the cross-attention blocks) that can be used (see past_key_values
input) to sp


FlaxBartForSequenceClassification

eed up sequential decoding.
hidden_states (tuple(jnp.ndarray), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple of jnp.ndarray (one for the output of the embeddings + one for the output of each layer) of shape
(batch_size, sequence_length, hidden_size).
Hidden-states of the model at the output of each layer plus the initial embedding outputs.
attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — 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 the self-attention
heads.
cross_attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True and config.add_cross_attention=True is passed or when config.output_attentions=True) — Tuple of jnp.ndarray (one for each layer)


FlaxBartForSequenceClassification

 of shape (batch_size, num_heads, sequence_length, sequence_length).
Attentions weights of the decoder’s cross-attention layer, after the attention softmax, used to compute the
weighted average in the cross-attention heads.
   Example:   Copied >>> import jax.numpy as jnp
>>> from transformers import AutoTokenizer, FlaxBartForConditionalGeneration
>>> model = FlaxBartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
>>> text = "My friends are cool but they eat too many carbs."
>>> inputs = tokenizer(text, max_length=1024, return_tensors="jax")
>>> encoder_outputs = model.encode(**inputs)
>>> decoder_start_token_id = model.config.decoder_start_token_id
>>> decoder_input_ids = jnp.ones((inputs.input_ids.shape[0], 1), dtype="i4") * decoder_start_token_id
>>> outputs = model.decode(decoder_input_ids, encoder_outputs)
>>> last_decoder_hidden_states = outputs.last_hidden_stat


FlaxBartForQuestionAnswering

class transformers.FlaxBartForQuestionAnswering  < source > ( config: BartConfig input_shape: typing.Tuple[int] = (1, 1) seed: int = 0 dtype: dtype = <class 'jax.numpy.float32'> _do_init: bool = True **kwargs  )   Parameters   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
configuration. Check out the from_pretrained() method to load the model weights.   dtype (jax.numpy.dtype, optional, defaults to jax.numpy.float32) —
The data type of the computation. Can be one of jax.numpy.float32, jax.numpy.float16 (on 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
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
parameters.
If you wish to ch


FlaxBartForQuestionAnswering

ange the dtype of the model parameters, see to_fp16() and
to_bf16().    BART Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
layer on top of the hidden-states output to compute span start logits and span end logits). 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 embeddings, pruning heads
etc.) This model is also a Flax Linen
flax.nn.Module subclass. Use it as a
regular Flax 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: Just-In-Time (JIT) compilation Automatic Differentiation Vectorization Parallelization  __call__  < source > ( input_ids: Array attention_mask: typing.Optional[jax.Array] = None decoder_input_ids: typing.Optional[jax.Array] = None decoder_atte


FlaxBartForQuestionAnswering

ntion_mask: typing.Optional[jax.Array] = None position_ids: typing.Optional[jax.Array] = None decoder_position_ids: typing.Optional[jax.Array] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None train: bool = False params: dict = None dropout_rng: PRNGKey = None  ) → transformers.modeling_flax_outputs.FlaxSeq2SeqQuestionAnsweringModelOutput or tuple(torch.FloatTensor)  Parameters   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
it.
Indices can be obtained using AutoTokenizer. See PreTrainedTokenizer.encode() and
PreTrainedTokenizer.call() for details.
What are input IDs?   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]:
1 for tokens


FlaxBartForQuestionAnswering

that are not masked,
0 for tokens that are masked.
What are attention masks?   decoder_input_ids (jnp.ndarray of shape (batch_size, target_sequence_length), optional) —
Indices of decoder input sequence tokens in the vocabulary.
Indices can be obtained using AutoTokenizer. See PreTrainedTokenizer.encode() and
PreTrainedTokenizer.call() for details.
What are decoder input IDs?
For translation and summarization training, decoder_input_ids should be provided. If no
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.   decoder_attention_mask (jnp.ndarray of shape (batch_size, target_sequence_length), optional) —
Default behavior: generate a tensor that ignores pad tokens in decoder_input_ids. Causal mask will also
be used by default.
If you want to change padding behavior, you should modify to your needs. See diagram 1 in the
paper for more information on the default strategy.


FlaxBartForQuestionAnswering

  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, config.max_position_embeddings - 1].   decoder_position_ids (numpy.ndarray of shape (batch_size, sequence_length), optional) —
Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the
range [0, config.max_position_embeddings - 1].   output_attentions (bool, optional) —
Whether or not to return the attentions tensors of all attention layers. See attentions under returned
tensors for more detail.   output_hidden_states (bool, optional) —
Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for
more detail.   return_dict (bool, optional) —
Whether or not to return a ModelOutput instead of a plain tuple.   Returns
transformers.modeling_flax_outputs.FlaxSeq2SeqQuestionAnsweringModelOutput or tuple(torc


FlaxBartForQuestionAnswering

h.FloatTensor)

A transformers.modeling_flax_outputs.FlaxSeq2SeqQuestionAnsweringModelOutput or a tuple of
torch.FloatTensor (if return_dict=False is passed or when config.return_dict=False) comprising various
elements depending on the configuration (BartConfig) and inputs.
start_logits (jnp.ndarray of shape (batch_size, sequence_length)) — Span-start scores (before SoftMax).
end_logits (jnp.ndarray of shape (batch_size, sequence_length)) — Span-end scores (before SoftMax).
past_key_values (tuple(tuple(jnp.ndarray)), optional, returned when use_cache=True is passed or when config.use_cache=True) — Tuple of tuple(jnp.ndarray) of length config.n_layers, with each tuple having 2 tensors of shape
(batch_size, num_heads, sequence_length, embed_size_per_head)) and 2 additional tensors of shape
(batch_size, num_heads, encoder_sequence_length, embed_size_per_head).
Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
blocks


FlaxBartForQuestionAnswering

) that can be used (see past_key_values input) to speed up sequential decoding.
decoder_hidden_states (tuple(jnp.ndarray), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple of jnp.ndarray (one for the output of the embeddings + one for the output of each layer) of shape
(batch_size, sequence_length, hidden_size).
Hidden-states of the decoder at the output of each layer plus the initial embedding outputs.
decoder_attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — Tuple of jnp.ndarray (one for each layer) of shape (batch_size, num_heads, sequence_length, sequence_length).
Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the
self-attention heads.
cross_attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — Tupl


FlaxBartForQuestionAnswering

e of jnp.ndarray (one for each layer) of shape (batch_size, num_heads, sequence_length, sequence_length).
Attentions weights of the decoder’s cross-attention layer, after the attention softmax, used to compute the
weighted average in the cross-attention heads.
encoder_last_hidden_state (jnp.ndarray of shape (batch_size, sequence_length, hidden_size), optional) — Sequence of hidden-states at the output of the last layer of the encoder of the model.
encoder_hidden_states (tuple(jnp.ndarray), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple of jnp.ndarray (one for the output of the embeddings + one for the output of each layer) of shape
(batch_size, sequence_length, hidden_size).
Hidden-states of the encoder at the output of each layer plus the initial embedding outputs.
encoder_attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — Tuple of


FlaxBartForQuestionAnswering

 jnp.ndarray (one for each layer) of shape (batch_size, num_heads, sequence_length, sequence_length).
Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the
self-attention heads.
  The FlaxBartPreTrainedModel forward method, overrides the __call__ special method. Although the recipe for forward pass needs to be defined within this function, one should call the Module
instance afterwards instead of this since the former takes care of running the pre and post processing steps while
the latter silently ignores them.  Example:   Copied >>> from transformers import AutoTokenizer, FlaxBartForQuestionAnswering
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
>>> model = FlaxBartForQuestionAnswering.from_pretrained("facebook/bart-base")
>>> question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"
>>> inputs = tokenizer(question, text, return_tensors="jax")
>>> outputs = model(**inputs)
>>> start


FlaxBartForQuestionAnswering

_scores = outputs.start_logits
>>> end_scores = outputs.end_logits  encode  < source > ( input_ids: Array attention_mask: typing.Optional[jax.Array] = None position_ids: typing.Optional[jax.Array] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None train: bool = False params: dict = None dropout_rng: PRNGKey = None  ) → transformers.modeling_flax_outputs.FlaxBaseModelOutput or tuple(torch.FloatTensor)  Parameters   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
it.
Indices can be obtained using AutoTokenizer. See PreTrainedTokenizer.encode() and
PreTrainedTokenizer.call() for details.
What are input IDs?   attention_mask (jnp.ndarray of shape (batch_size, sequence_length), optional) —
Mask to avoid performing attention on padding token indices. Mask values


FlaxBartForQuestionAnswering

 selected in [0, 1]:
1 for tokens that are not masked,
0 for tokens that are masked.
What are attention masks?   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, config.max_position_embeddings - 1].   output_attentions (bool, optional) —
Whether or not to return the attentions tensors of all attention layers. See attentions under returned
tensors for more detail.   output_hidden_states (bool, optional) —
Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for
more detail.   return_dict (bool, optional) —
Whether or not to return a ModelOutput instead of a plain tuple.   Returns
transformers.modeling_flax_outputs.FlaxBaseModelOutput or tuple(torch.FloatTensor)

A transformers.modeling_flax_outputs.FlaxBaseModelOutput or a tuple of
torch.FloatTensor (if return_dict=False is passed or when co


FlaxBartForQuestionAnswering

nfig.return_dict=False) comprising various
elements depending on the configuration (<class 'transformers.models.bart.configuration_bart.BartConfig'>) and inputs.
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.
hidden_states (tuple(jnp.ndarray), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple of jnp.ndarray (one for the output of the embeddings + one for the output of each layer) of shape
(batch_size, sequence_length, hidden_size).
Hidden-states of the model at the output of each layer plus the initial embedding outputs.
attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — Tuple of jnp.ndarray (one for each layer) of shape (batch_size, num_heads, sequence_length, sequence_length).
Attentions weights after the attention softmax,


FlaxBartForQuestionAnswering

used to compute the weighted average in the self-attention
heads.
   Example:   Copied >>> from transformers import AutoTokenizer, FlaxBartForConditionalGeneration
>>> model = FlaxBartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
>>> text = "My friends are cool but they eat too many carbs."
>>> inputs = tokenizer(text, max_length=1024, return_tensors="jax")
>>> encoder_outputs = model.encode(**inputs)  decode  < source > ( decoder_input_ids encoder_outputs encoder_attention_mask: typing.Optional[jax.Array] = None decoder_attention_mask: typing.Optional[jax.Array] = None decoder_position_ids: typing.Optional[jax.Array] = None past_key_values: dict = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None train: bool = False params: dict = None dropout_rng: PRNGKey = None  ) → transformers.mo


FlaxBartForQuestionAnswering

deling_flax_outputs.FlaxBaseModelOutputWithPastAndCrossAttentions or tuple(torch.FloatTensor)  Parameters   decoder_input_ids (jnp.ndarray of shape (batch_size, target_sequence_length)) —
Indices of decoder input sequence tokens in the vocabulary.
Indices can be obtained using AutoTokenizer. See PreTrainedTokenizer.encode() and
PreTrainedTokenizer.call() for details.
What are decoder input IDs?
For translation and summarization training, decoder_input_ids should be provided. If no
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.   encoder_outputs (tuple(tuple(jnp.ndarray)) —
Tuple consists of (last_hidden_state, optional: hidden_states, 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 encoder. Used in the cross-attention of the decoder.   encoder_at


FlaxBartForQuestionAnswering

tention_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]:
1 for tokens that are not masked,
0 for tokens that are masked.
What are attention masks?   decoder_attention_mask (jnp.ndarray of shape (batch_size, target_sequence_length), optional) —
Default behavior: generate a tensor that ignores pad tokens in decoder_input_ids. Causal mask will also
be used by default.
If you want to change padding behavior, you should modify to your needs. See diagram 1 in the
paper for more information on the default strategy.   decoder_position_ids (numpy.ndarray of shape (batch_size, sequence_length), optional) —
Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the
range [0, config.max_position_embeddings - 1].   past_key_values (Dict[str, np.ndarray], optional, returned by init_cache or when passing previous past_key_values) —


FlaxBartForQuestionAnswering

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].   output_attentions (bool, optional) —
Whether or not to return the attentions tensors of all attention layers. See attentions under returned
tensors for more detail.   output_hidden_states (bool, optional) —
Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for
more detail.   return_dict (bool, optional) —
Whether or not to return a ModelOutput instead of a plain tuple.   Returns
transformers.modeling_flax_outputs.FlaxBaseModelOutputWithPastAndCrossAttentions or tuple(torch.FloatTensor)

A transformers.modeling_flax_outputs.FlaxBaseModelOutputWithPastAndCrossAttentions or a tuple of
torch.FloatTensor (if return_dict=False is passed or when config.return_dict=False) comprising various
elements depending on the


FlaxBartForQuestionAnswering

 configuration (<class 'transformers.models.bart.configuration_bart.BartConfig'>) and inputs.
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.
If past_key_values is used only the last hidden-state of the sequences of shape (batch_size, 1, hidden_size) is output.
past_key_values (tuple(tuple(jnp.ndarray)), optional, returned when use_cache=True is passed or when config.use_cache=True) — Tuple of tuple(jnp.ndarray) of length config.n_layers, with each tuple having 2 tensors of shape
(batch_size, num_heads, sequence_length, embed_size_per_head)) and optionally if
config.is_encoder_decoder=True 2 additional tensors of shape (batch_size, num_heads, encoder_sequence_length, embed_size_per_head).
Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if
config.is_encoder_decoder=True in the cross-attention blocks) that can be used


FlaxBartForQuestionAnswering

(see past_key_values
input) to speed up sequential decoding.
hidden_states (tuple(jnp.ndarray), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple of jnp.ndarray (one for the output of the embeddings + one for the output of each layer) of shape
(batch_size, sequence_length, hidden_size).
Hidden-states of the model at the output of each layer plus the initial embedding outputs.
attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — 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 the self-attention
heads.
cross_attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True and config.add_cross_attention=True is passed or when config.output_attentions=True) — Tuple of jnp.ndarray


FlaxBartForQuestionAnswering

(one for each layer) of shape (batch_size, num_heads, sequence_length, sequence_length).
Attentions weights of the decoder’s cross-attention layer, after the attention softmax, used to compute the
weighted average in the cross-attention heads.
   Example:   Copied >>> import jax.numpy as jnp
>>> from transformers import AutoTokenizer, FlaxBartForConditionalGeneration
>>> model = FlaxBartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
>>> text = "My friends are cool but they eat too many carbs."
>>> inputs = tokenizer(text, max_length=1024, return_tensors="jax")
>>> encoder_outputs = model.encode(**inputs)
>>> decoder_start_token_id = model.config.decoder_start_token_id
>>> decoder_input_ids = jnp.ones((inputs.input_ids.shape[0], 1), dtype="i4") * decoder_start_token_id
>>> outputs = model.decode(decoder_input_ids, encoder_outputs)
>>> last_decoder_hidden_states = outputs.last_hidd


FlaxBartForCausalLM

class transformers.FlaxBartForCausalLM  < source > ( config: BartConfig input_shape: typing.Tuple[int] = (1, 1) seed: int = 0 dtype: dtype = <class 'jax.numpy.float32'> _do_init: bool = True **kwargs  )   Parameters   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
configuration. Check out the from_pretrained() method to load the model weights.   dtype (jax.numpy.dtype, optional, defaults to jax.numpy.float32) —
The data type of the computation. Can be one of jax.numpy.float32, jax.numpy.float16 (on 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
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
parameters.
If you wish to


FlaxBartForCausalLM

 change the dtype of the model parameters, see to_fp16() and
to_bf16().    Bart Decoder Model with a language modeling head on top (linear layer with weights tied to the input embeddings)
e.g for autoregressive tasks. 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 embeddings, pruning heads
etc.) This model is also a Flax Linen
flax.nn.Module subclass. Use it as a
regular Flax 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: Just-In-Time (JIT) compilation Automatic Differentiation Vectorization Parallelization  __call__  < source > ( input_ids: Array attention_mask: typing.Optional[jax.Array] = None position_ids: typing.Optional[jax.Array] = None encoder_hidden_states: typing.Optional[jax.Array] = None encod


FlaxBartForCausalLM

er_attention_mask: typing.Optional[jax.Array] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None train: bool = False params: dict = None past_key_values: dict = None dropout_rng: PRNGKey = None  ) → transformers.modeling_flax_outputs.FlaxCausalLMOutputWithCrossAttentions or tuple(torch.FloatTensor)  Parameters   decoder_input_ids (jnp.ndarray of shape (batch_size, target_sequence_length)) —
Indices of decoder input sequence tokens in the vocabulary.
Indices can be obtained using AutoTokenizer. See PreTrainedTokenizer.encode() and
PreTrainedTokenizer.call() for details.
What are decoder input IDs?
For translation and summarization training, decoder_input_ids should be provided. If no
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.   encoder_outputs (tuple(tuple(jn


FlaxBartForCausalLM

p.ndarray)) —
Tuple consists of (last_hidden_state, optional: hidden_states, 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 encoder. Used in the cross-attention of the decoder.   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]:
1 for tokens that are not masked,
0 for tokens that are masked.
What are attention masks?   decoder_attention_mask (jnp.ndarray of shape (batch_size, target_sequence_length), optional) —
Default behavior: generate a tensor that ignores pad tokens in decoder_input_ids. Causal mask will also
be used by default.
If you want to change padding behavior, you should modify to your needs. See diagram 1 in the
paper for more information on the default strategy.   decoder_position_ids (numpy.n


FlaxBartForCausalLM

darray of shape (batch_size, sequence_length), optional) —
Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the
range [0, config.max_position_embeddings - 1].   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
auto-regressive decoding. Pre-computed key and value hidden-states are of shape [batch_size, max_length].   output_attentions (bool, optional) —
Whether or not to return the attentions tensors of all attention layers. See attentions under returned
tensors for more detail.   output_hidden_states (bool, optional) —
Whether or not to return the hidden states of all layers. See hidden_states under returned tensors for
more detail.   return_dict (bool, optional) —
Whether or not to return a ModelOutput instead of a plain tuple.   Returns


FlaxBartForCausalLM

transformers.modeling_flax_outputs.FlaxCausalLMOutputWithCrossAttentions or tuple(torch.FloatTensor)

A transformers.modeling_flax_outputs.FlaxCausalLMOutputWithCrossAttentions or a tuple of
torch.FloatTensor (if return_dict=False is passed or when config.return_dict=False) comprising various
elements depending on the configuration (BartConfig) and inputs.
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).
hidden_states (tuple(jnp.ndarray), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — Tuple of jnp.ndarray (one for the output of the embeddings + one for the output of each layer) of shape
(batch_size, sequence_length, hidden_size).
Hidden-states of the model at the output of each layer plus the initial embedding outputs.
attentions (tuple(jnp.ndarray), optional, returned wh


FlaxBartForCausalLM

en output_attentions=True is passed or when config.output_attentions=True) — 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 the self-attention
heads.
cross_attentions (tuple(jnp.ndarray), optional, returned when output_attentions=True is passed or when config.output_attentions=True) — Tuple of jnp.ndarray (one for each layer) of shape (batch_size, num_heads, sequence_length, sequence_length).
Cross attentions weights after the attention softmax, used to compute the weighted average in the
cross-attention heads.
past_key_values (tuple(tuple(jnp.ndarray)), optional, returned when use_cache=True is passed or when config.use_cache=True) — Tuple of jnp.ndarray tuples of length config.n_layers, with each tuple containing the cached key, value
states of the self-attention and the cross-attention layers if model is used


FlaxBartForCausalLM

 in encoder-decoder setting.
Only relevant if config.is_decoder = True.
Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see
past_key_values input) to speed up sequential decoding.
  The FlaxBartDecoderPreTrainedModel forward method, overrides the __call__ special method. Although the recipe for forward pass needs to be defined within this function, one should call the Module
instance afterwards instead of this since the former takes care of running the pre and post processing steps while
the latter silently ignores them.  Example:   Copied >>> from transformers import AutoTokenizer, FlaxBartForCausalLM
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
>>> model = FlaxBartForCausalLM.from_pretrained("facebook/bart-base")
>>> inputs = tokenizer("Hello, my dog is cute", return_tensors="np")
>>> outputs = model(**inputs)
>>> # retrieve logts for next token
>>> next_token_logits = outputs.logits[:, -1


Overview

The Llama2 model was proposed in LLaMA: Open Foundation and Fine-Tuned Chat Models by Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushka rMishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subram


Overview

anian, Xiaoqing EllenTan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, Thomas Scialom. It is a collection of foundation language models ranging from 7B to 70B parameters, with checkpoints finetuned for chat application!
The abstract from the paper is the following:
In this work, we develop and release Llama 2, a collection of pretrained and fine-tuned large language models (LLMs) ranging in scale from 7 billion to 70 billion parameters. Our fine-tuned LLMs, called Llama 2-Chat, are optimized for dialogue use cases. Our models outperform open-source chat models on most benchmarks we tested, and based on our human evaluations for helpfulness and safety, may be a suitable substitute for closed-source model


Overview

s. We provide a detailed description of our approach to fine-tuning and safety improvements of Llama 2-Chat in order to enable the community to build on our work and contribute to the responsible development of LLMs.
Checkout all Llama2 models here
The Llama2 models were trained using bfloat16, but the original inference uses float16. The checkpoints uploaded on the hub use torch_dtype = ‘float16’which will be used by theAutoModelAPI to cast the checkpoints fromtorch.float32totorch.float16`. The dtype of the online weights is mostly irrelevant, unless you are using torch_dtype="auto" when initializing a model using model = AutoModelForCausalLM.from_pretrained("path", torch_dtype = "auto"). The reason is that the model will first be downloaded ( using the dtype of the checkpoints online) then it will be casted to the default dtype of torch (becomes torch.flo


Overview

at32) and finally, if there is a torch_dtype provided in the config, it will be used. Training the model in float16 is not recommended and known to produce nan, as such the model should be trained in bfloat16.
Tips:
Weights for the Llama2 models can be obtained by filling out this form The architecture is very similar to the first Llama, with the addition of Grouped Query Attention (GQA) following this paper Setting config.pretraining_tp to a value different than 1 will activate the more accurate but slower computation of the linear layers, which should better match the original logits. The original model uses pad_id = -1 which means that there is no padding token. We can’t have the same logic, make sure to add a padding token using tokenizer.add_special_tokens({"pad_token":"<pad>"}) and resize the token embedding accordingly. You should also set the model.


Overview

config.pad_token_id. The embed_tokens layer of the model is initialized with self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.config.padding_idx), which makes sure that encoding the padding token will output zeros, so passing it when initializing is recommended. After filling out the form and gaining access to the model checkpoints, you should be able to use the already converted checkpoints. Otherwise, if you are converting your own model, feel free to use the conversion script. The script can be called with the following (example) command:
  Copied python src/transformers/models/llama/convert_llama_weights_to_hf.py \
    --input_dir /path/to/downloaded/llama/weights --model_size 7B --output_dir /output/path
After conversion, the model and tokenizer can be loaded via:
  Copied from transformers import LlamaForCausalLM, LlamaToke


Overview

nizer
tokenizer = LlamaTokenizer.from_pretrained("/output/path")
model = LlamaForCausalLM.from_pretrained("/output/path")
Note that executing the script requires enough CPU RAM to host the whole model in float16 precision (even if the biggest versions
come in several checkpoints they each contain a part of each weight of the model, so we need to load them all in RAM). For the 75B model, it’s thus 145GB of RAM needed.
The LLaMA tokenizer is a BPE model based on sentencepiece. One quirk of sentencepiece is that when decoding a sequence, if the first token is the start of the word (e.g. “Banana”), the tokenizer does not prepend the prefix space to the string.
This model was contributed by Arthur Zucker with contributions from Lysandre Debut. The code of the implementation in Hugging Face is based on GPT-NeoX here. The original code of the authors can be found


Resources

A list of official Hugging Face and community (indicated by 🌎) resources to help you get started with LLaMA2. If you’re interested in submitting a resource to be included here, please feel free to open a Pull Request and we’ll review it! The resource should ideally demonstrate something new instead of duplicating an existing resource.
Llama 2 is here - get it on Hugging Face, a blog post about Llama 2 and how to use it with 🤗 Transformers and 🤗 PEFT. LLaMA 2 - Every Resource you need, a compilation of relevant resources to learn about LLaMA 2 and how to get started quickly.
 Text Generation
A notebook on how to fine-tune Llama 2 in Google Colab using QLoRA and 4-bit precision. 🌎 A notebook on how to fine-tune the “Llama-v2-7b-guanaco” model with 4-bit QLoRA and generate Q&A datasets from PDFs. 🌎
 Text Classification
A notebook on how to fine-tune the Llama 2 model with QLoRa, TRL, and Korean text classi


Resources

fication dataset. 🌎🇰🇷
⚗ Optimization
Fine-tune Llama 2 with DPO, a guide to using the TRL library’s DPO method to fine tune Llama 2 on a specific dataset. Extended Guide: Instruction-tune Llama 2, a guide to training Llama 2 to generate instructions from inputs, transforming the model from instruction-following to instruction-giving. A notebook on how to fine-tune the Llama 2 model on a personal computer using QLoRa and TRL. 🌎
⚡ Inference
A notebook on how to quantize the Llama 2 model using GPTQ from the AutoGPTQ library. 🌎 A notebook on how to run the Llama 2 Chat Model with 4-bit quantization on a local computer or Google Colab. 🌎
🚀 Deploy
Fine-tune LLaMA 2 (7-70B) on Amazon SageMaker, a complete guide from setup to QLoRA fine-tuning and deployment on Amazon SageMaker. Deploy Llama 2 7B/13B/70B on Amazon SageMaker, a guide on using Hugging Face’s LLM DLC container for secure and scalable deployment


LlamaConfig

class transformers.LlamaConfig  < source > ( vocab_size = 32000 hidden_size = 4096 intermediate_size = 11008 num_hidden_layers = 32 num_attention_heads = 32 num_key_value_heads = None hidden_act = 'silu' max_position_embeddings = 2048 initializer_range = 0.02 rms_norm_eps = 1e-06 use_cache = True pad_token_id = None bos_token_id = 1 eos_token_id = 2 pretraining_tp = 1 tie_word_embeddings = False rope_theta = 10000.0 rope_scaling = None attention_bias = False **kwargs  )   Parameters   vocab_size (int, optional, defaults to 32000) —
Vocabulary size of the LLaMA model. Defines the number of different tokens that can be represented by the
inputs_ids passed when calling LlamaModel   hidden_size (int, optional, defaults to 4096) —
Dimension of the hidden representations.   intermediate_size (int, optional, defaults to 11008) —
Dimension of the MLP representations.   num_hidden_layers (int, optional, defaults to 32


LlamaConfig

) —
Number of hidden layers in the Transformer encoder.   num_attention_heads (int, optional, defaults to 32) —
Number of attention heads for each attention layer in the Transformer encoder.   num_key_value_heads (int, optional) —
This is the number of key_value heads that should be used to implement Grouped Query Attention. If
num_key_value_heads=num_attention_heads, the model will use Multi Head Attention (MHA), if
num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed by meanpooling all the original heads within that group. For more details checkout [this paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to num_attention_heads`.   pretraining_tp (int, optional, defaults to 1) —
Experimental feature. Tensor parallelism rank use


LlamaConfig

d during pretraining. Please refer to this
document to understand more about it. This value is
necessary to ensure exact reproducibility of the pretraining results. Please refer to this
issue.   hidden_act (str or function, optional, defaults to "silu") —
The non-linear activation function (function or string) in the decoder.   max_position_embeddings (int, optional, defaults to 2048) —
The maximum sequence length that this model might ever be used with. Llama 1 supports up to 2048 tokens,
Llama 2 up to 4096, CodeLlama up to 16384.   initializer_range (float, optional, defaults to 0.02) —
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.   rms_norm_eps (float, optional, defaults to 1e-12) —
The epsilon used by the rms normalization layers.   use_cache (bool, optional, defaults to True) —
Whether or not the model should return the last key/values attentions (not u


LlamaConfig

sed by all models). Only
relevant if config.is_decoder=True.   tie_word_embeddings(bool, optional, defaults to False) —
Whether to tie weight embeddings   rope_theta (float, optional, defaults to 10000.0) —
The base period of the RoPE embeddings.   rope_scaling (Dict, optional) —
Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
strategies: linear and dynamic. Their scaling factor must be an float greater than 1. The expected format
is {"type": strategy name, "factor": scaling factor}. When using this flag, don’t update
max_position_embeddings to the expected new maximum. See the following thread for more information on how
these scaling strategies behave:
https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
experimental feature, subject to breaking API changes in future versions.   attention_bias (boo


LlamaConfig

l, defaults to False) —
Whether to use a bias in the query, key, value and output projection layers during self-attention.
Example —    This is the configuration class to store the configuration of a LlamaModel. It is used to instantiate an LLaMA
model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
defaults will yield a similar configuration to that of the LLaMA-7B. Configuration objects inherit from PretrainedConfig and can be used to control the model outputs. Read the
documentation from PretrainedConfig for more information.    Copied >>> from transformers import LlamaModel, LlamaConfig
>>> # Initializing a LLaMA llama-7b style configuration
>>> configuration = LlamaConfig()
>>> # Initializing a model from the llama-7b style configuration
>>> model = LlamaModel(configuration)
>>> # Accessing the model configuration
>>> configuration = model.con


LlamaTokenizer

class transformers.LlamaTokenizer  < source > ( vocab_file unk_token = '<unk>' bos_token = '<s>' eos_token = '</s>' pad_token = None sp_model_kwargs: typing.Union[typing.Dict[str, typing.Any], NoneType] = None add_bos_token = True add_eos_token = False clean_up_tokenization_spaces = False use_default_system_prompt = True spaces_between_special_tokens = False legacy = None **kwargs  )   Parameters   vocab_file (str) —
Path to the vocabulary file.   legacy (bool, optional) —
Whether or not the legacy behavior of the tokenizer should be used. Legacy is before the merge of #24622
and #25224 which includes fixes to properly handle tokens that appear after special tokens. A simple
example:
legacy=True:
    Construct a Llama tokenizer. Based on byte-level Byte-Pair-Encoding. The default padding token is unset as there is
no padding token in the original model.  build_inputs_with_special_token


LlamaTokenizer

s  < source > ( token_ids_0 token_ids_1 = None  )     get_special_tokens_mask  < source > ( token_ids_0: typing.List[int] token_ids_1: typing.Optional[typing.List[int]] = None already_has_special_tokens: bool = False  ) → List[int]  Parameters   token_ids_0 (List[int]) —
List of IDs.   token_ids_1 (List[int], optional) —
Optional second list of IDs for sequence pairs.   already_has_special_tokens (bool, optional, defaults to False) —
Whether or not the token list is already formatted with special tokens for the model.   Returns
List[int]

A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
  Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
special tokens using the tokenizer prepare_for_model method.  create_token_type_ids_from_sequences  < source > ( token_ids_0: typing.List[int] token_ids_


LlamaTokenizer

1: typing.Optional[typing.List[int]] = None  ) → List[int]  Parameters   token_ids_0 (List[int]) —
List of ids.   token_ids_1 (List[int], optional) —
Optional second list of IDs for sequence pairs.   Returns
List[int]

List of token type IDs according to the given sequence(s).
  Creates a mask from the two sequences passed to be used in a sequence-pair classification task. An ALBERT  sequence pair mask has the following format:   Copied 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
| first sequence    | second sequence | if token_ids_1 is None, only returns the first portion of the mask (0s).  save_vocabulary  < source > ( save_directory filename_prefix: typing.Optional[str] = None  ) → Tuple(str)  Parameters   save_directory (str) —
The directory in which to save the vocabulary.   Returns
Tuple(str)

Paths to the files saved.
  Save the vocabulary and special tokens file to a director


LlamaTokenizerFast

class transformers.LlamaTokenizerFast  < source > ( vocab_file = None tokenizer_file = None clean_up_tokenization_spaces = False unk_token = '<unk>' bos_token = '<s>' eos_token = '</s>' add_bos_token = True add_eos_token = False use_default_system_prompt = True **kwargs  )   Parameters   vocab_file (str) —
SentencePiece file (generally has a .model extension) that
contains the vocabulary necessary to instantiate a tokenizer.   tokenizer_file (str) —
tokenizers file (generally has a .json extension) that
contains everything needed to load the tokenizer.   clean_up_tokenization_spaces (str, optional, defaults to False) —
Wether to cleanup spaces after decoding, cleanup consists in removing potential artifacts like extra
spaces.   bos_token (str, optional, defaults to "<s>") —
The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.   eos_token (str, optional, defaul


LlamaTokenizerFast

ts to "</s>") —
The end of sequence token.   unk_token (str, optional, defaults to "<unk>") —
The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
token instead.    Construct a Llama tokenizer. Based on byte-level Byte-Pair-Encoding. This uses notably ByteFallback and no normalization.    Copied from transformers import LlamaTokenizerFast
tokenizer = LlamaTokenizerFast.from_pretrained("hf-internal-testing/llama-tokenizer")
tokenizer.encode("Hello this is a test")
>>> [1, 15043, 445, 338, 263, 1243] If you want to change the bos_token or the eos_token, make sure to specify them when initializing the model, or
call tokenizer.update_post_processor() to make sure that the post-processing is correctly done (otherwise the
values of the first token and final token of an encoded sequence will not be correct). For more details, checkout
[post-processors] (https://hugg


LlamaTokenizerFast

ingface.co/docs/tokenizers/api/post-processors) documentation. This tokenizer inherits from PreTrainedTokenizerFast which contains most of the main methods. Users should
refer to this superclass for more information regarding those methods.  build_inputs_with_special_tokens  < source > ( token_ids_0 token_ids_1 = None  )     get_special_tokens_mask  < source > ( token_ids_0: typing.List[int] token_ids_1: typing.Optional[typing.List[int]] = None already_has_special_tokens: bool = False  ) → A list of integers in the range [0, 1]  Parameters   token_ids_0 (List[int]) —
List of ids of the first sequence.   token_ids_1 (List[int], optional) —
List of ids of the second sequence.   already_has_special_tokens (bool, optional, defaults to False) —
Whether or not the token list is already formatted with special tokens for the model.   Returns
A list of integers in the range [0, 1]

1 for a special token, 0 for a sequen


LlamaTokenizerFast

ce token.
  Retrieves sequence ids from a token list that has no special tokens added. This method is called when adding
special tokens using the tokenizer prepare_for_model or encode_plus methods.  create_token_type_ids_from_sequences  < source > ( token_ids_0: typing.List[int] token_ids_1: typing.Optional[typing.List[int]] = None  ) → List[int]  Parameters   token_ids_0 (List[int]) — The first tokenized sequence.   token_ids_1 (List[int], optional) — The second tokenized sequence.   Returns
List[int]

The token type ids.
  Create the token type IDs corresponding to the sequences passed. What are token type
IDs? Should be overridden in a subclass if the model has a special way of building those.  update_post_processor  < source > ( )    Updates the underlying post processor with the current bos_token and eos_token.  save_vocabulary  < source > ( save_directory: str filename_prefix: typing.Optional[str] = None


LlamaModel

class transformers.LlamaModel  < source > ( config: LlamaConfig  )   Parameters   config (LlamaConfig) —
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 configuration. Check out the
from_pretrained() method to load the model weights.
config — LlamaConfig    The bare LLaMA Model outputting raw hidden-states without any specific head on top.
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, pruning heads
etc.) This model is also a PyTorch 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 behavior. Transformer decoder consisting of config.num_hidden_layers layers. Each layer is a LlamaDecoderLayer  forward  < s


LlamaModel

ource > ( input_ids: LongTensor = None attention_mask: typing.Optional[torch.Tensor] = None position_ids: typing.Optional[torch.LongTensor] = None past_key_values: typing.Optional[typing.List[torch.FloatTensor]] = None inputs_embeds: typing.Optional[torch.FloatTensor] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None  )   Parameters   input_ids (torch.LongTensor of shape (batch_size, sequence_length)) —
Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
it.
Indices can be obtained using AutoTokenizer. See PreTrainedTokenizer.encode() and
PreTrainedTokenizer.call() for details.
What are input IDs?   attention_mask (torch.Tensor of shape (batch_size, sequence_length), optional) —
Mask to avoid performing attention on padding token indices. Mask values selected in [0, 1]:
1


LlamaModel

 for tokens that are not masked,
0 for tokens that are masked.
What are attention masks?
Indices can be obtained using AutoTokenizer. See PreTrainedTokenizer.encode() and
PreTrainedTokenizer.call() for details.
If past_key_values is used, optionally only the last input_ids have to be input (see
past_key_values).
If you want to change padding behavior, you should read modeling_opt._prepare_decoder_attention_mask
and modify to your needs. See diagram 1 in the paper for more
information on the default strategy.
1 indicates the head is not masked,
0 indicates the head is masked.
   position_ids (torch.LongTensor of shape (batch_size, sequence_length), optional) —
Indices of positions of each input sequence tokens in the position embeddings. Selected in the range [0, config.n_positions - 1].
What are position IDs?   past_key_values (tuple(tuple(torch.FloatTensor)), optional, returned when use_cache=True is passed or when config.use_cache=True) —
Tuple of tuple(torc


LlamaModel

h.FloatTensor) of length config.n_layers, with each tuple having 2 tensors of shape
(batch_size, num_heads, sequence_length, embed_size_per_head)) and 2 additional tensors of shape
(batch_size, num_heads, encoder_sequence_length, embed_size_per_head).
Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
blocks) that can be used (see past_key_values input) to speed up sequential decoding.
If past_key_values are used, the user can optionally input only the last input_ids (those that don’t
have their past key value states given to this model) of shape (batch_size, 1) instead of all input_ids
of shape (batch_size, sequence_length).   inputs_embeds (torch.FloatTensor of shape (batch_size, sequence_length, hidden_size), optional) —
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 input_ids indices into