""" Configuration class for FalconTST model. This module defines the configuration for FalconTST, a large-scale time series foundation model that utilizes Mixture of Experts (MoE) architecture with multiple patch tokenizers. """ from typing import List, Optional, Union from transformers import PretrainedConfig import torch class FalconTSTConfig(PretrainedConfig): """ Configuration class for FalconTST model. FalconTST is a time series foundation model that uses Mixture of Experts architecture with multiple patch tokenizers for efficient time series forecasting. This configuration inherits from [`PretrainedConfig`] and can be used to control the model output. Read the documentation from [`PretrainedConfig`] for more information. Args: hidden_size (`int`, *optional*, defaults to 1024): Dimensionality of the encoder layers and the pooler layer. ffn_hidden_size (`int`, *optional*, defaults to 4096): Dimensionality of the feed-forward networks in the transformer layers. seq_length (`int`, *optional*, defaults to 2880): Maximum sequence length that the model can handle. add_bias_linear (`bool`, *optional*, defaults to `False`): Whether to add bias in linear layers. rope_theta (`int`, *optional*, defaults to 10000): The base period of the RoPE embeddings. num_hidden_layers (`int`, *optional*, defaults to 3): Number of hidden layers in the transformer encoder. num_attention_heads (`int`, *optional*, defaults to 16): Number of attention heads for each attention layer in the transformer encoder. mask_pad_value (`float`, *optional*, defaults to 255.0): Value used for padding/masking in input sequences. expert_num_layers (`int`, *optional*, defaults to 4): Number of transformer layers within each expert. shared_patch_size (`int`, *optional*, defaults to 64): Size of patches for the shared expert. patch_size_list (`List[int]`, *optional*, defaults to [96, 64, 48, 24]): List of patch sizes for different experts. multi_forecast_head_list (`List[int]`, *optional*, defaults to [24, 96, 336]): List of forecast lengths for multi-head prediction. is_revin (`bool`, *optional*, defaults to `True`): Whether to use RevIN (Reversible Instance Normalization). params_dtype (`str`, *optional*, defaults to "bfloat16"): Data type for model parameters. use_cpu_initialization (`bool`, *optional*, defaults to `False`): Whether to initialize model parameters on CPU. rotary_interleaved (`bool`, *optional*, defaults to `False`): Whether to use interleaved rotary position embeddings. do_expert_forecast (`bool`, *optional*, defaults to `True`): Whether experts perform forecasting. residual_backcast (`bool`, *optional*, defaults to `True`): Whether to use residual connections for backcast. do_base_forecast (`bool`, *optional*, defaults to `False`): Whether to use base forecasting. heterogeneous_moe_layer (`bool`, *optional*, defaults to `True`): Whether to use heterogeneous MoE layers. test_data_seq_len (`int`, *optional*, defaults to 2880): Sequence length for test data. test_data_test_len (`int`, *optional*, defaults to 720): Test length for test data. autoregressive_step_list (`List[int]`, *optional*, defaults to [2, 4, 1]): List of autoregressive steps for different forecast heads. multi_forecast_head_type (`str`, *optional*, defaults to "single"): Type of multi-forecast head aggregation. num_experts (`int`, *optional*, defaults to 4): Number of experts in the MoE layer. moe_router_topk (`int`, *optional*, defaults to 2): Number of top experts to route each token to. moe_ffn_hidden_size (`int`, *optional*, defaults to 4096): Hidden size for MoE feed-forward networks. moe_shared_expert_intermediate_size (`int`, *optional*, defaults to 4096): Intermediate size for shared experts. init_method_std (`float`, *optional*, defaults to 0.06): Standard deviation for weight initialization. initializer_range (`float`, *optional*, defaults to 0.02): Range for weight initialization. moe_router_enable_expert_bias (`bool`, *optional*, defaults to `False`): Whether to enable expert bias in routing. moe_expert_final_layernorm (`bool`, *optional*, defaults to `True`): Whether to apply layer normalization at the end of each expert. transformer_input_layernorm (`bool`, *optional*, defaults to `True`): Whether to apply layer normalization to transformer inputs. moe_router_pre_softmax (`bool`, *optional*, defaults to `True`): Whether to apply softmax before routing. q_layernorm (`bool`, *optional*, defaults to `False`): Whether to apply layer normalization to query vectors. k_layernorm (`bool`, *optional*, defaults to `False`): Whether to apply layer normalization to key vectors. moe_router_score_function (`str`, *optional*, defaults to "softmax"): Score function for MoE routing ("softmax" or "sigmoid"). tie_word_embeddings (`bool`, *optional*, defaults to `False`): Whether to tie word embeddings. """ model_type = "FalconTST" def __init__( self, # model configs add_bias_linear: bool = False, num_hidden_layers: int = 3, hidden_size: int = 1024, ffn_hidden_size: int = 4096, num_attention_heads: int = 16, seq_length: int = 2880, mask_pad_value: float = 255.0, is_revin: bool = True, shared_patch_size: int = 32, patch_size_list: Optional[List[int]] = None, residual_backcast: bool = True, do_base_forecast: bool = False, do_expert_forecast: bool = True, heterogeneous_moe_layer: bool = False, expert_num_layers: int = 4, multi_forecast_head_list: Optional[List[int]] = None, multi_forecast_head_type: str = "single", rope_theta: int = 1000000, rotary_interleaved: bool = False, block_input_layernorm: bool = True, # moe configs num_experts: int = 4, moe_router_topk: int = 2, moe_router_pre_softmax: bool = True, moe_router_score_function: str = "softmax", moe_ffn_hidden_size: int = 4096, moe_shared_expert_intermediate_size: int = 4096, moe_router_enable_expert_bias: bool = False, moe_expert_final_layernorm: bool = True, # initial configs use_cpu_initialization: bool = False, init_method_std: float = 0.06, initializer_range: float = 0.02, # test configs test_data_seq_len: int = 2880, test_data_test_len: int = 720, autoregressive_step_list: Optional[List[int]] = None, **kwargs, ): """Initialize FalconTST configuration.""" # model configs self.add_bias_linear = add_bias_linear self.num_hidden_layers = num_hidden_layers self.hidden_size = hidden_size self.ffn_hidden_size = ffn_hidden_size self.num_attention_heads = num_attention_heads self.kv_channels = self.hidden_size // self.num_attention_heads self.seq_length = seq_length self.mask_pad_value = mask_pad_value self.is_revin = is_revin self.shared_patch_size = shared_patch_size if patch_size_list is None: patch_size_list = [96, 64, 48, 24] self.patch_size_list = patch_size_list self.residual_backcast = residual_backcast self.do_base_forecast = do_base_forecast self.do_expert_forecast = do_expert_forecast self.heterogeneous_moe_layer = heterogeneous_moe_layer self.expert_num_layers = expert_num_layers if multi_forecast_head_list is None: multi_forecast_head_list = [24, 96, 336] self.multi_forecast_head_list = multi_forecast_head_list self.pred_length = max(self.multi_forecast_head_list) self.multi_forecast_head_type = multi_forecast_head_type self.rotary_base = rope_theta self.rotary_interleaved = rotary_interleaved self.block_input_layernorm = block_input_layernorm # moe configs self.num_moe_experts = num_experts self.moe_router_topk = moe_router_topk self.moe_router_input_size = self.seq_length self.moe_router_pre_softmax = moe_router_pre_softmax self.moe_router_score_function = moe_router_score_function self.moe_ffn_hidden_size = moe_ffn_hidden_size self.moe_shared_expert_intermediate_size = moe_shared_expert_intermediate_size self.moe_router_enable_expert_bias = moe_router_enable_expert_bias self.moe_expert_final_layernorm = moe_expert_final_layernorm # initial configs self.use_cpu_initialization = use_cpu_initialization self.init_method_std = init_method_std self.initializer_range = initializer_range # test configs self.test_data_seq_len = test_data_seq_len self.inference_length = test_data_test_len if autoregressive_step_list is None: autoregressive_step_list = [2, 4, 1] self.autoregressive_step_list = autoregressive_step_list self.use_cache = True super().__init__( **kwargs, )