codingutf8 2021 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license classes to support visionencodertextdecoder architectures import os from typing import optional tuple union import flax linen as nn import jax import jax numpy as jnp from flax core frozendict import frozendict freeze unfreeze from flax traverseutil import flattendict unflattendict from jax import lax from jax random import prngkey from modelingflaxoutputs import flaxbasemodeloutput flaxcausallmoutputwithcrossattentions flaxseq2seqlmoutput from modelingflaxutils import flaxpretrainedmodel from utils import addstartdocstrings addstartdocstringstomodelforward logging replacereturndocstrings from auto configurationauto import autoconfig from auto modelingflaxauto import flaxautomodel flaxautomodelforcausallm from configurationvisionencoderdecoder import visionencoderdecoderconfig logger logging getloggername configfordoc visionencoderdecoderconfig visionencoderdecoderstartdocstring r this class can be used to initialize an imagetotextsequence model with any pretrained vision autoencoding model as the encoder and any pretrained text autoregressive model as the decoder the encoder is loaded via automodel frompretrained function and the decoder is loaded via automodelforcausallm frompretrained function crossattention layers are automatically added to the decoder and should be finetuned on a downstream generative task like image captioning the effectiveness of initializing sequencetosequence models with pretrained checkpoints for sequence generation tasks was shown in leveraging pretrained checkpoints for sequence generation taskshttps arxiv orgabs1907 12461 by sascha rothe shashi narayan aliaksei severyn michael matena yanqi zhou wei li peter j liu additionally in trocr transformerbased optical character recognition with pretrained modelshttps arxiv orgabs2109 10282 it is shown how leveraging large pretrained vision models for optical character recognition ocr yields a significant performance improvement after such a visionencodertextdecoder model has been trainedfinetuned it can be savedloaded just like any other models see the examples for more information this model inherits from flaxpretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a flax linen flax nn modulehttps flax readthedocs ioenlatestautosummaryflax nn module html subclass use it as a regular flax module and refer to the flax documentation for all matter related to general usage and behavior parameters config visionencoderdecoderconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the flaxpretrainedmodel frompretrained method to load the model weights dtype jax numpy dtype optional defaults to jax numpy float32 the data type of the computation can be one of jax numpy float32 jax numpy float16 on gpus and jax numpy bfloat16 on tpus this can be used to enable mixedprecision training or halfprecision inference on gpus or tpus if specified all the computation will be performed with the given dtype note that this only specifies the dtype of the computation and does not influence the dtype of model parameters if you wish to change the dtype of the model parameters see flaxpretrainedmodel tofp16 and flaxpretrainedmodel tobf16 visionencoderdecoderencodeinputsdocstring r args pixelvalues jnp ndarray of shape batchsize numchannels height width pixel values pixel values can be obtained using the vision model s image processor for example using autoimageprocessor see vitimageprocessor call for details outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail returndict bool optional if set to true the model will return a utils flaxbasemodeloutput instead of a plain tuple class flaxvisionencoderdecodermodulenn module config visionencoderdecoderconfig dtype jnp dtype jnp float32 def setupself encoderconfig self config encoder decoderconfig self config decoder copied from modelinghybridclip py with modifications from models auto modelingflaxauto import flaxmodelforcausallmmapping flaxmodelmapping encodermodule flaxmodelmappingencoderconfig class moduleclass decodermodule flaxmodelforcausallmmappingdecoderconfig class moduleclass self encoder encodermoduleencoderconfig dtypeself dtype self decoder decodermoduledecoderconfig dtypeself dtype encoder outputs might need to be projected to different dimension for decoder if self encoder config hiddensize self decoder config hiddensize and self decoder config crossattentionhiddensize is none self enctodecproj nn dense self decoder config hiddensize kernelinitjax nn initializers normalself decoder config initializerrange dtypeself dtype else self enctodecproj none def getencodermoduleself return self encoder def getprojectionmoduleself return self enctodecproj def getdecodermoduleself return self decoder def call self pixelvalues decoderinputids decoderattentionmask decoderpositionids outputattentions bool false outputhiddenstates bool false returndict bool true deterministic bool true encoderoutputs self encoder pixelvaluespixelvalues outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict deterministicdeterministic encoderhiddenstates encoderoutputs0 optionally project encoderhiddenstates if self enctodecproj is not none encoderhiddenstates self enctodecprojencoderhiddenstates the advantage of explicitly setting this is tpu xla compiler knows as soon as possible what shape this variable has and can better optimize also passing none can lead to some problems when jitting the model in flaxjax we only want to pass none for nontensor function inputs for all tensor function inputs we should always pass a tensor and not none batchsize sequencelength encoderhiddenstates shape 2 encoderattentionmask jnp onesbatchsize sequencelength decoderoutputs self decoder inputidsdecoderinputids attentionmaskdecoderattentionmask positionidsdecoderpositionids encoderhiddenstatesencoderhiddenstates encoderattentionmaskencoderattentionmask outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict deterministicdeterministic if not returndict return decoderoutputs encoderoutputs return flaxseq2seqlmoutput logitsdecoderoutputs logits decoderhiddenstatesdecoderoutputs hiddenstates decoderattentionsdecoderoutputs attentions crossattentionsdecoderoutputs crossattentions encoderlasthiddenstateencoderoutputs lasthiddenstate encoderhiddenstatesencoderoutputs hiddenstates encoderattentionsencoderoutputs attentions addstartdocstringsvisionencoderdecoderstartdocstring class flaxvisionencoderdecodermodelflaxpretrainedmodel r flaxvisionencoderdecodermodel is a generic model class that will be instantiated as a transformer architecture with the module flax nn module of one of the base vision model classes of the library as encoder module and another one as decoder module when created with the methtransformers flaxautomodel frompretrained class method for the encoder and methtransformers flaxautomodelforcausallm frompretrained class method for the decoder init input tensors init input variables to retrieve cache decoderinputids jnp onesbatchsize maxlength dtypei4 decoderattentionmask jnp oneslikedecoderinputids decoderpositionids jnp broadcastto jnp arangejnp atleast2ddecoderinputids shape1 decoderinputids shape def decoderforwardmodule decoderinputids decoderattentionmask decoderpositionids kwargs decodermodule module getdecodermodule return decodermodule inputidsdecoderinputids attentionmaskdecoderattentionmask positionidsdecoderpositionids kwargs initvariables self module init jax random prngkey0 decoderinputidsdecoderinputids decoderattentionmaskdecoderattentionmask decoderpositionidsdecoderpositionids encoderhiddenstatesencoderoutputs0 initcachetrue methoddecoderforward we only need to call the decoder to init the cache return unfreezeinitvariablescache addstartdocstringsvisionencoderdecoderencodeinputsdocstring replacereturndocstringsoutputtypeflaxbasemodeloutput configclassconfigfordoc def encode self pixelvalues jnp ndarray outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none train bool false params dict none dropoutrng prngkey none r returns example python from transformers import autoimageprocessor flaxvisionencoderdecodermodel from pil import image import requests url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw imageprocessor autoimageprocessor frompretrainedgooglevitbasepatch16224in21k initialize a vitgpt2 from pretrained vit and gpt2 models note that the crossattention layers will be randomly initialized model flaxvisionencoderdecodermodel fromencoderdecoderpretrained googlevitbasepatch16224in21k gpt2 pixelvalues imageprocessorimagesimage returntensorsnp pixelvalues encoderoutputs model encodepixelvalues outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config returndict flaxvitmodel expects channel first format but flaxvitmodule expects channel last format currently we assume this holds for all flax vision models and perform a transpose here pixelvalues jnp transposepixelvalues 0 2 3 1 handle any prng if needed rngs if dropoutrng is not none rngsdropout dropoutrng def encoderforwardmodule pixelvalues kwargs encodemodule module getencodermodule return encodemodulepixelvalues kwargs outputs self module apply params params or self params pixelvaluesjnp arraypixelvalues dtypeself dtype outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict deterministicnot train rngsrngs methodencoderforward if returndict outputs flaxbasemodeloutput lasthiddenstateoutputs lasthiddenstate hiddenstatesoutputs hiddenstates attentionsoutputs attentions return outputs addstartdocstringsvisionencoderdecoderdecodeinputsdocstring replacereturndocstringsoutputtypeflaxcausallmoutputwithcrossattentions configclassconfigfordoc def decode self decoderinputids encoderoutputs decoderattentionmask optionaljnp ndarray none decoderpositionids optionaljnp ndarray none pastkeyvalues dict none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none train bool false params dict none dropoutrng prngkey none r returns example python from transformers import autoimageprocessor flaxvisionencoderdecodermodel import jax numpy as jnp from pil import image import requests url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw imageprocessor autoimageprocessor frompretrainedgooglevitbasepatch16224in21k initialize a vitgpt2 from pretrained vit and gpt2 models note that the crossattention layers will be randomly initialized model flaxvisionencoderdecodermodel fromencoderdecoderpretrained googlevitbasepatch16224in21k gpt2 pixelvalues imageprocessorimagesimage returntensorsnp pixelvalues encoderoutputs model encodepixelvalues decoderstarttokenid model config decoder bostokenid decoderinputids jnp onespixelvalues shape0 1 dtypei4 decoderstarttokenid outputs model decodedecoderinputids encoderoutputs logits outputs logits outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config returndict encoderhiddenstates encoderoutputs0 batchsize sequencelength encoderhiddenstates shape 2 encoderattentionmask jnp onesbatchsize sequencelength batchsize sequencelength decoderinputids shape if decoderattentionmask is none decoderattentionmask jnp onesbatchsize sequencelength if decoderpositionids is none if pastkeyvalues is not none raise valueerrormake sure to provide decoderpositionids when passing pastkeyvalues decoderpositionids jnp broadcastto jnp arangesequencelengthnone batchsize sequencelength handle any prng if needed rngs if dropoutrng is not none rngsdropout dropoutrng inputs params params or self params if pastkeyvalues are passed then cache is already initialized a private flag initcache has to be passed down to ensure cache is used it has to be made sure that cache is marked as mutable so that it can be changed by flaxbartattention module if pastkeyvalues inputscache pastkeyvalues mutable cache else mutable false def decoderforward module decoderinputids decoderattentionmask decoderpositionids encoderhiddenstates kwargs projectionmodule module getprojectionmodule decodermodule module getdecodermodule optionally project encoderhiddenstates if projectionmodule is not none encoderhiddenstates projectionmoduleencoderhiddenstates return decodermodule decoderinputids decoderattentionmask decoderpositionids encoderhiddenstates kwargs outputs self module apply inputs decoderinputidsjnp arraydecoderinputids dtypei4 decoderattentionmaskjnp arraydecoderattentionmask dtypei4 decoderpositionidsjnp arraydecoderpositionids dtypei4 encoderhiddenstatesencoderhiddenstates encoderattentionmaskjnp arrayencoderattentionmask dtypei4 outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict deterministicnot train rngsrngs mutablemutable methoddecoderforward add updated cache to model output if pastkeyvalues is not none and returndict outputs past outputs outputspastkeyvalues unfreezepastcache return outputs elif pastkeyvalues is not none and not returndict outputs past outputs outputs outputs 1 unfreezepastcache outputs1 return outputs addstartdocstringstomodelforwardvisionencoderdecoderinputsdocstring replacereturndocstringsoutputtypeflaxseq2seqlmoutput configclassconfigfordoc def call self pixelvalues jnp ndarray decoderinputids optionaljnp ndarray none decoderattentionmask optionaljnp ndarray none decoderpositionids optionaljnp ndarray none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none train bool false params dict none dropoutrng prngkey none r returns examples python from transformers import flaxvisionencoderdecodermodel autoimageprocessor autotokenizer from pil import image import requests url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw imageprocessor autoimageprocessor frompretrainedgooglevitbasepatch16224in21k load output tokenizer tokenizeroutput autotokenizer frompretrainedgpt2 initialize a vitgpt2 from pretrained vit and gpt2 models note that the crossattention layers will be randomly initialized model flaxvisionencoderdecodermodel fromencoderdecoderpretrained googlevitbasepatch16224in21k gpt2 pixelvalues imageprocessorimagesimage returntensorsnp pixelvalues use gpt2 s eostoken as the pad as well as eos token model config eostokenid model config decoder eostokenid model config padtokenid model config eostokenid generation sequences model generatepixelvalues numbeams4 maxlength12 sequences captions tokenizeroutput batchdecodesequences skipspecialtokenstrue outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config returndict prepare encoder inputs flaxvitmodel expects channel first format but flaxvitmodule expects channel last format currently we assume this holds for all flax vision models and perform a transpose here pixelvalues jnp transposepixelvalues 0 2 3 1 prepare decoder inputs if decoderinputids is none raise valueerrordecoderinputids can t be none if decoderattentionmask is none decoderattentionmask jnp oneslikedecoderinputids if decoderpositionids is none batchsize sequencelength decoderinputids shape decoderpositionids jnp broadcastto jnp arangesequencelengthnone batchsize sequencelength handle any prng if needed rngs dropout dropoutrng if dropoutrng is not none else return self module apply params params or self params pixelvaluesjnp arraypixelvalues dtypeself dtype decoderinputidsjnp arraydecoderinputids dtypei4 decoderattentionmaskjnp arraydecoderattentionmask dtypei4 decoderpositionidsjnp arraydecoderpositionids dtypei4 outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict deterministicnot train rngsrngs def prepareinputsforgeneration self decoderinputids maxlength decoderattentionmask optionaljax array none encoderoutputsnone kwargs initializing the cache batchsize seqlength decoderinputids shape pastkeyvalues self initcachebatchsize maxlength encoderoutputs note that usually one would have to put 0 s in the attentionmask for x inputids shape1 and x cachelength but since the decoder uses a causal mask those positions are masked anyways thus we can create a single static attentionmask here which is more efficient for compilation extendedattentionmask jnp onesbatchsize maxlength dtypei4 if decoderattentionmask is not none decoderpositionids decoderattentionmask cumsumaxis1 1 extendedattentionmask lax dynamicupdatesliceextendedattentionmask decoderattentionmask 0 0 else decoderpositionids jnp broadcastto jnp arangeseqlength dtypei4none batchsize seqlength return pastkeyvalues pastkeyvalues encoderoutputs encoderoutputs decoderattentionmask extendedattentionmask decoderpositionids decoderpositionids def updateinputsforgenerationself modeloutputs modelkwargs modelkwargspastkeyvalues modeloutputs pastkeyvalues modelkwargsdecoderpositionids modelkwargsdecoderpositionids 1 1 return modelkwargs classmethod def fromencoderdecoderpretrained cls encoderpretrainedmodelnameorpath optionalunionstr os pathlike none decoderpretrainedmodelnameorpath optionalunionstr os pathlike none modelargs kwargs flaxpretrainedmodel r instantiate an encoder and a decoder from one or two base classes of the library from pretrained model checkpoints params encoderpretrainedmodelnameorpath unionstr os pathlike optional information necessary to initiate the encoder can be either a string the model id of a pretrained model hosted inside a model repo on huggingface co an example is googlevitbasepatch16224in21k a path to a directory containing model weights saved using flaxpretrainedmodel savepretrained e g mymodeldirectory decoderpretrainedmodelnameorpath unionstr os pathlike optional defaults to none information necessary to initiate the decoder can be either a string the model id of a pretrained model hosted inside a model repo on huggingface co valid model ids can be located at the rootlevel like bertbaseuncased or namespaced under a user or organization name like dbmdzbertbasegermancased a path to a directory containing model weights saved using flaxpretrainedmodel savepretrained e g mymodeldirectory modelargs remaining positional arguments optional all remaning positional arguments will be passed to the underlying model s init method kwargs remaining dictionary of keyword arguments optional can be used to update the configuration object after it being loaded and initiate the model e g outputattentionstrue to update the encoder configuration use the prefix encoder for each configuration parameter to update the decoder configuration use the prefix decoder for each configuration parameter to update the parent model configuration do not use a prefix for each configuration parameter behaves differently depending on whether a config is provided or automatically loaded example python from transformers import flaxvisionencoderdecodermodel initialize a vitgpt2 from a pretrained vit and a pretrained gpt2 model note that the crossattention layers will be randomly initialized model flaxvisionencoderdecodermodel fromencoderdecoderpretrained googlevitbasepatch16224in21k gpt2 saving model after finetuning model savepretrained vitgpt2 load finetuned model model flaxvisionencoderdecodermodel frompretrained vitgpt2 kwargsencoder argumentlenencoder value for argument value in kwargs items if argument startswithencoder kwargsdecoder argumentlendecoder value for argument value in kwargs items if argument startswithdecoder remove encoder decoder kwargs from kwargs for key in kwargsencoder keys del kwargsencoder key for key in kwargsdecoder keys del kwargsdecoder key load and initialize the encoder and decoder the distinction between encoder and decoder at the model level is made by the value of the flag isdecoder that we need to set correctly encoder kwargsencoder popmodel none if encoder is none if encoderpretrainedmodelnameorpath is none raise valueerror if encodermodel is not defined as an argument a encoderpretrainedmodelnameorpath has to be defined if config not in kwargsencoder encoderconfig autoconfig frompretrainedencoderpretrainedmodelnameorpath if encoderconfig isdecoder is true or encoderconfig addcrossattention is true logger info finitializing encoderpretrainedmodelnameorpath as a encoder model from a decoder model crossattention and casual mask are disabled encoderconfig isdecoder false encoderconfig addcrossattention false kwargsencoderconfig encoderconfig encoder flaxautomodel frompretrained encoderpretrainedmodelnameorpath modelargs kwargsencoder decoder kwargsdecoder popmodel none if decoder is none if decoderpretrainedmodelnameorpath is none raise valueerror if decodermodel is not defined as an argument a decoderpretrainedmodelnameorpath has to be defined if config not in kwargsdecoder decoderconfig autoconfig frompretraineddecoderpretrainedmodelnameorpath if decoderconfig isdecoder is false or decoderconfig addcrossattention is false logger info finitializing decoderpretrainedmodelnameorpath as a decoder model cross attention f layers are added to decoderpretrainedmodelnameorpath and randomly initialized if f decoderpretrainedmodelnameorpath s architecture allows for cross attention layers decoderconfig isdecoder true decoderconfig addcrossattention true kwargsdecoderconfig decoderconfig if kwargsdecoderconfig isdecoder is false or kwargsdecoderconfig addcrossattention is false logger warning fdecoder model decoderpretrainedmodelnameorpath is not initialized as a decoder fin order to initialize decoderpretrainedmodelnameorpath as a decoder make sure that the attributes isdecoder and addcrossattention of decoderconfig passed to fromencoderdecoderpretrained are set to true or do not pass a decoderconfig to fromencoderdecoderpretrained decoder flaxautomodelforcausallm frompretraineddecoderpretrainedmodelnameorpath kwargsdecoder instantiate config with corresponding kwargs dtype kwargs popdtype jnp float32 config visionencoderdecoderconfig fromencoderdecoderconfigsencoder config decoder config kwargs init model model clsconfig dtypedtype model paramsencoder encoder params model paramsdecoder decoder params return model coding utf 8 2021 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license classes to support vision encoder text decoder architectures this class can be used to initialize an image to text sequence model with any pretrained vision autoencoding model as the encoder and any pretrained text autoregressive model as the decoder the encoder is loaded via automodel from_pretrained function and the decoder is loaded via automodelforcausallm from_pretrained function cross attention layers are automatically added to the decoder and should be fine tuned on a downstream generative task like image captioning the effectiveness of initializing sequence to sequence models with pretrained checkpoints for sequence generation tasks was shown in leveraging pre trained checkpoints for sequence generation tasks https arxiv org abs 1907 12461 by sascha rothe shashi narayan aliaksei severyn michael matena yanqi zhou wei li peter j liu additionally in trocr transformer based optical character recognition with pre trained models https arxiv org abs 2109 10282 it is shown how leveraging large pretrained vision models for optical character recognition ocr yields a significant performance improvement after such a vision encoder text decoder model has been trained fine tuned it can be saved loaded just like any other models see the examples for more information this model inherits from flaxpretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a flax linen flax nn module https flax readthedocs io en latest _autosummary flax nn module html subclass use it as a regular flax module and refer to the flax documentation for all matter related to general usage and behavior parameters config visionencoderdecoderconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the flaxpretrainedmodel from_pretrained method to load the model weights dtype jax numpy dtype optional defaults to jax numpy float32 the data type of the computation can be one of jax numpy float32 jax numpy float16 on gpus and jax numpy bfloat16 on tpus this can be used to enable mixed precision training or half precision inference on gpus or tpus if specified all the computation will be performed with the given dtype note that this only specifies the dtype of the computation and does not influence the dtype of model parameters if you wish to change the dtype of the model parameters see flaxpretrainedmodel to_fp16 and flaxpretrainedmodel to_bf16 args pixel_values jnp ndarray of shape batch_size num_channels height width pixel values pixel values can be obtained using the vision model s image processor for example using autoimageprocessor see vitimageprocessor __call__ for details decoder_input_ids jnp ndarray of shape batch_size target_sequence_length optional indices of decoder input sequence tokens in the vocabulary indices can be obtained using pretrainedtokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are decoder input ids glossary decoder input ids decoder_attention_mask jnp ndarray of shape batch_size target_sequence_length optional default behavior generate a tensor that ignores pad tokens in decoder_input_ids causal mask will also be used by default decoder_position_ids jnp ndarray of shape batch_size sequence_length optional indices of positions of each decoder input sequence tokens in the position embeddings selected in the range 0 config decoder max_position_embeddings 1 output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional if set to true the model will return a utils flaxseq2seqlmoutput instead of a plain tuple args pixel_values jnp ndarray of shape batch_size num_channels height width pixel values pixel values can be obtained using the vision model s image processor for example using autoimageprocessor see vitimageprocessor __call__ for details output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional if set to true the model will return a utils flaxbasemodeloutput instead of a plain tuple args decoder_input_ids jnp ndarray of shape batch_size target_sequence_length optional indices of decoder input sequence tokens in the vocabulary indices can be obtained using pretrainedtokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are decoder input ids glossary decoder input ids if past_key_values is used optionally only the last decoder_input_ids have to be input see past_key_values for sequence to sequence training decoder_input_ids should be provided if no decoder_input_ids is provided the model will create this tensor by shifting the input_ids to the right for denoising pre training encoder_outputs tuple tuple jnp ndarray tuple consists of last_hidden_state optional hidden_states optional attentions last_hidden_state of shape batch_size sequence_length hidden_size optional is a sequence of hidden states at the output of the last layer of the encoder used in the cross attention of the decoder decoder_attention_mask jnp ndarray of shape batch_size target_sequence_length optional default behavior generate a tensor that ignores pad tokens in decoder_input_ids causal mask will also be used by default decoder_position_ids jnp ndarray of shape batch_size sequence_length optional indices of positions of each decoder input sequence tokens in the position embeddings selected in the range 0 config decoder max_position_embeddings 1 past_key_values dict str jnp ndarray optional returned by init_cache or when passing previous past_key_values dictionary of pre computed hidden states key and values in the attention blocks that can be used for fast auto regressive decoding pre computed key and value hidden states are of shape batch_size max_length output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional if set to true the model will return a utils flaxcausallmoutputwithcrossattentions instead of a plain tuple copied from modeling_hybrid_clip py with modifications encoder outputs might need to be projected to different dimension for decoder optionally project encoder_hidden_states the advantage of explicitly setting this is tpu xla compiler knows as soon as possible what shape this variable has and can better optimize also passing none can lead to some problems when jitting the model in flax jax we only want to pass none for non tensor function inputs for all tensor function inputs we should always pass a tensor and not none flaxvisionencoderdecodermodel is a generic model class that will be instantiated as a transformer architecture with the module flax nn module of one of the base vision model classes of the library as encoder module and another one as decoder module when created with the meth transformers flaxautomodel from_pretrained class method for the encoder and meth transformers flaxautomodelforcausallm from_pretrained class method for the decoder init input tensors args batch_size int batch_size used for fast auto regressive decoding defines the batch size of the initialized cache max_length int maximum possible length for auto regressive decoding defines the sequence length of the initialized cache encoder_outputs union flaxbasemodeloutput tuple tuple jnp ndarray encoder_outputs consists of last_hidden_state optional hidden_states optional attentions last_hidden_state of shape batch_size sequence_length hidden_size optional is a sequence of hidden states at the output of the last layer of the encoder used in the cross attention of the decoder init input variables to retrieve cache we only need to call the decoder to init the cache returns example python from transformers import autoimageprocessor flaxvisionencoderdecodermodel from pil import image import requests url http images cocodataset org val2017 000000039769 jpg image image open requests get url stream true raw image_processor autoimageprocessor from_pretrained google vit base patch16 224 in21k initialize a vit gpt2 from pretrained vit and gpt2 models note that the cross attention layers will be randomly initialized model flaxvisionencoderdecodermodel from_encoder_decoder_pretrained google vit base patch16 224 in21k gpt2 pixel_values image_processor images image return_tensors np pixel_values encoder_outputs model encode pixel_values flaxvitmodel expects channel first format but flaxvitmodule expects channel last format currently we assume this holds for all flax vision models and perform a transpose here handle any prng if needed returns example python from transformers import autoimageprocessor flaxvisionencoderdecodermodel import jax numpy as jnp from pil import image import requests url http images cocodataset org val2017 000000039769 jpg image image open requests get url stream true raw image_processor autoimageprocessor from_pretrained google vit base patch16 224 in21k initialize a vit gpt2 from pretrained vit and gpt2 models note that the cross attention layers will be randomly initialized model flaxvisionencoderdecodermodel from_encoder_decoder_pretrained google vit base patch16 224 in21k gpt2 pixel_values image_processor images image return_tensors np pixel_values encoder_outputs model encode pixel_values decoder_start_token_id model config decoder bos_token_id decoder_input_ids jnp ones pixel_values shape 0 1 dtype i4 decoder_start_token_id outputs model decode decoder_input_ids encoder_outputs logits outputs logits handle any prng if needed if past_key_values are passed then cache is already initialized a private flag init_cache has to be passed down to ensure cache is used it has to be made sure that cache is marked as mutable so that it can be changed by flaxbartattention module optionally project encoder_hidden_states add updated cache to model output returns examples python from transformers import flaxvisionencoderdecodermodel autoimageprocessor autotokenizer from pil import image import requests url http images cocodataset org val2017 000000039769 jpg image image open requests get url stream true raw image_processor autoimageprocessor from_pretrained google vit base patch16 224 in21k load output tokenizer tokenizer_output autotokenizer from_pretrained gpt2 initialize a vit gpt2 from pretrained vit and gpt2 models note that the cross attention layers will be randomly initialized model flaxvisionencoderdecodermodel from_encoder_decoder_pretrained google vit base patch16 224 in21k gpt2 pixel_values image_processor images image return_tensors np pixel_values use gpt2 s eos_token as the pad as well as eos token model config eos_token_id model config decoder eos_token_id model config pad_token_id model config eos_token_id generation sequences model generate pixel_values num_beams 4 max_length 12 sequences captions tokenizer_output batch_decode sequences skip_special_tokens true prepare encoder inputs flaxvitmodel expects channel first format but flaxvitmodule expects channel last format currently we assume this holds for all flax vision models and perform a transpose here prepare decoder inputs handle any prng if needed initializing the cache note that usually one would have to put 0 s in the attention_mask for x input_ids shape 1 and x cache_length but since the decoder uses a causal mask those positions are masked anyways thus we can create a single static attention_mask here which is more efficient for compilation instantiate an encoder and a decoder from one or two base classes of the library from pretrained model checkpoints params encoder_pretrained_model_name_or_path union str os pathlike optional information necessary to initiate the encoder can be either a string the model id of a pretrained model hosted inside a model repo on huggingface co an example is google vit base patch16 224 in21k a path to a directory containing model weights saved using flaxpretrainedmodel save_pretrained e g my_model_directory decoder_pretrained_model_name_or_path union str os pathlike optional defaults to none information necessary to initiate the decoder can be either a string the model id of a pretrained model hosted inside a model repo on huggingface co valid model ids can be located at the root level like bert base uncased or namespaced under a user or organization name like dbmdz bert base german cased a path to a directory containing model weights saved using flaxpretrainedmodel save_pretrained e g my_model_directory model_args remaining positional arguments optional all remaning positional arguments will be passed to the underlying model s __init__ method kwargs remaining dictionary of keyword arguments optional can be used to update the configuration object after it being loaded and initiate the model e g output_attentions true to update the encoder configuration use the prefix encoder_ for each configuration parameter to update the decoder configuration use the prefix decoder_ for each configuration parameter to update the parent model configuration do not use a prefix for each configuration parameter behaves differently depending on whether a config is provided or automatically loaded example python from transformers import flaxvisionencoderdecodermodel initialize a vit gpt2 from a pretrained vit and a pretrained gpt2 model note that the cross attention layers will be randomly initialized model flaxvisionencoderdecodermodel from_encoder_decoder_pretrained google vit base patch16 224 in21k gpt2 saving model after fine tuning model save_pretrained vit gpt2 load fine tuned model model flaxvisionencoderdecodermodel from_pretrained vit gpt2 remove encoder decoder kwargs from kwargs load and initialize the encoder and decoder the distinction between encoder and decoder at the model level is made by the value of the flag is_decoder that we need to set correctly instantiate config with corresponding kwargs init model
import os
from typing import Optional, Tuple, Union

import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax import lax
from jax.random import PRNGKey

from ...modeling_flax_outputs import FlaxBaseModelOutput, FlaxCausalLMOutputWithCrossAttentions, FlaxSeq2SeqLMOutput
from ...modeling_flax_utils import FlaxPreTrainedModel
from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
from ..auto.configuration_auto import AutoConfig
from ..auto.modeling_flax_auto import FlaxAutoModel, FlaxAutoModelForCausalLM
from .configuration_vision_encoder_decoder import VisionEncoderDecoderConfig


logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "VisionEncoderDecoderConfig"

VISION_ENCODER_DECODER_START_DOCSTRING = r

VISION_ENCODER_DECODER_INPUTS_DOCSTRING = r

VISION_ENCODER_DECODER_ENCODE_INPUTS_DOCSTRING = r

VISION_ENCODER_DECODER_DECODE_INPUTS_DOCSTRING = r


class FlaxVisionEncoderDecoderModule(nn.Module):
    config: VisionEncoderDecoderConfig
    dtype: jnp.dtype = jnp.float32

    def setup(self):
        encoder_config = self.config.encoder
        decoder_config = self.config.decoder

        
        from ...models.auto.modeling_flax_auto import FLAX_MODEL_FOR_CAUSAL_LM_MAPPING, FLAX_MODEL_MAPPING

        encoder_module = FLAX_MODEL_MAPPING[encoder_config.__class__].module_class
        decoder_module = FLAX_MODEL_FOR_CAUSAL_LM_MAPPING[decoder_config.__class__].module_class

        self.encoder = encoder_module(encoder_config, dtype=self.dtype)
        self.decoder = decoder_module(decoder_config, dtype=self.dtype)

        
        if (
            self.encoder.config.hidden_size != self.decoder.config.hidden_size
            and self.decoder.config.cross_attention_hidden_size is None
        ):
            self.enc_to_dec_proj = nn.Dense(
                self.decoder.config.hidden_size,
                kernel_init=jax.nn.initializers.normal(self.decoder.config.initializer_range),
                dtype=self.dtype,
            )
        else:
            self.enc_to_dec_proj = None

    def _get_encoder_module(self):
        return self.encoder

    def _get_projection_module(self):
        return self.enc_to_dec_proj

    def _get_decoder_module(self):
        return self.decoder

    def __call__(
        self,
        pixel_values,
        decoder_input_ids,
        decoder_attention_mask,
        decoder_position_ids,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
        deterministic: bool = True,
    ):
        encoder_outputs = self.encoder(
            pixel_values=pixel_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            deterministic=deterministic,
        )

        encoder_hidden_states = encoder_outputs[0]

        
        if self.enc_to_dec_proj is not None:
            encoder_hidden_states = self.enc_to_dec_proj(encoder_hidden_states)

        
        
        
        
        batch_size, sequence_length = encoder_hidden_states.shape[:2]
        encoder_attention_mask = jnp.ones((batch_size, sequence_length))

        decoder_outputs = self.decoder(
            input_ids=decoder_input_ids,
            attention_mask=decoder_attention_mask,
            position_ids=decoder_position_ids,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            deterministic=deterministic,
        )

        if not return_dict:
            return decoder_outputs + encoder_outputs

        return FlaxSeq2SeqLMOutput(
            logits=decoder_outputs.logits,
            decoder_hidden_states=decoder_outputs.hidden_states,
            decoder_attentions=decoder_outputs.attentions,
            cross_attentions=decoder_outputs.cross_attentions,
            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
            encoder_hidden_states=encoder_outputs.hidden_states,
            encoder_attentions=encoder_outputs.attentions,
        )


@add_start_docstrings(VISION_ENCODER_DECODER_START_DOCSTRING)
class FlaxVisionEncoderDecoderModel(FlaxPreTrainedModel):
    r

    config_class = VisionEncoderDecoderConfig
    base_model_prefix = "vision_encoder_decoder"
    main_input_name = "pixel_values"
    module_class = FlaxVisionEncoderDecoderModule

    def __init__(
        self,
        config: VisionEncoderDecoderConfig,
        input_shape: Optional[Tuple] = None,
        seed: int = 0,
        dtype: jnp.dtype = jnp.float32,
        _do_init: bool = True,
        **kwargs,
    ):
        if not _do_init:
            raise ValueError(
                "`FlaxVisionEncoderDecoderModel` cannot be created without initializing, `_do_init` must be `True`."
            )

        if input_shape is None:
            num_channels = getattr(config.encoder, "num_channels", 3)
            input_shape = (
                (1, config.encoder.image_size, config.encoder.image_size, num_channels),
                (1, 1),
            )

        if config.decoder.cross_attention_hidden_size is not None:
            if config.decoder.cross_attention_hidden_size != config.encoder.hidden_size:
                raise ValueError(
                    "If `cross_attention_hidden_size` is specified in the decoder's configuration, it has to be equal"
                    f" to the encoder's `hidden_size`. Got {config.decoder.cross_attention_hidden_size} for"
                    f" `config.decoder.cross_attention_hidden_size` and {config.encoder.hidden_size} for"
                    " `config.encoder.hidden_size`."
                )

        module = self.module_class(config=config, dtype=dtype, **kwargs)
        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)

    def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict:
        encoder_input_shape, decoder_input_shape = input_shape

        
        pixel_values = jnp.zeros(encoder_input_shape, dtype=self.dtype)
        decoder_input_ids = jnp.zeros(decoder_input_shape, dtype="i4")
        decoder_attention_mask = jnp.ones_like(decoder_input_ids)

        batch_size, _, _, _ = pixel_values.shape
        decoder_batch_size, decoder_sequence_length = decoder_input_ids.shape
        if not decoder_batch_size == batch_size:
            raise ValueError(
                f"The inputs of encoder and decoder should have the same batch size, but got {batch_size} for encoder "
                f"and {decoder_batch_size} for decoder."
            )
        decoder_position_ids = jnp.broadcast_to(
            jnp.arange(decoder_sequence_length)[None, :], (decoder_batch_size, decoder_sequence_length)
        )

        params_rng, dropout_rng = jax.random.split(rng)
        rngs = {"params": params_rng, "dropout": dropout_rng}

        random_params = self.module.init(
            rngs,
            pixel_values,
            decoder_input_ids,
            decoder_attention_mask,
            decoder_position_ids,
        )["params"]

        if params is not None:
            random_params = flatten_dict(unfreeze(random_params))
            params = flatten_dict(unfreeze(params))
            for missing_key in self._missing_keys:
                params[missing_key] = random_params[missing_key]
            self._missing_keys = set()
            return freeze(unflatten_dict(params))
        else:
            return random_params

    def init_cache(self, batch_size, max_length, encoder_outputs):
        r
        
        decoder_input_ids = jnp.ones((batch_size, max_length), dtype="i4")
        decoder_attention_mask = jnp.ones_like(decoder_input_ids)
        decoder_position_ids = jnp.broadcast_to(
            jnp.arange(jnp.atleast_2d(decoder_input_ids).shape[-1]), decoder_input_ids.shape
        )

        def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):
            decoder_module = module._get_decoder_module()
            return decoder_module(
                input_ids=decoder_input_ids,
                attention_mask=decoder_attention_mask,
                position_ids=decoder_position_ids,
                **kwargs,
            )

        init_variables = self.module.init(
            jax.random.PRNGKey(0),
            decoder_input_ids=decoder_input_ids,
            decoder_attention_mask=decoder_attention_mask,
            decoder_position_ids=decoder_position_ids,
            encoder_hidden_states=encoder_outputs[0],
            init_cache=True,
            method=_decoder_forward,  
        )
        return unfreeze(init_variables["cache"])

    @add_start_docstrings(VISION_ENCODER_DECODER_ENCODE_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=FlaxBaseModelOutput, config_class=_CONFIG_FOR_DOC)
    def encode(
        self,
        pixel_values: jnp.ndarray,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        train: bool = False,
        params: dict = None,
        dropout_rng: PRNGKey = None,
    ):
        r
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        
        
        pixel_values = jnp.transpose(pixel_values, (0, 2, 3, 1))

        
        rngs = {}
        if dropout_rng is not None:
            rngs["dropout"] = dropout_rng

        def _encoder_forward(module, pixel_values, **kwargs):
            encode_module = module._get_encoder_module()
            return encode_module(pixel_values, **kwargs)

        outputs = self.module.apply(
            {"params": params or self.params},
            pixel_values=jnp.array(pixel_values, dtype=self.dtype),
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            deterministic=not train,
            rngs=rngs,
            method=_encoder_forward,
        )

        if return_dict:
            outputs = FlaxBaseModelOutput(
                last_hidden_state=outputs.last_hidden_state,
                hidden_states=outputs.hidden_states,
                attentions=outputs.attentions,
            )

        return outputs

    @add_start_docstrings(VISION_ENCODER_DECODER_DECODE_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=FlaxCausalLMOutputWithCrossAttentions, config_class=_CONFIG_FOR_DOC)
    def decode(
        self,
        decoder_input_ids,
        encoder_outputs,
        decoder_attention_mask: Optional[jnp.ndarray] = None,
        decoder_position_ids: Optional[jnp.ndarray] = None,
        past_key_values: dict = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        train: bool = False,
        params: dict = None,
        dropout_rng: PRNGKey = None,
    ):
        r
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        encoder_hidden_states = encoder_outputs[0]

        batch_size, sequence_length = encoder_hidden_states.shape[:2]
        encoder_attention_mask = jnp.ones((batch_size, sequence_length))

        batch_size, sequence_length = decoder_input_ids.shape
        if decoder_attention_mask is None:
            decoder_attention_mask = jnp.ones((batch_size, sequence_length))

        if decoder_position_ids is None:
            if past_key_values is not None:
                raise ValueError("Make sure to provide `decoder_position_ids` when passing `past_key_values`.")

            decoder_position_ids = jnp.broadcast_to(
                jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
            )

        
        rngs = {}
        if dropout_rng is not None:
            rngs["dropout"] = dropout_rng

        inputs = {"params": params or self.params}

        
        
        
        if past_key_values:
            inputs["cache"] = past_key_values
            mutable = ["cache"]
        else:
            mutable = False

        def _decoder_forward(
            module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, encoder_hidden_states, **kwargs
        ):
            projection_module = module._get_projection_module()
            decoder_module = module._get_decoder_module()

            
            if projection_module is not None:
                encoder_hidden_states = projection_module(encoder_hidden_states)

            return decoder_module(
                decoder_input_ids,
                decoder_attention_mask,
                decoder_position_ids,
                encoder_hidden_states,
                **kwargs,
            )

        outputs = self.module.apply(
            inputs,
            decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),
            decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),
            decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=jnp.array(encoder_attention_mask, dtype="i4"),
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            deterministic=not train,
            rngs=rngs,
            mutable=mutable,
            method=_decoder_forward,
        )

        
        if past_key_values is not None and return_dict:
            outputs, past = outputs
            outputs["past_key_values"] = unfreeze(past["cache"])
            return outputs
        elif past_key_values is not None and not return_dict:
            outputs, past = outputs
            outputs = outputs[:1] + (unfreeze(past["cache"]),) + outputs[1:]

        return outputs

    @add_start_docstrings_to_model_forward(VISION_ENCODER_DECODER_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=FlaxSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)
    def __call__(
        self,
        pixel_values: jnp.ndarray,
        decoder_input_ids: Optional[jnp.ndarray] = None,
        decoder_attention_mask: Optional[jnp.ndarray] = None,
        decoder_position_ids: Optional[jnp.ndarray] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        train: bool = False,
        params: dict = None,
        dropout_rng: PRNGKey = None,
    ):
        r

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        

        
        
        pixel_values = jnp.transpose(pixel_values, (0, 2, 3, 1))

        
        if decoder_input_ids is None:
            raise ValueError("`decoder_input_ids` can't be `None`.")
        if decoder_attention_mask is None:
            decoder_attention_mask = jnp.ones_like(decoder_input_ids)
        if decoder_position_ids is None:
            batch_size, sequence_length = decoder_input_ids.shape
            decoder_position_ids = jnp.broadcast_to(
                jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
            )

        
        rngs = {"dropout": dropout_rng} if dropout_rng is not None else {}

        return self.module.apply(
            {"params": params or self.params},
            pixel_values=jnp.array(pixel_values, dtype=self.dtype),
            decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),
            decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),
            decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            deterministic=not train,
            rngs=rngs,
        )

    def prepare_inputs_for_generation(
        self,
        decoder_input_ids,
        max_length,
        decoder_attention_mask: Optional[jax.Array] = None,
        encoder_outputs=None,
        **kwargs,
    ):
        
        batch_size, seq_length = decoder_input_ids.shape

        past_key_values = self.init_cache(batch_size, max_length, encoder_outputs)
        
        
        
        extended_attention_mask = jnp.ones((batch_size, max_length), dtype="i4")
        if decoder_attention_mask is not None:
            decoder_position_ids = decoder_attention_mask.cumsum(axis=-1) - 1
            extended_attention_mask = lax.dynamic_update_slice(extended_attention_mask, decoder_attention_mask, (0, 0))
        else:
            decoder_position_ids = jnp.broadcast_to(
                jnp.arange(seq_length, dtype="i4")[None, :], (batch_size, seq_length)
            )

        return {
            "past_key_values": past_key_values,
            "encoder_outputs": encoder_outputs,
            "decoder_attention_mask": extended_attention_mask,
            "decoder_position_ids": decoder_position_ids,
        }

    def update_inputs_for_generation(self, model_outputs, model_kwargs):
        model_kwargs["past_key_values"] = model_outputs.past_key_values
        model_kwargs["decoder_position_ids"] = model_kwargs["decoder_position_ids"][:, -1:] + 1
        return model_kwargs

    @classmethod
    def from_encoder_decoder_pretrained(
        cls,
        encoder_pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None,
        decoder_pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None,
        *model_args,
        **kwargs,
    ) -> FlaxPreTrainedModel:
        r

        kwargs_encoder = {
            argument[len("encoder_") :]: value for argument, value in kwargs.items() if argument.startswith("encoder_")
        }

        kwargs_decoder = {
            argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
        }

        
        for key in kwargs_encoder.keys():
            del kwargs["encoder_" + key]
        for key in kwargs_decoder.keys():
            del kwargs["decoder_" + key]

        
        
        
        encoder = kwargs_encoder.pop("model", None)
        if encoder is None:
            if encoder_pretrained_model_name_or_path is None:
                raise ValueError(
                    "If `encoder_model` is not defined as an argument, a `encoder_pretrained_model_name_or_path` has "
                    "to be defined."
                )

            if "config" not in kwargs_encoder:
                encoder_config = AutoConfig.from_pretrained(encoder_pretrained_model_name_or_path)
                if encoder_config.is_decoder is True or encoder_config.add_cross_attention is True:
                    logger.info(
                        f"Initializing {encoder_pretrained_model_name_or_path} as a encoder model "
                        "from a decoder model. Cross-attention and casual mask are disabled."
                    )
                    encoder_config.is_decoder = False
                    encoder_config.add_cross_attention = False

                kwargs_encoder["config"] = encoder_config

            encoder = FlaxAutoModel.from_pretrained(
                encoder_pretrained_model_name_or_path, *model_args, **kwargs_encoder
            )

        decoder = kwargs_decoder.pop("model", None)
        if decoder is None:
            if decoder_pretrained_model_name_or_path is None:
                raise ValueError(
                    "If `decoder_model` is not defined as an argument, a `decoder_pretrained_model_name_or_path` has "
                    "to be defined."
                )

            if "config" not in kwargs_decoder:
                decoder_config = AutoConfig.from_pretrained(decoder_pretrained_model_name_or_path)
                if decoder_config.is_decoder is False or decoder_config.add_cross_attention is False:
                    logger.info(
                        f"Initializing {decoder_pretrained_model_name_or_path} as a decoder model. Cross attention"
                        f" layers are added to {decoder_pretrained_model_name_or_path} and randomly initialized if"
                        f" {decoder_pretrained_model_name_or_path}'s architecture allows for cross attention layers."
                    )
                    decoder_config.is_decoder = True
                    decoder_config.add_cross_attention = True

                kwargs_decoder["config"] = decoder_config

            if kwargs_decoder["config"].is_decoder is False or kwargs_decoder["config"].add_cross_attention is False:
                logger.warning(
                    f"Decoder model {decoder_pretrained_model_name_or_path} is not initialized as a decoder. "
                    f"In order to initialize {decoder_pretrained_model_name_or_path} as a decoder, "
                    "make sure that the attributes `is_decoder` and `add_cross_attention` of `decoder_config` "
                    "passed to `.from_encoder_decoder_pretrained(...)` are set to `True` or do not pass a "
                    "`decoder_config` to `.from_encoder_decoder_pretrained(...)`"
                )

            decoder = FlaxAutoModelForCausalLM.from_pretrained(decoder_pretrained_model_name_or_path, **kwargs_decoder)

        
        dtype = kwargs.pop("dtype", jnp.float32)
        config = VisionEncoderDecoderConfig.from_encoder_decoder_configs(encoder.config, decoder.config, **kwargs)

        
        model = cls(config, dtype=dtype)
        model.params["encoder"] = encoder.params
        model.params["decoder"] = decoder.params

        return model