a hJ @sdZddlZddlmZddlmZmZmZddlZddlm Z ddl m Z ddl m Z dd lmZdd lmZmZdd lmZdd lmZmZmZdd lmZmZmZddlmZee Z!dfe j"ej#ej#ej#eej#ee$e$eej#dddZ%Gddde j"Z&Gddde j"Z'dgej#e$ee(e)e*dddZ+dhej#ee(e*fee(e*dddZ,Gdd d e j"Z-Gd!d"d"e j"Z.Gd#d$d$e j"Z/eGd%d&d&eZ0Gd'd(d(e j"Z1Gd)d*d*e j"Z2Gd+d,d,e0Z3eed-d.Gd/d0d0eZ4eed1d.Gd2d3d3eZ5eed4d.Gd5d6d6eZ6eed7d.Gd8d9d9eZ7eed:d.Gd;d<dd?d?eZ9ej:j;ej#ej#d@dAdBZGdHdIdIe j"Z?GdJdKdKe j"Z@GdLdMdMe j"ZAeGdNdOdOe0ZBGdPdQdQe j"ZCedRd.GdSdTdTe0ZDGdUdVdVe j"ZEedWd.GdXdYdYe0ZFedZd.Gd[d\d\e j"ZGed]d.Gd^d_d_e0ZHGd`dadae j"ZIedbd.Gdcdddde0ZJgdeZKdS)jzPyTorch PatchTST model.N) dataclass)CallableOptionalUnion)nn)ACT2CLS)FlashAttentionKwargs)BaseModelOutput)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)NegativeBinomialOutput NormalOutputStudentTOutput) ModelOutputauto_docstringlogging)PatchTSTConfig)modulequerykeyvalueattention_maskscalingdropout head_maskc Ks|dur|dd}t||dd|} |dur>| |} tjj| dd} |durj| |dddd} tjj| ||j d} t| |} | dd } | | fS)Nrdimr)ptraining) sizetorchmatmul transposerZ functionalZsoftmaxviewrr% contiguous) rrrrrrrrkwargs attn_weights attn_outputr/j/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/patchtst/modeling_patchtst.pyeager_attention_forward&s  r1c seZdZdZd eeeeeeeedfdd Z d e j ee j ee j ee j eee e ee j ee j eee j fd d d ZZS)PatchTSTAttentionz=Multi-headed attention from 'Attention Is All You Need' paperrFTN) embed_dim num_headsr is_decoderbias is_causalconfigcst||_||_||_|||_||_|j||jkrTtd|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 r6)super__init__r3r4rhead_dimr8 ValueErrorrr5r7rLineark_projv_projq_projout_proj)selfr3r4rr5r6r7r8 __class__r/r0r<Hs&    zPatchTSTAttention.__init__) hidden_stateskey_value_statesrlayer_head_maskoutput_attentionsr,returncKs|du}|jdd\}} |r(|jdn| } || d|jf} || d|jf} ||j| dd} |rh|n|}||j| dd}||j| dd}t}|jj dkrt |jj }||| |||f|j sdn|j |j ||d|\}}||| d}||}||dfS)z#Input shape: Batch x Time x ChannelNrrr!eagerr)rrrJr)shaper=rBr*r)r@rAr1r8Z_attn_implementationr r%rrreshaper+rC)rDrGrHrrIrJr,Zis_cross_attentionZbszZtgt_lenZsrc_lenZ q_input_shapeZkv_input_shapeZ query_statesZcurrent_statesZ key_statesZ value_statesZattention_interfacer.r-r/r/r0forwardgs:      zPatchTSTAttention.forward)rFTFN)NNNF)__name__ __module__ __qualname____doc__intfloatboolrrr<r'Tensorr r tuplerO __classcell__r/r/rEr0r2Es8"r2cs6eZdZdZedfdd ZejdddZZ S)PatchTSTBatchNormzP Compute batch normalization over the sequence length (time) dimension. r8cs"ttj|j|jd|_dS)Neps)r;r<rZ BatchNorm1dd_modelnorm_eps batchnormrDr8rEr/r0r<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)r`)rDrboutputr/r/r0rOs  zPatchTSTBatchNorm.forward rPrQrRrSrr<r'rWrOrYr/r/rEr0rZsrZFrb mask_ratiounmasked_channel_indiceschannel_consistent_masking mask_valuecCs,|dks|dkr td|d|j\}}}}|j} t|d|} |rjtj|d|| d} | d|d} ntj|||| 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>rMrkrTr'ZrandrepeatZonesZargsortZgather unsqueeze masked_fillrV)rbrfrgrhri batch_size num_channelssequence_length num_featuresrkZlen_keepnoisemaskZ ids_shuffleZ ids_restore inputs_maskr/r/r0random_maskings& rxrbnum_forecast_mask_patchesrgricCst|tr|g}dd|D}|j\}}}}tj||||jd} g} d} t|} t||D]P\} }| dksr| |krtd| dt||| }| | ||g| |7} qZt | ddd } | |kr| dd || | dd <n&| |kr| d d | || d d <d}| D]4\}}}||}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)rr/.0_r/r/r0 z$forecast_masking..rjrznum_forecast_mask_patches z6 should be greater than 0 and less than total patches.cSs|dS)Nr!r/)xr/r/r0rz"forecast_masking..)rr!rrNrm) isinstancerTrMr'zerosrksumzipr>appendsortedZrandpermrornrprV)rbrzrgriZforecast_mask_ratiosrqrrrsrtrvZt_list total_length total_ratio patch_lengthratioZtemp_lenZbatch1Z patch_lenr}Zbatch2permrwr/r/r0forecast_maskingsB     rcs6eZdZdZedfdd ZejdddZZ S)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|jkrHtd|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_lengthrsr patch_strider>max num_patchessequence_start)rDr8Znew_sequence_lengthrEr/r0r<9s   zPatchTSTPatchify.__init__) past_valuescCsp|jd}||jkr,td|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 (r9N) dimensionr&step) rMrsr>rZunfoldrrr)r+)rDrrsrcr/r/r0rOJs  zPatchTSTPatchify.forwardrdr/r/rEr0r1srcs6eZdZdZedfdd ZejdddZZ S)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[csTt|j|_|j|_|j|_|j|_|j|_|j|_|jdurPt|j|_dSN) r;r<random_mask_ratiorh mask_typerzrgrirrarEr/r0r<ns  zPatchTSTMasking.__init__ patch_inputcCsr|jdkr*t||j|j|j|jd\}}n8|jdkrPt||j|j|jd\}}ntd|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 randomreZforecastryzInvalid mask type .) rrxrrgrhrirrzr>rV)rDrZ masked_inputrvr/r/r0rOys$    zPatchTSTMasking.forwardrdr/r/rEr0ras  rcs>eZdZdZedfdd Zd ejee dddZ Z S) PatchTSTEncoderLayerz PatchTST encoder layer r[c st|j|_t|j|j|j|d|_|jdkr@t |jnt |_ |j dkr`t||_n0|j dkrt j|j|jd|_nt|j d|jr|jdkrt |jnt |_|j dkrt||_n0|j dkrt j|j|jd|_nt|j dt t j|j|j|jdt|j|jdkr6t |jnt t j|j|j|jd|_|jdkrnt |jnt |_|j dkrt||_n2|j dkrt j|j|jd|_nt|j d|j|_dS)N)r3r4rr8rr`Z layernormr\z$ is not a supported norm layer type.r:) r;r<channel_attentionr2r^Znum_attention_headsZattention_dropout self_attnZ path_dropoutrDropoutIdentity dropout_path1Z norm_typerZnorm_sublayer1 LayerNormr_r> dropout_path2norm_sublayer2Z Sequentialr?Zffn_dimr6rZactivation_functionZ ff_dropoutff dropout_path3norm_sublayer3pre_normrarEr/r0r<sD          "   zPatchTSTEncoderLayer.__init__N hidden_staterJc Cs|j\}}}}|||||}|jrP|j|||d\}}} |||}n(|j||d\}}} ||||}|||||}|jr*|dd }|||||}|jr|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)` )rGrJr!r)rMr*rrrrrNrr)r+rrrrr) rDrrJrqnum_input_channelsrsr^r.r-r}Zchannel_attn_weightsoutputsr/r/r0rOsF     zPatchTSTEncoderLayer.forward)N) rPrQrRrSrr<r'rWrrVrOrYr/r/rEr0rs2rc@s<eZdZUeed<dZdZdZej dddZ d dd Z d S) PatchTSTPreTrainedModelr8modelrF)rcCst|trdt|jj|jj|jj|jjd}|jjrRtj j |j dd|d7}| |j||_ nt|tjr|jj|jjdnbt|tr|jjj|jjjdn8t|tjr|jjj d|jjd|jdur|jjdS)z$ Initialize weights rg{Gz?)std?r)meanrN)rPatchTSTPositionalEncodingrr8rrr use_cls_tokenrinitZnormal_ cls_token_init_pe position_encrr6dataZzero_weightZfill_rZr`r?Zinit_std)rDrrr/r/r0 _init_weights/s(      z%PatchTSTPreTrainedModel._init_weightscCst|tr||_dSr)rPatchTSTEncodergradient_checkpointing)rDrrr/r/r0_set_gradient_checkpointingIs z3PatchTSTPreTrainedModel._set_gradient_checkpointingN)F) rPrQrRr__annotations__Zbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingrModulerrr/r/r/r0r(s rcs2eZdZedfdd ZejdddZZS)PatchTSTEmbeddingr[csjt|j|_|j|_|jr4t|j|j|_n2t |_t |jD]}|j t|j|jqHdSr) r;r<rshare_embeddingrr?rr^input_embedding ModuleListranger)rDr8r}rEr/r0r<Os  zPatchTSTEmbedding.__init__rcshjd}|jkr,tdjd|djr>}n&fddt|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)rr|irrDr/r0r~mrz-PatchTSTEmbedding.forward..r")rMrr>rrrr'stack)rDrrZ embeddingsr/rr0rO[s    zPatchTSTEmbedding.forward rPrQrRrr<r'rWrOrYr/r/rEr0rNs rcsPeZdZdZeedfdd Zeeeej dddZ e j dd d Z ZS) rz' Class for positional encoding r8rcstt|j|_|j|_|jrBttddd|j|_ |d7}| |||_ |j dkrft |j nt|_ dS)Nrr)r;r<rrr Parameterr'rr^rrrpositional_dropoutrrrDr8rrEr/r0r<ws z#PatchTSTPositionalEncoding.__init__)r8rrKcCs|jdkr$tjt||jdd}n|jdkrt||j}td|d}t td|jdt d|j }t |||dddddf<t |||dddddf<||}||d }tj|d d}nt|jd |S) NrTZ requires_gradZsincosrrr!g@ FzN is not a valid positional encoder. Available types are 'random' and 'sincos'.)Zpositional_encoding_typerrr'Zrandnr^rZarangeroexpmathlogsincosrrr>)r8rrpositionZdiv_termr/r/r0rs  (    z#PatchTSTPositionalEncoding._init_percCs|jrn|||jddddf}|j|jddddf}||jd|jdd}tj||fdd}n|||j}|S)Nrrrr!r") rrrrexpandrMrr'cat)rDrrZ cls_tokensrr/r/r0rOs z"PatchTSTPositionalEncoding.forward)rPrQrRrSrrTr< staticmethodrrrr'rWrOrYr/r/rEr0rrs rcsHeZdZdZeedfdd Zd eje e e e e dddZ Z S) rz PatchTST Encoder rcsTtd|_t|_t||_tfddt j D|_ | dS)NFcsg|] }tqSr/)rrr[r/r0r~rz,PatchTSTEncoder.__init__..) r;r<rrembedderrpositional_encoderrrrZnum_hidden_layerslayers post_initrrEr[r0r<s     zPatchTSTEncoder.__init__N)routput_hidden_statesrJrKc Cs|dur |n|jj}|dur |n|jj}||}||}|rDdnd}|rPdnd}|jD]8}|rl||f}|||d}|d}|rZ||df}qZt|||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` Nr/rrr)last_hidden_staterG attentions)r8rJrrrrr ) rDrrrJrencoder_statesZall_attentionsZ encoder_layerZ layer_outputsr/r/r0rOs       zPatchTSTEncoder.forward)NN)rPrQrRrSrrTr<r'rWrrVr rOrYr/r/rEr0rsrzG Base class for model's outputs, with potential hidden states. )Z 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 NrrGrrvlocscaler)rPrQrRrSrrr' FloatTensorrrGrXrrvrrrr/r/r/r0rs rz4 Output type of [`PatchTSTForPretraining`]. c@sbeZdZUdZdZeejed<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_outputrGr) rPrQrRrSrrr'rrrrGrXrr/r/r/r0r s rz3 Output type of [`PatchTSTForRegression`]. c@sbeZdZUdZdZeejed<dZ eejed<dZ ee ejed<dZ ee ejed<dS)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. Nrregression_outputsrGr) rPrQrRrSrrr'rrrrGrXrr/r/r/r0rs rz3 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 Nrprediction_outputsrGrrr)rPrQrRrSrrr'rrrrGrXrrrr/r/r/r0r1s rz7 Output type of [`PatchTSTForClassification`]. c@sbeZdZUdZdZeejed<dZ eejed<dZ ee ejed<dZ ee ejed<dS)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). Nrprediction_logitsrGr) rPrQrRrSrrr'rrrrGrXrr/r/r/r0rQs rz 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) rPrQrRrSrrr'rrr/r/r/r0rfs r)inputtargetrKcCs || S)zc Computes the negative log likelihood loss from input distribution with respect to target. )Zlog_prob)rrr/r/r0nllwsr) input_tensorweightsrKcCsr|durbt|dk||t|}tj|r8|j|dn|dd}|rV|j|dn||S|j|dSdS)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)rrr#Zweighted_tensorZ sum_weightsr/r/r0weighted_averages "rcsLeZdZdZedfdd ZejejeejejejfdddZ 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[csPtt|dr|jnd|_t|dr0|jnd|_t|drF|jnd|_dS)N scaling_dimrkeepdimT minimum_scalegh㈵>)r;r<hasattrrr#rrrarEr/r0r<s zPatchTSTStdScaler.__init__robserved_indicatorrKcCsz|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)`) rrr!)rr#rZ clamp_minr'sqrtr)rDrr denominatorrZvariancerr/r/r0rOs  "zPatchTSTStdScaler.forward rPrQrRrSrr<r'rWrXrOrYr/r/rEr0rs rcsLeZdZdZedfdd ZejejeejejejfdddZ Z S)PatchTSTMeanScalerz Computes a scaling factor as the weighted average absolute value along the first dimension, and scales the data accordingly. r[csftt|dr|jnd|_t|dr0|jnd|_t|drF|jnd|_t|dr\|jnd|_dS)NrrrTr绽|= default_scale)r;r<rrr#rrrrarEr/r0r<s  zPatchTSTMeanScaler.__init__rc Cs||j|jdd}|j|jdd}|tj|dd}|jdurt|jdd}tj|ddd}t||}n|jt|}t|dk||}tj||j d}||} |j s|j|jd}| t ||fS)r Tr rrNrr") absrr#r'rrZsqueeze ones_likerrrr) rDrrZts_sumZ num_observedrZ batch_sumZbatch_observationsrZ scaled_datar/r/r0rOs  zPatchTSTMeanScaler.forwardr r/r/rEr0rs rcsReZdZdZedfdd Zd ejeeje ejejejfdddZ 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|dr0|jnd|_dS)NrrrT)r;r<rrr#rrarEr/r0r<s zPatchTSTNOPScaler.__init__NrcCsBtj|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#r)r'rrr#rr)rDrrrrr/r/r0rOs zPatchTSTNOPScaler.forward)N) rPrQrRrSrr<r'rWrrXrOrYr/r/rEr0rs rcsHeZdZedfdd ZejejeejejejfdddZZ S)PatchTSTScalerr[csNt|jdks|jdur*t||_n |jdkr@t||_n t||_dS)NrTr)r;r<rrscalerrrrarEr/r0r<s     zPatchTSTScaler.__init__rcCs|||\}}}|||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)rDrrrrr/r/r0rOszPatchTSTScaler.forward) rPrQrRrr<r'rWrXrOrYr/r/rEr0rs rc s`eZdZedfdd Zdejeejeejeeeeeee e e fdddZ Z S) PatchTSTModelr[csft|t||_t||_|j|_|jj}|jrBt||_ n t |_ t ||d|_ |dS)N)r)r;r<rrr patchifier do_mask_inputrrmaskingrrrencoderrrrEr/r0r<,s     zPatchTSTModel.__init__Nrpast_observed_mask future_valuesrrJ return_dictrKc Cs|dur |n|jj}|dur |n|jj}|dur4|n|jj}|durNt|}|||\}}} ||} |jr| | \} } n| | d} } |j | ||d} |s| 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)rrrJcss|]}|dur|VqdSrr/)r|vr/r/r0 rz(PatchTSTModel.forward..)rrGrrvrrr)r8use_return_dictrJrr'rrrrrrrrGrrXr)rDrrrrrJrZscaled_past_valuesrrZpatched_valuesZ masked_valuesrvZencoder_outputrr/r/r0rO>s66  zPatchTSTModel.forward)NNNNN)rPrQrRrr<r'rWrrVrrXrrOrYr/r/rEr0r*s rcs:eZdZdZedfdd ZejejdddZZ S)PatchTSTMaskPretrainHeadz- Pretraining head for mask modelling r[csHt|jdkr t|jnt|_t|j|j |_ |j |_ dS)Nr) r;r< head_dropoutrrrrr?r^rlinearrrarEr/r0r<s  z!PatchTSTMaskPretrainHead.__init__) embeddingrKcCs:|||}|jr6|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)r$rr)rDr%r/r/r0rOs  z PatchTSTMaskPretrainHead.forwardrdr/r/rEr0r"sr"z* The PatchTST for pretrain model. c sXeZdZedfdd Zdejeejeeeeeee e e fdddZ Z S) PatchTSTForPretrainingr[cs4t|d|_t|d|_t||_|dS)NTr[)r;r<rrrr"headrrarEr/r0r<s    zPatchTSTForPretraining.__init__N)rrrrJrrKc Cs|dur |n|jj}|j||||dd}||j}tjdd}|||j} | jdd|j  |j d} |j } |s|f|d d } | dur| f| n| } | 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() ```NTrrrrJrnoneZ reductionrr"rr)rrrGr)r8r!rr'rrMSELossrrrvrrGrr) rDrrrrJr model_outputZx_hatrloss_valZ masked_lossrrr/r/r0rOs(E   $ zPatchTSTForPretraining.forward)NNNN)rPrQrRrr<r'rWrrVrrXrrOrYr/r/rEr0r&s  r&cs2eZdZedfdd ZejdddZZS)PatchTSTClassificationHeadr[csdt|j|_|j|_tjdd|_|jdkr>t|jnt |_ t |j |j |j|_dSNrZ start_dimr)r;r<r pooling_typerFlattenflattenr#rrrr?rr^ num_targetsr$rarEr/r0r<,s   z#PatchTSTClassificationHead.__init__r%cCs|jr$|dddddddf}nD|jdkr<|jdd}n,|jdkrV|jddj}ntd|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) rr2rrvaluesr>r4r$rrDr%pooled_embeddingrcr/r/r0rO4s    z"PatchTSTClassificationHead.forwardrr/r/rEr0r/+sr/z0 The PatchTST for classification model. c sbeZdZedfdd Zedejeejee ee ee ee e e e fdddZ ZS) PatchTSTForClassificationr[csBt||jr"tdd|_t||_t||_| dS)N+Setting `do_mask_input` parameter to False.F) r;r<rloggerwarningrrr/r'rrarEr/r0r<Vs    z"PatchTSTForClassification.__init__Nr target_valuesrrrJrrKc Cs|dur |n|jj}|j||||dd}||j}d} |durRt} | ||} |s|f|dd} | durz| f| n| } | 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 ```NTr(rr)rrrGr) r8r!rr'rrZCrossEntropyLossrrGr) rDrrArrrJrr-y_hatr.rrr/r/r0rOds.,  z!PatchTSTForClassification.forward)NNNNN)rPrQrRrr<rr'rWrrVrrXrrOrYr/r/rEr0r<Ps  r<z, The PatchTST for regression Model. cs6eZdZdeedfdd ZejdddZZ S) PatchTSTPredictionHeadNrcsFt|j|_|j|_|j|_|j|_|js6|jr>|j}n |j|}|jst|_ t|_ t|_ t |jD]p}|j tjdd|dur|j t||jn|j |||j |jdkrt|jntqvnXtjdd|_|durt||j|_n |||_|jdkr8t|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!r1Nr)r;r<share_projectionrrr2r^rr projectionsdropoutsflattensrrr3r?prediction_lengthget_parameter_projectionr#rrr4 projectionr)rDr8rdistribution_outputr=rrEr/r0r<s.      *  zPatchTSTPredictionHead.__init__r6cCs|jr$|dddddddf}n6|jdkr<|jdd}n|jdkrV|jddj}n|}|jsg}t|jD]J}|j||dd|ddf}|j ||}|j ||}| |qnt j |dd}n||}||}||}t|trtdd |D}n |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))r|zr/r/r0r  rz1PatchTSTPredictionHead.forward..)rr2rrr9rDrrrGrFrErr'rr4rrJrrXr))rDr%r;rcrr/r/r0rOs,          zPatchTSTPredictionHead.forward)N) rPrQrRrrTr<r'rWrOrYr/r/rEr0rCs+rCz, The PatchTST for prediction model. c seZdZedfdd Zd ejeejeejeeeeeee e e fdddZ e d ejeejedd d ZZS) PatchTSTForPredictionr[cst||jr"tdd|_t||_|jdkr>d|_n^|jdkrXt |j d|_nD|jdkrrt |j d|_n*|jdkrt |j d|_nt d|jt||jjj|jd |_|dS) Nr=Fmse student_tr"normalnegative_binomialUnknown distribution output )rK)r;r<rr>r?rrrrKrrHrrr>rCrrr'rrarEr/r0r<s$       zPatchTSTForPrediction.__init__NrcCs|dur |n|jj}|j||||dd}||j}d} |jrD|} n||j|j} |dur|jr|jj||j|jd} t | |} t | } nt j dd} | | |} |j} |j}|s| f|dd}| dur| f|n|}|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 ```NTr(rrrr*rr)rrrGrrr)r8r!rr'rrKrr distributionrrrr,rrGr)rDrrrrrJrr-rBr.Z y_hat_outrTrrrrr/r/r0rO4sH=      zPatchTSTForPrediction.forwardrrrKcsr|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)rrrrrScsg|] }qSr/sampler{rTr/r0r~rz2PatchTSTForPrediction.generate..rr"r) r8num_parallel_samplesrKrTrrrrr'rrorrDrrrZrZsamplesr/rXr0generates  zPatchTSTForPrediction.generate)NNNNN)N)rPrQrRrr<r'rWrrVrrXrrOno_gradrr\rYr/r/rEr0rMs,  mrMcs8eZdZdZd edfdd ZejdddZZ S) PatchTSTRegressionHeadz Regression head Nr[cst|j|_|j|_|j|_||_|j|j}t j dd|_ |j dkrXt |j nt |_|dur|t ||j|_n |||_dSr0)r;r<Z output_rangey_rangerr2rKrr^rr3r4r#rrrr?r5rJrI)rDr8rKr=rEr/r0r<s   zPatchTSTRegressionHead.__init__r6cCs|jr$|dddddddf}nD|jdkr<|jdd}n,|jdkrV|jddj}ntd|jd|||}||}|j du|j du@rt ||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"rr7r8r) rr2rrr9r>rr4rJrKr_r'Zsigmoidr:r/r/r0rOs    (zPatchTSTRegressionHead.forward)Nrdr/r/rEr0r^sr^z, The PatchTST for regression model. c seZdZedfdd Zed ejeejeejee ee ee e e e fdddZ ed ejeejedd d ZZS) PatchTSTForRegressionr[cst||jr"tdd|_t||_|jdkr>d|_n^|jdkrXt |j d|_nD|jdkrrt |j d|_n*|jdkrt |j d|_nt d|jt||j|_|dS) Nr=FrNrOr"rPrQrR)r;r<rr>r?rrrrKrr5rrr>r^r'rrarEr/r0r<s        zPatchTSTForRegression.__init__Nr@c s|dur |njj}j||||dd}|j}d} |durjr|j|} tfdd|D}t| |} t | } nt j dd} | ||} |s|f|dd } | dur| f| n| } | 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 ```NTr(c3s|]}|djjVqdS)rN)r*r8r5)r|itemrDr/r0r _rz0PatchTSTForRegression.forward..rr*rr)rrrGr)r8r!rr'rrKrTrXrrrr,rrGr) rDrrArrrJrr-rBrrTrr/rbr0rO)s8%      zPatchTSTForRegression.forwardrUcsb|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)rrArrcsg|] }qSr/rVr{rXr/r0r~rz2PatchTSTForRegression.generate..rr"rrY) r8rZrKrTrrr'rr*r5rr[r/rXr0r\sszPatchTSTForRegression.generate)NNNNN)N)rPrQrRrr<rr'rWrrVrrXrrOr]rr\rYr/r/rEr0r` s. Ir`)rrrMr&r`r<)NrN)NFr)Nr)NN)LrSr dataclassesrtypingrrrr'rZ activationsrZmodeling_flash_attention_utilsr Zmodeling_outputsr Zmodeling_utilsr r Zprocessing_utilsr Ztime_series_utilsrrrutilsrrrZconfiguration_patchtstrZ get_loggerrPr>rrWrUr1r2rZlistrVrTrxrrrrrrrrrrrrrr distributions Distributionrrrrrrrr"r&r/r<rCrMr^r`__all__r/r/r/r0s         X =  D0< %$8> $7po%W`=7