codingutf8 2022 the salesforce team s and the huggingface team licensed under the bsd3clause license the license you may not use this file except in compliance with the license you may obtain a copy of the license at https opensource orglicensesbsd3clause 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 adapted from https github comsalesforceblipblobmainmodelsmed pyl52 construct the embeddings from word and position embeddings def initself config super init self wordembeddings nn embeddingconfig vocabsize config hiddensize paddingidxconfig padtokenid self positionembeddings nn embeddingconfig maxpositionembeddings config hiddensize self layernorm is not snakecased to stick with tensorflow model variable name and be able to load any tensorflow checkpoint file self layernorm nn layernormconfig hiddensize epsconfig layernormeps self dropout nn dropoutconfig hiddendropoutprob positionids 1 len position emb is contiguous in memory and exported when serialized self registerbuffer positionids torch arangeconfig maxpositionembeddings expand1 1 persistentfalse self positionembeddingtype getattrconfig positionembeddingtype absolute self config config def forward self inputids optionaltorch longtensor none positionids optionaltorch longtensor none inputsembeds optionaltorch floattensor none pastkeyvalueslength int 0 torch tensor if inputids is not none inputshape inputids size else inputshape inputsembeds size 1 seqlength inputshape1 if positionids is none positionids self positionids pastkeyvalueslength seqlength pastkeyvalueslength if inputsembeds is none inputids inputids toself wordembeddings weight device inputsembeds self wordembeddingsinputids embeddings inputsembeds if self positionembeddingtype absolute positionembeddings self positionembeddingspositionids embeddings positionembeddings embeddings self layernormembeddings embeddings self dropoutembeddings return embeddings adapted from https github comsalesforceblipblobmainmodelsmed pyl97 class bliptextselfattentionnn module def initself config iscrossattention super init self config config if config hiddensize config numattentionheads 0 and not hasattrconfig embeddingsize raise valueerror the hidden size d is not a multiple of the number of attention heads d config hiddensize config numattentionheads self numattentionheads config numattentionheads self attentionheadsize intconfig hiddensize config numattentionheads self allheadsize self numattentionheads self attentionheadsize self query nn linearconfig hiddensize self allheadsize if iscrossattention self key nn linearconfig encoderhiddensize self allheadsize self value nn linearconfig encoderhiddensize self allheadsize else self key nn linearconfig hiddensize self allheadsize self value nn linearconfig hiddensize self allheadsize self dropout nn dropoutconfig attentionprobsdropoutprob self positionembeddingtype getattrconfig positionembeddingtype absolute if self positionembeddingtype relativekey or self positionembeddingtype relativekeyquery self maxpositionembeddings config maxpositionembeddings self distanceembedding nn embedding2 config maxpositionembeddings 1 self attentionheadsize def saveattngradientsself attngradients self attngradients attngradients def getattngradientsself return self attngradients def saveattentionmapself attentionmap self attentionmap attentionmap def getattentionmapself return self attentionmap def transposeforscoresself x newxshape x size 1 self numattentionheads self attentionheadsize x x viewnewxshape return x permute0 2 1 3 def forward self hiddenstates torch tensor attentionmask optionaltorch floattensor none headmask optionaltorch floattensor none encoderhiddenstates optionaltorch floattensor none encoderattentionmask optionaltorch floattensor none pastkeyvalue optionaltupletupletorch floattensor none outputattentions optionalbool false tupletorch tensor mixedquerylayer self queryhiddenstates if this is instantiated as a crossattention module the keys and values come from an encoder the attention mask needs to be such that the encoder s padding tokens are not attended to iscrossattention encoderhiddenstates is not none if iscrossattention keylayer self transposeforscoresself keyencoderhiddenstates valuelayer self transposeforscoresself valueencoderhiddenstates attentionmask encoderattentionmask elif pastkeyvalue is not none keylayer self transposeforscoresself keyhiddenstates valuelayer self transposeforscoresself valuehiddenstates keylayer torch catpastkeyvalue0 keylayer dim2 valuelayer torch catpastkeyvalue1 valuelayer dim2 else keylayer self transposeforscoresself keyhiddenstates valuelayer self transposeforscoresself valuehiddenstates querylayer self transposeforscoresmixedquerylayer pastkeyvalue keylayer valuelayer take the dot product between query and key to get the raw attention scores attentionscores torch matmulquerylayer keylayer transpose1 2 if self positionembeddingtype relativekey or self positionembeddingtype relativekeyquery seqlength hiddenstates size1 positionidsl torch arangeseqlength dtypetorch long devicehiddenstates device view1 1 positionidsr torch arangeseqlength dtypetorch long devicehiddenstates device view1 1 distance positionidsl positionidsr positionalembedding self distanceembeddingdistance self maxpositionembeddings 1 positionalembedding positionalembedding todtypequerylayer dtype fp16 compatibility if self positionembeddingtype relativekey relativepositionscores torch einsumbhld lrdbhlr querylayer positionalembedding attentionscores attentionscores relativepositionscores elif self positionembeddingtype relativekeyquery relativepositionscoresquery torch einsumbhld lrdbhlr querylayer positionalembedding relativepositionscoreskey torch einsumbhrd lrdbhlr keylayer positionalembedding attentionscores attentionscores relativepositionscoresquery relativepositionscoreskey attentionscores attentionscores math sqrtself attentionheadsize if attentionmask is not none apply the attention mask is precomputed for all layers in bliptextmodel forward function attentionscores attentionscores attentionmask toattentionscores device normalize the attention scores to probabilities attentionprobs nn softmaxdim1attentionscores this is actually dropping out entire tokens to attend to which might seem a bit unusual but is taken from the original transformer paper attentionprobsdropped self dropoutattentionprobs mask heads if we want to if headmask is not none attentionprobsdropped attentionprobsdropped headmask contextlayer torch matmulattentionprobsdropped valuelayer contextlayer contextlayer permute0 2 1 3 contiguous newcontextlayershape contextlayer size 2 self allheadsize contextlayer contextlayer viewnewcontextlayershape outputs contextlayer attentionprobs if outputattentions else contextlayer outputs outputs pastkeyvalue return outputs copied from transformers models bert modelingbert bertselfoutput with bert bliptext class bliptextselfoutputnn module def initself config super init self dense nn linearconfig hiddensize config hiddensize self layernorm nn layernormconfig hiddensize epsconfig layernormeps self dropout nn dropoutconfig hiddendropoutprob def forwardself hiddenstates torch tensor inputtensor torch tensor torch tensor hiddenstates self densehiddenstates hiddenstates self dropouthiddenstates hiddenstates self layernormhiddenstates inputtensor return hiddenstates adapted from https github comsalesforceblipblobmainmodelsmed py242 class bliptextattentionnn module def initself config iscrossattentionfalse super init self self bliptextselfattentionconfig iscrossattention self output bliptextselfoutputconfig self prunedheads set def pruneheadsself heads if lenheads 0 return heads index findpruneableheadsandindices heads self self numattentionheads self self attentionheadsize self prunedheads prune linear layers self self query prunelinearlayerself self query index self self key prunelinearlayerself self key index self self value prunelinearlayerself self value index self output dense prunelinearlayerself output dense index dim1 update hyper params and store pruned heads self self numattentionheads self self numattentionheads lenheads self self allheadsize self self attentionheadsize self self numattentionheads self prunedheads self prunedheads unionheads def forward self hiddenstates torch tensor attentionmask optionaltorch floattensor none headmask optionaltorch floattensor none encoderhiddenstates optionaltorch floattensor none encoderattentionmask optionaltorch floattensor none pastkeyvalue optionaltupletupletorch floattensor none outputattentions optionalbool false tupletorch tensor selfoutputs self self hiddenstates attentionmask headmask encoderhiddenstates encoderattentionmask pastkeyvalue outputattentions attentionoutput self outputselfoutputs0 hiddenstates outputs attentionoutput selfoutputs1 add attentions if we output them return outputs copied from transformers models bert modelingbert bertintermediate with bert bliptext class bliptextintermediatenn module def initself config super init self dense nn linearconfig hiddensize config intermediatesize if isinstanceconfig hiddenact str self intermediateactfn act2fnconfig hiddenact else self intermediateactfn config hiddenact def forwardself hiddenstates torch tensor torch tensor hiddenstates self densehiddenstates hiddenstates self intermediateactfnhiddenstates return hiddenstates copied from transformers models bert modelingbert bertoutput with bert bliptext class bliptextoutputnn module def initself config super init self dense nn linearconfig intermediatesize config hiddensize self layernorm nn layernormconfig hiddensize epsconfig layernormeps self dropout nn dropoutconfig hiddendropoutprob def forwardself hiddenstates torch tensor inputtensor torch tensor torch tensor hiddenstates self densehiddenstates hiddenstates self dropouthiddenstates hiddenstates self layernormhiddenstates inputtensor return hiddenstates class bliptextlayernn module def initself config layernum super init self config config self chunksizefeedforward config chunksizefeedforward self seqlendim 1 self attention bliptextattentionconfig self layernum layernum if self config isdecoder self crossattention bliptextattentionconfig iscrossattentionself config isdecoder self intermediate bliptextintermediateconfig self output bliptextoutputconfig def forward self hiddenstates torch tensor attentionmask optionaltorch floattensor none headmask optionaltorch floattensor none encoderhiddenstates optionaltorch floattensor none encoderattentionmask optionaltorch floattensor none pastkeyvalue optionaltupletupletorch floattensor none outputattentions optionalbool false tupletorch tensor decoder unidirectional selfattention cached keyvalues tuple is at positions 1 2 selfattnpastkeyvalue pastkeyvalue 2 if pastkeyvalue is not none else none selfattentionoutputs self attention hiddenstates attentionmask headmask outputattentionsoutputattentions pastkeyvalueselfattnpastkeyvalue attentionoutput selfattentionoutputs0 outputs selfattentionoutputs1 1 presentkeyvalue selfattentionoutputs1 if encoderhiddenstates is not none crossattentionoutputs self crossattention attentionoutput attentionmask headmask encoderhiddenstates encoderattentionmask outputattentionsoutputattentions attentionoutput crossattentionoutputs0 outputs outputs crossattentionoutputs1 1 add cross attentions if we output attention weights layeroutput applychunkingtoforward self feedforwardchunk self chunksizefeedforward self seqlendim attentionoutput outputs layeroutput outputs outputs outputs presentkeyvalue return outputs def feedforwardchunkself attentionoutput intermediateoutput self intermediateattentionoutput layeroutput self outputintermediateoutput attentionoutput return layeroutput adapted from https github comsalesforceblipblobmainmodelsmed pyl386 class bliptextencodernn module def initself config super init self config config self layer nn modulelistbliptextlayerconfig i for i in rangeconfig numhiddenlayers self gradientcheckpointing false def forward self hiddenstates torch tensor attentionmask optionaltorch floattensor none headmask optionaltorch floattensor none encoderhiddenstates optionaltorch floattensor none encoderattentionmask optionaltorch floattensor none pastkeyvalues optionaltupletupletorch floattensor none usecache optionalbool none outputattentions optionalbool false outputhiddenstates optionalbool false returndict optionalbool true uniontupletorch tensor basemodeloutputwithpastandcrossattentions if self gradientcheckpointing and self training if usecache logger warning usecachetrue is incompatible with gradient checkpointing setting usecachefalse usecache false allhiddenstates if outputhiddenstates else none allselfattentions if outputattentions else none allcrossattentions if outputattentions and self config isdecoder else none nextdecodercache if usecache else none for i in rangeself config numhiddenlayers layermodule self layeri if outputhiddenstates allhiddenstates allhiddenstates hiddenstates layerheadmask headmaski if headmask is not none else none pastkeyvalue pastkeyvaluesi if pastkeyvalues is not none else none if self gradientcheckpointing and self training layeroutputs self gradientcheckpointingfunc layermodule call hiddenstates attentionmask layerheadmask encoderhiddenstates encoderattentionmask pastkeyvalue outputattentions else layeroutputs layermodule hiddenstates attentionmask layerheadmask encoderhiddenstates encoderattentionmask pastkeyvalue outputattentions hiddenstates layeroutputs0 if usecache nextdecodercache layeroutputs1 if outputattentions allselfattentions allselfattentions layeroutputs1 allcrossattentions allcrossattentions layeroutputs2 if outputhiddenstates allhiddenstates allhiddenstates hiddenstates if not returndict return tuple v for v in hiddenstates nextdecodercache allhiddenstates allselfattentions allcrossattentions if v is not none return basemodeloutputwithpastandcrossattentions lasthiddenstatehiddenstates pastkeyvaluesnextdecodercache hiddenstatesallhiddenstates attentionsallselfattentions crossattentionsallcrossattentions copied from transformers models bert modelingbert bertpooler with bertbliptext class bliptextpoolernn module def initself config super init self dense nn linearconfig hiddensize config hiddensize self activation nn tanh def forwardself hiddenstates torch tensor torch tensor we pool the model by simply taking the hidden state corresponding to the first token firsttokentensor hiddenstates 0 pooledoutput self densefirsttokentensor pooledoutput self activationpooledoutput return pooledoutput copied from transformers models bert modelingbert bertpredictionheadtransform with bertbliptext class bliptextpredictionheadtransformnn module def initself config super init self dense nn linearconfig hiddensize config hiddensize if isinstanceconfig hiddenact str self transformactfn act2fnconfig hiddenact else self transformactfn config hiddenact self layernorm nn layernormconfig hiddensize epsconfig layernormeps def forwardself hiddenstates torch tensor torch tensor hiddenstates self densehiddenstates hiddenstates self transformactfnhiddenstates hiddenstates self layernormhiddenstates return hiddenstates copied from transformers models bert modelingbert bertlmpredictionhead with bertbliptext class bliptextlmpredictionheadnn module def initself config super init self transform bliptextpredictionheadtransformconfig the output weights are the same as the input embeddings but there is an outputonly bias for each token self decoder nn linearconfig hiddensize config vocabsize biasfalse self bias nn parametertorch zerosconfig vocabsize need a link between the two variables so that the bias is correctly resized with resizetokenembeddings self decoder bias self bias def forwardself hiddenstates hiddenstates self transformhiddenstates hiddenstates self decoderhiddenstates return hiddenstates copied from transformers models bert modelingbert bertonlymlmhead with bertbliptext class bliptextonlymlmheadnn module def initself config super init self predictions bliptextlmpredictionheadconfig def forwardself sequenceoutput torch tensor torch tensor predictionscores self predictionssequenceoutput return predictionscores adapted from https github comsalesforceblipblobmainmodelsmed pyl548 class bliptextpretrainedmodelpretrainedmodel configclass bliptextconfig basemodelprefix bert def initweightsself module slightly different from the tf version which uses truncatednormal for initialization cf https github compytorchpytorchpull5617 adapted from https github comsalesforceblipblob3a29b7410476bf5f2ba0955827390eb6ea1f4f9dmodelsmed pyl571 the model can behave as an encoder with only selfattention as well as a decoder in which case a layer of crossattention is added between the selfattention layers following the architecture described in attention is all you needhttps arxiv orgabs1706 03762 by ashish vaswani noam shazeer niki parmar jakob uszkoreit llion jones aidan n gomez lukasz kaiser and illia polosukhin argument and isdecoder set to true an encoderhiddenstates is then expected as an input to the forward pass copied from transformers models bert modelingbert bertmodel pruneheads prunes heads of the model headstoprune dict of layernum list of heads to prune in this layer see base class pretrainedmodel makes broadcastable attention and causal masks so that future and masked tokens are ignored arguments attentionmask torch tensor mask with ones indicating tokens to attend to zeros for tokens to ignore inputshape tupleint the shape of the input to the model device torch device the device of the input to the model returns torch tensor the extended attention mask with a the same dtype as attentionmask dtype 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 provided a padding mask of dimensions batchsize seqlength if the model is a decoder apply a causal mask in addition to the padding mask if the model is an encoder make the mask broadcastable to batchsize numheads seqlength seqlength in case pastkeyvalues are used we need to add a prefix ones mask to the causal mask causal and attention masks must have same type with pytorch version 1 3 since attentionmask is 1 0 for positions we want to attend and 0 0 for masked positions this operation will create a tensor which is 0 0 for positions we want to attend and 10000 0 for masked positions since we are adding it to the raw scores before the softmax this is effectively the same as removing these entirely 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 isdecoder usecache usecache if usecache is not none else self config usecache else usecache false if inputids is not none and inputsembeds is not none raise valueerroryou cannot specify both inputids and inputsembeds at the same time elif inputids is not none self warnifpaddingandnoattentionmaskinputids attentionmask inputshape inputids size batchsize seqlength inputshape device inputids device elif inputsembeds is not none inputshape inputsembeds size 1 batchsize seqlength inputshape device inputsembeds device elif encoderembeds is not none inputshape encoderembeds size 1 batchsize seqlength inputshape device encoderembeds device else raise valueerroryou have to specify either inputids or inputsembeds or encoderembeds pastkeyvalueslength pastkeyvalueslength pastkeyvalues00 shape2 if pastkeyvalues is not none else 0 if attentionmask is none attentionmask torch onesbatchsize seqlength pastkeyvalueslength todevice 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 extendedattentionmask torch tensor self getextendedattentionmask attentionmask inputshape device isdecoder if a 2d or 3d attention mask is provided for the crossattention we need to make broadcastable to batchsize numheads seqlength seqlength if encoderhiddenstates is not none if isinstanceencoderhiddenstates list encoderbatchsize encodersequencelength encoderhiddenstates0 size else encoderbatchsize encodersequencelength encoderhiddenstates size encoderhiddenshape encoderbatchsize encodersequencelength if isinstanceencoderattentionmask list encoderextendedattentionmask self invertattentionmaskmask for mask in encoderattentionmask elif encoderattentionmask is none encoderattentionmask torch onesencoderhiddenshape devicedevice encoderextendedattentionmask self invertattentionmaskencoderattentionmask else encoderextendedattentionmask self invertattentionmaskencoderattentionmask else encoderextendedattentionmask none 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 if encoderembeds is none embeddingoutput self embeddings inputidsinputids positionidspositionids inputsembedsinputsembeds pastkeyvalueslengthpastkeyvalueslength else embeddingoutput encoderembeds encoderoutputs self encoder embeddingoutput attentionmaskextendedattentionmask headmaskheadmask encoderhiddenstatesencoderhiddenstates encoderattentionmaskencoderextendedattentionmask pastkeyvaluespastkeyvalues usecacheusecache outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict sequenceoutput encoderoutputs0 pooledoutput self poolersequenceoutput if self pooler is not none else none if not returndict return sequenceoutput pooledoutput encoderoutputs1 return basemodeloutputwithpoolingandcrossattentions lasthiddenstatesequenceoutput pooleroutputpooledoutput pastkeyvaluesencoderoutputs pastkeyvalues hiddenstatesencoderoutputs hiddenstates attentionsencoderoutputs attentions crossattentionsencoderoutputs crossattentions adapted from https github comsalesforceblipblobmainmodelsmed pyl811 class bliptextlmheadmodelbliptextpretrainedmodel def initself config super initconfig self bert bliptextmodelconfig addpoolinglayerfalse self cls bliptextonlymlmheadconfig def getoutputembeddingsself return self cls predictions decoder def setoutputembeddingsself newembeddings self cls predictions decoder newembeddings def forward self inputids optionaltorch tensor none attentionmask optionaltorch tensor none positionids optionaltorch tensor none headmask optionaltorch tensor none inputsembeds optionaltorch tensor none encoderhiddenstates optionaltorch tensor none encoderattentionmask optionaltorch tensor none labels optionaltorch tensor none pastkeyvalues optionallisttorch tensor none usecache optionalbool none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none returnlogits optionalbool false isdecoder optionalbool true reduction optionalstr mean uniontupletorch tensor causallmoutputwithcrossattentions r encoderhiddenstates torch floattensor optional sequence of hiddenstates at the output of the last layer of the encoder used in the crossattention if the model is configured as a decoder encoderattentionmask torch floattensor optional mask to avoid performing attention on the padding token indices of the encoder input this mask is used in the crossattention if the model is configured as a decoder mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked labels torch longtensor optional labels for computing the lefttoright language modeling loss next word prediction indices should be in 100 0 config vocabsize see inputids docstring tokens with indices set to 100 are ignored masked the loss is only computed for the tokens with labels n 0 config vocabsize pastkeyvalues tupletupletorch floattensor optional contains precomputed key and value hidden states of the attention blocks 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 usecache bool optional if set to true pastkeyvalues key value states are returned and can be used to speed up decoding see pastkeyvalues we are doing nexttoken prediction shift prediction scores and input ids by one if model is used as a decoder in encoderdecoder model the decoder attention mask is created on the fly 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 2022 the salesforce team s and the huggingface team licensed under the bsd 3 clause license the license you may not use this file except in compliance with the license you may obtain a copy of the license at https opensource org licenses bsd 3 clause 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 adapted from https github com salesforce blip blob main models med py l52 construct the embeddings from word and position embeddings self layernorm is not snake cased to stick with tensorflow model variable name and be able to load any tensorflow checkpoint file position_ids 1 len position emb is contiguous in memory and exported when serialized adapted from https github com salesforce blip blob main models med py l97 if this is instantiated as a cross attention module the keys and values come from an encoder the attention mask needs to be such that the encoder s padding tokens are not attended to take the dot product between query and key to get the raw attention scores fp16 compatibility apply the attention mask is precomputed for all layers in bliptextmodel forward function normalize the attention scores to probabilities this is actually dropping out entire tokens to attend to which might seem a bit unusual but is taken from the original transformer paper mask heads if we want to copied from transformers models bert modeling_bert bertselfoutput with bert bliptext adapted from https github com salesforce blip blob main models med py 242 prune linear layers update hyper params and store pruned heads add attentions if we output them copied from transformers models bert modeling_bert bertintermediate with bert bliptext copied from transformers models bert modeling_bert bertoutput with bert bliptext decoder uni directional self attention cached key values tuple is at positions 1 2 add cross attentions if we output attention weights adapted from https github com salesforce blip blob main models med py l386 copied from transformers models bert modeling_bert bertpooler with bert bliptext we pool the model by simply taking the hidden state corresponding to the first token copied from transformers models bert modeling_bert bertpredictionheadtransform with bert bliptext copied from transformers models bert modeling_bert bertlmpredictionhead with bert bliptext the output weights are the same as the input embeddings but there is an output only bias for each token need a link between the two variables so that the bias is correctly resized with resize_token_embeddings copied from transformers models bert modeling_bert bertonlymlmhead with bert bliptext adapted from https github com salesforce blip blob main models med py l548 an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights slightly different from the tf version which uses truncated_normal for initialization cf https github com pytorch pytorch pull 5617 adapted from https github com salesforce blip blob 3a29b7410476bf5f2ba0955827390eb6ea1f4f9d models med py l571 the model can behave as an encoder with only self attention as well as a decoder in which case a layer of cross attention is added between the self attention layers following the architecture described in attention is all you need https arxiv org abs 1706 03762 by ashish vaswani noam shazeer niki parmar jakob uszkoreit llion jones aidan n gomez lukasz kaiser and illia polosukhin argument and is_decoder set to true an encoder_hidden_states is then expected as an input to the forward pass copied from transformers models bert modeling_bert bertmodel _prune_heads prunes heads of the model heads_to_prune dict of layer_num list of heads to prune in this layer see base class pretrainedmodel makes broadcastable attention and causal masks so that future and masked tokens are ignored arguments attention_mask torch tensor mask with ones indicating tokens to attend to zeros for tokens to ignore input_shape tuple int the shape of the input to the model device torch device the device of the input to the model returns torch tensor the extended attention mask with a the same dtype as attention_mask dtype 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 provided a padding mask of dimensions batch_size seq_length if the model is a decoder apply a causal mask in addition to the padding mask if the model is an encoder make the mask broadcastable to batch_size num_heads seq_length seq_length in case past_key_values are used we need to add a prefix ones mask to the causal mask causal and attention masks must have same type with pytorch version 1 3 since attention_mask is 1 0 for positions we want to attend and 0 0 for masked positions this operation will create a tensor which is 0 0 for positions we want to attend and 10000 0 for masked positions since we are adding it to the raw scores before the softmax this is effectively the same as removing these entirely fp16 compatibility encoder_hidden_states torch floattensor optional sequence of hidden states at the output of the last layer of the encoder used in the cross attention if the model is configured as a decoder encoder_attention_mask torch floattensor optional mask to avoid performing attention on the padding token indices of the encoder input this mask is used in the cross attention if the model is configured as a decoder mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked past_key_values tuple tuple torch floattensor optional contains precomputed key and value hidden states of the attention blocks 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 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 past_key_values_length 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 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 adapted from https github com salesforce blip blob main models med py l811 encoder_hidden_states torch floattensor optional sequence of hidden states at the output of the last layer of the encoder used in the cross attention if the model is configured as a decoder encoder_attention_mask torch floattensor optional mask to avoid performing attention on the padding token indices of the encoder input this mask is used in the cross attention if the model is configured as a decoder mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked labels torch longtensor optional labels for computing the left to right language modeling loss next word prediction indices should be in 100 0 config vocab_size see input_ids docstring tokens with indices set to 100 are ignored masked the loss is only computed for the tokens with labels n 0 config vocab_size past_key_values tuple tuple torch floattensor optional contains precomputed key and value hidden states of the attention blocks 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 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 we are doing next token prediction shift prediction scores and input ids by one if model is used as a decoder in encoder decoder model the decoder attention mask is created on the fly 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 List, Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import Tensor, device, nn
from torch.nn import CrossEntropyLoss

from ...activations import ACT2FN
from ...modeling_outputs import (
    BaseModelOutputWithPastAndCrossAttentions,
    BaseModelOutputWithPoolingAndCrossAttentions,
    CausalLMOutputWithCrossAttentions,
)
from ...modeling_utils import (
    PreTrainedModel,
    apply_chunking_to_forward,
    find_pruneable_heads_and_indices,
    prune_linear_layer,
)
from ...utils import logging
from .configuration_blip import BlipTextConfig


logger = logging.get_logger(__name__)



class BlipTextEmbeddings(nn.Module):
    

    def __init__(self, config):
        super().__init__()
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)

        
        
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

        
        self.register_buffer(
            "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
        )
        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")

        self.config = config

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        past_key_values_length: int = 0,
    ) -> torch.Tensor:
        if input_ids is not None:
            input_shape = input_ids.size()
        else:
            input_shape = inputs_embeds.size()[:-1]

        seq_length = input_shape[1]

        if position_ids is None:
            position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]

        if inputs_embeds is None:
            input_ids = input_ids.to(self.word_embeddings.weight.device)
            inputs_embeds = self.word_embeddings(input_ids)

        embeddings = inputs_embeds

        if self.position_embedding_type == "absolute":
            position_embeddings = self.position_embeddings(position_ids)
            embeddings += position_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings



class BlipTextSelfAttention(nn.Module):
    def __init__(self, config, is_cross_attention):
        super().__init__()
        self.config = config
        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention heads (%d)"
                % (config.hidden_size, config.num_attention_heads)
            )

        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        if is_cross_attention:
            self.key = nn.Linear(config.encoder_hidden_size, self.all_head_size)
            self.value = nn.Linear(config.encoder_hidden_size, self.all_head_size)
        else:
            self.key = nn.Linear(config.hidden_size, self.all_head_size)
            self.value = nn.Linear(config.hidden_size, self.all_head_size)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
            self.max_position_embeddings = config.max_position_embeddings
            self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)

    def save_attn_gradients(self, attn_gradients):
        self.attn_gradients = attn_gradients

    def get_attn_gradients(self):
        return self.attn_gradients

    def save_attention_map(self, attention_map):
        self.attention_map = attention_map

    def get_attention_map(self):
        return self.attention_map

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.Tensor]:
        mixed_query_layer = self.query(hidden_states)

        
        
        
        is_cross_attention = encoder_hidden_states is not None

        if is_cross_attention:
            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
            attention_mask = encoder_attention_mask
        elif past_key_value is not None:
            key_layer = self.transpose_for_scores(self.key(hidden_states))
            value_layer = self.transpose_for_scores(self.value(hidden_states))
            key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
            value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
        else:
            key_layer = self.transpose_for_scores(self.key(hidden_states))
            value_layer = self.transpose_for_scores(self.value(hidden_states))

        query_layer = self.transpose_for_scores(mixed_query_layer)

        past_key_value = (key_layer, value_layer)

        
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
            seq_length = hidden_states.size()[1]
            position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
            position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
            distance = position_ids_l - position_ids_r
            positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
            positional_embedding = positional_embedding.to(dtype=query_layer.dtype)  

            if self.position_embedding_type == "relative_key":
                relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
                attention_scores = attention_scores + relative_position_scores
            elif self.position_embedding_type == "relative_key_query":
                relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
                relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
                attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key

        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        if attention_mask is not None:
            
            attention_scores = attention_scores + attention_mask.to(attention_scores.device)

        
        attention_probs = nn.Softmax(dim=-1)(attention_scores)

        
        
        attention_probs_dropped = self.dropout(attention_probs)

        
        if head_mask is not None:
            attention_probs_dropped = attention_probs_dropped * head_mask

        context_layer = torch.matmul(attention_probs_dropped, value_layer)

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)

        outputs = outputs + (past_key_value,)
        return outputs



class BlipTextSelfOutput(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states



class BlipTextAttention(nn.Module):
    def __init__(self, config, is_cross_attention=False):
        super().__init__()
        self.self = BlipTextSelfAttention(config, is_cross_attention)
        self.output = BlipTextSelfOutput(config)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        heads, index = find_pruneable_heads_and_indices(
            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
        )

        
        self.self.query = prune_linear_layer(self.self.query, index)
        self.self.key = prune_linear_layer(self.self.key, index)
        self.self.value = prune_linear_layer(self.self.value, index)
        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)

        
        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
        self.pruned_heads = self.pruned_heads.union(heads)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.Tensor]:
        self_outputs = self.self(
            hidden_states,
            attention_mask,
            head_mask,
            encoder_hidden_states,
            encoder_attention_mask,
            past_key_value,
            output_attentions,
        )
        attention_output = self.output(self_outputs[0], hidden_states)
        outputs = (attention_output,) + self_outputs[1:]  
        return outputs



class BlipTextIntermediate(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
        if isinstance(config.hidden_act, str):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)
        return hidden_states



class BlipTextOutput(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states


class BlipTextLayer(nn.Module):
    def __init__(self, config, layer_num):
        super().__init__()
        self.config = config
        self.chunk_size_feed_forward = config.chunk_size_feed_forward
        self.seq_len_dim = 1
        self.attention = BlipTextAttention(config)
        self.layer_num = layer_num
        if self.config.is_decoder:
            self.crossattention = BlipTextAttention(config, is_cross_attention=self.config.is_decoder)
        self.intermediate = BlipTextIntermediate(config)
        self.output = BlipTextOutput(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.Tensor]:
        
        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
        self_attention_outputs = self.attention(
            hidden_states,
            attention_mask,
            head_mask,
            output_attentions=output_attentions,
            past_key_value=self_attn_past_key_value,
        )
        attention_output = self_attention_outputs[0]

        outputs = self_attention_outputs[1:-1]
        present_key_value = self_attention_outputs[-1]

        if encoder_hidden_states is not None:
            cross_attention_outputs = self.crossattention(
                attention_output,
                attention_mask,
                head_mask,
                encoder_hidden_states,
                encoder_attention_mask,
                output_attentions=output_attentions,
            )
            attention_output = cross_attention_outputs[0]
            outputs = outputs + cross_attention_outputs[1:-1]  
        layer_output = apply_chunking_to_forward(
            self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
        )
        outputs = (layer_output,) + outputs

        outputs = outputs + (present_key_value,)

        return outputs

    def feed_forward_chunk(self, attention_output):
        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output)
        return layer_output



class BlipTextEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.layer = nn.ModuleList([BlipTextLayer(config, i) for i in range(config.num_hidden_layers)])
        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = False,
        output_hidden_states: Optional[bool] = False,
        return_dict: Optional[bool] = True,
    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
        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
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None
        all_cross_attentions = () if output_attentions and self.config.is_decoder else None

        next_decoder_cache = () if use_cache else None

        for i in range(self.config.num_hidden_layers):
            layer_module = self.layer[i]
            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
            past_key_value = past_key_values[i] if past_key_values 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,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value,
                    output_attentions,
                )
            else:
                layer_outputs = layer_module(
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value,
                    output_attentions,
                )

            hidden_states = layer_outputs[0]
            if use_cache:
                next_decoder_cache += (layer_outputs[-1],)
            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)
                all_cross_attentions = all_cross_attentions + (layer_outputs[2],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(
                v
                for v in [
                    hidden_states,
                    next_decoder_cache,
                    all_hidden_states,
                    all_self_attentions,
                    all_cross_attentions,
                ]
                if v is not None
            )
        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=next_decoder_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
            cross_attentions=all_cross_attentions,
        )



class BlipTextPooler(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        
        
        first_token_tensor = hidden_states[:, 0]
        pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(pooled_output)
        return pooled_output



class BlipTextPredictionHeadTransform(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        if isinstance(config.hidden_act, str):
            self.transform_act_fn = ACT2FN[config.hidden_act]
        else:
            self.transform_act_fn = config.hidden_act
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.transform_act_fn(hidden_states)
        hidden_states = self.LayerNorm(hidden_states)
        return hidden_states



class BlipTextLMPredictionHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.transform = BlipTextPredictionHeadTransform(config)

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

        self.bias = nn.Parameter(torch.zeros(config.vocab_size))

        
        self.decoder.bias = self.bias

    def forward(self, hidden_states):
        hidden_states = self.transform(hidden_states)
        hidden_states = self.decoder(hidden_states)
        return hidden_states



class BlipTextOnlyMLMHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.predictions = BlipTextLMPredictionHead(config)

    def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
        prediction_scores = self.predictions(sequence_output)
        return prediction_scores



class BlipTextPreTrainedModel(PreTrainedModel):
    

    config_class = BlipTextConfig
    base_model_prefix = "bert"

    def _init_weights(self, module):
        
        if isinstance(module, (nn.Linear, nn.Embedding)):
            
            
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
        if isinstance(module, nn.Linear) and module.bias is not None:
            module.bias.data.zero_()



class BlipTextModel(BlipTextPreTrainedModel):
    

    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)
        self.config = config

        self.embeddings = BlipTextEmbeddings(config)
        self.encoder = BlipTextEncoder(config)
        self.pooler = BlipTextPooler(config) if add_pooling_layer else None

        self.post_init()

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

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    
    def _prune_heads(self, heads_to_prune):
        
        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

    def get_extended_attention_mask(
        self, attention_mask: Tensor, input_shape: Tuple[int], device: device, is_decoder: bool
    ) -> Tensor:
        
        
        
        if attention_mask.dim() == 3:
            extended_attention_mask = attention_mask[:, None, :, :]
        elif attention_mask.dim() == 2:
            
            
            
            if is_decoder:
                batch_size, seq_length = input_shape

                seq_ids = torch.arange(seq_length, device=device)
                causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None]
                
                
                causal_mask = causal_mask.to(attention_mask.dtype)

                if causal_mask.shape[1] < attention_mask.shape[1]:
                    prefix_seq_len = attention_mask.shape[1] - causal_mask.shape[1]
                    causal_mask = torch.cat(
                        [
                            torch.ones(
                                (batch_size, seq_length, prefix_seq_len), device=device, dtype=causal_mask.dtype
                            ),
                            causal_mask,
                        ],
                        axis=-1,
                    )

                extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
            else:
                extended_attention_mask = attention_mask[:, None, None, :]
        else:
            raise ValueError(
                "Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
                    input_shape, attention_mask.shape
                )
            )

        
        
        
        
        
        extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)  
        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
        return extended_attention_mask

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        encoder_embeds: Optional[torch.Tensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        is_decoder: Optional[bool] = False,
    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:
        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 is_decoder:
            use_cache = use_cache if use_cache is not None else self.config.use_cache
        else:
            use_cache = False

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
            input_shape = input_ids.size()
            batch_size, seq_length = input_shape
            device = input_ids.device
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
            batch_size, seq_length = input_shape
            device = inputs_embeds.device
        elif encoder_embeds is not None:
            input_shape = encoder_embeds.size()[:-1]
            batch_size, seq_length = input_shape
            device = encoder_embeds.device
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds or encoder_embeds")

        
        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0

        if attention_mask is None:
            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length))).to(device)

        
        
        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(
            attention_mask, input_shape, device, is_decoder
        )

        
        
        if encoder_hidden_states is not None:
            if isinstance(encoder_hidden_states, list):
                encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states[0].size()
            else:
                encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)

            if isinstance(encoder_attention_mask, list):
                encoder_extended_attention_mask = [self.invert_attention_mask(mask) for mask in encoder_attention_mask]
            elif encoder_attention_mask is None:
                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
                encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
            else:
                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_hidden_layers)

        if encoder_embeds is None:
            embedding_output = self.embeddings(
                input_ids=input_ids,
                position_ids=position_ids,
                inputs_embeds=inputs_embeds,
                past_key_values_length=past_key_values_length,
            )
        else:
            embedding_output = encoder_embeds

        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=extended_attention_mask,
            head_mask=head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_extended_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

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

        return BaseModelOutputWithPoolingAndCrossAttentions(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            past_key_values=encoder_outputs.past_key_values,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
            cross_attentions=encoder_outputs.cross_attentions,
        )



class BlipTextLMHeadModel(BlipTextPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)

        self.bert = BlipTextModel(config, add_pooling_layer=False)
        self.cls = BlipTextOnlyMLMHead(config)

    def get_output_embeddings(self):
        return self.cls.predictions.decoder

    def set_output_embeddings(self, new_embeddings):
        self.cls.predictions.decoder = new_embeddings

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[torch.Tensor]] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        return_logits: Optional[bool] = False,
        is_decoder: Optional[bool] = True,
        reduction: Optional[str] = "mean",
    ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:
        r
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        if labels is not None:
            use_cache = False

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            is_decoder=is_decoder,
        )

        sequence_output = outputs[0]
        prediction_scores = self.cls(sequence_output)

        if return_logits:
            return prediction_scores[:, :-1, :].contiguous()

        lm_loss = None
        if labels is not None:
            
            shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous()
            labels = labels[:, 1:].contiguous().to(shifted_prediction_scores.device)
            loss_fct = CrossEntropyLoss(reduction=reduction, label_smoothing=0.1)
            lm_loss = loss_fct(shifted_prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
            if reduction == "none":
                lm_loss = lm_loss.view(prediction_scores.size(0), -1).sum(1)

        if not return_dict:
            output = (prediction_scores,) + outputs[2:]
            return ((lm_loss,) + output) if lm_loss is not None else output

        return CausalLMOutputWithCrossAttentions(
            loss=lm_loss,
            logits=prediction_scores,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            cross_attentions=outputs.cross_attentions,
        )

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **model_kwargs):
        input_shape = input_ids.shape
        
        if attention_mask is None:
            attention_mask = input_ids.new_ones(input_shape)

        
        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 {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "past_key_values": past_key_values,
            "encoder_hidden_states": model_kwargs.get("encoder_hidden_states", None),
            "encoder_attention_mask": model_kwargs.get("encoder_attention_mask", None),
            "is_decoder": True,
        }

    def _reorder_cache(self, past_key_values, beam_idx):
        reordered_past = ()
        for layer_past in past_key_values:
            reordered_past += (
                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
            )
        return reordered_past