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 Blenderbot 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)39from ...modeling_flax_utils import (40 ACT2FN,41 FlaxPreTrainedModel,42 append_call_sample_docstring,43 append_replace_return_docstrings,44 overwrite_call_docstring,45)46from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings47from .configuration_blenderbot import BlenderbotConfig48 49 50logger = logging.get_logger(__name__)51 52_CONFIG_FOR_DOC = "BlenderbotConfig"53_CHECKPOINT_FOR_DOC = "facebook/blenderbot-400M-distill"54 55 56BLENDERBOT_START_DOCSTRING = r"""57 This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the generic methods the58 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads59 etc.)60 61 This model is also a Flax Linen62 [flax.nn.Module](https://flax.readthedocs.io/en/latest/_autosummary/flax.nn.module.html) subclass. Use it as a63 regular Flax Module and refer to the Flax documentation for all matter related to general usage and behavior.64 65 Finally, this model supports inherent JAX features such as:66 67 - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)68 - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)69 - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)70 - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)71 72 Parameters:73 config ([`BlenderbotConfig`]): Model configuration class with all the parameters of the model.74 Initializing with a config file does not load the weights associated with the model, only the75 configuration. Check out the [`~FlaxPreTrainedModel.from_pretrained`] method to load the model weights.76"""77 78BLENDERBOT_INPUTS_DOCSTRING = r"""79 Args:80 input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):81 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide82 it.83 84 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and85 [`PreTrainedTokenizer.__call__`] for details.86 87 [What are input IDs?](../glossary#input-ids)88 attention_mask (`jnp.ndarray` of shape `(batch_size, sequence_length)`, *optional*):89 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:90 91 - 1 for tokens that are **not masked**,92 - 0 for tokens that are **masked**.93 94 [What are attention masks?](../glossary#attention-mask)95 decoder_input_ids (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):96 Indices of decoder input sequence tokens in the vocabulary.97 98 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and99 [`PreTrainedTokenizer.__call__`] for details.100 101 [What are decoder input IDs?](../glossary#decoder-input-ids)102 103 For translation and summarization training, `decoder_input_ids` should be provided. If no104 `decoder_input_ids` is provided, the model will create this tensor by shifting the `input_ids` to the right105 for denoising pre-training following the paper.106 decoder_attention_mask (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):107 Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also108 be used by default.109 110 If you want to change padding behavior, you should modify to your needs. See diagram 1 in [the111 paper](https://huggingface.co/papers/1910.13461) for more information on the default strategy.112 position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):113 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,114 config.max_position_embeddings - 1]`.115 decoder_position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):116 Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the117 range `[0, config.max_position_embeddings - 1]`.118 output_attentions (`bool`, *optional*):119 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned120 tensors for more detail.121 output_hidden_states (`bool`, *optional*):122 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for123 more detail.124 return_dict (`bool`, *optional*):125 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.126"""127 128 129BLENDERBOT_ENCODE_INPUTS_DOCSTRING = r"""130 Args:131 input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):132 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide133 it.134 135 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and136 [`PreTrainedTokenizer.__call__`] for details.137 138 [What are input IDs?](../glossary#input-ids)139 attention_mask (`jnp.ndarray` of shape `(batch_size, sequence_length)`, *optional*):140 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:141 142 - 1 for tokens that are **not masked**,143 - 0 for tokens that are **masked**.144 145 [What are attention masks?](../glossary#attention-mask)146 position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):147 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,148 config.max_position_embeddings - 1]`.149 output_attentions (`bool`, *optional*):150 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned151 tensors for more detail.152 output_hidden_states (`bool`, *optional*):153 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for154 more detail.155 return_dict (`bool`, *optional*):156 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.157"""158 159BLENDERBOT_DECODE_INPUTS_DOCSTRING = r"""160 Args:161 decoder_input_ids (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`):162 Indices of decoder input sequence tokens in the vocabulary.163 164 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and165 [`PreTrainedTokenizer.__call__`] for details.166 167 [What are decoder input IDs?](../glossary#decoder-input-ids)168 169 For translation and summarization training, `decoder_input_ids` should be provided. If no170 `decoder_input_ids` is provided, the model will create this tensor by shifting the `input_ids` to the right171 for denoising pre-training following the paper.172 encoder_outputs (`tuple(tuple(jnp.ndarray)`):173 Tuple consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*: `attentions`)174 `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*) is a sequence of175 hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder.176 encoder_attention_mask (`jnp.ndarray` of shape `(batch_size, sequence_length)`, *optional*):177 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:178 179 - 1 for tokens that are **not masked**,180 - 0 for tokens that are **masked**.181 182 [What are attention masks?](../glossary#attention-mask)183 decoder_attention_mask (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):184 Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also185 be used by default.186 187 If you want to change padding behavior, you should modify to your needs. See diagram 1 in [the188 paper](https://huggingface.co/papers/1910.13461) for more information on the default strategy.189 decoder_position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):190 Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the191 range `[0, config.max_position_embeddings - 1]`.192 past_key_values (`dict[str, np.ndarray]`, *optional*, returned by `init_cache` or when passing previous `past_key_values`):193 Dictionary of pre-computed hidden-states (key and values in the attention blocks) that can be used for fast194 auto-regressive decoding. Pre-computed key and value hidden-states are of shape *[batch_size, max_length]*.195 output_attentions (`bool`, *optional*):196 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned197 tensors for more detail.198 output_hidden_states (`bool`, *optional*):199 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for200 more detail.201 return_dict (`bool`, *optional*):202 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.203"""204 205 206# Copied from transformers.models.bart.modeling_flax_bart.shift_tokens_right207def shift_tokens_right(input_ids: jnp.ndarray, pad_token_id: int, decoder_start_token_id: int) -> jnp.ndarray:208 """209 Shift input ids one token to the right.210 """211 shifted_input_ids = jnp.zeros_like(input_ids)212 shifted_input_ids = shifted_input_ids.at[:, 1:].set(input_ids[:, :-1])213 shifted_input_ids = shifted_input_ids.at[:, 0].set(decoder_start_token_id)214 215 shifted_input_ids = jnp.where(shifted_input_ids == -100, pad_token_id, shifted_input_ids)216 return shifted_input_ids217 218 219# Copied from transformers.models.bart.modeling_flax_bart.FlaxBartAttention with Bart->Blenderbot220class FlaxBlenderbotAttention(nn.Module):221 config: BlenderbotConfig222 embed_dim: int223 num_heads: int224 dropout: float = 0.0225 causal: bool = False226 bias: bool = True227 dtype: jnp.dtype = jnp.float32 # the dtype of the computation228 229 def setup(self) -> None:230 self.head_dim = self.embed_dim // self.num_heads231 if self.head_dim * self.num_heads != self.embed_dim:232 raise ValueError(233 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"234 f" and `num_heads`: {self.num_heads})."235 )236 237 dense = partial(238 nn.Dense,239 self.embed_dim,240 use_bias=self.bias,241 dtype=self.dtype,242 kernel_init=jax.nn.initializers.normal(self.config.init_std),243 )244 245 self.q_proj, self.k_proj, self.v_proj = dense(), dense(), dense()246 self.out_proj = dense()247 248 self.dropout_layer = nn.Dropout(rate=self.dropout)249 250 if self.causal:251 self.causal_mask = make_causal_mask(252 jnp.ones((1, self.config.max_position_embeddings), dtype="bool"), dtype="bool"253 )254 255 def _split_heads(self, hidden_states):256 return hidden_states.reshape(hidden_states.shape[:2] + (self.num_heads, self.head_dim))257 258 def _merge_heads(self, hidden_states):259 return hidden_states.reshape(hidden_states.shape[:2] + (self.embed_dim,))260 261 @nn.compact262 def _concatenate_to_cache(self, key, value, query, attention_mask):263 """264 This function takes projected key, value states from a single input token and concatenates the states to cached265 states from previous steps. This function is slightly adapted from the official Flax repository:266 https://github.com/google/flax/blob/491ce18759622506588784b4fca0e4bf05f8c8cd/flax/linen/attention.py#L252267 """268 # detect if we're initializing by absence of existing cache data.269 is_initialized = self.has_variable("cache", "cached_key")270 cached_key = self.variable("cache", "cached_key", jnp.zeros, key.shape, key.dtype)271 cached_value = self.variable("cache", "cached_value", jnp.zeros, value.shape, value.dtype)272 cache_index = self.variable("cache", "cache_index", lambda: jnp.array(0, dtype=jnp.int32))273 274 if is_initialized:275 *batch_dims, max_length, num_heads, depth_per_head = cached_key.value.shape276 # update key, value caches with our new 1d spatial slices277 cur_index = cache_index.value278 indices = (0,) * len(batch_dims) + (cur_index, 0, 0)279 key = lax.dynamic_update_slice(cached_key.value, key, indices)280 value = lax.dynamic_update_slice(cached_value.value, value, indices)281 cached_key.value = key282 cached_value.value = value283 num_updated_cache_vectors = query.shape[1]284 cache_index.value = cache_index.value + num_updated_cache_vectors285 # 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.286 pad_mask = jnp.broadcast_to(287 jnp.arange(max_length) < cur_index + num_updated_cache_vectors,288 tuple(batch_dims) + (1, num_updated_cache_vectors, max_length),289 )290 attention_mask = combine_masks(pad_mask, attention_mask)291 return key, value, attention_mask292 293 def __call__(294 self,295 hidden_states: jnp.ndarray,296 key_value_states: Optional[jnp.ndarray] = None,297 attention_mask: Optional[jnp.ndarray] = None,298 init_cache: bool = False,299 deterministic: bool = True,300 ) -> tuple[jnp.ndarray]:301 """Input shape: Batch x Time x Channel"""302 303 # if key_value_states are provided this layer is used as a cross-attention layer304 # for the decoder305 is_cross_attention = key_value_states is not None306 batch_size = hidden_states.shape[0]307 308 # get query proj309 query_states = self.q_proj(hidden_states)310 # get key, value proj311 if is_cross_attention:312 # cross_attentions313 key_states = self.k_proj(key_value_states)314 value_states = self.v_proj(key_value_states)315 else:316 # self_attention317 key_states = self.k_proj(hidden_states)318 value_states = self.v_proj(hidden_states)319 320 query_states = self._split_heads(query_states)321 key_states = self._split_heads(key_states)322 value_states = self._split_heads(value_states)323 324 # handle cache prepare causal attention mask325 if self.causal:326 query_length, key_length = query_states.shape[1], key_states.shape[1]327 if self.has_variable("cache", "cached_key"):328 mask_shift = self.variables["cache"]["cache_index"]329 max_decoder_length = self.variables["cache"]["cached_key"].shape[1]330 causal_mask = lax.dynamic_slice(331 self.causal_mask, (0, 0, mask_shift, 0), (1, 1, query_length, max_decoder_length)332 )333 else:334 causal_mask = self.causal_mask[:, :, :query_length, :key_length]335 causal_mask = jnp.broadcast_to(causal_mask, (batch_size,) + causal_mask.shape[1:])336 337 # combine masks if needed338 if attention_mask is not None and self.causal:339 attention_mask = jnp.broadcast_to(jnp.expand_dims(attention_mask, axis=(-3, -2)), causal_mask.shape)340 attention_mask = combine_masks(attention_mask, causal_mask)341 elif self.causal:342 attention_mask = causal_mask343 elif attention_mask is not None:344 attention_mask = jnp.expand_dims(attention_mask, axis=(-3, -2))345 346 # During fast autoregressive decoding, we feed one position at a time,347 # and cache the keys and values step by step.348 if self.causal and (self.has_variable("cache", "cached_key") or init_cache):349 key_states, value_states, attention_mask = self._concatenate_to_cache(350 key_states, value_states, query_states, attention_mask351 )352 353 # Convert the boolean attention mask to an attention bias.354 if attention_mask is not None:355 # attention mask in the form of attention bias356 attention_bias = lax.select(357 attention_mask > 0,358 jnp.full(attention_mask.shape, 0.0).astype(self.dtype),359 jnp.full(attention_mask.shape, jnp.finfo(self.dtype).min).astype(self.dtype),360 )361 else:362 attention_bias = None363 364 dropout_rng = None365 if not deterministic and self.dropout > 0.0:366 dropout_rng = self.make_rng("dropout")367 368 attn_weights = dot_product_attention_weights(369 query_states,370 key_states,371 bias=attention_bias,372 dropout_rng=dropout_rng,373 dropout_rate=self.dropout,374 broadcast_dropout=True,375 deterministic=deterministic,376 dtype=self.dtype,377 precision=None,378 )379 380 attn_output = jnp.einsum("...hqk,...khd->...qhd", attn_weights, value_states)381 attn_output = self._merge_heads(attn_output)382 attn_output = self.out_proj(attn_output)383 384 return attn_output, attn_weights385 386 387# Copied from transformers.models.mbart.modeling_flax_mbart.FlaxMBartEncoderLayer with MBart->Blenderbot388class FlaxBlenderbotEncoderLayer(nn.Module):389 config: BlenderbotConfig390 dtype: jnp.dtype = jnp.float32391 392 def setup(self) -> None:393 self.embed_dim = self.config.d_model394 self.self_attn = FlaxBlenderbotAttention(395 config=self.config,396 embed_dim=self.embed_dim,397 num_heads=self.config.encoder_attention_heads,398 dropout=self.config.attention_dropout,399 dtype=self.dtype,400 )401 self.self_attn_layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)402 self.dropout_layer = nn.Dropout(rate=self.config.dropout)403 self.activation_fn = ACT2FN[self.config.activation_function]404 self.activation_dropout_layer = nn.Dropout(rate=self.config.activation_dropout)405 self.fc1 = nn.Dense(406 self.config.encoder_ffn_dim,407 dtype=self.dtype,408 kernel_init=jax.nn.initializers.normal(self.config.init_std),409 )410 self.fc2 = nn.Dense(411 self.embed_dim, dtype=self.dtype, kernel_init=jax.nn.initializers.normal(self.config.init_std)412 )413 self.final_layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)414 415 def __call__(416 self,417 hidden_states: jnp.ndarray,418 attention_mask: jnp.ndarray,419 output_attentions: bool = True,420 deterministic: bool = True,421 ) -> tuple[jnp.ndarray]:422 residual = hidden_states423 hidden_states = self.self_attn_layer_norm(hidden_states)424 hidden_states, attn_weights = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask)425 hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)426 hidden_states = residual + hidden_states427 428 residual = hidden_states429 hidden_states = self.final_layer_norm(hidden_states)430 hidden_states = self.activation_fn(self.fc1(hidden_states))431 hidden_states = self.activation_dropout_layer(hidden_states, deterministic=deterministic)432 hidden_states = self.fc2(hidden_states)433 hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)434 hidden_states = residual + hidden_states435 436 outputs = (hidden_states,)437 438 if output_attentions:439 outputs += (attn_weights,)440 441 return outputs442 443 444# Copied from transformers.models.bart.modeling_flax_bart.FlaxBartEncoderLayerCollection with Bart->Blenderbot445class FlaxBlenderbotEncoderLayerCollection(nn.Module):446 config: BlenderbotConfig447 dtype: jnp.dtype = jnp.float32 # the dtype of the computation448 449 def setup(self):450 self.layers = [451 FlaxBlenderbotEncoderLayer(self.config, name=str(i), dtype=self.dtype)452 for i in range(self.config.encoder_layers)453 ]454 self.layerdrop = self.config.encoder_layerdrop455 456 def __call__(457 self,458 hidden_states,459 attention_mask,460 deterministic: bool = True,461 output_attentions: bool = False,462 output_hidden_states: bool = False,463 return_dict: bool = True,464 ):465 all_attentions = () if output_attentions else None466 all_hidden_states = () if output_hidden_states else None467 468 for encoder_layer in self.layers:469 if output_hidden_states:470 all_hidden_states = all_hidden_states + (hidden_states,)471 # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)472 dropout_probability = random.uniform(0, 1)473 if not deterministic and (dropout_probability < self.layerdrop): # skip the layer474 layer_outputs = (None, None)475 else:476 layer_outputs = encoder_layer(477 hidden_states,478 attention_mask,479 output_attentions,480 deterministic,481 )482 hidden_states = layer_outputs[0]483 if output_attentions:484 all_attentions = all_attentions + (layer_outputs[1],)485 486 if output_hidden_states:487 all_hidden_states += (hidden_states,)488 489 outputs = (hidden_states, all_hidden_states, all_attentions)490 491 if not return_dict:492 return tuple(v for v in outputs if v is not None)493 494 return FlaxBaseModelOutput(495 last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions496 )497 498 499# Copied from transformers.models.mbart.modeling_flax_mbart.FlaxMBartDecoderLayer with MBart->Blenderbot500class FlaxBlenderbotDecoderLayer(nn.Module):501 config: BlenderbotConfig502 dtype: jnp.dtype = jnp.float32503 504 def setup(self) -> None:505 self.embed_dim = self.config.d_model506 self.self_attn = FlaxBlenderbotAttention(507 config=self.config,508 embed_dim=self.embed_dim,509 num_heads=self.config.decoder_attention_heads,510 dropout=self.config.attention_dropout,511 causal=True,512 dtype=self.dtype,513 )514 self.dropout_layer = nn.Dropout(rate=self.config.dropout)515 self.activation_fn = ACT2FN[self.config.activation_function]516 self.activation_dropout_layer = nn.Dropout(rate=self.config.activation_dropout)517 518 self.self_attn_layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)519 self.encoder_attn = FlaxBlenderbotAttention(520 config=self.config,521 embed_dim=self.embed_dim,522 num_heads=self.config.decoder_attention_heads,523 dropout=self.config.attention_dropout,524 dtype=self.dtype,525 )526 self.encoder_attn_layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)527 self.fc1 = nn.Dense(528 self.config.decoder_ffn_dim,529 dtype=self.dtype,530 kernel_init=jax.nn.initializers.normal(self.config.init_std),531 )532 self.fc2 = nn.Dense(533 self.embed_dim, dtype=self.dtype, kernel_init=jax.nn.initializers.normal(self.config.init_std)534 )535 self.final_layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)536 537 def __call__(538 self,539 hidden_states: jnp.ndarray,540 attention_mask: jnp.ndarray,541 encoder_hidden_states: Optional[jnp.ndarray] = None,542 encoder_attention_mask: Optional[jnp.ndarray] = None,543 init_cache: bool = False,544 output_attentions: bool = True,545 deterministic: bool = True,546 ) -> tuple[jnp.ndarray]:547 residual = hidden_states548 hidden_states = self.self_attn_layer_norm(hidden_states)549 550 # Self Attention551 hidden_states, self_attn_weights = self.self_attn(552 hidden_states=hidden_states, attention_mask=attention_mask, init_cache=init_cache553 )554 hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)555 hidden_states = residual + hidden_states556 557 # Cross-Attention Block558 cross_attn_weights = None559 if encoder_hidden_states is not None:560 residual = hidden_states561 562 hidden_states = self.encoder_attn_layer_norm(hidden_states)563 hidden_states, cross_attn_weights = self.encoder_attn(564 hidden_states=hidden_states,565 key_value_states=encoder_hidden_states,566 attention_mask=encoder_attention_mask,567 )568 hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)569 hidden_states = residual + hidden_states570 571 # Fully Connected572 residual = hidden_states573 hidden_states = self.final_layer_norm(hidden_states)574 hidden_states = self.activation_fn(self.fc1(hidden_states))575 hidden_states = self.activation_dropout_layer(hidden_states, deterministic=deterministic)576 hidden_states = self.fc2(hidden_states)577 hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)578 hidden_states = residual + hidden_states579 580 outputs = (hidden_states,)581 582 if output_attentions:583 outputs += (self_attn_weights, cross_attn_weights)584 585 return outputs586 587 588# Copied from transformers.models.bart.modeling_flax_bart.FlaxBartDecoderLayerCollection with Bart->Blenderbot589class FlaxBlenderbotDecoderLayerCollection(nn.Module):590 config: BlenderbotConfig591 dtype: jnp.dtype = jnp.float32 # the dtype of the computation592 593 def setup(self):594 self.layers = [595 FlaxBlenderbotDecoderLayer(self.config, name=str(i), dtype=self.dtype)596 for i in range(self.config.decoder_layers)597 ]598 self.layerdrop = self.config.decoder_layerdrop599 600 def __call__(601 self,602 hidden_states,603 attention_mask,604 encoder_hidden_states: Optional[jnp.ndarray] = None,605 encoder_attention_mask: Optional[jnp.ndarray] = None,606 deterministic: bool = True,607 init_cache: bool = False,608 output_attentions: bool = False,609 output_hidden_states: bool = False,610 return_dict: bool = True,611 ):612 # decoder layers613 all_hidden_states = () if output_hidden_states else None614 all_self_attns = () if output_attentions else None615 all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None616 617 for decoder_layer in self.layers:618 if output_hidden_states:619 all_hidden_states += (hidden_states,)620 # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)621 dropout_probability = random.uniform(0, 1)622 if not deterministic and (dropout_probability < self.layerdrop):623 layer_outputs = (None, None, None)624 else:625 layer_outputs = decoder_layer(626 hidden_states,627 attention_mask=attention_mask,628 encoder_hidden_states=encoder_hidden_states,629 encoder_attention_mask=encoder_attention_mask,630 init_cache=init_cache,631 output_attentions=output_attentions,632 deterministic=deterministic,633 )634 635 hidden_states = layer_outputs[0]636 if output_attentions:637 all_self_attns += (layer_outputs[1],)638 639 if encoder_hidden_states is not None:640 all_cross_attentions += (layer_outputs[2],)641 642 # add hidden states from the last decoder layer643 if output_hidden_states:644 all_hidden_states += (hidden_states,)645 646 outputs = [hidden_states, all_hidden_states, all_self_attns, all_cross_attentions]647 648 if not return_dict:649 return tuple(v for v in outputs if v is not None)650 651 return FlaxBaseModelOutputWithPastAndCrossAttentions(652 last_hidden_state=hidden_states,653 hidden_states=all_hidden_states,654 attentions=all_self_attns,655 cross_attentions=all_cross_attentions,656 )657 658 659class FlaxBlenderbotEncoder(nn.Module):660 config: BlenderbotConfig661 embed_tokens: nn.Embed662 dtype: jnp.dtype = jnp.float32 # the dtype of the computation663 664 def setup(self):665 self.dropout_layer = nn.Dropout(rate=self.config.dropout)666 667 embed_dim = self.config.d_model668 self.padding_idx = self.config.pad_token_id669 self.max_source_positions = self.config.max_position_embeddings670 self.embed_scale = math.sqrt(embed_dim) if self.config.scale_embedding else 1.0671 672 self.embed_positions = nn.Embed(673 self.config.max_position_embeddings,674 embed_dim,675 embedding_init=jax.nn.initializers.normal(self.config.init_std),676 )677 self.layers = FlaxBlenderbotEncoderLayerCollection(self.config, self.dtype)678 self.layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)679 680 def __call__(681 self,682 input_ids,683 attention_mask,684 position_ids,685 output_attentions: bool = False,686 output_hidden_states: bool = False,687 return_dict: bool = True,688 deterministic: bool = True,689 ):690 input_shape = input_ids.shape691 input_ids = input_ids.reshape(-1, input_shape[-1])692 693 inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale694 695 embed_pos = self.embed_positions(position_ids)696 697 hidden_states = inputs_embeds + embed_pos698 hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)699 700 outputs = self.layers(701 hidden_states,702 attention_mask,703 deterministic=deterministic,704 output_attentions=output_attentions,705 output_hidden_states=output_hidden_states,706 return_dict=return_dict,707 )708 last_hidden_states = outputs[0]709 last_hidden_states = self.layer_norm(last_hidden_states)710 711 # update the last element in `hidden_states` after applying `layernorm` above712 hidden_states = None713 if output_hidden_states:714 hidden_states = outputs[1]715 hidden_states = hidden_states[:-1] + (last_hidden_states,)716 717 if not return_dict:718 outputs = (last_hidden_states, hidden_states) + (outputs[2:] if output_hidden_states else outputs[1:])719 return tuple(v for v in outputs if v is not None)720 721 return FlaxBaseModelOutput(722 last_hidden_state=last_hidden_states,723 hidden_states=hidden_states,724 attentions=outputs.attentions,725 )726 727 728class FlaxBlenderbotDecoder(nn.Module):729 config: BlenderbotConfig730 embed_tokens: nn.Embed731 dtype: jnp.dtype = jnp.float32 # the dtype of the computation732 733 def setup(self):734 self.dropout_layer = nn.Dropout(rate=self.config.dropout)735 736 embed_dim = self.config.d_model737 self.padding_idx = self.config.pad_token_id738 self.max_target_positions = self.config.max_position_embeddings739 self.embed_scale = math.sqrt(self.config.d_model) if self.config.scale_embedding else 1.0740 741 self.embed_positions = nn.Embed(742 self.config.max_position_embeddings,743 embed_dim,744 embedding_init=jax.nn.initializers.normal(self.config.init_std),745 )746 747 self.layers = FlaxBlenderbotDecoderLayerCollection(self.config, self.dtype)748 self.layer_norm = nn.LayerNorm(dtype=self.dtype, epsilon=1e-05)749 750 def __call__(751 self,752 input_ids,753 attention_mask,754 position_ids,755 encoder_hidden_states: Optional[jnp.ndarray] = None,756 encoder_attention_mask: Optional[jnp.ndarray] = None,757 init_cache: bool = False,758 output_attentions: bool = False,759 output_hidden_states: bool = False,760 return_dict: bool = True,761 deterministic: bool = True,762 ):763 input_shape = input_ids.shape764 input_ids = input_ids.reshape(-1, input_shape[-1])765 766 inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale767 768 # embed positions769 positions = self.embed_positions(position_ids)770 771 hidden_states = inputs_embeds + positions772 hidden_states = self.dropout_layer(hidden_states, deterministic=deterministic)773 774 outputs = self.layers(775 hidden_states,776 attention_mask,777 encoder_hidden_states,778 encoder_attention_mask,779 deterministic=deterministic,780 init_cache=init_cache,781 output_attentions=output_attentions,782 output_hidden_states=output_hidden_states,783 return_dict=return_dict,784 )785 786 last_hidden_states = outputs[0]787 last_hidden_states = self.layer_norm(last_hidden_states)788 789 # update the last element in `hidden_states` after applying `layernorm` above790 hidden_states = None791 if output_hidden_states:792 hidden_states = outputs[1]793 hidden_states = hidden_states[:-1] + (last_hidden_states,)794 795 if not return_dict:796 outputs = (last_hidden_states, hidden_states) + (outputs[2:] if output_hidden_states else outputs[1:])797 return tuple(v for v in outputs if v is not None)798 799 return FlaxBaseModelOutputWithPastAndCrossAttentions(800 last_hidden_state=last_hidden_states,801 hidden_states=hidden_states,802 attentions=outputs.attentions,803 cross_attentions=outputs.cross_attentions,804 )805 806 807# Copied from transformers.models.bart.modeling_flax_bart.FlaxBartModule with Bart->Blenderbot808class FlaxBlenderbotModule(nn.Module):809 config: BlenderbotConfig810 dtype: jnp.dtype = jnp.float32 # the dtype of the computation811 812 def setup(self):813 self.shared = nn.Embed(814 self.config.vocab_size,815 self.config.d_model,816 embedding_init=jax.nn.initializers.normal(self.config.init_std),817 dtype=self.dtype,818 )819 820 self.encoder = FlaxBlenderbotEncoder(self.config, dtype=self.dtype, embed_tokens=self.shared)821 self.decoder = FlaxBlenderbotDecoder(self.config, dtype=self.dtype, embed_tokens=self.shared)822 823 def _get_encoder_module(self):824 return self.encoder825 826 def _get_decoder_module(self):827 return self.decoder828 829 def __call__(830 self,831 input_ids,832 attention_mask,833 decoder_input_ids,834 decoder_attention_mask,835 position_ids,836 decoder_position_ids,837 output_attentions: bool = False,838 output_hidden_states: bool = False,839 return_dict: bool = True,840 deterministic: bool = True,841 ):842 encoder_outputs = self.encoder(843 input_ids=input_ids,844 attention_mask=attention_mask,845 position_ids=position_ids,846 output_attentions=output_attentions,847 output_hidden_states=output_hidden_states,848 return_dict=return_dict,849 deterministic=deterministic,850 )851 852 decoder_outputs = self.decoder(853 input_ids=decoder_input_ids,854 attention_mask=decoder_attention_mask,855 position_ids=decoder_position_ids,856 encoder_hidden_states=encoder_outputs[0],857 encoder_attention_mask=attention_mask,858 output_attentions=output_attentions,859 output_hidden_states=output_hidden_states,860 return_dict=return_dict,861 deterministic=deterministic,862 )863 864 if not return_dict:865 return decoder_outputs + encoder_outputs866 867 return FlaxSeq2SeqModelOutput(868 last_hidden_state=decoder_outputs.last_hidden_state,869 decoder_hidden_states=decoder_outputs.hidden_states,870 decoder_attentions=decoder_outputs.attentions,871 cross_attentions=decoder_outputs.cross_attentions,872 encoder_last_hidden_state=encoder_outputs.last_hidden_state,873 encoder_hidden_states=encoder_outputs.hidden_states,874 encoder_attentions=encoder_outputs.attentions,875 )876 877 878class FlaxBlenderbotPreTrainedModel(FlaxPreTrainedModel):879 config_class = BlenderbotConfig880 base_model_prefix: str = "model"881 module_class: nn.Module = None882 883 def __init__(884 self,885 config: BlenderbotConfig,886 input_shape: tuple[int] = (1, 1),887 seed: int = 0,888 dtype: jnp.dtype = jnp.float32,889 _do_init: bool = True,890 **kwargs,891 ):892 module = self.module_class(config=config, dtype=dtype, **kwargs)893 super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)894 895 def init_weights(self, rng: jax.random.PRNGKey, input_shape: tuple, params: FrozenDict = None) -> FrozenDict:896 # init input tensors897 input_ids = jnp.zeros(input_shape, dtype="i4")898 # make sure initialization pass will work for FlaxBlenderbotForSequenceClassificationModule899 input_ids = input_ids.at[(..., -1)].set(self.config.eos_token_id)900 attention_mask = jnp.ones_like(input_ids)901 decoder_input_ids = input_ids902 decoder_attention_mask = jnp.ones_like(input_ids)903 904 batch_size, sequence_length = input_ids.shape905 position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))906 decoder_position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))907 908 params_rng, dropout_rng = jax.random.split(rng)909 rngs = {"params": params_rng, "dropout": dropout_rng}910 911 random_params = self.module.init(912 rngs,913 input_ids,914 attention_mask,915 decoder_input_ids,916 decoder_attention_mask,917 position_ids,918 decoder_position_ids,919 )["params"]920 921 if params is not None:922 random_params = flatten_dict(unfreeze(random_params))923 params = flatten_dict(unfreeze(params))924 for missing_key in self._missing_keys:925 params[missing_key] = random_params[missing_key]926 self._missing_keys = set()927 return freeze(unflatten_dict(params))928 else:929 return random_params930 931 def init_cache(self, batch_size, max_length, encoder_outputs):932 r"""933 Args:934 batch_size (`int`):935 batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.936 max_length (`int`):937 maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized938 cache.939 encoder_outputs (`Union[FlaxBaseModelOutput, tuple(tuple(jnp.ndarray)]`):940 `encoder_outputs` consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*:941 `attentions`). `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*)942 is a sequence of hidden-states at the output of the last layer of the encoder. Used in the943 cross-attention of the decoder.944 """945 # init input variables to retrieve cache946 decoder_input_ids = jnp.ones((batch_size, max_length), dtype="i4")947 decoder_attention_mask = jnp.ones_like(decoder_input_ids)948 decoder_position_ids = jnp.broadcast_to(949 jnp.arange(jnp.atleast_2d(decoder_input_ids).shape[-1]), decoder_input_ids.shape950 )951 952 def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):953 decoder_module = module._get_decoder_module()954 return decoder_module(955 decoder_input_ids,956 decoder_attention_mask,957 decoder_position_ids,958 **kwargs,959 )960 961 init_variables = self.module.init(962 jax.random.PRNGKey(0),963 decoder_input_ids=decoder_input_ids,964 decoder_attention_mask=decoder_attention_mask,965 decoder_position_ids=decoder_position_ids,966 encoder_hidden_states=encoder_outputs[0],967 init_cache=True,968 method=_decoder_forward, # we only need to call the decoder to init the cache969 )970 return unfreeze(init_variables["cache"])971 972 @add_start_docstrings(BLENDERBOT_ENCODE_INPUTS_DOCSTRING)973 @replace_return_docstrings(output_type=FlaxBaseModelOutput, config_class=BlenderbotConfig)974 def encode(975 self,976 input_ids: jnp.ndarray,977 attention_mask: Optional[jnp.ndarray] = None,978 position_ids: Optional[jnp.ndarray] = None,979 output_attentions: Optional[bool] = None,980 output_hidden_states: Optional[bool] = None,981 return_dict: Optional[bool] = None,982 train: bool = False,983 params: Optional[dict] = None,984 dropout_rng: PRNGKey = None,985 ):986 r"""987 Returns:988 989 Example:990 991 ```python992 >>> from transformers import AutoTokenizer, FlaxBlenderbotForConditionalGeneration993 994 >>> model = FlaxBlenderbotForConditionalGeneration.from_pretrained("facebook/blenderbot-400M-distill")995 >>> tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot-400M-distill")996 997 >>> text = "My friends are cool but they eat too many carbs."998 >>> inputs = tokenizer(text, max_length=1024, return_tensors="jax")999 >>> encoder_outputs = model.encode(**inputs)1000 ```"""1001 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1002 output_hidden_states = (1003 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1004 )1005 return_dict = return_dict if return_dict is not None else self.config.return_dict1006 1007 if attention_mask is None:1008 attention_mask = jnp.ones_like(input_ids)1009 if position_ids is None:1010 batch_size, sequence_length = input_ids.shape1011 position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))1012 1013 # Handle any PRNG if needed1014 rngs = {}1015 if dropout_rng is not None:1016 rngs["dropout"] = dropout_rng1017 1018 def _encoder_forward(module, input_ids, attention_mask, position_ids, **kwargs):1019 encode_module = module._get_encoder_module()1020 return encode_module(input_ids, attention_mask, position_ids, **kwargs)1021 1022 return self.module.apply(1023 {"params": params or self.params},1024 input_ids=jnp.array(input_ids, dtype="i4"),1025 attention_mask=jnp.array(attention_mask, dtype="i4"),1026 position_ids=jnp.array(position_ids, dtype="i4"),1027 output_attentions=output_attentions,1028 output_hidden_states=output_hidden_states,1029 return_dict=return_dict,1030 deterministic=not train,1031 rngs=rngs,1032 method=_encoder_forward,1033 )1034 1035 @add_start_docstrings(BLENDERBOT_DECODE_INPUTS_DOCSTRING)1036 @replace_return_docstrings(1037 output_type=FlaxBaseModelOutputWithPastAndCrossAttentions, config_class=BlenderbotConfig1038 )1039 def decode(1040 self,1041 decoder_input_ids,1042 encoder_outputs,1043 encoder_attention_mask: Optional[jnp.ndarray] = None,1044 decoder_attention_mask: Optional[jnp.ndarray] = None,1045 decoder_position_ids: Optional[jnp.ndarray] = None,1046 past_key_values: Optional[dict] = None,1047 output_attentions: Optional[bool] = None,1048 output_hidden_states: Optional[bool] = None,1049 return_dict: Optional[bool] = None,1050 train: bool = False,1051 params: Optional[dict] = None,1052 dropout_rng: PRNGKey = None,1053 ):1054 r"""1055 Returns:1056 1057 Example:1058 1059 ```python1060 >>> import jax.numpy as jnp1061 >>> from transformers import AutoTokenizer, FlaxBlenderbotForConditionalGeneration1062 1063 >>> model = FlaxBlenderbotForConditionalGeneration.from_pretrained("facebook/blenderbot-400M-distill")1064 >>> tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot-400M-distill")1065 1066 >>> text = "My friends are cool but they eat too many carbs."1067 >>> inputs = tokenizer(text, max_length=1024, return_tensors="jax")1068 >>> encoder_outputs = model.encode(**inputs)1069 1070 >>> decoder_start_token_id = model.config.decoder_start_token_id1071 >>> decoder_input_ids = jnp.ones((inputs.input_ids.shape[0], 1), dtype="i4") * decoder_start_token_id1072 1073 >>> outputs = model.decode(decoder_input_ids, encoder_outputs)1074 >>> last_decoder_hidden_states = outputs.last_hidden_state1075 ```"""1076 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1077 output_hidden_states = (1078 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1079 )1080 return_dict = return_dict if return_dict is not None else self.config.return_dict1081 1082 encoder_hidden_states = encoder_outputs[0]1083 if encoder_attention_mask is None:1084 batch_size, sequence_length = encoder_hidden_states.shape[:2]1085 encoder_attention_mask = jnp.ones((batch_size, sequence_length))1086 1087 batch_size, sequence_length = decoder_input_ids.shape1088 if decoder_attention_mask is None:1089 decoder_attention_mask = jnp.ones((batch_size, sequence_length))1090 1091 if decoder_position_ids is None:1092 if past_key_values is not None:1093 raise ValueError("Make sure to provide `decoder_position_ids` when passing `past_key_values`.")1094 1095 decoder_position_ids = jnp.broadcast_to(1096 jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)1097 )1098 1099 # Handle any PRNG if needed1100 rngs = {}1101 if dropout_rng is not None:1102 rngs["dropout"] = dropout_rng1103 1104 inputs = {"params": params or self.params}1105 1106 # if past_key_values are passed then cache is already initialized a private flag init_cache has to be1107 # passed down to ensure cache is used. It has to be made sure that cache is marked as mutable so that1108 # it can be changed by FlaxBlenderbotAttention module1109 if past_key_values:1110 inputs["cache"] = past_key_values1111 mutable = ["cache"]1112 else:1113 mutable = False1114 1115 def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):1116 decoder_module = module._get_decoder_module()1117 return decoder_module(1118 decoder_input_ids,1119 decoder_attention_mask,1120 decoder_position_ids,1121 **kwargs,1122 )1123 1124 outputs = self.module.apply(1125 inputs,1126 decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),1127 decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),1128 decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),1129 encoder_hidden_states=encoder_hidden_states,1130 encoder_attention_mask=jnp.array(encoder_attention_mask, dtype="i4"),1131 output_attentions=output_attentions,1132 output_hidden_states=output_hidden_states,1133 return_dict=return_dict,1134 deterministic=not train,1135 rngs=rngs,1136 mutable=mutable,1137 method=_decoder_forward,1138 )1139 1140 # add updated cache to model output1141 if past_key_values is not None and return_dict:1142 outputs, past = outputs1143 outputs["past_key_values"] = unfreeze(past["cache"])1144 return outputs1145 elif past_key_values is not None and not return_dict:1146 outputs, past = outputs1147 outputs = outputs[:1] + (unfreeze(past["cache"]),) + outputs[1:]1148 1149 return outputs1150 1151 @add_start_docstrings_to_model_forward(BLENDERBOT_INPUTS_DOCSTRING)1152 def __call__(1153 self,1154 input_ids: jnp.ndarray,1155 attention_mask: Optional[jnp.ndarray] = None,1156 decoder_input_ids: Optional[jnp.ndarray] = None,1157 decoder_attention_mask: Optional[jnp.ndarray] = None,1158 position_ids: Optional[jnp.ndarray] = None,1159 decoder_position_ids: Optional[jnp.ndarray] = None,1160 output_attentions: Optional[bool] = None,1161 output_hidden_states: Optional[bool] = None,1162 return_dict: Optional[bool] = None,1163 train: bool = False,1164 params: Optional[dict] = None,1165 dropout_rng: PRNGKey = None,1166 ):1167 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1168 output_hidden_states = (1169 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1170 )1171 return_dict = return_dict if return_dict is not None else self.config.return_dict1172 1173 # prepare encoder inputs1174 if attention_mask is None:1175 attention_mask = jnp.ones_like(input_ids)1176 if position_ids is None:1177 batch_size, sequence_length = input_ids.shape1178 position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))1179 1180 # prepare decoder inputs1181 if decoder_input_ids is None:1182 decoder_input_ids = shift_tokens_right(1183 input_ids, self.config.pad_token_id, decoder_start_token_id=self.config.decoder_start_token_id1184 )1185 if decoder_attention_mask is None:1186 decoder_attention_mask = jnp.ones_like(decoder_input_ids)1187 if decoder_position_ids is None:1188 batch_size, sequence_length = decoder_input_ids.shape1189 decoder_position_ids = jnp.broadcast_to(1190 jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)1191 )1192 1193 # Handle any PRNG if needed1194 rngs = {"dropout": dropout_rng} if dropout_rng is not None else {}1195 1196 return self.module.apply(1197 {"params": params or self.params},1198 input_ids=jnp.array(input_ids, dtype="i4"),1199 attention_mask=jnp.array(attention_mask, dtype="i4"),1200 position_ids=jnp.array(position_ids, dtype="i4"),