o i@sdZddlZddlZddlmZddlmZddlm Z m Z ddl Z ddl m Z ddl mZmZmZdd lmZdd lmZdd lmZmZdd lmZdd lmZmZddlmZmZm Z m!Z!m"Z"ddl#m$Z$e!%e&Z'dZ(dZ)eGdddeZ*eGdddeZ+eGdddeZ,dNddZ-dOddZ.dPd d!Z/Gd"d#d#e j0Z1Gd$d%d%e j0Z2Gd&d'd'e j0Z3Gd(d)d)e j0Z4Gd*d+d+e j0Z5Gd,d-d-e j0Z6Gd.d/d/e j0Z7Gd0d1d1e j0Z8Gd2d3d3e j0Z9Gd4d5d5eZ:Gd6d7d7e j0Z;Gd8d9d9eZedd>e<Z?Gd?d@d@e j0Z@edAe=GdBdCdCe<ZAGdDdEdEe j0ZBGdFdGdGe j0ZCGdHdIdIe j0ZDedJe=GdKdLdLe<ZEgdMZFdS)QzPyTorch TVLT model.N)deepcopy) dataclass)OptionalUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)GradientCheckpointingLayer)BaseModelOutputSequenceClassifierOutput)PreTrainedModel) find_pruneable_heads_and_indicesprune_linear_layer) ModelOutputadd_start_docstrings%add_start_docstrings_to_model_forwardloggingreplace_return_docstrings) TvltConfigrzZinengTang/tvlt-basec@seZdZUdZdZeejed<dZ eejed<dZ eejed<dZ eej ed<dZ eej ed<dZeej ed<dZeej ed <dZeeejd fed <dZeeejd fed <dS) TvltModelOutputa Class for TvltModel's outputs, with potential hidden states and attentions. Args: last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. last_pixel_hidden_state (`torch.FloatTensor` of shape `(batch_size, pixel_sequence_length, hidden_size)`): Pixel sequence of hidden-states at the output of the last layer of the model. last_audio_hidden_state (`torch.FloatTensor` of shape `(batch_size, audio_sequence_length, hidden_size)`): Audio sequence of hidden-states at the output of the last layer of the model. pixel_label_masks (`torch.FloatTensor` of shape `(batch_size, pixel_patch_length)`): Tensor indicating which pixel patches are masked (1) and which are not (0). audio_label_masks (`torch.FloatTensor` of shape `(batch_size, audio_patch_length)`): Tensor indicating which audio patches are masked (1) and which are not (0). pixel_ids_restore (`torch.LongTensor` of shape `(batch_size, pixel_patch_length)`): Tensor containing the ids permutation of pixel masking. audio_ids_restore (`torch.LongTensor` of shape `(batch_size, audio_patch_length)`): Tensor containing the ids permutation of audio masking. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings and one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. Nlast_hidden_statelast_pixel_hidden_statelast_audio_hidden_statepixel_label_masksaudio_label_maskspixel_ids_restoreaudio_ids_restore. hidden_states attentions)__name__ __module__ __qualname____doc__rrtorch FloatTensor__annotations__rrr LongTensorrrr r!tupler"r,r,e/home/ubuntu/.local/lib/python3.10/site-packages/transformers/models/deprecated/tvlt/modeling_tvlt.pyr0s rc@sXeZdZUdZdZeejed<dZ ee ejdfed<dZ ee ejdfed<dS)TvltDecoderOutputaM Class for TvltDecoder's outputs, with potential hidden states and attentions. Args: logits (`torch.FloatTensor` of shape `(batch_size, patch_size ** 2 * num_channels)`): Pixel reconstruction logits. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings and one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. Nlogits.r!r") r#r$r%r&r/rr'r(r)r!r+r"r,r,r,r-r.Ys r.c@seZdZUdZdZeejed<dZ eejed<dZ eejed<dZ eejed<dZ ee ejdfed<dZee ejdfed <dS) TvltForPreTrainingOutputa  Class for TvltForPreTraining's outputs, with potential hidden states and attentions. Args: loss (`torch.FloatTensor` of shape `(1,)`): Pixel reconstruction loss. matching_logits (`torch.FloatTensor` of shape `(batch_size, 1)`): Matching objective logits. pixel_logits (`torch.FloatTensor` of shape `(batch_size, pixel_patch_length, image_patch_size ** 3 * pixel_num_channels)`): Pixel reconstruction logits. audio_logits (`torch.FloatTensor` of shape `(batch_size, audio_patch_length, image_patch_size[0] * image_patch_size[1])`): Audio reconstruction logits. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings and one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. Nlossmatching_logits pixel_logits audio_logits.r!r")r#r$r%r&r1rr'r(r)r2r3r4r!r+r"r,r,r,r-r0ps r0?cCs>|jdd\}}tj||f|jd}t|d|}||fS)!Generate noise for audio masking.Ndevicer)shaper'randr9int) pixel_values pixel_mask mask_ratio batch_sizeseq_lennoiselen_keepr,r,r-generate_pixel_mask_noisesrD patch-levelc Cs|jdd\}}|dkr'||}tj|||jdddd|||}n |dkr4tj|||jd}t|d|} || fS)r6Nr7z frame-levelr8rrE)r:r'r;r9 unsqueezerepeatviewr<) audio_values audio_maskr? mask_typefreq_lenr@rAnum_time_patchesrBrCr,r,r-generate_audio_mask_noises rPc Cs|j\}}}tj|dd}tj|dd}|ddd|f} tj|d| ddd|d} tj||g|jd} d| ddd|f<tj| d|d} |durZ| |9} tj|d| d}| || |fS)z Perform random masking by per-sample shuffling on frame-level. Per-sample shuffling is done by argsort random noise. sequence: [batch_size, seq_len, hidden_dim], sequence rdimNrGrRindexr8r)r:r'argsortgatherrHrIonesr9) sequencerBrCattention_masksr@rA hidden_dim ids_shuffle ids_restoreids_keepsequence_masked label_masksr,r,r-random_maskings   r`c*eZdZdZfddZdddZZS)TvltPixelEmbeddings,Construct the patch and position embeddings.csttt||_|jj|_ttdd|j |_ ttd|j |j |_ ttd|j|j |_ ||_dSNr)super__init__TvltPixelPatchEmbeddingspatch_embeddingsnum_patches_per_imager Parameterr'zeros hidden_size type_embed_v num_framestemporal_embed pos_embed_vconfigselfrq __class__r,r-rfs    zTvltPixelEmbeddings.__init__Nc Csh|j\}}}}}||}||jd|d7}|tj|jddd|f|jdd7}||j7}||fSNrrQ) r:rhrprIr'repeat_interleaverorirm) rsr=rYr@rn num_channelsheightwidth embeddingsr,r,r-forwards  ( zTvltPixelEmbeddings.forwardNr#r$r%r&rfr| __classcell__r,r,rtr-rbs  rbcra)TvltAudioEmbeddingsrccstt||_|jj|_ttdd|j |_ |j |j d|_ ttd|j|j |j |_ttd|j |j |_|j |j d|_ ||_dSrd)rerfTvltAudioPatchEmbeddingsrh num_patchesrrjr'rkrl type_embed_afrequency_lengthaudio_patch_sizenum_freq_patches pos_embed_a freq_embedrqrrrtr,r-rfs     zTvltAudioEmbeddings.__init__NcCsh||}|d|j}||jd|d7}|tj|jddd|f|jdd7}||j7}||fSrv) rhsizerrrIr'rwrr)rsrKrYr{rOr,r,r-r|s ( zTvltAudioEmbeddings.forwardr}r~r,r,rtr-rs rc6eZdZdZfddZdejdejfddZZS)rgz This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a Transformer. cst|j|j}}|j|j}}t|tjj r|n||f}t|tjj r)|n||f}|d|d|d|d}||_||_ ||_ ||_ ||_t j||||d|_dSNrr) kernel_sizestride)rerf image_sizeimage_patch_sizenum_image_channelsrl isinstance collectionsabcIterable patch_sizerxrirConv2d projection)rsrqrrrxrlrirtr,r-rf s  z!TvltPixelPatchEmbeddings.__init__r=returnc Cs|j\}}}}}||jkrtd||jdks||jdkr6td|d|d|jdd|jdd ||||||}||ddd}||||j|j }|S) NeMake sure that the channel dimension of the pixel values match with the one set in the configuration.rrzInput image size (*) doesn't match model ().r7) r:rx ValueErrorrreshaperflatten transposerirl)rsr=r@rnrxryrzr{r,r,r-r|s (z TvltPixelPatchEmbeddings.forward r#r$r%r&rfr'Tensorr|rr,r,rtr-rgs rgcr)rz This class turns `audio_values` of shape `(batch_size, num_channels, height, width)` into the initial `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a Transformer. c st|j|j|j}}}|j|j}}||f}t|tj j r$|n||f}|d|d|d|d}|d|d|d|df} ||_ ||_ ||_ ||_| |_tj||||d|_dSr)rerfspectrogram_lengthrrnum_audio_channelsrlrrrrspectrogram_sizerrxr patch_shaperrr) rsrqrrrrxrlrrrrtr,r-rf2s    z!TvltAudioPatchEmbeddings.__init__rKrc Cs|j\}}}}||jkrtd||jdks||jdkr5td|d|d|jdd|jdd ||ddd}|S) NrrrzInput audio size (rrrr7)r:rxrrrrr)rsrKr@rxryrzr{r,r,r-r|Gs z TvltAudioPatchEmbeddings.forwardrr,r,rtr-r+s rc.eZdZfddZddZd ddZZS) TvltSelfAttentioncst|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)Nrembedding_sizezThe hidden size z4 is not a multiple of the number of attention heads .bias)rerfrlnum_attention_headshasattrrr<attention_head_size all_head_sizerLinearqkv_biasquerykeyvalueDropoutattention_probs_dropout_probdropoutrrrtr,r-rfXs  zTvltSelfAttention.__init__cCs6|dd|j|jf}|j|}|ddddS)NrGrr7r)rrrrJpermute)rsx new_x_shaper,r,r-transpose_for_scoresjs z&TvltSelfAttention.transpose_for_scoresNFcCs||}|||}|||}||}t||dd} | t|j } |dur4| |} t j dd| } | | } |durI| |} t| |} | dddd} | dd|jf} | j| } |rr| | f} | S| f} | S)NrGrQrr7rr)rrrrr'matmulrmathsqrtrrSoftmaxrr contiguousrrrJ)rsr!attention_mask head_maskoutput_attentionsmixed_query_layer key_layer value_layer query_layerattention_scoresattention_probs context_layernew_context_layer_shapeoutputsr,r,r-r|os(      zTvltSelfAttention.forwardNNF)r#r$r%rfrr|rr,r,rtr-rWs rcsFeZdZdZdeddffdd Zdejdejdejfd d ZZ S) TvltSelfOutputz The residual connection is defined in TvltLayer instead of here (as is the case with other models), due to the layernorm applied before each block. rqrNcs.tt|j|j|_t|j|_dSr}) rerfrrrldenserhidden_dropout_probrrrrtr,r-rf zTvltSelfOutput.__init__r! input_tensorcC||}||}|Sr}rrrsr!rr,r,r-r|  zTvltSelfOutput.forward) r#r$r%r&rrfr'rr|rr,r,rtr-rs$rcr) TvltAttentioncs*tt||_t||_t|_dSr})rerfr attentionroutputset pruned_headsrrrtr,r-rfs    zTvltAttention.__init__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)NrrrQ)lenrrrrrrrrrrrrunion)rsheadsrTr,r,r- prune_headss zTvltAttention.prune_headsNFcCs6|||||}||d|}|f|dd}|S)Nrr)rr)rsr!rrr self_outputsattention_outputrr,r,r-r|szTvltAttention.forwardr)r#r$r%rfrr|rr,r,rtr-rs rcs<eZdZdeddffdd ZdejdejfddZZS) TvltIntermediaterqrNcsDtt|j|j|_t|jt rt |j|_ dS|j|_ dSr}) rerfrrrlintermediate_sizerr hidden_actstrr intermediate_act_fnrrrtr,r-rfs   zTvltIntermediate.__init__r!cCrr})rrrsr!r,r,r-r|rzTvltIntermediate.forward r#r$r%rrfr'rr|rr,r,rtr-rsrcsBeZdZdeddffdd Zdejdejdejfdd ZZS) TvltOutputrqrNcs.tt|j|j|_t|j|_ dSr}) rerfrrrrlrrrrrrrtr,r-rfrzTvltOutput.__init__r!rcCs ||}||}||}|Sr}rrr,r,r-r|s  zTvltOutput.forwardrr,r,rtr-rs$rcs*eZdZdZfddZdddZZS) TvltLayerz?This corresponds to the Block class in the timm implementation.csbt|j|_d|_t||_t||_t||_ t j |j |j d|_t j |j |j d|_dSNreps)rerfchunk_size_feed_forward seq_len_dimrrr intermediaterrr LayerNormrllayer_norm_epslayernorm_beforelayernorm_afterrrrtr,r-rfs    zTvltLayer.__init__NFc Csj|j|||||d}|d}|dd}|||j}||}||}|||}|f|}|S)Nrrr)rrtor9rrr) rsr!rrrself_attention_outputsrr layer_outputr,r,r-r|s     zTvltLayer.forwardrr~r,r,rtr-rs  rcs0eZdZfddZ     dddZZS) TvltEncodercs:t|_tfddtjD|_d|_dS)Ncg|]}tqSr,r.0_rqr,r- z(TvltEncoder.__init__..F) rerfrqr ModuleListrangenum_hidden_layerslayergradient_checkpointingrrrtrr-rf s   zTvltEncoder.__init__NFTc Cs|rdnd}|r dnd}t|jD])\} } |r||f}|dur$|| nd} | ||| |} | d}|r:|| df}q|rB||f}|sPtdd|||fDSt|||dS)Nr,rrcs|] }|dur|VqdSr}r,rvr,r,r- 0z&TvltEncoder.forward..)rr!r") enumeraterr+r ) rsr!rrroutput_hidden_states return_dictall_hidden_statesall_self_attentionsi layer_modulelayer_head_mask layer_outputsr,r,r-r|s(   zTvltEncoder.forward)NNFFTr#r$r%rfr|rr,r,rtr-r s  rc@s.eZdZUdZeed<dZdZdZddZ dS) TvltPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. rqtvltr=TcCstt|tjtjfr#|jjjd|jjd|j dur!|j j dSdSt|tj r8|j j |jj ddSdS)zInitialize the weights)meanstdNg?) rrrrweightdatanormal_rqinitializer_rangerzero_rfill_)rsmoduler,r,r- _init_weightsCs   z!TvltPreTrainedModel._init_weightsN) r#r$r%r&rr)base_model_prefixmain_input_namesupports_gradient_checkpointingr%r,r,r,r-r8s  raF 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 ([`TvltConfig`]): 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 Args: pixel_values (`torch.FloatTensor` of shape `(batch_size, num_frames, num_channels, height, width)`): Pixel values. Pixel values can be obtained using [`TvltProcessor`]. See [`TvltProcessor.__call__`] for details. audio_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Audio values. Audio values can be obtained using [`TvltProcessor`]. See [`TvltProcessor.__call__`] for details. pixel_mask (`torch.FloatTensor` of shape `(batch_size, num_pixel_patches)`): Pixel masks. Pixel masks can be obtained using [`TvltProcessor`]. See [`TvltProcessor.__call__`] for details. audio_mask (`torch.FloatTensor` of shape `(batch_size, num_audio_patches)`): Audio masks. Audio masks can be obtained using [`TvltProcessor`]. See [`TvltProcessor.__call__`] for details. pixel_values_mixed (`torch.FloatTensor` of shape `(batch_size, num_frames, num_channels, height, width)`): Pixel values that mix positive and negative samples in Tvlt vision-audio matching. Pixel values mixed can be obtained using [`TvltProcessor`]. See [`TvltProcessor.__call__`] for details. pixel_mask_mixed (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Pixel masks of pixel_values_mixed. Pixel masks mixed can be obtained using [`TvltProcessor`]. See [`TvltProcessor.__call__`] for details. mask_pixel (`bool`, *optional*): Whether to mask pixel for MAE tasks. Only set to True in TvltForPreTraining. mask_audio (`bool`, *optional*): Whether to mask audio for MAE tasks. Only set to True in TvltForPreTraining. 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. z^The bare TVLT Model transformer outputting raw hidden-states without any specific head on top.cseZdZfddZddZddZeeee e d     dd e j d e j d e e j d e e j dedede ede ede edeee j e ffddZZS) TvltModelcsvt|||_t||_t||_t||_t t dd|j |_|jr+d|_n t j|j |jd|_|dSr)rerfrqrbpixel_embeddingsraudio_embeddingsrencoderrrjr'rkrl cls_embeddinguse_mean_pooling layernormrr post_initrrrtr,r-rfs     zTvltModel.__init__cCs|jj|jjfSr})r*rhr+)rsr,r,r-get_input_embeddingsszTvltModel.get_input_embeddingscCs*|D]\}}|jj|j|qdS)z Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base class PreTrainedModel N)itemsr,rrr)rsheads_to_prunerrr,r,r- _prune_headsszTvltModel._prune_heads output_type config_classNFr=rKr>rL mask_pixel mask_audiorrrrc  Cs|dur|n|jj}|dur|n|jj}| dur| n|jj} |||\} }|||\} }d} d} |rKt| ||jjd\}}t| |||d\} }} } d}d}|rv|jj |jj d}t | ||jj |jj |d\}}t| |||d\} }}}|d}t|j|dd| | gd}| d}d}|dur|durt|ddddf||gd}|}d}|dur|||}|j||||| d}|d}|jdur||}|dddd|f}|ddd|df}| s|||| || |f|ddSt|||| || ||j|jd S) a Returns: Examples: ```python >>> from transformers import TvltProcessor, TvltModel >>> import numpy as np >>> import torch >>> num_frames = 8 >>> images = list(np.random.randn(num_frames, 3, 224, 224)) >>> audio = list(np.random.randn(10000)) >>> processor = TvltProcessor.from_pretrained("ZinengTang/tvlt-base") >>> model = TvltModel.from_pretrained("ZinengTang/tvlt-base") >>> input_dict = processor(images, audio, sampling_rate=44100, return_tensors="pt") >>> outputs = model(**input_dict) >>> loss = outputs.loss ```N)r>r?)rYr)rLr?rMrNr)rrrr) rrrrrrr r!r")rqrruse_return_dictr*r+rDpixel_mask_ratior`rrrPaudio_mask_ratioaudio_mask_typerr'catr-rIget_extended_attention_maskr,r/rr!r")rsr=rKr>rLr8r9rrrpixel_embedding_outputaudio_embedding_outputrrpixel_mask_noisepixel_len_keeprr raudio_mask_noiseaudio_len_keepr@embedding_outputmasked_pixel_lenr input_shapeextended_attention_maskencoder_outputssequence_outputpixel_sequence_outputaudio_sequence_outputr,r,r-r|s%     "     zTvltModel.forward)NNFFNNN)r#r$r%rfr1r4rTVLT_INPUTS_DOCSTRINGrr_CONFIG_FOR_DOCr'r(rboolrr+r|rr,r,rtr-r)sD     r)cs,eZdZfddZ   dddZZS) TvltDecodercsvtt||j_|j_|j_|j _ t fddt |jD|_t j|j|jd|_d|_||_dS)Ncrr,rrdecoder_configr,r-r9rz(TvltDecoder.__init__..rF)rerfrdecoder_hidden_sizerldecoder_num_hidden_layersrdecoder_num_attention_headsrdecoder_intermediate_sizerrrrdecoder_layersrrr/r rqrrrtrRr-rf0s  zTvltDecoder.__init__FTc Cs|rdnd}|r dnd}t|jD]\}}|r||f}|||d} | d}|r/|| df}q|r7||f}||} |sJtdd| ||fDSt| ||dS)Nr,rrrcsr r}r,r r,r,r-r ]rz&TvltDecoder.forward..)r/r!r")rrXr/r+r.) rsr!rrrrrrrrr/r,r,r-r|As      zTvltDecoder.forward)FFTrr,r,rtr-rQ/s  rQzTThe TVLT Model transformer with the decoder on top for self-supervised pre-training.cseZdZfddZddZddZddZd d Zd d Ze e e e e d         ddejdejdeejdeejdeejdeejdeejdeedeedeedeeeje ffddZZS)TvltForPreTrainingc st|||_|j|_|j|_|js|jstdt||_|jr(t||_ |jrt j |j |j dd|_t tdd|j |_t tdd|j |_t||_|j }|j}|jjj}t td|||_t td|j||_t tdd||_|jjj}|j|jd}t td||||_ t td|||_!t tdd||_"|jj#dd|jj$}t%|||_&|jjd|jjd|jj'}t%|||_(||_||_||_)|j#|_#|j|_|*dS)Nz;Must set at least one of matching task and MAE task to trueTrrrr7)+rerfrq task_matchingtask_maerr)rTvltMatchingHead matching_headrrrlrTencoder_to_decoderrjr'rkpixel_mask_tokenaudio_mask_tokenrQdecoderrnr*ridecoder_pixel_pos_embeddecoder_temporal_embeddecoder_pixel_type_embedr+rrrdecoder_audio_pos_embeddecoder_freq_embeddecoder_audio_type_embedrr TvltMAEHeadpixel_mae_headraudio_mae_headrr0) rsrqrTrnrinum_audio_patchesrpixel_mae_output_dimaudio_mae_output_dimrtr,r-rffsL          zTvltForPreTraining.__init__c Cs|j\}}}}}|jd|jd}|jd|jd}|j|||||jd||jdfd} td| } | j|||||jd|jd|fd} | S)zJ pixel_values: [batch_size, num_frames, 3, height, width] rrr rr:zntchpwq->nthwpqc)r:rrr'einsum) rsr=r@rnrxryrznum_patches_heightnum_patches_widthpatchified_pixel_valuesr,r,r-patchify_pixels*  z!TvltForPreTraining.patchify_pixelc Cs|j\}}}}||jd}||jd}|j||||jd||jdfd}td|}|j||||jd|jd|fd}|S)z> audio_values: [batch_size, 1, height, width] rrrnznchpwq->nhwpqc)r:rrr'ro) rsrKr@rxryrzrprqpatchified_audio_valuesr,r,r-patchify_audios( z!TvltForPreTraining.patchify_audiocC:||}||d}|jdd}|||}|SNr7rGrQ)rsrsum)rsr=pixel_predictionsmaskrrr1r,r,r-pixel_mae_loss   z!TvltForPreTraining.pixel_mae_losscCrvrw)rurrx)rsrKaudio_predictionsrzrtr1r,r,r-audio_mae_lossr|z!TvltForPreTraining.audio_mae_lossc CsZ|j\}}}|||jd|d}tj||gdd}tj|d|ddd|d}|S)NrrQrGrS)r:rIr'r>rVrH) rs mask_tokenrXr\r@ seq_lengthrR mask_tokenspadded_sequencer,r,r-concatenate_masks z#TvltForPreTraining.concatenate_maskr5Nr=rKr>rLlabelspixel_values_mixedpixel_mask_mixedrrrrc Cs| dur| n|jj} d} |jrF|durtd|durtd|j|||||| | d} | d} || }t}||d|d}| |7} d}d}|jr+|j r+|j||||dd|| | d } | re| j n| d }| rn| j n| d }| rw| j n| d }| r| j n| d }| r| jn| d}| r| jn| d}||}||}|d }||j||}||jd |d }|tj|jddd|f|jd d}||j}||}||j}||j||}|d |j}||j d |d }|tj|j!ddd|f|jd d}||j"}||}|#|j}|$||||%|||}| |7} | sE|||f| dd}|durC| f|S|St&| |||| j'| j(dS)aF pixel_values_mixed (`torch.FloatTensor` of shape `(batch_size, num_frames, num_channels, height, width)`): Pixel values that mix positive and negative samples in Tvlt vision-audio matching. Audio values can be obtained using [`TvltProcessor`]. See [`TvltProcessor.__call__`] for details. pixel_mask_mixed (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Pixel masks of pixel_values_mixed. Pixel values mixed can be obtained using [`TvltProcessor`]. See [`TvltProcessor.__call__`] for details. labels (`torch.LongTensor` of shape `(batch_size, num_labels)`, *optional*): Labels for computing the vision audio matching loss. Indices should be in `[0, 1]`. num_labels has to be 1. Return: Examples: ```python >>> from transformers import TvltProcessor, TvltForPreTraining >>> import numpy as np >>> import torch >>> num_frames = 8 >>> images = list(np.random.randn(num_frames, 3, 224, 224)) >>> images_mixed = list(np.random.randn(num_frames, 3, 224, 224)) >>> audio = list(np.random.randn(10000)) >>> processor = TvltProcessor.from_pretrained("ZinengTang/tvlt-base") >>> model = TvltForPreTraining.from_pretrained("ZinengTang/tvlt-base") >>> input_dict = processor( ... images, audio, images_mixed, sampling_rate=44100, mask_pixel=True, mask_audio=True, return_tensors="pt" ... ) >>> outputs = model(**input_dict) >>> loss = outputs.loss ```NrzMatching task requires labelsz)Matching task requires pixel_values_mixedr>rLrrrrrGT)r>rLr8r9rrrrr7rr rQ)r1r2r3r4r!r"))rqr:rZrrr]rrJr[trainingrrrrrr r^rrr_rbrIr'rwrcrirdrarir/r`rrfrergrjr{r~r0r!r") rsr=rKr>rLrrrrrr total_lossrrKr2loss_fctr1r3r4rLrMrrrr pixel_decoder_inputaudio_decoder_inputrnpixel_decoder_outputsrOaudio_decoder_outputsrr,r,r-r|s1         zTvltForPreTraining.forward)NNNNNNNN)r#r$r%rfrsrur{r~rrrNrr0rOr'r(rr*rPrr+r|rr,r,rtr-rYasP 6     rYc$eZdZfddZddZZS) TvltPoolercs*tt|j|j|_t|_dSr})rerfrrrlrTanh activationrrrtr,r-rfxs zTvltPooler.__init__cCs(|dddf}||}||}|S)Nr)rr)rsr!first_token_tensor pooled_outputr,r,r-r|}s  zTvltPooler.forwardrr,r,rtr-rw rcr)r\cs(tt||_t|jd|_dSrd)rerfrpoolerrrrlfcrrrtr,r-rfs  zTvltMatchingHead.__init__cCs|||}|Sr})rrrr,r,r-r|szTvltMatchingHead.forwardrr,r,rtr-r\rr\cs&eZdZdfdd ZddZZS)rhNcs$t||_t|j||_dSr})rerfrqrrrTra)rsrq output_dimrtr,r-rfs zTvltMAEHead.__init__cCs||}|Sr})rarr,r,r-r|s zTvltMAEHead.forwardr}rr,r,rtr-rhsrhz Tvlt Model transformer with a classifier head on top (an MLP on top of the final hidden state of the [CLS] token) for audiovisual classification tasks, e.g. CMU-MOSEI Sentiment Analysis and Audio to Video Retrieval. cseZdZfddZeeeeed      dde j de j de e j de e j d e e d e e d e e d e e j d eee j effddZZS) TvltForAudioVisualClassificationc spt|t||_tt|j|jdtj|jd|j dt t|jd|j |_ ||_ |dS)Nr7r)rerfr)rr SequentialrrlrrGELU num_labels classifierrqr0rrrtr,r-rfs   z)TvltForAudioVisualClassification.__init__r5Nr=rKr>rLrrrrrc  Cs|dur|n|jj}|j|||||||d} | ddddf} || } d} |durH|jjdkr:t} | | |} n|jjdkrHt} | | |} |s^| f| dd}| dur\| f|S|St| | | j| j dS)a labels (`torch.LongTensor` of shape `(batch_size, num_labels)`, *optional*): Labels for computing the audiovisual loss. Indices should be in `[0, ..., num_classes-1]` where num_classes refers to the number of classes in audiovisual tasks. Return: Examples: ```python >>> from transformers import TvltProcessor, TvltForAudioVisualClassification >>> import numpy as np >>> import torch >>> num_frames = 8 >>> images = list(np.random.randn(num_frames, 3, 224, 224)) >>> audio = list(np.random.randn(10000)) >>> processor = TvltProcessor.from_pretrained("ZinengTang/tvlt-base") >>> model = TvltForAudioVisualClassification.from_pretrained("ZinengTang/tvlt-base") >>> input_dict = processor(images, audio, sampling_rate=44100, return_tensors="pt") >>> outputs = model(**input_dict) >>> loss = outputs.loss ```Nrr regressionclassificationr )r1r/r!r") rqr:rr loss_typer rrr!r")rsr=rKr>rLrrrrrrKr/r1rrr,r,r-r|s:$     z(TvltForAudioVisualClassification.forward)NNNNNN)r#r$r%rfrrNrrrOr'r(rrPr*rr+r|rr,r,rtr-rs:    r)r)rYrr)Nr5)Nr5rErFr})Gr&collections.abcrrcopyr dataclassesrtypingrrr'rtorch.nnrrr activationsr modeling_layersr modeling_outputsr rmodeling_utilsr pytorch_utilsrrutilsrrrrrconfiguration_tvltr get_loggerr#loggerrO_CHECKPOINT_FOR_DOCrr.r0rDrPr`ModulerbrrgrrrrrrrrrTVLT_START_DOCSTRINGrNr)rQrYrr\rhr__all__r,r,r,r-s        ( ! ),<"&, -$2    Y