o wiQ@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/sommelier/.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 e j dee j dee j dede ede e j ee j ee e 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_statespast_key_valuerlayer_head_maskoutput_attentionsr.returncKs|du}|jdd\} } |r|jdn| } | | d|jf} | | d|jf} ||j| dd}|rK|durK|djd|jdkrK|d}|d}nf|rf||j| dd}||j| dd}nK|dur||j| dd}||j| dd}tj|d|gdd}tj|d|gdd}n||j| dd}||j| dd}|j r||f}t }|j j dkrt |j j }||||||f|jsdn|j|j||d |\}}|| | d}||}|||fS) z#Input shape: Batch x Time x ChannelNrrr!rr"eagerr)rrrMr)shaper?rDr,r)rBrCr'catr7r3r:_attn_implementationr r%rrreshaper-rE)rFrIrJrKrrLrMr.is_cross_attentionbsztgt_lensrc_len q_input_shapekv_input_shape query_states key_states value_statesattention_interfacer0r/r1r1r2forwardgsX      zPatchTSTAttention.forward)rFTFN)NNNNF)__name__ __module__ __qualname____doc__intfloatboolrrr>r'Tensortupler r r^ __classcell__r1r1rGr2r4EsX"   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)rp)rFrroutputr1r1r2r^s   zPatchTSTBatchNorm.forward r_r`rarbrr>r'rfr^rhr1r1rGr2rjsrjFrr 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@rPrzrcr'randrepeatonesargsortgather unsqueeze masked_fillre)rrrurvrwrx batch_size num_channelssequence_length num_featuresrzlen_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..ryrznum_forecast_mask_patches z6 should be greater than 0 and less than total patches.cSs|dS)Nr!r1)xr1r1r2;sz"forecast_masking..)rr!rrNr|) isinstancercrPr'zerosrzsumzipr@appendsortedrandpermrr~rre)rrrrvrxforecast_mask_ratiosrrrrrt_list total_length total_ratio patch_lengthratiotemp_lenbatch1 patch_lenrbatch2permrr1r1r2forecast_maskingsB    rcri)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>[s   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) rPrr@runfoldrrr)r-)rFrrrsr1r1r2r^ls zPatchTSTPatchify.forwardrtr1r1rGr2rSsrcri)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_ratiorw mask_typerrvrxrrqrGr1r2r>s  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)rrrurvrwrxforecast)rrrrvrxzInvalid mask type .) rrrrvrwrxrrr@re)rFr masked_inputrr1r1r2r^s$   zPatchTSTMasking.forwardrtr1r1rGr2rs  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:rrp layernormrkz$ is not a supported norm layer type.r<) r=r>channel_attentionr4rnnum_attention_headsattention_dropout self_attn path_dropoutrDropoutIdentity dropout_path1 norm_typerjnorm_sublayer1 LayerNormror@ dropout_path2norm_sublayer2 SequentialrAffn_dimr8ractivation_function ff_dropoutff dropout_path3norm_sublayer3pre_normrqrGr1r2r>sD               zPatchTSTEncoderLayer.__init__N hidden_staterMc 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)` )rIrMr!r)rPr,rrrrrSrr)r-rrrrr) rFrrMrnum_input_channelsrrnr0r/rchannel_attn_weightsoutputsr1r1r2r^sF       zPatchTSTEncoderLayer.forwardr) r_r`rarbrr>r'rfrrer^rhr1r1rGr2rs"2rc@s.eZdZeZdZdZdZddZd ddZ dS) PatchTSTPreTrainedModelmodelrFcCst|tr&|jjrtjj|jdd|jjdkr$tjj|j ddddSdSt|tj r;|j j |jj ddSt|trQ|jj j |jjj ddSt|tjtjfrr|jj jd|jjd|j durt|j j dSdSdS) z$ Initialize weights g{Gz?)stdrrg?)meanr?N)rPatchTSTPositionalEncodingr: use_cls_tokenrinitnormal_ cls_tokenpositional_encoding_type position_encrr8datazero_weightfill_rjrprAConv1dinit_std)rFrr1r1r2 _init_weightsQs$      z%PatchTSTPreTrainedModel._init_weightscCst|tr ||_dSdSr)rPatchTSTEncodergradient_checkpointing)rFrrr1r1r2_set_gradient_checkpointinggs  z3PatchTSTPreTrainedModel._set_gradient_checkpointingN)F) r_r`rar config_classbase_model_prefixmain_input_namesupports_gradient_checkpointingrrr1r1r1r2rJsrc2eZdZdeffdd ZdejfddZZS)PatchTSTEmbeddingr:cslt|j|_|j|_|jrt|j|j|_dSt |_t |jD]}|j t|j|jq%dSr) r=r>rshare_embeddingrrArrninput_embedding ModuleListranger)rFr:rrGr1r2r>ms  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)rrirrFr1r2rs2z-PatchTSTEmbedding.forward..r")rPrr@rrrr'stack)rFrr embeddingsr1rr2r^ys   zPatchTSTEmbedding.forward r_r`rarr>r'rfr^rhr1r1rGr2rls 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'rrnr_init_perpositional_dropoutrrrFr:rrGr1r2r>s z#PatchTSTPositionalEncoding.__init__rNcCs|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'.)rrrr'randnrnrarangerexpmathlogsincosrrr@)r:rrpositiondiv_termr1r1r2rs  (    z#PatchTSTPositionalEncoding._init_percCs|jr8|||jddddf}|j|jddddf}||jd|jdd}tj||fdd}|S|||j}|S)Nrrrr!r") rrrrexpandrPrr'rQ)rFrr cls_tokensrr1r1r2r^s z"PatchTSTPositionalEncoding.forward)r_r`rarbrrcr> staticmethodrrrr'rfr^rhr1r1rGr2rs 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_statesrMrNc 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)rrMrr)last_hidden_staterI attentions)r:rMrrrrr ) rFrrrMrencoder_statesall_attentions encoder_layer layer_outputsr1r1r2r^s        zPatchTSTEncoder.forwardNN)r_r`rarbrrcr>r'rfrrer r^rhr1r1rGr2rsrzG 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`rarbrrr' FloatTensor__annotations__rIrgrrr&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`rarbr,rr'r(r)r-rIrgrr1r1r1r2r+' 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`rarbr,rr'r(r)r0rIrgrr1r1r1r2r/;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`rarbr,rr'r(r)r2rIrgrr&r'r1r1r1r2r1Os r1z7 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`rarbr,rr'r(r)r4rIrgrr1r1r1r2r3os r3z 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`rarbr6rr'r(r)r1r1r1r2r5s r5inputtargetrNcCs || S)zc Computes the negative log likelihood loss from input distribution with respect to target. )log_prob)r7r8r1r1r2nlls r: 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 "  rDc 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>hasattrrGr#rHrIrqrGr1r2r>s  zPatchTSTStdScaler.__init__robserved_indicatorrNcCsz|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)`) rHrr!)rr#rH clamp_minr'sqrtrI)rFrrK denominatorr&variancer'r1r1r2r^s  "zPatchTSTStdScaler.forward r_r`rarbrr>r'rfrgr^rhr1r1rGr2rFsrFc rE) 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)NrGrrHTrI绽|= default_scale)r=r>rJrGr#rHrIrUrqrGr1r2r>s  zPatchTSTMeanScaler.__init__rrKrNc 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)rLTrMrr=Nrr") absrr#r'rArUsqueeze ones_liker?rIrHr@) rFrrKts_sum num_observedr' batch_sumbatch_observationsrU scaled_datar1r1r2r^s  zPatchTSTMeanScaler.forwardrRr1r1rGr2rSsrSc 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)NrGrrHT)r=r>rJrGr#rHrqrGr1r2r>s  zPatchTSTNOPScaler.__init__NrrKrNcCsBtj|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#rH)r'rXrr#rHr@)rFrrKr'r&r1r1r2r^s  zPatchTSTNOPScaler.forwardr) r_r`rarbrr>r'rfrrgr^rhr1r1rGr2r^sr^c sLeZdZdeffdd ZdejdejdeejejejffddZZ S) PatchTSTScalerr:csRt|jdks|jdurt||_dS|jdkr"t||_dSt||_dS)NrTr)r=r>rrSscalerrFr^rqrGr1r2r>-s  zPatchTSTScaler.__init__rrKrNcCs|||\}}}|||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`)rFrrKr&r'r1r1r2r^6s zPatchTSTScaler.forward) r_r`rarr>r'rfrgr^rhr1r1rGr2r_,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>Js      zPatchTSTModel.__init__Nrpast_observed_mask future_valuesrrM return_dictrNc 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)rrrMcss|] }|dur|VqdSrr1)rvr1r1r2 z(PatchTSTModel.forward..)rrIrrr&r'r)r:use_return_dictrMrr'rXr`rbrcrdrerrIrrgr%)rFrrfrgrrMrhscaled_past_valuesr&r'patched_values masked_valuesrencoder_outputrr1r1r2r^\s66  zPatchTSTModel.forwardNNNNN)r_r`rarr>r'rfrrerrgr%r^rhr1r1rGr2raHs, racs<eZdZdZdeffdd ZdejdejfddZZ S) PatchTSTMaskPretrainHeadz- Pretraining head for mask modelling r:csHt|jdkrt|jnt|_t|j|j |_ |j |_ dSNr) r=r> head_dropoutrrrrrArnrlinearrrqrGr1r2r>s   z!PatchTSTMaskPretrainHead.__init__ embeddingrNcCs:|||}|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)rurr)rFrvr1r1r2r^s  z PatchTSTMaskPretrainHead.forwardrtr1r1rGr2rrsrrz* 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>rcrarrrheadrrqrGr1r2r>s    zPatchTSTForPretraining.__init__NrrfrrMrhrNc 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() ```NTrrfrrMrhnone reductionrr"rTr)r,r-rIr)r:rlrrxrrMSELossrrrrrIr+r) rFrrfrrMrh model_outputx_hatr,loss_val masked_lossrrr1r1r2r^s,E  $ zPatchTSTForPretraining.forward)NNNN)r_r`rarr>r'rfrrerrgr+r^rhr1r1rGr2rws&  rwcr)PatchTSTClassificationHeadr:csdt|j|_|j|_tjdd|_|jdkrt|jnt |_ t |j |j |j|_dSNr start_dimr)r=r>r pooling_typerFlattenflattenrtrrrrArrn num_targetsrurqrGr1r2r>Js  z#PatchTSTClassificationHead.__init__rvcCs|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@rrurrFrvpooled_embeddingrsr1r1r2r^Rs    z"PatchTSTClassificationHead.forwardrr1r1rGr2rIsrz0 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>rcloggerwarningrarrrxrrqrGr1r2r>ts     z"PatchTSTForClassification.__init__Nr target_valuesrfrrMrhrNc 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 ```NTryrr)r,r4rIr) r:rlrrxrrCrossEntropyLossr3rIr) rFrrrfrrMrhry_hatrr,rr1r1r2r^s2,  z!PatchTSTForClassification.forwardrq)r_r`rarr>rr'rfrrerrgr3r^rhr1r1rGr2rns. 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_projectionrrrrnrr projectionsdropoutsflattensrrrrAprediction_lengthget_parameter_projectionrtrrr projectionr)rFr:rdistribution_outputr?rrGr1r2r>s0      (  $zPatchTSTPredictionHead.__init__rvcCs|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))rzr1r1r2rj)rkz1PatchTSTPredictionHead.forward..)rrrrrrrrrrrrr'rrrrrrgr))rFrvrrsrr1r1r2r^s.          zPatchTSTPredictionHead.forwardr) r_r`rarrcr>r'rfr^rhr1r1rGr2rs+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>rcrrrarr,rrrrrr@rrbrrxrrqrGr1r2r>5s$        zPatchTSTForPrediction.__init__NrrfrgrrMrhrNcCs|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 ```NTryr&r'rr{rr)r,r2rIrr&r')r:rlrrxrrr'r& distributionr:rDrr~r1rIr)rFrrfrgrrMrhrrr y_hat_outrr,r&r'rr1r1r2r^RsL=      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)rrgrfrrcg|]}qSr1samplerrr1r2rrz2PatchTSTForPrediction.generate..rr"r6) r:num_parallel_samplesrrr2r&r'rr'rrr5rFrrfrrsamplesr1rr2generates   zPatchTSTForPrediction.generaterqr)r_r`rarr>r'rfrrerrgr1r^no_gradr5rrhr1r1rGr2r/s>   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_rangerrrrrnrrrrtrrrrArrr)rFr:rr?rGr1r2r>s   zPatchTSTRegressionHead.__init__rvcCs|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.forwardrrtr1r1rGr2rsrz, 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>rcrrrarr,rrrrrr@rrxrrqrGr1r2r>-s        zPatchTSTForRegression.__init__NrrrfrrMrhrNc 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 ```NTrycsg|] }|djjqS)r)r,r:r)ritemrFr1r2r}sz1PatchTSTForRegression.forward..rr{rr)r,r0rIr)r:rlrrxrrrrgr:rDrr~r/rIr) rFrrrfrrMrhrrr,rrr1rr2r^Gs<%      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)rrrfrcrr1rrrr1r2rrz2PatchTSTForRegression.generate..rr"rr) r:rrrr0rr'rr,rr5rr1rr2rs zPatchTSTForRegression.generaterqr)r_r`rarr>rr'rfrrerrgr/r^rr5rrhr1r1rGr2r's@  Ir)rarrrwrr)NrN)NFrrsr#)Lrbr  dataclassesrtypingrrrr'r activationsrmodeling_flash_attention_utilsr modeling_outputsr modeling_utilsr r processing_utilsr time_series_utilsrrrutilsrrrconfiguration_patchtstr get_loggerr_rModulerfrdr3r4rjlistrercrrrrrrrrrr%r+r/r1r3r5 distributions Distributionr:rDrFrSr^r_rarrrwrrrrrr__all__r1r1r1r2s          z =  D0< !$8> "$7po%W`=7