o h@sdZddlZddlmZddlmZmZddlZddlZddlm Z ddl m Z ddl m Z mZmZdd lmZmZmZdd lmZdd lmZdd lmZd dlmZddlmZeeZ ddgZ!eGdddeZ"Gddde j#Z$Gddde j#Z%Gddde j#Z&Gddde j#Z'Gddde j#Z(Gdd d e j#Z)Gd!d"d"e j#Z*Gd#d$d$e j#Z+Gd%d&d&e j#Z,Gd'd(d(e j#Z-Gd)d*d*eZ.d+Z/d,Z0Gd-d.d.e j#Z1Gd/d0d0e j#Z2e1e2d1Z3e d2e/Gd3d4d4e.Z4Gd5d6d6e j#Z5e d7e/Gd8d9d9e.Z6dS):zPyTorch TVP ModelN) dataclass)OptionalTuple)nn)ACT2FN)add_start_docstrings%add_start_docstrings_to_model_forwardreplace_return_docstrings)BaseModelOutputBaseModelOutputWithPooling ModelOutput)PreTrainedModel)prune_linear_layer)logging) AutoBackbone) TvpConfigzIntel/tvp-basezIntel/tvp-base-ANetc@s^eZdZUdZdZeejed<dZ ejed<dZ ee ejed<dZ ee ejed<dS)TvpVideoGroundingOutputa Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`): Temporal-Distance IoU loss for video grounding. logits (`torch.FloatTensor` of shape `(batch_size, 2)`): Contains start_time/duration and end_time/duration. It is the time slot of the videos corresponding to the input texts. 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, if the model has an embedding layer, + 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 optional 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)`. Nlosslogits hidden_states attentions) __name__ __module__ __qualname____doc__rrtorch FloatTensor__annotations__rrrrr!r!Z/var/www/html/ai/venv/lib/python3.10/site-packages/transformers/models/tvp/modeling_tvp.pyr,s rcs@eZdZdZfddZddZddZdd Zd d ZZ S) TvpLossa~ This class computes the losses for `TvpForVideoGrounding`. The process happens in two steps: 1) we compute hungarian assignment between ground truth boxes and the outputs of the model 2) we supervise each pair of matched ground-truth / prediction (supervise class and box). Args: losses (`List[str]`): List of all the losses to be applied. csLt|j|j|jd|_|D]}||jvr td|dq||_dS)NioudistancedurationzLoss z not supported)super__init__loss_iou loss_distance loss_durationloss_map ValueErrorlosses)selfr/r __class__r!r"r)Os   zTvpLoss.__init__c CsHt||t||}t||t||}d|jdd|}|S)z6 Measure the intersection over union. rrmin)rr4maxclamp) r0 start_timeend_timecandidates_start_timecandidates_end_timer'interunionr%r!r!r"r*\szTvpLoss.loss_iouc CsTtt||d}tt||d}tt||t|||jdd}|S)z5 Measure the distance of mid points. g@g?r3)rdivaddr5r4r6) r0r7r8r9r:r'mid_candidatesmid_groundtruth distance_diffr!r!r"r+fszTvpLoss.loss_distancec CsBt||}t||}ttt|||}|jdd}|S)z5 Measure the difference of duration. g?r3)rsubsquarer=r6) r0r7r8r9r:r'duration_candidatesduration_groundtruth duration_diffr!r!r"r,rs   zTvpLoss.loss_durationc Cst|\}}}t||}|dddf|dddf}}i} |jD]} | | |j| |||||iq%| S)am This performs the loss computation. Args: logits (`torch.FloatTensor`): The output logits of head module. labels (`List[torch.FloatTensor]`): List of tensors ([start, end, duration]), which contains start time, end time of the video corresponding to the text, and also the duration. Nrr)rmulfloatr/updater-) r0rlabelsr'r7r8 candidatesr9r: losses_dictrr!r!r"forward}s * zTvpLoss.forward) rrrrr)r*r+r,rM __classcell__r!r!r1r"r#Ds     r#c$eZdZfddZddZZS)TvpVisionModelc s@tt|j|_tj|jjd|j dddddd|_ dS)NrrF) kernel_sizestridepaddinggroupsbias) r(r)r from_configbackbone_configbackbonerConv2d hidden_sizes hidden_sizegrid_encoder_convr0configr1r!r"r)s   zTvpVisionModel.__init__c Cs|j\}}}}}||||||}||dd}||}tjj|ddd}tjj|dd}|jdd\} } } |||| | | }|ddd d d}|S) N feature_mapsrr)rRrST)inplacerr) shapeviewrYr]r functional max_pool2drelupermute) r0 pixel_values batch_size num_frames num_channelsheightwidthgrid_feat_outputsgrid new_channel new_height new_widthr!r!r"rMs zTvpVisionModel.forwardrrrr)rMrNr!r!r1r"rPs  rPcs0eZdZdZfddZddZddZZS)TvpVisualInputEmbeddingz; Takes input of both image and video (multi-frame) csvtt|j|j|_t|j|j|_t|j |j|_ td|j|_ tj |j|j d|_t|j|_dS)Nreps)r(r)r Embeddingmax_position_embeddingsr\position_embeddings max_grid_row_position_embeddingsrow_position_embeddings max_grid_col_position_embeddingscol_position_embeddingstoken_type_embeddings LayerNormlayer_norm_eps layer_normDropouthidden_dropout_probdropoutr^r1r!r"r)s z TvpVisualInputEmbedding.__init__c Cs|j\}}}}tj|tj|jd}||}dt|jd|d|f}||j|}tj|tj|jd} || } |d||f} || j| S)z Args: grid: (batch_size, height, width, hidden_dim) Returns: grid + col_position_embeddings.view(*col_shape): (batch_size, *, height, width, hidden_dim) dtypedevice)rrr) rdrarangelongrr}lenrer) r0rqrkrnro hidden_dimrow_position_idsr} row_shapecol_position_idsr col_shaper!r!r"add_2d_positional_embeddingss   z4TvpVisualInputEmbedding.add_2d_positional_embeddingsc Cs|j\}}}}}|d}||}||d|}|jdd}|j} tj|tj| d} || } || } | | } | | } | S)a} Args: grid: Array of shape (batch_size, num_frames, height, width, num_channels). It contains processed frames extracted from videos, and is generated by Tvp image preprocessor. Note, num_frames can be 1 Returns: embeddings: The embedding of grid with size (batch_size, height*width, num_channels) rrQNr) rdmeanrrerrzerosrrrr) r0rqrkrlrnrorm visual_tokensvisual_tokens_shapertoken_type_idsr embeddingsr!r!r"rMs     zTvpVisualInputEmbedding.forward)rrrrr)rrMrNr!r!r1r"rvs   rvcs*eZdZdZfddZdddZZS)TvpTextInputEmbeddingszGConstruct the embeddings from word, position and token_type embeddings.cslttj|j|j|jd|_t|j|j|_ t|j |j|_ tj |j|j d|_t|j|_dS)N) padding_idxrw)r(r)rry vocab_sizer\ pad_token_idword_embeddingsrzr{type_vocab_sizerrrrrrrr^r1r!r"r)s zTvpTextInputEmbeddings.__init__Nc Cs|dur |}n|dd}|d}|dur|jn|j}|dur4tj|tj|d}|d|}|durAtj|tj|d}|durJ||}| |}| |} ||| } | | } | | } | S)NrQrrr) sizerrrr unsqueezeexpandrrr{rrr) r0 input_idsr position_ids inputs_embeds input_shape seq_lengthrr{rrr!r!r"rMs$       zTvpTextInputEmbeddings.forward)NNNNrrrrr)rMrNr!r!r1r"rs rcsVeZdZfddZddZdejdedefdd Z dd e e fd d Z Z S) TvpAttentioncst|j|jdkrt|dstd|jd|j|j|_t|j|j|_|j|j|_t |j|j|_ t |j|j|_ t |j|j|_ t |j|_t |j|j|_t j|j|jd|_t |j|_t|_dS)Nrembedding_sizezThe hidden size z4 is not a multiple of the number of attention heads rw)r(r)r\num_attention_headshasattrr.intattention_head_size all_head_sizerLinearquerykeyvaluerattention_probs_dropout_prob attn_dropoutdenserrrrrset pruned_headsr^r1r!r"r)s  zTvpAttention.__init__cst|dkrdSt|j|j}t||j}|D]tfdd|jDd|<q|d d}t t|| }t |j||_t |j||_t |j||_t |j|dd|_|jt||_|j|j|_|j||_dS)Nrc3s |] }|kr dndVqdS)rrNr!).0hheadr!r" 5sz+TvpAttention.prune_heads..rQrdim)rronesrrrrsumre contiguouseqrrrrrrrrr<)r0headsmaskindexr!rr" prune_heads.s  zTvpAttention.prune_headstensorsequence_lengthrkcCs ||||j|jddS)Nrr)rerr transposer)r0rrrkr!r!r"_reshapeEszTvpAttention._reshapeNoutput_attentionscCs |jdd\}}||}||}||} ||||} ||||} || ||} t| | dd} | t |j } |durG| |} t j j | dd}||}|dur\||}t|| }|dd}||||j}||}||}|||}|r||f}|S|f}|S)NrrQrr)rdrrrrrmatmulrmathsqrtrrrfsoftmaxrrreshaperrrr)r0rattention_mask head_maskrrkrmixed_query_layermixed_key_layermixed_value_layer query_layer key_layer value_layerattention_scoresattention_probs attn_outputoutputsr!r!r"rMLs2        zTvpAttention.forwardNNN) rrrr)rrTensorrrrboolrMrNr!r!r1r"rs  rc2eZdZfddZdejdejfddZZS)TvpIntermediatecsDtt|j|j|_t|jt rt |j|_ dS|j|_ dSN) r(r)rrr\intermediate_sizer isinstance hidden_actstrrintermediate_act_fnr^r1r!r"r)|s   zTvpIntermediate.__init__rreturncCs||}||}|Sr)rr)r0rr!r!r"rMs  zTvpIntermediate.forwardrrrr)rrrMrNr!r!r1r"r{s rcs8eZdZfddZdejdejdejfddZZS)TvpOutputLayercsBtt|j|j|_tj|j|jd|_ t |j |_ dS)Nrw) r(r)rrrr\rrrrrrrr^r1r!r"r)s zTvpOutputLayer.__init__r input_tensorrcCs&||}||}|||}|Sr)rrr)r0rrr!r!r"rMs  zTvpOutputLayer.forwardrr!r!r1r"rs $rcs6eZdZfddZ   ddeefddZZS)TvpEncodeLayercs,tt||_t||_t||_dSr)r(r)r attentionr intermediateroutputr^r1r!r"r)s   zTvpEncodeLayer.__init__Nrc CsJ|j||||d}|d}|dd}||}|||} | f|}|S)N)rrr)rrr) r0rrrrself_attention_outputsattention_outputrintermediate_output layer_outputr!r!r"rMs    zTvpEncodeLayer.forwardr)rrrr)rrrMrNr!r!r1r"rs  rc sTeZdZfddZ     d deejdeedeedeefdd ZZ S) TvpEncodercs:t|_tfddtjD|_d|_dS)Ncsg|]}tqSr!)r)r_r_r!r" sz'TvpEncoder.__init__..F) r(r)r_r ModuleListrangenum_hidden_layerslayergradient_checkpointingr^r1rr"r)s   zTvpEncoder.__init__Nrroutput_hidden_states return_dictc Cs|dur|n|jj}|dur|n|jj}|dur|n|jj}d}d}t|jD]:\} } |r2||f}|jrK|jrK|| j |||durF|| nd|} n | |||| |} | d}|ra|| df}q'|ri||f}|s~|f} |ru| |f} |r|| |f} | St ||r|nd|r|dSddS)Nr!rr)last_hidden_staterr) r_rrr enumeraterrtraining_gradient_checkpointing_func__call__r ) r0rrrrrrall_hidden_statesall_attentionsi layer_module layer_outputsrr!r!r"rMsL       zTvpEncoder.forward)NNNNN) rrrr)rrrrrMrNr!r!r1r"rs   rcr) TvpPoolercs*tt|j|j|_t|_dSr)r(r)rrr\rTanh activationr^r1r!r"r)s zTvpPooler.__init__rrcCs(|dddf}||}||}|S)Nr)rr )r0rfirst_token_tensor pooled_outputr!r!r"rMs  zTvpPooler.forwardrr!r!r1r"r s r c@s$eZdZdZeZdZdZddZdS)TvpPreTrainedModelzAn abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. modelTcCst|tjtjfr|jjjd|jjdnt|tj r(|j j |jj dt|tjr9|j dur9|j j t|tj rXtjj|jddd|j durZtj|j ddSdSdS) zInitialize the weights)rstdg?Nfan_outrh)mode nonlinearityr)rrrryweightdatanormal_r_initializer_rangerrVzero_fill_rZinitkaiming_normal_ constant_)r0moduler!r!r" _init_weights s     z TvpPreTrainedModel._init_weightsN) rrrrr config_classbase_model_prefixsupports_gradient_checkpointingr r!r!r!r"rs  raE 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 ([`TvpConfig`]): 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: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_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) pixel_values (`torch.FloatTensor` of shape `(batch_size, num_frames, num_channels, height, width)`): Pixel values. Pixel values can be obtained using [`TvpImageProcessor`]. See [`TvpImageProcessor.__call__`] for details. attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. 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) head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. 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. c(eZdZdZfddZddZZS)TvpFrameDownPadPrompterz> Pad frames extracted from videos only at the bottom. c sb|jdvr tdt|j|_|j|_|j|_|j|_tt d|jd|j|jg|_ dS)Nr>replaceremove9`visual_prompter_apply` must be in (add, replace, remove)rr) visual_prompter_applyr.r(r)visual_prompt_size frame_num max_img_sizer Parameterrrandnpad_downr^r1r!r"r)Qs   z TvpFrameDownPadPrompter.__init__cCs|jdkr&tj|j|jg|j|jd}d||j|j|jddf<||9}|jdkrctj|jd|jdd|j|jg|jd}|j|j}|j |dddddd||jddf<|| |j7}|S) Nr>rrr(rrrr) r*rrr-rrr+rrdr0to)r0rjvisual_prompt_maskprompt start_pointr!r!r"rM_s   *zTvpFrameDownPadPrompter.forwardrr!r!r1r"r%Ls r%cr$)TvpFramePadPrompterz? Pad frames extracted from videos in the surroundings. c s|jdvr tdt|j|_|j|_|j|_|j|jd|_t t d|jd|j|jg|_ t t d|jd|j|jg|_ t t d|jd|j|jd|jg|_t t d|jd|j|jd|jg|_dS)Nr&r)rrr)r*r.r(r)rlr-r+ base_sizerr.rr/pad_upr0pad_left pad_rightr^r1r!r"r)vsB    zTvpFramePadPrompter.__init__cCs|jdvr td|j|jdvr$tj|j|jg|j|jd}||9}|jdvrctjd|jd|j |j |jd}tj |j ||j gd d }tj |j ||jgdd }t |d |g}|||j7}|S) N)r>r(r'z$Invalid visual_prompter_apply value )r'r(r)r'r>rrr1rcrr)r*r.rrr-rrrrlr7catr9r:r8r0rr2)r0rjr3baser4r!r!r"rMs   zTvpFramePadPrompter.forwardrr!r!r1r"r6qs 'r6) framedownpadframepadzmThe bare Tvp Model transformer outputting BaseModelOutputWithPooling object without any specific head on top.cseZdZfddZddZddZddZeee e e d  dd e e jd e e jd e e jde e jde ede ede efddZZS)TvpModelcst|||_t||_t||_t||_t ||_ t ||_ t tdd|jg|_t |j|_|jtvr?tdt|j||_|dS)Nr z:`visual_prompter_type` must be in (framedownpad, framepad))r(r)r_rP vision_modelrrrvvisual_embeddingsrencoderr poolerrr.rr/r\ text_promptrrrvisual_prompter_typeTVP_PROMPTER_CLASSES_MAPPINGr.visual_prompter post_initr^r1r!r"r)s        zTvpModel.__init__cCs|jjSrrr)r0r!r!r"get_input_embeddingsszTvpModel.get_input_embeddingscCs ||j_dSrrJ)r0rr!r!r"set_input_embeddingss zTvpModel.set_input_embeddingscCs*|D]\}}|jj|j|qdS)zPrunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base class PreTrainedModel N)itemsrCrrr)r0heads_to_prunerrr!r!r" _prune_headsszTvpModel._prune_heads output_typer!NrrjrrrrrcCsJ|dur|n|jj}|||}|j|d}||} |durQ|| jdd} t |jddj |j |j d} tj | || gdd}||| |j }|j|jddd} tj | || gd d} |j| ||||jj|||d }|r||jn|d}||}||}||}|s||f|d dSt|||j|jd S) a( Returns: Examples: ```python >>> import torch >>> from transformers import AutoConfig, AutoTokenizer, TvpModel >>> model = TvpModel.from_pretrained("Jiqing/tiny-random-tvp") >>> tokenizer = AutoTokenizer.from_pretrained("Jiqing/tiny-random-tvp") >>> pixel_values = torch.rand(1, 1, 3, 448, 448) >>> text_inputs = tokenizer("This is an example input", return_tensors="pt") >>> output = model(text_inputs.input_ids, pixel_values, text_inputs.attention_mask) ```N)rrrr@)rrrQrr)rrrrr)r pooler_outputrr)r_rrArHrrBnew_onesrdrrr2rrr;get_extended_attention_maskrrErrC get_head_maskrrrDrr rr)r0rrjrrrrrtext_embedding_outputvisual_embedding_outputvisual_attention_maskpt_maskrEembedding_outputencoder_outputsrrr!r!r"rMsB     zTvpModel.forward)NNNNNNN)rrrr)rKrLrOr TVP_INPUTS_DOCSTRINGr r rrr LongTensorrrrMrNr!r!r1r"r?s:  r?crO)TvpVideoGroundingHeadcsLtt|j|jd|_t|jdd|_t|_t |_ dS)Nr) r(r)rrr\layer_0layer_1ReLU activation_0Sigmoid activation_1r^r1r!r"r) s  zTvpVideoGroundingHead.__init__cCs$|||}|||}|Sr)rbr_rdr`)r0rRrr!r!r"rM'szTvpVideoGroundingHead.forwardrur!r!r1r"r^s r^zb Tvp Model with a video grounding head on top computing IoU, distance, and duration loss. cseZdZfddZeeeeed        dde e j de e j de e j de e jd e e j d e ed e ed e efd dZZS)TvpForVideoGroundingcs2t|||_t||_t||_|dSr)r(r)r_r?rr^video_grounding_headrIr^r1r!r"r)4s    zTvpForVideoGrounding.__init__rPNrrjrrJrrrrc  Cs|dur|n|jj}|j|||||||d} | d} || } d} |durJtgd} | |j| | |}|d|jj|d|jj|d} |s`| f| dd} | dur^| f| } | St | | | j | j d S) a labels (`torch.FloatTensor` of shape `(batch_size, 3)`, *optional*): The labels contains duration, start time, and end time of the video corresponding to the text. Returns: Examples: ```python >>> import torch >>> from transformers import AutoConfig, AutoTokenizer, TvpForVideoGrounding >>> model = TvpForVideoGrounding.from_pretrained("Jiqing/tiny-random-tvp") >>> tokenizer = AutoTokenizer.from_pretrained("Jiqing/tiny-random-tvp") >>> pixel_values = torch.rand(1, 1, 3, 448, 448) >>> text_inputs = tokenizer("This is an example input", return_tensors="pt") >>> output = model(text_inputs.input_ids, pixel_values, text_inputs.attention_mask) ```N)rrrrrr$r%r&r'r)rrrr) r_rrrfr#r2rdistance_loss_weightduration_loss_weightrrr)r0rrjrrJrrrrrrRrr criterion loss_dictr!r!r"rM<sD     zTvpForVideoGrounding.forward)NNNNNNNN)rrrr)r r\r rrrrr]rrrrrMrNr!r!r1r"re-s:   re)7rr dataclassesrtypingrrrtorch.utils.checkpointr activationsr file_utilsrr r modeling_outputsr r r modeling_utilsr pytorch_utilsrutilsrautorconfiguration_tvpr get_loggerrlogger!TVP_PRETRAINED_MODEL_ARCHIVE_LISTrModuler#rPrvrrrrrrr rTVP_START_DOCSTRINGr\r%r6rGr?r^rer!r!r!r"sd         PB$c? #%>g