o Zhs@sdZddlZddlZddlmZddlmZddlmZmZm Z m Z m Z m Z m Z ddlZddlZddlmZddlmZdd lmZmZdd lmZmZmZdd lmZmZmZmZd d lm Z m!Z!m"Z"m#Z#m$Z$e%e&Z'dZ(dZ)dZ*e e$e"e#fZ+eGdddeZ,eGdddeZ-eGdddeZ.Gdddej/Z0Gdddej/Z1Gdddej/Z2Gdddej/Z3Gdddej/Z4Gd d!d!ej/Z5Gd"d#d#ej/Z6Gd$d%d%ej/Z7Gd&d'd'ej/Z8Gd(d)d)ej/Z9Gd*d+d+ej/Z:eGd,d-d-eZ;eGd.d/d/e;ZeGd4d5d5e;Z?Gd6d7d7ej/Z@Gd8d9d9ej/ZAGd:d;d;ej/ZBedd?d?e;ZCGd@dAdAej/ZDGdBdCdCej/ZEGdDdEdEej/ZFGdFdGdGej/ZGedHd=GdIdJdJe;ZHgdKZIdS)LzPyTorch FLAVA model.N) OrderedDict) dataclass)AnyDictListOptionalSetTupleUnion)nn)ACT2FN)BaseModelOutputBaseModelOutputWithPooling)PreTrainedModel find_pruneable_heads_and_indicesprune_linear_layer) ModelOutputauto_docstringlogging torch_int) FlavaConfigFlavaImageCodebookConfigFlavaImageConfigFlavaMultimodalConfigFlavaTextConfigzfacebook/flava-image-codebookg$(~k@c@seZdZUdZdZeejed<dZ ee ed<dZ eejed<dZ ee ed<dZ eejed<dZee ed<d eefd 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 |nt|VqdS))r!rr#Ngetattrto_tuple.0kselfW/var/www/auris/lib/python3.10/site-packages/transformers/models/flava/modeling_flava.py Rs  z,FlavaModelOutput.to_tuple..tuplekeysr+r-r+r.r'Qs zFlavaModelOutput.to_tuple)__name__ __module__ __qualname____doc__rrtorch FloatTensor__annotations__rrr r!r"r#r rr'r-r-r-r.r2s rc@seZdZUdZdZeejed<dZ eejed<dZ eejed<dZ eejed<dZ eejed<dZ eejed<d efd 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] }|durd}|Sq|S)NTF)values)r,all_nonevr-r-r.rBts zFlavaLosses.all_none)r3r4r5r6r;rr7r8r9r<r=r>r?r@boolrBr-r-r-r.r:Xs 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<dee fddZ!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_inforrr 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!rr#rKrIrMc3s.|]}|vr |nt|VqdSNr%r(r,Ztransformer_outputsr-r.r/s,z5FlavaForPreTrainingOutput.to_tuple..r0r+r-rVr.r'sz"FlavaForPreTrainingOutput.to_tuple)"r3r4r5r6rFrr7r8r9rGr:rrrr r!r"r#rHrIrJrKrLrMrNrOrPrQrRrSrTr rr'r-r-r-r.rE}s0 @ rEc sxeZdZdZddededdffdd Zd ejd e d e dejfd d Z  ddejde ej dedejfddZ ZS)FlavaImageEmbeddingszb Construct the CLS token, position and patch embeddings. Optionally, also the mask token. Fconfiguse_mask_tokenr$Ncst|p |j}ttdd|j|_|r#ttdd|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 Parameterr7zeros hidden_size cls_tokenPatchEmbeddingsrZr[r\patch_embeddings num_patchesposition_embeddingsDropouthidden_dropout_probdropoutrX)r,rXrYrg __class__r-r.r_s    zFlavaImageEmbeddings.__init__ embeddingsheightwidthc 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?rr ZbicubicF)sizemodeZ align_cornersdim)shaperhr7Zjit 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}cCs|j\}}}}|j||d}|\}} } |durB|j|| d} |dkr0||dd}|d| } |d| | | }|j |dd} t j | |fdd}|r_|| |||}n||j }||}|S)N)r}rqr r?rru)rwrfrsr`expandrvr{Z unsqueezeZtype_asrdr7r|r}rhrk)r,r~rr} batch_sizer\rorprnseq_len_Z mask_tokensmask cls_tokensr-r-r.forward#s    zFlavaImageEmbeddings.forwardF)NF)r3r4r5r6rrDr_r7Tensorintr}r BoolTensorr __classcell__r-r-rlr.rWs+rWc sheZdZdZ    ddedeeeeeffded effd d Zdd ej de dej fddZ Z S)rez# Image to Patch Embedding. r rZr[r\r]cstt|tjjs||f}t|tjjs||f}|d|d|d|d}||_||_||_t j ||||d|_ dS)Nrr) kernel_sizeZstride) r^r_ isinstance collectionsabcIterablerZr[rgr Conv2d projection)r,rZr[r\r]rgrlr-r.r_Ls  zPatchEmbeddings.__init__Fr~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.r_szPatchEmbeddings.forward)rrr rr) r3r4r5r6rr r r_r7rrDrrr-r-rlr.reGs $recsPeZdZdZfddZ   d deejdeejdeejfdd ZZ 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)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_bufferr7arangerrbrrslongr,rXrlr-r.r_ns  zFlavaTextEmbeddings.__init__N input_idsrrc Cs|}|d}|dur|jddd|f}|durAt|dr6|jddd|f}||d|}|}n tj|tj|jjd}| |}| |} || } |j dkr]| |} | | 7} | | } || } | S)Nrrr)rdevicer)rsrhasattrrrr7rbrrrrrrhrrk) r,rrr input_shape seq_lengthZbuffered_token_type_idsZ buffered_token_type_ids_expandedZ inputs_embedsrrnrhr-r-r.rs&       zFlavaTextEmbeddings.forward)NNN) r3r4r5r6r_rr7rrrr-r-rlr.rks rcseZdZdeddffdd ZdejdejfddZ   dd ejd eejd eejd e de e ejejfe ejff ddZ Z S)FlavaSelfAttentionrXr$Ncst|j|jdkrt|dstd|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_sizer LinearZqkv_biasquerykeyvalueriZattention_probs_dropout_probrkrrlr-r.r_s  zFlavaSelfAttention.__init__rcCs6|dd|j|jf}|j|}|ddddS)Nrqrrrrr )rsrrr{ry)r,rZ new_x_shaper-r-r.transpose_for_scoress z'FlavaSelfAttention.transpose_for_scoresF hidden_statesattention_mask head_maskoutput_attentionscCs||}|||}|||}||}t||dd} | t|j } |dur4| |} t j j | dd} | | } |durI| |} t| |} | dddd} | dd|jf} | j| } |rr| | f} | S| f} | S)Nrqrurrrrr )rrrrr7matmulrmathsqrtrr rzZsoftmaxrkry contiguousrsrr{)r,rrrrZmixed_query_layerZ key_layerZ value_layerZ query_layerZattention_scoresZattention_probsZ context_layerZnew_context_layer_shapeoutputsr-r-r.rs(      zFlavaSelfAttention.forwardNNF)r3r4r5FlavaPossibleConfigsr_r7rrrrDr r rrr-r-rlr.rs"rcsFeZdZdZdeddffdd Zdejdejdejfd d ZZ 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. rXr$Ncs.tt|j|j|_t|j|_dSrU) r^r_r rrcdenserirjrkrrlr-r.r_ zFlavaSelfOutput.__init__r input_tensorcC||}||}|SrUrrkr,rrr-r-r.r  zFlavaSelfOutput.forward) r3r4r5r6rr_r7rrrr-r-rlr.rs$rcseZdZdeddffdd ZdeeddfddZ   dd ej d e ej d e ej d e de e ej ej fe ej ff ddZZS)FlavaAttentionrXr$Ncs*tt||_t||_t|_dSrU)r^r_r attentionroutputset pruned_headsrrlr-r.r_s    zFlavaAttention.__init__headscCst|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_headss zFlavaAttention.prune_headsFrrrrcCs8|j||||d}||d|}|f|dd}|SN)rrrrr)rr)r,rrrrZ self_outputsattention_outputrr-r-r.rs zFlavaAttention.forwardr)r3r4r5rr_rrrr7rrrDr r rrr-r-rlr.rs"rcs<eZdZdeddffdd ZdejdejfddZZS) FlavaIntermediaterXr$NcsDtt|j|j|_t|jt rt |j|_ dS|j|_ dSrU) r^r_r rrcintermediate_sizerr hidden_actstrr intermediate_act_fnrrlr-r.r_$s   zFlavaIntermediate.__init__rcCrrU)rrr,rr-r-r.r-rzFlavaIntermediate.forward r3r4r5rr_r7rrrr-r-rlr.r#s rcsBeZdZdeddffdd Zdejdejdejfdd ZZS) FlavaOutputrXr$Ncs.tt|j|j|_t|j|_ dSrU) r^r_r rrrcrrirjrkrrlr-r.r_5rzFlavaOutput.__init__rrcCs ||}||}||}|SrUrrr-r-r.r;s  zFlavaOutput.forwardrr-r-rlr.r4s$rcsxeZdZdZdeddffdd Z   ddejd eejd eejd e de e ejejfe ejff d d Z Z S) FlavaLayerz?This corresponds to the Block class in the timm implementation.rXr$Ncsbt|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 intermediaterrr rrcrlayernorm_beforelayernorm_afterrrlr-r.r_Gs    zFlavaLayer.__init__Frrrrc Csb|j|||||d}|d}|dd}||}||}||}|||}|f|}|Sr)rrrrr) r,rrrrZself_attention_outputsrrZ layer_outputr-r-r.rSs     zFlavaLayer.forwardr)r3r4r5r6rr_r7rrrDr r rrr-r-rlr.rDs"rcsneZdZdeddffdd Z     ddejd eejd eejd ed ed ede e e ffddZ Z S) FlavaEncoderrXr$Ncs:t|_tfddtjD|_d|_dS)Ncsg|]}tqSr-)rr)rrXr-r. vsz)FlavaEncoder.__init__..F) r^r_rXr Z ModuleListrangenum_hidden_layerslayergradient_checkpointingrrlrr.r_ss   zFlavaEncoder.__init__FTrrrroutput_hidden_states return_dictc Cs|rdnd}|r dnd}t|jD]:\} } |r||f}|dur$|| nd} |jr7|jr7|| j||| |} n| ||| |} | d}|rK|| df}q|rS||f}|satdd|||fDSt|||dS)Nr-rrcss|] }|dur|VqdSrUr-)r)rCr-r-r.r/z'FlavaEncoder.forward..)last_hidden_stater attentions) enumeraterrZtrainingZ_gradient_checkpointing_func__call__r1r) r,rrrrrrZall_hidden_statesZall_self_attentionsiZ layer_moduleZlayer_head_maskZ layer_outputsr-r-r.rys4    zFlavaEncoder.forward)NNFFT)r3r4r5rr_r7rrrDr r1rrrr-r-rlr.rrs,  rcs2eZdZdeffdd ZdejfddZZS) FlavaPoolerrXcs*tt|j|j|_t|_dSrU)r^r_r rrcrZTanh 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:eZdZeZdZdZdeej ej ej fddfddZ dS)FlavaPreTrainedModelflavaTmoduler$NcCsXt|tjtjfr#|jjjd|jjd|j dur!|j j dSdSt|tj rF|jjjd|jjd|j durD|jj|j  dSdSt|tj r[|j j |jjddSt|trh|j j dSt|tr|jj |jj |jdur|jj dSdSt|tr|jr|jj dSdSt|tr|jj|jjdSdS)zInitialize the weightsg)meanZstdNr)rr rrweightdataZnormal_rXZinitializer_rangerZzero_rrrZfill_FlavaMaskedPredictionHeadrWrdrhr`FlavaMultimodalModel use_cls_token FlavaModel logit_scalelogit_scale_init_value)r,rr-r-r. _init_weightss8            z"FlavaPreTrainedModel._init_weights) r3r4r5r config_classbase_model_prefixsupports_gradient_checkpointingr r rrrr r-r-r-r.rs &rcseZdZeZdZdZddedeffdd Zde j fd d Z d e j fd d Z de eeefddfddZe        ddeejdeejdeedeejdeejdeedeedeedeeeffddZZS)FlavaImageModelzflava.image_modelr~TrXadd_pooling_layercXt|||_t||_t||_tj|j |j d|_ |r#t |nd|_ |dSv add_pooling_layer (bool, *optional*, defaults to `True`): Whether to add a pooling layer rN)r^r_rXrWrnrencoderr rrcr layernormrpooler post_initr,rXrrlr-r.r_s    zFlavaImageModel.__init__r$cC|jjSrUrnrfr+r-r-r.get_input_embeddingsz$FlavaImageModel.get_input_embeddingsrcC ||j_dSrUrr,rr-r-r.set_input_embeddings z$FlavaImageModel.set_input_embeddingsheads_to_pruneNcC*|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,r!rrr-r-r. _prune_headszFlavaImageModel._prune_headsrr}rrrrrc Cs|dur|n|jj}|dur|n|jj}|dur|n|jj}|dur&td|||jj}|j|||d} |j| |||||d} | d} | | } |j durT| | nd} |sb| | 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)rXrruse_return_dictr get_head_maskrrnrrrrrr) r,r~rr}rrrrrembedding_outputencoder_outputssequence_outputrr-r-r.rs: zFlavaImageModel.forwardTNNNNNNNN)r3r4r5rr r main_input_namerDr_r Modulerrrrrr'rrr7rrr r1rrrr-r-rlr.rsH  rcseZdZeZdZddedeffdd Zdefdd Z d e j fd d Z d e eeefddfddZe        ddeejdeejdeejdeejdeejdeedeedeedeeeffddZZS)FlavaTextModelzflava.text_modelTrXrcrr)r^r_rXrrnrrr rrcrrrrrrrlr-r.r_>s    zFlavaTextModel.__init__r$cCrrUrnrr+r-r-r.rNrz#FlavaTextModel.get_input_embeddingsrcCrrUr5rr-r-r.rQr z#FlavaTextModel.set_input_embeddingsr!NcCr"r#r$r&r-r-r.r'Tr(zFlavaTextModel._prune_headsrrrrrrrrc Cs|dur|n|jj}|dur|n|jj}|dur|n|jj}|dur&td|} |dur6tj| |jd}| ||jj }| || |j} |j |||d} |j | | ||||d} | d} || } |jdurl|| nd}|sz| |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_idsr)rrrr)rrr*)rXrrr+rrsr7onesrr,rget_extended_attention_maskrnrrrrrr)r,rrrrrrrrrextended_attention_maskr-r.r/rr-r-r.r\sJ zFlavaTextModel.forwardr0r1)r3r4r5rr r rDr_rerr r3rrrrr'rrr7rr r1rrrr-r-rlr.r48sF  r4cseZdZeZdZdZddeffdd Zdee e e fdd fd d Z e dde jd ee jd ee jdeedeedeedeeeffddZZS)rzflava.multimodal_modelrTrXcsvt|||_|jj|_|jrttdd|j|_ t ||_ tj |j|j d|_|r2t|nd|_|dS)rrrN)r^r_rXrr rar7rbrcrdrrrrrrrrrrlr-r.r_s    zFlavaMultimodalModel.__init__r!r$NcCr"r#r$r&r-r-r.r'r(z!FlavaMultimodalModel._prune_headsrrrrrcCs&|dur|n|jj}|dur|n|jj}|dur|n|jj}|\}}} |jr=|j|dd} tj | |fdd}|d7}|durKtj ||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. Nrqrrur6r)rr*)rXrrr+rsrrdrr7r|r7rr,rr8rrrrrr)r,rrrrrrrrrrr9r.r/rr-r-r.rsD  zFlavaMultimodalModel.forwardr0)NNNNN)r3r4r5rr r r2r_rrrr'rr7rrrDr r1rrrr-r-rlr.rs6 rcseZdZeZdeffdd Ze       ddeej deej deej deej d ee d ee d ee d ej fd dZ e        ddeej deej dee deej deej d ee d ee d ee d ej fddZe           ddeejdeej deej deej deej deejdeej dee d ee d e d ee d eeeffddZZS)rrXcs0t|t|jtstdt|jdt|jts(tdt|jdt|j t s;tddt|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_sizer4 text_modelr image_modelrmultimodal_modelr rimage_projectiontext_projectionrar7ZtensorrXr r image_to_mm_projectiontext_to_mm_projectionr)r,rXr:r=r>rlr-r.r_ sF        zFlavaModel.__init__Nrrrrrrrr$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)r?rC) r,rrrrrrrZ text_outputsrZ text_featuresr-r-r.get_text_features8s) zFlavaModel.get_text_featuresr~rr}rc 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)r@rB) r,r~rr}rrrrrZ image_outputsrZimage_featuresr-r-r.get_image_featuresps& zFlavaModel.get_image_featuresTimage_attention_maskskip_multimodal_encoderc  Csz| dur| n|jj} | stdd} d} d}d}|dur7|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}| s| |||||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). skip_multimodal_encoder (*bool*, *optional*): Skip any calculations for multimodal encoder. Useful if multimodal encoding is not going to be used. 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**). 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)rrrrrrrrr6ru)rr)rrr r!r"r#)rXrrr@rDr?rErwrArr7r7rr|r)r,rr~rrrrrHrIrrrrZ image_statesZimage_mm_projectionrr Z text_statesZtext_mm_projectionr!r"r#rrrZattention_mask_imageZattention_multimodalZmultimodal_inputr-r-r.rsC   zFlavaModel.forward)NNNNNNNr1) NNNNNNNNNTN)r3r4r5rr r_rrr7rrDr8rFrrG LongTensorr r rrrr-r-rlr.r s+ 7  5     rcs<eZdZdedeffdd ZdejdejfddZZS) FlavaImageCodebookResPathin_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)r^r_rr ReLUr Sequentialpath)r,rLrMkwargsZhid_sizerSrlr-r.r_6s     z"FlavaImageCodebookResPath.__init__rr$cC ||SrU)rSr,rr-r-r.rF z!FlavaImageCodebookResPath.forward r3r4r5rr_r7rrrr-r-rlr.rK5srKcs@eZdZdededeffdd Zdejdejfdd ZZS) FlavaImageCodebookBlockrLrM num_layersc sPtd|d|_||krtj||ddd|_nt|_t|||_dS)NrrrrrO) r^r_ post_gainr rid_pathZIdentityrKres_path)r,rLrMrZrTrlr-r.r_Ks  z FlavaImageCodebookBlock.__init__rr$cCs|||j||SrU)r\r[r]rVr-r-r.rWszFlavaImageCodebookBlock.forwardrXr-r-rlr.rYJs rYc sJeZdZd dededededef fdd Zd ejd ejfd d ZZ S)FlavaImageCodebookLayerGroupT num_blocksrZrLrMuse_poolcstt}t|D]!}|dkr t||||d|d<q t||||d|d<q |r8tjdd|d<t||_dS)NrZblock_rrr)rpool) r^r_rrrYr Z MaxPool2drRgroup)r,r_rZrLrMr`blocksrrlr-r.r_\s  z%FlavaImageCodebookLayerGroup.__init__rr$cCrUrU)rbrVr-r-r.rjrWz$FlavaImageCodebookLayerGroup.forwardr0) r3r4r5rrDr_r7rrrr-r-rlr.r^[s$r^a" 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. )Z custom_introcsxeZdZdZeZdZdZdedeffdd Z de j de j fd d Z de j de j fd d Z de jde j fd dZZS)FlavaImageCodebookr~FrXrTc 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|_qdSdS)NZrelurrrOconvr inputZgroup_1rrZgroup_2rNZgroup_3F)r`Zgroup_4r)r^r_rXZ num_groupsinput_channelsZnum_blocks_per_grouprcrrr rQrr^rRrcrfreeze parametersZ requires_grad)r,rXrTrZZ output_blocksrcparamrlr-r.r_|sB      zFlavaImageCodebook.__init__r$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)format_CHECKPOINT_FOR_CODEBOOK_DOCrcr7Zargmaxr,r~Zz_logitsr-r-r.get_codebook_indicess z'FlavaImageCodebook.get_codebook_indicescCs||}tjdd|S)Nrru)rcr ZSoftmaxrpr-r-r.get_codebook_probss z%FlavaImageCodebook.get_codebook_probscCs`dtt|jdkrtd|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("{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) ``` rNz input shape z is not 4drz input has z channels but model built for )rnrorrwrrjrc)r,r~r-r-r.rs zFlavaImageCodebook.forward)r3r4r5r rr r2rrr_r7rrqrrr8rrr-r-rlr.rdos,rdc$eZdZfddZddZZS)FlavaPredictionHeadTransformcsVtt|j|j|_t|jtrt |j|_ n|j|_ tj |j|j d|_ dS)Nr) r^r_r rrcrrrrr transform_act_fnrrrrlr-r.r_s  z%FlavaPredictionHeadTransform.__init__cCs"||}||}||}|SrU)rrurrr-r-r.rs   z$FlavaPredictionHeadTransform.forwardr3r4r5r_rrr-r-rlr.rts  rtcs.eZdZdfdd ZddZddZZS) rNcsbt||_t||_tj|j|jdd|_ t t |j|_ |dur*||j _|j |j _ dS)NFr)r^r_rXrt transformr rrcrdecoderrar7rbrr)r,rXrrlr-r.r_s  z"FlavaMaskedPredictionHead.__init__cCs|j|j_dSrU)rrxr+r-r-r. _tie_weights sz&FlavaMaskedPredictionHead._tie_weightscCrrU)rwrxrVr-r-r.r   z!FlavaMaskedPredictionHead.forwardrU)r3r4r5r_ryrrr-r-rlr.rs rcrs) FlavaITMHeadcs.t||_t||_t|jd|_dS)Nrr) r^r_rXrrr rrcseq_relationshiprrlr-r.r_s  zFlavaITMHead.__init__cCrrU)rr|rVr-r-r.rrzzFlavaITMHead.forwardrvr-r-rlr.r{s r{crs)FlavaGlobalContrastiveHeadcst||_|j|_dSrU)r^r_rXglobal_backprop_contrastiverrlr-r.r_ s  z#FlavaGlobalContrastiveHead.__init__c s2t|}tjrtjs!tjdjd}g}g}nQd}tj} |j r?tjj j }tjj j }n$fddt | D}fddt | D}tj |tj ||tjtj|jd}t|}t|}t|dd|} t|dd|} | | |fS)Nrr6cg|]}tqSr-r7Z zeros_liker)r r-r.r5z6FlavaGlobalContrastiveHead.forward..crr-rr)rr-r.r6rr)r7exp distributedZ is_availableZis_initializedrrsrZget_world_sizer~r rzZ all_gatherrZget_rankr|rr) r,rr r Z temperaturelabelsZimage_embeddings_allZtext_embeddings_allZlocal_batch_sizeZ world_sizelogits_per_imagelogits_per_textr-)rr r.r%s,      z"FlavaGlobalContrastiveHead.forwardrvr-r-rlr.r}s r}zk The FLAVA model for pretraining which outputs losses, embeddings, logits and transformer outputs. c(s&eZdZgdZddedeejffdd Zde j fdd Z e                d d ee j d ee j d ee jdee jdee j dee j dee j dee j dee j deedee j dee j dee j deededeedeedeee j eff$ddZZS)!FlavaForPreTraining)zmmm_text_head.decoder.biaszmmm_image_head.decoder.biaszmlm_head.decoder.biaszmim_head.decoder.biasNrXimage_codebookcst|t||_||_|jdur|jrt|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_rrrZ init_codebookrdZimage_codebook_configrr=mim_headr:mlm_headr{itm_headmmm_image_head mmm_text_headr}global_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,rXrrlr-r.r_Us,            zFlavaForPreTraining.__init__rcCs"|dkr||dd}|S)Nrrrrq)rvr{rsrVr-r-r. _resize_to_2dxs z!FlavaForPreTraining._resize_to_2dTrinput_ids_maskedr~codebook_pixel_valuesrrrrrHr mlm_labels mim_labels itm_labelsrrr return_lossr$c6Cs|dur|n|jj}|dur|n|jj}| dur| n|j} |dur,|dur,td|}|j|||||| | ||dd }|j||||| |||dd }d}|j}|j}|j}|j}|j }d}}}}}}} d}!}"}#}$d}%}&}'|dus~|dur| dur|r|j durt d|durt d|j |} |jdkr |dur |dur |}(| dur|| } ||}|j| |d<|(dd| d  dddf}(| |j})| |)}*|(|)ddf}(||(}!|rtj|!d |j|*d }||j9}n||(}!|jdkrj|durj|durj|}+| dure|| } |+dd| d  dddf}+| |j})| |)},|+|)ddf}+||+}"|rdtj|"d |j|,d }||j9}n||+}"|jdkr|dur||}%| dur| d}-t|- |-|-!dg}|rtj|%| } | |j9} |dur||}| dur| |} | dur| |} ||}|dur-|j"dkr-|}(|d d }.|(ddd d |.ddf}(| dur(|| } ||}|j| |d<| |j})| |)}*|(|)ddf}(|#|(}$|r'tj|$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|r tj|&|1}2tj|'|1}3|2|3d }||j&9}t0||| |||d }4|r&|41s&t2dd|43D}|s||j4dur5|j45nd||j6durB|j65nd|j |j7durP|j75nd||j4dur]|j45nd||j6durj|j65nd||j7durw|j75nd|!|"|%|&|&|$|#f}5|r|41s||4f|5}5t8dd|5DSt9d&id|d|4d|d|j4d|d|j6d|j d|j7d|d|j4d|d|j6d|d|j7d|!d |"d!|%d"|&d#|'d$|$d%|#S)'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) 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. 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. 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~rrrrHrIrrr) rr~rrrHrrrrz`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)r;r<r=r>r?r@css |] }|dur |ndVqdSrr-)r)rFr-r-r.r/sz.FlavaForPreTraining.forward..css|] }|dur|VqdSrUr-)r)rr-r-r.r/rrFrGrrr r!r"r#rHrIrJrKrLrMrNrOrPrQrRrSrTr-):rXr+rrloggerwarningrrr r"r RuntimeErrorrrqrrrnersrr rzZ cross_entropyr{rrrrrrr7whereanynewrrrrrrC normalizerBr rZclamp_LOGIT_SCALE_CLAMP_MINLOGIT_SCALE_CLAMP_MAXrr:rBsumrArr'r!r#r1rE)6r,rrr~rrrrrrHrrrrrrrrZ flava_outputZflava_masked_outputZpos_maskrr rHrJrLZ total_lossZmim_lossZmlm_lossZ mmm_text_lossZmmm_image_lossZgc_lossZitm_lossrNrOrTrSrPrrZsequence_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.r}s_         "       "                  "                   zFlavaForPreTraining.forwardrU)NNNNNNNNNNNNNNTNN)r3r4r5Z_tied_weights_keysrrr r3r_r7rrrrJr8rDr r rErrr-r-rlr.rGsv#     r)rrdrrrrr4)Jr6rrr dataclassesrtypingrrrrrr r r7Ztorch.utils.checkpointr Z activationsr Zmodeling_outputsrrZmodeling_utilsrrrutilsrrrrZconfiguration_flavarrrrrZ get_loggerr3rrorrrrr:rEr3rWrerrrrrrrrrrrr4rrrKrYr^rdrtrr{r}r__all__r-r-r-r.s  $   %$ed$9C*.3"`p_-u ( #