from typing import Optional, Tuple, List, Union import torch from torch import nn import torch.nn.functional as F from transformers import PreTrainedModel, Cache, DynamicCache from transformers.activations import ACT2FN from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask from transformers.modeling_outputs import MoeModelOutputWithPast, MoeCausalLMOutputWithPast from .configuration_timer import TimerConfig from .ts_generation_mixin import TSGenerationMixin def rotate_half(x): x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2:] return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): cos = cos[position_ids].unsqueeze(unsqueeze_dim) sin = sin[position_ids].unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed class TimerPatchEmbedding(nn.Module): def __init__(self, config: TimerConfig): super().__init__() self.input_token_len = config.input_token_len self.emb = nn.Linear(config.input_token_len, config.hidden_size, bias=False) def forward(self, hidden_state: torch.Tensor): hidden_state = hidden_state.unfold( dimension=-1, size=self.input_token_len, step=self.input_token_len) return self.emb(hidden_state) class TimerPointEmbedding(nn.Module): def __init__(self, config: TimerConfig): super().__init__() self.emb_layer = nn.Linear( config.input_token_len, config.hidden_size, bias=False) self.gate_layer = nn.Linear( config.input_token_len, config.hidden_size, bias=False) self.act_fn = ACT2FN[config.hidden_act] def forward(self, x): emb = self.act_fn(self.gate_layer(x)) * self.emb_layer(x) return emb class TimeMoeRotaryEmbedding(torch.nn.Module): def __init__(self, dim, max_position_embeddings=10000, base=10000, device=None): super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) # Build here to make `torch.jit.trace` work. self._set_cos_sin_cache( seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype() ) def _set_cos_sin_cache(self, seq_len, device, dtype): self.max_seq_len_cached = seq_len t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq) freqs = torch.outer(t, self.inv_freq) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer( "cos_cached", emb.cos().to(dtype), persistent=False) self.register_buffer( "sin_cached", emb.sin().to(dtype), persistent=False) def forward(self, x, seq_len=None): # x: [bs, num_attention_heads, seq_len, head_size] if seq_len > self.max_seq_len_cached: self._set_cos_sin_cache( seq_len=seq_len, device=x.device, dtype=x.dtype) return ( self.cos_cached[:seq_len].to(dtype=x.dtype), self.sin_cached[:seq_len].to(dtype=x.dtype), ) class TimerAttention(nn.Module): def __init__(self, config: TimerConfig, layer_idx: Optional[int] = None): super().__init__() self.layer_idx = layer_idx self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = self.hidden_size // self.num_heads self.attention_dropout = config.attention_dropout self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=True) self.k_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=True) self.v_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=True) self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False) self.rotary_emb = TimeMoeRotaryEmbedding( self.head_dim, max_position_embeddings=config.max_position_embeddings) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: bool = False, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) query_states = query_states.view( bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = key_states.view( bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) value_states = value_states.view( bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) kv_seq_len = key_states.shape[-2] if past_key_value is not None: kv_seq_len += past_key_value.get_usable_length( kv_seq_len, self.layer_idx) cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) query_states, key_states = apply_rotary_pos_emb( query_states, key_states, cos, sin, position_ids) if past_key_value is not None: key_states, value_states = past_key_value.update( key_states, value_states, self.layer_idx) attn_output = F.scaled_dot_product_attention( query_states, key_states, value_states, attention_mask, dropout_p=self.attention_dropout) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) attn_output = self.o_proj(attn_output) if not output_attentions: attn_weights = None return attn_output, attn_weights, past_key_value class TimerMLP(nn.Module): def __init__(self, hidden_size: int, intermediate_size: int, hidden_act: str): super().__init__() self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.gate_proj = nn.Linear( self.hidden_size, self.intermediate_size, bias=False) self.up_proj = nn.Linear( self.hidden_size, self.intermediate_size, bias=False) self.down_proj = nn.Linear( self.intermediate_size, self.hidden_size, bias=False) self.act_fn = ACT2FN[hidden_act] def forward(self, hidden_state): return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)) class TimerDecoderLayer(nn.Module): def __init__(self, config: TimerConfig, layer_idx: int): super().__init__() self.self_attn = TimerAttention(config, layer_idx) self.ffn_layer = TimerMLP( hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, ) self.norm1 = torch.nn.LayerNorm(config.hidden_size) self.norm2 = torch.nn.LayerNorm(config.hidden_size) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, **kwargs, ) -> Tuple[torch.FloatTensor, torch.FloatTensor, Optional[torch.FloatTensor], Optional[torch.FloatTensor]]: residual = hidden_states # Self Attention hidden_states, self_attn_weights, present_key_value = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, ) hidden_states = residual + hidden_states hidden_states = self.norm1(hidden_states) # Fully Connected residual = hidden_states hidden_states = self.ffn_layer(hidden_states) hidden_states = residual + hidden_states hidden_states = self.norm2(hidden_states) if not output_attentions: self_attn_weights = None if not use_cache: present_key_value = None return hidden_states, self_attn_weights, present_key_value class TimerPreTrainedModel(PreTrainedModel): config_class = TimerConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["TimeMoeDecoderLayer"] _skip_keys_device_placement = "past_key_values" _supports_flash_attn_2 = True _supports_sdpa = False _supports_cache_class = True def _init_weights(self, module): std = self.config.initializer_range if isinstance(module, torch.nn.Linear): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, torch.nn.Embedding): module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() class TimerModel(TimerPreTrainedModel): def __init__(self, config: TimerConfig): super().__init__(config) self.embed_layer = TimerPatchEmbedding(config) self.layers = nn.ModuleList( [TimerDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = torch.nn.LayerNorm(config.hidden_size) self.gradient_checkpointing = False def forward( self, input_ids: torch.FloatTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple, MoeModelOutputWithPast]: # input_ids is the input of time series, its shape is [batch_size, seq_len] output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = return_dict if return_dict is not None else self.config.use_return_dict # retrieve input_ids and inputs_embeds if input_ids is not None and inputs_embeds is not None: raise ValueError( "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") elif input_ids is not None: batch_size, seq_length = input_ids.shape elif inputs_embeds is not None: batch_size, seq_length, _ = inputs_embeds.shape else: raise ValueError( "You have to specify either decoder_input_ids or decoder_inputs_embeds") if inputs_embeds is None: inputs_embeds = self.embed_layer(input_ids) seq_length = inputs_embeds.shape[1] if self.gradient_checkpointing and self.training: if use_cache: use_cache = False past_key_values_length = 0 if use_cache: use_legacy_cache = not isinstance(past_key_values, Cache) if use_legacy_cache: past_key_values = DynamicCache.from_legacy_cache( past_key_values) past_key_values_length = past_key_values.get_usable_length( seq_length) if position_ids is None: device = input_ids.device if input_ids is not None else inputs_embeds.device position_ids = torch.arange( past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device ) # position_ids = position_ids.unsqueeze(0).view(-1, seq_length) position_ids = position_ids.view(-1, seq_length) else: position_ids = position_ids.view(-1, seq_length).long() # 4d mask is passed through the layers attention_mask = _prepare_4d_causal_attention_mask( attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length, sliding_window=None, ) hidden_states = inputs_embeds # decoder layers all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None next_decoder_cache = None for decoder_layer in self.layers: if output_hidden_states: all_hidden_states += (hidden_states,) if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( decoder_layer.__call__, hidden_states, attention_mask, position_ids, past_key_values, output_attentions, use_cache, ) else: layer_outputs = decoder_layer( hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_values, output_attentions=output_attentions, use_cache=use_cache, ) hidden_states = layer_outputs[0] if output_attentions: all_self_attns += (layer_outputs[1],) if use_cache: next_decoder_cache = layer_outputs[2] hidden_states = self.norm(hidden_states) # add hidden states from the last decoder layer if output_hidden_states: all_hidden_states += (hidden_states,) next_cache = None if use_cache: next_cache = next_decoder_cache.to_legacy_cache( ) if use_legacy_cache else next_decoder_cache if not return_dict: return tuple( v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None ) return MoeModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, hidden_states=all_hidden_states, attentions=all_self_attns, ) class TimerForPrediction(TimerPreTrainedModel, TSGenerationMixin): def __init__(self, config: TimerConfig): super().__init__(config) self.config = config self.model = TimerModel(self.config) lm_head_list = [] self.output_token_len_map = {} for i, output_token_len in enumerate(self.config.output_token_lens): lm_head_list.append( nn.Linear(self.config.hidden_size, output_token_len, bias=False)) self.output_token_len_map[output_token_len] = i self.lm_heads = nn.ModuleList(lm_head_list) self.loss_function = torch.nn.MSELoss(reduction='none') self.post_init() def set_decoder(self, decoder): self.model = decoder def get_decoder(self): return self.model def forward( self, input_ids: torch.FloatTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.FloatTensor] = None, loss_masks: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, max_output_length: Optional[int] = None, revin: Optional[bool] = False, ) -> Union[Tuple, MoeCausalLMOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict if revin: mean, std = input_ids.mean(dim=-1, keepdim=True), input_ids.std(dim=-1, keepdim=True) input_ids = (input_ids - mean) / std outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state predictions = None loss = None if labels is not None: ar_loss = 0.0 for lm_head, output_token_len in zip(self.lm_heads, self.config.output_token_lens): one_predictions = lm_head(hidden_states) one_loss = self.calc_ar_loss( one_predictions, labels, loss_masks, output_token_len) ar_loss += one_loss if predictions is None: predictions = one_predictions loss = ar_loss / len(self.config.output_token_lens) else: if max_output_length is None: output_token_len = self.config.output_token_lens[0] max_output_length = output_token_len else: output_token_len = self.config.output_token_lens[0] for h in self.config.output_token_lens[1:]: if h > max_output_length: break else: output_token_len = h lm_head = self.lm_heads[self.output_token_len_map[output_token_len]] predictions = lm_head(hidden_states)[:, -1, :] if output_token_len > max_output_length: predictions = predictions[:, :max_output_length] if revin: predictions = predictions * std + mean if not return_dict: output = (predictions,) + outputs[1:] return (loss) + output if loss is not None else output return MoeCausalLMOutputWithPast( loss=loss, logits=predictions, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) def calc_ar_loss(self, predictions, labels, loss_masks, output_token_len): seq_len = predictions.shape[1] * self.config.input_token_len labels = labels[:, :seq_len - self.config.input_token_len + output_token_len] shift_labels = labels.unfold( dimension=-1, size=output_token_len, step=self.config.input_token_len) # Calculate loss with mask losses = self.loss_function(predictions, shift_labels).mean(dim=-1) if loss_masks is not None: losses = losses * loss_masks loss = losses.sum() / loss_masks.sum() else: loss = torch.mean(losses) return loss def prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, revin=True, **kwargs ): # Omit tokens covered by past_key_values if past_key_values is not None: if isinstance(past_key_values, Cache): cache_length = past_key_values.get_seq_length() if isinstance(past_key_values, DynamicCache): past_length = past_key_values.seen_tokens else: past_length = cache_length max_cache_length = past_key_values.get_max_length() else: cache_length = past_length = past_key_values[0][0].shape[2] max_cache_length = None # Keep only the unprocessed tokens: # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as # input) if attention_mask is not None and attention_mask.shape[1] > (input_ids.shape[1] // self.config.input_token_len): input_ids = input_ids[:, - (attention_mask.shape[1] - past_length):] # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard # input_ids based on the past_length. elif past_length < (input_ids.shape[1] // self.config.input_token_len): input_ids = input_ids[:, past_length * self.config.input_token_len:] # 3 - Otherwise (past_length >= (input_ids.shape[1] // self.config.input_token_len)), let's assume input_ids only has unprocessed tokens. # If we are about to go beyond the maximum cache length, we need to crop the input attention mask. if ( max_cache_length is not None and attention_mask is not None and cache_length + (input_ids.shape[1] // self.config.input_token_len) > max_cache_length ): attention_mask = attention_mask[:, -max_cache_length:] position_ids = kwargs.get("position_ids", None) if attention_mask is not None and position_ids is None: # create position_ids on the fly for batch generation position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) if past_key_values: position_ids = position_ids[:, - (input_ids.shape[1] // self.config.input_token_len):] # if `inputs_embeds` are passed, we only want to use them in the 1st generation step if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: model_inputs = {"input_ids": input_ids} model_inputs.update( { "position_ids": position_ids, "past_key_values": past_key_values, "use_cache": kwargs.get("use_cache"), "attention_mask": attention_mask, "revin": revin } ) return model_inputs