a hXr@s8dZddlZddlZddlmZddlmZddlmZmZm Z ddl Z ddl Z ddl m Z ddl mZdd lmZdd lmZmZdd lmZdd lmZmZdd lmZmZmZmZddlmZm Z m!Z!m"Z"m#Z#e$e%Z&dZ'dZ(dZ)e e#e!e"fZ*eeddGdddeZ+eeddGdddeZ,eeddGdddeZ-Gddde j.Z/Gddde j.Z0Gd d!d!e j.Z1Gd"d#d#e j.Z2Gd$d%d%e j.Z3Gd&d'd'e j.Z4Gd(d)d)e j.Z5Gd*d+d+e j.Z6Gd,d-d-eZ7Gd.d/d/e j.Z8Gd0d1d1e j.Z9eGd2d3d3eZ:eGd4d5d5e:Z;eGd6d7d7e:ZGdd?d?e j.Z@Gd@dAdAe j.ZAedBdGdCdDdDe:ZBGdEdFdFe j.ZCGdGdHdHe j.ZDGdIdJdJe j.ZEGdKdLdLe j.ZFedMdGdNdOdOe:ZGgdPZHdS)QzPyTorch FLAVA model.N) OrderedDict) dataclass)AnyOptionalUnion)nn)ACT2FN)GradientCheckpointingLayer)BaseModelOutputBaseModelOutputWithPooling)PreTrainedModel) find_pruneable_heads_and_indicesprune_linear_layer) ModelOutputauto_docstringlogging torch_int) FlavaConfigFlavaImageCodebookConfigFlavaImageConfigFlavaMultimodalConfigFlavaTextConfigzfacebook/flava-image-codebookg$(~k@a 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. )Z custom_introc@seZdZUdZdZeejed<dZ ee ed<dZ eejed<dZ ee ed<dZ eejed<dZee ed<eed d d ZdS) FlavaModelOutputa 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))rrr Ngetattrto_tuple.0kselfd/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/flava/modeling_flava.py Vsz,FlavaModelOutput.to_tuple..tuplekeysr)r+r)r,r%Us zFlavaModelOutput.to_tuple)__name__ __module__ __qualname____doc__rrtorch FloatTensor__annotations__rr rrrr r/rr%r+r+r+r,r4s rz@ Class representing pretraining losses from FLAVA model c@seZdZUdZdZeejed<dZ eejed<dZ eejed<dZ eejed<dZ eejed<dZ eejed<ed d d ZdS) FlavaLossesa 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)r*all_nonevr+r+r,r@{s  zFlavaLosses.all_none)r1r2r3r4r9rr5r6r7r:r;r<r=r>boolr@r+r+r+r,r8\s r8a 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. 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)FlavaForPreTrainingOutputay 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. 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. 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. Nloss loss_inforrrrrr 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)rrr rIrGrKc3s,|]$}|vr|n t|VqdSNr#r&r*Ztransformer_outputsr+r,r-z5FlavaForPreTrainingOutput.to_tuple..r.r)r+rTr,r%sz"FlavaForPreTrainingOutput.to_tuple)"r1r2r3r4rDrr5r6r7rEr8rrr rrrr rFrGrHrIrJrKrLrMrNrOrPrQrRr/rr%r+r+r+r,rCs. 7 rCcsdeZdZdZd 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|_|j |_ ||_dS)Nr image_size patch_size num_channels embed_dim)super__init__ mask_tokenr Parameterr5zeros hidden_size cls_tokenPatchEmbeddingsrZr[r\patch_embeddings num_patchesposition_embeddingsDropouthidden_dropout_probdropoutrW)r*rWrXrg __class__r+r,r_s   zFlavaImageEmbeddings.__init__) embeddingsheightwidthr"c Cs|jdd}|jjdd}tjs>||kr>||kr>|jS|jddddf}|jddddf}|jd}||j} ||j} t|d} |d| | |}|dddd}t j j || | fdd d }|dddd dd|}tj ||fdd S) a This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution images. This method is also adapted to support torch.jit tracing. Adapted from: - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211 rNg?rrZbicubicF)sizemodeZ align_cornersdim)shaperhr5Zjit is_tracingr[rZreshapepermuter functionalZ interpolateviewcat) r*rnrorprgZ num_positionsZclass_pos_embedZpatch_pos_embedrvZ new_heightZ new_widthZsqrt_num_positionsr+r+r,interpolate_pos_encodings(     z-FlavaImageEmbeddings.interpolate_pos_encoding) pixel_valuesbool_masked_posr}r"cCs|j\}}}}|j||d}|\}} } |dur|j|| d} |dkr`||dd}|d| } |d| | | }|j |dd} t j | |fdd}|r|| |||}n ||j }||}|S)N)r}rqrr?rru)rwrfrsr`expandrvr{Z unsqueezeZtype_asrdr5r|r}rhrk)r*r~rr} batch_sizer\rorprnseq_len_Z mask_tokensmask cls_tokensr+r+r,forward*s    zFlavaImageEmbeddings.forward)F)NF)r1r2r3r4rrBr_r5Tensorintr}r BoolTensorr __classcell__r+r+rlr,rVs+rVcsVeZdZdZd eeeeeeffeedfdd Zdej e ej d d d Z Z S)rez# Image to Patch Embedding. rrYcstt|tjjs ||f}t|tjjs6||f}|d|d|d|d}||_||_||_t j ||||d|_ dS)Nrr) kernel_sizeZstride) r^r_ isinstance collectionsabcIterablerZr[rgrConv2d projection)r*rZr[r\r]rgrlr+r,r_Ss  zPatchEmbeddings.__init__F)r~r}r"c Csx|j\}}}}|s\||jdks.||jdkr\td|d|d|jdd|jdd ||ddd}|S)NrrzInput image size (*z) doesn't match model (z).rr)rwrZ ValueErrorrflatten transpose)r*r~r}rr\rorpxr+r+r,rfszPatchEmbeddings.forward)rrrr)F) r1r2r3r4rrr/r_r5rrBrrr+r+rlr,reNsrecsFeZdZdZfddZdeejeejeejdddZZ 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_idxepsposition_embedding_typeabsolute position_ids)rrqF) persistenttoken_type_ids)dtype)r^r_r Embedding vocab_sizercZ pad_token_idword_embeddingsZmax_position_embeddingsrhZtype_vocab_sizetoken_type_embeddings LayerNormlayer_norm_epsrirjrkr$rZregister_bufferr5arangerrbrrslongr*rWrlr+r,r_us zFlavaTextEmbeddings.__init__N input_idsrrc Cs|}|d}|dur.|jddd|f}|durt|drl|jddd|f}||d|}|}ntj|tj|jjd}| |}| |} || } |j dkr| |} | | 7} | | } || } | S)Nrrr)rdevicer)rsrhasattrrrr5rbrrrrrrhrrk) r*rrr input_shape seq_lengthZbuffered_token_type_idsZ buffered_token_type_ids_expandedZ inputs_embedsrrnrhr+r+r,rs&       zFlavaTextEmbeddings.forward)NNN) r1r2r3r4r_rr5rrrr+r+rlr,rrs rc sdeZdZeddfdd Zd ejeejeejee e ejejfe ejfdddZ Z S) FlavaSelfAttentionNrWr"cst|j|jdkr>t|ds>td|jd|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)r^r_rcnum_attention_headsrrrattention_head_size all_head_sizerLinearZqkv_biasquerykeyvalueriZattention_probs_dropout_probrkrrlr+r,r_s  zFlavaSelfAttention.__init__F hidden_statesattention_mask head_maskoutput_attentionsr"cCs*|j\}}}|||d|j|jdd}|||d|j|jdd} |||d|j|jdd} t || dd} | t |j} |dur| |} t j j| dd} || } |dur| |} t | | } | dddd} | dd|jf}| j|} |r | | fn| f}|S)Nrqrrrrurr)rwrr{rrrrrr5matmulmathsqrtrrzZsoftmaxrkry contiguousrsr)r*rrrrrrrZ query_layerZ key_layerZ value_layerZattention_scoresZattention_probsZ context_layerZnew_context_layer_shapeoutputsr+r+r,rsB          zFlavaSelfAttention.forward)NNF) r1r2r3FlavaPossibleConfigsr_r5rrrBrr/rrr+r+rlr,rsrcs@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|_dSrS) r^r_rrrcdenserirjrkrrlr+r,r_s zFlavaSelfOutput.__init__r input_tensorr"cCs||}||}|SrSrrkr*rrr+r+r,rs  zFlavaSelfOutput.forward) r1r2r3r4rr_r5rrrr+r+rlr,rsrc sxeZdZeddfdd ZeeddddZd ej e ej e ej e e e ej ej fe ej fd d d ZZS) FlavaAttentionNrcs*tt||_t||_t|_dSrS)r^r_r attentionroutputset pruned_headsrrlr+r,r_s   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)Nrrru)lenrrrrrrrrrrrrunion)r*rindexr+r+r, prune_heads s zFlavaAttention.prune_headsFrcCs8|j||||d}||d|}|f|dd}|SN)rrrrr)rr)r*rrrrZ self_outputsattention_outputrr+r+r,rs zFlavaAttention.forward)NNF)r1r2r3rr_rrrr5rrrBrr/rrr+r+rlr,rsrcs8eZdZeddfdd ZejejdddZZS)FlavaIntermediateNrcsBtt|j|j|_t|jt r6t |j|_ n|j|_ dSrS) r^r_rrrcintermediate_sizerr hidden_actstrr intermediate_act_fnrrlr+r,r_1s   zFlavaIntermediate.__init__)rr"cCs||}||}|SrS)rrr*rr+r+r,r:s  zFlavaIntermediate.forward r1r2r3rr_r5rrrr+r+rlr,r0s rcs<eZdZeddfdd ZejejejdddZZS) FlavaOutputNrcs.tt|j|j|_t|j|_ dSrS) r^r_rrrrcrrirjrkrrlr+r,r_Bs zFlavaOutput.__init__rcCs ||}||}||}|SrSrrr+r+r,rHs  zFlavaOutput.forwardrr+r+rlr,rAsrc 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|_dS)Nrr)r^r_Zchunk_size_feed_forwardZ seq_len_dimrrr intermediaterrrrrcrlayernorm_beforelayernorm_afterrrlr+r,r_Ts    zFlavaLayer.__init__Frc Csb|j|||||d}|d}|dd}||}||}||}|||}|f|}|Sr)rrrrr) r*rrrrZself_attention_outputsrrZ layer_outputr+r+r,r`s     zFlavaLayer.forward)NNF)r1r2r3r4rr_r5rrrBrr/rrr+r+rlr,rQsrc sVeZdZeddfdd Zd ejeejeejeeee e e fddd Z Z S) FlavaEncoderNrcs:t|_tfddtjD|_d|_dS)Ncsg|] }tqSr+)rr'rrWr+r, rUz)FlavaEncoder.__init__..F) r^r_rWrZ ModuleListrangenum_hidden_layerslayerZgradient_checkpointingrrlrr,r_s  zFlavaEncoder.__init__FTrrrroutput_hidden_states return_dictr"c Cs|rdnd}|rdnd}t|jD]R\} } |r8||f}|durH|| nd} | ||| |} | d}|r"|| df}q"|r||f}|stdd|||fDSt|||dS)Nr+rrcss|]}|dur|VqdSrSr+)r'rAr+r+r,r-rUz'FlavaEncoder.forward..)last_hidden_stater attentions) enumeraterr/r ) r*rrrrrrZall_hidden_statesZall_self_attentionsiZ layer_moduleZlayer_head_maskZ layer_outputsr+r+r,rs"    zFlavaEncoder.forward)NNFFT)r1r2r3rr_r5rrrBrr/r rrr+r+rlr,rs  rcs2eZdZedfdd ZejdddZZS) FlavaPoolerrcs*tt|j|j|_t|_dSrS)r^r_rrrcrZTanh activationrrlr+r,r_s zFlavaPooler.__init__)rcCs(|dddf}||}||}|SNr)rr)r*rZfirst_token_tensor pooled_outputr+r+r,rs  zFlavaPooler.forwardrr+r+rlr,rsrc@s>eZdZUeed<dZdZeej ej ej fddddZ dS)FlavaPreTrainedModelrWflavaTN)moduler"cCs>t|tjtjfr@|jjjd|jjd|j dur>|j j nt|tj r|jjjd|jjd|j dur~|jj|j  nt|tj r|j j |jjdnt|tr|j j nzt|tr|jj |jj |jdur|jj n@t|tr|jr:|jj nt|tr:|jj|jjdS)zInitialize the weightsg)meanZstdNr)rrrrweightdataZnormal_rWZinitializer_rangerZzero_rrrZfill_FlavaMaskedPredictionHeadrVrdrhr`FlavaMultimodalModel use_cls_token FlavaModel logit_scalelogit_scale_init_value)r*rr+r+r, _init_weightss.            z"FlavaPreTrainedModel._init_weights) r1r2r3rr7base_model_prefixsupports_gradient_checkpointingrrrrrrr+r+r+r,rs rc seZdZUeed<dZdZdeedfdd Ze j dd d Z e j d d d Z e eeefddddZedeejeejeeeejeejeeeeeeeeefd ddZZS)FlavaImageModelrWzflava.image_modelr~TrWadd_pooling_layercsXt|||_t||_t||_tj|j |j d|_ |rFt |nd|_ |dSv add_pooling_layer (bool, *optional*, defaults to `True`): Whether to add a pooling layer rN)r^r_rWrVrnrencoderrrrcr layernormrpooler post_initr*rWr rlr+r,r_s   zFlavaImageModel.__init__r!cCs|jjSrSrnrfr)r+r+r,get_input_embeddingssz$FlavaImageModel.get_input_embeddingsrcCs ||j_dSrSrr*rr+r+r,set_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 Nitemsrrrrr*rrrr+r+r, _prune_headsszFlavaImageModel._prune_heads r~rr}rrrrrr"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)z 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). Nz You have to specify pixel_values)rr}rrrrrrrrZ pooler_outputrr)rWrruse_return_dictr get_head_maskrrnrrrr rr) r*r~rr}rrrrrembedding_outputencoder_outputssequence_outputrr+r+r,rs: zFlavaImageModel.forward)T)NNNNNNNN)r1r2r3rr7r main_input_namerBr_rModulerrdictrlistr rrr5rrrr/r rrr+r+rlr,r s6  r c seZdZUeed<dZdeedfdd Zeddd Z e j d d d Z e eeefd dddZedeejeejeejeejeejeeeeeeeeefd ddZZS)FlavaTextModelrWzflava.text_modelTr csXt|||_t||_t||_tj|j |j d|_ |rFt |nd|_ |dSr)r^r_rWrrnrrrrrcrrrrrrrlr+r,r_Bs   zFlavaTextModel.__init__r!cCs|jjSrSrnrr)r+r+r,rRsz#FlavaTextModel.get_input_embeddingsrcCs ||j_dSrSr.rr+r+r,rUsz#FlavaTextModel.set_input_embeddingsNrcCs*|D]\}}|jj|j|qdSrrrr+r+r,r XszFlavaTextModel._prune_heads) 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) a input_ids (`torch.LongTensor` of shape `(batch_size, text_seq_length)`): 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 `(batch_size, text_seq_length)`, *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) NzYou have to specify input_idsrrr"rrr#)rWrrr$rrsr5onesrr%rget_extended_attention_maskrnrrrr rr)r*rrrrrrrrrextended_attention_maskr&r'r(rr+r+r,r`sJ zFlavaTextModel.forward)T)NNNNNNNN)r1r2r3rr7r rBr_rerrr*rr+rr,r rrr5rrr/r rrr+r+rlr,r-<s4  r-c seZdZUeed<dZdZdedfdd Zee e e fdd d d Z e de jee jee jeeeeeeeeefd d dZZS)rrWzflava.multimodal_modelrTrcsvt|||_|jj|_|jr:ttdd|j|_ t ||_ tj |j|j d|_|rdt|nd|_|dS)rrrN)r^r_rWrrrar5rbrcrdrrrrrrrrrrlr+r,r_s   zFlavaMultimodalModel.__init__NrcCs*|D]\}}|jj|j|qdSrrrr+r+r,r sz!FlavaMultimodalModel._prune_headsrcCs(|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) z hidden_states (`torch.FloatTensor` of shape `(batch_size, image_num_patches + text_seq_len, hidden_size)`): The concatenated hidden states of unimodal encoders. Nrqrrur/r"rr#)rWrrr$rsrrdrr5r|r0rr%rr1rrrr rr)r*rrrrrrrrrrr2r'r(rr+r+r,rsD  zFlavaMultimodalModel.forward)T)NNNNN)r1r2r3rr7r r)r_r+rr,r rr5rrrBrr/r rrr+r+rlr,rs(  rcs*eZdZUeed<edfdd Zedeej eej eej eej ee ee ee ej dddZ edeej eej ee eej eej ee ee ee ej d d d Zedeejeej eej eej eej eejeej ee ee e ee eeefd ddZZS)rrWrcs0t|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 )#r^r_r text_configr TypeErrortype image_configrmultimodal_configrZprojection_dimrcZtext_hidden_sizeZimage_hidden_sizeZmm_hidden_sizer- text_modelr  image_modelrmultimodal_modelrrimage_projectiontext_projectionrar5ZtensorrWrrimage_to_mm_projectiontext_to_mm_projectionr)r*rWr3r6r7rlr+r,r_sF       zFlavaModel.__init__N)rrrrrrrr"c Cs.|j|||||||d}|d} || } | S)a input_ids (`torch.LongTensor` of shape `(batch_size, text_seq_length)`): 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 `(batch_size, text_seq_length)`, *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) 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)r8r<) r*rrrrrrrZ text_outputsrZ text_featuresr+r+r,get_text_features<s)  zFlavaModel.get_text_featuresr!c Cs0|j||||||||d} | d} || } | S)a 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). 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) ``` )r~rrrrrr}rr)r9r;) r*r~rr}rrrrrZ image_outputsrZimage_featuresr+r+r,get_image_featurests&  zFlavaModel.get_image_featuresT) rr~rrrrimage_attention_maskskip_multimodal_encoderrrrr"c  Cs| 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\|s\|dur.|j\}}}|jj r|d7}t j |||j d }t j ||gdd }nd}t j ||gdd }|j||| d }|d}| sr| |||||fSt| |||||d S) a input_ids (`torch.LongTensor` of shape `(batch_size, image_num_patches + text_seq_len)`): 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 `(batch_size, image_num_patches + text_seq_len)`, *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) 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). image_attention_mask (`torch.Tensor` of shape `(batch_size, image_num_patches)`, *optional*): Mask to avoid performing attention on padding pixel values for image inputs. Mask values selected in `[0, 1]`: - 1 for pixel values that are real (i.e., **not masked**), - 0 for pixel values that are padding (i.e., **masked**). skip_multimodal_encoder (*bool*, *optional*): Skip any calculations for multimodal encoder. Useful if multimodal encoding is not going to be used. 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`)r~rrrrrrrrrq)rrrrrrrrr/ru)rr)rrrrrr )rWrrr9r=r8r>rwr:rr5r0rr|r)r*rr~rrrrrArBrrrrZ image_statesZimage_mm_projectionrrZ text_statesZtext_mm_projectionrrr rrrZattention_mask_imageZattention_multimodalZmultimodal_inputr+r+r,rsC     zFlavaModel.forward)NNNNNNN)NNNNNNNN) NNNNNNNNNTN)r1r2r3rr7r_rrr5rrBr6r?rr@ LongTensorrr/rrrr+r+rlr,r s +75 rcs8eZdZeedfdd 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_1rrrpaddingZconv_1Zrelu_2Zconv_2Zrelu_3Zconv_3Zrelu_4rZconv_4)r^r_rrReLUr Sequentialpath)r*rErFkwargsZhid_sizerLrlr+r,r_:s     z"FlavaImageCodebookResPath.__init__rr"cCs ||SrS)rLr*rr+r+r,rJsz!FlavaImageCodebookResPath.forward r1r2r3rr_r5rrrr+r+rlr,rD9srDcs:eZdZeeedfdd ZejejdddZZS)FlavaImageCodebookBlock)rErF num_layersc sPtd|d|_||kr6tj||ddd|_n t|_t|||_dS)NrrrrrH) r^r_ post_gainrrid_pathZIdentityrDres_path)r*rErFrRrMrlr+r,r_Os   z FlavaImageCodebookBlock.__init__rNcCs|||j||SrS)rTrSrUrOr+r+r,r[szFlavaImageCodebookBlock.forwardrPr+r+rlr,rQNs rQcs@eZdZdeeeeedfdd ZejejdddZZ S) FlavaImageCodebookLayerGroupT) num_blocksrRrErFuse_poolcstt}t|D]B}|dkr@t||||d|d<qt||||d|d<q|rptjdd|d<t||_dS)NrZblock_rrr)rpool) r^r_rrrQrZ MaxPool2drKgroup)r*rWrRrErFrXblocksrrlr+r,r_`s  z%FlavaImageCodebookLayerGroup.__init__rNcCs ||SrS)rZrOr+r+r,rnsz$FlavaImageCodebookLayerGroup.forward)T) r1r2r3rrBr_r5rrrr+r+rlr,rV_srVa" 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. csveZdZUdZeed<dZdZeedfdd Z e j e j dd d Z e j e j dd d Z e je j dd dZZS)FlavaImageCodebookrWr~F)rWrMc 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)NZrelurrrHconvrinputZgroup_1rrZgroup_2rGZgroup_3F)rXZgroup_4r)r^r_rWZ num_groupsinput_channelsZnum_blocks_per_grouprcrrrrJrrVrKr[rfreeze parametersZ requires_grad)r*rWrMrRZ output_blocksr[paramrlr+r,r_s>       zFlavaImageCodebook.__init__)r~r"cCs*dtdtd||}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("E") >>> image_processor = AutoImageProcessor.from_pretrained("a") >>> 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)_CHECKPOINT_FOR_CODEBOOK_DOCr[r5Zargmaxr*r~Zz_logitsr+r+r,get_codebook_indicess   z'FlavaImageCodebook.get_codebook_indicescCs||}tjdd|S)Nrru)r[rZSoftmaxrhr+r+r,get_codebook_probss z%FlavaImageCodebook.get_codebook_probscCshdtdtdt|jdkr2td|jd|jd|jkr^td|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("rfa") >>> 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) ``` rGz input shape z is not 4drz input has z channels but model built for )rgrrwrrbr[)r*r~r+r+r,rs  zFlavaImageCodebook.forward)r1r2r3r rr7r)r rr_r5rrirjr6rrr+r+rlr,r\ss ,r\cs$eZdZfddZddZZS)FlavaPredictionHeadTransformcsVtt|j|j|_t|jtr6t |j|_ n|j|_ tj |j|j d|_ dS)Nr) r^r_rrrcrrrrr transform_act_fnrrrrlr+r,r_s   z%FlavaPredictionHeadTransform.__init__cCs"||}||}||}|SrS)rrlrrr+r+r,rs   z$FlavaPredictionHeadTransform.forwardr1r2r3r_rrr+r+rlr,rks rkcs.eZdZdfdd ZddZddZZS) rNcsbt||_t||_tj|j|jdd|_ t t |j|_ |durT||j _|j |j _ dS)NFr)r^r_rWrk transformrrrcrdecoderrar5rbrr)r*rWrrlr+r,r_s  z"FlavaMaskedPredictionHead.__init__cCs|j|j_dSrS)rror)r+r+r, _tie_weights sz&FlavaMaskedPredictionHead._tie_weightscCs||}||}|SrS)rnrorOr+r+r,rs  z!FlavaMaskedPredictionHead.forward)N)r1r2r3r_rprrr+r+rlr,rs rcs$eZdZfddZddZZS) FlavaITMHeadcs.t||_t||_t|jd|_dS)Nrr) r^r_rWrrrrrcseq_relationshiprrlr+r,r_s  zFlavaITMHead.__init__cCs||}||}|SrS)rrrrOr+r+r,rs  zFlavaITMHead.forwardrmr+r+rlr,rqs rqcs$eZdZfddZddZZS)FlavaGlobalContrastiveHeadcst||_|j|_dSrS)r^r_rWglobal_backprop_contrastiverrlr+r,r_$s 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|]}tqSr+r5Z zeros_liker)rr+r,r9rUz6FlavaGlobalContrastiveHead.forward..csg|]}tqSr+rur)rr+r,r:rUr)r5exp distributedZ is_availableZis_initializedrrsrZget_world_sizertrrzZ all_gatherrZget_rankr|rr) r*rrrZ temperaturelabelsZimage_embeddings_allZtext_embeddings_allZlocal_batch_sizeZ world_sizelogits_per_imagelogits_per_textr+)rrr,r)s,     z"FlavaGlobalContrastiveHead.forwardrmr+r+rlr,rs#s rszk The FLAVA model for pretraining which outputs losses, embeddings, logits and transformer outputs. cseZdZgdZd eeejdfdd Ze j dddZ e dee j ee j ee jee jee j ee j ee j ee j ee j eeee j ee j ee j eeeeeeeeee j efd d d ZZS)FlavaForPreTraining)zmmm_text_head.decoder.biaszmmm_image_head.decoder.biaszmlm_head.decoder.biaszmim_head.decoder.biasN)rWimage_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|_|dS)z 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. N)r^r_rrr|Z init_codebookr\Zimage_codebook_configrr6mim_headr3mlm_headrqitm_headmmm_image_head mmm_text_headrsglobal_contrastive_headrimage_vocab_sizetext_vocab_size mlm_weight mim_weightglobal_contrastive_weightce_ignore_index itm_weightmmm_image_weightmmm_text_weight skip_unmasked_multimodal_encoderr)r*rWr|rlr+r,r_Ys,           zFlavaForPreTraining.__init__)rcCs"|dkr||dd}|S)Nrrrrq)rvr{rsrOr+r+r, _resize_to_2d|s z!FlavaForPreTraining._resize_to_2dT)rinput_ids_maskedr~codebook_pixel_valuesrrrrrAr mlm_labels mim_labels itm_labelsrrr return_lossr"c6Csz|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| |} ||}|durd|j"dkrd|}(|d d }.|(ddd d |.ddf}(| durZ|| } ||}|j| |d<| |j})| |)}*|(|)ddf}(|#|(}$|rdtj|$d |j|*d }||j"9}n |#|(}$|dur|j$dkr|}+|+dd|d  dddf}+| dur|| } | |j})| |)},|+|)ddf}+|%|+}#|rtj|#d |j|,d }||j$9}n |%|+}#|dur|dur|j&dkr|j'|dddddf}/tjj(|/d d }/|j)|dddddf}0tjj(|0d d }0|jj*j+,t-t.|/|0|/|jj*\}&}'}1|dur|&|}&|'|}'|1|}1|rtj|&|1}2tj|'|1}3|2|3d }||j&9}t0||| |||d }4|rV|41sVt2dd|43D}|s:||j4durt|j45nd||j6dur|j65nd|j |j7dur|j75nd||j4dur|j45nd||j6dur|j65nd||j7dur|j75nd|!|"|%|&|&|$|#f}5|r(|41s(||4f|5}5t8dd|5DSt9||4||j4||j6|j |j7||j4||j6||j7|!|"|%|&|'|$|#dS)a input_ids (`torch.LongTensor` of shape `(batch_size, text_seq_len)`): 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) input_ids_masked (`torch.LongTensor` of shape `(batch_size, text_seq_len)`): 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) codebook_pixel_values (`torch.FloatTensor` of shape `(batch_size, num_image_patches, patch_size, patch_size, 3)`, *optional*): Pixel values for image patches that are used to compute the image codebook labels for masked image modeling. token_type_ids (`torch.LongTensor` of shape `(batch_size, text_seq_len)`, *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) 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). image_attention_mask (`torch.FloatTensor` of shape `(batch_size, image_num_patches)`, *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. 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) ``` 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) rr~rrrrArBrrr) rr~rrrArrrrz`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 Truerrrqrrru)r9r:r;r<r=r>css|]}|dur|ndVqdSrr+)r'rDr+r+r,r-rUz.FlavaForPreTraining.forward..css|]}|dur|VqdSrSr+)r'rr+r+r,r-rU)rDrErrrrrr rFrGrHrIrJrKrLrMrNrOrPrQrR):rWr$rrloggerwarningrrrrr| RuntimeErrorrrirrrnersr}rrzZ cross_entropyr{rrr~rrrr5whereanynewrrrrrr< normalizer;rrZclamp_LOGIT_SCALE_CLAMP_MINLOGIT_SCALE_CLAMP_MAXrr8r@sumr?rr%rr r/rC)6r*rrr~rrrrrrArrrrrrrrZ flava_outputZflava_masked_outputZpos_maskrrrFrHrJZ total_lossZmim_lossZmlm_lossZ mmm_text_lossZmmm_image_lossZgc_lossZitm_lossrLrMrRrQrNryrzZsequence_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_lossesrr+r+r,rs_         "       "                  "            zFlavaForPreTraining.forward)N)NNNNNNNNNNNNNNTNN)r1r2r3Z_tied_weights_keysrrrr*r_r5rrrrCr6rBrr/rCrrr+r+rlr,r{KsR#r{)r{r\r rrrr-)Ir4rrr dataclassesrtypingrrrr5Ztorch.utils.checkpointrZ activationsr Zmodeling_layersr Zmodeling_outputsr r Zmodeling_utilsr Z pytorch_utilsrrutilsrrrrZconfiguration_flavarrrrrZ get_loggerr1rrgrrrrr8rCr*rVrerrrrrrrrrrr r-rrrDrQrVr\rkrrqrsr{__all__r+r+r+r,s        " \d$9I*.*"`p_-u (#