U ,-e?@sdZddlZddlZddlmZddlmZddlm Z m Z m Z ddl Z ddl Z ddl mZddlmZmZmZdd lmZdd lmZmZdd lmZdd lmZmZdd lmZmZm Z m!Z!m"Z"ddl#m$Z$e!%e&Z'dZ(dZ)dgZ*eGdddeZ+eGdddeZ,eGdddeZ-dLddZ.dMddZ/dNdd Z0Gd!d"d"ej1Z2Gd#d$d$ej1Z3Gd%d&d&ej1Z4Gd'd(d(ej1Z5Gd)d*d*ej1Z6Gd+d,d,ej1Z7Gd-d.d.ej1Z8Gd/d0d0ej1Z9Gd1d2d2ej1Z:Gd3d4d4ej1Z;Gd5d6d6ej1Zd:Z?ed;e>Gdd?d?ej1ZAed@e>GdAdBdBe=ZBGdCdDdDej1ZCGdEdFdFej1ZDGdGdHdHej1ZEedIe>GdJdKdKe=ZFdS)Oz PyTorch TVLT model.N)deepcopy) dataclass)OptionalTupleUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)BaseModelOutputSequenceClassifierOutput)PreTrainedModel) find_pruneable_heads_and_indicesprune_linear_layer) ModelOutputadd_start_docstrings%add_start_docstrings_to_model_forwardloggingreplace_return_docstrings) TvltConfigrzZinengTang/tvlt-basec@seZdZUdZdZejed<dZejed<dZ ejed<dZ ej ed<dZ ej ed<dZ ej ed<dZej ed <dZeeejed <dZeeejed <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__rtorch FloatTensor__annotations__rrr LongTensorrrr r!rrr"r+r+g/var/www/html/Darija-Ai-Train/env/lib/python3.8/site-packages/transformers/models/tvlt/modeling_tvlt.pyr6s rc@sLeZdZUdZdZejed<dZe e ejed<dZ e e ejed<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. Nlogitsr!r") r#r$r%r&r.r'r(r)r!rrr"r+r+r+r,r-_s r-c@szeZdZUdZdZeejed<dZ ejed<dZ ejed<dZ ejed<dZ ee ejed<dZee ejed<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_logitsr!r")r#r$r%r&r0rr'r(r)r1r2r3r!rr"r+r+r+r,r/vs r/?cCs>|jdd\}}tj||f|jd}t|d|}||fS)!Generate noise for audio masking.Ndevicer)shaper'randr8int) pixel_values pixel_mask mask_ratio batch_sizeseq_lennoiselen_keepr+r+r,generate_pixel_mask_noisesrC patch-levelc Cs|jdd\}}|dkrN||}tj|||jdddd|||}n|dkrhtj|||jd}t|d|} || fS)r5Nr6z frame-levelr7rrD)r9r'r:r8 unsqueezerepeatviewr;) audio_values audio_maskr> mask_typefreq_lenr?r@num_time_patchesrArBr+r+r,generate_audio_mask_noises$rOc Cs|j\}}}tj|dd}tj|dd}|ddd|f} tj|d| ddd|d} tj||g|jd} d| ddd|f<tj| d|d} |dk r| |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 rdimNrFrQindexr7r)r9r'ZargsortgatherrGrHZonesr8) sequencerArBattention_masksr?r@Z hidden_dimZ ids_shuffle ids_restoreZids_keepZsequence_maskedZ label_masksr+r+r,random_maskings  rXcs*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_vconfigselfrh __class__r+r,r]s   zTvltPixelEmbeddings.__init__Nc Csh|j\}}}}}||}||jd|d7}|tj|jddd|f|jdd7}||j7}||fSNrrP) r9r_rgrHr'repeat_interleaverfr`rd) rjr<rVr?re num_channelsheightwidth embeddingsr+r+r,forwards  ( zTvltPixelEmbeddings.forward)Nr#r$r%r&r]rs __classcell__r+r+rkr,rYs rYcs*eZdZdZfddZdddZZS)TvltAudioEmbeddingsrZcstt||_|jj|_ttdd|j |_ |j |j d|_ ttd|j|j |j |_ttd|j |j |_|j |j d|_ ||_dSr[)r\r]TvltAudioPatchEmbeddingsr_ num_patchesrrar'rbrc type_embed_afrequency_lengthaudio_patch_sizenum_freq_patches pos_embed_a freq_embedrhrirkr+r,r]s    zTvltAudioEmbeddings.__init__NcCsh||}|d|j}||jd|d7}|tj|jddd|f|jdd7}||j7}||fSrm) r_sizer|r~rHr'rnr}ry)rjrJrVrrrNr+r+r,rss  ( zTvltAudioEmbeddings.forward)Nrtr+r+rkr,rvs rvcs4eZdZdZfddZejejdddZZS)r^z 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 r8|n||f}t|tjj rR|n||f}|d|d|d|d}||_||_ ||_ ||_ ||_t j||||d|_dSNrr)Z kernel_sizeZstride)r\r] image_sizeimage_patch_sizenum_image_channelsrc isinstance collectionsabcIterable patch_sizeror`rConv2d projection)rjrhrrrorcr`rkr+r,r]s  z!TvltPixelPatchEmbeddings.__init__)r<returnc Cs|j\}}}}}||jkr"td||jdks>||jdkrltd|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 ().r6) r9ro ValueErrorrreshaperflatten transposer`rc)rjr<r?rerorprqrrr+r+r,rss (z TvltPixelPatchEmbeddings.forward r#r$r%r&r]r'Tensorrsrur+r+rkr,r^s r^cs4eZdZdZfddZejejdddZZS)rwz 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 rH|n||f}|d|d|d|d}|d|d|d|df} ||_ ||_ ||_ ||_| |_tj||||d|_dSr)r\r]spectrogram_lengthrzr{num_audio_channelsrcrrrrspectrogram_sizerrorx patch_shaperrr) rjrhrrzrrorcrrxrrkr+r,r]8s     z!TvltAudioPatchEmbeddings.__init__)rJrc Cs|j\}}}}||jkr td||jdks<||jdkrjtd|d|d|jdd|jdd ||ddd}|S) NrrrzInput audio size (rrrr6)r9rorrrrr)rjrJr?rorprqrrr+r+r,rsMs (z TvltAudioPatchEmbeddings.forwardrr+r+rkr,rw1s rwcs.eZdZfddZddZd ddZZS) TvltSelfAttentioncst|j|jdkr@t|ds@td|jfd|jd|j|_t|j|j|_|j|j|_t j |j|j|j d|_ t j |j|j|j d|_ t j |j|j|j d|_t |j|_dS)NrZembedding_sizezThe hidden size z4 is not a multiple of the number of attention heads .bias)r\r]rcnum_attention_headshasattrrr;attention_head_size all_head_sizerLinearZqkv_biasquerykeyvalueDropoutZattention_probs_dropout_probdropoutrirkr+r,r]_s zTvltSelfAttention.__init__cCs6|dd|j|jf}|j|}|ddddS)NrFrr6rr )rrrrIpermute)rjxZ new_x_shaper+r+r,transpose_for_scoresqs z&TvltSelfAttention.transpose_for_scoresNFcCs||}|||}|||}||}t||dd} | t|j } |dk rh| |} t j dd| } | | } |dk r| |} t| |} | dddd} | dd|jf} | j| } |r| | fn| f} | S)NrFrPrr6rr )rrrrr'matmulrmathsqrtrrZSoftmaxrr contiguousrrrI)rjr!attention_mask head_maskoutput_attentionsZmixed_query_layerZ key_layerZ value_layerZ query_layerZattention_scoresZattention_probsZ context_layerZnew_context_layer_shapeoutputsr+r+r,rsvs$     zTvltSelfAttention.forward)NNF)r#r$r%r]rrsrur+r+rkr,r^s rcs@eZdZdZeddfdd ZejejejdddZZ 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. Nrhrcs.tt|j|j|_t|j|_dSN) r\r]rrrcdenserhidden_dropout_probrrirkr+r,r]s zTvltSelfOutput.__init__r! input_tensorrcCs||}||}|Srrrrjr!rr+r+r,rss  zTvltSelfOutput.forward) r#r$r%r&rr]r'rrsrur+r+rkr,rsrcs.eZdZfddZddZd ddZZS) TvltAttentioncs*tt||_t||_t|_dSr)r\r]r attentionroutputset pruned_headsrirkr+r,r]s   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)NrrrP)lenrrrrrrrrrrrrunion)rjheadsrSr+r+r, prune_headss zTvltAttention.prune_headsNFcCs6|||||}||d|}|f|dd}|S)Nrr)rr)rjr!rrrZ self_outputsattention_outputrr+r+r,rsszTvltAttention.forward)NNF)r#r$r%r]rrsrur+r+rkr,rs rcs8eZdZeddfdd ZejejdddZZS)TvltIntermediateNrcsBtt|j|j|_t|jt r6t |j|_ n|j|_ dSr) r\r]rrrcintermediate_sizerrZ hidden_actstrr intermediate_act_fnrirkr+r,r]s   zTvltIntermediate.__init__)r!rcCs||}||}|Sr)rrrjr!r+r+r,rss  zTvltIntermediate.forward r#r$r%rr]r'rrsrur+r+rkr,rsrcs<eZdZeddfdd ZejejejdddZZS) TvltOutputNrcs.tt|j|j|_t|j|_ dSr) r\r]rrrrcrrrrrirkr+r,r]s zTvltOutput.__init__rcCs ||}||}||}|Srrrr+r+r,rss  zTvltOutput.forwardrr+r+rkr,rsrcs*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|_dSNrZeps)r\r]Zchunk_size_feed_forwardZ seq_len_dimrrr intermediaterrr LayerNormrclayer_norm_epslayernorm_beforelayernorm_afterrirkr+r,r]s    zTvltLayer.__init__NFc Csj|j|||||d}|d}|dd}|||j}||}||}|||}|f|}|S)Nrrr)rrtor8rrr) rjr!rrrZself_attention_outputsrrZ layer_outputr+r+r,rss     zTvltLayer.forward)NNFrtr+r+rkr,rs rcs&eZdZfddZdddZZS) TvltEncodercs:t|_tfddtjD|_d|_dS)Ncsg|] }tqSr+r.0_rhr+r, sz(TvltEncoder.__init__..F) r\r]rhr ModuleListrangenum_hidden_layerslayergradient_checkpointingrirkrr,r]s  zTvltEncoder.__init__NFTcs|rdnd}rdnd}t|jD]\} } |r8||f}|dk rH|| nd} |jr~|jr~fdd} tjj| | ||| } n| ||| } | d}r"|| df}q"|r||f}|stdd|||fDSt|||dS) Nr+csfdd}|S)Ncs|fSrr+inputsmodulerr+r,custom_forward5szJTvltEncoder.forward..create_custom_forward..custom_forwardr+rrrrr,create_custom_forward4sz2TvltEncoder.forward..create_custom_forwardrrcss|]}|dk r|VqdSrr+rvr+r+r, Lsz&TvltEncoder.forward..)rr!r") enumeraterrtrainingr'utils checkpointtupler )rjr!rrroutput_hidden_states return_dictall_hidden_statesall_self_attentionsi layer_moduleZlayer_head_maskr layer_outputsr+rr,rs s6      zTvltEncoder.forward)NNFFTr#r$r%r]rsrur+r+rkr,rs rc@s2eZdZdZeZdZdZdZddZ d dd Z d S) TvltPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. tvltr<TcCsjt|tjtjfr@|jjjd|jjd|j dk rf|j j n&t|tj rf|j j |jj ddS)zInitialize the weights)meanZstdNg?) rrrrweightdataZnormal_rhZinitializer_rangerZzero_rZfill_)rjrr+r+r, _init_weights_s   z!TvltPreTrainedModel._init_weightsFcCst|tr||_dSr)rrr)rjrrr+r+r,_set_gradient_checkpointingks z/TvltPreTrainedModel._set_gradient_checkpointingN)F) r#r$r%r&r config_classZbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingrrr+r+r+r,rTs 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 dd e j e j e e j e e j eee ee ee eeee j e fd d d ZZS) TvltModelcsvt|||_t||_t||_t||_t t dd|j |_|jrVd|_nt j|j |jd|_|dSr)r\r]rhrYpixel_embeddingsrvaudio_embeddingsrencoderrrar'rbrc cls_embeddingZuse_mean_pooling layernormrr post_initrirkr+r,r]s    zTvltModel.__init__cCs|jj|jjfSr)rr_r)rjr+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)itemsrrrr)rjZheads_to_prunerrr+r+r, _prune_headsszTvltModel._prune_heads output_typerNF) r<rJr=rK mask_pixel mask_audiorrrrc  Cs|dk r |n|jj}|dk r |n|jj}| dk r4| n|jj} |||\} }|||\} }d} d} |rt| ||jjd\}}t| |||d\} }} } d}d}|r|jj |jj d}t | ||jj |jj |d\}}t| |||d\} }}}|d}t|j|dd| | gd}| d}d}|dk rX|dk rXt|ddddf||gd}|}d}|dk rz|||}|j||||| d}|d}|jdk r||}|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>)rVr)rKr>rLrMr)rrrr) rrrrrrr r!r")rhrruse_return_dictrrrCZpixel_mask_ratiorXrzr{rOZaudio_mask_ratioZaudio_mask_typerr'catr rHZget_extended_attention_maskrr rr!r")rjr<rJr=rKrrrrrZpixel_embedding_outputZaudio_embedding_outputrrZpixel_mask_noiseZpixel_len_keeprr r|Zaudio_mask_noiseZaudio_len_keepr?Zembedding_outputZmasked_pixel_lenrZ input_shapeZextended_attention_maskZencoder_outputssequence_outputpixel_sequence_outputaudio_sequence_outputr+r+r,rss%    "      zTvltModel.forward)NNFFNNN)r#r$r%r]r rrTVLT_INPUTS_DOCSTRINGrr_CONFIG_FOR_DOCr'r(rboolrrrsrur+r+rkr,rs0  rcs&eZdZfddZdddZZS) TvltDecodercsvtt||j_|j_|j_|j _ t fddt |jD|_t j|j|jd|_d|_||_dS)Ncsg|] }tqSr+rrZdecoder_configr+r,rYsz(TvltDecoder.__init__..rF)r\r]rdecoder_hidden_sizercZdecoder_num_hidden_layersrZdecoder_num_attention_headsrZdecoder_intermediate_sizerrrrdecoder_layersrrr rrhrirkrr,r]Ps zTvltDecoder.__init__FTc s|rdnd}rdnd}t|jD]l\}}|r8||f}|jrh|jrhfdd} tjj| ||d} n ||d} | d}r"|| df}q"|r||f}||} |stdd| ||fDSt | ||d S) Nr+csfdd}|S)Ncs|fSrr+rrr+r,rrszJTvltDecoder.forward..create_custom_forward..custom_forwardr+rrrr,rqsz2TvltDecoder.forward..create_custom_forwardrrrcss|]}|dk r|VqdSrr+rr+r+r,rsz&TvltDecoder.forward..)r.r!r") rrrrr'rrr rr-) rjr!rrrrrrrrrr.r+rr,rsas,        zTvltDecoder.forward)FFTrr+r+rkr,rOs  rzTThe TVLT Model transformer with the decoder on top for self-supervised pre-training.cseZdZfddZddZddZddZd d Zd d Ze e e e e d dejejeejeejeejeejeejeeeeeeeeeje fd ddZZS)TvltForPreTrainingc st|||_|j|_|j|_|js6|js6tdt||_|jrPt||_ |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 trueTrrrr6)+r\r]rh task_matchingtask_maerrrTvltMatchingHead matching_headrrrcrencoder_to_decoderrar'rbpixel_mask_tokenaudio_mask_tokenrdecoderrerr`decoder_pixel_pos_embeddecoder_temporal_embeddecoder_pixel_type_embedrrxrzr{decoder_audio_pos_embeddecoder_freq_embeddecoder_audio_type_embedrr TvltMAEHeadpixel_mae_headraudio_mae_headr|r ) rjrhrrer`Znum_audio_patchesr|Zpixel_mae_output_dimZaudio_mae_output_dimrkr+r,r]sL         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] r rrr9zntchpwq->nthwpqc)r9rrr'einsum) rjr<r?rerorprqnum_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] rrr2znchpwq->nhwpqc)r9r{rr'r3) rjrJr?rorprqr4r5patchified_audio_valuesr+r+r,patchify_audios( z!TvltForPreTraining.patchify_audiocCs:||}||d}|jdd}|||}|SNr6rFrP)r7rsum)rjr<Zpixel_predictionsmaskr6r0r+r+r,pixel_mae_losss    z!TvltForPreTraining.pixel_mae_losscCs:||}||d}|jdd}|||}|Sr:)r9rr;)rjrJZaudio_predictionsr<r8r0r+r+r,audio_mae_losss    z!TvltForPreTraining.audio_mae_lossc CsZ|j\}}}|||jd|d}tj||gdd}tj|d|ddd|d}|S)NrrPrFrR)r9rHr'rrTrG) rjZ mask_tokenrUrWr?Z seq_lengthrQZ mask_tokensZpadded_sequencer+r+r,concatenate_masks z#TvltForPreTraining.concatenate_maskrN) r<rJr=rKlabelspixel_values_mixedpixel_mask_mixedrrrrc Cs| dk r | n|jj} d} |jr|dkr.td|dkr>td|j|||||| | d} | d} || }t}||d|d}| |7} d}d}|jr\|j r\|j||||dd|| | d } | r| j n| d }| r| j n| d }| r| 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} | s|||f| dd}|dk r| 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=rKrrrrrFT)r=rKrrrrrrr6r r1rP)r0r1r2r3r!r"))rhrr rrr#rrIr!rrrrrrr r$rr?r%r(rHr'rnr)r`r*r'r/r.r&r|r,r+r-r0r=r>r/r!r") rjr<rJr=rKr@rArBrrrZ total_lossrrr1loss_fctr0r2r3rrrrrr Zpixel_decoder_inputZaudio_decoder_inputreZpixel_decoder_outputsrNZaudio_decoder_outputsrr+r+r,rss1          zTvltForPreTraining.forward)NNNNNNNN)r#r$r%r]r7r9r=r>r?rrrr/rr'r(rr*rrrrsrur+r+rkr,rs: 6  rcs$eZdZfddZddZZS) TvltPoolercs*tt|j|j|_t|_dSr)r\r]rrrcrZTanh activationrirkr+r,r]s zTvltPooler.__init__cCs(|dddf}||}||}|S)Nr)rrI)rjr!Zfirst_token_tensorZ pooled_outputr+r+r,rss  zTvltPooler.forwardrr+r+rkr,rHs rHcs$eZdZfddZddZZS)r"cs(tt||_t|jd|_dSr[)r\r]rHpoolerrrrcfcrirkr+r,r]s  zTvltMatchingHead.__init__cCs|||}|Sr)rKrJrr+r+r,rsszTvltMatchingHead.forwardrr+r+rkr,r"s r"cs&eZdZdfdd ZddZZS)r.Ncs$t||_t|j||_dSr)r\r]rhrrrr')rjrhZ output_dimrkr+r,r]s zTvltMAEHead.__init__cCs||}|Sr)r'rr+r+r,rss zTvltMAEHead.forward)Nrr+r+rkr,r.sr.z 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. cseZdZfddZeeeeedde j e j e e j e e j e e e e e e e e j eee j efd ddZZS) TvltForAudioVisualClassificationc spt|t||_tt|j|jdtj|jd|j dt t|jd|j |_ ||_ |dS)Nr6r)r\r]rrrZ SequentialrrcrrZGELUZ num_labels classifierrhr rirkr+r,r]s  z)TvltForAudioVisualClassification.__init__rN) r<rJr=rKrrrr@rc  Cs|dk r |n|jj}|j|||||||d} | ddddf} || } d} |dk r|jjdkrtt} | | |} n|jjdkrt} | | |} |s| f| dd}| dk r| 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 ```NrCrZ regressionZclassificationr1)r0r.r!r") rhrrrMZ loss_typer r rr!r")rjr<rJr=rKrrrr@rrr.r0rGrr+r+r,rss:$      z(TvltForAudioVisualClassification.forward)NNNNNN)r#r$r%r]rrrrrr'r(rrr*rrrsrur+r+rkr,rLs(  rL)Nr4)Nr4rDrE)N)Gr&collections.abcrrcopyr dataclassesrtypingrrrr'Ztorch.utils.checkpointrZtorch.nnrr r Z activationsr Zmodeling_outputsr rZmodeling_utilsrZ pytorch_utilsrrrrrrrrZconfiguration_tvltrZ get_loggerr#loggerrZ_CHECKPOINT_FOR_DOCZ"TVLT_PRETRAINED_MODEL_ARCHIVE_LISTrr-r/rCrOrXModulerYrvr^rwrrrrrrrrZTVLT_START_DOCSTRINGrrrrrHr"r.rLr+r+r+r,s       (!  )-=#'; -$@