o iJ@sdZddlZddlmZddlmZmZmZddlZddlm Z ddl m Z ddl m Z dd lmZdd lmZmZdd lmZdd lmZmZmZdd lmZmZmZddlmZee Z!   dtde j"dej#dej#dej#deej#dee$de$deej#fddZ%Gddde j"Z&Gddde j"Z'   dud ej#d!e$d"ee(d#e)d$e*f d%d&Z+  dvd ej#d'ee(e*fd"ee(d$e*fd(d)Z,Gd*d+d+e j"Z-Gd,d-d-e j"Z.Gd.d/d/e j"Z/eGd0d1d1eZ0Gd2d3d3e j"Z1Gd4d5d5e j"Z2Gd6d7d7e0Z3eed8d9Gd:d;d;eZ4eedd>eZ5eed?d9Gd@dAdAeZ6eedBd9GdCdDdDeZ7eedEd9GdFdGdGeZ8eedHd9GdIdJdJeZ9dKej:j;dLej#dMej#fdNdOZGdVdWdWe j"Z?GdXdYdYe j"Z@GdZd[d[e j"ZAeGd\d]d]e0ZBGd^d_d_e j"ZCed`d9Gdadbdbe0ZDGdcdddde j"ZEeded9Gdfdgdge0ZFedhd9Gdidjdje j"ZGedkd9Gdldmdme0ZHGdndodoe j"ZIedpd9Gdqdrdre0ZJgdsZKdS)xzPyTorch PatchTST model.N) dataclass)CallableOptionalUnion)nn)ACT2CLS)FlashAttentionKwargs)BaseModelOutput)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)NegativeBinomialOutput NormalOutputStudentTOutput) ModelOutputauto_docstringlogging)PatchTSTConfigmodulequerykeyvalueattention_maskscalingdropout head_maskc Ks|dur |dd}t||dd|} |dur| |} tjj| dd} |dur5| |dddd} tjj| ||j d} t| |} | dd } | | fS)Nrdimr)ptraining) sizetorchmatmul transposer functionalsoftmaxviewrr% contiguous) rrrrrrrrkwargs attn_weights attn_outputr1k/home/ubuntu/veenaModal/venv/lib/python3.10/site-packages/transformers/models/patchtst/modeling_patchtst.pyeager_attention_forward&s  r3cseZdZdZ     ddededed ed ed ed eeffd d Z    dde j dee j dee j dee j deede e dee j ee j eee j ffddZZS)PatchTSTAttentionz=Multi-headed attention from 'Attention Is All You Need' paperrFTN embed_dim num_headsr is_decoderbias is_causalconfigcst||_||_||_|||_||_|j||jkr*td|jd|d|jd|_||_ ||_ t j |||d|_ t j |||d|_t j |||d|_t j |||d|_dS)Nz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: ).r r8)super__init__r5r6rhead_dimr: ValueErrorrr7r9rLineark_projv_projq_projout_proj)selfr5r6rr7r8r9r: __class__r1r2r>Hs&    zPatchTSTAttention.__init__ hidden_stateskey_value_statesrlayer_head_maskoutput_attentionsr.returncKs|du}|jdd\}} |r|jdn| } || d|jf} || d|jf} ||j| dd} |r4|n|}||j| dd}||j| dd}t}|jj dkr\t |jj }||| |||f|j shdn|j |j ||d|\}}||| d}||}||dfS)z#Input shape: Batch x Time x ChannelNrrr!eagerr)rrrLr)shaper?rDr,r)rBrCr3r:_attn_implementationr r%rrreshaper-rE)rFrIrJrrKrLr.is_cross_attentionbsztgt_lensrc_len q_input_shapekv_input_shape query_statescurrent_states key_states value_statesattention_interfacer0r/r1r1r2forwardgs:      zPatchTSTAttention.forward)rFTFN)NNNF)__name__ __module__ __qualname____doc__intfloatboolrrr>r'Tensorr r tupler] __classcell__r1r1rGr2r4EsR"  r4c6eZdZdZdeffdd ZdejfddZZ S)PatchTSTBatchNormzP Compute batch normalization over the sequence length (time) dimension. r:cs"ttj|j|jd|_dS)Neps)r=r>r BatchNorm1dd_modelnorm_eps batchnormrFr:rGr1r2r>s zPatchTSTBatchNorm.__init__inputscCs"|dd}||}|ddS)a Parameters: inputs (`torch.Tensor` of shape `(batch_size, sequence_length, d_model)`): input for Batch norm calculation Returns: `torch.Tensor` of shape `(batch_size, sequence_length, d_model)` rr!)r)ro)rFrqoutputr1r1r2r]s   zPatchTSTBatchNorm.forward r^r_r`rarr>r'rer]rgr1r1rGr2risriFrq mask_ratiounmasked_channel_indiceschannel_consistent_masking mask_valuecCs*|dks|dkrtd|d|j\}}}}|j} t|d|} |r5tj|d|| d} | d|d} n tj|||| d} tj|||| d} d| ddddd| f<tj| dd} tj| dd}tj | d|d } | dddd|} |durd| dd|ddddf<| | |}|| d fS) arandom_masking: Mask the input considering the control variables. Args: inputs (`torch.Tensor` of shape `(batch_size, num_channels, sequence_length, num_features)`): The input tensor to mask. mask_ratio (`float`): Masking ratio applied to mask the input data during random pretraining. It is the number between 0 and 1. unmasked_channel_indices (list, *optional*): Indices of channels that will not be masked. channel_consistent_masking (bool, *optional*, defaults to `False`): When true, masking will be same across all channels of a timeseries. Otherwise, masking positions will vary across channels. mask_value (int, *optional*, defaults to 0): Define the value of masked patches for pretraining. Returns: `tuple(torch.Tensor)`: inputs_mask, masked input, same shape as input Tensor and mask tensor of shape [bs x c x n] rrz Mask ratio z has to be between 0 and 1.deviceNrr")r#index.r) r@rOryrbr'randrepeatonesargsortgather unsqueeze masked_fillrd)rqrtrurvrw batch_size num_channelssequence_length num_featuresrylen_keepnoisemask ids_shuffle ids_restore inputs_maskr1r1r2random_maskings& rnum_forecast_mask_patchescCst|tr|g}dd|D}|j\}}}}tj||||jd} g} d} t|} t||D](\} }| dks9| |krAtd| dt||| }| | ||g| |7} q-t | ddd } | |krq| dd || | dd <n| |kr| d d | || d d <d}| D]\}}}||}d | ||d d | d f<|}qt | jd}| |} | d  d d d |} |d urd| d d |d d d d f<|| |}|| dfS)aForecast masking that masks the last K patches where K is from the num_forecast_mask_patches. If num_forecast_mask_patches is a list, samples in the batch will be randomly masked by numbers defined in the list. Parameters: inputs (`torch.Tensor`): Input of shape `(bs, num_channels, num_patch, patch_length)` num_forecast_mask_patches (`list`): Number of patches to be masked at the end of each batch sample. e.g. 4 or [3, 5]. unmasked_channel_indices (`list`, *optional*): Indices of channels that are not masked. mask_value (`int`, *optional*, defaults to 0): Values in the masked patches will be filled by `mask_value`. Returns: `tuple(torch.Tensor)`: inputs_mask, masked input, same shape as inputs Tensor and Mask tensor of shape `(bs, num_channels , num_patch)` or `(bs, tsg1, tsg2, num_channels, num_patch)` cSsg|]}dqS)rr1.0_r1r1r2 sz$forecast_masking..rxrznum_forecast_mask_patches z6 should be greater than 0 and less than total patches.cSs|dS)Nr!r1)xr1r1r2sz"forecast_masking..)rr!rrNr{) isinstancerbrOr'zerosrysumzipr@appendsortedrandpermrr}rrd)rqrrurwforecast_mask_ratiosrrrrrt_list total_length total_ratio patch_lengthratiotemp_lenbatch1 patch_lenrbatch2permrr1r1r2forecast_maskingsB    rcrh)PatchTSTPatchifyz A class to patchify the time series sequence into different patches Returns: `torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)` r:cst|j|_|j|_|j|_|j|jkr$td|jd|jdt|j|j|j|jd|_|j|j|jd}|j||_ dS)NzSequence length (z+) has to be greater than the patch length ()r) r=r>context_lengthrr patch_strider@max num_patchessequence_start)rFr:new_sequence_lengthrGr1r2r>9s   zPatchTSTPatchify.__init__ past_valuescCsp|jd}||jkrtd|d|jd|dd|jdddf}|jd|j|jd}|dd}|S)a! Parameters: past_values (`torch.Tensor` of shape `(batch_size, sequence_length, num_channels)`, *required*): Input for patchification Returns: `torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)` zInput sequence length (z%) doesn't match model configuration (r;N) dimensionr&step) rOrr@runfoldrrr)r-)rFrrrrr1r1r2r]Js zPatchTSTPatchify.forwardrsr1r1rGr2r1srcrh)PatchTSTMaskinga Class to perform random or forecast masking. Parameters: config (`PatchTSTConfig`): model config Returns: x_mask (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`) Masked patched input mask (`torch.Tensor` of shape `(batch_size, num_channels, num_patches)`) Bool tensor indicating True on masked points r:csXt|j|_|j|_|j|_|j|_|j|_|j|_|jdur*t|j|_dSdSN) r=r>random_mask_ratiorv mask_typerrurwrrprGr1r2r>ns  zPatchTSTMasking.__init__ patch_inputcCsr|jdkrt||j|j|j|jd\}}n|jdkr(t||j|j|jd\}}n td|jd| }||fS)a Parameters: patch_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`, *required*): Patch input Return: masked_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`) Masked patched input mask (`torch.Tensor` of shape `(batch_size, num_channels, num_patches)`) Bool tensor indicating True on masked points random)rqrtrurvrwforecast)rqrrurwzInvalid mask type .) rrrrurvrwrrr@rd)rFr masked_inputrr1r1r2r]ys$   zPatchTSTMasking.forwardrsr1r1rGr2ras  rcs@eZdZdZdeffdd Zd dejdee fdd Z Z S) PatchTSTEncoderLayerz PatchTST encoder layer r:c st|j|_t|j|j|j|d|_|jdkr t |jnt |_ |j dkr0t||_n|j dkr@t j|j|jd|_nt|j d|jr~|jdkrVt |jnt |_|j dkrft||_n|j dkrvt j|j|jd|_nt|j dt t j|j|j|jdt|j|jdkrt |jnt t j|j|j|jd|_|jdkrt |jnt |_|j dkrt||_n|j dkrt j|j|jd|_nt|j d|j|_dS)N)r5r6rr:rro layernormrjz$ is not a supported norm layer type.r<) r=r>channel_attentionr4rmnum_attention_headsattention_dropout self_attn path_dropoutrDropoutIdentity dropout_path1 norm_typerinorm_sublayer1 LayerNormrnr@ dropout_path2norm_sublayer2 SequentialrAffn_dimr8ractivation_function ff_dropoutff dropout_path3norm_sublayer3pre_normrprGr1r2r>sD               zPatchTSTEncoderLayer.__init__N hidden_staterLc Cs|j\}}}}|||||}|jr(|j|||d\}}} |||}n|j||d\}}} ||||}|||||}|jr|dd }|||||}|jrp|j| ||d\}} } || |}n|j||d\}} } | || |}|||||}|dd }|||||}|jr|| | ||}n ||| | |}|||||}|f} |r| |jr|| fn|f7} | S)a Parameters: hidden_state (`torch.Tensor` of shape `(batch_size, num_channels, sequence_length, d_model)`, *required*): Past values of the time series output_attentions (`bool`, *optional*): Whether or not to return the output attention of all layers Return: `torch.Tensor` of shape `(batch_size, num_channels, sequence_length, d_model)` )rIrLr!r)rOr,rrrrrQrr)r-rrrrr) rFrrLrnum_input_channelsrrmr0r/rchannel_attn_weightsoutputsr1r1r2r]sF       zPatchTSTEncoderLayer.forwardr) r^r_r`rarr>r'rerrdr]rgr1r1rGr2rs"2rc@s<eZdZUeed<dZdZdZdej fddZ d dd Z d S) PatchTSTPreTrainedModelr:modelrFrcCst|tr3t|jj|jj|jj|jjd}|jjr)tj j |j dd|d7}| |j||_ dSt|tjrH|jj|jjddSt|tr^|jjj|jjjddSt|tjr||jjj d|jjd|jdur~|jjdSdSdS)z$ Initialize weights rg{Gz?)std?r)meanrN)rPatchTSTPositionalEncodingrr:rrr use_cls_tokenrinitnormal_ cls_token_init_pe position_encrr8datazero_weightfill_rirorAinit_std)rFrrr1r1r2 _init_weights/s,      z%PatchTSTPreTrainedModel._init_weightscCst|tr ||_dSdSr)rPatchTSTEncodergradient_checkpointing)rFrrr1r1r2_set_gradient_checkpointingIs  z3PatchTSTPreTrainedModel._set_gradient_checkpointingN)F) r^r_r`r__annotations__base_model_prefixmain_input_namesupports_gradient_checkpointingrModulerrr1r1r1r2r(s rc2eZdZdeffdd ZdejfddZZS)PatchTSTEmbeddingr:cslt|j|_|j|_|jrt|j|j|_dSt |_t |jD]}|j t|j|jq%dSr) r=r>rshare_embeddingrrArrminput_embedding ModuleListranger)rFr:rrGr1r2r>Os  zPatchTSTEmbedding.__init__rcsjjd}|jkrtdjd|djr }|Sfddt|D}tj|dd}|S)a% Parameters: patch_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`, *required*): Patch input for embedding return: `torch.Tensor` of shape `(batch_size, num_channels, num_patches, d_model)` rz&The defined number of input channels (zQ) in the config has to be the same as the number of channels in the batch input (rc s2g|]}j|dd|ddddfqSr)rrirrFr1r2rms2z-PatchTSTEmbedding.forward..r")rOrr@rrrr'stack)rFrr embeddingsr1rr2r][s   zPatchTSTEmbedding.forward r^r_r`rr>r'rer]rgr1r1rGr2rNs rcsVeZdZdZdedeffdd Zedededej fddZ d e j fd d Z ZS) rz' Class for positional encoding r:rcszt|j|_|j|_|jr!ttddd|j|_ |d7}| |||_ |j dkr6t |j |_ dSt|_ dS)Nrr)r=r>rrr Parameterr'rrmrrrpositional_dropoutrrrFr:rrGr1r2r>ws z#PatchTSTPositionalEncoding.__init__rMcCs|jdkrtjt||jdd}|S|jdkrst||j}td|d}t td|jdt d|j }t |||dddddf<t |||dddddf<||}||d }tj|d d}|St|jd ) NrT requires_gradsincosrrr!g@ FzN is not a valid positional encoder. Available types are 'random' and 'sincos'.)positional_encoding_typerrr'randnrmrarangerexpmathlogsincosrrr@)r:rrpositiondiv_termr1r1r2rs  (    z#PatchTSTPositionalEncoding._init_percCs|jr8|||jddddf}|j|jddddf}||jd|jdd}tj||fdd}|S|||j}|S)Nrrrr!r") rrrrexpandrOrr'cat)rFrr cls_tokensrr1r1r2r]s z"PatchTSTPositionalEncoding.forward)r^r_r`rarrbr> staticmethodrrrr'rer]rgr1r1rGr2rrs rc sTeZdZdZdedeffdd Z  d dejde e d e e d e fd d Z Z S)rz PatchTST Encoder r:rcsTtd|_t|_t||_tfddt j D|_ | dS)NFcsg|]}tqSr1)rrr:r1r2rz,PatchTSTEncoder.__init__..) r=r>rrembedderrpositional_encoderrrrnum_hidden_layerslayers post_initrrGrr2r>s     zPatchTSTEncoder.__init__Nroutput_hidden_statesrLrMc Cs|dur|n|jj}|dur|n|jj}||}||}|r"dnd}|r(dnd}|jD]}|r6||f}|||d}|d}|rI||df}q-t|||dS)a Parameters: patch_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`, *required*): Past values of the time series output_hidden_states (bool, optional): Indicates if hidden states should be outputted. output_attentions (bool, optional): Indicates if attentions should be outputted. return: `BaseModelOutput` Nr1)rrLrr)last_hidden_staterI attentions)r:rLrrrrr ) rFrrrLrencoder_statesall_attentions encoder_layer layer_outputsr1r1r2r]s        zPatchTSTEncoder.forwardNN)r^r_r`rarrbr>r'rerrdr r]rgr1r1rGr2rsrzG Base class for model's outputs, with potential hidden states. ) custom_introc@seZdZUdZdZeejed<dZ ee ejed<dZ ee ejed<dZ eejed<dZ eejed<dZeejed<dZeejed <dS) PatchTSTModelOutputa> last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches, patch_length)`): Sequence of hidden-states at the output of the last layer of the model. 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, num_channels, height, width)`. Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. mask (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches)`, *optional*): Bool masked tensor indicating which patches are masked loc (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`, *optional*): Mean of the input data (batch_size, sequence_length, num_channels) over the sequence_length scale (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`, *optional*): Std of the input data (batch_size, sequence_length, num_channels) over the sequence_length patch_input (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches, patch_length)`): Patched input to the Transformer NrrIrrlocscaler)r^r_r`rarrr' FloatTensorrrIrfrrr&r'rr1r1r1r2r%s r%z4 Output type of [`PatchTSTForPretraining`]. c@beZdZUdZdZeejed<dZ eejed<dZ ee ejed<dZ ee ejed<dS)PatchTSTForPretrainingOutputa loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`): MSE loss. prediction_output (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction outputs of the time series modeling heads. Nlossprediction_outputrIr) r^r_r`rar+rr'r(rr,rIrfrr1r1r1r2r*  r*z3 Output type of [`PatchTSTForRegression`]. c@r))PatchTSTForRegressionOutputa loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`): MSE loss. regression_outputs (`torch.FloatTensor` of shape `(batch_size, num_targets)`): Regression outputs of the time series modeling heads. Nr+regression_outputsrIr) r^r_r`rar+rr'r(rr/rIrfrr1r1r1r2r.r-r.z3 Output type of [`PatchTSTForPrediction`]. c@seZdZUdZdZeejed<dZ eejed<dZ ee ejed<dZ ee ejed<dZ eejed<dZeejed<dS) PatchTSTForPredictionOutputa! loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`): MSE loss. prediction_outputs (`torch.FloatTensor` of shape `(batch_size, prediction_length, -1)`): Prediction outputs of the time series modeling heads. 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. loc: (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`, *optional*) Mean of the input data (batch_size, sequence_length, num_channels) over the sequence_length scale: (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`, *optional*) Std of the input data (batch_size, sequence_length, num_channels) over the sequence_length Nr+prediction_outputsrIrr&r')r^r_r`rar+rr'r(rr1rIrfrr&r'r1r1r1r2r01s r0z7 Output type of [`PatchTSTForClassification`]. c@r))PatchTSTForClassificationOutputa loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`): Total loss as the sum of the masked language modeling loss and the next sequence prediction (classification) loss. prediction_logits (`torch.FloatTensor` of shape `(batch_size, num_targets)`): Prediction scores of the PatchTST modeling head (scores before SoftMax). Nr+prediction_logitsrIr) r^r_r`rar+rr'r(rr3rIrfrr1r1r1r2r2Qs r2z Base class for time series model's predictions outputs that contains the sampled values from the chosen distribution. c@s$eZdZUdZdZeejed<dS)SamplePatchTSTOutputz sequences (`torch.FloatTensor` of shape `(batch_size, num_samples, prediction_length, num_targets)`): Sampled values from the chosen distribution. N sequences) r^r_r`rar5rr'r(rr1r1r1r2r4fs r4inputtargetrMcCs || S)zc Computes the negative log likelihood loss from input distribution with respect to target. )log_prob)r6r7r1r1r2nllws r9 input_tensorweightscCsr|dur3t|dk||t|}tj|r|j|dn|dd}|r-|j|d|S||S|j|dS)aj Computes the weighted average of a given tensor across a given `dim`, masking values associated with weight zero, meaning instead of `nan * 0 = nan` you will get `0 * 0 = 0`. Args: input_tensor (`torch.FloatTensor`): Input tensor, of which the average must be computed. weights (`torch.FloatTensor`, *optional*): Weights tensor, of the same shape as `input_tensor`. dim (`int`, *optional*): The dim along which to average `input_tensor`. Returns: `torch.FloatTensor`: The tensor with values averaged along the specified `dim`. Nrr"rmin)r'where zeros_likeclamprr)r:r;r#weighted_tensor sum_weightsr1r1r2weighted_averages "  rCc PeZdZdZdeffdd Zdejdejdeejejejffdd Z Z S) PatchTSTStdScalerz Standardize features by calculating the mean and scaling along the first dimension, and then normalizes it by subtracting from the mean and dividing by the standard deviation. r:csVtt|dr |jnd|_t|dr|jnd|_t|dr&|j|_dSd|_dS)N scaling_dimrkeepdimT minimum_scalegh㈵>)r=r>hasattrrFr#rGrHrprGr1r2r>s  zPatchTSTStdScaler.__init__robserved_indicatorrMcCsz|j|j|jd}|d}||j|j|jd|}|||dj|j|jd|}t||j}|||||fS)C Parameters: data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`): input for Batch norm calculation observed_indicator (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`): Calculating the scale on the observed indicator. Returns: tuple of `torch.Tensor` of shapes (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`, `(batch_size, 1, num_input_channels)`) rGrr!)rr#rG clamp_minr'sqrtrH)rFrrJ denominatorr&variancer'r1r1r2r]s  "zPatchTSTStdScaler.forward r^r_r`rarr>r'rerfr]rgr1r1rGr2rEsrEc rD) PatchTSTMeanScalerz Computes a scaling factor as the weighted average absolute value along the first dimension, and scales the data accordingly. r:csltt|dr |jnd|_t|dr|jnd|_t|dr#|jnd|_t|dr1|j|_dSd|_dS)NrFrrGTrH绽|= default_scale)r=r>rIrFr#rGrHrTrprGr1r2r>s  zPatchTSTMeanScaler.__init__rrJrMc Cs||j|jdd}|j|jdd}|tj|dd}|jdur:|jdd}tj|ddd}t||}n|jt|}t|dk||}tj||j d}||} |j sa|j|jd}| t ||fS)rKTrLrr<Nrr") absrr#r'r@rTsqueeze ones_liker>rHrGr?) rFrrJts_sum num_observedr' batch_sumbatch_observationsrT scaled_datar1r1r2r]s  zPatchTSTMeanScaler.forwardrQr1r1rGr2rRsrRc sXeZdZdZdeffdd Z d dejdeejde ejejejffd d Z Z S) PatchTSTNOPScalerz| Assigns a scaling factor equal to 1 along the first dimension, and therefore applies no scaling to the input data. r:cs@tt|dr |jnd|_t|dr|j|_dSd|_dS)NrFrrGT)r=r>rIrFr#rGrprGr1r2r>s  zPatchTSTNOPScaler.__init__NrrJrMcCsBtj|ddj|j|jd}tj|ddj|j|jd}|||fS)a Parameters: data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`): input for Batch norm calculation Returns: tuple of `torch.Tensor` of shapes (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`, `(batch_size, 1, num_input_channels)`) Fr)r#rG)r'rWrr#rGr?)rFrrJr'r&r1r1r2r]s  zPatchTSTNOPScaler.forwardr) r^r_r`rarr>r'rerrfr]rgr1r1rGr2r]sr]c sLeZdZdeffdd ZdejdejdeejejejffddZZ S) PatchTSTScalerr:csRt|jdks|jdurt||_dS|jdkr"t||_dSt||_dS)NrTr)r=r>rrRscalerrEr]rprGr1r2r>s  zPatchTSTScaler.__init__rrJrMcCs|||\}}}|||fS)a> Parameters: data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`): Input for scaler calculation observed_indicator (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`): Calculating the scale on the observed indicator. Returns: tuple of `torch.Tensor` of shapes (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`, `(batch_size, 1, um_input_channels)`) )r_)rFrrJr&r'r1r1r2r]s zPatchTSTScaler.forward) r^r_r`rr>r'rerfr]rgr1r1rGr2r^s r^csveZdZdeffdd Z     ddejdeejdeejdeed eed eed e e e ffd d Z Z S) PatchTSTModelr:csft|t||_t||_|j|_|jj}|jr!t||_ nt |_ t ||d|_ |dS)N)r)r=r>r^r_r patchifier do_mask_inputrrmaskingrrrencoderrrrGr1r2r>,s      zPatchTSTModel.__init__Nrpast_observed_mask future_valuesrrL return_dictrMc Cs|dur|n|jj}|dur|n|jj}|dur|n|jj}|dur't|}|||\}}} ||} |jr@| | \} } n| | d} } |j | ||d} |sk| j | j | j f}|| || | f}tdd|DSt| j | j | j | || | dS)a Parameters: past_values (`torch.Tensor` of shape `(bs, sequence_length, num_input_channels)`, *required*): Input sequence to the model past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*): Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in `[0, 1]`: - 1 for values that are **observed**, - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). future_values (`torch.BoolTensor` of shape `(batch_size, prediction_length, num_input_channels)`, *optional*): Future target values associated with the `past_values` output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers output_attentions (`bool`, *optional*): Whether or not to return the output attention of all layers return_dict (`bool`, *optional*): Whether or not to return a `ModelOutput` instead of a plain tuple. Returns: `PatchTSTModelOutput` or tuple of `torch.Tensor` (if `return_dict`=False or `config.return_dict`=False) Examples: ```python >>> from huggingface_hub import hf_hub_download >>> import torch >>> from transformers import PatchTSTModel >>> file = hf_hub_download( ... repo_id="hf-internal-testing/etth1-hourly-batch", filename="train-batch.pt", repo_type="dataset" ... ) >>> batch = torch.load(file) >>> model = PatchTSTModel.from_pretrained("namctin/patchtst_etth1_pretrain") >>> # during training, one provides both past and future values >>> outputs = model( ... past_values=batch["past_values"], ... future_values=batch["future_values"], ... ) >>> last_hidden_state = outputs.last_hidden_state ```N)rrrLcss|] }|dur|VqdSrr1)rvr1r1r2 z(PatchTSTModel.forward..)rrIrrr&r'r)r:use_return_dictrLrr'rWr_rarbrcrdrrIrrfr%)rFrrerfrrLrgscaled_past_valuesr&r'patched_values masked_valuesrencoder_outputrr1r1r2r]>s66  zPatchTSTModel.forwardNNNNN)r^r_r`rr>r'rerrdrrfr%r]rgr1r1rGr2r`*s, r`cs<eZdZdZdeffdd ZdejdejfddZZ S) PatchTSTMaskPretrainHeadz- Pretraining head for mask modelling r:csHt|jdkrt|jnt|_t|j|j |_ |j |_ dSNr) r=r> head_dropoutrrrrrArmrlinearrrprGr1r2r>s   z!PatchTSTMaskPretrainHead.__init__ embeddingrMcCs:|||}|jr|ddddddddf}|S)a Parameters: embedding (`torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True, *required*): Embedding from the model Returns: `torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True Nr)rtrr)rFrur1r1r2r]s  z PatchTSTMaskPretrainHead.forwardrsr1r1rGr2rqsrqz* The PatchTST for pretrain model. csjeZdZdeffdd Z    d dejdeejdeedeed eed e e e ff d d Z Z S)PatchTSTForPretrainingr:cs4t|d|_t|d|_t||_|dS)NTr)r=r>rbr`rrqheadrrprGr1r2r>s    zPatchTSTForPretraining.__init__NrrerrLrgrMc Cs|dur|n|jj}|j||||dd}||j}tjdd}|||j} | jdd|j  |j d} |j } |sU|f|d d } | durQ| f| } | S| } | St | || |j d S) a Parameters: past_values (`torch.Tensor` of shape `(bs, sequence_length, num_input_channels)`, *required*): Input sequence to the model past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*): Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in `[0, 1]`: - 1 for values that are **observed**, - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers output_attentions (`bool`, *optional*): Whether or not to return the output attention of all layers return_dict (`bool`, *optional*): Whether or not to return a `ModelOutput` instead of a plain tuple. Returns: `PatchTSTForPretrainingOutput` or tuple of `torch.Tensor` (if `return_dict`=False or `config.return_dict`=False) Examples: ```python >>> from huggingface_hub import hf_hub_download >>> import torch >>> from transformers import PatchTSTConfig, PatchTSTForPretraining >>> file = hf_hub_download( ... repo_id="hf-internal-testing/etth1-hourly-batch", filename="train-batch.pt", repo_type="dataset" ... ) >>> batch = torch.load(file) >>> # Config for random mask pretraining >>> config = PatchTSTConfig( ... num_input_channels=7, ... context_length=512, ... patch_length=12, ... stride=12, ... mask_type='random', ... random_mask_ratio=0.4, ... use_cls_token=True, ... ) >>> # Config for forecast mask pretraining >>> config = PatchTSTConfig( ... num_input_channels=7, ... context_length=512, ... patch_length=12, ... stride=12, ... mask_type='forecast', ... num_forecast_mask_patches=5, ... use_cls_token=True, ... ) >>> model = PatchTSTForPretraining(config) >>> # during training, one provides both past and future values >>> outputs = model(past_values=batch["past_values"]) >>> loss = outputs.loss >>> loss.backward() ```NTrrerrLrgnone reductionrr"rSr)r+r,rIr)r:rkrrwrrMSELossrrrrrIr*r) rFrrerrLrg model_outputx_hatr+loss_val masked_lossrrr1r1r2r]s,E  $ zPatchTSTForPretraining.forward)NNNN)r^r_r`rr>r'rerrdrrfr*r]rgr1r1rGr2rvs&  rvcr)PatchTSTClassificationHeadr:csdt|j|_|j|_tjdd|_|jdkrt|jnt |_ t |j |j |j|_dSNr start_dimr)r=r>r pooling_typerFlattenflattenrsrrrrArrm num_targetsrtrprGr1r2r>,s  z#PatchTSTClassificationHead.__init__rucCs|jr|dddddddf}n"|jdkr|jdd}n|jdkr+|jddj}n td|jd||}|||}|S) a[ Parameters: embedding (`torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True, *required*): Embedding from the model Returns: `torch.Tensor` of shape `(bs, num_targets)` Nrrr!r"rpooling operator  is not implemented yet) rrrrvaluesr@rrtrrFrupooled_embeddingrrr1r1r2r]4s    z"PatchTSTClassificationHead.forwardrr1r1rGr2r+srz0 The PatchTST for classification model. csxeZdZdeffdd Ze     ddejdeejdee dee d ee d ee d e e e ffd d Z ZS)PatchTSTForClassificationr:csBt||jrtdd|_t||_t||_| dS)N+Setting `do_mask_input` parameter to False.F) r=r>rbloggerwarningr`rrrwrrprGr1r2r>Vs     z"PatchTSTForClassification.__init__Nr target_valuesrerrLrgrMc Cs|dur|n|jj}|j||||dd}||j}d} |dur)t} | ||} |sC|f|dd} | dur?| f| } | S| } | St| ||j|j dS)ac past_values (`torch.Tensor` of shape `(bs, sequence_length, num_input_channels)`, *required*): Input sequence to the model target_values (`torch.Tensor`, *optional*): Labels associates with the `past_values` past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*): Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in `[0, 1]`: - 1 for values that are **observed**, - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). Examples: ```python >>> from transformers import PatchTSTConfig, PatchTSTForClassification >>> # classification task with two input channel2 and 3 classes >>> config = PatchTSTConfig( ... num_input_channels=2, ... num_targets=3, ... context_length=512, ... patch_length=12, ... stride=12, ... use_cls_token=True, ... ) >>> model = PatchTSTForClassification(config=config) >>> # during inference, one only provides past values >>> past_values = torch.randn(20, 512, 2) >>> outputs = model(past_values=past_values) >>> labels = outputs.prediction_logits ```NTrxrr)r+r3rIr) r:rkrrwrrCrossEntropyLossr2rIr) rFrrrerrLrgr~y_hatrr+rr1r1r2r]ds2,  z!PatchTSTForClassification.forwardrp)r^r_r`rr>rr'rerrdrrfr2r]rgr1r1rGr2rPs. rz, The PatchTST for regression Model. cs8eZdZd dedeffdd ZdejfddZZ S) PatchTSTPredictionHeadNr:rcsDt|j|_|j|_|j|_|j|_|js|jr|j}n|j|}|jsvt|_ t|_ t|_ t |jD]8}|j tjdd|durW|j t||jn |j |||j |jdkrnt|jntq;dStjdd|_|durt||j|_n|||_|jdkrt|jnt|_dS)a num_patches (`int`): The number of patches in the input sequence. distribution_output (`DistributionOutput`, *optional*): The distribution output layer for probabilistic forecasting. If None, a linear output layer is used. r!rNr)r=r>share_projectionrrrrmrr projectionsdropoutsflattensrrrrAprediction_lengthget_parameter_projectionrsrrr projectionr)rFr:rdistribution_outputr?rrGr1r2r>s0      (  $zPatchTSTPredictionHead.__init__rucCs|jr|dddddddf}n|jdkr|jdd}n|jdkr+|jddj}n|}|jseg}t|jD]%}|j||dd|ddf}|j ||}|j ||}| |q7t j |dd}n||}||}||}t|trtdd |D}|S|dd}|S) aj Parameters: embedding (`torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True, *required*): Embedding from the model Returns: `torch.Tensor` of shape `(bs, forecast_len, num_channels)` Nrrr!r"rrcss|] }|ddVqdS)r!rN)r))rzr1r1r2ri rjz1PatchTSTPredictionHead.forward..)rrrrrrrrrrrrr'rrrrrrfr))rFrurrrrr1r1r2r]s.          zPatchTSTPredictionHead.forwardr) r^r_r`rrbr>r'rer]rgr1r1rGr2rs+rz, The PatchTST for prediction model. cseZdZdeffdd Z     ddejdeejdeejdeed eed eed e e e ffd d Z e ddejdeejd efddZZS)PatchTSTForPredictionr:cst||jrtdd|_t||_|jdkrd|_n/|jdkr,t |j d|_n"|jdkr9t |j d|_n|jdkrFt |j d|_nt d|jt||jjj|jd |_|dS) NrFmse student_tr"normalnegative_binomialUnknown distribution output )r)r=r>rbrrr`rr+rrrrrr@rrarrwrrprGr1r2r>s$        zPatchTSTForPrediction.__init__NrrerfrrLrgrMcCs|dur|n|jj}|j||||dd}||j}d} |jr"|} n||j|j} |durQ|jrF|jj||j|jd} t | |} t | } n t j dd} | | |} |j} |j}|sq| f|dd}| durm| f|}|S|}|St | | |j|j| |d S) aV Parameters: past_values (`torch.Tensor` of shape `(bs, sequence_length, num_input_channels)`, *required*): Input sequence to the model past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*): Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in `[0, 1]`: - 1 for values that are **observed**, - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). future_values (`torch.Tensor` of shape `(bs, forecast_len, num_input_channels)`, *optional*): Future target values associated with the `past_values` output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers output_attentions (`bool`, *optional*): Whether or not to return the output attention of all layers return_dict (`bool`, *optional*): Whether or not to return a `ModelOutput` instead of a plain tuple. Returns: `PatchTSTForPredictionOutput` or tuple of `torch.Tensor` (if `return_dict`=False or `config.return_dict`=False) Examples: ```python >>> from huggingface_hub import hf_hub_download >>> import torch >>> from transformers import PatchTSTConfig, PatchTSTForPrediction >>> file = hf_hub_download( ... repo_id="hf-internal-testing/etth1-hourly-batch", filename="train-batch.pt", repo_type="dataset" ... ) >>> batch = torch.load(file) >>> # Prediction task with 7 input channels and prediction length is 96 >>> model = PatchTSTForPrediction.from_pretrained("namctin/patchtst_etth1_forecast") >>> # during training, one provides both past and future values >>> outputs = model( ... past_values=batch["past_values"], ... future_values=batch["future_values"], ... ) >>> loss = outputs.loss >>> loss.backward() >>> # during inference, one only provides past values, the model outputs future values >>> outputs = model(past_values=batch["past_values"]) >>> prediction_outputs = outputs.prediction_outputs ```NTrxr&r'rrzrr)r+r1rIrr&r')r:rkrrwrrr'r& distributionr9rCrr}r0rIr)rFrrerfrrLrgr~rr y_hat_outrr+r&r'rr1r1r2r]4sL=      zPatchTSTForPrediction.forwardcsr|jj}||d|dd}|jr.|jj|j|j|jdfddt|D}tj |dd}n|j d}t |d S) a  Generate sequences of sample predictions from a model with a probability distribution head. Parameters: past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_input_channels)`): Past values of the time series that serves as context in order to predict the future. past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*): Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in `[0, 1]`: - 1 for values that are **observed**, - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). Return: [`SamplePatchTSTOutput`] where the outputs `sequences` tensor will have shape `(batch_size, number of samples, prediction_length, 1)` or `(batch_size, number of samples, prediction_length, num_input_channels)` for multivariate predictions. NF)rrfrerrcg|]}qSr1samplerrr1r2rrz2PatchTSTForPrediction.generate..rr"r5) r:num_parallel_samplesrrr1r&r'rr'rrr4rFrrerrsamplesr1rr2generates   zPatchTSTForPrediction.generaterpr)r^r_r`rr>r'rerrdrrfr0r]no_gradr4rrgr1r1rGr2rs>   mrcs8eZdZdZd deffdd ZdejfddZZ S) PatchTSTRegressionHeadz Regression head Nr:cst|j|_|j|_|j|_||_|j|j}t j dd|_ |j dkr,t |j nt |_|dur?t ||j|_dS|||_dSr)r=r> output_rangey_rangerrrrrmrrrrsrrrrArrr)rFr:rr?rGr1r2r>s   zPatchTSTRegressionHead.__init__rucCs|jr|dddddddf}n"|jdkr|jdd}n|jdkr+|jddj}n td|jd|||}||}|j du|j du@r_t ||j d |j d|j d}|S) aY Parameters: embedding (`torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True, *required*): Embedding from the model Returns: `torch.Tensor` of shape `(bs, output_dim)` Nrrr!r"rrrr) rrrrrr@rrrrrr'sigmoidrr1r1r2r]s    (zPatchTSTRegressionHead.forwardrrsr1r1rGr2rsrz, The PatchTST for regression model. cseZdZdeffdd Ze     ddejdeejdeejdee d ee d ee d e e e ffd d Z e ddejdeejd efddZZS)PatchTSTForRegressionr:cst||jrtdd|_t||_|jdkrd|_n/|jdkr,t |j d|_n"|jdkr9t |j d|_n|jdkrFt |j d|_nt d|jt||j|_|dS) NrFrrr"rrr)r=r>rbrrr`rr+rrrrrr@rrwrrprGr1r2r>s        zPatchTSTForRegression.__init__NrrrerrLrgrMc s|dur|njj}j||||dd}|j}d} |durIjr>j|} tfdd|D}t| |} t | } n t j dd} | ||} |sc|f|dd } | dur_| f| } | S| } | St | ||j |jd S) a# past_values (`torch.Tensor` of shape `(bs, sequence_length, num_input_channels)`, *required*): Input sequence to the model target_values (`torch.Tensor` of shape `(bs, num_input_channels)`): Target values associates with the `past_values` past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*): Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in `[0, 1]`: - 1 for values that are **observed**, - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). Whether or not to return a `ModelOutput` instead of a plain tuple. Examples: ```python >>> from transformers import PatchTSTConfig, PatchTSTForRegression >>> # Regression task with 6 input channels and regress 2 targets >>> model = PatchTSTForRegression.from_pretrained("namctin/patchtst_etth1_regression") >>> # during inference, one only provides past values, the model outputs future values >>> past_values = torch.randn(20, 512, 6) >>> outputs = model(past_values=past_values) >>> regression_outputs = outputs.regression_outputs ```NTrxc3s |] }|djjVqdS)rN)r,r:r)ritemrFr1r2ri_sz0PatchTSTForRegression.forward..rrzrr)r+r/rIr)r:rkrrwrrrrfr9rCrr}r.rIr) rFrrrerrLrgr~rr+rrr1rr2r])s<%      zPatchTSTForRegression.forwardcsb|jj}||d|dd}|j|jfddt|D}tj|ddd||jj }t |d S) a Generate sequences of sample predictions from a model with a probability distribution head. Parameters: past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_input_channels)`): Past values of the time series that serves as context in order to predict the future. past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*): Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in `[0, 1]`: - 1 for values that are **observed**, - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros). Return: [`SamplePatchTSTOutput`] where the outputs `sequences` tensor will have shape `(batch_size, number of samples, num_targets)`. NF)rrrercrr1rrrr1r2rrz2PatchTSTForRegression.generate..rr"rr) r:rrrr/rr'rr,rr4rr1rr2rss zPatchTSTForRegression.generaterpr)r^r_r`rr>rr'rerrdrrfr.r]rr4rrgr1r1rGr2r s@  Ir)r`rrrvrr)NrN)NFrrrr#)Lrar  dataclassesrtypingrrrr'r activationsrmodeling_flash_attention_utilsr modeling_outputsr modeling_utilsr r processing_utilsr time_series_utilsrrrutilsrrrconfiguration_patchtstr get_loggerr^rrrercr3r4rilistrdrbrrrrrrrrrr%r*r.r0r2r4 distributions Distributionr9rCrErRr]r^r`rqrvrrrrrr__all__r1r1r1r2s          X =  D0< %$8> "$7po%W`=7