a  dy@sdZddlZddlZddlmZddlmZddlmZmZm Z m Z m Z m Z m Z ddlZddlZddlmZddlmZdd lmZmZdd lmZmZmZdd lmZmZmZmZmZm Z d d l!m"Z"m#Z#m$Z$m%Z%m&Z&e'e(Z)dZ*dZ+dZ,dZ-dZ.gdZ/dgZ0dgZ1dZ2dZ3e e&e$e%fZ4eGdddeZ5eGdddeZ6eGdddeZ7Gdddej8Z9Gdddej8Z:Gdd d ej8Z;Gd!d"d"ej8ZGd'd(d(ej8Z?Gd)d*d*ej8Z@Gd+d,d,ej8ZAGd-d.d.ej8ZBGd/d0d0ej8ZCd1ZDd2ZEd3ZFeFeEZGd4ZHeHeEZId5eEZJd6ZKeFeHeEeKZLd7eHeFd8eEZMd9ZNGd:d;d;eZOedd?d?eOZQed@eDjPdd=GdAdBdBeOZRedCeDjPdd=GdDdEdEeOZSedFeDjPdGd=GdHdIdIeOZTGdJdKdKej8ZUGdLdMdMej8ZVGdNdOdOej8ZWedPeDjPdQd=GdRdSdSeOZXGdTdUdUej8ZYGdVdWdWej8ZZGdXdYdYej8Z[GdZd[d[ej8Z\ed\eDjPdGd=eNGd]d^d^eOZ]dS)_z PyTorch FLAVA model.N) OrderedDict) dataclass)AnyDictListOptionalSetTupleUnion)nn)ACT2FN)BaseModelOutputBaseModelOutputWithPooling)PreTrainedModel find_pruneable_heads_and_indicesprune_linear_layer) ModelOutputadd_code_sample_docstringsadd_start_docstrings%add_start_docstrings_to_model_forwardloggingreplace_return_docstrings) FlavaConfigFlavaImageCodebookConfigFlavaImageConfigFlavaMultimodalConfigFlavaTextConfigzfacebook/flava-fullzfacebook/flava-image-codebookrrr)rg$(~k@c@seZdZUdZdZeejed<dZ ee ed<dZ eejed<dZ ee ed<dZ eejed<dZee ed<eed d d ZdS) FlavaModelOutputa Output from FlavaModel containing embeddings and outputs from individual encoders. Note that `image_embeddings` and `text_embeddigns` returned are similar to pooled output returned from a transformer. If you want embeddings for contrastive loss or retrieval use a FLAVA model's `image_projection` and `text_projection` layers on `image_embeddings` and `text_embeddings` respectively. Args: image_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `pixel_values` are present): The image embeddings which are basically the pooled output of [`FlavaImageModel`]. image_output (`BaseModelOutputWithPooling`, *optional*, returned when `pixel_values` are present): The output of the [`FlavaImageModel`]. text_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids` are present): The text embeddings which are basically the pooled output of [`FlavaTextModel`]. text_output (`BaseModelOutputWithPooling`, *optional*, returned when `input_ids` are present): The output of the [`FlavaTextModel`]. multimodal_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids` and `pixel_values` are present and `skip_multimodal_encoder` is `None` or `False`): The multimodal embeddings which are basically the pooled output of [`FlavaTextModel`]. multimodal_output (`BaseModelOutputWithPooling`, returned when `input_ids` and `pixel_values` are present and `skip_multimodal_encoder` is `None` or `False`): The output of the [`FlavaMultimodalModel`]. Nimage_embeddings image_outputtext_embeddings text_outputmultimodal_embeddingsmultimodal_outputreturncstfddDS)Nc3s,|]$}|dvr|n t|VqdS))r%r#r'Ngetattrto_tuple.0kselfq/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/transformers/models/flava/modeling_flava.py esz,FlavaModelOutput.to_tuple..tuplekeysr0r2r0r3r,ds zFlavaModelOutput.to_tuple)__name__ __module__ __qualname____doc__r"rtorch FloatTensor__annotations__r#rr$r%r&r'r rr,r2r2r2r3r!Es r!c@seZdZUdZdZeejed<dZ eejed<dZ eejed<dZ eejed<dZ eejed<dZ eejed<ed d d ZdS) FlavaLossesa"Class representing pretraining losses from FLAVA model Args: mim (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `mim_labels` and `pixel_values` are present, `input_ids_masked` is absent and `mim_weight` > 0.: Masked Image Modeling loss as used in BeIT calculated only for unimodal image data. mlm (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `mlm_labels` and `input_ids_masked` are present, `pixel_values` is absent and `mlm_weight` > 0.: Masked Language Modeling loss as used in BERT calculated only for unimodal text data. itm (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `itm_labels`, `input_ids_masked`, `pixel_values` are present and `itm_weight` > 0.: Image Text Matching (ITM) loss calculated for paired image-text data. Note that ITM loss is calculated on masked pairs in FLAVA. global_contrastive (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `input_ids` and `pixel_values` are present and `global_contrastive_weight` > 0.: Contrastive loss for image-text similarity similar to CLIP but calculated globally for paired image-text data. This is calculated on unmasked images and texts. mmm_image (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `mim_labels`, `pixel_values` and `input_ids_masked` are present and `mmm_image_weight` > 0.: Masked Multimodal Modeling loss's image component calculated on paired image-text data. mmm_text (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `mlm_labels`, `pixel_values` and `input_ids_masked` are present and `mmm_text_weight` > 0.: Masked Multimodal Modeling loss's text component calculated on paired image-text data. Nmimmlmitmglobal_contrastive mmm_imagemmm_textr(cCs&d}|D]}|dur d}q"q |S)NTF)values)r1all_nonevr2r2r3rGs  zFlavaLosses.all_none)r8r9r:r;r@rr<r=r>rArBrCrDrEboolrGr2r2r2r3r?ks r?c@seZdZUdZdZeejed<dZ e ed<dZ eejed<dZ ee ed<dZeejed<dZee ed<dZeejed <dZee ed <dZeejed <dZee ed <dZeejed <dZee ed<dZeejed<dZee ed<dZeejed<dZeejed<dZeejed<dZeejed<dZeejed<dZeejed<dZeejed<ee dddZ!dS)FlavaForPreTrainingOutputa Output from FlavaForPreTraining containing embeddings, and outputs from individual encoders. Note that `image_embeddings` and `text_embeddings` returned are similar to pooled output returned from a transformer. If you want embeddings for contrastive loss or retrieval use a FLAVA model's `image_projection` and `text_projection` layers on `image_embeddings` and `text_embeddings` respectively. Args: loss (`torch.FloatTensor`, *optional*, returned when `return_loss` is True): Total loss calculated for this model. loss_info (`FlavaLosses`): Detailed info for FLAVA Pretraining losses. Check `FlavaLosses` class description for the information on the keys. image_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `pixel_values` are present): The image embeddings which are basically the pooled output of [`FlavaImageModel`]. image_output (`BaseModelOutputWithPooling`, *optional*, returned when `pixel_values` are present): The output of the [`FlavaImageModel`]. text_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids` are present): The text embeddings which are basically the pooled output of [`FlavaTextModel`]. text_output (`BaseModelOutputWithPooling`, *optional*, returned when `input_ids` are present): The output of the [`FlavaTextModel`]. multimodal_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids` and `pixel_values` are present and `skip_unmasked_multimodal_encoder` is `None` or `False`): The multimodal embeddings which are basically the pooled output of [`FlavaTextModel`]. multimodal_output (`BaseModelOutputWithPooling`, returned when `input_ids` and `pixel_values` are present and `skip_unmasked_multimodal_encoder` is `None` or `False`): The output of the [`FlavaMultimodalModel`]. image_masked_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `pixel_values` are present): The image embeddings which are basically the pooled output of [`FlavaImageModel`]. Uses `bool_masked_pos` to create masked images. image_masked_output (`BaseModelOutputWithPooling`, *optional*, returned when `pixel_values` are present): The output of the [`FlavaImageModel`]. Uses `bool_masked_pos` to create masked images. text_masked_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids_masked` are present): The text embeddings which are basically the pooled output of [`FlavaTextModel`]. text_masked_output (`BaseModelOutputWithPooling`, *optional*, returned when `input_ids_masked` are present): The output of the [`FlavaTextModel`]. multimodal_masked_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids` and `pixel_values` are present): The multimodal embeddings which are basically the pooled output of [`FlavaTextModel`]. multimodal_masked_output (`BaseModelOutputWithPooling`, returned when `input_ids_masked` and `pixel_values` are present): The output of the [`FlavaMultimodalModel`]. mim_logits (`torch.FloatTensor` of shape `(batch_size, num_image_patches, image_vocab_size)` or of shape `(total_masked_patches, image_vocab_size)` , *optional*, returned when `pixel_values` are present and `input_ids_masked` are not): The logits for MIM unimodal loss. Uses `book_masked_pos` to get masked patches. The flattened output is returned when `bool_masked_pos` has some of the patches masked. mlm_logits (`torch.FloatTensor` of shape `(batch_size, text_seq_length, text_vocab_size)` or of shape `(total_masked_seq_length, text_vocab_size)`, *optional*, returned when `input_ids_masked` are present and `pixel_values` are not): The logits for MLM unimodal loss. The flattened output is returned when `input_ids_masked` has some of the tokens masked. itm_logits (`torch.FloatTensor` of shape `(batch_size, 2)`, *optional*, returned when `input_ids_masked` and `pixel_values` are present): The logits for ITM loss. Note that ITM loss is calculated on masked pairs in FLAVA. mmm_image_logits (`torch.FloatTensor` of shape `(batch_size, num_image_patches, image_vocab_size)` or of shape`(total_masked_patches, image_vocab_size)`, *optional*, returned when `pixel_values` and `input_ids_masked` are present): The logits for MMM image multimodal loss. Uses `book_masked_pos` to get masked patches. The flattened output is returned when `bool_masked_pos` has some of the patches masked. mmm_text_logits (`torch.FloatTensor` of shape `(batch_size, text_seq_length, text_vocab_size)` or of shape `(`(total_masked_seq_length, text_vocab_size)`), *optional*, returned when `pixel_values` and `input_ids_masked` are present): The logits for MMM text multimodal loss. The flattened output is returned when `input_ids_masked` has some of the tokens masked. contrastive_logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`): The scaled dot product scores between `image_embeddings` and `text_embeddings` but passed through FLAVA's `image_projection` and `text_projection` layers respectively. This represents the image-text similarity scores. This is calculated on unmasked images and texts. contrastive_logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`): The scaled dot product scores between `text_embeddings` and `image_embeddings` but passed through FLAVA's `text_projection` and `image_projection` layers respectively. This is calculated on unmasked images and texts. Nloss loss_infor"r#r$r%r&r'image_masked_embeddingsimage_masked_outputtext_masked_embeddingstext_masked_outputmultimodal_masked_embeddingsmultimodal_masked_output mim_logits mlm_logits itm_logitscontrastive_logits_per_imagecontrastive_logits_per_textmmm_image_logitsmmm_text_logitsr(cs$gdtfddDS)N)r%r#r'rPrNrRc3s,|]$}|vr|n t|VqdSNr*r-r1Ztransformer_outputsr2r3r4z5FlavaForPreTrainingOutput.to_tuple..r5r0r2r[r3r,sz"FlavaForPreTrainingOutput.to_tuple)"r8r9r:r;rKrr<r=r>rLr?r"r#rr$r%r&r'rMrNrOrPrQrRrSrTrUrVrWrXrYr rr,r2r2r2r3rJs. @ rJcsdeZdZdZd eeddfdd Zeje e ejddd Z deje ej eejd d d Z ZS)FlavaImageEmbeddingszb Construct the CLS token, position and patch embeddings. Optionally, also the mask token. FN)configuse_mask_tokenr)cst|p|j}ttdd|j|_|rFttdd|jnd|_t |j |j |j |jd|_ |j j}ttd|d|j|_t|j|_||_dS)Nr image_size patch_size num_channels embed_dim)super__init__ mask_tokenr Parameterr<zeros hidden_size cls_tokenPatchEmbeddingsrarbrcpatch_embeddings num_patchesposition_embeddingsDropouthidden_dropout_probdropoutr^)r1r^r_rn __class__r2r3rfs   zFlavaImageEmbeddings.__init__) embeddingsheightwidthr)c Cst|jdd}|jjdd}||kr4||kr4|jS|jdddf}|jddddf}|jd}||jj} ||jj} | d| d} } tjj|dtt |tt || dddd| t || t |fdd d }t| |jd kst| |jdkrBt d t| t| fd |jd |jdfd| dddd dd|}tj|d|fddS)a" This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution images. Source: https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/image_transformer.py#L174 rNrg?r ZbicubicF)Z scale_factormodeZ align_cornerszNumber of patches for images (z/) don't match the shape of position embedding ()dim)shaperor^rbr functionalZ interpolateZreshapeintmathsqrtpermute ValueErrorviewr<cat unsqueeze) r1rurvrwZnpatchnum_posZclass_pos_embedZpatch_pos_embedr~Z num_h_patchesZ num_w_patchesr2r2r3interpolate_pos_encoding s0    .(z-FlavaImageEmbeddings.interpolate_pos_encoding) pixel_valuesbool_masked_posrr)cCs|j\}}}}|j||d}|\}} } |dur|j|| d} |dkr`||dd}|d| } |d| | | }|j |dd} t j | |fdd}|r|| |||}n ||j }||}|S)N)rrxr r?rr})rrmsizergexpandr~rrZtype_asrkr<rrrorr)r1rrr batch_sizercrvrwruZseq_len_Z mask_tokensmask cls_tokensr2r2r3forward/s    zFlavaImageEmbeddings.forward)F)NF)r8r9r:r;rrIrfr<Tensorrrr BoolTensorr __classcell__r2r2rsr3r]s&r]csVeZdZdZd eeeeeeffeedfdd Zdej e ej d d d Z Z S)rlz# Image to Patch Embedding. r r r`cstt|tjjs ||f}t|tjjs6||f}|d|d|d|d}||_||_||_t j ||||d|_ dS)Nrr) kernel_sizeZstride) rerf isinstance collectionsabcIterablerarbrnr Conv2d projection)r1rarbrcrdrnrsr2r3rfXs  zPatchEmbeddings.__init__F)rrr)c Csx|j\}}}}|s\||jdks.||jdkr\td|d|d|jdd|jdd ||ddd}|S)NrrzInput image size (*z) doesn't match model (z).ry)rrarrflatten transpose)r1rrrrcrvrwxr2r2r3rkszPatchEmbeddings.forward)rrr r )F) r8r9r:r;rr r rfr<rrIrrr2r2rsr3rlSsrlcsFeZdZdZfddZdeejeejeejdddZZ S) FlavaTextEmbeddingszGConstruct the embeddings from word, position and token_type embeddings.csttj|j|j|jd|_t|j|j|_ t|j |j|_ tj |j|j d|_ t|j|_t|dd|_|dt|jd|jdtj|jtjdd d dS) N) padding_idxZepsposition_embedding_typeabsolute position_ids)rrxtoken_type_ids)dtypeF) persistent)rerfr Embedding vocab_sizerjZ pad_token_idword_embeddingsZmax_position_embeddingsroZtype_vocab_sizetoken_type_embeddings LayerNormlayer_norm_epsrprqrrr+rZregister_bufferr<arangerrirrlongr1r^rsr2r3rfzs zFlavaTextEmbeddings.__init__N input_idsrrc Cs|}|d}|dur.|jddd|f}|durt|drl|jddd|f}||d|}|}ntj|tj|jjd}| |}| |} || } |j dkr| |} | | 7} | | } || } | S)Nrrr)rdevicer)rrhasattrrrr<rirrrrrrorrr) r1rrr input_shape seq_lengthZbuffered_token_type_idsZ buffered_token_type_ids_expandedZ inputs_embedsrruror2r2r3rs&       zFlavaTextEmbeddings.forward)NNN) r8r9r:r;rfrr<rrrr2r2rsr3rws rc sxeZdZeddfdd ZejejdddZd ejeejeeje e e ejejfe ejfd d d Z Z S) FlavaSelfAttentionNr^r)cst|j|jdkr@t|ds@td|jfd|jd|j|_t|j|j|_|j|j|_t j |j|j|j d|_ t j |j|j|j d|_ t j |j|j|j d|_t |j|_dS)NrZembedding_sizezThe hidden size z4 is not a multiple of the number of attention heads .bias)rerfrjnum_attention_headsrrrattention_head_size all_head_sizer LinearZqkv_biasquerykeyvaluerpZattention_probs_dropout_probrrrrsr2r3rfs  zFlavaSelfAttention.__init__rr)cCs6|dd|j|jf}|j|}|ddddS)Nrxrryrr )rrrrr)r1rZ new_x_shaper2r2r3transpose_for_scoress z'FlavaSelfAttention.transpose_for_scoresF hidden_statesattention_mask head_maskoutput_attentionsr)cCs||}|||}|||}||}t||dd} | t|j } |durh| |} t j j | dd} t j j | dd} | | } |dur| |} t| |} | dddd} | dd|jf} | j| } |r| | fn| f} | S)Nrxr{r}rryrr )rrrrr<matmulrrrrr rZsoftmaxrrr contiguousrrr)r1rrrrZmixed_query_layerZ key_layerZ value_layerZ query_layerZattention_scoresZattention_probsZ context_layerZnew_context_layer_shapeoutputsr2r2r3rs&     zFlavaSelfAttention.forward)NNF)r8r9r:FlavaPossibleConfigsrfr<rrrrIr r rrr2r2rsr3rsrcs@eZdZdZeddfdd ZejejejdddZZ S) FlavaSelfOutputz The residual connection is defined in FlavaLayer (same as ViTLayer) instead of here (as is the case with other models), due to the layernorm applied before each block. Nrcs.tt|j|j|_t|j|_dSrZ) rerfr rrjdenserprqrrrrsr2r3rfs zFlavaSelfOutput.__init__r input_tensorr)cCs||}||}|SrZrrrr1rrr2r2r3rs  zFlavaSelfOutput.forward) r8r9r:r;rrfr<rrrr2r2rsr3rsrc sxeZdZeddfdd ZeeddddZd ej e ej e ej e e e ej ej fe ej fd d d ZZS) FlavaAttentionNrcs*tt||_t||_t|_dSrZ)rerfr attentionroutputset pruned_headsrrsr2r3rfs   zFlavaAttention.__init__)headsr)cCst|dkrdSt||jj|jj|j\}}t|jj||j_t|jj||j_t|jj ||j_ t|j j |dd|j _ |jjt||j_|jj|jj|j_ |j ||_dS)Nrrr})lenrrrrrrrrrrrrunion)r1rindexr2r2r3 prune_heads s zFlavaAttention.prune_headsFrcCs8|j||||d}||d|}|f|dd}|SN)rrrrr)rr)r1rrrrZ self_outputsattention_outputrr2r2r3rs zFlavaAttention.forward)NNF)r8r9r:rrfrrrr<rrrIr r rrr2r2rsr3rsrcs8eZdZeddfdd ZejejdddZZS)FlavaIntermediateNrcsBtt|j|j|_t|jt r6t |j|_ n|j|_ dSrZ) rerfr rrjintermediate_sizerr hidden_actstrr intermediate_act_fnrrsr2r3rf0s   zFlavaIntermediate.__init__)rr)cCs||}||}|SrZ)rrr1rr2r2r3r9s  zFlavaIntermediate.forward r8r9r:rrfr<rrrr2r2rsr3r/s rcs<eZdZeddfdd ZejejejdddZZS) FlavaOutputNrcs.tt|j|j|_t|j|_ dSrZ) rerfr rrrjrrprqrrrrsr2r3rfAs zFlavaOutput.__init__rcCs ||}||}||}|SrZrrr2r2r3rGs  zFlavaOutput.forwardrr2r2rsr3r@src sheZdZdZeddfdd Zd ejeejeeje e e ejejfe ejfddd Z Z S) FlavaLayerz?This corresponds to the Block class in the timm implementation.Nrcsbt|j|_d|_t||_t||_t||_ t j |j |j d|_t j |j |j d|_dSNrr)rerfZchunk_size_feed_forwardZ seq_len_dimrrr intermediaterrr rrjrlayernorm_beforelayernorm_afterrrsr2r3rfSs    zFlavaLayer.__init__Frc Csb|j|||||d}|d}|dd}||}||}||}|||}|f|}|Sr)rrrrr) r1rrrrZself_attention_outputsrrZ layer_outputr2r2r3r_s     zFlavaLayer.forward)NNF)r8r9r:r;rrfr<rrrIr r rrr2r2rsr3rPsrc sVeZdZeddfdd Zd ejeejeejeeee e e fddd Z Z S) FlavaEncoderNrcs:t|_tfddtjD|_d|_dS)Ncsg|] }tqSr2)rr.rr^r2r3 r\z)FlavaEncoder.__init__..F) rerfr^r Z ModuleListrangenum_hidden_layerslayergradient_checkpointingrrsrr3rfs  zFlavaEncoder.__init__FTrrrroutput_hidden_states return_dictr)cs|rdnd}rdnd}t|jD]\} } |r8||f}|durH|| nd} |jr~|jr~fdd} tjj| | ||| } n| ||| } | d}r"|| df}q"|r||f}|stdd|||fDSt|||dS) Nr2csfdd}|S)Ncsg|RSrZr2)inputs)modulerr2r3custom_forwardszKFlavaEncoder.forward..create_custom_forward..custom_forwardr2)rrr)rr3create_custom_forwardsz3FlavaEncoder.forward..create_custom_forwardrrcss|]}|dur|VqdSrZr2)r.rHr2r2r3r4r\z'FlavaEncoder.forward..)last_hidden_stater attentions) enumeraterrZtrainingr<utils checkpointr6r)r1rrrrrrZall_hidden_statesZall_self_attentionsiZ layer_moduleZlayer_head_maskrZ layer_outputsr2rr3rs2      zFlavaEncoder.forward)NNFFT)r8r9r:rrfr<rrrIr r6rrrr2r2rsr3r~s  rcs2eZdZedfdd ZejdddZZS) FlavaPoolerrcs*tt|j|j|_t|_dSrZ)rerfr rrjrZTanh activationrrsr2r3rfs zFlavaPooler.__init__)rcCs(|dddf}||}||}|SNr)rr )r1rZfirst_token_tensor pooled_outputr2r2r3rs  zFlavaPooler.forwardrr2r2rsr3r sr aD 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 ([`{config}`]): 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. a attention_mask (`torch.FloatTensor` of shape `({0})`, *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) 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. a; Args: pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See [`FlavaImageProcessor.__call__`] for details. bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, image_num_patches)`): Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). interpolate_pos_encoding (`bool`, *optional*): Whether to interpolate the pre-trained position encodings. a Args: input_ids (`torch.LongTensor` of shape `({0})`): Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*): Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, 1]`: - 0 corresponds to a *sentence A* token, - 1 corresponds to a *sentence B* token. [What are token type IDs?](../glossary#token-type-ids) z Args: hidden_states (`torch.FloatTensor` of shape `(batch_size, image_num_patches + text_seq_len, hidden_size)`): The concatenated hidden states of unimodal encoders. z Args: skip_multimodal_encoder (*bool*, *optional*): Skip any calculations for multimodal encoder. Useful if multimodal encoding is not going to be used. a Args: input_ids_masked (`torch.LongTensor` of shape `({0})`): Indices of input sequence tokens in the vocabulary. These ones are the masked version of the original task to be used with MLM. Indices can be obtained using [`AutoTokenizer`] along with [`DataCollatorForMaskedLanguageModeling`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) a image_attention_mask (`torch.FloatTensor` of shape `({1})`, *optional*): Mask to avoid performing attention on padding token indices specifically for images. 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) skip_unmasked_multimodal_encoder (*bool*, *optional*): Skip any calculations for multimodal encoder for unmasked inputs. FLAVA pretraining doesn't need unmasked multimodal embeddings or outputs as of now. mlm_labels (`torch.LongTensor` of shape `(batch_size, text_seq_len)`, *optional*): Labels for computing the left-to-right language and multimodal masked modeling loss (next word prediction). Indices should be in `[-100, 0, ..., text_config.vocab_size - 1]` (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., text_config.vocab_size - 1]`. mim_labels (`torch.LongTensor` of shape `(batch_size, image_num_patches)`, *optional*): Labels for computing the image and multimodal masked modeling loss. Indices should be in `[-100, 0, ..., image_config.vocab_size - 1]`. Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., image_config.vocab_size - 1]`. If not passed, they are generated automatically using the image codebook assigned to the model. By default, it uses [`FlavaImageCodebook`]. See [`FlavaImageCodebook`] to understand how to generate mim_labels. itm_labels (`torch.LongTensor` of shape `(batch_size, 1)`, *optional*): Labels for computing the image-text matching loss. 0 means the pairs don't match and 1 means they match. The pairs with 0 will be skipped for calculation of MMM and global contrastive losses as well. return_loss (`bool`, *optional*, default to None): Whether to return calculated loss or not. z Parameters: image_codebook ([`nn.Module`]): If passed, the image codebook will be set to this. Otherwise. it will be initialized using the image_codebook_config defined in the config first as the first parameter. c@sPeZdZdZeZdZdZee j e j e j fddddZ d eedd d d ZdS) FlavaPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. flavaTN)rr)cCst|tjtjfr@|jjjd|jjd|j dur|j j nft|tj r|jjjd|jjd|j dur|jj|j  n&t|tj r|j j |jjddS)zInitialize the weightsg)meanZstdNr)rr rrweightdataZnormal_r^Zinitializer_rangerZzero_rrrZfill_)r1rr2r2r3 _init_weights`s     z"FlavaPreTrainedModel._init_weightsF)rrr)cCst|tr||_dSrZ)rrr)r1rrr2r2r3_set_gradient_checkpointingps z0FlavaPreTrainedModel._set_gradient_checkpointing)F)r8r9r:r;r config_classbase_model_prefixsupports_gradient_checkpointingr r rrrrrrIrr2r2r2r3r Vs  r zeThe bare FLAVA Image Model transformer outputting raw hidden-states without any specific head on top.rcseZdZeZdZdZdeedfdd Ze j ddd Z e j d d d Z e eeefd dddZeedeeeededdeejeejeeeejeejeeeeeeeeefd ddZZ S)FlavaImageModelzflava.image_modelrTr^add_pooling_layercsXt|||_t||_t||_tj|j |j d|_ |rFt |nd|_ |dSNr)rerfr^r]rurencoderr rrjr layernormr pooler post_initr1r^rrsr2r3rfs   zFlavaImageModel.__init__r(cCs|jjSrZrurmr0r2r2r3get_input_embeddingssz$FlavaImageModel.get_input_embeddingsrcCs ||j_dSrZr r1rr2r2r3set_input_embeddingssz$FlavaImageModel.set_input_embeddingsNheads_to_pruner)cCs*|D]\}}|jj|j|qdSz Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base class PreTrainedModel Nitemsrrrrr1r&rrr2r2r3 _prune_headsszFlavaImageModel._prune_headsbatch_size, image_num_patchesZvision)r output_typerZmodalityZexpected_output rrrrrrrrr)c Cs|dur |n|jj}|dur |n|jj}|dur4|n|jj}|durLtd|||jj}|j|||d} |j| |||||d} | d} | | } |j dur| | nd} |s| | f| ddSt | | | j | j dS)Nz You have to specify pixel_values)rrrrrrrrrrZ pooler_outputrr)r^rruse_return_dictr get_head_maskrrurrrrrr) r1rrrrrrrrembedding_outputencoder_outputssequence_outputr r2r2r3rs: zFlavaImageModel.forward)T)NNNNNNNN)!r8r9r:rrrmain_input_namerIrfr Moduler!r$rrrr+rFLAVA_IMAGE_INPUTS_DOCSTRINGformatr_CHECKPOINT_FOR_DOCr!_CONFIG_CLASS_FOR_IMAGE_MODEL_DOC_EXPECTED_IMAGE_OUTPUT_SHAPErr<rrr r6rrr2r2rsr3rusD    rzdThe bare FLAVA Text Model transformer outputting raw hidden-states without any specific head on top.cseZdZeZdZdeedfdd ZedddZ e j d d d Z e eeefd d ddZeedeeeeddeejeejeejeejeejeeeeeeeeefd ddZZS)FlavaTextModelzflava.text_modelTrcsXt|||_t||_t||_tj|j |j d|_ |rFt |nd|_ |dSr)rerfr^rrurrr rrjrrr rrrrsr2r3rfs   zFlavaTextModel.__init__r(cCs|jjSrZrurr0r2r2r3r!sz#FlavaTextModel.get_input_embeddingsr"cCs ||j_dSrZr>r#r2r2r3r$sz#FlavaTextModel.set_input_embeddingsNr%cCs*|D]\}}|jj|j|qdSr'r(r*r2r2r3r+szFlavaTextModel._prune_headsbatch_size, text_seq_lengthrr-r) rrrrrrrrr)c Cs|dur |n|jj}|dur |n|jj}|dur4|n|jj}|durLtd|} |durltj| |jd}| ||jj }| || |j} |j |||d} |j | | ||||d} | d} || } |jdur|| nd}|s| |f| ddSt| || j| jdS)NzYou have to specify input_idsrrr/rrr0)r^rrr1rrr<onesrr2rget_extended_attention_maskrurrrrrr)r1rrrrrrrrrextended_attention_maskr3r4r5r r2r2r3rsJ zFlavaTextModel.forward)T)NNNNNNNN)r8r9r:rrrrIrfrlr!r r7r$rrrr+rFLAVA_TEXT_INPUTS_DOCSTRINGr9rr:r _CONFIG_CLASS_FOR_TEXT_MODEL_DOCrr<rr r6rrr2r2rsr3r=s>   r=zjThe bare FLAVA Multimodal Model transformer outputting raw hidden-states without any specific head on top.c seZdZeZdZdZdedfdd Zee e e fddd d Z e e d eeeed dejeejeejeeeeeeeeefd ddZZS)FlavaMultimodalModelzflava.multimodal_modelrTrcsvt|||_|jj|_|jr:ttdd|j|_ t ||_ tj |j|j d|_|rdt|nd|_|dSr)rerfr^ use_cls_tokenr rhr<rirjrkrrrrrr rrrrsr2r3rfKs   zFlavaMultimodalModel.__init__Nr%cCs*|D]\}}|jj|j|qdSr'r(r*r2r2r3r+Ysz!FlavaMultimodalModel._prune_heads,batch_size, image_num_patches + text_seq_lenr@rcCs(|dur |n|jj}|dur |n|jj}|dur4|n|jj}|\}}} |jrz|j|dd} tj | |fdd}|d7}|durtj ||f|j d}| ||jj }||||f|j } |j|| ||||d} | d} || } |jdur|| nd}|s| |f| ddSt| || j| jdS)Nrxrr}rAr/rr0)r^rrr1rrHrkrr<rrBrr2rrCrrrrrr)r1rrrrrrrrrrrDr4r5r r2r2r3rasD  zFlavaMultimodalModel.forward)T)NNNNN)r8r9r:rrrr6rfrrrr+r!FLAVA_MULTIMODAL_INPUTS_DOCSTRINGr9rr:r&_CONFIG_CLASS_FOR_MULTIMODAL_MODEL_DOCr<rrrIr r6rrr2r2rsr3rGAs6 rGz_The bare FLAVA Model transformer outputting raw hidden-states without any specific head on top.rcsNeZdZeZedfdd Zeedde e j e e j e e j e e j e e e e e e e j dddZeed de e j e e je e e e j e e j e e e e e e e j d d d Zeed eeedde e je e j e e j e e j e e j e e je e j e e e e e e e eeefd ddZZS) FlavaModelrcs4t|t|jts.tdt|jdt|jtsPtdt|jdt|j t svtddt|j d|j}|j}|j }|j |_ |j |_ |j |_|j |_t||_t||_t||_t|j|j |_t|j |j |_ttg|jj|_t|j|j|_ t|j |j|_!|"dS)NzLconfig.text_config is expected to be of type FlavaTextConfig but is of type rzNconfig.image_config is expected to be of type FlavaImageConfig but is of type zMconfig.multimodal_config is expected to be of type FlavaMultimodalConfig but z is of type )#rerfr text_configrrtype image_configrmultimodal_configrZprojection_dimrjZtext_hidden_sizeZimage_hidden_sizeZmm_hidden_sizer= text_modelr image_modelrGmultimodal_modelr rimage_projectiontext_projectionrhr<rBr^Zlogit_scale_init_value logit_scaleimage_to_mm_projectiontext_to_mm_projectionr)r1r^rMrOrPrsr2r3rfsF       zFlavaModel.__init__r?N)rrrrrrrr)c Cs8dt|j|||||||d}|d} || } | S)Na Returns: text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by applying the projection layer to the pooled output of [`FlavaTextModel`]. Examples: ```python >>> from transformers import AutoProcessor, FlavaModel >>> model = FlavaModel.from_pretrained("{0}") >>> processor = AutoProcessor.from_pretrained("{0}") >>> inputs = processor( ... text=["a photo of a cat", "a photo of a dog"], max_length=77, padding="max_length", return_tensors="pt" ... ) >>> text_features = model.get_text_features(**inputs) ```)rrrrrrrr)r9r:rQrU) r1rrrrrrrZ text_outputsr Z text_featuresr2r2r3get_text_featuress   zFlavaModel.get_text_featuresr,r.c Cs:dt|j||||||||d} | d} || } | S)Na Returns: image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by applying the projection layer to the pooled output of [`FlavaImageModel`]. Examples: ```python >>> from PIL import Image >>> import requests >>> from transformers import AutoProcessor, FlavaModel >>> model = FlavaModel.from_pretrained("{0}") >>> processor = AutoProcessor.from_pretrained("{0}") >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> inputs = processor(images=image, return_tensors="pt") >>> image_features = model.get_image_features(**inputs) ```)rrrrrrrrr)r9r:rRrT) r1rrrrrrrrZ image_outputsr Zimage_featuresr2r2r3get_image_featuress   zFlavaModel.get_image_featuresrIr-rT) rrrrrrimage_attention_maskskip_multimodal_encoderrrrr)c  Cs2| dur | n|jj} | s tdd} d} d}d}|durn|j|||| | | d}|d|d} } || d}d}d}d}d}|dur|j||||| | | d}|d|d}}||d}d}d}|dur|dur|stj||gdd }|j || d }|d}| s| |||||fSt | |||||d S) a\ Returns: Examples: ```python >>> from PIL import Image >>> import requests >>> from transformers import AutoProcessor, FlavaModel >>> model = FlavaModel.from_pretrained("facebook/flava-full") >>> processor = AutoProcessor.from_pretrained("facebook/flava-full") >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> inputs = processor(text=["a photo of a cat"], images=image, return_tensors="pt", padding=True) >>> outputs = model(**inputs) >>> logits_per_image = outputs.contrastive_logits_per_image # this is the image-text similarity score >>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities ``` NzRFLAVA model requires hidden states to work. Please set `output_hidden_states=True`)rrrrrrrryrx)rrrrrrrrr})r)r"r#r$r%r&r') r^rrrRrWrQrXr<rrSr!)r1rrrrrrr\r]rrrr"Z image_statesZimage_mm_projectionr#r$Z text_statesZtext_mm_projectionr%r&r'Zmultimodal_inputr2r2r3r7sp*  zFlavaModel.forward)NNNNNNN)NNNNNNNN) NNNNNNNNNTN)r8r9r:rrrfrrEr9rr<rrIr=rYr8rrZFLAVA_MODEL_INPUTS_DOCSTRINGrr! LongTensorr r rrrr2r2rsr3rLs+ - 3  rLcs8eZdZeedfdd ZejejdddZZS)FlavaImageCodebookResPath)in_sizeout_sizec st|d}t}t|d<tj||ddd|d<t|d<tj||ddd|d<t|d <tj||ddd|d <t|d <tj||dd d|d <t||_dS)NZrelu_1r rrpaddingZconv_1Zrelu_2Zconv_2Zrelu_3Zconv_3Zrelu_4rZconv_4)rerfrr ReLUr Sequentialpath)r1rarbkwargsZhid_sizerhrsr2r3rfs     z"FlavaImageCodebookResPath.__init__rcCs ||SrZ)rhr1rr2r2r3rsz!FlavaImageCodebookResPath.forward r8r9r:rrfr<rrrr2r2rsr3r`sr`cs:eZdZeeedfdd ZejejdddZZS)FlavaImageCodebookBlock)rarb num_layersc sPtd|d|_||kr6tj||ddd|_n t|_t|||_dS)Nrryrrd) rerf post_gainr rid_pathZIdentityr`res_path)r1rarbrmrirsr2r3rfs   z FlavaImageCodebookBlock.__init__rcCs|||j||SrZ)rornrprjr2r2r3rszFlavaImageCodebookBlock.forwardrkr2r2rsr3rls rlcs@eZdZdeeeeedfdd ZejejdddZZ S) FlavaImageCodebookLayerGroupT) num_blocksrmrarbuse_poolcstt}t|D]B}|dkr@t||||d|d<qt||||d|d<q|rptjdd|d<t||_dS)NrZblock_rry)rpool) rerfrrrlr Z MaxPool2drggroup)r1rrrmrarbrsblocksrrsr2r3rfs  z%FlavaImageCodebookLayerGroup.__init__rcCs ||SrZ)rurjr2r2r3rsz$FlavaImageCodebookLayerGroup.forward)T) r8r9r:rrIrfr<rrrr2r2rsr3rqsrqa" The FLAVA's image codebook model inspired from DALL-E's original encoder. Outputs raw hidden states and can be used to generate image tokens for an image based on DALL-E's vocab. Used to generate labels for MIM. Use `get_codebook_indices` to get image tokens for an image. rcspeZdZdZeZdZdZeedfdd Z e j e j ddd Z e j e j dd d Z e je j dd d ZZS)FlavaImageCodebookrF)r^ric sdt|||_|j|_|j|_|j|_|j|_|j|_|j|j}t}t |d<t j d|j|jddd|d<t}t j |jd|jddd|d <t |j|d|jd|j|d <t |j|d|jd |j|d <t |j|d |jd |j|d<t |j|d |jd|jdd|d<t ||d<t ||_||jjr`|D] }d|_qRdS)NZrelurrrdconvr inputZgroup_1ryZgroup_2rcZgroup_3F)rsZgroup_4r)rerfr^Z num_groupsinput_channelsZnum_blocks_per_grouprjrrr rfrrqrgrvrfreeze parametersZ requires_grad)r1r^rirmZ output_blocksrvparamrsr2r3rfs>       zFlavaImageCodebook.__init__)rr)cCs"dt||}tj|ddS)Na Args: pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Pixel values. Codebook pixel values can be obtained using [`AutoImageProcessor`] by passing `return_codebook_pixels=True`. See [`FlavaImageProcessor.__call__`] for details. Examples: ```python >>> from PIL import Image >>> import requests >>> from transformers import AutoImageProcessor, FlavaImageCodebook >>> model = FlavaImageCodebook.from_pretrained("{0}") >>> image_processor = AutoImageProcessor.from_pretrained("{0}") >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> inputs = image_processor([image], return_codebook_pixels=True, return_tensors="pt") >>> inputs = dict(pixel_values=inputs.codebook_pixel_values) >>> outputs = model.get_codebook_indices(**inputs) ``` r)Zaxis)r9_CHECKPOINT_FOR_CODEBOOK_DOCrvr<Zargmaxr1rZz_logitsr2r2r3get_codebook_indicess  z'FlavaImageCodebook.get_codebook_indicescCs||}tjdd|S)Nrr})rvr ZSoftmaxrr2r2r3get_codebook_probs5s z%FlavaImageCodebook.get_codebook_probscCs`dtt|jdkr*td|jd|jd|jkrVtd|jdd|j||S)Na Args: pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Pixel values. Codebook pixel values can be obtained using [`AutoImageProcessor`] by passing `return_codebook_pixels=True`. See [`FlavaImageProcessor.__call__`] for details. Examples: ```python >>> from PIL import Image >>> import requests >>> from transformers import AutoImageProcessor, FlavaImageCodebook >>> model = FlavaImageCodebook.from_pretrained("{0}") >>> image_processor = AutoImageProcessor.from_pretrained("{0}") >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> inputs = image_processor([image], return_codebook_pixels=True, return_tensors="pt") >>> inputs = dict(pixel_values=inputs.codebook_pixel_values) >>> outputs = model(**inputs) >>> print(outputs.shape) (1, 196) ``` rcz input shape z is not 4drz input has z channels but model built for )r9rrrrr}rv)r1rr2r2r3r9szFlavaImageCodebook.forward)r8r9r:rrrr6rrrfr<rrrr=rrr2r2rsr3rws ,rwcs$eZdZfddZddZZS)FlavaPredictionHeadTransformcsVtt|j|j|_t|jtr6t |j|_ n|j|_ tj |j|j d|_ dSr) rerfr rrjrrrrr transform_act_fnrrrrsr2r3rf_s   z%FlavaPredictionHeadTransform.__init__cCs"||}||}||}|SrZ)rrrrr2r2r3rhs   z$FlavaPredictionHeadTransform.forwardr8r9r:rfrrr2r2rsr3r^s rcs&eZdZdfdd ZddZZS)FlavaMaskedPredictionHeadNcsbt||_t||_tj|j|jdd|_ t t |j|_ |durT||j _|j |j _ dS)NFr)rerfr^r transformr rrjrdecoderrhr<rirr)r1r^rrsr2r3rfps  z"FlavaMaskedPredictionHead.__init__cCs||}||}|SrZ)rrrjr2r2r3r|s  z!FlavaMaskedPredictionHead.forward)Nrr2r2rsr3ros rcs$eZdZfddZddZZS) FlavaITMHeadcs.t||_t||_t|jd|_dS)Nry) rerfr^r rr rrjseq_relationshiprrsr2r3rfs  zFlavaITMHead.__init__cCs||}||}|SrZ)rrrjr2r2r3rs  zFlavaITMHead.forwardrr2r2rsr3rs rcs$eZdZfddZddZZS)FlavaGlobalContrastiveHeadcst||_|j|_dSrZ)rerfr^global_backprop_contrastiverrsr2r3rfs z#FlavaGlobalContrastiveHead.__init__c s2t|}tjrtjsBtjdjd}g}g}nd}tj} |j r~tjj j }tjj j }nHfddt | D}fddt | D}tj |tj ||tjtj|jd}t|}t|}t|dd|} t|dd|} | | |fS)NrrAcsg|]}tqSr2r<Z zeros_liker)r$r2r3rr\z6FlavaGlobalContrastiveHead.forward..csg|]}tqSr2rr)r"r2r3rr\r)r<expZ distributedZ is_availableZis_initializedrrrZget_world_sizerr rZ all_gatherrZget_rankrrr) r1r"r$rVZ temperaturelabelsZimage_embeddings_allZtext_embeddings_allZlocal_batch_sizeZ world_sizelogits_per_imagelogits_per_textr2)r"r$r3rs,     z"FlavaGlobalContrastiveHead.forwardrr2r2rsr3rs rzk The FLAVA model for pretraining which outputs losses, embeddings, logits and transformer outputs. cseZdZgdZdeeejdfdd Ze j dddZ e e d d eeed dee jee jee jee jee j ee j ee j ee jee j eee j ee j ee j eeeeeeeeee j efd ddZZS)FlavaForPreTraining)zmmm_text_head.decoder.biaszmmm_image_head.decoder.biaszmlm_head.decoder.biaszmim_head.decoder.biasN)r^image_codebookcst|t||_||_|jdur8|jr8t|j|_t|j |_ t|j |_ t ||_t|j |_t|j |_t||_|j j|_|j j|_|j|_|j|_|j|_|j|_|j|_|j|_|j|_|j|_|dSrZ)rerfrLrrZ init_codebookrwZimage_codebook_configrrOmim_headrMmlm_headritm_headmmm_image_head mmm_text_headrglobal_contrastive_headrimage_vocab_sizetext_vocab_size mlm_weight mim_weightglobal_contrastive_weightce_ignore_index itm_weightmmm_image_weightmmm_text_weight skip_unmasked_multimodal_encoderr)r1r^rrsr2r3rfs,           zFlavaForPreTraining.__init__)rcCs"|dkr||dd}|S)Nryrrx)r~rrrjr2r2r3 _resize_to_2ds z!FlavaForPreTraining._resize_to_2dzbatch_size, text_seq_lenr,r[T)rinput_ids_maskedrcodebook_pixel_valuesrrrrr\r mlm_labels mim_labels itm_labelsrrr return_lossr)c6Cs|dur |n|jj}|dur |n|jj}| dur4| n|j} |durX|durXtd|}|j|||||| | ||dd }|j||||| |||dd }d}|j}|j}|j}|j}|j }d}}}}}}} d}!}"}#}$d}%}&}'|dus|dur@| dur@|r@|j dur"t d|dur4t d|j |} |jdkr|dur|dur|}(| dur|| } ||}|j| |d<|(dd| d  dddf}(| |j})| |)}*|(|)ddf}(||(}!|rtj|!d |j|*d }||j9}n ||(}!|jdkr|dur|dur|}+| dur|| } |+dd| d  dddf}+| |j})| |)},|+|)ddf}+||+}"|rtj|"d |j|,d }||j9}n ||+}"|jdkr|dur||}%| dur| d}-t|- |-|-!dg}|rJtj|%| } | |j9} |dur\||}| durn| |} | dur| |} |durn|j"dkrn|}(|d d }.|(ddd d |.ddf}(|dur|(|}(| durd|| } ||}|j| |d<| |j})| |)}*|(|)ddf}(|#|(}$|rntj|$d |j|*d }||j"9}n |#|(}$|dur8|j$dkr8|}+|+dd|d  dddf}+|dur|+|}+| dur.|| } | |j})| |)},|+|)ddf}+|%|+}#|r8tj|#d |j|,d }||j$9}n |%|+}#|dur8|dur8|j&dkr8|j'|dddddf}/tjj(|/d d }/|j)|dddddf}0tjj(|0d d }0|jj*j+,t-t.|/|0|/|jj*\}&}'}1|dur|&|}&|'|}'|1|}1|r8tj|&|1}2tj|'|1}3|2|3d }||j&9}t0||| |||d }4|rr|41srt2dd|43D}|sV||j4dur|j45nd||j6dur|j65nd|j |j7dur|j75nd||j4dur|j45nd||j6dur|j65nd||j7dur|j75nd|!|"|%|&|&|$|#f}5|rD|41sD||4f|5}5t8dd|5DSt9||4||j4||j6|j |j7||j4||j6||j7|!|"|%|&|'|$|#dS)ai Examples: ```python >>> from PIL import Image >>> import requests >>> from transformers import FlavaForPreTraining, AutoProcessor >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> model = FlavaForPreTraining.from_pretrained("facebook/flava-full") >>> processor = AutoProcessor.from_pretrained("facebook/flava-full") >>> text = ["a photo of a cat"] >>> inputs = processor( ... images=[image], ... text=text, ... return_masks=True, ... return_codebook_pixels=True, ... padding=True, ... max_length=77, ... return_tensors="pt", ... ) >>> output = model(**inputs) ``` Return: Nz`input_ids_masked` isn't passed which means MLM loss won't be calculated correctlySetting it to `input_ids` so that model can work. Please pass it if this is unintentional. This is usually OKAY if you are doing inference on unmasked text...T) rrrrrr\r]rrr) rrrrr\rrrrz`return_loss` is set to True but the image codebook is not initialized and no `mim_labels` have been passed. Reinstantiate the model with `init_codebook` set to True or pass in your custom `mim_labels`z`codebook_pixel_value` are required to generate `mim_labels` if loss is expected. Call `AutoProcessor` with `return_codebook_pixels` set to Truerrrxryr})r@rArBrCrDrEcss|]}|dur|ndVqdSr r2)r.rKr2r2r3r4r\z.FlavaForPreTraining.forward..css|]}|dur|VqdSrZr2)r.rr2r2r3r4r\)rKrLr"r#r$r%r&r'rMrNrOrPrQrRrSrTrUrVrWrXrY):r^r1rrloggerwarningrr"r$r&r RuntimeErrorrrrrrnerrr rZ cross_entropyrrrrrrrr<whereanynewrrrrrrU normalizerTrVrZclamp_LOGIT_SCALE_CLAMP_MINLOGIT_SCALE_CLAMP_MAXrr?rGsumrFr#r,r%r'r6rJ)6r1rrrrrrrrr\rrrrrrrrZ flava_outputZflava_masked_outputZpos_maskr"r$rMrOrQZ total_lossZmim_lossZmlm_lossZ mmm_text_lossZmmm_image_lossZgc_lossZitm_lossrSrTrYrXrUrrZsequence_for_imageZ masked_tokensZmim_labels_filteredZsequence_for_textZmlm_labels_filteredZ pos_pairsZ end_indexZtext_embeddingZimage_embeddingZ gc_labelsZ gc_loss_imageZ gc_loss_textZ flava_lossesrr2r2r3rs8         "       "                   "             zFlavaForPreTraining.forward)N)NNNNNNNNNNNNNNTNN)r8r9r:Z_keys_to_ignore_on_load_missingrrr r7rfr<rrr"FLAVA_PRETRAINING_INPUTS_DOCSTRINGr9rrJr_r=rIr r rrr2r2rsr3rsX  r)^r;rrrZ dataclassesrtypingrrrrrr r r<Ztorch.utils.checkpointr Z activationsr Zmodeling_outputsrrZmodeling_utilsrrrrrrrrrrZconfiguration_flavarrrrrZ get_loggerr8rr:rr;rFrKr<Z#FLAVA_PRETRAINED_MODEL_ARCHIVE_LISTZ,FLAVA_CODEBOOK_PRETRAINED_MODEL_ARCHIVE_LISTrrrr!r?rJr7r]rlrrrrrrrrr ZFLAVA_START_DOCSTRINGZFLAVA_INPUTS_DOCSTRING_COMMONZ!FLAVA_IMAGE_INPUTS_DOCSTRING_BASEr8Z FLAVA_TEXT_INPUTS_DOCSTRING_BASErErJZ!FLAVA_MODEL_INPUTS_DOCSTRING_BASEr^rZ'FLAVA_PRETRAINING_START_DOCSTRING_EXTRAr r9rr=rGrLr`rlrqrwrrrrrr2r2r2r3s  $    %$e]$7E*.9       +/ _ e ^ } y (