Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2021 The Fairseq Authors and The Google Flax Team Authors And The HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""Flax Bart model."""16 17import math18import random19from functools import partial20from typing import Callable, Optional21 22import flax.linen as nn23import jax24import jax.numpy as jnp25from flax.core.frozen_dict import FrozenDict, freeze, unfreeze26from flax.linen import combine_masks, make_causal_mask27from flax.linen.attention import dot_product_attention_weights28from flax.traverse_util import flatten_dict, unflatten_dict29from jax import lax30from jax.random import PRNGKey31 32from ...modeling_flax_outputs import (33 FlaxBaseModelOutput,34 FlaxBaseModelOutputWithPastAndCrossAttentions,35 FlaxCausalLMOutputWithCrossAttentions,36 FlaxSeq2SeqLMOutput,37 FlaxSeq2SeqModelOutput,38 FlaxSeq2SeqQuestionAnsweringModelOutput,39 FlaxSeq2SeqSequenceClassifierOutput,40)41from ...modeling_flax_utils import (42 ACT2FN,43 FlaxPreTrainedModel,44 append_call_sample_docstring,45 append_replace_return_docstrings,46 overwrite_call_docstring,47)48from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings49from .configuration_bart import BartConfig50 51 52logger = logging.get_logger(__name__)53 54_CHECKPOINT_FOR_DOC = "facebook/bart-base"55_CONFIG_FOR_DOC = "BartConfig"56 57 58BART_START_DOCSTRING = r"""59 This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the generic methods the60 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads61 etc.)62 63 This model is also a Flax Linen64 [flax.nn.Module](https://flax.readthedocs.io/en/latest/_autosummary/flax.nn.module.html) subclass. Use it as a65 regular Flax Module and refer to the Flax documentation for all matter related to general usage and behavior.66 67 Finally, this model supports inherent JAX features such as:68 69 - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)70 - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)71 - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)72 - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)73 74 Parameters:75 config ([`BartConfig`]): Model configuration class with all the parameters of the model.76 Initializing with a config file does not load the weights associated with the model, only the77 configuration. Check out the [`~FlaxPreTrainedModel.from_pretrained`] method to load the model weights.78 dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):79 The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on GPUs) and80 `jax.numpy.bfloat16` (on TPUs).81 82 This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If83 specified all the computation will be performed with the given `dtype`.84 85 **Note that this only specifies the dtype of the computation and does not influence the dtype of model86 parameters.**87 88 If you wish to change the dtype of the model parameters, see [`~FlaxPreTrainedModel.to_fp16`] and89 [`~FlaxPreTrainedModel.to_bf16`].90"""91 92BART_INPUTS_DOCSTRING = r"""93 Args:94 input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):95 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide96 it.97 98 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and99 [`PreTrainedTokenizer.__call__`] for details.100 101 [What are input IDs?](../glossary#input-ids)102 attention_mask (`jnp.ndarray` of shape `(batch_size, sequence_length)`, *optional*):103 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:104 105 - 1 for tokens that are **not masked**,106 - 0 for tokens that are **masked**.107 108 [What are attention masks?](../glossary#attention-mask)109 decoder_input_ids (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):110 Indices of decoder input sequence tokens in the vocabulary.111 112 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and113 [`PreTrainedTokenizer.__call__`] for details.114 115 [What are decoder input IDs?](../glossary#decoder-input-ids)116 117 For translation and summarization training, `decoder_input_ids` should be provided. If no118 `decoder_input_ids` is provided, the model will create this tensor by shifting the `input_ids` to the right119 for denoising pre-training following the paper.120 decoder_attention_mask (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):121 Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also122 be used by default.123 124 If you want to change padding behavior, you should modify to your needs. See diagram 1 in [the125 paper](https://huggingface.co/papers/1910.13461) for more information on the default strategy.126 position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):127 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,128 config.max_position_embeddings - 1]`.129 decoder_position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):130 Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the131 range `[0, config.max_position_embeddings - 1]`.132 output_attentions (`bool`, *optional*):133 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned134 tensors for more detail.135 output_hidden_states (`bool`, *optional*):136 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for137 more detail.138 return_dict (`bool`, *optional*):139 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.140"""141 142 143BART_ENCODE_INPUTS_DOCSTRING = r"""144 Args:145 input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):146 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide147 it.148 149 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and150 [`PreTrainedTokenizer.__call__`] for details.151 152 [What are input IDs?](../glossary#input-ids)153 attention_mask (`jnp.ndarray` of shape `(batch_size, sequence_length)`, *optional*):154 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:155 156 - 1 for tokens that are **not masked**,157 - 0 for tokens that are **masked**.158 159 [What are attention masks?](../glossary#attention-mask)160 position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):161 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,162 config.max_position_embeddings - 1]`.163 output_attentions (`bool`, *optional*):164 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned165 tensors for more detail.166 output_hidden_states (`bool`, *optional*):167 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for168 more detail.169 return_dict (`bool`, *optional*):170 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.171"""172 173BART_DECODE_INPUTS_DOCSTRING = r"""174 Args:175 decoder_input_ids (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`):176 Indices of decoder input sequence tokens in the vocabulary.177 178 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and179 [`PreTrainedTokenizer.__call__`] for details.180 181 [What are decoder input IDs?](../glossary#decoder-input-ids)182 183 For translation and summarization training, `decoder_input_ids` should be provided. If no184 `decoder_input_ids` is provided, the model will create this tensor by shifting the `input_ids` to the right185 for denoising pre-training following the paper.186 encoder_outputs (`tuple(tuple(jnp.ndarray)`):187 Tuple consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*: `attentions`)188 `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*) is a sequence of189 hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder.190 encoder_attention_mask (`jnp.ndarray` of shape `(batch_size, sequence_length)`, *optional*):191 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:192 193 - 1 for tokens that are **not masked**,194 - 0 for tokens that are **masked**.195 196 [What are attention masks?](../glossary#attention-mask)197 decoder_attention_mask (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):198 Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also199 be used by default.200 201 If you want to change padding behavior, you should modify to your needs. See diagram 1 in [the202 paper](https://huggingface.co/papers/1910.13461) for more information on the default strategy.203 decoder_position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):204 Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the205 range `[0, config.max_position_embeddings - 1]`.206 past_key_values (`dict[str, np.ndarray]`, *optional*, returned by `init_cache` or when passing previous `past_key_values`):207 Dictionary of pre-computed hidden-states (key and values in the attention blocks) that can be used for fast208 auto-regressive decoding. Pre-computed key and value hidden-states are of shape *[batch_size, max_length]*.209 output_attentions (`bool`, *optional*):210 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned211 tensors for more detail.212 output_hidden_states (`bool`, *optional*):213 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for214 more detail.215 return_dict (`bool`, *optional*):216 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.217"""218 219 220def shift_tokens_right(input_ids: jnp.ndarray, pad_token_id: int, decoder_start_token_id: int) -> jnp.ndarray:221 """222 Shift input ids one token to the right.223 """224 shifted_input_ids = jnp.zeros_like(input_ids)225 shifted_input_ids = shifted_input_ids.at[:, 1:].set(input_ids[:, :-1])226 shifted_input_ids = shifted_input_ids.at[:, 0].set(decoder_start_token_id)227 228 shifted_input_ids = jnp.where(shifted_input_ids == -100, pad_token_id, shifted_input_ids)229 return shifted_input_ids230 231 232class FlaxBartAttention(nn.Module):233 config: BartConfig234 embed_dim: int235 num_heads: int236 dropout: float = 0.0237 causal: bool = False238 bias: bool = True239 dtype: jnp.dtype = jnp.float32 # the dtype of the computation240 241 def setup(self) -> None:242 self.head_dim = self.embed_dim // self.num_heads243 if self.head_dim * self.num_heads != self.embed_dim:244 raise ValueError(245 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"246 f" and `num_heads`: {self.num_heads})."247 )248 249 dense = partial(250 nn.Dense,251 self.embed_dim,252 use_bias=self.bias,253 dtype=self.dtype,254 kernel_init=jax.nn.initializers.normal(self.config.init_std),255 )256 257 self.q_proj, self.k_proj, self.v_proj = dense(), dense(), dense()258 self.out_proj = dense()259 260 self.dropout_layer = nn.Dropout(rate=self.dropout)261 262 if self.causal:263 self.causal_mask = make_causal_mask(264 jnp.ones((1, self.config.max_position_embeddings), dtype="bool"), dtype="bool"265 )266 267 def _split_heads(self, hidden_states):268 return hidden_states.reshape(hidden_states.shape[:2] + (self.num_heads, self.head_dim))269 270 def _merge_heads(self, hidden_states):271 return hidden_states.reshape(hidden_states.shape[:2] + (self.embed_dim,))272 273 @nn.compact274 def _concatenate_to_cache(self, key, value, query, attention_mask):275 """276 This function takes projected key, value states from a single input token and concatenates the states to cached277 states from previous steps. This function is slightly adapted from the official Flax repository:278 https://github.com/google/flax/blob/491ce18759622506588784b4fca0e4bf05f8c8cd/flax/linen/attention.py#L252279 """280 # detect if we're initializing by absence of existing cache data.281 is_initialized = self.has_variable("cache", "cached_key")282 cached_key = self.variable("cache", "cached_key", jnp.zeros, key.shape, key.dtype)283 cached_value = self.variable("cache", "cached_value", jnp.zeros, value.shape, value.dtype)284 cache_index = self.variable("cache", "cache_index", lambda: jnp.array(0, dtype=jnp.int32))285 286 if is_initialized:287 *batch_dims, max_length, num_heads, depth_per_head = cached_key.value.shape288 # update key, value caches with our new 1d spatial slices289 cur_index = cache_index.value290 indices = (0,) * len(batch_dims) + (cur_index, 0, 0)291 key = lax.dynamic_update_slice(cached_key.value, key, indices)292 value = lax.dynamic_update_slice(cached_value.value, value, indices)293 cached_key.value = key294 cached_value.value = value295 num_updated_cache_vectors = query.shape[1]296 cache_index.value = cache_index.value + num_updated_cache_vectors297 # causal mask for cached decoder self-attention: our single query position should only attend to those key positions that have already been generated and cached, not the remaining zero elements.298 pad_mask = jnp.broadcast_to(299 jnp.arange(max_length) < cur_index + num_updated_cache_vectors,300 tuple(batch_dims) + (1, num_updated_cache_vectors, max_length),301 )302 attention_mask = combine_masks(pad_mask, attention_mask)303 return key, value, attention_mask304 305 def __call__(306 self,307 hidden_states: jnp.ndarray,308 key_value_states: Optional[jnp.ndarray] = None,309 attention_mask: Optional[jnp.ndarray] = None,310 init_cache: bool = False,311 deterministic: bool = True,312 ) -> tuple[jnp.ndarray]:313 """Input shape: Batch x Time x Channel"""314 315 # if key_value_states are provided this layer is used as a cross-attention layer316 # for the decoder317 is_cross_attention = key_value_states is not None318 batch_size = hidden_states.shape[0]319 320 # get query proj321 query_states = self.q_proj(hidden_states)322 # get key, value proj323 if is_cross_attention:324 # cross_attentions325 key_states = self.k_proj(key_value_states)326 value_states = self.v_proj(key_value_states)327 else:328 # self_attention329 key_states = self.k_proj(hidden_states)330 value_states = self.v_proj(hidden_states)331 332 query_states = self._split_heads(query_states)333 key_states = self._split_heads(key_states)334 value_states = self._split_heads(value_states)335 336 # handle cache prepare causal attention mask337 if self.causal:338 query_length, key_length = query_states.shape[1], key_states.shape[1]339 if self.has_variable("cache", "cached_key"):340 mask_shift = self.variables["cache"]["cache_index"]341 max_decoder_length = self.variables["cache"]["cached_key"].shape[1]342 causal_mask = lax.dynamic_slice(343 self.causal_mask, (0, 0, mask_shift, 0), (1, 1, query_length, max_decoder_length)344 )345 else:346 causal_mask = self.causal_mask[:, :, :query_length, :key_length]347 causal_mask = jnp.broadcast_to(causal_mask, (batch_size,) + causal_mask.shape[1:])348 349 # combine masks if needed350 if attention_mask is not None and self.causal:351 attention_mask = jnp.broadcast_to(jnp.expand_dims(attention_mask, axis=(-3, -2)), causal_mask.shape)352 attention_mask = combine_masks(attention_mask, causal_mask)353 elif self.causal:354 attention_mask = causal_mask355 elif attention_mask is not None:356 attention_mask = jnp.expand_dims(attention_mask, axis=(-3, -2))357 358 # During fast autoregressive decoding, we feed one position at a time,359 # and cache the keys and values step by step.360 if self.causal and (self.has_variable("cache", "cached_key") or init_cache):361 key_states, value_states, attention_mask = self._concatenate_to_cache(362 key_states, value_states, query_states, attention_mask363 )364 365 # Convert the boolean attention mask to an attention bias.366 if attention_mask is not None:367 # attention mask in the form of attention bias368 attention_bias = lax.select(369 attention_mask > 0,370 jnp.full(attention_mask.shape, 0.0).astype(self.dtype),371 jnp.full(attention_mask.shape, jnp.finfo(self.dtype).min).astype(self.dtype),372 )373 else:374 attention_bias = None375 376 dropout_rng = None377 if not deterministic and self.dropout > 0.0:378 dropout_rng = self.make_rng("dropout")379 380 attn_weights = dot_product_attention_weights(381 query_states,382 key_states,383 bias=attention_bias,384 dropout_rng=dropout_rng,385 dropout_rate=self.dropout,386 broadcast_dropout=True,387 deterministic=deterministic,388 dtype=self.dtype,389 precision=None,390 )391 392 attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, value_states)393 attn_output = self._merge_heads(attn_output)394 attn_output = self.out_proj(attn_output)395 396 return attn_output, attn_weights397 398 399class FlaxBartEncoderLayer(nn.Module):400 config: BartConfig401 dtype: jnp.dtype = jnp.float32402 403 def setup(self) -> None:404 self.embed_dim = self.config.d_model405 self.self_attn = FlaxBartAttention(406 config=self.config,407 embed_dim=self.embed_dim,408 num_heads=self.config.encoder_attention_heads,409 dropout=self.config.attention_dropout,410 dtype=self.dtype,411 )412 self.self_attn_layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)413 self.dropout_layer = nn.Dropout(rate=self.config.dropout)414 self.activation_fn = ACT2FN[self.config.activation_function]415 self.activation_dropout_layer = nn.Dropout(rate=self.config.activation_dropout)416 self.fc1 = nn.Dense(417 self.config.encoder_ffn_dim,418 dtype=self.dtype,419 kernel_init=jax.nn.initializers.normal(self.config.init_std),420 )421 self.fc2 = nn.Dense(422 self.embed_dim, dtype=self.dtype, kernel_init=jax.nn.initializers.normal(self.config.init_std)423 )424 self.final_layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)425 426 def __call__(427 self,428 hidden_states: jnp.ndarray,429 attention_mask: jnp.ndarray,430 output_attentions: bool = True,431 deterministic: bool = True,432 ) -> tuple[jnp.ndarray]:433 residual = hidden_states434 hidden_states, attn_weights = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask)435 436 hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)437 hidden_states = residual + hidden_states438 hidden_states = self.self_attn_layer_norm(hidden_states)439 440 residual = hidden_states441 hidden_states = self.activation_fn(self.fc1(hidden_states))442 hidden_states = self.activation_dropout_layer(hidden_states, deterministic=deterministic)443 hidden_states = self.fc2(hidden_states)444 hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)445 hidden_states = residual + hidden_states446 hidden_states = self.final_layer_norm(hidden_states)447 448 outputs = (hidden_states,)449 450 if output_attentions:451 outputs += (attn_weights,)452 453 return outputs454 455 456class FlaxBartEncoderLayerCollection(nn.Module):457 config: BartConfig458 dtype: jnp.dtype = jnp.float32 # the dtype of the computation459 460 def setup(self):461 self.layers = [462 FlaxBartEncoderLayer(self.config, name=str(i), dtype=self.dtype) for i in range(self.config.encoder_layers)463 ]464 self.layerdrop = self.config.encoder_layerdrop465 466 def __call__(467 self,468 hidden_states,469 attention_mask,470 deterministic: bool = True,471 output_attentions: bool = False,472 output_hidden_states: bool = False,473 return_dict: bool = True,474 ):475 all_attentions = () if output_attentions else None476 all_hidden_states = () if output_hidden_states else None477 478 for encoder_layer in self.layers:479 if output_hidden_states:480 all_hidden_states = all_hidden_states + (hidden_states,)481 # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)482 dropout_probability = random.uniform(0, 1)483 if not deterministic and (dropout_probability < self.layerdrop): # skip the layer484 layer_outputs = (None, None)485 else:486 layer_outputs = encoder_layer(487 hidden_states,488 attention_mask,489 output_attentions,490 deterministic,491 )492 hidden_states = layer_outputs[0]493 if output_attentions:494 all_attentions = all_attentions + (layer_outputs[1],)495 496 if output_hidden_states:497 all_hidden_states += (hidden_states,)498 499 outputs = (hidden_states, all_hidden_states, all_attentions)500 501 if not return_dict:502 return tuple(v for v in outputs if v is not None)503 504 return FlaxBaseModelOutput(505 last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions506 )507 508 509class FlaxBartDecoderLayer(nn.Module):510 config: BartConfig511 dtype: jnp.dtype = jnp.float32512 513 def setup(self) -> None:514 self.embed_dim = self.config.d_model515 self.self_attn = FlaxBartAttention(516 config=self.config,517 embed_dim=self.embed_dim,518 num_heads=self.config.decoder_attention_heads,519 dropout=self.config.attention_dropout,520 causal=True,521 dtype=self.dtype,522 )523 self.dropout_layer = nn.Dropout(rate=self.config.dropout)524 self.activation_fn = ACT2FN[self.config.activation_function]525 self.activation_dropout_layer = nn.Dropout(rate=self.config.activation_dropout)526 527 self.self_attn_layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)528 self.encoder_attn = FlaxBartAttention(529 config=self.config,530 embed_dim=self.embed_dim,531 num_heads=self.config.decoder_attention_heads,532 dropout=self.config.attention_dropout,533 dtype=self.dtype,534 )535 self.encoder_attn_layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)536 self.fc1 = nn.Dense(537 self.config.decoder_ffn_dim,538 dtype=self.dtype,539 kernel_init=jax.nn.initializers.normal(self.config.init_std),540 )541 self.fc2 = nn.Dense(542 self.embed_dim, dtype=self.dtype, kernel_init=jax.nn.initializers.normal(self.config.init_std)543 )544 self.final_layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)545 546 def __call__(547 self,548 hidden_states: jnp.ndarray,549 attention_mask: jnp.ndarray,550 encoder_hidden_states: Optional[jnp.ndarray] = None,551 encoder_attention_mask: Optional[jnp.ndarray] = None,552 init_cache: bool = False,553 output_attentions: bool = True,554 deterministic: bool = True,555 ) -> tuple[jnp.ndarray]:556 residual = hidden_states557 558 # Self Attention559 hidden_states, self_attn_weights = self.self_attn(560 hidden_states=hidden_states, attention_mask=attention_mask, init_cache=init_cache561 )562 hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)563 hidden_states = residual + hidden_states564 hidden_states = self.self_attn_layer_norm(hidden_states)565 566 # Cross-Attention Block567 cross_attn_weights = None568 if encoder_hidden_states is not None:569 residual = hidden_states570 571 hidden_states, cross_attn_weights = self.encoder_attn(572 hidden_states=hidden_states,573 key_value_states=encoder_hidden_states,574 attention_mask=encoder_attention_mask,575 )576 hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)577 hidden_states = residual + hidden_states578 hidden_states = self.encoder_attn_layer_norm(hidden_states)579 580 # Fully Connected581 residual = hidden_states582 hidden_states = self.activation_fn(self.fc1(hidden_states))583 hidden_states = self.activation_dropout_layer(hidden_states, deterministic=deterministic)584 hidden_states = self.fc2(hidden_states)585 hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)586 hidden_states = residual + hidden_states587 hidden_states = self.final_layer_norm(hidden_states)588 589 outputs = (hidden_states,)590 591 if output_attentions:592 outputs += (self_attn_weights, cross_attn_weights)593 594 return outputs595 596 597class FlaxBartDecoderLayerCollection(nn.Module):598 config: BartConfig599 dtype: jnp.dtype = jnp.float32 # the dtype of the computation600 601 def setup(self):602 self.layers = [603 FlaxBartDecoderLayer(self.config, name=str(i), dtype=self.dtype) for i in range(self.config.decoder_layers)604 ]605 self.layerdrop = self.config.decoder_layerdrop606 607 def __call__(608 self,609 hidden_states,610 attention_mask,611 encoder_hidden_states: Optional[jnp.ndarray] = None,612 encoder_attention_mask: Optional[jnp.ndarray] = None,613 deterministic: bool = True,614 init_cache: bool = False,615 output_attentions: bool = False,616 output_hidden_states: bool = False,617 return_dict: bool = True,618 ):619 # decoder layers620 all_hidden_states = () if output_hidden_states else None621 all_self_attns = () if output_attentions else None622 all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None623 624 for decoder_layer in self.layers:625 if output_hidden_states:626 all_hidden_states += (hidden_states,)627 # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)628 dropout_probability = random.uniform(0, 1)629 if not deterministic and (dropout_probability < self.layerdrop):630 layer_outputs = (None, None, None)631 else:632 layer_outputs = decoder_layer(633 hidden_states,634 attention_mask=attention_mask,635 encoder_hidden_states=encoder_hidden_states,636 encoder_attention_mask=encoder_attention_mask,637 init_cache=init_cache,638 output_attentions=output_attentions,639 deterministic=deterministic,640 )641 642 hidden_states = layer_outputs[0]643 if output_attentions:644 all_self_attns += (layer_outputs[1],)645 646 if encoder_hidden_states is not None:647 all_cross_attentions += (layer_outputs[2],)648 649 # add hidden states from the last decoder layer650 if output_hidden_states:651 all_hidden_states += (hidden_states,)652 653 outputs = [hidden_states, all_hidden_states, all_self_attns, all_cross_attentions]654 655 if not return_dict:656 return tuple(v for v in outputs if v is not None)657 658 return FlaxBaseModelOutputWithPastAndCrossAttentions(659 last_hidden_state=hidden_states,660 hidden_states=all_hidden_states,661 attentions=all_self_attns,662 cross_attentions=all_cross_attentions,663 )664 665 666class FlaxBartClassificationHead(nn.Module):667 """Head for sentence-level classification tasks."""668 669 config: BartConfig670 inner_dim: int671 num_classes: int672 pooler_dropout: float673 dtype: jnp.dtype = jnp.float32674 675 def setup(self):676 self.dense = nn.Dense(677 self.inner_dim, dtype=self.dtype, kernel_init=jax.nn.initializers.normal(self.config.init_std)678 )679 self.dropout = nn.Dropout(rate=self.pooler_dropout)680 self.out_proj = nn.Dense(681 self.num_classes,682 dtype=self.dtype,683 kernel_init=jax.nn.initializers.normal(self.config.init_std),684 )685 686 def __call__(self, hidden_states: jnp.ndarray, deterministic: bool):687 hidden_states = self.dropout(hidden_states, deterministic=deterministic)688 hidden_states = self.dense(hidden_states)689 hidden_states = jnp.tanh(hidden_states)690 hidden_states = self.dropout(hidden_states, deterministic=deterministic)691 hidden_states = self.out_proj(hidden_states)692 return hidden_states693 694 695class FlaxBartEncoder(nn.Module):696 config: BartConfig697 embed_tokens: nn.Embed698 dtype: jnp.dtype = jnp.float32 # the dtype of the computation699 700 def setup(self):701 self.dropout_layer = nn.Dropout(rate=self.config.dropout)702 703 embed_dim = self.config.d_model704 self.padding_idx = self.config.pad_token_id705 self.max_source_positions = self.config.max_position_embeddings706 self.embed_scale = math.sqrt(embed_dim) if self.config.scale_embedding else 1.0707 708 # Bart is set up so that if padding_idx is specified then offset the embedding ids by 2709 # and adjust num_embeddings appropriately. Other models don't have this hack710 self.offset = 2711 self.embed_positions = nn.Embed(712 self.config.max_position_embeddings + self.offset,713 embed_dim,714 embedding_init=jax.nn.initializers.normal(self.config.init_std),715 dtype=self.dtype,716 )717 self.layers = FlaxBartEncoderLayerCollection(self.config, self.dtype)718 self.layernorm_embedding = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)719 720 def __call__(721 self,722 input_ids,723 attention_mask,724 position_ids,725 output_attentions: bool = False,726 output_hidden_states: bool = False,727 return_dict: bool = True,728 deterministic: bool = True,729 ):730 input_shape = input_ids.shape731 input_ids = input_ids.reshape(-1, input_shape[-1])732 733 inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale734 735 embed_pos = self.embed_positions(position_ids + self.offset)736 737 hidden_states = inputs_embeds + embed_pos738 hidden_states = self.layernorm_embedding(hidden_states)739 hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)740 741 outputs = self.layers(742 hidden_states,743 attention_mask,744 deterministic=deterministic,745 output_attentions=output_attentions,746 output_hidden_states=output_hidden_states,747 return_dict=return_dict,748 )749 750 if not return_dict:751 return outputs752 753 return FlaxBaseModelOutput(754 last_hidden_state=outputs.last_hidden_state,755 hidden_states=outputs.hidden_states,756 attentions=outputs.attentions,757 )758 759 760class FlaxBartDecoder(nn.Module):761 config: BartConfig762 embed_tokens: nn.Embed763 dtype: jnp.dtype = jnp.float32 # the dtype of the computation764 765 def setup(self):766 self.dropout_layer = nn.Dropout(rate=self.config.dropout)767 768 embed_dim = self.config.d_model769 self.padding_idx = self.config.pad_token_id770 self.max_target_positions = self.config.max_position_embeddings771 self.embed_scale = math.sqrt(self.config.d_model) if self.config.scale_embedding else 1.0772 773 # Bart is set up so that if padding_idx is specified then offset the embedding ids by 2774 # and adjust num_embeddings appropriately. Other models don't have this hack775 self.offset = 2776 self.embed_positions = nn.Embed(777 self.config.max_position_embeddings + self.offset,778 embed_dim,779 embedding_init=jax.nn.initializers.normal(self.config.init_std),780 dtype=self.dtype,781 )782 783 self.layers = FlaxBartDecoderLayerCollection(self.config, self.dtype)784 self.layernorm_embedding = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)785 786 def __call__(787 self,788 input_ids,789 attention_mask,790 position_ids,791 encoder_hidden_states: Optional[jnp.ndarray] = None,792 encoder_attention_mask: Optional[jnp.ndarray] = None,793 init_cache: bool = False,794 output_attentions: bool = False,795 output_hidden_states: bool = False,796 return_dict: bool = True,797 deterministic: bool = True,798 ):799 input_shape = input_ids.shape800 input_ids = input_ids.reshape(-1, input_shape[-1])801 802 inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale803 804 # embed positions805 positions = self.embed_positions(position_ids + self.offset)806 807 hidden_states = inputs_embeds + positions808 hidden_states = self.layernorm_embedding(hidden_states)809 810 hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)811 812 outputs = self.layers(813 hidden_states,814 attention_mask,815 encoder_hidden_states,816 encoder_attention_mask,817 deterministic=deterministic,818 init_cache=init_cache,819 output_attentions=output_attentions,820 output_hidden_states=output_hidden_states,821 return_dict=return_dict,822 )823 824 if not return_dict:825 return outputs826 827 return FlaxBaseModelOutputWithPastAndCrossAttentions(828 last_hidden_state=outputs.last_hidden_state,829 hidden_states=outputs.hidden_states,830 attentions=outputs.attentions,831 cross_attentions=outputs.cross_attentions,832 )833 834 835class FlaxBartModule(nn.Module):836 config: BartConfig837 dtype: jnp.dtype = jnp.float32 # the dtype of the computation838 839 def setup(self):840 self.shared = nn.Embed(841 self.config.vocab_size,842 self.config.d_model,843 embedding_init=jax.nn.initializers.normal(self.config.init_std),844 dtype=self.dtype,845 )846 847 self.encoder = FlaxBartEncoder(self.config, dtype=self.dtype, embed_tokens=self.shared)848 self.decoder = FlaxBartDecoder(self.config, dtype=self.dtype, embed_tokens=self.shared)849 850 def _get_encoder_module(self):851 return self.encoder852 853 def _get_decoder_module(self):854 return self.decoder855 856 def __call__(857 self,858 input_ids,859 attention_mask,860 decoder_input_ids,861 decoder_attention_mask,862 position_ids,863 decoder_position_ids,864 output_attentions: bool = False,865 output_hidden_states: bool = False,866 return_dict: bool = True,867 deterministic: bool = True,868 ):869 encoder_outputs = self.encoder(870 input_ids=input_ids,871 attention_mask=attention_mask,872 position_ids=position_ids,873 output_attentions=output_attentions,874 output_hidden_states=output_hidden_states,875 return_dict=return_dict,876 deterministic=deterministic,877 )878 879 decoder_outputs = self.decoder(880 input_ids=decoder_input_ids,881 attention_mask=decoder_attention_mask,882 position_ids=decoder_position_ids,883 encoder_hidden_states=encoder_outputs[0],884 encoder_attention_mask=attention_mask,885 output_attentions=output_attentions,886 output_hidden_states=output_hidden_states,887 return_dict=return_dict,888 deterministic=deterministic,889 )890 891 if not return_dict:892 return decoder_outputs + encoder_outputs893 894 return FlaxSeq2SeqModelOutput(895 last_hidden_state=decoder_outputs.last_hidden_state,896 decoder_hidden_states=decoder_outputs.hidden_states,897 decoder_attentions=decoder_outputs.attentions,898 cross_attentions=decoder_outputs.cross_attentions,899 encoder_last_hidden_state=encoder_outputs.last_hidden_state,900 encoder_hidden_states=encoder_outputs.hidden_states,901 encoder_attentions=encoder_outputs.attentions,902 )903 904 905class FlaxBartPreTrainedModel(FlaxPreTrainedModel):906 config_class = BartConfig907 base_model_prefix: str = "model"908 module_class: nn.Module = None909 910 def __init__(911 self,912 config: BartConfig,913 input_shape: tuple[int] = (1, 1),914 seed: int = 0,915 dtype: jnp.dtype = jnp.float32,916 _do_init: bool = True,917 **kwargs,918 ):919 module = self.module_class(config=config, dtype=dtype, **kwargs)920 super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)921 922 def init_weights(self, rng: jax.random.PRNGKey, input_shape: tuple, params: FrozenDict = None) -> FrozenDict:923 # init input tensors924 input_ids = jnp.zeros(input_shape, dtype="i4")925 # make sure initialization pass will work for FlaxBartForSequenceClassificationModule926 input_ids = input_ids.at[(..., -1)].set(self.config.eos_token_id)927 attention_mask = jnp.ones_like(input_ids)928 decoder_input_ids = input_ids929 decoder_attention_mask = jnp.ones_like(input_ids)930 931 batch_size, sequence_length = input_ids.shape932 position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))933 decoder_position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))934 935 params_rng, dropout_rng = jax.random.split(rng)936 rngs = {"params": params_rng, "dropout": dropout_rng}937 938 random_params = self.module.init(939 rngs,940 input_ids,941 attention_mask,942 decoder_input_ids,943 decoder_attention_mask,944 position_ids,945 decoder_position_ids,946 )["params"]947 948 if params is not None:949 random_params = flatten_dict(unfreeze(random_params))950 params = flatten_dict(unfreeze(params))951 for missing_key in self._missing_keys:952 params[missing_key] = random_params[missing_key]953 self._missing_keys = set()954 return freeze(unflatten_dict(params))955 else:956 return random_params957 958 def init_cache(self, batch_size, max_length, encoder_outputs):959 r"""960 Args:961 batch_size (`int`):962 batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.963 max_length (`int`):964 maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized965 cache.966 encoder_outputs (`Union[FlaxBaseModelOutput, tuple(tuple(jnp.ndarray)]`):967 `encoder_outputs` consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*:968 `attentions`). `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*)969 is a sequence of hidden-states at the output of the last layer of the encoder. Used in the970 cross-attention of the decoder.971 """972 # init input variables to retrieve cache973 decoder_input_ids = jnp.ones((batch_size, max_length), dtype="i4")974 decoder_attention_mask = jnp.ones_like(decoder_input_ids)975 decoder_position_ids = jnp.broadcast_to(976 jnp.arange(jnp.atleast_2d(decoder_input_ids).shape[-1]), decoder_input_ids.shape977 )978 979 def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):980 decoder_module = module._get_decoder_module()981 return decoder_module(982 decoder_input_ids,983 decoder_attention_mask,984 decoder_position_ids,985 **kwargs,986 )987 988 init_variables = self.module.init(989 jax.random.PRNGKey(0),990 decoder_input_ids=decoder_input_ids,991 decoder_attention_mask=decoder_attention_mask,992 decoder_position_ids=decoder_position_ids,993 encoder_hidden_states=encoder_outputs[0],994 init_cache=True,995 method=_decoder_forward, # we only need to call the decoder to init the cache996 )997 return unfreeze(init_variables["cache"])998 999 @add_start_docstrings(BART_ENCODE_INPUTS_DOCSTRING)1000 @replace_return_docstrings(output_type=FlaxBaseModelOutput, config_class=BartConfig)1001 def encode(1002 self,1003 input_ids: jnp.ndarray,1004 attention_mask: Optional[jnp.ndarray] = None,1005 position_ids: Optional[jnp.ndarray] = None,1006 output_attentions: Optional[bool] = None,1007 output_hidden_states: Optional[bool] = None,1008 return_dict: Optional[bool] = None,1009 train: bool = False,1010 params: Optional[dict] = None,1011 dropout_rng: PRNGKey = None,1012 ):1013 r"""1014 Returns:1015 1016 Example:1017 1018 ```python1019 >>> from transformers import AutoTokenizer, FlaxBartForConditionalGeneration1020 1021 >>> model = FlaxBartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")1022 >>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")1023 1024 >>> text = "My friends are cool but they eat too many carbs."1025 >>> inputs = tokenizer(text, max_length=1024, return_tensors="jax")1026 >>> encoder_outputs = model.encode(**inputs)1027 ```"""1028 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1029 output_hidden_states = (1030 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1031 )1032 return_dict = return_dict if return_dict is not None else self.config.return_dict1033 1034 if attention_mask is None:1035 attention_mask = jnp.ones_like(input_ids)1036 if position_ids is None:1037 batch_size, sequence_length = input_ids.shape1038 position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))1039 1040 # Handle any PRNG if needed1041 rngs = {}1042 if dropout_rng is not None:1043 rngs["dropout"] = dropout_rng1044 1045 def _encoder_forward(module, input_ids, attention_mask, position_ids, **kwargs):1046 encode_module = module._get_encoder_module()1047 return encode_module(input_ids, attention_mask, position_ids, **kwargs)1048 1049 return self.module.apply(1050 {"params": params or self.params},1051 input_ids=jnp.array(input_ids, dtype="i4"),1052 attention_mask=jnp.array(attention_mask, dtype="i4"),1053 position_ids=jnp.array(position_ids, dtype="i4"),1054 output_attentions=output_attentions,1055 output_hidden_states=output_hidden_states,1056 return_dict=return_dict,1057 deterministic=not train,1058 rngs=rngs,1059 method=_encoder_forward,1060 )1061 1062 @add_start_docstrings(BART_DECODE_INPUTS_DOCSTRING)1063 @replace_return_docstrings(output_type=FlaxBaseModelOutputWithPastAndCrossAttentions, config_class=BartConfig)1064 def decode(1065 self,1066 decoder_input_ids,1067 encoder_outputs,1068 encoder_attention_mask: Optional[jnp.ndarray] = None,1069 decoder_attention_mask: Optional[jnp.ndarray] = None,1070 decoder_position_ids: Optional[jnp.ndarray] = None,1071 past_key_values: Optional[dict] = None,1072 output_attentions: Optional[bool] = None,1073 output_hidden_states: Optional[bool] = None,1074 return_dict: Optional[bool] = None,1075 train: bool = False,1076 params: Optional[dict] = None,1077 dropout_rng: PRNGKey = None,1078 ):1079 r"""1080 Returns:1081 1082 Example:1083 1084 ```python1085 >>> import jax.numpy as jnp1086 >>> from transformers import AutoTokenizer, FlaxBartForConditionalGeneration1087 1088 >>> model = FlaxBartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")1089 >>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")1090 1091 >>> text = "My friends are cool but they eat too many carbs."1092 >>> inputs = tokenizer(text, max_length=1024, return_tensors="jax")1093 >>> encoder_outputs = model.encode(**inputs)1094 1095 >>> decoder_start_token_id = model.config.decoder_start_token_id1096 >>> decoder_input_ids = jnp.ones((inputs.input_ids.shape[0], 1), dtype="i4") * decoder_start_token_id1097 1098 >>> outputs = model.decode(decoder_input_ids, encoder_outputs)1099 >>> last_decoder_hidden_states = outputs.last_hidden_state1100 ```"""1101 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1102 output_hidden_states = (1103 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1104 )1105 return_dict = return_dict if return_dict is not None else self.config.return_dict1106 1107 encoder_hidden_states = encoder_outputs[0]1108 if encoder_attention_mask is None:1109 batch_size, sequence_length = encoder_hidden_states.shape[:2]1110 encoder_attention_mask = jnp.ones((batch_size, sequence_length))1111 1112 batch_size, sequence_length = decoder_input_ids.shape1113 if decoder_attention_mask is None:1114 decoder_attention_mask = jnp.ones((batch_size, sequence_length))1115 1116 if decoder_position_ids is None:1117 if past_key_values is not None:1118 raise ValueError("Make sure to provide `decoder_position_ids` when passing `past_key_values`.")1119 1120 decoder_position_ids = jnp.broadcast_to(1121 jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)1122 )1123 1124 # Handle any PRNG if needed1125 rngs = {}1126 if dropout_rng is not None:1127 rngs["dropout"] = dropout_rng1128 1129 inputs = {"params": params or self.params}1130 1131 # if past_key_values are passed then cache is already initialized a private flag init_cache has to be1132 # passed down to ensure cache is used. It has to be made sure that cache is marked as mutable so that1133 # it can be changed by FlaxBartAttention module1134 if past_key_values:1135 inputs["cache"] = past_key_values1136 mutable = ["cache"]1137 else:1138 mutable = False1139 1140 def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):1141 decoder_module = module._get_decoder_module()1142 return decoder_module(1143 decoder_input_ids,1144 decoder_attention_mask,1145 decoder_position_ids,1146 **kwargs,1147 )1148 1149 outputs = self.module.apply(1150 inputs,1151 decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),1152 decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),1153 decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),1154 encoder_hidden_states=encoder_hidden_states,1155 encoder_attention_mask=jnp.array(encoder_attention_mask, dtype="i4"),1156 output_attentions=output_attentions,1157 output_hidden_states=output_hidden_states,1158 return_dict=return_dict,1159 deterministic=not train,1160 rngs=rngs,1161 mutable=mutable,1162 method=_decoder_forward,1163 )1164 1165 # add updated cache to model output1166 if past_key_values is not None and return_dict:1167 outputs, past = outputs1168 outputs["past_key_values"] = unfreeze(past["cache"])1169 return outputs1170 elif past_key_values is not None and not return_dict:1171 outputs, past = outputs1172 outputs = outputs[:1] + (unfreeze(past["cache"]),) + outputs[1:]1173 1174 return outputs1175 1176 @add_start_docstrings_to_model_forward(BART_INPUTS_DOCSTRING)1177 def __call__(1178 self,1179 input_ids: jnp.ndarray,1180 attention_mask: Optional[jnp.ndarray] = None,1181 decoder_input_ids: Optional[jnp.ndarray] = None,1182 decoder_attention_mask: Optional[jnp.ndarray] = None,1183 position_ids: Optional[jnp.ndarray] = None,1184 decoder_position_ids: Optional[jnp.ndarray] = None,1185 output_attentions: Optional[bool] = None,1186 output_hidden_states: Optional[bool] = None,1187 return_dict: Optional[bool] = None,1188 train: bool = False,1189 params: Optional[dict] = None,1190 dropout_rng: PRNGKey = None,1191 ):1192 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1193 output_hidden_states = (1194 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1195 )1196 return_dict = return_dict if return_dict is not None else self.config.return_dict1197 1198 # prepare encoder inputs1199 if attention_mask is None:1200 attention_mask = jnp.ones_like(input_ids)