o ZhR@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)eGdddeZ*eGdddeZ+eGdddeZ,dMddZ-dNddZ.dOdd Z/Gd!d"d"ej0Z1Gd#d$d$ej0Z2Gd%d&d&ej0Z3Gd'd(d(ej0Z4Gd)d*d*ej0Z5Gd+d,d,ej0Z6Gd-d.d.ej0Z7Gd/d0d0ej0Z8Gd1d2d2ej0Z9Gd3d4d4ej0Z:Gd5d6d6ej0Z;Gd7d8d8eZed;e=Gdd?d?ej0Z@ed@e=GdAdBdBe<ZAGdCdDdDej0ZBGdEdFdFej0ZCGdGdHdHej0ZDedIe=GdJdKdKe<ZEgdLZFdS)PzPyTorch 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@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!rr"r+r+`/var/www/auris/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!rr"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&r0rr'r(r)r1r2r3r!rr"r+r+r+r,r/ps 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\}}|dkr'||}tj|||jdddd|||}n |dkr4tj|||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} |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 rdimNrFrQindexr7r)r9r'ZargsortgatherrGrHZonesr8) sequencerArBattention_masksr?r@Z hidden_dimZ ids_shuffle ids_restoreZids_keepZsequence_maskedZ label_masksr+r+r,random_maskings   rXc*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_vconfigselfri __class__r+r,r^s    zTvltPixelEmbeddings.__init__Nc Csh|j\}}}}}||}||jd|d7}|tj|jddd|f|jdd7}||j7}||fSNrrP) r9r`rhrHr'repeat_interleavergrare) rkr<rVr?rf num_channelsheightwidth embeddingsr+r+r,forwards  ( zTvltPixelEmbeddings.forwardNr#r$r%r&r^rt __classcell__r+r+rlr,rZs  rZcrY)TvltAudioEmbeddingsr[cstt||_|jj|_ttdd|j |_ |j |j d|_ ttd|j|j |j |_ttd|j |j |_|j |j d|_ ||_dSr\)r]r^TvltAudioPatchEmbeddingsr` num_patchesrrbr'rcrd type_embed_afrequency_lengthaudio_patch_sizenum_freq_patches pos_embed_a freq_embedrirjrlr+r,r^s     zTvltAudioEmbeddings.__init__NcCsh||}|d|j}||jd|d7}|tj|jddd|f|jdd7}||j7}||fSrn) r`sizer~rrHr'rorr{)rkrJrVrsrNr+r+r,rts ( zTvltAudioEmbeddings.forwardrurvr+r+rlr,rxs rxc6eZdZdZfddZdejdejfddZZS)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 r|n||f}t|tjj r)|n||f}|d|d|d|d}||_||_ ||_ ||_ ||_t j||||d|_dSNrr)Z kernel_sizeZstride)r]r^ image_sizeimage_patch_sizenum_image_channelsrd isinstance collectionsabcIterable patch_sizerprarConv2d projection)rkrirrrprdrarlr+r,r^ 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 ().r6) r9rp ValueErrorrreshaperflatten transposerard)rkr<r?rfrprqrrrsr+r+r,rts (z TvltPixelPatchEmbeddings.forward r#r$r%r&r^r'Tensorrtrwr+r+rlr,r_s r_cr)ryz 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)r]r^spectrogram_lengthr|r}num_audio_channelsrdrrrrspectrogram_sizerrprz patch_shaperrr) rkrirr|rrprdrrzrrlr+r,r^2s    z!TvltAudioPatchEmbeddings.__init__rJrc Cs|j\}}}}||jkrtd||jdks||jdkr5td|d|d|jdd|jdd ||ddd}|S) NrrrzInput audio size (rrrr6)r9rprrrrr)rkrJr?rprqrrrsr+r+r,rtGs z TvltAudioPatchEmbeddings.forwardrr+r+rlr,ry+s ryc.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)NrZembedding_sizezThe hidden size z4 is not a multiple of the number of attention heads .bias)r]r^rdnum_attention_headshasattrrr;attention_head_size all_head_sizerLinearZqkv_biasquerykeyvalueDropoutZattention_probs_dropout_probdropoutrjrlr+r,r^Xs  zTvltSelfAttention.__init__cCs6|dd|j|jf}|j|}|ddddS)NrFrr6r)rrrrIpermute)rkxZ 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)NrFrPrr6rr)rrrrr'matmulrmathsqrtrrZSoftmaxrr contiguousrrrI)rkr!attention_mask head_maskoutput_attentionsZmixed_query_layerZ key_layerZ value_layerZ query_layerZattention_scoresZattention_probsZ context_layerZnew_context_layer_shapeoutputsr+r+r,rtos(      zTvltSelfAttention.forwardNNF)r#r$r%r^rrtrwr+r+rlr,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. rirNcs.tt|j|j|_t|j|_dSru) r]r^rrrddenserhidden_dropout_probrrjrlr+r,r^ zTvltSelfOutput.__init__r! input_tensorcC||}||}|Srurrrkr!rr+r+r,rt  zTvltSelfOutput.forward) r#r$r%r&rr^r'rrtrwr+r+rlr,rs$rcr) TvltAttentioncs*tt||_t||_t|_dSru)r]r^r attentionroutputset pruned_headsrjrlr+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)rkheadsrSr+r+r, prune_headss zTvltAttention.prune_headsNFcCs6|||||}||d|}|f|dd}|S)Nrr)rr)rkr!rrrZ self_outputsattention_outputrr+r+r,rtszTvltAttention.forwardr)r#r$r%r^rrtrwr+r+rlr,rs rcs<eZdZdeddffdd ZdejdejfddZZS) TvltIntermediaterirNcsDtt|j|j|_t|jt rt |j|_ dS|j|_ dSru) r]r^rrrdintermediate_sizerrZ hidden_actstrr intermediate_act_fnrjrlr+r,r^s   zTvltIntermediate.__init__r!cCrru)rrrkr!r+r+r,rtrzTvltIntermediate.forward r#r$r%rr^r'rrtrwr+r+rlr,rsrcsBeZdZdeddffdd Zdejdejdejfdd ZZS) TvltOutputrirNcs.tt|j|j|_t|j|_ dSru) r]r^rrrrdrrrrrjrlr+r,r^rzTvltOutput.__init__r!rcCs ||}||}||}|Srurrr+r+r,rts  zTvltOutput.forwardrr+r+rlr,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|_dSNrZeps)r]r^Zchunk_size_feed_forwardZ seq_len_dimrrr intermediaterrr LayerNormrdlayer_norm_epslayernorm_beforelayernorm_afterrjrlr+r,r^s    zTvltLayer.__init__NFc Csj|j|||||d}|d}|dd}|||j}||}||}|||}|f|}|S)Nrrr)rrtor8rrr) rkr!rrrZself_attention_outputsrrZ layer_outputr+r+r,rts     zTvltLayer.forwardrrvr+r+rlr,rs  rcs0eZdZfddZ     dddZZS) TvltEncodercs:t|_tfddtjD|_d|_dS)Ncg|]}tqSr+r.0_rir+r, z(TvltEncoder.__init__..F) r]r^rir ModuleListrangenum_hidden_layerslayergradient_checkpointingrjrlrr,r^ s   zTvltEncoder.__init__NFTc Cs|rdnd}|r dnd}t|jD]:\} } |r||f}|dur$|| nd} |jr7|jr7|| j||| |} n| ||| |} | d}|rK|| df}q|rS||f}|satdd|||fDSt|||dS)Nr+rrcs|] }|dur|VqdSrur+rvr+r+r, 9z&TvltEncoder.forward..)rr!r") enumeraterrtraining_gradient_checkpointing_func__call__tupler ) rkr!rrroutput_hidden_states return_dictall_hidden_statesall_self_attentionsi layer_moduleZlayer_head_mask layer_outputsr+r+r,rts8    zTvltEncoder.forward)NNFFTr#r$r%r^rtrwr+r+rlr,r s  rc@s(eZdZdZeZdZdZdZddZ dS)TvltPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. tvltr<TcCstt|tjtjfr#|jjjd|jjd|j dur!|j j dSdSt|tj r8|j j |jj ddSdS)zInitialize the weights)meanZstdNg?) rrrrweightdataZnormal_riZinitializer_rangerZzero_rZfill_)rkmoduler+r+r, _init_weightsLs   z!TvltPreTrainedModel._init_weightsN) r#r$r%r&r config_classZbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingrr+r+r+r,rAs 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)r]r^rirZpixel_embeddingsrxaudio_embeddingsrencoderrrbr'rcrd cls_embeddingZuse_mean_pooling layernormrr post_initrjrlr+r,r^s     zTvltModel.__init__cCs|jj|jjfSru)rr`r )rkr+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)rkZheads_to_prunerrr+r+r, _prune_headsszTvltModel._prune_heads output_typerNFr<rJr=rK 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>)rVr)rKr>rLrMr)rrrr) rrrrrrr r!r")rirruse_return_dictrr rCZpixel_mask_ratiorXr|r}rOZaudio_mask_ratioZaudio_mask_typerr'catr rHZget_extended_attention_maskr r rr!r")rkr<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,rts%     "     zTvltModel.forward)NNFFNNN)r#r$r%r^rrrTVLT_INPUTS_DOCSTRINGrr_CONFIG_FOR_DOCr'r(rboolrrrtrwr+r+rlr,rsD     rcs,eZdZfddZ   dddZZS) TvltDecodercsvtt||j_|j_|j_|j _ t fddt |jD|_t j|j|jd|_d|_||_dS)Ncrr+rrZdecoder_configr+r,rBrz(TvltDecoder.__init__..rF)r]r^rdecoder_hidden_sizerdZdecoder_num_hidden_layersrZdecoder_num_attention_headsrZdecoder_intermediate_sizerrrrdecoder_layersrrr rrirjrlrr,r^9s  zTvltDecoder.__init__FTc Cs|rdnd}|r dnd}t|jD].\}}|r||f}|jr,|jr,||j|d|} n|||d} | d}|r?|| df}q|rG||f}||} |sZtdd| ||fDSt| ||dS)Nr+rrrcsrrur+rr+r+r,rnrz&TvltDecoder.forward..)r.r!r") rr rrrrr rr-) rkr!rrrrrrrrr.r+r+r,rtJs.       zTvltDecoder.forward)FFTrr+r+rlr,r8s  rzTThe 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 trueTrrrr6)+r]r^ri task_matchingtask_maerrrTvltMatchingHead matching_headrrrdrencoder_to_decoderrbr'rcpixel_mask_tokenaudio_mask_tokenrdecoderrfrradecoder_pixel_pos_embeddecoder_temporal_embeddecoder_pixel_type_embedr rzr|r}decoder_audio_pos_embeddecoder_freq_embeddecoder_audio_type_embedrr TvltMAEHeadpixel_mae_headraudio_mae_headr~r ) rkrirrfraZnum_audio_patchesr~Zpixel_mae_output_dimZaudio_mae_output_dimrlr+r,r^wsL          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 rr9zntchpwq->nthwpqc)r9rrr'einsum) rkr<r?rfrprqrrnum_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] rrr3znchpwq->nhwpqc)r9r}rr'r4) rkrJr?rprqrrr5r6patchified_audio_valuesr+r+r,patchify_audios( z!TvltForPreTraining.patchify_audiocC:||}||d}|jdd}|||}|SNr6rFrP)r8rsum)rkr<Zpixel_predictionsmaskr7r0r+r+r,pixel_mae_loss   z!TvltForPreTraining.pixel_mae_losscCr;r<)r:rr=)rkrJZaudio_predictionsr>r9r0r+r+r,audio_mae_lossr@z!TvltForPreTraining.audio_mae_lossc CsZ|j\}}}|||jd|d}tj||gdd}tj|d|ddd|d}|S)NrrPrFrR)r9rHr'rrTrG) rkZ mask_tokenrUrWr?Z seq_lengthrQZ mask_tokensZpadded_sequencer+r+r,concatenate_masks z#TvltForPreTraining.concatenate_maskrNr<rJr=rKlabelspixel_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=rKrrrrrFT)r=rKrrrrrrr6rr rP)r0r1r2r3r!r"))rirr"rrr%rrIr#rrrrrrr r&rrBr'r*rHr'ror+rar,r)r1r.r(r~r.r-r/r2r?rAr/r!r") rkr<rJr=rKrCrDrErrrZ total_lossrrr1loss_fctr0r2r3rrrrrr Zpixel_decoder_inputZaudio_decoder_inputrfZpixel_decoder_outputsrNZaudio_decoder_outputsrr+r+r,rts1         zTvltForPreTraining.forward)NNNNNNNN)r#r$r%r^r8r:r?rArBrrrr/rr'r(rr*rrrrtrwr+r+rlr,r!rsP 6     r!c$eZdZfddZddZZS) TvltPoolercs*tt|j|j|_t|_dSru)r]r^rrrdrZTanh activationrjrlr+r,r^s zTvltPooler.__init__cCs(|dddf}||}||}|S)Nr)rrM)rkr!Zfirst_token_tensorZ pooled_outputr+r+r,rts  zTvltPooler.forwardrr+r+rlr,rL rLcrK)r$cs(tt||_t|jd|_dSr\)r]r^rLpoolerrrrdfcrjrlr+r,r^s  zTvltMatchingHead.__init__cCs|||}|Sru)rPrOrr+r+r,rtszTvltMatchingHead.forwardrr+r+rlr,r$rNr$cs&eZdZdfdd ZddZZS)r0Ncs$t||_t|j||_dSru)r]r^rirrrr))rkriZ output_dimrlr+r,r^s zTvltMAEHead.__init__cCs||}|Sru)r)rr+r+r,rts zTvltMAEHead.forwardrurr+r+rlr,r0sr0z 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)Nr6r)r]r^rrrZ SequentialrrdrrZGELUZ num_labels classifierrir rjrlr+r,r^s   z)TvltForAudioVisualClassification.__init__rNr<rJr=rKrrrrCrc  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 ```NrFrZ regressionZclassificationr )r0r.r!r") rirrrRZ loss_typer r rr!r")rkr<rJr=rKrrrrCrrr.r0rJrr+r+r,rts:$     z(TvltForAudioVisualClassification.forward)NNNNNN)r#r$r%r^rrrrrr'r(rrr*rrrtrwr+r+rlr,rQs:    rQ)rr!rQr)Nr4)Nr4rDrEru)Gr&collections.abcrrcopyr dataclassesrtypingrrrr'Ztorch.utils.checkpointrZtorch.nnrr r Z activationsr Zmodeling_outputsr rZmodeling_utilsrZ pytorch_utilsrrutilsrrrrrZconfiguration_tvltrZ get_loggerr#loggerrZ_CHECKPOINT_FOR_DOCrr-r/rCrOrXModulerZrxr_ryrrrrrrrrZTVLT_START_DOCSTRINGrrrr!rLr$r0rQ__all__r+r+r+r,s       ( ! ),<"&5 -$:    Y