U 4Afy@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.d ddgZ/dZ0dZ1e e&e$e%fZ2eGdddeZ3eGdddeZ4eGdddeZ5Gdddej6Z7Gdddej6Z8Gd d!d!ej6Z9Gd"d#d#ej6Z:Gd$d%d%ej6Z;Gd&d'd'ej6ZGd,d-d-ej6Z?Gd.d/d/ej6Z@Gd0d1d1ej6ZAd2ZBd3ZCd4ZDeDeCZEd5ZFeFeCZGd6eCZHd7ZIeDeFeCeIZJd8eFeDd9eCZKd:ZLGd;d<dGd?d@d@eMZOedAeBjNdd>GdBdCdCeMZPedDeBjNdd>GdEdFdFeMZQedGeBjNdHd>GdIdJdJeMZRGdKdLdLej6ZSGdMdNdNej6ZTGdOdPdPej6ZUedQeBjNdRd>GdSdTdTeMZVGdUdVdVej6ZWGdWdXdXej6ZXGdYdZdZej6ZYGd[d\d\ej6ZZed]eBjNdHd>eLGd^d_d_eMZ[dS)`zPyTorch 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-codebookrrrg$(~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,|]$}|dkr|n t|VqdS))r%r#r'Ngetattrto_tuple.0kselfL/tmp/pip-unpacked-wheel-zw5xktn0/transformers/models/flava/modeling_flava.py asz,FlavaModelOutput.to_tuple..tuplekeysr0r2r0r3r,`s zFlavaModelOutput.to_tuple)__name__ __module__ __qualname____doc__r"rtorch FloatTensor__annotations__r#rr$r%r&r'r rr,r2r2r2r3r!As 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]}|dk r d}q"q |S)NTF)values)r1all_nonevr2r2r3rGs  zFlavaLosses.all_none)r8r9r:r;r@rr<r=r>rArBrCrDrEboolrGr2r2r2r3r?gs 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`, *optional*, 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,ddddddgtfddDS) Nr%r#r'rPrNrRc3s,|]$}|kr|n t|VqdSNr*r-r1Ztransformer_outputsr2r3r4sz5FlavaForPreTrainingOutput.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_tokenPatchEmbeddingsr`rarbpatch_embeddings num_patchesposition_embeddingsDropouthidden_dropout_probdropoutr])r1r]r^rm __class__r2r3res   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)shapernr]rar functionalZ interpolateZreshapeintmathsqrtpermute ValueErrorviewr<cat unsqueeze) r1rtrurvZnpatchnum_posZclass_pos_embedZpatch_pos_embedr}Z num_h_patchesZ num_w_patchesr2r2r3interpolate_pos_encodings,    .(,z-FlavaImageEmbeddings.interpolate_pos_encoding) pixel_valuesbool_masked_posrr)cCs|j\}}}}|j||d}|\}} } |dk r|j|| d} |dkr`||dd}|d| } |d| | | }|j |dd} t j | |fdd}|r|| |||}n ||j }||}|S)N)rrwr r?rr|)r~rlsizerfexpandr}rrZtype_asrjr<rrrnrq)r1rrr batch_sizerbrurvrtseq_len_Z mask_tokensmask cls_tokensr2r2r3forward+s    zFlavaImageEmbeddings.forward)F)NF)r8r9r:r;rrIrer<Tensorrrr BoolTensorr __classcell__r2r2rrr3r\s&r\csVeZdZdZd eeeeeeffeedfdd Zdej e ej d d d Z Z S)rkz# Image to Patch Embedding. r r r_cstt|tjjs ||f}t|tjjs6||f}|d|d|d|d}||_||_||_t j ||||d|_ dS)Nrr) kernel_sizeZstride) rdre isinstance collectionsabcIterabler`rarmr Conv2d projection)r1r`rarbrcrmrrr2r3reTs  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).rx)r~r`rrflatten transpose)r1rrrrbrurvxr2r2r3rgs(zPatchEmbeddings.forward)rrr r )F) r8r9r:r;rr r rer<rrIrrr2r2rrr3rkOsrkcsFeZdZdZfddZdeejeejeejdddZZ 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|_|jdt|jddd|jd tj|jtjd dddS) N) padding_idxZepsposition_embedding_typeabsolute position_ids)rrwF) persistenttoken_type_ids)dtype)rdrer Embedding vocab_sizeriZ pad_token_idword_embeddingsZmax_position_embeddingsrnZtype_vocab_sizetoken_type_embeddings LayerNormlayer_norm_epsrorprqr+rZregister_bufferr<arangerrhrrlongr1r]rrr2r3revs" zFlavaTextEmbeddings.__init__N input_idsrrc Cs|}|d}|dkr.|jddd|f}|dkrt|drl|jddd|f}||d|}|}ntj|tj|jjd}| |}| |} || } |j dkr| |} | | 7} | | } || } | S)Nrrr)rdevicer)rrhasattrrrr<rhrrrrrrnrrq) r1rrr input_shape seq_lengthZbuffered_token_type_idsZ buffered_token_type_ids_expandedZ inputs_embedsrrtrnr2r2r3rs&       zFlavaTextEmbeddings.forward)NNN) r8r9r:r;rerr<rrrr2r2rrr3rss 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)rdrerinum_attention_headsrrrattention_head_size all_head_sizer LinearZqkv_biasquerykeyvalueroZattention_probs_dropout_probrqrrrr2r3res zFlavaSelfAttention.__init__rr)cCs6|dd|j|jf}|j|}|ddddS)Nrwrrxrr )rrrrr)r1rZ new_x_shaper2r2r3transpose_for_scoress z'FlavaSelfAttention.transpose_for_scoresF hidden_statesattention_mask head_maskoutput_attentionsr)cCs||}|||}|||}||}t||dd} | t|j } |dk rh| |} t j j | dd} | | } |dk r| |} t| |} | dddd} | dd|jf} | j| } |r| | fn| f} | S)Nrwrzr|rrxrr )rrrrr<matmulrrrrr rZsoftmaxrqr contiguousrrr)r1rrrrZmixed_query_layerZ key_layerZ value_layerZ query_layerZattention_scoresZattention_probsZ context_layerZnew_context_layer_shapeoutputsr2r2r3rs$     zFlavaSelfAttention.forward)NNF)r8r9r:FlavaPossibleConfigsrer<rrrrIr r rrr2r2rrr3rsrcs@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) rdrer rridenserorprqrrrr2r3res zFlavaSelfOutput.__init__r input_tensorr)cCs||}||}|SrZrrqr1rrr2r2r3rs  zFlavaSelfOutput.forward) r8r9r:r;rrer<rrrr2r2rrr3rsrc sxeZdZeddfdd ZeeddddZd ej e ej e ej e e e ej ej fe ej fd d d ZZS) FlavaAttentionNrcs*tt||_t||_t|_dSrZ)rdrer attentionroutputset pruned_headsrrrr2r3res   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_headss zFlavaAttention.prune_headsFrcCs8|j||||d}||d|}|f|dd}|SN)rrrrr)rr)r1rrrrZ self_outputsattention_outputrr2r2r3rszFlavaAttention.forward)NNF)r8r9r:rrerrrr<rrrIr r rrr2r2rrr3rsrcs8eZdZeddfdd ZejejdddZZS)FlavaIntermediateNrcsBtt|j|j|_t|jt r6t |j|_ n|j|_ dSrZ) rdrer rriintermediate_sizerr hidden_actstrr intermediate_act_fnrrrr2r3re,s   zFlavaIntermediate.__init__)rr)cCs||}||}|SrZ)rrr1rr2r2r3r5s  zFlavaIntermediate.forward r8r9r:rrer<rrrr2r2rrr3r+s rcs<eZdZeddfdd ZejejejdddZZS) FlavaOutputNrcs.tt|j|j|_t|j|_ dSrZ) rdrer rrrirrorprqrrrr2r3re=s zFlavaOutput.__init__rcCs ||}||}||}|SrZrrr2r2r3rCs  zFlavaOutput.forwardrr2r2rrr3r<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)rdreZchunk_size_feed_forwardZ seq_len_dimrrr intermediaterrr rrirlayernorm_beforelayernorm_afterrrrr2r3reOs    zFlavaLayer.__init__Frc Csb|j|||||d}|d}|dd}||}||}||}|||}|f|}|Sr)rrrrr) r1rrrrZself_attention_outputsrrZ layer_outputr2r2r3r[s     zFlavaLayer.forward)NNF)r8r9r:r;rrer<rrrIr r rrr2r2rrr3rLsrc sVeZdZeddfdd Zd ejeejeejeeee e e fddd Z Z S) FlavaEncoderNrcs:t|_tfddtjD|_d|_dS)Ncsg|] }tqSr2)rr.rr]r2r3 ~sz)FlavaEncoder.__init__..F) rdrer]r Z ModuleListrangenum_hidden_layerslayergradient_checkpointingrrrrr3re{s  zFlavaEncoder.__init__FTrrrroutput_hidden_states return_dictr)c Cs|rdnd}|rdnd}t|jD]t\} } |r8||f}|dk rH|| nd} |jrn|jrn|| j||| |} n| ||| |} | d}|r"|| df}q"|r||f}|stdd|||fDSt|||dS)Nr2rrcss|]}|dk r|VqdSrZr2)r.rHr2r2r3r4sz'FlavaEncoder.forward..)last_hidden_stater attentions) enumeraterrZtrainingZ_gradient_checkpointing_func__call__r6r) r1rrrrrrZall_hidden_statesZall_self_attentionsiZ layer_moduleZlayer_head_maskZ layer_outputsr2r2r3rs6     zFlavaEncoder.forward)NNFFT)r8r9r:rrer<rrrIr r6rrrr2r2rrr3rzs  rcs2eZdZedfdd ZejdddZZS) FlavaPoolerrcs*tt|j|j|_t|_dSrZ)rdrer rrirZTanh activationrrrr2r3res zFlavaPooler.__init__)rcCs(|dddf}||}||}|SNr)rr)r1rZfirst_token_tensor pooled_outputr2r2r3rs  zFlavaPooler.forwardrr2r2rrr3rsraD 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@s<eZdZdZeZdZdZee j e j e j fddddZ dS)FlavaPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. flavaTN)moduler)cCst|tjtjfr@|jjjd|jjd|j dk r|j j nft|tj r|jjjd|jjd|j dk r|jj|j  n&t|tj r|j j |jjddS)zInitialize the weightsg)ZmeanZstdNr)rr rrweightdataZnormal_r]Zinitializer_rangerZzero_rrrZfill_)r1r r2r2r3 _init_weightsVs     z"FlavaPreTrainedModel._init_weights)r8r9r:r;r config_classbase_model_prefixsupports_gradient_checkpointingr r rrrr r2r2r2r3rLs rzeThe 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)rdrer]r\rtrencoderr rrir layernormrpooler post_initr1r]rrrr2r3reqs   zFlavaImageModel.__init__r(cCs|jjSrZrtrlr0r2r2r3get_input_embeddings~sz$FlavaImageModel.get_input_embeddingsrcCs ||j_dSrZrr1rr2r2r3set_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 Nitemsrrrrr1rrrr2r2r3 _prune_headsszFlavaImageModel._prune_headsbatch_size, image_num_patchesZvision) checkpoint output_typer ZmodalityZexpected_output rrrrrrrrr)c Cs|dk r |n|jj}|dk r |n|jj}|dk r4|n|jj}|dkrLtd|||jj}|j|||d} |j| |||||d} | d} | | } |j dk r| | nd} |s| | f| ddSt | | | j | j dS)Nz You have to specify pixel_values)rrrrrrrrrrZ pooler_outputrr)r]rruse_return_dictr get_head_maskrrtrrrrrr) r1rrrrrrrrembedding_outputencoder_outputssequence_outputrr2r2r3rs> zFlavaImageModel.forward)T)NNNNNNNN)!r8r9r:rr rmain_input_namerIrer Modulerrrrrr$rFLAVA_IMAGE_INPUTS_DOCSTRINGformatr_CHECKPOINT_FOR_DOCr!_CONFIG_CLASS_FOR_IMAGE_MODEL_DOC_EXPECTED_IMAGE_OUTPUT_SHAPErr<rrr r6rrr2r2rrr3rgsD    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)rdrer]rrtrrr rrirrrrrrrrr2r3res   zFlavaTextModel.__init__r(cCs|jjSrZrtrr0r2r2r3rsz#FlavaTextModel.get_input_embeddingsrcCs ||j_dSrZr8rr2r2r3rsz#FlavaTextModel.set_input_embeddingsNrcCs*|D]\}}|jj|j|qdSr r!r#r2r2r3r$szFlavaTextModel._prune_headsbatch_size, text_seq_lengthr&r'r ) rrrrrrrrr)c Cs|dk r |n|jj}|dk r |n|jj}|dk r4|n|jj}|dkrLtd|} |dkrltj| |jd}| ||jj }| || |j} |j |||d} |j | | ||||d} | d} || } |jdk r|| nd}|s| |f| ddSt| || j| jdS)NzYou have to specify input_idsrrr)rrr*)r]rrr+rrr<onesrr,rget_extended_attention_maskrtrrrrrr)r1rrrrrrrrrextended_attention_maskr-r.r/rr2r2r3rsN zFlavaTextModel.forward)T)NNNNNNNN)r8r9r:rr rrIrerkrr r1rrrrr$rFLAVA_TEXT_INPUTS_DOCSTRINGr3rr4r _CONFIG_CLASS_FOR_TEXT_MODEL_DOCrr<rr r6rrr2r2rrr3r7s>   r7zjThe 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)rdrer] use_cls_tokenr rgr<rhrirjrrrrrrrrrrrr2r3re=s   zFlavaMultimodalModel.__init__NrcCs*|D]\}}|jj|j|qdSr r!r#r2r2r3r$Ksz!FlavaMultimodalModel._prune_heads,batch_size, image_num_patches + text_seq_lenr:rcCs(|dk r |n|jj}|dk r |n|jj}|dk r4|n|jj}|\}}} |jrz|j|dd} tj | |fdd}|d7}|dkrtj ||f|j d}| ||jj }||||f|j } |j|| ||||d} | d} || } |jdk r|| nd}|s| |f| ddSt| || j| jdS)Nrwrr|r;r)rr*)r]rrr+rrBrjrr<rr<rr,rr=rrrrrr)r1rrrrrrrrrrr>r.r/rr2r2r3rSsH zFlavaMultimodalModel.forward)T)NNNNN)r8r9r:rr rr0rerrrr$r!FLAVA_MULTIMODAL_INPUTS_DOCSTRINGr3rr4r&_CONFIG_CLASS_FOR_MULTIMODAL_MODEL_DOCr<rrrIr r6rrr2r2rrr3rA3s6 rAz_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) FlavaModelrcs0t|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 |_tt|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 )#rdrer text_configr TypeErrortype image_configrmultimodal_configrZprojection_dimriZtext_hidden_sizeZimage_hidden_sizeZmm_hidden_sizer7 text_modelr image_modelrAmultimodal_modelr rimage_projectiontext_projectionrgr<Ztensorr]Zlogit_scale_init_value logit_scaleimage_to_mm_projectiontext_to_mm_projectionr)r1r]rGrJrKrrr2r3res>       zFlavaModel.__init__r9N)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)r3r4rLrP) r1rrrrrrrZ text_outputsrZ 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)r3r4rMrO) r1rrrrrrrrZ image_outputsrZimage_featuresr2r2r3get_image_featuress   zFlavaModel.get_image_featuresrCr'r T) rrrrrrimage_attention_maskskip_multimodal_encoderrrrr)c  Cs| dk r | n|jj} | s tdd} d} d}d}|dk rn|j|||| | | d}|d|d} } || d}d}d}d}d}|dk r|j||||| | | d}|d|d}}||d}d}d}|dk r\|dk r\|s\|dk r.|j\}}}|jj r|d7}t j |||j d }t j ||gdd }nd}t j ||gdd }|j||| d }|d}| sr| |||||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) >>> image_embeddings = outputs.image_embeddings >>> text_embeddings = outputs.text_embeddings >>> multimodal_embeddings = outputs.multimodal_embeddings >>> outputs.image_embeddings.shape torch.Size([1, 197, 768]) >>> text_embeddings.shape torch.Size([1, 7, 768]) >>> multimodal_embeddings.shape torch.Size([1, 205, 768]) ``` NzRFLAVA model requires hidden states to work. Please set `output_hidden_states=True`)rrrrrrrrxrw)rrrrrrrrr;r|)rr)r"r#r$r%r&r')r]rrrMrRrLrSr~rNrBr<r<rrr!)r1rrrrrrrWrXrrrr"Z image_statesZimage_mm_projectionr#r$Z text_statesZtext_mm_projectionr%r&r'rrrZattention_mask_imageZattention_multimodalZmultimodal_inputr2r2r3r%s5     zFlavaModel.forward)NNNNNNN)NNNNNNNN) NNNNNNNNNTN)r8r9r:rr rerr?r3rr<rrIr=rTr2rrUFLAVA_MODEL_INPUTS_DOCSTRINGrr! LongTensorr r rrrr2r2rrr3rFs+ + 1  rFcs8eZdZeedfdd 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)rdrerr ReLUr Sequentialpath)r1r\r]kwargsZhid_sizercrrr2r3res     z"FlavaImageCodebookResPath.__init__rcCs ||SrZ)rcr1rr2r2r3rsz!FlavaImageCodebookResPath.forward r8r9r:rrer<rrrr2r2rrr3r[sr[cs:eZdZeeedfdd ZejejdddZZS)FlavaImageCodebookBlock)r\r] num_layersc sPtd|d|_||kr6tj||ddd|_n t|_t|||_dS)Nrrxrr_) rdre post_gainr rid_pathZIdentityr[res_path)r1r\r]rhrdrrr2r3res   z FlavaImageCodebookBlock.__init__rcCs|||j||SrZ)rjrirkrer2r2r3rszFlavaImageCodebookBlock.forwardrfr2r2rrr3rgs rgcs@eZdZdeeeeedfdd ZejejdddZZ S) FlavaImageCodebookLayerGroupT) num_blocksrhr\r]use_poolcstt}t|D]B}|dkr@t||||d|d<qt||||d|d<q|rptjdd|d<t||_dS)NrZblock_rrx)rpool) rdrerrrgr Z MaxPool2drbgroup)r1rmrhr\r]rnblocksrrrr2r3res  z%FlavaImageCodebookLayerGroup.__init__rcCs ||SrZ)rprer2r2r3rsz$FlavaImageCodebookLayerGroup.forward)T) r8r9r:rrIrer<rrrr2r2rrr3rlsrla" 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]rdc 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)NZrelurrr_convr inputZgroup_1rxZgroup_2r^Zgroup_3F)rnZgroup_4r)rdrer]Z num_groupsinput_channelsZnum_blocks_per_grouprirrr rarrlrbrqrfreeze parametersZ requires_grad)r1r]rdrhZ output_blocksrqparamrrr2r3resX       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)r3_CHECKPOINT_FOR_CODEBOOK_DOCrqr<Zargmaxr1rZz_logitsr2r2r3get_codebook_indicess  z'FlavaImageCodebook.get_codebook_indicescCs||}tjdd|S)Nrr|)rqr ZSoftmaxr}r2r2r3get_codebook_probs6s 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) ``` r^z input shape z is not 4drz input has z channels but model built for )r3r|rr~rrxrq)r1rr2r2r3r:szFlavaImageCodebook.forward)r8r9r:rrr r0rrrer<rr~rr=rrr2r2rrr3rrs ,rrcs$eZdZfddZddZZS)FlavaPredictionHeadTransformcsVtt|j|j|_t|jtr6t |j|_ n|j|_ tj |j|j d|_ dSr) rdrer rrirrrrr transform_act_fnrrrrrr2r3re^s   z%FlavaPredictionHeadTransform.__init__cCs"||}||}||}|SrZ)rrrrr2r2r3rgs   z$FlavaPredictionHeadTransform.forwardr8r9r:rerrr2r2rrr3r]s rcs.eZdZdfdd ZddZddZZS) FlavaMaskedPredictionHeadNcsbt||_t||_tj|j|jdd|_ t t |j|_ |dk rT||j _|j |j _ dS)NFr)rdrer]r transformr rrirdecoderrgr<rhrr )r1r]r rrr2r3reos  z"FlavaMaskedPredictionHead.__init__cCs|j|j_dSrZ)rrr0r2r2r3 _tie_weights{sz&FlavaMaskedPredictionHead._tie_weightscCs||}||}|SrZ)rrrer2r2r3r~s  z!FlavaMaskedPredictionHead.forward)N)r8r9r:rerrrr2r2rrr3rns rcs$eZdZfddZddZZS) FlavaITMHeadcs.t||_t||_t|jd|_dS)Nrx) rdrer]rrr rriseq_relationshiprrrr2r3res  zFlavaITMHead.__init__cCs||}||}|SrZ)rrrer2r2r3rs  zFlavaITMHead.forwardrr2r2rrr3rs rcs$eZdZfddZddZZS)FlavaGlobalContrastiveHeadcst||_|j|_dSrZ)rdrer]global_backprop_contrastiverrrr2r3res 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)Nrr;csg|]}tqSr2r<Z zeros_liker)r$r2r3rsz6FlavaGlobalContrastiveHead.forward..csg|]}tqSr2rr)r"r2r3rsr)r<expZ distributedZ is_availableZis_initializedrrrZget_world_sizerr rZ all_gatherrZget_rankrrr) r1r"r$rQZ temperaturelabelsZimage_embeddings_allZtext_embeddings_allZlocal_batch_sizeZ world_sizelogits_per_imagelogits_per_textr2)r"r$r3rs.     z"FlavaGlobalContrastiveHead.forwardrr2r2rrr3rs rzk The FLAVA model for pretraining which outputs losses, embeddings, logits and transformer outputs. cseZdZddddgZdeeejdfdd Ze j d d d Z e e d d eeeddee jee jee jee jee j ee j ee j ee jee j eee j ee j ee j eeeeeeeeee j efdddZZS)FlavaForPreTrainingzmmm_text_head.decoder.biaszmmm_image_head.decoder.biaszmlm_head.decoder.biaszmim_head.decoder.biasN)r]image_codebookcst|t||_||_|jdkr8|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)rdrerFrrZ init_codebookrrZimage_codebook_configrrJmim_headrGmlm_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]rrrr2r3res,           zFlavaForPreTraining.__init__)rcCs"|dkr||dd}|S)Nrxrrw)r}rrrer2r2r3 _resize_to_2ds z!FlavaForPreTraining._resize_to_2dzbatch_size, text_seq_lenr%rVT)rinput_ids_maskedrcodebook_pixel_valuesrrrrrWr mlm_labels mim_labels itm_labelsrrr return_lossr)c6Csz|dk r |n|jj}|dk r |n|jj}| dk r4| n|j} |dkrX|dk rXtd|}|j|||||| | ||dd }|j||||| |||dd }d}|j}|j}|j}|j}|j }d}}}}}}} d}!}"}#}$d}%}&}'|dk s|dk r@| dkr@|r@|j dkr"t d|dkr4t d|j |} |jdkr|dk r|dkr|}(| dk r|| } ||}|j| |d<|(dd| d  dddf}(| |j})| |)}*|(|)ddf}(||(}!|rtj|!d |j|*d }||j9}n ||(}!|jdkr|dk r|dkr|}+| dk r|| } |+dd| d  dddf}+| |j})| |)},|+|)ddf}+||+}"|rtj|"d |j|,d }||j9}n ||+}"|jdkr|dk r||}%| dk r| d}-t|- |-|-!dg}|rJtj|%| } | |j9} |dk r\||}| dk rn| |} | dk r| |} ||}|dk rd|j"dkrd|}(|d d }.|(ddd d |.ddf}(| dk rZ|| } ||}|j| |d<| |j})| |)}*|(|)ddf}(|#|(}$|rdtj|$d |j|*d }||j"9}n |#|(}$|dk r|j$dkr|}+|+dd|d  dddf}+| dk r|| } | |j})| |)},|+|)ddf}+|%|+}#|rtj|#d |j|,d }||j$9}n |%|+}#|dk r|dk r|j&dkr|j'|dddddf}/tjj(|/d d }/|j)|dddddf}0tjj(|0d d }0|jj*j+,t-t.|/|0|/|jj*\}&}'}1|dk r|&|}&|'|}'|1|}1|rtj|&|1}2tj|'|1}3|2|3d }||j&9}t0||| |||d }4|rV|41sVt2dd|43D}|s:||j4dk rt|j45nd||j6dk r|j65nd|j |j7dk r|j75nd||j4dk r|j45nd||j6dk r|j65nd||j7dk r|j75nd|!|"|%|&|&|$|#f}5|r(|41s(||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) rrrrrrWrXrrr) rrrrrWrrrrz`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 Truerrrwrxr|)r@rArBrCrDrEcss|]}|dk r|ndVqdSrr2)r.rKr2r2r3r4sz.FlavaForPreTraining.forward..css|]}|dkr|VqdSrZr2)r.rr2r2r3r4s)rKrLr"r#r$r%r&r'rMrNrOrPrQrRrSrTrUrVrWrXrY):r]r+rrloggerwarningrr"r$r&r RuntimeErrorrr~rrrnerrr rZ cross_entropyrrrrrrrr<whereanynewrrrrrrP normalizerOrQr Zclamp_LOGIT_SCALE_CLAMP_MINLOGIT_SCALE_CLAMP_MAXrr?rGsumrFr#r,r%r'r6rJ)6r1rrrrrrrrrWrrrrrrrrZ 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_tied_weights_keysrrr r1rer<rrr"FLAVA_PRETRAINING_INPUTS_DOCSTRINGr3rrJrZr=rIr r rrr2r2rrr3rs`   r)\r;rrrZ dataclassesrtypingrrrrrr r r<Ztorch.utils.checkpointr Z activationsr Zmodeling_outputsrrZmodeling_utilsrrrutilsrrrrrrZconfiguration_flavarrrrrZ get_loggerr8rr4r|r5r@rEr6rrrr!r?rJr1r\rkrrrrrrrrrZFLAVA_START_DOCSTRINGZFLAVA_INPUTS_DOCSTRING_COMMONZ!FLAVA_IMAGE_INPUTS_DOCSTRING_BASEr2Z FLAVA_TEXT_INPUTS_DOCSTRING_BASEr?rDZ!FLAVA_MODEL_INPUTS_DOCSTRING_BASErYrZ'FLAVA_PRETRAINING_START_DOCSTRING_EXTRArr3rr7rArFr[rgrlrrrrrrrr2r2r2r3s  $     %$e]$9C*.3       +/ _ e ^  u (