qi~dZddlZddlmZmZddlmZddlZddlm Z ddl m cm Z ddl mZddlmZddlmZdd lmZmZdd lmZdd lmZmZmZd d lmZd dlm Z ddl!m"Z"ejFe$Z%eeGddeZ&eeGddeZ'Gdde jPZ)Gdde jPZ*GddeZ+Gdde jPZ,Gdde jPZ-Gdd e jPZ.eGd!d"eZ/eGd#d$e/Z0Gd%d&e/Z1gd'Z2y)(zPyTorch TimesFM model.N)CallableSequence) dataclass)initialization)FlashAttentionKwargs)BaseModelOutput)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)auto_docstringcan_return_tuplelogging) LlamaRMSNorm)simple_eager_attention_forward) TimesFmConfigcbeZdZUdZdZej dzed<dZej dzed<y) TimesFmOutputz loc (`torch.Tensor` of shape `(batch_size, )`): The mean of the time series inputs. scale (`torch.Tensor` of shape `(batch_size,)`): The scale of the time series inputs. Nlocscale) __name__ __module__ __qualname____doc__rtorchTensor__annotations__r]/opt/pipecat/venv/lib/python3.12/site-packages/transformers/models/timesfm/modular_timesfm.pyrr&s/ $C #!%E5<<$ %r!rceZdZUdZdZej dzed<dZej dzed<dZ ej e zdzed<y)TimesFmOutputForPredictiona mean_predictions (`torch.Tensor` of shape `(batch_size, sequence_length)`): The mean predictions of the time series. full_predictions (`torch.Tensor` of shape `(batch_size, sequence_length)`): The full predictions of the time series including the mean and the quantiles. loss (`torch.Tensor` of shape `(1,)`, *optional*, returned when `future_values` is provided): The loss of the TimesFM model. Nmean_predictionsfull_predictionsloss) rrrrr%rrrr&r'floatr r!r"r$r$4sI-1ellT)0,0ellT)0(,D%,,  %,r!r$c0eZdZdZdeffd ZddZxZS) TimesFmMLPzPax MLP in pytorch.configct||j}|j}t j |||_t j |||_t j|d|_ y)Ngư>)normalized_shapeeps) super__init__ hidden_sizeintermediate_sizennLinear gate_proj down_proj LayerNorm layer_norm)selfr+r1r2 __class__s r"r0zTimesFmMLP.__init__Hs] (( "44;0AB#4kB,, Nr!c|j|}|j|}tj|}|j |}||d|dddddfz z}||zS)N?)r8r5Frelur6)r9xpaddingsgate_inpgateoutputss r"forwardzTimesFmMLP.forwardQsc??1%~~h'vvd|..&  x1d ';!; .D.$00,,1"&((5+?% BV+V"WZ]^lop^prsZt"t   EIIell>&W[rZr&rs s r!cN| | td|Jtj|tj|jj j d}n'|jdk7rtd|j|jg|jd|jjdddz}tjtj|tj|gd}tj|ddd|jdzf}|S) aGenerates a Tensor of sinusoids with different frequencies. Args: seq_length: an optional Python int defining the output sequence length. if the `position` argument is specified. position: [B, seq_length], optional position for each token in the sequence, only required when the sequence is packed. Returns: [B, seqlen, D] if `position` is specified, else [1, seqlen, D] z.Either position or seq_length must be providedr_devicerrz*position must be 2-dimensional, got shape rdim) ValueErrorrrhrir]rn unsqueezendimshapeviewcatsincosr=padrb)r9 seq_lengthposition scaled_timesignals r"rDz"TimesFmPositionalEmbedding.forwards   2MN N  ||JemmDL_L_LfLfgqqrstH ]]a I(..IYZ[ [#hmm7X^^7Q7$:M:M:R:RSTVWY[:\\ EIIk2EIIk4JKQRSv1a)<)@P@PQ r!queryreturnctj|jjdt j |j z }||dddddfzS)Ng^$3eG?)r=softplusrmulrcsqrtr)r9rrs r" _scale_queryzTimesFmAttention._scale_querysJ 4<<(,,[499T]];S-STuT4q0111r!N hidden_statesattention_maskkwargsc |jdd}g|d|j}|j|j|j dd}|j |}|j |j|j dd}|j|j|j dd}tj|jjt} | |||||f|jsdn |jdd|\} } | jg|dj!} |j#| } | | fS)Nrorrr<)dropoutr)rurrrv transposerrrr get_interfacer+_attn_implementationrtrainingrreshape contiguousr) r9rrr input_shape hidden_shape query_states key_states value_statesattention_interface attn_output attn_weightss r"rDzTimesFmAttention.forwardsW $))#2. 88b8$--8 {{=166|DNNqRST ((6 [[/44\BLLQPQR {{=166|DNNqRST (?(M(M KK , ,.L) %8      % $}}C$2H2H %  % ! \*k));;;;FFH kk+. L((r!rE)rrrrrintr0rrrr rtuplerDrGrHs@r"rrsvR}RR(2%,,25<<2/3)||) t+)-. ) u||U\\D00 1 )r!rc eZdZdZdedeffd Z d dejdejdejde d e ejd zejff d Z xZ S) TimesFmDecoderLayerzTransformer layer.r+rct|t|||_t ||_t |j|j|_ y)N)r)r.) r/r0r self_attnr*mlprYr1 rms_norm_epsinput_layernormrs r"r0zTimesFmDecoderLayer.__init__sC )&IFf%-f.@.@fFYFYZr!rrr@output_attentionsrNc|}|j|}|j|||\}}||z}|j||}||fS)N)rrr)r@)rrr)r9rrr@rrWscoress r"rDzTimesFmDecoderLayer.forwardse!,,]; $')/!/!  v !=0 B }$$r!)F) rrrrrrr0rrboolrrDrGrHs@r"rrsy[}[[#( %||% %,, %  % u||d"ELL0 1 %r!rcheZdZUeed<dZdgZdZdZdZ e jfdZ xZ S)TimesFmPreTrainedModelr+timesfmr past_values)timeTc "t||t|tr t j |j yt|tr|jdz}|j|j}}tjt|t|z t|dz dz }t j|j |t#j$t#j&|t"j(| zzyy)Nrrr^)r/ _init_weights isinstancerinitones_rr[rbrar`rcrdr(recopy_r]rrgrhri)r9modulerjrar`rkr:s r"rz$TimesFmPreTrainedModel._init_weights s f% f. / JJv~~ &  : ;#22a7N+1+?+?AUAU=M&*hhu]/CeMFZ/Z&[^a"A_' # JJ%%))ELLu}}MQhPhhij  S>S S LLL$6$6|?R?R S  !% 7 7 U(3+=> ><"@bI ))-8  ))Lb9!< ;; / /'' (9(9!(<=GllG9{/@/@/C#CKG44_gNG 7 "K d#u $ 88*)//2%% '' 9 [[!@4;;#@#@A 8E$)+-("3 % !FM !%%f-#!((7 8 !, 0A A  $ ++):~a(   r!rNrr_rnrc|jrtj|jntj|j}|&|j |j dddd}||z}|rbtjtj||f|||zd}|j dd||}|tj||}|S|}|S)a Creates 4D attention mask and combines causal and padding masks if needed. Args: attention_mask: Optional tensor of shape (batch_size, seq_length) containing padding mask sequence_length: Length of the sequence dtype: Data type of the mask device: Device of the mask is_causal: Whether to apply causal masking Returns: 4D attention mask of shape (batch_size, 1, seq_length, seq_length) rrrorm)diagonal) is_floating_pointrfinforiinforvrutriuonesminimum)rrr_rnr min_value causal_masks r"rz'TimesFmModel._prepare_4d_attention_masks,/4.E.EEKK&**5;;W\K]KaKa   %+001E1Ea1H!QPRSN+i7N ** O_=USYZ]ffK&**1a/RK)!&~{!K"-r!paddingcFdtjfd}tjd|z d}||}tj|jd}|||ddf}|||ddf}d|z }tj|d} tj | d } tj||zd} | | z } || j d z |z} tj| dzd| z } tj | d } tj| }| |fS) aCalculates mean and standard deviation of `inputs` across axis 1. It excludes values where `padding` is 1. Args: inputs: A PyTorch tensor of shape [b, n, p]. padding: A PyTorch tensor of shape [b, n, p] with values 0 or 1. Returns: A tuple containing the mean and standard deviation. We return the statistics of the first patch with more than three non-padded values. arrc(tj|dk\jtjd}|dk\jtjj d}tj |dk(|j ddz |S)Nrrrpr)rargmaxtoint32sumrru)rindicesrow_sums r"_get_patch_indexz?TimesFmModel._timesfm_masked_mean_std.._get_patch_indexskllC1H==#=1EGaxmmEKK0444;G;;w!|SYYq\A-=wG Gr!rrrprNr<rror)rrrrhrurrsr)rrrpad_sum patch_indicesbidxsrrzmasknum_valid_elements masked_sum masked_meanmasked_centered_arr masked_var masked_stds r"rz%TimesFmModel._timesfm_masked_mean_stds  H%,, H ))AKQ/(1  V\\!_-UM1,-e]A-.3w#YYt3"[[);EYYsTzq1  #55  #[%:%:2%>>$FYY2A51=@RR [[5 ZZ + J&&r!rseqc|j\}}}|dk(}|jtjj d}d||j d<tj ||jjdddj|d|}||ddddfz |z}|jd|} | S)zShifts rows of seq based on the first 0 in each row of the mask. Args: mask: mask tensor of shape [B, N] seq: seq tensor of shape [B, N, P] Returns: The shifted sequence. rrrpro)rnN) rurrrranyrhrnrvexpandgather) rr# batch_sizenum_seq feature_dimnew_maskr idx_range shifted_idx shifted_seqs r"rz&TimesFmModel._timesfm_shift_padded_seqs,/99( G[%)QY++ekk*11a18)+!$$%LL<AA!RKRRS]_acno !71dD=#99WD jjK0 r!)FF)T)rrrrr0rrrrrr LongTensorrrrD staticmethodrr_rnrrrrGrHs@r"rrs}&$ll$27,,$ u||U5<<#=>> ?$  #(%* V \\V #--V ll V  V # V  V V p  + t+++{{+ +  +   ++Z,',' ,'QVW\WcWcejeqeqWqQr,','\ 5<<ELLr!rc`eZdZdZdeffd Z ddeejdee dzde dzde ejd ffd Z d ejd e ejejfdejfd Z dejdejdejfdZ ee ddeejdeeje zdzde dzdejdzde dzdedededzdedzdefdZedejde deejfdZxZS)TimesFmModelForPredictionz/TimesFM model for quantile and mean prediction.r+cJt||||_|j|_|j |_t||_t|j|j dt|jzz|j|_|jy)Nrr)r/r0r+context_length context_lenhorizon_length horizon_lenrdecoderrJr1len quantilesr2horizon_ff_layerr)r9r+r:s r"r0z"TimesFmModelForPrediction.__init__s   !00!00#F+ !5))--S9I9I5J1JK00!  r!Nrrr5r.cX| |j}gg}}|D]}|jd}tj||jz|j |j }||kr||z } tjtj| |j |j |gd}tjtj| |j |j |gd}n||kDr|| d}|||jz d}|j||j|tj|dtj|df} |E| tj|dt|tjjddfz} | S)aPad/truncate input time series to `context_len` and build a padding mask. Args: inputs: A list of 1d Tensors. Each Tensor is the context time series of a single forecast task. freq: Optional list of frequencies (returned as a tensor when provided). context_len: Optional context length override (defaults to `self.context_len`). Returns: Tuple of (padded_inputs, padding_mask) and optionally a freq tensor. Nrrmrpr^ror)r5rurzerosr7r_rnrwr rstackrr9rr) r9rrr5input_ts input_paddingts input_lenr num_front_padresults r" _preprocessz%TimesFmModelForPrediction._preprocess/s  **K"$b- *B Ikk)d.>.>">bhhWYW`W`aG;& +i 7 YY MRTR[R[ \^`aghi))UZZ RXXV]VdVd%egn$ouvw[( &!K$2B2B$B"C"EF OOB    ) *++hA. Mq0QR  u||D3v;,?u{{S[[\^`abddF r! model_outputrc|j|}|j\}}}|j|||jjt |jj dz}|\}}||dddddfz|dddddfzS)z*Postprocess output of stacked transformer.rN)r;rurvr+r6r9r:) r9rFr output_tsbn_rrs r"_postprocess_outputz-TimesFmModelForPrediction._postprocess_outputTs )),7 //1aNN1a)C)CSI^I^E_bcEcd  E5D$!4551dD$;N8OOOr! predictionstargetsc*g}t|jjD]M\}}||d|fz }tj|dz |z||z}|j |j Otj|j S)N.r) enumerater+r:rrermeanr>)r9rMrNlossesiqerrorsr's r"_quantile_lossz(TimesFmModelForPrediction._quantile_losscsdkk334 'DAq{3622F99a!ev-q6z:D MM$))+ & '{{6"''))r!r window_size future_valuesforecast_context_lenreturn_forecast_on_contexttruncate_negativerrc n | |j} n|} |dj} |D cgc]} | | d }} tjtj|D cgc]} tj| c} }|Yg}g}t |D]A\}} |j |j| ||*|j ||gdzC|}||}|$tjddgt|z}||jj}| |jj} |j||\}}}|j| }|j| }|j| }|}|j d}g}|j d|j d|j"zk7r8t%d|j dd|j dd|j"|jj&}|j"|zdz |z}t)|D]/}|ddd|j df}|dd| df}|dd| df}|j+||||| }|j-|j.|j0|j2f}|rl|dk(rg|dddd d|jj4ddf}|j7|j9dd |j9d }|j;||ddd d|df} |ddd d|ddf}|j;|tj<|| gd }2|rHtj<|d ddd||jj4z |j"zddf}n-tj<|d ddd|j"ddf}|dddddf}!|*|!dddd f|!dddd fz}!|dddd f|dddd fz}|dk\r.|r,tj>|!d}!tj>|d}d}"|9tAjB|!|}#|jE|ddddddf|}$|#|$z}"tGj.|r |jHnd| r |jJnd|!||"Scc} wcc} w)aa past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): Past values of the time series that serves as input to the model. freq (`torch.LongTensor` of shape `(batch_size,)`): Frequency indices for the time series data. window_size (`int`, *optional*): Window size of trend + residual decomposition. If None then we do not do decomposition. future_values (`torch.Tensor`, *optional*): Optional future time series values to be used for loss computation. forecast_context_len (`int`, *optional*): Optional max context length. return_forecast_on_context (`bool`, *optional*): True to return the forecast on the context when available, i.e. after the first input patch. truncate_negative (`bool`, *optional*): Truncate to only non-negative values if any of the contexts have non-negative values, otherwise do nothing. output_attentions (`bool`, *optional*): Whether to output the attentions. output_hidden_states (`bool`, *optional*): Whether to output the hidden states. Example: ```python >>> from transformers import TimesFmModelForPrediction >>> model = TimesFmModelForPrediction.from_pretrained("google/timesfm-2.0-500m-pytorch") >>> forecast_input = [torch.linspace(0, 20, 100).sin(), torch.linspace(0, 20, 200).sin(), torch.linspace(0, 20, 400).sin()] >>> frequency_input = torch.tensor([0, 1, 2], dtype=torch.long) >>> # Generate >>> with torch.no_grad(): >>> outputs = model(past_values=forecast_input, freq=frequency_input, return_dict=True) >>> point_forecast_conv = outputs.mean_predictions >>> quantile_forecast_conv = outputs.full_predictions ``` Nrrz6No frequency provided via `freq`. Default to high (0).rz=Length of paddings must match length of input + horizon_len: z != z + )rrrrrror)axis.r)rrrr%r&r')&r5rnrrr>rPextend_timesfm_moving_averageloggerinfor9r+rrrErrur7rrr6rr8rLrrrrrsizer concatenatemaximumr=mse_lossrVr$rr)%r9rrrWrXrYrZr[rrr fcontext_lenrnrArinp_min new_inputs new_freqsrSr?r@inp_freq final_outr5 full_outputsoutput_patch_lennum_decode_patches step_indexcurrent_paddingdecoder_output fprop_outputs new_full_tsnew_ts mean_outputsr're quantile_losss% r"rDz!TimesFmModelForPrediction.forwardks-j '++L/LQ&&0;;"l]^$;;))EKK(H22(HIJ  "JI"6* 42!!$">">r;"OP#$$d1gY]3 4 F  < KKP Q3V$D  $ $ = =   '#';;#C#C ,0,<,M!\\$$1"3%9 *N!4400##^%9%9:M *jAo,Ass4Ndkk6N6N4NPQ,QR )11+2B2B12Er;K[K[\]K^_ ##K0#1b*;+;*;Q#>?F'2/@0@/@!(CDK    ,))9f*=BGI= H@ & ,,\BPkDKK$<$<H // Q12X0F VMm+D),>>4E~004:N.66TX))   A<(Hs R-R2 rc4tj||dz dfdd}tj||j|j |z }tj |jddd|jdddj}|||z gS)zCCalculates the moving average using PyTorch's convolution function.rrconstantrmro) r=rzrr r_rnconv1drvsqueeze)rrW arr_paddedkernel smoothed_arrs r"r_z1TimesFmModelForPrediction._timesfm_moving_averagesUU3q! 4j!D KsyyL{Zxx 1b 96;;q!R;PQYY[ cL011r!r)NNNNFFNN)rrrrrr0rrrrrrErLrVrr rr$rDr0listr_rGrHs@r"r2r2s9}(lp#u||,#4rsI .! &B/F&>>/U0   H %  &O &  &  - -  -,!299!, \ ++\9)ryy9)x%"))%B_6y)yyxB2 6B2J Rr!