codingutf8 2022 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 flax speechencoderdecoder 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 configurationspeechencoderdecoder import speechencoderdecoderconfig logger logging getloggername configfordoc speechencoderdecoderconfig speechencoderdecoderstartdocstring r this class can be used to initialize a speechsequencetotextsequence model with any pretrained speech 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 summarization 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 largescale self and semisupervised learning for speech translationhttps arxiv orgabs2104 06678 it is shown how leveraging large pretrained speech models for speech translation yields a significant performance improvement after such an speechencoder decoder 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 speechencoderdecoderconfig 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 speechencoderdecoderencodeinputsdocstring r args inputs jnp ndarray of shape batchsize sequencelength or batchsize sequencelength featuredim optional float values of input raw speech waveform or speech features values can be obtained by loading a flac or wav audio file into an array of type listfloat or a numpy ndarray e g via the soundfile library pip install soundfile to prepare the array into inputs either the wav2vec2processor or speech2textprocessor should be used for padding and conversion into a tensor of type torch floattensor attentionmask jnp ndarray of shape batchsize sequencelength optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossaryattentionmask 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 flaxspeechencoderdecodermodulenn module config speechencoderdecoderconfig 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 getfeatextractoutputlengths self inputlengths unionjnp ndarray int addadapter optionalbool none addadapter self config encoder addadapter if addadapter is none else addadapter def convoutlengthinputlength kernelsize stride 1d convolutional layer output length formula taken from https pytorch orgdocsstablegeneratedtorch nn conv1d html return inputlength kernelsize stride 1 for kernelsize stride in zipself config encoder convkernel self config encoder convstride inputlengths convoutlengthinputlengths kernelsize stride if addadapter for in rangeself config encoder numadapterlayers inputlengths convoutlengthinputlengths 1 self config encoder adapterstride return inputlengths def getencodermoduleself return self encoder def getprojectionmoduleself return self enctodecproj def getdecodermoduleself return self decoder def call self inputs attentionmask decoderinputids decoderattentionmask decoderpositionids encoderoutputsnone outputattentions bool false outputhiddenstates bool false returndict bool true deterministic bool true freezefeatureencoder bool false if encoderoutputs is none encoderoutputs self encoder inputs attentionmaskattentionmask outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict deterministicdeterministic freezefeatureencoderfreezefeatureencoder encoderhiddenstates encoderoutputs0 optionally project encoderhiddenstates if self enctodecproj is not none encoderhiddenstates self enctodecprojencoderhiddenstates compute correct encoder attention mask if attentionmask is not none encoderattentionmask self encoder getfeaturevectorattentionmask encoderhiddenstates shape1 attentionmask else encoderattentionmask none flax script modelingflaxwav2vec2 py 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 encoderlasthiddenstateencoderhiddenstates encoderhiddenstatesencoderoutputs hiddenstates encoderattentionsencoderoutputs attentions addstartdocstringsspeechencoderdecoderstartdocstring class flaxspeechencoderdecodermodelflaxpretrainedmodel r flaxspeechencoderdecodermodel is a generic model class that will be instantiated as a transformer architecture with the module flax nn module of one of the base 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 raise valueerror or option to project enc to dec hiddensize eg encadapterlayer make sure input output embeddings are not tied speech encoders almost always downsample the sequence length dimension init input devicearrays 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 def getfeatextractoutputlengths self inputlengths unionjnp ndarray int addadapter optionalbool none return self module getfeatextractoutputlengthsinputlengths addadapteraddadapter addstartdocstringsspeechencoderdecoderencodeinputsdocstring replacereturndocstringsoutputtypeflaxbasemodeloutput configclassconfigfordoc def encode self inputs jnp ndarray attentionmask optionaljnp ndarray none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none train bool false freezefeatureencoder bool false params dict none dropoutrng prngkey none r returns example python from transformers import flaxspeechencoderdecodermodel initialize a wav2vec22bart from pretrained wav2vec2 and bart models note that the crossattention layers will be randomly initialized model flaxspeechencoderdecodermodel fromencoderdecoderpretrained facebookwav2vec2largelv60 facebookbartlarge inputs jnp ones2 5000 dtypejnp float32 encoderoutputs model encodeinputs 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 if attentionmask is none attentionmask jnp oneslikeinputs dtypei4 handle any prng if needed rngs if dropoutrng is not none rngsdropout dropoutrng def encoderforwardmodule inputs attentionmask kwargs encodemodule module getencodermodule return encodemoduleinputs attentionmask kwargs outputs self module apply params params or self params inputsjnp arrayinputs dtypef4 attentionmaskjnp arrayattentionmask dtypei4 outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict deterministicnot train freezefeatureencoderfreezefeatureencoder rngsrngs methodencoderforward if returndict outputs flaxbasemodeloutput lasthiddenstateoutputs lasthiddenstate hiddenstatesoutputs hiddenstates attentionsoutputs attentions return outputs addstartdocstringsspeechencoderdecoderdecodeinputsdocstring replacereturndocstringsoutputtypeflaxcausallmoutputwithcrossattentions configclassconfigfordoc def decode self decoderinputids encoderoutputs encoderattentionmask optionaljnp ndarray none 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 flaxspeechencoderdecodermodel import jax numpy as jnp initialize a wav2vec22bart from pretrained wav2vec2 and bart models note that the crossattention layers will be randomly initialized model flaxspeechencoderdecodermodel fromencoderdecoderpretrained facebookwav2vec2largelv60 facebookbartlarge inputs jnp ones2 5000 dtypejnp float32 encoderoutputs model encodeinputs decoderstarttokenid model config decoder bostokenid decoderinputids jnp onesinputs 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 if encoderattentionmask is none 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 params 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 paramscache 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 encoderhiddenstatesencoderhiddenstates kwargs outputs self module apply params 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 addstartdocstringstomodelforwardspeechencoderdecoderinputsdocstring replacereturndocstringsoutputtypeflaxseq2seqlmoutput configclassconfigfordoc def call self inputs jnp ndarray attentionmask optionaljnp ndarray none decoderinputids optionaljnp ndarray none decoderattentionmask optionaljnp ndarray none decoderpositionids optionaljnp ndarray none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none train bool false freezefeatureencoder bool false params dict none dropoutrng prngkey none r returns examples python from transformers import flaxspeechencoderdecodermodel autotokenizer load a finetuned wav2vec22bart model model flaxspeechencoderdecodermodel frompretrainedpatrickvonplatenwav2vec22bartlarge load output tokenizer tokenizeroutput autotokenizer frompretrainedfacebookbartlarge inputs jnp ones2 5000 dtypejnp float32 use bart s special bos pad and eos tokens model config decoderstarttokenid model decoder config bostokenid model config padtokenid model decoder config padtokenid model config eostokenid model decoder config eostokenid outputs model generateinputs assert something more interesting input dtype correct prepare encoder inputs prepare decoder inputs handle any prng if needed initializing the cache note that usually one would have to put 0 s in the attentionmask for x input 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 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 isdecoder that we need to set correctly instantiate config with corresponding kwargs make sure input output word embeddings are not tied init model coding utf 8 2022 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 flax speech encoder decoder architectures this class can be used to initialize a speech sequence to text sequence model with any pretrained speech 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 summarization 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 large scale self and semi supervised learning for speech translation https arxiv org abs 2104 06678 it is shown how leveraging large pretrained speech models for speech translation yields a significant performance improvement after such an speech encoder 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 speechencoderdecoderconfig 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 inputs jnp ndarray of shape batch_size sequence_length or batch_size sequence_length feature_dim optional float values of input raw speech waveform or speech features values can be obtained by loading a flac or wav audio file into an array of type list float or a numpy ndarray e g via the soundfile library pip install soundfile to prepare the array into inputs either the wav2vec2processor or speech2textprocessor should be used for padding and conversion into a tensor of type torch floattensor attention_mask jnp ndarray of shape batch_size sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask 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 input ids glossary 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 decoder_input_ids should be created outside of the model by shifting the labels to the right replacing 100 by the pad_token_id and prepending them with the decoder_start_token_id 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 numpy ndarray of shape batch_size sequence_length optional indices of positions of each decoder input sequence tokens in the position embeddings selected in the range 0 config decoder max_position_embeddings 1 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 inputs jnp ndarray of shape batch_size sequence_length or batch_size sequence_length feature_dim optional float values of input raw speech waveform or speech features values can be obtained by loading a flac or wav audio file into an array of type list float or a numpy ndarray e g via the soundfile library pip install soundfile to prepare the array into inputs either the wav2vec2processor or speech2textprocessor should be used for padding and conversion into a tensor of type torch floattensor attention_mask jnp ndarray of shape batch_size sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask 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 decoder_input_ids should be created outside of the model by shifting the labels to the right replacing 100 by the pad_token_id and prepending them with the decoder_start_token_id encoder_outputs tuple tuple jnp ndarray tuple consists of last_hidden_state optional hidden_states optional attentions last_hidden_state of shape batch_size sequence_length hidden_size optional is a sequence of hidden states at the output of the last layer of the encoder used in the cross attention of the decoder encoder_attention_mask jnp ndarray of shape batch_size sequence_length optional mask to avoid performing attention on padding token indices mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossary attention mask 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 numpy ndarray of shape batch_size sequence_length optional indices of positions of each decoder input sequence tokens in the position embeddings selected in the range 0 config decoder max_position_embeddings 1 past_key_values dict str np ndarray optional returned by init_cache or when passing previous past_key_values dictionary of pre computed hidden states key and values in the attention blocks that can be used for fast auto regressive decoding pre computed key and value hidden states are of shape batch_size max_length output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail return_dict bool optional 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 computes the output length of the convolutional layers 1d convolutional layer output length formula taken from https pytorch org docs stable generated torch nn conv1d html optionally project encoder_hidden_states compute correct encoder attention mask flax script modeling_flax_wav2vec2 py flaxspeechencoderdecodermodel is a generic model class that will be instantiated as a transformer architecture with the module flax nn module of one of the base 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 raise valueerror or option to project enc to dec hidden_size eg encadapterlayer make sure input output embeddings are not tied speech encoders almost always downsample the sequence length dimension init input devicearrays 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 flaxspeechencoderdecodermodel initialize a wav2vec2 2 bart from pretrained wav2vec2 and bart models note that the cross attention layers will be randomly initialized model flaxspeechencoderdecodermodel from_encoder_decoder_pretrained facebook wav2vec2 large lv60 facebook bart large inputs jnp ones 2 5000 dtype jnp float32 encoder_outputs model encode inputs handle any prng if needed returns example python from transformers import flaxspeechencoderdecodermodel import jax numpy as jnp initialize a wav2vec2 2 bart from pretrained wav2vec2 and bart models note that the cross attention layers will be randomly initialized model flaxspeechencoderdecodermodel from_encoder_decoder_pretrained facebook wav2vec2 large lv60 facebook bart large inputs jnp ones 2 5000 dtype jnp float32 encoder_outputs model encode inputs decoder_start_token_id model config decoder bos_token_id decoder_input_ids jnp ones inputs 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 flaxspeechencoderdecodermodel autotokenizer load a fine tuned wav2vec2 2 bart model model flaxspeechencoderdecodermodel from_pretrained patrickvonplaten wav2vec2 2 bart large load output tokenizer tokenizer_output autotokenizer from_pretrained facebook bart large inputs jnp ones 2 5000 dtype jnp float32 use bart s special bos pad and eos tokens model config decoder_start_token_id model decoder config bos_token_id model config pad_token_id model decoder config pad_token_id model config eos_token_id model decoder config eos_token_id outputs model generate inputs assert something more interesting input dtype correct prepare encoder inputs 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 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 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 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 flaxspeechencoderdecodermodel initialize a wav2vec2 2 bart from pretrained wav2vec2 and bart models note that the cross attention layers will be randomly initialized model flaxspeechencoderdecodermodel from_encoder_decoder_pretrained facebook wav2vec2 large lv60 facebook bart large saving model after fine tuning model save_pretrained wav2vec2 2 bart large load fine tuned model model flaxspeechencoderdecodermodel from_pretrained wav2vec2 2 bart large 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 make sure input output word embeddings are not tied 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_speech_encoder_decoder import SpeechEncoderDecoderConfig


logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "SpeechEncoderDecoderConfig"

SPEECH_ENCODER_DECODER_START_DOCSTRING = r

SPEECH_ENCODER_DECODER_INPUTS_DOCSTRING = r

SPEECH_ENCODER_DECODER_ENCODE_INPUTS_DOCSTRING = r

SPEECH_ENCODER_DECODER_DECODE_INPUTS_DOCSTRING = r


class FlaxSpeechEncoderDecoderModule(nn.Module):
    config: SpeechEncoderDecoderConfig
    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_feat_extract_output_lengths(
        self, input_lengths: Union[jnp.ndarray, int], add_adapter: Optional[bool] = None
    ):
        

        add_adapter = self.config.encoder.add_adapter if add_adapter is None else add_adapter

        def _conv_out_length(input_length, kernel_size, stride):
            
            
            return (input_length - kernel_size) // stride + 1

        for kernel_size, stride in zip(self.config.encoder.conv_kernel, self.config.encoder.conv_stride):
            input_lengths = _conv_out_length(input_lengths, kernel_size, stride)

        if add_adapter:
            for _ in range(self.config.encoder.num_adapter_layers):
                input_lengths = _conv_out_length(input_lengths, 1, self.config.encoder.adapter_stride)

        return input_lengths

    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,
        inputs,
        attention_mask,
        decoder_input_ids,
        decoder_attention_mask,
        decoder_position_ids,
        encoder_outputs=None,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
        deterministic: bool = True,
        freeze_feature_encoder: bool = False,
    ):
        if encoder_outputs is None:
            encoder_outputs = self.encoder(
                inputs,
                attention_mask=attention_mask,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
                deterministic=deterministic,
                freeze_feature_encoder=freeze_feature_encoder,
            )

        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)

        
        if attention_mask is not None:
            encoder_attention_mask = self.encoder._get_feature_vector_attention_mask(
                encoder_hidden_states.shape[1], attention_mask
            )
        else:
            encoder_attention_mask = None

        
        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_hidden_states,
            encoder_hidden_states=encoder_outputs.hidden_states,
            encoder_attentions=encoder_outputs.attentions,
        )


@add_start_docstrings(SPEECH_ENCODER_DECODER_START_DOCSTRING)
class FlaxSpeechEncoderDecoderModel(FlaxPreTrainedModel):
    r

    config_class = SpeechEncoderDecoderConfig
    base_model_prefix: str = "speech_encoder_decoder"
    module_class = FlaxSpeechEncoderDecoderModule

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

        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`."
                )

        
        config.tie_word_embeddings = False
        module = self.module_class(config=config, dtype=dtype, **kwargs)

        if input_shape is None:
            
            encoder_input_length = 1024
            decoder_input_length = module._get_feat_extract_output_lengths(encoder_input_length)
            input_shape = ((1, encoder_input_length), (1, decoder_input_length))

        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

        
        inputs = jnp.zeros(encoder_input_shape, dtype="f4")
        attention_mask = jnp.ones_like(inputs, dtype="i4")
        decoder_input_ids = jnp.zeros(decoder_input_shape, dtype="i4")
        decoder_attention_mask = jnp.ones_like(decoder_input_ids)

        batch_size, sequence_length = inputs.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,
            inputs,
            attention_mask,
            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"])

    def _get_feat_extract_output_lengths(
        self, input_lengths: Union[jnp.ndarray, int], add_adapter: Optional[bool] = None
    ):
        return self.module._get_feat_extract_output_lengths(input_lengths, add_adapter=add_adapter)

    @add_start_docstrings(SPEECH_ENCODER_DECODER_ENCODE_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=FlaxBaseModelOutput, config_class=_CONFIG_FOR_DOC)
    def encode(
        self,
        inputs: jnp.ndarray,
        attention_mask: Optional[jnp.ndarray] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        train: bool = False,
        freeze_feature_encoder: 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

        if attention_mask is None:
            attention_mask = jnp.ones_like(inputs, dtype="i4")

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

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

        outputs = self.module.apply(
            {"params": params or self.params},
            inputs=jnp.array(inputs, dtype="f4"),
            attention_mask=jnp.array(attention_mask, dtype="i4"),
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            deterministic=not train,
            freeze_feature_encoder=freeze_feature_encoder,
            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(SPEECH_ENCODER_DECODER_DECODE_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=FlaxCausalLMOutputWithCrossAttentions, config_class=_CONFIG_FOR_DOC)
    def decode(
        self,
        decoder_input_ids,
        encoder_outputs,
        encoder_attention_mask: Optional[jnp.ndarray] = None,
        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]
        if encoder_attention_mask is None:
            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

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

        
        
        
        if past_key_values:
            params["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=encoder_hidden_states,
                **kwargs,
            )

        outputs = self.module.apply(
            params,
            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(SPEECH_ENCODER_DECODER_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=FlaxSeq2SeqLMOutput, config_class=_CONFIG_FOR_DOC)
    def __call__(
        self,
        inputs: jnp.ndarray,
        attention_mask: Optional[jnp.ndarray] = None,
        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,
        freeze_feature_encoder: 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

        
        if attention_mask is None:
            attention_mask = jnp.ones_like(inputs, dtype="i4")

        
        if decoder_input_ids is None:
            raise ValueError(
                "`decoder_input_ids` cannot be `None`. For sequence to sequence training, `decoder_position_ids` must"
                " be specified as an input argument."
            )
        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},
            inputs=jnp.array(inputs, dtype="f4"),
            attention_mask=jnp.array(attention_mask, dtype="i4"),
            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,
            freeze_feature_encoder=freeze_feature_encoder,
            rngs=rngs,
        )

    def prepare_inputs_for_generation(
        self,
        decoder_input_ids,
        max_length,
        attention_mask: Optional[jax.Array] = None,
        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,
            "encoder_attention_mask": attention_mask,
            "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, kwargs_encoder = AutoConfig.from_pretrained(
                    encoder_pretrained_model_name_or_path, **kwargs_encoder, return_unused_kwargs=True
                )
                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, kwargs_decoder = AutoConfig.from_pretrained(
                    decoder_pretrained_model_name_or_path, **kwargs_decoder, return_unused_kwargs=True
                )
                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 = SpeechEncoderDecoderConfig.from_encoder_decoder_configs(encoder.config, decoder.config, **kwargs)

        
        config.tie_word_embeddings = False

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

        return model