Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2021 The HuggingFace Inc. team.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"""Classes to support Flax Encoder-Decoder architectures"""16 17import os18from typing import Optional, Union19 20import flax.linen as nn21import jax22import jax.numpy as jnp23from flax.core.frozen_dict import FrozenDict, freeze, unfreeze24from flax.traverse_util import flatten_dict, unflatten_dict25from jax import lax26from jax.random import PRNGKey27 28from ...modeling_flax_outputs import FlaxBaseModelOutput, FlaxCausalLMOutputWithCrossAttentions, FlaxSeq2SeqLMOutput29from ...modeling_flax_utils import FlaxPreTrainedModel30from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings31from ..auto.configuration_auto import AutoConfig32from ..auto.modeling_flax_auto import FlaxAutoModel, FlaxAutoModelForCausalLM33from .configuration_encoder_decoder import EncoderDecoderConfig34 35 36logger = logging.get_logger(__name__)37 38_CONFIG_FOR_DOC = "EncoderDecoderConfig"39 40ENCODER_DECODER_START_DOCSTRING = r"""41 This class can be used to initialize a sequence-to-sequence model with any pretrained autoencoding model as the42 encoder and any pretrained autoregressive model as the decoder. The encoder is loaded via43 [`~AutoModel.from_pretrained`] function and the decoder is loaded via [`~AutoModelForCausalLM.from_pretrained`]44 function. Cross-attention layers are automatically added to the decoder and should be fine-tuned on a downstream45 generative task, like summarization.46 47 The effectiveness of initializing sequence-to-sequence models with pretrained checkpoints for sequence generation48 tasks was shown in [Leveraging Pre-trained Checkpoints for Sequence Generation49 Tasks](https://huggingface.co/papers/1907.12461) by Sascha Rothe, Shashi Narayan, Aliaksei Severyn. Michael Matena, Yanqi50 Zhou, Wei Li, Peter J. Liu.51 52 After such an Encoder Decoder model has been trained/fine-tuned, it can be saved/loaded just like any other models53 (see the examples for more information).54 55 This model inherits from [`FlaxPreTrainedModel`]. Check the superclass documentation for the generic methods the56 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads57 etc.)58 59 This model is also a Flax Linen60 [flax.nn.Module](https://flax.readthedocs.io/en/latest/_autosummary/flax.nn.module.html) subclass. Use it as a61 regular Flax Module and refer to the Flax documentation for all matter related to general usage and behavior.62 63 Parameters:64 config ([`EncoderDecoderConfig`]): Model configuration class with all the parameters of the model.65 Initializing with a config file does not load the weights associated with the model, only the66 configuration. Check out the [`~FlaxPreTrainedModel.from_pretrained`] method to load the model weights.67 dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):68 The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on GPUs) and69 `jax.numpy.bfloat16` (on TPUs).70 71 This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If72 specified all the computation will be performed with the given `dtype`.73 74 **Note that this only specifies the dtype of the computation and does not influence the dtype of model75 parameters.**76 77 If you wish to change the dtype of the model parameters, see [`~FlaxPreTrainedModel.to_fp16`] and78 [`~FlaxPreTrainedModel.to_bf16`].79"""80 81ENCODER_DECODER_INPUTS_DOCSTRING = r"""82 Args:83 input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):84 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide85 it.86 87 Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and88 [`PreTrainedTokenizer.__call__`] for details.89 90 [What are input IDs?](../glossary#input-ids)91 attention_mask (`jnp.ndarray` of shape `(batch_size, sequence_length)`, *optional*):92 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:93 94 - 1 for tokens that are **not masked**,95 - 0 for tokens that are **masked**.96 97 [What are attention masks?](../glossary#attention-mask)98 decoder_input_ids (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):99 Indices of decoder input sequence tokens in the vocabulary.100 101 Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and102 [`PreTrainedTokenizer.__call__`] for details.103 104 [What are decoder input IDs?](../glossary#decoder-input-ids)105 106 For sequence to sequence training, `decoder_input_ids` should be provided. `decoder_input_ids` should be107 created outside of the model by shifting the `labels` to the right, replacing -100 by the `pad_token_id`108 and prepending them with the `decoder_start_token_id`.109 decoder_attention_mask (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):110 Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also111 be used by default.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.encoder.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.decoder.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 If set to `True`, the model will return a [`~utils.FlaxSeq2SeqLMOutput`] instead of a plain tuple.126"""127 128ENCODER_DECODER_ENCODE_INPUTS_DOCSTRING = r"""129 Args:130 input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):131 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide132 it.133 134 Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and135 [`PreTrainedTokenizer.__call__`] for details.136 137 [What are input IDs?](../glossary#input-ids)138 attention_mask (`jnp.ndarray` of shape `(batch_size, sequence_length)`, *optional*):139 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:140 141 - 1 for tokens that are **not masked**,142 - 0 for tokens that are **masked**.143 144 [What are attention masks?](../glossary#attention-mask)145 position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):146 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,147 config.encoder.max_position_embeddings - 1]`.148 output_attentions (`bool`, *optional*):149 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned150 tensors for more detail.151 output_hidden_states (`bool`, *optional*):152 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for153 more detail.154 return_dict (`bool`, *optional*):155 If set to `True`, the model will return a [`~utils.FlaxBaseModelOutput`] instead of a plain tuple.156"""157 158ENCODER_DECODER_DECODE_INPUTS_DOCSTRING = r"""159 Args:160 decoder_input_ids (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):161 Indices of decoder input sequence tokens in the vocabulary.162 163 Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and164 [`PreTrainedTokenizer.__call__`] for details.165 166 [What are decoder input IDs?](../glossary#decoder-input-ids)167 168 If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see169 `past_key_values`).170 171 For sequence to sequence training, `decoder_input_ids` should be provided. `decoder_input_ids` should be172 created outside of the model by shifting the `labels` to the right, replacing -100 by the `pad_token_id`173 and prepending them with the `decoder_start_token_id`.174 encoder_outputs (`tuple(tuple(jnp.ndarray)`):175 Tuple consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*: `attentions`)176 `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*) is a sequence of177 hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder.178 encoder_attention_mask (`jnp.ndarray` of shape `(batch_size, sequence_length)`, *optional*):179 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:180 181 - 1 for tokens that are **not masked**,182 - 0 for tokens that are **masked**.183 184 [What are attention masks?](../glossary#attention-mask)185 decoder_attention_mask (`jnp.ndarray` of shape `(batch_size, target_sequence_length)`, *optional*):186 Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also187 be used by default.188 decoder_position_ids (`numpy.ndarray` of shape `(batch_size, sequence_length)`, *optional*):189 Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the190 range `[0, config.decoder.max_position_embeddings - 1]`.191 past_key_values (`dict[str, np.ndarray]`, *optional*, returned by `init_cache` or when passing previous `past_key_values`):192 Dictionary of pre-computed hidden-states (key and values in the attention blocks) that can be used for fast193 auto-regressive decoding. Pre-computed key and value hidden-states are of shape *[batch_size, max_length]*.194 output_attentions (`bool`, *optional*):195 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned196 tensors for more detail.197 output_hidden_states (`bool`, *optional*):198 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for199 more detail.200 return_dict (`bool`, *optional*):201 If set to `True`, the model will return a [`~utils.FlaxCausalLMOutputWithCrossAttentions`] instead of a202 plain tuple.203"""204 205 206class FlaxEncoderDecoderModule(nn.Module):207 config: EncoderDecoderConfig208 dtype: jnp.dtype = jnp.float32209 210 def setup(self):211 encoder_config = self.config.encoder212 decoder_config = self.config.decoder213 214 # Copied from `modeling_hybrid_clip.py` with modifications.215 from ...models.auto.modeling_flax_auto import FLAX_MODEL_FOR_CAUSAL_LM_MAPPING, FLAX_MODEL_MAPPING216 217 encoder_module = FLAX_MODEL_MAPPING[encoder_config.__class__].module_class218 decoder_module = FLAX_MODEL_FOR_CAUSAL_LM_MAPPING[decoder_config.__class__].module_class219 220 self.encoder = encoder_module(encoder_config, dtype=self.dtype)221 self.decoder = decoder_module(decoder_config, dtype=self.dtype)222 223 # encoder outputs might need to be projected to different dimension for decoder224 if (225 self.encoder.config.hidden_size != self.decoder.config.hidden_size226 and self.decoder.config.cross_attention_hidden_size is None227 ):228 self.enc_to_dec_proj = nn.Dense(229 self.decoder.config.hidden_size,230 kernel_init=jax.nn.initializers.normal(self.decoder.config.initializer_range),231 dtype=self.dtype,232 )233 else:234 self.enc_to_dec_proj = None235 236 def _get_encoder_module(self):237 return self.encoder238 239 def _get_projection_module(self):240 return self.enc_to_dec_proj241 242 def _get_decoder_module(self):243 return self.decoder244 245 def __call__(246 self,247 input_ids,248 attention_mask,249 decoder_input_ids,250 decoder_attention_mask,251 position_ids,252 decoder_position_ids,253 output_attentions: bool = False,254 output_hidden_states: bool = False,255 return_dict: bool = True,256 deterministic: bool = True,257 ):258 encoder_outputs = self.encoder(259 input_ids=input_ids,260 attention_mask=attention_mask,261 position_ids=position_ids,262 output_attentions=output_attentions,263 output_hidden_states=output_hidden_states,264 return_dict=return_dict,265 deterministic=deterministic,266 )267 268 encoder_hidden_states = encoder_outputs[0]269 270 # optionally project encoder_hidden_states271 if self.enc_to_dec_proj is not None:272 encoder_hidden_states = self.enc_to_dec_proj(encoder_hidden_states)273 274 decoder_outputs = self.decoder(275 input_ids=decoder_input_ids,276 attention_mask=decoder_attention_mask,277 position_ids=decoder_position_ids,278 encoder_hidden_states=encoder_hidden_states,279 encoder_attention_mask=attention_mask,280 output_attentions=output_attentions,281 output_hidden_states=output_hidden_states,282 return_dict=return_dict,283 deterministic=deterministic,284 )285 286 if not return_dict:287 return decoder_outputs + encoder_outputs288 289 return FlaxSeq2SeqLMOutput(290 logits=decoder_outputs.logits,291 decoder_hidden_states=decoder_outputs.hidden_states,292 decoder_attentions=decoder_outputs.attentions,293 cross_attentions=decoder_outputs.cross_attentions,294 encoder_last_hidden_state=encoder_outputs.last_hidden_state,295 encoder_hidden_states=encoder_outputs.hidden_states,296 encoder_attentions=encoder_outputs.attentions,297 )298 299 300@add_start_docstrings(ENCODER_DECODER_START_DOCSTRING)301class FlaxEncoderDecoderModel(FlaxPreTrainedModel):302 r"""303 [`FlaxEncoderDecoderModel`] is a generic model class that will be instantiated as a transformer architecture with304 the module (flax.nn.Module) of one of the base model classes of the library as encoder module and another one as305 decoder module when created with the :meth*~transformers.FlaxAutoModel.from_pretrained* class method for the306 encoder and :meth*~transformers.FlaxAutoModelForCausalLM.from_pretrained* class method for the decoder.307 """308 309 config_class = EncoderDecoderConfig310 base_model_prefix = "encoder_decoder"311 module_class = FlaxEncoderDecoderModule312 313 def __init__(314 self,315 config: EncoderDecoderConfig,316 input_shape: Optional[tuple] = None,317 seed: int = 0,318 dtype: jnp.dtype = jnp.float32,319 _do_init: bool = True,320 **kwargs,321 ):322 if input_shape is None:323 input_shape = ((1, 1), (1, 1))324 325 if not _do_init:326 raise ValueError(327 "`FlaxEncoderDecoderModel` cannot be created without initializing, `_do_init` must be `True`."328 )329 330 if config.decoder.cross_attention_hidden_size is not None:331 if config.decoder.cross_attention_hidden_size != config.encoder.hidden_size:332 raise ValueError(333 "If `cross_attention_hidden_size` is specified in the decoder's configuration, it has to be equal"334 f" to the encoder's `hidden_size`. Got {config.decoder.cross_attention_hidden_size} for"335 f" `config.decoder.cross_attention_hidden_size` and {config.encoder.hidden_size} for"336 " `config.encoder.hidden_size`."337 )338 339 module = self.module_class(config=config, dtype=dtype, **kwargs)340 super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)341 342 def init_weights(self, rng: jax.random.PRNGKey, input_shape: tuple, params: FrozenDict = None) -> FrozenDict:343 encoder_input_shape, decoder_input_shape = input_shape344 345 # init input tensors346 input_ids = jnp.zeros(encoder_input_shape, dtype="i4")347 attention_mask = jnp.ones_like(input_ids)348 decoder_input_ids = jnp.zeros(decoder_input_shape, dtype="i4")349 decoder_attention_mask = jnp.ones_like(decoder_input_ids)350 351 batch_size, sequence_length = input_ids.shape352 position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))353 354 decoder_batch_size, decoder_sequence_length = decoder_input_ids.shape355 if not decoder_batch_size == batch_size:356 raise ValueError(357 f"The inputs of encoder and decoder should have the same batch size, but got {batch_size} for encoder"358 f" and {decoder_batch_size} for decoder."359 )360 decoder_position_ids = jnp.broadcast_to(361 jnp.arange(decoder_sequence_length)[None, :], (decoder_batch_size, decoder_sequence_length)362 )363 364 params_rng, dropout_rng = jax.random.split(rng)365 rngs = {"params": params_rng, "dropout": dropout_rng}366 367 random_params = self.module.init(368 rngs,369 input_ids,370 attention_mask,371 decoder_input_ids,372 decoder_attention_mask,373 position_ids,374 decoder_position_ids,375 )["params"]376 377 if params is not None:378 random_params = flatten_dict(unfreeze(random_params))379 params = flatten_dict(unfreeze(params))380 for missing_key in self._missing_keys:381 params[missing_key] = random_params[missing_key]382 self._missing_keys = set()383 return freeze(unflatten_dict(params))384 else:385 return random_params386 387 def init_cache(self, batch_size, max_length, encoder_outputs):388 r"""389 Args:390 batch_size (`int`):391 batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.392 max_length (`int`):393 maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized394 cache.395 encoder_outputs (`Union[FlaxBaseModelOutput, tuple(tuple(jnp.ndarray)]`):396 `encoder_outputs` consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*:397 `attentions`). `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*)398 is a sequence of hidden-states at the output of the last layer of the encoder. Used in the399 cross-attention of the decoder.400 """401 # init input variables to retrieve cache402 decoder_input_ids = jnp.ones((batch_size, max_length), dtype="i4")403 decoder_attention_mask = jnp.ones_like(decoder_input_ids)404 decoder_position_ids = jnp.broadcast_to(405 jnp.arange(jnp.atleast_2d(decoder_input_ids).shape[-1]), decoder_input_ids.shape406 )407 408 def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):409 decoder_module = module._get_decoder_module()410 return decoder_module(411 input_ids=decoder_input_ids,412 attention_mask=decoder_attention_mask,413 position_ids=decoder_position_ids,414 **kwargs,415 )416 417 init_variables = self.module.init(418 jax.random.PRNGKey(0),419 decoder_input_ids=decoder_input_ids,420 decoder_attention_mask=decoder_attention_mask,421 decoder_position_ids=decoder_position_ids,422 encoder_hidden_states=encoder_outputs[0],423 init_cache=True,424 method=_decoder_forward, # we only need to call the decoder to init the cache425 )426 return unfreeze(init_variables["cache"])427 428 @add_start_docstrings(ENCODER_DECODER_ENCODE_INPUTS_DOCSTRING)429 @replace_return_docstrings(output_type=FlaxBaseModelOutput, config_class=_CONFIG_FOR_DOC)430 def encode(431 self,432 input_ids: jnp.ndarray,433 attention_mask: Optional[jnp.ndarray] = None,434 position_ids: Optional[jnp.ndarray] = None,435 output_attentions: Optional[bool] = None,436 output_hidden_states: Optional[bool] = None,437 return_dict: Optional[bool] = None,438 train: bool = False,439 params: Optional[dict] = None,440 dropout_rng: PRNGKey = None,441 ):442 r"""443 Returns:444 445 Example:446 447 ```python448 >>> from transformers import FlaxEncoderDecoderModel, BertTokenizer449 450 >>> # initialize a bert2gpt2 from pretrained BERT and GPT2 models. Note that the cross-attention layers will be randomly initialized451 >>> model = FlaxEncoderDecoderModel.from_encoder_decoder_pretrained("google-bert/bert-base-cased", "openai-community/gpt2")452 453 >>> tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-cased")454 455 >>> text = "My friends are cool but they eat too many carbs."456 >>> input_ids = tokenizer.encode(text, return_tensors="np")457 >>> encoder_outputs = model.encode(input_ids)458 ```"""459 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions460 output_hidden_states = (461 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states462 )463 return_dict = return_dict if return_dict is not None else self.config.return_dict464 465 if attention_mask is None:466 attention_mask = jnp.ones_like(input_ids)467 if position_ids is None:468 batch_size, sequence_length = input_ids.shape469 position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))470 471 # Handle any PRNG if needed472 rngs = {}473 if dropout_rng is not None:474 rngs["dropout"] = dropout_rng475 476 def _encoder_forward(module, input_ids, attention_mask, position_ids, **kwargs):477 encode_module = module._get_encoder_module()478 return encode_module(input_ids, attention_mask, position_ids, **kwargs)479 480 outputs = self.module.apply(481 {"params": params or self.params},482 input_ids=jnp.array(input_ids, dtype="i4"),483 attention_mask=jnp.array(attention_mask, dtype="i4"),484 position_ids=jnp.array(position_ids, dtype="i4"),485 output_attentions=output_attentions,486 output_hidden_states=output_hidden_states,487 return_dict=return_dict,488 deterministic=not train,489 rngs=rngs,490 method=_encoder_forward,491 )492 493 if return_dict:494 outputs = FlaxBaseModelOutput(495 last_hidden_state=outputs.last_hidden_state,496 hidden_states=outputs.hidden_states,497 attentions=outputs.attentions,498 )499 500 return outputs501 502 @add_start_docstrings(ENCODER_DECODER_DECODE_INPUTS_DOCSTRING)503 @replace_return_docstrings(output_type=FlaxCausalLMOutputWithCrossAttentions, config_class=_CONFIG_FOR_DOC)504 def decode(505 self,506 decoder_input_ids,507 encoder_outputs,508 encoder_attention_mask: Optional[jnp.ndarray] = None,509 decoder_attention_mask: Optional[jnp.ndarray] = None,510 decoder_position_ids: Optional[jnp.ndarray] = None,511 past_key_values: Optional[dict] = None,512 output_attentions: Optional[bool] = None,513 output_hidden_states: Optional[bool] = None,514 return_dict: Optional[bool] = None,515 train: bool = False,516 params: Optional[dict] = None,517 dropout_rng: PRNGKey = None,518 ):519 r"""520 Returns:521 522 Example:523 524 ```python525 >>> from transformers import FlaxEncoderDecoderModel, BertTokenizer526 >>> import jax.numpy as jnp527 528 >>> # initialize a bert2gpt2 from pretrained BERT and GPT2 models. Note that the cross-attention layers will be randomly initialized529 >>> model = FlaxEncoderDecoderModel.from_encoder_decoder_pretrained("google-bert/bert-base-cased", "openai-community/gpt2")530 531 >>> tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-cased")532 533 >>> text = "My friends are cool but they eat too many carbs."534 >>> input_ids = tokenizer.encode(text, max_length=1024, return_tensors="np")535 >>> encoder_outputs = model.encode(input_ids)536 537 >>> decoder_start_token_id = model.config.decoder.bos_token_id538 >>> decoder_input_ids = jnp.ones((input_ids.shape[0], 1), dtype="i4") * decoder_start_token_id539 540 >>> outputs = model.decode(decoder_input_ids, encoder_outputs)541 >>> logits = outputs.logits542 ```"""543 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions544 output_hidden_states = (545 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states546 )547 return_dict = return_dict if return_dict is not None else self.config.return_dict548 549 encoder_hidden_states = encoder_outputs[0]550 if encoder_attention_mask is None:551 batch_size, sequence_length = encoder_hidden_states.shape[:2]552 encoder_attention_mask = jnp.ones((batch_size, sequence_length))553 554 batch_size, sequence_length = decoder_input_ids.shape555 if decoder_attention_mask is None:556 decoder_attention_mask = jnp.ones((batch_size, sequence_length))557 558 if decoder_position_ids is None:559 if past_key_values is not None:560 raise ValueError("Make sure to provide `decoder_position_ids` when passing `past_key_values`.")561 562 decoder_position_ids = jnp.broadcast_to(563 jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)564 )565 566 # Handle any PRNG if needed567 rngs = {}568 if dropout_rng is not None:569 rngs["dropout"] = dropout_rng570 571 inputs = {"params": params or self.params}572 573 # if past_key_values are passed then cache is already initialized a private flag init_cache has to be574 # passed down to ensure cache is used. It has to be made sure that cache is marked as mutable so that575 # it can be changed by FlaxBartAttention module576 if past_key_values:577 inputs["cache"] = past_key_values578 mutable = ["cache"]579 else:580 mutable = False581 582 def _decoder_forward(583 module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, encoder_hidden_states, **kwargs584 ):585 projection_module = module._get_projection_module()586 decoder_module = module._get_decoder_module()587 588 # optionally project encoder_hidden_states589 if projection_module is not None:590 encoder_hidden_states = projection_module(encoder_hidden_states)591 592 return decoder_module(593 decoder_input_ids,594 decoder_attention_mask,595 decoder_position_ids,596 encoder_hidden_states=encoder_hidden_states,597 **kwargs,598 )599 600 outputs = self.module.apply(601 inputs,602 decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),603 decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),604 decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),605 encoder_hidden_states=encoder_hidden_states,606 encoder_attention_mask=jnp.array(encoder_attention_mask, dtype="i4"),607 output_attentions=output_attentions,608 output_hidden_states=output_hidden_states,609 return_dict=return_dict,610 deterministic=not train,611 rngs=rngs,612 mutable=mutable,613 method=_decoder_forward,614 )615 616 # add updated cache to model output617 if past_key_values is not None and return_dict:618 outputs, past = outputs619 outputs["past_key_values"] = unfreeze(past["cache"])620 return outputs621 elif past_key_values is not None and not return_dict:622 outputs, past = outputs623 outputs = outputs[:1] + (unfreeze(past["cache"]),) + outputs[1:]624 625 return outputs626 627 @add_start_docstrings_to_model_forward(ENCODER_DECODER_INPUTS_DOCSTRING)628 @replace_return_docstrings(output_type=FlaxSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)629 def __call__(630 self,631 input_ids: jnp.ndarray,632 attention_mask: Optional[jnp.ndarray] = None,633 decoder_input_ids: Optional[jnp.ndarray] = None,634 decoder_attention_mask: Optional[jnp.ndarray] = None,635 position_ids: Optional[jnp.ndarray] = None,636 decoder_position_ids: Optional[jnp.ndarray] = None,637 output_attentions: Optional[bool] = None,638 output_hidden_states: Optional[bool] = None,639 return_dict: Optional[bool] = None,640 train: bool = False,641 params: Optional[dict] = None,642 dropout_rng: PRNGKey = None,643 ):644 r"""645 Returns:646 647 Examples:648 649 ```python650 >>> from transformers import FlaxEncoderDecoderModel, BertTokenizer, GPT2Tokenizer651 652 >>> # load a fine-tuned bert2gpt2 model653 >>> model = FlaxEncoderDecoderModel.from_pretrained("patrickvonplaten/bert2gpt2-cnn_dailymail-fp16")654 >>> # load input & output tokenizer655 >>> tokenizer_input = BertTokenizer.from_pretrained("google-bert/bert-base-cased")656 >>> tokenizer_output = GPT2Tokenizer.from_pretrained("openai-community/gpt2")657 658 >>> article = '''Sigma Alpha Epsilon is under fire for a video showing party-bound fraternity members659 >>> singing a racist chant. SAE's national chapter suspended the students,660 >>> but University of Oklahoma President David Boren took it a step further,661 >>> saying the university's affiliation with the fraternity is permanently done.'''662 663 >>> input_ids = tokenizer_input(article, add_special_tokens=True, return_tensors="np").input_ids664 665 >>> # use GPT2's eos_token as the pad as well as eos token666 >>> model.config.eos_token_id = model.config.decoder.eos_token_id667 >>> model.config.pad_token_id = model.config.eos_token_id668 669 >>> sequences = model.generate(input_ids, num_beams=4, max_length=12).sequences670 671 >>> summary = tokenizer_output.batch_decode(sequences, skip_special_tokens=True)[0]672 >>> assert summary == "SAS Alpha Epsilon suspended Sigma Alpha Epsilon members"673 ```674 """675 676 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions677 output_hidden_states = (678 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states679 )680 return_dict = return_dict if return_dict is not None else self.config.return_dict681 682 # prepare encoder inputs683 if attention_mask is None:684 attention_mask = jnp.ones_like(input_ids)685 if position_ids is None:686 batch_size, sequence_length = input_ids.shape687 position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))688 689 # prepare decoder inputs690 if decoder_input_ids is None:691 raise ValueError(692 "`decoder_input_ids` cannot be `None`. For sequence to sequence training, `decoder_position_ids` must"693 " be specified as an input argument."694 )695 if decoder_attention_mask is None:696 decoder_attention_mask = jnp.ones_like(decoder_input_ids)697 if decoder_position_ids is None:698 batch_size, sequence_length = decoder_input_ids.shape699 decoder_position_ids = jnp.broadcast_to(700 jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)701 )702 703 # Handle any PRNG if needed704 rngs = {"dropout": dropout_rng} if dropout_rng is not None else {}705 706 return self.module.apply(707 {"params": params or self.params},708 input_ids=jnp.array(input_ids, dtype="i4"),709 attention_mask=jnp.array(attention_mask, dtype="i4"),710 decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),711 decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),712 position_ids=jnp.array(position_ids, dtype="i4"),713 decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),714 output_attentions=output_attentions,715 output_hidden_states=output_hidden_states,716 return_dict=return_dict,717 deterministic=not train,718 rngs=rngs,719 )720 721 def prepare_inputs_for_generation(722 self,723 decoder_input_ids,724 max_length,725 attention_mask: Optional[jax.Array] = None,726 decoder_attention_mask: Optional[jax.Array] = None,727 encoder_outputs=None,728 **kwargs,729 ):730 # initializing the cache731 batch_size, seq_length = decoder_input_ids.shape732 733 past_key_values = self.init_cache(batch_size, max_length, encoder_outputs)734 # Note that usually one would have to put 0's in the attention_mask for x > input_ids.shape[-1] and x < cache_length.735 # But since the decoder uses a causal mask, those positions are masked anyways.736 # Thus we can create a single static attention_mask here, which is more efficient for compilation737 extended_attention_mask = jnp.ones((batch_size, max_length), dtype="i4")738 if decoder_attention_mask is not None:739 decoder_position_ids = decoder_attention_mask.cumsum(axis=-1) - 1740 extended_attention_mask = lax.dynamic_update_slice(extended_attention_mask, decoder_attention_mask, (0, 0))741 else:742 decoder_position_ids = jnp.broadcast_to(743 jnp.arange(seq_length, dtype="i4")[None, :], (batch_size, seq_length)744 )745 746 return {747 "past_key_values": past_key_values,748 "encoder_outputs": encoder_outputs,749 "encoder_attention_mask": attention_mask,750 "decoder_attention_mask": extended_attention_mask,751 "decoder_position_ids": decoder_position_ids,752 }753 754 def update_inputs_for_generation(self, model_outputs, model_kwargs):755 model_kwargs["past_key_values"] = model_outputs.past_key_values756 model_kwargs["decoder_position_ids"] = model_kwargs["decoder_position_ids"][:, -1:] + 1757 return model_kwargs758 759 @classmethod760 def from_encoder_decoder_pretrained(761 cls,762 encoder_pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None,763 decoder_pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None,764 *model_args,765 **kwargs,766 ) -> FlaxPreTrainedModel:767 r"""768 Instantiate an encoder and a decoder from one or two base classes of the library from pretrained model769 checkpoints.770 771 Params:772 encoder_pretrained_model_name_or_path (`Union[str, os.PathLike]`, *optional*):773 Information necessary to initiate the encoder. Can be either:774 775 - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.776 - A path to a *directory* containing model weights saved using777 [`~FlaxPreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.778 779 decoder_pretrained_model_name_or_path (`Union[str, os.PathLike]`, *optional*, defaults to `None`):780 Information necessary to initiate the decoder. Can be either:781 782 - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.783 - A path to a *directory* containing model weights saved using784 [`~FlaxPreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.785 786 model_args (remaining positional arguments, *optional*):787 All remaining positional arguments will be passed to the underlying model's `__init__` method.788 789 kwargs (remaining dictionary of keyword arguments, *optional*):790 Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,791 `output_attentions=True`).792 793 - To update the encoder configuration, use the prefix *encoder_* for each configuration parameter.794 - To update the decoder configuration, use the prefix *decoder_* for each configuration parameter.795 - To update the parent model configuration, do not use a prefix for each configuration parameter.796 797 Behaves differently depending on whether a `config` is provided or automatically loaded.798 799 Example:800 801 ```python802 >>> from transformers import FlaxEncoderDecoderModel803 804 >>> # initialize a bert2gpt2 from pretrained BERT and GPT2 models. Note that the cross-attention layers will be randomly initialized805 >>> model = FlaxEncoderDecoderModel.from_encoder_decoder_pretrained("google-bert/bert-base-cased", "openai-community/gpt2")806 >>> # saving model after fine-tuning807 >>> model.save_pretrained("./bert2gpt2")808 >>> # load fine-tuned model809 >>> model = FlaxEncoderDecoderModel.from_pretrained("./bert2gpt2")810 ```"""811 812 kwargs_encoder = {813 argument[len("encoder_") :]: value for argument, value in kwargs.items() if argument.startswith("encoder_")814 }815 816 kwargs_decoder = {817 argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")818 }819 820 # remove encoder, decoder kwargs from kwargs821 for key in kwargs_encoder:822 del kwargs["encoder_" + key]823 for key in kwargs_decoder:824 del kwargs["decoder_" + key]825 826 # Load and initialize the encoder and decoder827 # The distinction between encoder and decoder at the model level is made828 # by the value of the flag `is_decoder` that we need to set correctly.829 encoder = kwargs_encoder.pop("model", None)830 if encoder is None:831 if encoder_pretrained_model_name_or_path is None:832 raise ValueError(833 "If `encoder_model` is not defined as an argument, a `encoder_pretrained_model_name_or_path` has "834 "to be defined."835 )836 837 if "config" not in kwargs_encoder:838 encoder_config, kwargs_encoder = AutoConfig.from_pretrained(839 encoder_pretrained_model_name_or_path, **kwargs_encoder, return_unused_kwargs=True840 )841 if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:842 logger.info(843 f"Initializing {encoder_pretrained_model_name_or_path} as a encoder model "844 "from a decoder model. Cross-attention and causal mask are disabled."845 )846 encoder_config.is_decoder = False847 encoder_config.add_cross_attention = False848 849 kwargs_encoder["config"] = encoder_config850 851 encoder = FlaxAutoModel.from_pretrained(852 encoder_pretrained_model_name_or_path, *model_args, **kwargs_encoder853 )854 855 decoder = kwargs_decoder.pop("model", None)856 if decoder is None:857 if decoder_pretrained_model_name_or_path is None:858 raise ValueError(859 "If `decoder_model` is not defined as an argument, a `decoder_pretrained_model_name_or_path` has "860 "to be defined."861 )862 863 if "config" not in kwargs_decoder:864 decoder_config, kwargs_decoder = AutoConfig.from_pretrained(865 decoder_pretrained_model_name_or_path, **kwargs_decoder, return_unused_kwargs=True866 )867 if decoder_config.is_decoder is False or decoder_config.add_cross_attention is False:868 logger.info(869 f"Initializing {decoder_pretrained_model_name_or_path} as a decoder model. Cross attention"870 f" layers are added to {decoder_pretrained_model_name_or_path} and randomly initialized if"871 f" {decoder_pretrained_model_name_or_path}'s architecture allows for cross attention layers."872 )873 decoder_config.is_decoder = True874 decoder_config.add_cross_attention = True875 876 kwargs_decoder["config"] = decoder_config877 878 if kwargs_decoder["config"].is_decoder is False or kwargs_decoder["config"].add_cross_attention is False:879 logger.warning(880 f"Decoder model {decoder_pretrained_model_name_or_path} is not initialized as a decoder. "881 f"In order to initialize {decoder_pretrained_model_name_or_path} as a decoder, "882 "make sure that the attributes `is_decoder` and `add_cross_attention` of `decoder_config` "883 "passed to `.from_encoder_decoder_pretrained(...)` are set to `True` or do not pass a "884 "`decoder_config` to `.from_encoder_decoder_pretrained(...)`"885 )886 887 decoder = FlaxAutoModelForCausalLM.from_pretrained(decoder_pretrained_model_name_or_path, **kwargs_decoder)888 889 # instantiate config with corresponding kwargs890 dtype = kwargs.pop("dtype", jnp.float32)891 config = EncoderDecoderConfig.from_encoder_decoder_configs(encoder.config, decoder.config, **kwargs)892 893 # init model894 model = cls(config, dtype=dtype)895 model.params["encoder"] = encoder.params896 model.params["decoder"] = decoder.params897 898 return model899 900 901__all__ = ["FlaxEncoderDecoderModel"]902