codingutf8 2023 the huggingface inc google 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 pix2struct modeling file import math from typing import dict list optional tuple union import torch import torch utils checkpoint from torch import nn from activations import act2fn from modelingoutputs import basemodeloutput basemodeloutputwithpooling causallmoutputwithcrossattentions seq2seqlmoutput seq2seqmodeloutput from modelingutils import pretrainedmodel from pytorchutils import alllayernormlayers from utils import dummyinputs dummymask addstartdocstrings addstartdocstringstomodelforward istorchfxproxy logging replacereturndocstrings from configurationpix2struct import pix2structconfig pix2structtextconfig pix2structvisionconfig logger logging getloggername general docstring configfordoc pix2structconfig pix2structpretrainedmodelarchivelist googlepix2structtextcapsbase googlepix2structtextcapslarge googlepix2structbase googlepix2structlarge googlepix2structai2dbase googlepix2structai2dlarge googlepix2structwidgetcaptioningbase googlepix2structwidgetcaptioninglarge googlepix2structscreen2wordsbase googlepix2structscreen2wordslarge googlepix2structdocvqabase googlepix2structdocvqalarge googlepix2structocrvqabase googlepix2structocrvqalarge googlepix2structchartqabase googlepix2structinforgraphicsvqabase googlepix2structinforgraphicsvqalarge see all pix2structvision models at https huggingface comodels filterpix2struct adapted from transformers models t5 modelingt5 t5layernorm with t5pix2struct class pix2structlayernormnn module def initself hiddensize eps1e6 super init self weight nn parametertorch oneshiddensize self varianceepsilon eps def forwardself hiddenstates t5 uses a layernorm which only scales and doesn t shift which is also known as root mean square layer normalization https arxiv orgabs1910 07467 thus varience is calculated wo mean and there is no bias additionally we want to make sure that the accumulation for halfprecision inputs is done in fp32 variance hiddenstates totorch float32 pow2 mean1 keepdimtrue hiddenstates hiddenstates torch rsqrtvariance self varianceepsilon convert into halfprecision if necessary if self weight dtype in torch float16 torch bfloat16 hiddenstates hiddenstates toself weight dtype return self weight hiddenstates try from apex normalization import fusedrmsnorm pix2structlayernorm fusedrmsnorm noqa logger infodiscovered apex normalization fusedrmsnorm will use it instead of pix2structlayernorm except importerror using the normal pix2structlayernorm pass except exception logger warningdiscovered apex but it failed to load falling back to pix2structlayernorm pass alllayernormlayers appendpix2structlayernorm class pix2structvisionembeddingsnn module r construct the embeddings from patch in pix2struct the input is different from classic visiontransformer models here the input is a sequence of seqlen flattened patches that also combines padding patches tokens each patch is represented by a vector of hiddensize values the row and column indices are stored in the first and second position of the flattenedpatches flattenedpatches batchsize seqlen hiddensize 2 sum all embeddings together mesh tensorflow initialization to avoid scaling before softmax selfattention block input is batchsize seqlength dim mask is batchsize keylength noncausal or batchsize keylength keylength pastkeyvalue0 is batchsize nheads qlen 1 dimperhead projection return states contiguous viewbatchsize 1 self nheads self keyvalueprojdim transpose1 2 get query states batchsize nheads seqlength dimperhead querystates toprojectionshapeself queryhiddenstates get keyvalue states keystates toprojectionshapeself keyhiddenstates valuestates toprojectionshapeself valuehiddenstates compute scores equivalent of torch einsumbnqd bnkdbnqk querystates keystates compatible with onnx op9 scores torch matmulquerystates keystates transpose3 2 if positionbias is none positionbias torch zeros 1 self nheads seqlength seqlength devicescores device dtypescores dtype if self gradientcheckpointing and self training positionbias requiresgrad true if attentionmask is none attentionmask torch onesbatchsize seqlength devicescores device dtypescores dtype if attentionmask dim 2 positionbias positionbias attentionmask none none topositionbias device else batchsize nheads seqlength keylength positionbias positionbias attentionmask topositionbias device positionbias 1 positionbias positionbiasmasked positionbias maskedfillpositionbias 1 torch finfoscores dtype min scores positionbiasmasked scores torch maxscores torch tensortorch finfoscores dtype min batchsize nheads seqlength keylength attnweights nn functional softmaxscores dim1 dtypetorch float32 typeasscores batchsize nheads seqlength keylength attnweights nn functional dropoutattnweights pself dropout trainingself training mask heads if we want to if layerheadmask is not none attnweights attnweights layerheadmask attnoutput torch matmulattnweights valuestates batchsize seqlength dim attnoutput attnoutput transpose1 2 contiguous viewbatchsize 1 self innerdim attnoutput self outputattnoutput outputs attnoutput positionbias if outputattentions outputs outputs attnweights return outputs copied from transformers models t5 modelingt5 t5densegatedactdense with t5densegatedactdensepix2structvisionmlp t5configpix2structvisionconfig config dmodelconfig hiddensize dropoutratedropoutrate class pix2structvisionmlpnn module def initself config pix2structvisionconfig super init self wi0 nn linearconfig hiddensize config dff biasfalse self wi1 nn linearconfig hiddensize config dff biasfalse self wo nn linearconfig dff config hiddensize biasfalse self dropout nn dropoutconfig dropoutrate self act act2fnconfig denseactfn def forwardself hiddenstates hiddengelu self actself wi0hiddenstates hiddenlinear self wi1hiddenstates hiddenstates hiddengelu hiddenlinear hiddenstates self dropouthiddenstates to make 8bit quantization work for googleflant5xxl self wo is kept in float32 see https github comhuggingfacetransformersissues20287 we also make sure the weights are not in int8 in case users will force keepinfp32modules to be none if isinstanceself wo weight torch tensor and hiddenstates dtype self wo weight dtype and self wo weight dtype torch int8 hiddenstates hiddenstates toself wo weight dtype hiddenstates self wohiddenstates return hiddenstates class pix2structvisionlayernn module def initself config pix2structconfig none super init self chunksizefeedforward config chunksizefeedforward self seqlendim 1 self attention pix2structvisionattentionconfig self mlp pix2structvisionmlpconfig self premlplayernorm pix2structlayernormconfig hiddensize epsconfig layernormeps self preattentionlayernorm pix2structlayernormconfig hiddensize epsconfig layernormeps def forward self hiddenstates torch tensor attentionmask optionaltorch tensor none headmask optionaltorch tensor none outputattentions bool false uniontupletorch tensor torch tensor tupletorch tensor residual hiddenstates in pix2structvision layernorm is applied before selfattention hiddenstates self preattentionlayernormhiddenstates selfattentionoutputs self attention hiddenstates attentionmaskattentionmask layerheadmaskheadmask outputattentionsoutputattentions attentionoutput selfattentionoutputs0 outputs selfattentionoutputs1 add self attentions if we output attention weights first residual connection hiddenstates attentionoutput residual in pix2structvision layernorm is also applied after selfattention layeroutput self premlplayernormhiddenstates layeroutput self mlplayeroutput hiddenstates second residual connection outputs layeroutput outputs return outputs class pix2structvisionencodernn module def initself config pix2structconfig none super init self config config self layer nn modulelistpix2structvisionlayerconfig for in rangeconfig numhiddenlayers self gradientcheckpointing false def forward self hiddenstates torch tensor attentionmask optionaltorch tensor none headmask optionaltorch tensor none outputattentions bool false outputhiddenstates bool false returndict bool true uniontuple basemodeloutput allhiddenstates if outputhiddenstates else none allselfattentions if outputattentions else none for i layermodule in enumerateself layer if outputhiddenstates allhiddenstates allhiddenstates hiddenstates layerheadmask headmaski if headmask is not none else none if self gradientcheckpointing and self training layeroutputs self gradientcheckpointingfunc layermodule call hiddenstates attentionmask layerheadmask outputattentions else layeroutputs layermodulehiddenstates attentionmask layerheadmask outputattentions hiddenstates layeroutputs0 if outputattentions allselfattentions allselfattentions layeroutputs1 if outputhiddenstates allhiddenstates allhiddenstates hiddenstates if not returndict return tuplev for v in hiddenstates allhiddenstates allselfattentions if v is not none return basemodeloutput lasthiddenstatehiddenstates hiddenstatesallhiddenstates attentionsallselfattentions class pix2structpretrainedmodelpretrainedmodel configclass pix2structconfig property def dummyinputsself inputids torch tensordummyinputs inputmask torch tensordummymask dummyinputs decoderinputids inputids inputids inputids decoderattentionmask inputmask return dummyinputs def initweightsself module mesh tensorflow attention initialization to avoid scaling before softmax see https github comtensorflowmeshblobfa19d69eafc9a482aff0b59ddd96b025c0cb207dmeshtensorflowtransformerattention pyl136 upcast the input in fp32 and cast it back to desired dtype to avoid truncnormalcpu not implemented in half issues copied from transformers models t5 modelingt5 t5pretrainedmodel shiftright with t5pix2struct shift inputs to the right item assignment is not supported natively for proxies replace possible 100 values in labels by padtokenid pix2structvisioninputsdocstring r args flattenedpatches torch floattensor of shape batchsize sequencelength numchannels x patchheight x patchwidth flattened and padded pixel values these values can be obtained using autoimageprocessor see pix2structvisionimageprocessor call for details check the original paperhttps arxiv orgabs2210 03347 figure 5 for more details attentionmask torch floattensor of shape batchsize sequencelength optional mask to avoid performing attention on padding pixel values mask values selected in 0 1 headmask torch floattensor of shape numheads or numlayers numheads optional mask to nullify selected heads of the selfattention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked 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 whether or not to return a utils modeloutput instead of a plain tuple initialize weights and apply final processing prunes heads of the model headstoprune dict of layernum list of heads to prune in this layer see base class pretrainedmodel 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 usereturndict if flattenedpatches is none raise valueerroryou have to specify flattenedpatches if attentionmask is none check where flattenedpatches is not 0 attentionmask flattenedpatches sumdim1 0 float prepare head mask if needed 1 0 in headmask indicate we keep the head attentionprobs has shape bsz x nheads x n x n input headmask has shape numheads or numhiddenlayers x numheads and headmask is converted to shape numhiddenlayers x batch x numheads x seqlength x seqlength headmask self getheadmaskheadmask self config numhiddenlayers embeddingoutput self embeddingsflattenedpatches encoderoutputs self encoder embeddingoutput attentionmaskattentionmask headmaskheadmask outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict sequenceoutput encoderoutputs0 sequenceoutput self layernormsequenceoutput if not returndict headoutputs sequenceoutput return headoutputs encoderoutputs1 return basemodeloutput lasthiddenstatesequenceoutput hiddenstatesencoderoutputs hiddenstates attentionsencoderoutputs attentions copied from transformers models t5 modelingt5 t5densegatedactdense with t5pix2structtext dmodelhiddensize class pix2structtextdensegatedactdensenn module def initself config pix2structtextconfig super init self wi0 nn linearconfig hiddensize config dff biasfalse self wi1 nn linearconfig hiddensize config dff biasfalse self wo nn linearconfig dff config hiddensize biasfalse self dropout nn dropoutconfig dropoutrate self act act2fnconfig denseactfn def forwardself hiddenstates hiddengelu self actself wi0hiddenstates hiddenlinear self wi1hiddenstates hiddenstates hiddengelu hiddenlinear hiddenstates self dropouthiddenstates to make 8bit quantization work for googleflant5xxl self wo is kept in float32 see https github comhuggingfacetransformersissues20287 we also make sure the weights are not in int8 in case users will force keepinfp32modules to be none if isinstanceself wo weight torch tensor and hiddenstates dtype self wo weight dtype and self wo weight dtype torch int8 hiddenstates hiddenstates toself wo weight dtype hiddenstates self wohiddenstates return hiddenstates class pix2structtextlayerffnn module def initself config pix2structtextconfig super init self densereludense pix2structtextdensegatedactdenseconfig self layernorm pix2structlayernormconfig hiddensize epsconfig layernormepsilon self dropout nn dropoutconfig dropoutrate copied from transformers models t5 modelingt5 t5layerff forward def forwardself hiddenstates forwardedstates self layernormhiddenstates forwardedstates self densereludenseforwardedstates hiddenstates hiddenstates self dropoutforwardedstates return hiddenstates class pix2structtextattentionnn module def initself config pix2structtextconfig hasrelativeattentionbiasfalse super init self hasrelativeattentionbias hasrelativeattentionbias self relativeattentionnumbuckets config relativeattentionnumbuckets self relativeattentionmaxdistance config relativeattentionmaxdistance self hiddensize config hiddensize self keyvalueprojdim config dkv self nheads config numheads self dropout config dropoutrate self innerdim self nheads self keyvalueprojdim mesh tensorflow initialization to avoid scaling before softmax self query nn linearself hiddensize self hiddensize biasfalse self key nn linearself hiddensize self hiddensize biasfalse self value nn linearself hiddensize self hiddensize biasfalse self output nn linearself hiddensize self hiddensize biasfalse if self hasrelativeattentionbias self relativeattentionbias nn embeddingself relativeattentionnumbuckets self nheads self prunedheads set self gradientcheckpointing false staticmethod copied from transformers models t5 modelingt5 t5attention relativepositionbucket def relativepositionbucketrelativeposition bidirectionaltrue numbuckets32 maxdistance128 relativebuckets 0 if bidirectional numbuckets 2 relativebuckets relativeposition 0 totorch long numbuckets relativeposition torch absrelativeposition else relativeposition torch minrelativeposition torch zeroslikerelativeposition now relativeposition is in the range 0 inf half of the buckets are for exact increments in positions maxexact numbuckets 2 issmall relativeposition maxexact the other half of the buckets are for logarithmically bigger bins in positions up to maxdistance relativepositioniflarge maxexact torch logrelativeposition float maxexact math logmaxdistance maxexact numbuckets maxexact totorch long relativepositioniflarge torch min relativepositioniflarge torch fulllikerelativepositioniflarge numbuckets 1 relativebuckets torch whereissmall relativeposition relativepositioniflarge return relativebuckets adapted from transformers models t5 modelingt5 t5attention computebias def computebiasself querylength keylength devicenone selfattention if keyvaluestates is none or attention over source sentence provided by keyvaluestates input is batchsize seqlength dim mask is batchsize keylength noncausal or batchsize keylength keylength pastkeyvalue0 is batchsize nheads qlen 1 dimperhead projection return states contiguous viewbatchsize 1 self nheads self keyvalueprojdim transpose1 2 def projecthiddenstates projlayer keyvaluestates pastkeyvalue selfattn batchsize nheads seqlength dimperhead crossattn batchsize nheads seqlength dimperhead selfattn batchsize nheads keylength dimperhead checking that the sequencelength of the pastkeyvalue is the same as the provided keyvaluestates to support prefix tuning crossattn batchsize nheads seqlength dimperhead crossattn get query states batchsize nheads seqlength dimperhead get keyvalue states compute scores if key and values are already calculated we want only the last query position bias batchsize nheads seqlength keylength batchsize nheads seqlength keylength mask heads if we want to batchsize seqlength dim copied from transformers models t5 modelingt5 t5layerselfattention with t5layernormpix2structlayernorm t5attentionpix2structtextattention self selfattentionself attention config dmodelconfig hiddensize copied from transformers models t5 modelingt5 t5layercrossattention with t5layernormpix2structlayernorm t5attentionpix2structtextattention self encdecattentionself attention config dmodelconfig hiddensize clamp inf values to enable fp16 training the actual query length is unknown for cross attention if using past key value states need to inject it here clamp inf values to enable fp16 training combine self attn and cross attn key value states keep crossattention outputs and relative position weights apply feed forward layer clamp inf values to enable fp16 training pix2structtextinputsdocstring r args inputids torch longtensor of shape batchsize sequencelength indices of input sequence tokens in the vocabulary pix2structtext is a model with relative position embeddings so you should be able to pad the inputs on both the right and the left indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer call for detail what are input ids glossaryinputids to know more on how to prepare inputids for pretraining take a look a pix2structtext training t5training attentionmask torch floattensor 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 decoderinputids torch longtensor of shape batchsize targetsequencelength optional indices of decoder input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer call for details what are decoder input ids glossarydecoderinputids pix2structtext uses the padtokenid as the starting token for decoderinputids generation if pastkeyvalues is used optionally only the last decoderinputids have to be input see pastkeyvalues to know more on how to prepare decoderinputids for pretraining take a look at pix2structtext training t5training decoderattentionmask torch booltensor of shape batchsize targetsequencelength optional default behavior generate a tensor that ignores pad tokens in decoderinputids causal mask will also be used by default headmask torch floattensor of shape numheads or numlayers numheads optional mask to nullify selected heads of the selfattention modules in the encoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked decoderheadmask torch floattensor of shape numheads or numlayers numheads optional mask to nullify selected heads of the selfattention modules in the decoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked crossattnheadmask torch tensor of shape numheads or numlayers numheads optional mask to nullify selected heads of the crossattention modules in the decoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked encoderoutputs tupletupletorch floattensor optional tuple consists of lasthiddenstate optional hiddenstates optional attentions lasthiddenstate of shape batchsize sequencelength hiddensize is a sequence of hidden states at the output of the last layer of the encoder used in the crossattention of the decoder pastkeyvalues tupletupletorch floattensor of length config nlayers with each tuple having 4 tensors of shape batchsize numheads sequencelength 1 embedsizeperhead contains precomputed key and value hidden states of the attention layers can be used to speed up decoding if pastkeyvalues are used the user can optionally input only the last decoderinputids those that don t have their past key value states given to this model of shape batchsize 1 instead of all decoderinputids of shape batchsize sequencelength inputsembeds torch floattensor of shape batchsize sequencelength hiddensize optional optionally instead of passing inputids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert inputids indices into associated vectors than the model s internal embedding lookup matrix decoderinputsembeds torch floattensor of shape batchsize targetsequencelength hiddensize optional optionally instead of passing decoderinputids you can choose to directly pass an embedded representation if pastkeyvalues is used optionally only the last decoderinputsembeds have to be input see pastkeyvalues this is useful if you want more control over how to convert decoderinputids indices into associated vectors than the model s internal embedding lookup matrix if decoderinputids and decoderinputsembeds are both unset decoderinputsembeds takes the value of inputsembeds usecache bool optional if set to true pastkeyvalues key value states are returned and can be used to speed up decoding see pastkeyvalues 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 whether or not to return a utils modeloutput instead of a plain tuple addstartdocstrings the standalone text decoder of pix2struct pix2structstartdocstring class pix2structtextmodelpix2structpretrainedmodel configclass pix2structtextconfig nosplitmodules pix2structtextblock tiedweightskeys lmhead weight supportsgradientcheckpointing true def initself config super initconfig self embedtokens nn embeddingconfig vocabsize config hiddensize self layer nn modulelist pix2structtextblockconfig hasrelativeattentionbiasbooli 0 for i in rangeconfig numlayers self finallayernorm pix2structlayernormconfig hiddensize epsconfig layernormepsilon self dropout nn dropoutconfig dropoutrate self lmhead nn linearconfig hiddensize config vocabsize biasfalse initialize weights and apply final processing self postinit self gradientcheckpointing false copied from transformers models t5 modelingt5 t5pretrainedmodel reordercache def reordercacheself pastkeyvalues beamidx if decoder past is not included in output speedy decoding is disabled and no need to reorder if pastkeyvalues is none logger warningyou might want to consider setting usecachetrue to speed up decoding return pastkeyvalues reordereddecoderpast for layerpaststates in pastkeyvalues get the correct batch idx from layer past batch dim batch dim of past is at 2nd position reorderedlayerpaststates for layerpaststate in layerpaststates need to set correct past for each of the four key value states reorderedlayerpaststates reorderedlayerpaststates layerpaststate indexselect0 beamidx tolayerpaststate device if reorderedlayerpaststates0 shape layerpaststates0 shape raise valueerror freorderedlayerpaststates0 shape reorderedlayerpaststates0 shape and layerpaststates0 shape layerpaststates0 shape mismatched if lenreorderedlayerpaststates lenlayerpaststates raise valueerror flength of reorderedlayerpaststates lenreorderedlayerpaststates and length of layerpaststates lenlayerpaststates mismatched reordereddecoderpast reordereddecoderpast reorderedlayerpaststates return reordereddecoderpast def getinputembeddingsself return self embedtokens def setinputembeddingsself newembeddings self embedtokens newembeddings def getoutputembeddingsself return self lmhead def setoutputembeddingsself newembeddings self lmhead newembeddings addstartdocstringstomodelforwardpix2structtextinputsdocstring replacereturndocstringsoutputtypecausallmoutputwithcrossattentions configclassconfigfordoc def forward self inputids optionaltorch longtensor none attentionmask optionaltorch floattensor none encoderhiddenstates optionaltorch floattensor none encoderattentionmask optionaltorch floattensor none inputsembeds optionaltorch longtensor none headmask optionaltorch floattensor none crossattnheadmask optionaltorch tensor none pastkeyvalues optionaltupletupletorch floattensor none usecache optionalbool none outputattentions optionalbool none outputhiddenstates optionalbool none labels optionaltorch longtensor none returndict optionalbool none kwargs uniontupletorch floattensor causallmoutputwithcrossattentions r returns example python from transformers import autoprocessor pix2structtextmodel processor autoprocessor frompretrainedgooglepix2structtextcapsbase model pix2structtextmodel frompretrainedgooglepix2structtextcapsbase inputs processortexthello my dog is cute returntensorspt outputs modelinputs loss outputs loss required mask seq length can be calculated via length of past initialize pastkeyvalues with none if past does not exist we can provide a selfattention mask of dimensions batchsize fromseqlength toseqlength ourselves in which case we just need to make it broadcastable to all heads if a 2d or 3d attention mask is provided for the crossattention we need to make broadcastable to batchsize numheads seqlength seqlength prepare head mask if needed layeroutputs is a tuple with hiddenstates keyvaluestates selfattention position bias selfattention weights crossattention position bias crossattention weights we share the position biases between the layers the first layer store them layeroutputs hiddenstates keyvaluestates selfattention position bias selfattention weights crossattention position bias crossattention weights append next layer key value states add last layer move labels to correct device to enable model parallelism initialize weights and apply final processing update vocab size encode if needed training first prediction pass get decoder inputs from shifting lm labels to the right always attend to the first token decode cut decoderinputids if pastkeyvalues is used some generation methods already pass only the last input id default to old behavior keep only final id coding utf 8 2023 the huggingface inc google 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 pix2struct modeling file general docstring see all pix2structvision models at https huggingface co models filter pix2struct adapted from transformers models t5 modeling_t5 t5layernorm with t5 pix2struct construct a layernorm module in the t5 style no bias and no subtraction of mean t5 uses a layer_norm which only scales and doesn t shift which is also known as root mean square layer normalization https arxiv org abs 1910 07467 thus varience is calculated w o mean and there is no bias additionally we want to make sure that the accumulation for half precision inputs is done in fp32 convert into half precision if necessary noqa using the normal pix2structlayernorm construct the embeddings from patch in pix2struct the input is different from classic vision transformer models here the input is a sequence of seq_len flattened patches that also combines padding patches tokens each patch is represented by a vector of hidden_size values the row and column indices are stored in the first and second position of the flattened_patches flattened_patches batch_size seq_len hidden_size 2 sum all embeddings together mesh tensorflow initialization to avoid scaling before softmax self attention block input is batch_size seq_length dim mask is batch_size key_length non causal or batch_size key_length key_length past_key_value 0 is batch_size n_heads q_len 1 dim_per_head projection get query states batch_size n_heads seq_length dim_per_head get key value states compute scores equivalent of torch einsum bnqd bnkd bnqk query_states key_states compatible with onnx op 9 batch_size n_heads seq_length key_length batch_size n_heads seq_length key_length batch_size n_heads seq_length key_length mask heads if we want to batch_size seq_length dim copied from transformers models t5 modeling_t5 t5densegatedactdense with t5densegatedactdense pix2structvisionmlp t5config pix2structvisionconfig config d_model config hidden_size dropout_rate dropout_rate to make 8bit quantization work for google flan t5 xxl self wo is kept in float32 see https github com huggingface transformers issues 20287 we also make sure the weights are not in int8 in case users will force _keep_in_fp32_modules to be none in pix2structvision layernorm is applied before self attention add self attentions if we output attention weights first residual connection in pix2structvision layernorm is also applied after self attention second residual connection an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights used for testing weights initialization mesh tensorflow attention initialization to avoid scaling before softmax see https github com tensorflow mesh blob fa19d69eafc9a482aff0b59ddd96b025c0cb207d mesh_tensorflow transformer attention py l136 upcast the input in fp32 and cast it back to desired dtype to avoid trunc_normal_cpu not implemented in half issues copied from transformers models t5 modeling_t5 t5pretrainedmodel _shift_right with t5 pix2struct shift inputs to the right item assignment is not supported natively for proxies replace possible 100 values in labels by pad_token_id this model is a pytorch torch nn module https pytorch org docs stable nn html torch nn module subclass use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config pix2structconfig 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 pretrainedmodel from_pretrained method to load the model weights args flattened_patches torch floattensor of shape batch_size sequence_length num_channels x patch_height x patch_width flattened and padded pixel values these values can be obtained using autoimageprocessor see pix2structvisionimageprocessor __call__ for details check the original paper https arxiv org abs 2210 03347 figure 5 for more details attention_mask torch floattensor of shape batch_size sequence_length optional mask to avoid performing attention on padding pixel values mask values selected in 0 1 head_mask torch floattensor of shape num_heads or num_layers num_heads optional mask to nullify selected heads of the self attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked 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 whether or not to return a utils modeloutput instead of a plain tuple initialize weights and apply final processing prunes heads of the model heads_to_prune dict of layer_num list of heads to prune in this layer see base class pretrainedmodel returns example python import requests from pil import image from transformers import autoprocessor pix2structvisionmodel image_processor autoprocessor from_pretrained google pix2struct textcaps base model pix2structvisionmodel from_pretrained google pix2struct textcaps base url https www ilankelman org stopsigns australia jpg image image open requests get url stream true raw inputs image_processor images image return_tensors pt with torch no_grad outputs model inputs last_hidden_states outputs last_hidden_state list last_hidden_states shape 1 2048 768 check where flattened_patches is not 0 prepare head mask if needed 1 0 in head_mask indicate we keep the head attention_probs has shape bsz x n_heads x n x n input head_mask has shape num_heads or num_hidden_layers x num_heads and head_mask is converted to shape num_hidden_layers x batch x num_heads x seq_length x seq_length copied from transformers models t5 modeling_t5 t5densegatedactdense with t5 pix2structtext d_model hidden_size to make 8bit quantization work for google flan t5 xxl self wo is kept in float32 see https github com huggingface transformers issues 20287 we also make sure the weights are not in int8 in case users will force _keep_in_fp32_modules to be none copied from transformers models t5 modeling_t5 t5layerff forward mesh tensorflow initialization to avoid scaling before softmax copied from transformers models t5 modeling_t5 t5attention _relative_position_bucket adapted from mesh tensorflow https github com tensorflow mesh blob 0cb87fe07da627bf0b7e60475d59f95ed6b5be3d mesh_tensorflow transformer transformer_layers py l593 translate relative position to a bucket number for relative attention the relative position is defined as memory_position query_position i e the distance in tokens from the attending position to the attended to position if bidirectional false then positive relative positions are invalid we use smaller buckets for small absolute relative_position and larger buckets for larger absolute relative_positions all relative positions max_distance map to the same bucket all relative positions max_distance map to the same bucket this should allow for more graceful generalization to longer sequences than the model has been trained on args relative_position an int32 tensor bidirectional a boolean whether the attention is bidirectional num_buckets an integer max_distance an integer returns a tensor with the same shape as relative_position containing int32 values in the range 0 num_buckets now relative_position is in the range 0 inf half of the buckets are for exact increments in positions the other half of the buckets are for logarithmically bigger bins in positions up to max_distance adapted from transformers models t5 modeling_t5 t5attention compute_bias compute binned relative position bias shape query_length key_length shape query_length key_length shape query_length key_length num_heads shape 1 num_heads query_length key_length self attention if key_value_states is none or attention over source sentence provided by key_value_states input is batch_size seq_length dim mask is batch_size key_length non causal or batch_size key_length key_length past_key_value 0 is batch_size n_heads q_len 1 dim_per_head projection projects hidden states correctly to key query states self attn batch_size n_heads seq_length dim_per_head cross attn batch_size n_heads seq_length dim_per_head self attn batch_size n_heads key_length dim_per_head checking that the sequence_length of the past_key_value is the same as the provided key_value_states to support prefix tuning cross attn batch_size n_heads seq_length dim_per_head cross attn get query states batch_size n_heads seq_length dim_per_head get key value states compute scores equivalent of torch einsum bnqd bnkd bnqk query_states key_states compatible with onnx op 9 if key and values are already calculated we want only the last query position bias batch_size n_heads seq_length key_length batch_size n_heads seq_length key_length batch_size n_heads seq_length key_length mask heads if we want to batch_size seq_length dim copied from transformers models t5 modeling_t5 t5layerselfattention with t5layernorm pix2structlayernorm t5attention pix2structtextattention self selfattention self attention config d_model config hidden_size add attentions if we output them copied from transformers models t5 modeling_t5 t5layercrossattention with t5layernorm pix2structlayernorm t5attention pix2structtextattention self encdecattention self attention config d_model config hidden_size add attentions if we output them keep self attention outputs and relative position weights clamp inf values to enable fp16 training the actual query length is unknown for cross attention if using past key value states need to inject it here clamp inf values to enable fp16 training combine self attn and cross attn key value states keep cross attention outputs and relative position weights apply feed forward layer clamp inf values to enable fp16 training the pix2struct model was proposed in pix2struct screenshot parsing as pretraining for visual language understanding https arxiv org abs 2210 03347 by kenton lee mandar joshi iulia turc hexiang hu fangyu liu julian eisenschlos urvashi khandelwal peter shaw ming wei chang kristina toutanova it s an encoder decoder transformer pre trained in a image to text setting this model inherits from pretrainedmodel 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 pytorch torch nn module https pytorch org docs stable nn html torch nn module subclass use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config union pix2structconfig pix2structtextconfig 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 pretrainedmodel from_pretrained method to load the model weights args input_ids torch longtensor of shape batch_size sequence_length indices of input sequence tokens in the vocabulary pix2structtext is a model with relative position embeddings so you should be able to pad the inputs on both the right and the left indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for detail what are input ids glossary input ids to know more on how to prepare input_ids for pretraining take a look a pix2structtext training t5 training attention_mask torch floattensor 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 torch longtensor of shape batch_size target_sequence_length optional indices of decoder input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are decoder input ids glossary decoder input ids pix2structtext uses the pad_token_id as the starting token for decoder_input_ids generation if past_key_values is used optionally only the last decoder_input_ids have to be input see past_key_values to know more on how to prepare decoder_input_ids for pretraining take a look at pix2structtext training t5 training decoder_attention_mask torch booltensor 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 head_mask torch floattensor of shape num_heads or num_layers num_heads optional mask to nullify selected heads of the self attention modules in the encoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked decoder_head_mask torch floattensor of shape num_heads or num_layers num_heads optional mask to nullify selected heads of the self attention modules in the decoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked cross_attn_head_mask torch tensor of shape num_heads or num_layers num_heads optional mask to nullify selected heads of the cross attention modules in the decoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked encoder_outputs tuple tuple torch floattensor optional tuple consists of last_hidden_state optional hidden_states optional attentions last_hidden_state of shape batch_size sequence_length hidden_size is a sequence of hidden states at the output of the last layer of the encoder used in the cross attention of the decoder past_key_values tuple tuple torch floattensor of length config n_layers with each tuple having 4 tensors of shape batch_size num_heads sequence_length 1 embed_size_per_head contains precomputed key and value hidden states of the attention layers can be used to speed up decoding if past_key_values are used the user can optionally input only the last decoder_input_ids those that don t have their past key value states given to this model of shape batch_size 1 instead of all decoder_input_ids of shape batch_size sequence_length inputs_embeds torch floattensor of shape batch_size sequence_length hidden_size optional optionally instead of passing input_ids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert input_ids indices into associated vectors than the model s internal embedding lookup matrix decoder_inputs_embeds torch floattensor of shape batch_size target_sequence_length hidden_size optional optionally instead of passing decoder_input_ids you can choose to directly pass an embedded representation if past_key_values is used optionally only the last decoder_inputs_embeds have to be input see past_key_values this is useful if you want more control over how to convert decoder_input_ids indices into associated vectors than the model s internal embedding lookup matrix if decoder_input_ids and decoder_inputs_embeds are both unset decoder_inputs_embeds takes the value of inputs_embeds use_cache bool optional if set to true past_key_values key value states are returned and can be used to speed up decoding see past_key_values 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 whether or not to return a utils modeloutput instead of a plain tuple args flattened_patches torch floattensor of shape batch_size seq_length hidden_size flattened pixel patches the hidden_size is obtained by the following formula hidden_size num_channels patch_size patch_size the process of flattening the pixel patches is done by pix2structprocessor attention_mask torch floattensor 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 torch longtensor of shape batch_size target_sequence_length optional indices of decoder input sequence tokens in the vocabulary indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details what are decoder input ids glossary decoder input ids pix2structtext uses the pad_token_id as the starting token for decoder_input_ids generation if past_key_values is used optionally only the last decoder_input_ids have to be input see past_key_values to know more on how to prepare decoder_input_ids for pretraining take a look at pix2structtext training t5 training decoder_attention_mask torch booltensor 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 head_mask torch floattensor of shape num_heads or num_layers num_heads optional mask to nullify selected heads of the self attention modules in the encoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked decoder_head_mask torch floattensor of shape num_heads or num_layers num_heads optional mask to nullify selected heads of the self attention modules in the decoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked cross_attn_head_mask torch tensor of shape num_heads or num_layers num_heads optional mask to nullify selected heads of the cross attention modules in the decoder mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked encoder_outputs tuple tuple torch floattensor optional tuple consists of last_hidden_state optional hidden_states optional attentions last_hidden_state of shape batch_size sequence_length hidden_size is a sequence of hidden states at the output of the last layer of the encoder used in the cross attention of the decoder past_key_values tuple tuple torch floattensor of length config n_layers with each tuple having 4 tensors of shape batch_size num_heads sequence_length 1 embed_size_per_head contains precomputed key and value hidden states of the attention layers can be used to speed up decoding if past_key_values are used the user can optionally input only the last decoder_input_ids those that don t have their past key value states given to this model of shape batch_size 1 instead of all decoder_input_ids of shape batch_size sequence_length decoder_inputs_embeds torch floattensor of shape batch_size target_sequence_length hidden_size optional optionally instead of passing decoder_input_ids you can choose to directly pass an embedded representation if past_key_values is used optionally only the last decoder_inputs_embeds have to be input see past_key_values this is useful if you want more control over how to convert decoder_input_ids indices into associated vectors than the model s internal embedding lookup matrix if decoder_input_ids and decoder_inputs_embeds are both unset decoder_inputs_embeds takes the value of inputs_embeds labels torch longtensor of shape batch_size sequence_length optional labels for computing the masked language modeling loss for the decoder use_cache bool optional if set to true past_key_values key value states are returned and can be used to speed up decoding see past_key_values 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 whether or not to return a utils modeloutput instead of a plain tuple initialize weights and apply final processing copied from transformers models t5 modeling_t5 t5pretrainedmodel _reorder_cache if decoder past is not included in output speedy decoding is disabled and no need to reorder get the correct batch idx from layer past batch dim batch dim of past is at 2nd position need to set correct past for each of the four key value states returns example python from transformers import autoprocessor pix2structtextmodel processor autoprocessor from_pretrained google pix2struct textcaps base model pix2structtextmodel from_pretrained google pix2struct textcaps base inputs processor text hello my dog is cute return_tensors pt outputs model inputs loss outputs loss required mask seq length can be calculated via length of past initialize past_key_values with none if past does not exist we can provide a self attention mask of dimensions batch_size from_seq_length to_seq_length ourselves in which case we just need to make it broadcastable to all heads if a 2d or 3d attention mask is provided for the cross attention we need to make broadcastable to batch_size num_heads seq_length seq_length prepare head mask if needed past_key_value is always none with gradient checkpointing layer_outputs is a tuple with hidden states key value states self attention position bias self attention weights cross attention position bias cross attention weights we share the position biases between the layers the first layer store them layer_outputs hidden states key value states self attention position bias self attention weights cross attention position bias cross attention weights append next layer key value states add last layer move labels to correct device to enable model parallelism initialize weights and apply final processing update vocab size returns example inference python from pil import image import requests from transformers import autoprocessor pix2structforconditionalgeneration processor autoprocessor from_pretrained google pix2struct textcaps base model pix2structforconditionalgeneration from_pretrained google pix2struct textcaps base url https www ilankelman org stopsigns australia jpg image image open requests get url stream true raw inputs processor images image return_tensors pt autoregressive generation generated_ids model generate inputs max_new_tokens 50 generated_text processor batch_decode generated_ids skip_special_tokens true 0 print generated_text a stop sign is on a street corner conditional generation text a picture of inputs processor text text images image return_tensors pt add_special_tokens false generated_ids model generate inputs max_new_tokens 50 generated_text processor batch_decode generated_ids skip_special_tokens true 0 print generated_text a picture of a stop sign with a red stop sign training python from pil import image import requests from transformers import autoprocessor pix2structforconditionalgeneration processor autoprocessor from_pretrained google pix2struct base model pix2structforconditionalgeneration from_pretrained google pix2struct base url https www ilankelman org stopsigns australia jpg image image open requests get url stream true raw text a stop sign is on the street corner inputs processor images image return_tensors pt labels processor text text return_tensors pt input_ids forward pass outputs model inputs labels labels loss outputs loss print f loss item 5f 5 94282 encode if needed training first prediction pass get decoder inputs from shifting lm labels to the right always attend to the first token decode cut decoder_input_ids if past_key_values is used some generation methods already pass only the last input id default to old behavior keep only final id
import math
from typing import Dict, List, Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn

from ...activations import ACT2FN
from ...modeling_outputs import (
    BaseModelOutput,
    BaseModelOutputWithPooling,
    CausalLMOutputWithCrossAttentions,
    Seq2SeqLMOutput,
    Seq2SeqModelOutput,
)
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import ALL_LAYERNORM_LAYERS
from ...utils import (
    DUMMY_INPUTS,
    DUMMY_MASK,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    is_torch_fx_proxy,
    logging,
    replace_return_docstrings,
)
from .configuration_pix2struct import Pix2StructConfig, Pix2StructTextConfig, Pix2StructVisionConfig


logger = logging.get_logger(__name__)


_CONFIG_FOR_DOC = "Pix2StructConfig"


PIX2STRUCT_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "google/pix2struct-textcaps-base",
    "google/pix2struct-textcaps-large",
    "google/pix2struct-base",
    "google/pix2struct-large",
    "google/pix2struct-ai2d-base",
    "google/pix2struct-ai2d-large",
    "google/pix2struct-widget-captioning-base",
    "google/pix2struct-widget-captioning-large",
    "google/pix2struct-screen2words-base",
    "google/pix2struct-screen2words-large",
    "google/pix2struct-docvqa-base",
    "google/pix2struct-docvqa-large",
    "google/pix2struct-ocrvqa-base",
    "google/pix2struct-ocrvqa-large",
    "google/pix2struct-chartqa-base",
    "google/pix2struct-inforgraphics-vqa-base",
    "google/pix2struct-inforgraphics-vqa-large",
    
]



class Pix2StructLayerNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        
        
        
        

        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)

        
        if self.weight.dtype in [torch.float16, torch.bfloat16]:
            hidden_states = hidden_states.to(self.weight.dtype)

        return self.weight * hidden_states


try:
    from apex.normalization import FusedRMSNorm

    Pix2StructLayerNorm = FusedRMSNorm  

    logger.info("Discovered apex.normalization.FusedRMSNorm - will use it instead of Pix2StructLayerNorm")
except ImportError:
    
    pass
except Exception:
    logger.warning("Discovered apex but it failed to load, falling back to Pix2StructLayerNorm")
    pass

ALL_LAYERNORM_LAYERS.append(Pix2StructLayerNorm)


class Pix2StructVisionEmbeddings(nn.Module):
    r

    def __init__(self, config: Pix2StructConfig) -> None:
        super().__init__()
        self.patch_projection = nn.Linear(config.patch_embed_hidden_size, config.hidden_size)

        self.row_embedder = nn.Embedding(config.seq_len, config.hidden_size)
        self.column_embedder = nn.Embedding(config.seq_len, config.hidden_size)

        self.dropout = nn.Dropout(config.dropout_rate)

    def forward(self, flattened_patches: torch.Tensor) -> torch.Tensor:
        
        
        row_indices = flattened_patches[:, :, 0].long()
        col_indices = flattened_patches[:, :, 1].long()

        flattened_patches = flattened_patches[:, :, 2:]

        embeddings = self.patch_projection(flattened_patches)
        row_embeddings = self.row_embedder(row_indices)
        col_embeddings = self.column_embedder(col_indices)

        
        embeddings = embeddings + row_embeddings + col_embeddings

        embeddings = self.dropout(embeddings)

        return embeddings


class Pix2StructVisionAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.key_value_proj_dim = config.d_kv
        self.n_heads = config.num_attention_heads
        self.dropout = config.attention_dropout
        self.inner_dim = self.n_heads * self.key_value_proj_dim

        
        self.query = nn.Linear(self.hidden_size, self.inner_dim, bias=False)
        self.key = nn.Linear(self.hidden_size, self.inner_dim, bias=False)
        self.value = nn.Linear(self.hidden_size, self.inner_dim, bias=False)
        self.output = nn.Linear(self.inner_dim, self.hidden_size, bias=False)

        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        position_bias=None,
        layer_head_mask=None,
        output_attentions=False,
    ):
        
        
        
        
        batch_size, seq_length = hidden_states.shape[:2]

        def to_projection_shape(states):
            
            return states.contiguous().view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)

        
        
        query_states = to_projection_shape(self.query(hidden_states))

        
        key_states = to_projection_shape(self.key(hidden_states))
        value_states = to_projection_shape(self.value(hidden_states))

        
        
        scores = torch.matmul(query_states, key_states.transpose(3, 2))

        if position_bias is None:
            position_bias = torch.zeros(
                (1, self.n_heads, seq_length, seq_length), device=scores.device, dtype=scores.dtype
            )
            if self.gradient_checkpointing and self.training:
                position_bias.requires_grad = True

            if attention_mask is None:
                attention_mask = torch.ones((batch_size, seq_length), device=scores.device, dtype=scores.dtype)

            if attention_mask.dim() == 2:
                position_bias = position_bias + attention_mask[:, None, None, :].to(position_bias.device)
            else:
                
                position_bias = position_bias + attention_mask.to(position_bias.device)
            position_bias = 1 - position_bias

        position_bias_masked = position_bias.masked_fill(position_bias == 1, torch.finfo(scores.dtype).min)
        scores += position_bias_masked
        scores = torch.max(scores, torch.tensor(torch.finfo(scores.dtype).min))

        
        attn_weights = nn.functional.softmax(scores, dim=-1, dtype=torch.float32).type_as(scores)

        
        attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)

        
        if layer_head_mask is not None:
            attn_weights = attn_weights * layer_head_mask

        attn_output = torch.matmul(attn_weights, value_states)

        
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.inner_dim)

        attn_output = self.output(attn_output)

        outputs = (attn_output,) + (position_bias,)

        if output_attentions:
            outputs = outputs + (attn_weights,)
        return outputs



class Pix2StructVisionMlp(nn.Module):
    def __init__(self, config: Pix2StructVisionConfig):
        super().__init__()
        self.wi_0 = nn.Linear(config.hidden_size, config.d_ff, bias=False)
        self.wi_1 = nn.Linear(config.hidden_size, config.d_ff, bias=False)
        self.wo = nn.Linear(config.d_ff, config.hidden_size, bias=False)
        self.dropout = nn.Dropout(config.dropout_rate)
        self.act = ACT2FN[config.dense_act_fn]

    def forward(self, hidden_states):
        hidden_gelu = self.act(self.wi_0(hidden_states))
        hidden_linear = self.wi_1(hidden_states)
        hidden_states = hidden_gelu * hidden_linear
        hidden_states = self.dropout(hidden_states)

        
        
        
        if (
            isinstance(self.wo.weight, torch.Tensor)
            and hidden_states.dtype != self.wo.weight.dtype
            and self.wo.weight.dtype != torch.int8
        ):
            hidden_states = hidden_states.to(self.wo.weight.dtype)

        hidden_states = self.wo(hidden_states)
        return hidden_states


class Pix2StructVisionLayer(nn.Module):
    def __init__(self, config: Pix2StructConfig) -> None:
        super().__init__()
        self.chunk_size_feed_forward = config.chunk_size_feed_forward
        self.seq_len_dim = 1
        self.attention = Pix2StructVisionAttention(config)
        self.mlp = Pix2StructVisionMlp(config)
        self.pre_mlp_layer_norm = Pix2StructLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.pre_attention_layer_norm = Pix2StructLayerNorm(config.hidden_size, eps=config.layer_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
    ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
        residual = hidden_states

        
        hidden_states = self.pre_attention_layer_norm(hidden_states)

        self_attention_outputs = self.attention(
            hidden_states,
            attention_mask=attention_mask,
            layer_head_mask=head_mask,
            output_attentions=output_attentions,
        )
        attention_output = self_attention_outputs[0]
        outputs = self_attention_outputs[1:]  

        
        hidden_states = attention_output + residual

        
        layer_output = self.pre_mlp_layer_norm(hidden_states)
        layer_output = self.mlp(layer_output) + hidden_states  

        outputs = (layer_output,) + outputs

        return outputs


class Pix2StructVisionEncoder(nn.Module):
    def __init__(self, config: Pix2StructConfig) -> None:
        super().__init__()
        self.config = config
        self.layer = nn.ModuleList([Pix2StructVisionLayer(config) for _ in range(config.num_hidden_layers)])
        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
    ) -> Union[tuple, BaseModelOutput]:
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None

        for i, layer_module in enumerate(self.layer):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_head_mask = head_mask[i] if head_mask is not None else None

            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(
                    layer_module.__call__,
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    output_attentions,
                )
            else:
                layer_outputs = layer_module(hidden_states, attention_mask, layer_head_mask, output_attentions)

            hidden_states = layer_outputs[0]

            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
        return BaseModelOutput(
            last_hidden_state=hidden_states,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
        )


class Pix2StructPreTrainedModel(PreTrainedModel):
    

    config_class = Pix2StructConfig

    @property
    def dummy_inputs(self):
        input_ids = torch.tensor(DUMMY_INPUTS)
        input_mask = torch.tensor(DUMMY_MASK)
        dummy_inputs = {
            "decoder_input_ids": input_ids,
            "input_ids": input_ids,
            "decoder_attention_mask": input_mask,
        }
        return dummy_inputs

    def _init_weights(self, module):
        
        factor = self.config.initializer_factor  
        if isinstance(module, Pix2StructLayerNorm):
            module.weight.data.fill_(factor * 1.0)
        elif isinstance(module, Pix2StructTextDenseGatedActDense):
            hidden_size = (
                self.config.text_config.hidden_size
                if isinstance(self.config, Pix2StructConfig)
                else self.config.hidden_size
            )
            d_ff = self.config.text_config.d_ff if isinstance(self.config, Pix2StructConfig) else self.config.d_ff

            module.wi_0.weight.data.normal_(mean=0.0, std=factor * ((hidden_size) ** -0.5))
            if hasattr(module.wi_0, "bias") and module.wi_0.bias is not None:
                module.wi_0.bias.data.zero_()
            module.wi_1.weight.data.normal_(mean=0.0, std=factor * ((hidden_size) ** -0.5))
            if hasattr(module.wi_1, "bias") and module.wi_1.bias is not None:
                module.wi_1.bias.data.zero_()
            module.wo.weight.data.normal_(mean=0.0, std=factor * ((d_ff) ** -0.5))
            if hasattr(module.wo, "bias") and module.wo.bias is not None:
                module.wo.bias.data.zero_()
        elif isinstance(module, Pix2StructTextAttention):
            
            
            hidden_size = (
                self.config.text_config.hidden_size
                if isinstance(self.config, Pix2StructConfig)
                else self.config.hidden_size
            )
            key_value_proj_dim = (
                self.config.text_config.d_kv if isinstance(self.config, Pix2StructConfig) else self.config.hidden_size
            )
            n_heads = (
                self.config.text_config.num_heads
                if isinstance(self.config, Pix2StructConfig)
                else self.config.num_heads
            )

            module.query.weight.data.normal_(mean=0.0, std=factor * ((hidden_size * key_value_proj_dim) ** -0.5))
            module.key.weight.data.normal_(mean=0.0, std=factor * (hidden_size**-0.5))
            module.value.weight.data.normal_(mean=0.0, std=factor * (hidden_size**-0.5))
            module.output.weight.data.normal_(mean=0.0, std=factor * ((n_heads * key_value_proj_dim) ** -0.5))
            if module.has_relative_attention_bias:
                module.relative_attention_bias.weight.data.normal_(mean=0.0, std=factor * ((hidden_size) ** -0.5))
        elif isinstance(module, nn.Embedding):
            hidden_size = (
                self.config.text_config.hidden_size
                if isinstance(self.config, Pix2StructConfig)
                else self.config.hidden_size
            )

            module.weight.data.normal_(mean=0.0, std=factor * ((hidden_size) ** -0.5))
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, Pix2StructTextModel):
            hidden_size = (
                self.config.text_config.hidden_size
                if isinstance(self.config, Pix2StructConfig)
                else self.config.hidden_size
            )

            module.lm_head.weight.data.normal_(mean=0.0, std=factor * ((hidden_size) ** -0.5))
        elif isinstance(module, (nn.Linear, nn.Conv2d)):
            
            
            module.weight.data = nn.init.trunc_normal_(
                module.weight.data.to(torch.float32), mean=0.0, std=self.config.initializer_range
            ).to(module.weight.dtype)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, Pix2StructLayerNorm):
            if module.weight is not None:
                module.weight.data.fill_(1.0)
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()

    
    def _shift_right(self, input_ids):
        decoder_start_token_id = self.config.decoder_start_token_id
        pad_token_id = self.config.pad_token_id

        if decoder_start_token_id is None:
            raise ValueError(
                "self.model.config.decoder_start_token_id has to be defined. In Pix2Struct it is usually set to the pad_token_id. "
                "See Pix2Struct docs for more information."
            )

        
        if is_torch_fx_proxy(input_ids):
            
            shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), decoder_start_token_id)
            shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1)
        else:
            shifted_input_ids = input_ids.new_zeros(input_ids.shape)
            shifted_input_ids[..., 1:] = input_ids[..., :-1].clone()
            shifted_input_ids[..., 0] = decoder_start_token_id

        if pad_token_id is None:
            raise ValueError("self.model.config.pad_token_id has to be defined.")
        
        shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)

        return shifted_input_ids


PIX2STRUCT_VISION_START_DOCSTRING = r

PIX2STRUCT_VISION_INPUTS_DOCSTRING = r


@add_start_docstrings(
    "The bare Pix2StructVision Model transformer outputting raw hidden-states without any specific head on top.",
    PIX2STRUCT_VISION_START_DOCSTRING,
)
class Pix2StructVisionModel(Pix2StructPreTrainedModel):
    config_class = Pix2StructVisionConfig
    main_input_name = "flattened_patches"
    supports_gradient_checkpointing = True
    _no_split_modules = ["Pix2StructVisionLayer"]

    def __init__(self, config: Pix2StructConfig):
        super().__init__(config)
        self.config = config

        self.embeddings = Pix2StructVisionEmbeddings(config)
        self.encoder = Pix2StructVisionEncoder(config)

        self.layernorm = Pix2StructLayerNorm(config.hidden_size, eps=config.layer_norm_eps)

        
        self.post_init()

    def get_input_embeddings(self):
        return self.embeddings.patch_projection

    def _prune_heads(self, heads_to_prune: Dict[int, List[int]]) -> None:
        
        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

    @add_start_docstrings_to_model_forward(PIX2STRUCT_VISION_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        flattened_patches: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPooling]:
        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.use_return_dict

        if flattened_patches is None:
            raise ValueError("You have to specify flattened_patches")

        if attention_mask is None:
            
            attention_mask = (flattened_patches.sum(dim=-1) != 0).float()

        
        
        
        
        
        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        embedding_output = self.embeddings(flattened_patches)

        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=attention_mask,
            head_mask=head_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
        sequence_output = self.layernorm(sequence_output)

        if not return_dict:
            head_outputs = (sequence_output,)
            return head_outputs + encoder_outputs[1:]

        return BaseModelOutput(
            last_hidden_state=sequence_output,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )



class Pix2StructTextDenseGatedActDense(nn.Module):
    def __init__(self, config: Pix2StructTextConfig):
        super().__init__()
        self.wi_0 = nn.Linear(config.hidden_size, config.d_ff, bias=False)
        self.wi_1 = nn.Linear(config.hidden_size, config.d_ff, bias=False)
        self.wo = nn.Linear(config.d_ff, config.hidden_size, bias=False)
        self.dropout = nn.Dropout(config.dropout_rate)
        self.act = ACT2FN[config.dense_act_fn]

    def forward(self, hidden_states):
        hidden_gelu = self.act(self.wi_0(hidden_states))
        hidden_linear = self.wi_1(hidden_states)
        hidden_states = hidden_gelu * hidden_linear
        hidden_states = self.dropout(hidden_states)

        
        
        
        if (
            isinstance(self.wo.weight, torch.Tensor)
            and hidden_states.dtype != self.wo.weight.dtype
            and self.wo.weight.dtype != torch.int8
        ):
            hidden_states = hidden_states.to(self.wo.weight.dtype)

        hidden_states = self.wo(hidden_states)
        return hidden_states


class Pix2StructTextLayerFF(nn.Module):
    def __init__(self, config: Pix2StructTextConfig):
        super().__init__()
        self.DenseReluDense = Pix2StructTextDenseGatedActDense(config)

        self.layer_norm = Pix2StructLayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
        self.dropout = nn.Dropout(config.dropout_rate)

    
    def forward(self, hidden_states):
        forwarded_states = self.layer_norm(hidden_states)
        forwarded_states = self.DenseReluDense(forwarded_states)
        hidden_states = hidden_states + self.dropout(forwarded_states)
        return hidden_states


class Pix2StructTextAttention(nn.Module):
    def __init__(self, config: Pix2StructTextConfig, has_relative_attention_bias=False):
        super().__init__()
        self.has_relative_attention_bias = has_relative_attention_bias
        self.relative_attention_num_buckets = config.relative_attention_num_buckets
        self.relative_attention_max_distance = config.relative_attention_max_distance
        self.hidden_size = config.hidden_size
        self.key_value_proj_dim = config.d_kv
        self.n_heads = config.num_heads
        self.dropout = config.dropout_rate
        self.inner_dim = self.n_heads * self.key_value_proj_dim

        
        self.query = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        self.key = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        self.value = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        self.output = nn.Linear(self.hidden_size, self.hidden_size, bias=False)

        if self.has_relative_attention_bias:
            self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets, self.n_heads)
        self.pruned_heads = set()
        self.gradient_checkpointing = False

    @staticmethod
    
    def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128):
        
        relative_buckets = 0
        if bidirectional:
            num_buckets //= 2
            relative_buckets += (relative_position > 0).to(torch.long) * num_buckets
            relative_position = torch.abs(relative_position)
        else:
            relative_position = -torch.min(relative_position, torch.zeros_like(relative_position))
        

        
        max_exact = num_buckets // 2
        is_small = relative_position < max_exact

        
        relative_position_if_large = max_exact + (
            torch.log(relative_position.float() / max_exact)
            / math.log(max_distance / max_exact)
            * (num_buckets - max_exact)
        ).to(torch.long)
        relative_position_if_large = torch.min(
            relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1)
        )

        relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)
        return relative_buckets

    
    def compute_bias(self, query_length, key_length, device=None):
        
        if device is None:
            device = self.relative_attention_bias.weight.device
        context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
        memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
        relative_position = memory_position - context_position  
        relative_position_bucket = self._relative_position_bucket(
            relative_position,  
            bidirectional=False,
            num_buckets=self.relative_attention_num_buckets,
            max_distance=self.relative_attention_max_distance,
        )
        values = self.relative_attention_bias(relative_position_bucket)  
        values = values.permute([2, 0, 1]).unsqueeze(0)  
        return values

    def forward(
        self,
        hidden_states,
        mask=None,
        key_value_states=None,
        position_bias=None,
        past_key_value=None,
        layer_head_mask=None,
        query_length=None,
        use_cache=False,
        output_attentions=False,
    ):
        
        
        
        
        batch_size, seq_length = hidden_states.shape[:2]

        real_seq_length = seq_length

        if past_key_value is not None:
            if len(past_key_value) != 2:
                raise ValueError(
                    f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states"
                )
            real_seq_length += past_key_value[0].shape[2] if query_length is None else query_length

        key_length = real_seq_length if key_value_states is None else key_value_states.shape[1]

        def to_projection_shape(states):
            
            return states.contiguous().view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)

        def project(hidden_states, proj_layer, key_value_states, past_key_value):
            
            if key_value_states is None:
                
                
                hidden_states = to_projection_shape(proj_layer(hidden_states))
            elif past_key_value is None:
                
                
                hidden_states = to_projection_shape(proj_layer(key_value_states))

            if past_key_value is not None:
                if key_value_states is None:
                    
                    
                    hidden_states = torch.cat([past_key_value, hidden_states], dim=2)
                elif past_key_value.shape[2] != key_value_states.shape[1]:
                    
                    
                    
                    
                    hidden_states = to_projection_shape(proj_layer(key_value_states))
                else:
                    
                    hidden_states = past_key_value
            return hidden_states

        
        
        query_states = to_projection_shape(self.query(hidden_states))

        
        key_states = project(
            hidden_states, self.key, key_value_states, past_key_value[0] if past_key_value is not None else None
        )
        value_states = project(
            hidden_states, self.value, key_value_states, past_key_value[1] if past_key_value is not None else None
        )

        
        scores = torch.matmul(
            query_states, key_states.transpose(3, 2)
        )  

        if position_bias is None:
            if not self.has_relative_attention_bias:
                position_bias = torch.zeros(
                    (1, self.n_heads, real_seq_length, key_length), device=scores.device, dtype=scores.dtype
                )
                if self.gradient_checkpointing and self.training:
                    position_bias.requires_grad = True
            else:
                position_bias = self.compute_bias(real_seq_length, key_length, device=scores.device)

            
            
            if past_key_value is not None:
                position_bias = position_bias[:, :, -hidden_states.size(1) :, :]

            if mask is not None:
                position_bias = position_bias + mask  

        if self.pruned_heads:
            mask = torch.ones(position_bias.shape[1])
            mask[list(self.pruned_heads)] = 0
            position_bias_masked = position_bias[:, mask.bool()]
        else:
            position_bias_masked = position_bias

        scores += position_bias_masked
        
        attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores)

        
        attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)

        
        if layer_head_mask is not None:
            attn_weights = attn_weights * layer_head_mask

        attn_output = torch.matmul(attn_weights, value_states)
        
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.inner_dim)

        attn_output = self.output(attn_output)

        present_key_value_state = (key_states, value_states) if use_cache else None
        outputs = (attn_output,) + (present_key_value_state,) + (position_bias,)

        if output_attentions:
            outputs = outputs + (attn_weights,)
        return outputs



class Pix2StructTextLayerSelfAttention(nn.Module):
    def __init__(self, config, has_relative_attention_bias=False):
        super().__init__()
        self.attention = Pix2StructTextAttention(config, has_relative_attention_bias=has_relative_attention_bias)
        self.layer_norm = Pix2StructLayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
        self.dropout = nn.Dropout(config.dropout_rate)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        position_bias=None,
        layer_head_mask=None,
        past_key_value=None,
        use_cache=False,
        output_attentions=False,
    ):
        normed_hidden_states = self.layer_norm(hidden_states)
        attention_output = self.attention(
            normed_hidden_states,
            mask=attention_mask,
            position_bias=position_bias,
            layer_head_mask=layer_head_mask,
            past_key_value=past_key_value,
            use_cache=use_cache,
            output_attentions=output_attentions,
        )
        hidden_states = hidden_states + self.dropout(attention_output[0])
        outputs = (hidden_states,) + attention_output[1:]  
        return outputs



class Pix2StructTextLayerCrossAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attention = Pix2StructTextAttention(config, has_relative_attention_bias=False)
        self.layer_norm = Pix2StructLayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
        self.dropout = nn.Dropout(config.dropout_rate)

    def forward(
        self,
        hidden_states,
        key_value_states,
        attention_mask=None,
        position_bias=None,
        layer_head_mask=None,
        past_key_value=None,
        use_cache=False,
        query_length=None,
        output_attentions=False,
    ):
        normed_hidden_states = self.layer_norm(hidden_states)
        attention_output = self.attention(
            normed_hidden_states,
            mask=attention_mask,
            key_value_states=key_value_states,
            position_bias=position_bias,
            layer_head_mask=layer_head_mask,
            past_key_value=past_key_value,
            use_cache=use_cache,
            query_length=query_length,
            output_attentions=output_attentions,
        )
        layer_output = hidden_states + self.dropout(attention_output[0])
        outputs = (layer_output,) + attention_output[1:]  
        return outputs


class Pix2StructTextBlock(nn.Module):
    def __init__(self, config, has_relative_attention_bias=False):
        super().__init__()

        self.self_attention = Pix2StructTextLayerSelfAttention(
            config, has_relative_attention_bias=has_relative_attention_bias
        )

        self.encoder_decoder_attention = Pix2StructTextLayerCrossAttention(config)

        self.mlp = Pix2StructTextLayerFF(config)

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        position_bias=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        encoder_decoder_position_bias=None,
        layer_head_mask=None,
        cross_attn_layer_head_mask=None,
        past_key_value=None,
        use_cache=False,
        output_attentions=False,
        return_dict=True,
    ):
        if past_key_value is not None:
            expected_num_past_key_values = 2 if encoder_hidden_states is None else 4

            if len(past_key_value) != expected_num_past_key_values:
                raise ValueError(
                    f"There should be {expected_num_past_key_values} past states. "
                    f"{'2 (past / key) for cross attention. ' if expected_num_past_key_values == 4 else ''}"
                    f"Got {len(past_key_value)} past key / value states"
                )

            self_attn_past_key_value = past_key_value[:2]
            cross_attn_past_key_value = past_key_value[2:]
        else:
            self_attn_past_key_value, cross_attn_past_key_value = None, None

        self_attention_outputs = self.self_attention(
            hidden_states,
            attention_mask=attention_mask,
            position_bias=position_bias,
            layer_head_mask=layer_head_mask,
            past_key_value=self_attn_past_key_value,
            use_cache=use_cache,
            output_attentions=output_attentions,
        )
        hidden_states, present_key_value_state = self_attention_outputs[:2]
        attention_outputs = self_attention_outputs[2:]  

        
        if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any():
            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)

        do_cross_attention = encoder_hidden_states is not None
        if do_cross_attention:
            
            
            if present_key_value_state is not None:
                query_length = present_key_value_state[0].shape[2]
            else:
                query_length = None

            cross_attention_outputs = self.encoder_decoder_attention(
                hidden_states,
                key_value_states=encoder_hidden_states,
                attention_mask=encoder_attention_mask,
                position_bias=encoder_decoder_position_bias,
                layer_head_mask=cross_attn_layer_head_mask,
                past_key_value=cross_attn_past_key_value,
                query_length=query_length,
                use_cache=use_cache,
                output_attentions=output_attentions,
            )
            hidden_states = cross_attention_outputs[0]

            
            if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any():
                clamp_value = torch.finfo(hidden_states.dtype).max - 1000
                hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)

            
            if present_key_value_state is not None:
                present_key_value_state = present_key_value_state + cross_attention_outputs[1]

            
            attention_outputs = attention_outputs + cross_attention_outputs[2:]

        
        hidden_states = self.mlp(hidden_states)

        
        if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any():
            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)

        outputs = (hidden_states,)

        if use_cache:
            outputs = outputs + (present_key_value_state,) + attention_outputs
        else:
            outputs = outputs + attention_outputs

        return outputs


PIX2STRUCT_START_DOCSTRING = r

PIX2STRUCT_TEXT_INPUTS_DOCSTRING = r

PIX2STRUCT_INPUTS_DOCSTRING = r


@add_start_docstrings(
    "The standalone text decoder of Pix2Struct",
    PIX2STRUCT_START_DOCSTRING,
)
class Pix2StructTextModel(Pix2StructPreTrainedModel):
    config_class = Pix2StructTextConfig
    _no_split_modules = ["Pix2StructTextBlock"]
    _tied_weights_keys = ["lm_head.weight"]
    supports_gradient_checkpointing = True

    def __init__(self, config):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)

        self.layer = nn.ModuleList(
            [Pix2StructTextBlock(config, has_relative_attention_bias=bool(i == 0)) for i in range(config.num_layers)]
        )
        self.final_layer_norm = Pix2StructLayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
        self.dropout = nn.Dropout(config.dropout_rate)

        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        
        self.post_init()
        self.gradient_checkpointing = False

    
    def _reorder_cache(self, past_key_values, beam_idx):
        
        
        if past_key_values is None:
            logger.warning("You might want to consider setting `use_cache=True` to speed up decoding")
            return past_key_values

        reordered_decoder_past = ()
        for layer_past_states in past_key_values:
            
            
            reordered_layer_past_states = ()
            for layer_past_state in layer_past_states:
                
                reordered_layer_past_states = reordered_layer_past_states + (
                    layer_past_state.index_select(0, beam_idx.to(layer_past_state.device)),
                )

            if reordered_layer_past_states[0].shape != layer_past_states[0].shape:
                raise ValueError(
                    f"reordered_layer_past_states[0] shape {reordered_layer_past_states[0].shape} and layer_past_states[0] shape {layer_past_states[0].shape} mismatched"
                )
            if len(reordered_layer_past_states) != len(layer_past_states):
                raise ValueError(
                    f"length of reordered_layer_past_states {len(reordered_layer_past_states)} and length of layer_past_states {len(layer_past_states)} mismatched"
                )

            reordered_decoder_past = reordered_decoder_past + (reordered_layer_past_states,)
        return reordered_decoder_past

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, new_embeddings):
        self.embed_tokens = new_embeddings

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    @add_start_docstrings_to_model_forward(PIX2STRUCT_TEXT_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=CausalLMOutputWithCrossAttentions, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        inputs_embeds: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        labels: Optional[torch.LongTensor] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple[torch.FloatTensor, ...], CausalLMOutputWithCrossAttentions]:
        r
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        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.use_return_dict

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = input_ids.size()
            input_ids = input_ids.view(-1, input_shape[-1])
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")

        if inputs_embeds is None:
            assert self.embed_tokens is not None, "You have to initialize the model with valid token embeddings"
            inputs_embeds = self.embed_tokens(input_ids)

        batch_size, seq_length = input_shape

        
        mask_seq_length = past_key_values[0][0].shape[2] + seq_length if past_key_values is not None else seq_length

        if attention_mask is None:
            attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device)
        if encoder_attention_mask is None and encoder_hidden_states is not None:
            encoder_seq_length = encoder_hidden_states.shape[1]
            encoder_attention_mask = torch.ones(
                batch_size, encoder_seq_length, device=inputs_embeds.device, dtype=torch.long
            )

        
        if past_key_values is None:
            past_key_values = [None] * len(self.layer)

        
        
        extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape)

        
        
        if encoder_hidden_states is not None:
            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
            if encoder_attention_mask is None:
                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=inputs_embeds.device)
            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
        else:
            encoder_extended_attention_mask = None

        
        head_mask = self.get_head_mask(head_mask, self.config.num_layers)
        cross_attn_head_mask = self.get_head_mask(cross_attn_head_mask, self.config.num_layers)
        present_key_value_states = () if use_cache else None
        all_hidden_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None
        all_cross_attentions = () if (output_attentions) else None
        position_bias = None
        encoder_decoder_position_bias = None

        hidden_states = self.dropout(inputs_embeds)

        for i, (layer_module, past_key_value) in enumerate(zip(self.layer, past_key_values)):
            layer_head_mask = head_mask[i]
            cross_attn_layer_head_mask = cross_attn_head_mask[i]
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            if self.gradient_checkpointing and self.training:
                if use_cache:
                    logger.warning(
                        "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
                    )
                    use_cache = False
                layer_outputs = self._gradient_checkpointing_func(
                    layer_module.forward,
                    hidden_states,
                    extended_attention_mask,
                    position_bias,
                    encoder_hidden_states,
                    encoder_extended_attention_mask,
                    encoder_decoder_position_bias,
                    layer_head_mask,
                    cross_attn_layer_head_mask,
                    None,  
                    use_cache,
                    output_attentions,
                )
            else:
                layer_outputs = layer_module(
                    hidden_states,
                    attention_mask=extended_attention_mask,
                    position_bias=position_bias,
                    encoder_hidden_states=encoder_hidden_states,
                    encoder_attention_mask=encoder_extended_attention_mask,
                    encoder_decoder_position_bias=encoder_decoder_position_bias,
                    layer_head_mask=layer_head_mask,
                    cross_attn_layer_head_mask=cross_attn_layer_head_mask,
                    past_key_value=past_key_value,
                    use_cache=use_cache,
                    output_attentions=output_attentions,
                )

            
            
            if use_cache is False:
                layer_outputs = layer_outputs[:1] + (None,) + layer_outputs[1:]

            hidden_states, present_key_value_state = layer_outputs[:2]

            
            
            
            position_bias = layer_outputs[2]
            if encoder_hidden_states is not None:
                encoder_decoder_position_bias = layer_outputs[4 if output_attentions else 3]
            
            if use_cache:
                present_key_value_states = present_key_value_states + (present_key_value_state,)

            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[3],)
                if encoder_hidden_states is not None:
                    all_cross_attentions = all_cross_attentions + (layer_outputs[5],)

        hidden_states = self.final_layer_norm(hidden_states)
        hidden_states = self.dropout(hidden_states)

        logits = self.lm_head(hidden_states)

        
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        loss = None
        if labels is not None:
            
            labels = labels.to(logits.device)
            loss_fct = nn.CrossEntropyLoss(ignore_index=-100, reduction="mean")

            loss = loss_fct(logits.contiguous().view(-1, logits.size(-1)), labels.contiguous().view(-1))

        if not return_dict:
            return tuple(
                v
                for v in [
                    loss,
                    logits,
                    present_key_value_states,
                    all_hidden_states,
                    all_attentions,
                    all_cross_attentions,
                ]
                if v is not None
            )
        return CausalLMOutputWithCrossAttentions(
            loss=loss,
            logits=logits,
            past_key_values=present_key_value_states,
            hidden_states=all_hidden_states,
            attentions=all_attentions,
            cross_attentions=all_cross_attentions,
        )


@add_start_docstrings(
    "A conditional generation model with a language modeling head. Can be used for sequence generation tasks.",
    PIX2STRUCT_START_DOCSTRING,
)
class Pix2StructForConditionalGeneration(Pix2StructPreTrainedModel):
    config_class = Pix2StructConfig
    main_input_name = "flattened_patches"
    _tied_weights_keys = ["decoder.lm_head.weight"]

    def __init__(self, config: Pix2StructConfig):
        super().__init__(config)

        self.encoder = Pix2StructVisionModel(config.vision_config)
        self.decoder = Pix2StructTextModel(config.text_config)

        self.is_vqa = config.is_vqa

        
        self.post_init()

    def get_input_embeddings(self):
        return self.decoder.get_input_embeddings()

    def set_input_embeddings(self, new_embeddings):
        self.decoder.set_input_embeddings(new_embeddings)

    def get_output_embeddings(self) -> nn.Module:
        return self.decoder.get_output_embeddings()

    def set_output_embeddings(self, new_embeddings):
        self.decoder.set_output_embeddings(new_embeddings)

    def resize_token_embeddings(self, new_num_tokens: Optional[int] = None) -> nn.Embedding:
        model_embeds = self.decoder.resize_token_embeddings(new_num_tokens)

        
        self.config.text_config.vocab_size = new_num_tokens

        return model_embeds

    def get_decoder(self):
        return self.decoder

    def get_encoder(self):
        return self.encoder

    @add_start_docstrings_to_model_forward(PIX2STRUCT_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=Seq2SeqModelOutput, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        flattened_patches: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        decoder_input_ids: Optional[torch.LongTensor] = None,
        decoder_attention_mask: Optional[torch.BoolTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        decoder_head_mask: Optional[torch.FloatTensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        labels: Optional[torch.LongTensor] = None,
        decoder_inputs_embeds: Optional[torch.Tensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.FloatTensor], Seq2SeqModelOutput]:
        r
        use_cache = use_cache if use_cache is not None else self.config.text_config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        
        if encoder_outputs is None:
            encoder_outputs = self.encoder(
                flattened_patches=flattened_patches,
                attention_mask=attention_mask,
                head_mask=head_mask,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
            )
        elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
            encoder_outputs = BaseModelOutput(
                last_hidden_state=encoder_outputs[0],
                hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
                attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
            )

        hidden_states = encoder_outputs[0]

        if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None:
            
            decoder_input_ids = self._shift_right(labels)
            decoder_attention_mask = (
                decoder_attention_mask
                if decoder_attention_mask is not None
                else decoder_input_ids.ne(self.config.pad_token_id).float()
            )
            
            decoder_attention_mask[:, 0] = 1

        
        decoder_outputs = self.decoder(
            input_ids=decoder_input_ids,
            attention_mask=decoder_attention_mask,
            inputs_embeds=decoder_inputs_embeds,
            past_key_values=past_key_values,
            encoder_hidden_states=hidden_states,
            encoder_attention_mask=attention_mask,
            head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            labels=labels,
            return_dict=return_dict,
        )

        if not return_dict:
            return decoder_outputs + encoder_outputs

        return Seq2SeqLMOutput(
            loss=decoder_outputs.loss,
            logits=decoder_outputs.logits,
            past_key_values=decoder_outputs.past_key_values,
            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,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids,
        flattened_patches: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        decoder_attention_mask: Optional[torch.BoolTensor] = None,
        past_key_values=None,
        head_mask=None,
        decoder_head_mask=None,
        cross_attn_head_mask=None,
        use_cache=None,
        encoder_outputs=None,
        **kwargs,
    ):
        if decoder_attention_mask is None:
            decoder_attention_mask = torch.ones_like(input_ids).to(input_ids.device)

        
        if past_key_values is not None:
            past_length = past_key_values[0][0].shape[2]

            
            if input_ids.shape[1] > past_length:
                remove_prefix_length = past_length
            else:
                
                remove_prefix_length = input_ids.shape[1] - 1

            input_ids = input_ids[:, remove_prefix_length:]

        return {
            "flattened_patches": flattened_patches,
            "decoder_input_ids": input_ids,
            "past_key_values": past_key_values,
            "encoder_outputs": encoder_outputs,
            "attention_mask": attention_mask,
            "decoder_attention_mask": decoder_attention_mask,
            "head_mask": head_mask,
            "decoder_head_mask": decoder_head_mask,
            "cross_attn_head_mask": cross_attn_head_mask,
            "use_cache": use_cache,
        }