diff --git "a/modeling_interns2_preview.py" "b/modeling_interns2_preview.py" --- "a/modeling_interns2_preview.py" +++ "b/modeling_interns2_preview.py" @@ -19,37 +19,47 @@ # limitations under the License. import math -from collections.abc import Callable +from collections.abc import Callable, Sequence from dataclasses import dataclass from typing import Any, Optional +import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from transformers import initialization as init from transformers.activations import ACT2FN -from transformers.cache_utils import Cache, EncoderDecoderCache +from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache from transformers.conversion_mapping import register_checkpoint_conversion_mapping from transformers.core_model_loading import Concatenate, MergeModulelist, WeightConverter from transformers.generation import GenerationMixin from transformers.integrations import use_experts_implementation, use_kernelized_func -from transformers.masking_utils import create_causal_mask +from transformers.masking_utils import create_bidirectional_mask, create_causal_mask from transformers.modeling_flash_attention_utils import FlashAttentionKwargs from transformers.modeling_layers import GradientCheckpointingLayer -from transformers.modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + BaseModelOutputWithPastAndCrossAttentions, + BaseModelOutputWithPooling, + BaseModelOutputWithPoolingAndCrossAttentions, + ModelOutput, +) from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from transformers.processing_utils import Unpack +from transformers.pytorch_utils import apply_chunking_to_forward from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check from transformers.utils.generic import is_flash_attention_requested, maybe_autocast, merge_with_config_defaults from transformers.utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available from transformers.utils.output_capturing import OutputRecorder, capture_outputs from .configuration_interns2_preview import ( + InternS2PreviewBertConfig, InternS2PreviewConfig, InternS2PreviewTextConfig, InternS2PreviewTimeSeriesConfig, + InternS2PreviewTimeSeriesForecasterConfig, InternS2PreviewVisionConfig, ) @@ -1823,14 +1833,30 @@ class InternS2PreviewTimeSeriesEncoder(PreTrainedModel): attention_mask = matrix.to(inputs_embeds.device) return attention_mask + def prepare_padding_mask(self, attention_mask, input_shape, inputs_embeds, input_lens): + + matrix_size = input_shape[1] + matrix_list = [] + + for i in range(input_shape[0]): + padding_matrix = torch.ones(matrix_size, matrix_size) * -65504 + padding_matrix[: input_lens[i], : input_lens[i]] = 0 + matrix_list.append(padding_matrix) + + attention_mask = torch.stack(matrix_list).unsqueeze(1).to(inputs_embeds.dtype).to(inputs_embeds.device) + + return attention_mask + def forward( self, input_features, + input_length=None, attention_mask=None, head_mask=None, output_attentions=None, output_hidden_states=None, return_dict=None, + causal=True, ): # (N, T, C) -> (T, N, C) -> (N, C, T) input_features = input_features.permute(1, 0, 2) @@ -1867,11 +1893,16 @@ class InternS2PreviewTimeSeriesEncoder(PreTrainedModel): input_shape = inputs_embeds.size()[:-1] past_key_values_length = 0 attention_mask = None - if self.mask_type == "chunk": - attention_mask = self.prepare_chunk_attention_mask(attention_mask, input_shape, inputs_embeds) + if causal: + if self.mask_type == "chunk": + attention_mask = self.prepare_chunk_attention_mask(attention_mask, input_shape, inputs_embeds) + else: + attention_mask = self._prepare_decoder_attention_mask( + attention_mask, input_shape, inputs_embeds, past_key_values_length + ) else: - attention_mask = self._prepare_decoder_attention_mask( - attention_mask, input_shape, inputs_embeds, past_key_values_length + attention_mask = self.prepare_padding_mask( + attention_mask, input_shape, inputs_embeds, (input_length + 1) // 2 ) if head_mask is not None: @@ -1935,24 +1966,8 @@ class InternS2PreviewTimeSeriesEncoder(PreTrainedModel): return ModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions) -class InternS2PreviewTimeSeriesConcatSubsampling(nn.Module): - def __init__(self, in_channels: int, concat_size: int): - super().__init__() - self.in_channels = in_channels - self.out_channels = in_channels * concat_size - - def forward(self, ts_signals: torch.Tensor, ts_lens: torch.Tensor): - if ts_signals.shape[1] % 2 != 0: - ts_signals = ts_signals[:, :-1, :] - even_frames = ts_signals[:, ::2, :] - odd_frames = ts_signals[:, 1::2, :] - ts_signals = torch.cat((even_frames, odd_frames), dim=2) - ts_lens = ts_lens // 2 - return ts_signals, ts_lens - - class InternS2PreviewTimeSeriesFixPositionalEncoding(nn.Module): - def __init__(self, d_model: int, max_len: int = 20000): + def __init__(self, d_model: int, max_len: int = 200000): super().__init__() pe = torch.zeros(max_len, d_model, dtype=torch.float) position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) @@ -1964,272 +1979,2763 @@ class InternS2PreviewTimeSeriesFixPositionalEncoding(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: # x: (seq_len, batch_size, d_model) - x = x + self.pe[: x.size(0), :] + x = x + self.pe[: x.size(0), :].to(device=x.device, dtype=x.dtype) return x.clone() -class InternS2PreviewTimeSeriesMultiChannelAdaptiveSubsampling(nn.Module): - def __init__(self, hidden_dim: int = 128, nhead: int = 8, num_encoder_layers: int = 1): +class InternS2PreviewBertSelfAttention(nn.Module): + def __init__(self, config, is_causal=False, layer_idx=None): super().__init__() - self.conv = nn.Conv1d(in_channels=1, out_channels=hidden_dim, kernel_size=5, stride=1, padding=2) - encoder_layers = nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=nhead) - self.transformer_encoder = nn.TransformerEncoder(encoder_layers, num_encoder_layers) - self.pos_encoder = InternS2PreviewTimeSeriesFixPositionalEncoding(d_model=hidden_dim) - self.subsampling = InternS2PreviewTimeSeriesConcatSubsampling(128, 2) + if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"): + raise ValueError( + f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " + f"heads ({config.num_attention_heads})" + ) + self.config = config + + self.num_attention_heads = config.num_attention_heads + self.attention_head_size = int(config.hidden_size / config.num_attention_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + self.scaling = self.attention_head_size**-0.5 + + self.query = nn.Linear(config.hidden_size, self.all_head_size) + self.key = nn.Linear(config.hidden_size, self.all_head_size) + self.value = nn.Linear(config.hidden_size, self.all_head_size) + + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) + + self.is_decoder = config.is_decoder + self.is_causal = is_causal + self.layer_idx = layer_idx def forward( - self, inputs: torch.Tensor, input_lens: torch.Tensor, sr: torch.Tensor, channels: torch.LongTensor - ) -> tuple[torch.Tensor, torch.Tensor]: - features, feature_lens = self.forward_patch(inputs, input_lens, sr, channels) - outputs = features - output_lens = feature_lens - return outputs, output_lens + self, + hidden_states: torch.Tensor, + attention_mask: torch.FloatTensor | None = None, + past_key_values: Cache | None = None, + cache_position: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.attention_head_size) - def forward_patch( - self, inputs: torch.Tensor, input_lens: torch.Tensor, sr: torch.Tensor, channels: torch.LongTensor - ) -> tuple[torch.Tensor, torch.Tensor]: - sr = sr.float() - strides = torch.floor(160 / ((1 + torch.exp(-sr / 100)) ** 6)) - patch_sizes = strides * 2 - patched_outputs = [] - output_lens = [] + # get all proj + query_layer = self.query(hidden_states).view(*hidden_shape).transpose(1, 2) + key_layer = self.key(hidden_states).view(*hidden_shape).transpose(1, 2) + value_layer = self.value(hidden_states).view(*hidden_shape).transpose(1, 2) - for i in range(len(inputs)): - le = input_lens[i] - channel = channels[i] - seq = inputs[i, :le, :channel] # [seq_len, channel] - ps = patch_sizes[i].item() - st = strides[i].item() + if past_key_values is not None: + # decoder-only intern_s2_preview_bert can have a simple dynamic cache for example + current_past_key_values = past_key_values + if isinstance(past_key_values, EncoderDecoderCache): + current_past_key_values = past_key_values.self_attention_cache + + # save all key/value_layer to cache to be re-used for fast auto-regressive generation + key_layer, value_layer = current_past_key_values.update( + key_layer, + value_layer, + self.layer_idx, + {"cache_position": cache_position}, + ) - output_len = torch.ceil((le - ps) / st) + 1 - pad_len = ((output_len - 1) * st + ps - le).long().item() - if seq.ndim == 1: - seq = seq.unsqueeze(-1) - seq = nn.functional.pad(seq, (0, 0, 0, pad_len), "constant", 0) - assert output_len > 0, (seq.shape, ps, st, le, output_len) - output_lens.append(output_len) - indices = (torch.arange(0, output_len * st, st).unsqueeze(1) + torch.arange(ps)).long() - patched = seq[indices] + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) - output = self.forward_encoder(patched) # [num_patch, D] - patched_outputs.append(output) + attn_output, attn_weights = attention_interface( + self, + query_layer, + key_layer, + value_layer, + attention_mask, + dropout=0.0 if not self.training else self.dropout.p, + scaling=self.scaling, + **kwargs, + ) + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + return attn_output, attn_weights - outputs = nn.utils.rnn.pad_sequence(patched_outputs, batch_first=True) - output_lens = torch.tensor(output_lens).squeeze().to(outputs.device).long() - if output_lens.ndim == 0: - output_lens = output_lens.unsqueeze(0) - outputs, output_lens = self.subsampling(outputs.clone(), output_lens.clone()) - return outputs, output_lens +class InternS2PreviewBertCrossAttention(nn.Module): + def __init__(self, config, is_causal=False, layer_idx=None): + super().__init__() + if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"): + raise ValueError( + f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention " + f"heads ({config.num_attention_heads})" + ) + self.config = config - def forward_encoder(self, x: torch.Tensor) -> torch.Tensor: - num_patch, patch_len, C = x.shape - # conv1 - x = x.reshape(num_patch * C, 1, patch_len) # each channel is treated as an independent sample into conv1 - x = nn.functional.relu(self.conv(x)) # [B*C, D1, L] - x = x.permute(2, 0, 1) # [L, B*C, D1] + self.num_attention_heads = config.num_attention_heads + self.attention_head_size = int(config.hidden_size / config.num_attention_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + self.scaling = self.attention_head_size**-0.5 - x = self.pos_encoder(x) # [L, B*C, D1] - x = self.transformer_encoder(x.to(torch.bfloat16)) - x = x.mean(0) + self.query = nn.Linear(config.hidden_size, self.all_head_size) + self.key = nn.Linear(config.hidden_size, self.all_head_size) + self.value = nn.Linear(config.hidden_size, self.all_head_size) - x = x.reshape(num_patch, C, -1) + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) - return x.mean(1) + self.is_causal = is_causal + self.layer_idx = layer_idx + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.FloatTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + past_key_values: EncoderDecoderCache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor]: + # determine input shapes + bsz, tgt_len = hidden_states.shape[:-1] + src_len = encoder_hidden_states.shape[1] -class InternS2PreviewTimeSeriesProjector(nn.Module): - def __init__(self, config: InternS2PreviewTimeSeriesConfig): - super().__init__() - self.layer_norm = nn.LayerNorm(config.ts_hidden_dim) - self.linear_1 = nn.Linear(config.ts_hidden_dim, config.out_hidden_size) - self.act = ACT2FN[config.activation_function] - self.linear_2 = nn.Linear(config.out_hidden_size, config.out_hidden_size) + q_input_shape = (bsz, tgt_len, -1, self.attention_head_size) + kv_input_shape = (bsz, src_len, -1, self.attention_head_size) - def forward(self, ts_features): - hidden_states = self.layer_norm(ts_features) - hidden_states = self.linear_1(hidden_states) - hidden_states = self.act(hidden_states) - hidden_states = self.linear_2(hidden_states) - return hidden_states + # get query proj + query_layer = self.query(hidden_states).view(*q_input_shape).transpose(1, 2) + is_updated = past_key_values.is_updated.get(self.layer_idx) if past_key_values is not None else False + if past_key_values is not None and is_updated: + # reuse k,v, cross_attentions + key_layer = past_key_values.cross_attention_cache.layers[self.layer_idx].keys + value_layer = past_key_values.cross_attention_cache.layers[self.layer_idx].values + else: + key_layer = self.key(encoder_hidden_states).view(*kv_input_shape).transpose(1, 2) + value_layer = self.value(encoder_hidden_states).view(*kv_input_shape).transpose(1, 2) -@auto_docstring -class InternS2PreviewTimeSeriesModel(PreTrainedModel): - main_input_name = "time_series_signals" - _supports_flash_attn_2 = False - config_class = InternS2PreviewTimeSeriesConfig - _no_split_modules = ["InternS2PreviewTimeSeriesEncoderLayer"] + if past_key_values is not None: + # save all states to the cache + key_layer, value_layer = past_key_values.cross_attention_cache.update( + key_layer, value_layer, self.layer_idx + ) + # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls + past_key_values.is_updated[self.layer_idx] = True - def __init__(self, config: InternS2PreviewTimeSeriesConfig): - super().__init__(config) - self.config = config - self.encoder_embed = InternS2PreviewTimeSeriesMultiChannelAdaptiveSubsampling() - self.encoder = InternS2PreviewTimeSeriesEncoder(config) - self.projector = InternS2PreviewTimeSeriesProjector(config) + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) - def get_input_embeddings(self): - return self.encoder_embed + attn_output, attn_weights = attention_interface( + self, + query_layer, + key_layer, + value_layer, + attention_mask, + dropout=0.0 if not self.training else self.dropout.p, + scaling=self.scaling, + **kwargs, + ) + attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous() + return attn_output, attn_weights - def make_pad_mask(self, lengths: torch.Tensor) -> torch.Tensor: - r""" - lengths (`torch.Tensor` of shape `(batch_size,)`): - A 1-D tensor containing sentence lengths. - Returns: - A 2-D bool tensor, where masked positions are filled with `True` and non-masked positions are filled with `False`. - Example: - ```python - >>> lengths = torch.tensor([1, 3, 2, 5]) - >>> self.make_pad_mask(lengths) - tensor([[False, True, True, True, True], - [False, False, False, True, True], - [False, False, True, True, True], - [False, False, False, False, False]]) - ``` - """ - assert lengths.ndim == 1, lengths.ndim - max_len = lengths.max() - n = lengths.size(0) - seq_range = torch.arange(0, max_len, device=lengths.device) - expaned_lengths = seq_range.unsqueeze(0).expand(n, max_len) - return expaned_lengths >= lengths.unsqueeze(-1) +class InternS2PreviewBertSelfOutput(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.LayerNorm(hidden_states + input_tensor) + return hidden_states + + +class InternS2PreviewBertAttention(nn.Module): + def __init__(self, config, is_causal=False, layer_idx=None, is_cross_attention=False): + super().__init__() + self.is_cross_attention = is_cross_attention + attention_class = InternS2PreviewBertCrossAttention if is_cross_attention else InternS2PreviewBertSelfAttention + self.self = attention_class(config, is_causal=is_causal, layer_idx=layer_idx) + self.output = InternS2PreviewBertSelfOutput(config) def forward( self, - time_series_signals: torch.FloatTensor = None, - ts_lens: torch.Tensor = None, - sr: torch.Tensor = None, - channels: torch.LongTensor | None = None, - output_hidden_states: bool = None, - return_dict: bool = None, - time_series_embeds: torch.FloatTensor = None, - ): - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + hidden_states: torch.Tensor, + attention_mask: torch.FloatTensor | None = None, + encoder_hidden_states: torch.FloatTensor | None = None, + encoder_attention_mask: torch.FloatTensor | None = None, + past_key_values: Cache | None = None, + cache_position: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor]: + attention_mask = attention_mask if not self.is_cross_attention else encoder_attention_mask + attention_output, attn_weights = self.self( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + cache_position=cache_position, + **kwargs, ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict + attention_output = self.output(attention_output, hidden_states) + return attention_output, attn_weights - if time_series_signals is None and time_series_embeds is None: - raise ValueError("You have to specify time_series_signals or time_series_embeds") - if ( - time_series_embeds is not None - and len(time_series_embeds.shape) == 3 - and time_series_embeds.shape[-1] == self.config.ts_adapt_in_dim - ): - time_series_embeds = time_series_embeds +class InternS2PreviewBertIntermediate(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.intermediate_size) + if isinstance(config.hidden_act, str): + self.intermediate_act_fn = ACT2FN[config.hidden_act] else: - if time_series_signals.ndim == 3: - time_series_embeds, ts_lens = self.encoder_embed(time_series_signals, ts_lens, sr, channels) - else: - raise ValueError(f"wrong time_series_signals size: {time_series_signals[0].shape}") + self.intermediate_act_fn = config.hidden_act - # [B, 64000, 1] -> [B, 200, 256] -> [B, 100, 1024] - encoder_outputs = self.encoder( - input_features=time_series_embeds, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.intermediate_act_fn(hidden_states) + return hidden_states - # ts_lens after encoder - ts_lens = (ts_lens + 1) // 2 - assert torch.all(ts_lens > 0), f"The length of time_series_embeds is so small. ts_lens: {ts_lens}" - src_key_padding_mask = self.make_pad_mask(ts_lens) - last_hidden_state = encoder_outputs.last_hidden_state +class InternS2PreviewBertOutput(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.intermediate_size, config.hidden_size) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.LayerNorm(hidden_states + input_tensor) + return hidden_states - ts_pad_mask = src_key_padding_mask - ts_embeds = self.projector(last_hidden_state) - return ts_embeds, ts_pad_mask +class InternS2PreviewBertLayer(GradientCheckpointingLayer): + def __init__(self, config, layer_idx=None): + super().__init__() + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = InternS2PreviewBertAttention(config, is_causal=config.is_decoder, layer_idx=layer_idx) + self.is_decoder = config.is_decoder + self.add_cross_attention = config.add_cross_attention + self.intermediate = InternS2PreviewBertIntermediate(config) + self.output = InternS2PreviewBertOutput(config) + if self.add_cross_attention: + if not self.is_decoder: + raise ValueError(f"{self} should be used as a decoder model if cross attention is added") + self.crossattention = InternS2PreviewBertAttention( + config, + is_causal=False, + layer_idx=layer_idx, + is_cross_attention=True, + ) + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.FloatTensor | None = None, + encoder_hidden_states: torch.FloatTensor | None = None, + encoder_attention_mask: torch.FloatTensor | None = None, + past_key_values: Cache | None = None, + cache_position: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor]: + self_attention_output, _ = self.attention( + hidden_states, + attention_mask, + past_key_values=past_key_values, + cache_position=cache_position, + **kwargs, + ) + attention_output = self_attention_output -@dataclass -@auto_docstring( - custom_intro=""" - Base class for Llava outputs, with hidden states and attentions. - """ -) -class InternS2PreviewModelOutputWithPast(ModelOutput): - r""" - past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): - It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). + if self.is_decoder and encoder_hidden_states is not None: + if not hasattr(self, "crossattention"): + raise ValueError( + f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers" + " by setting `config.add_cross_attention=True`" + ) - Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see - `past_key_values` input) to speed up sequential decoding. - rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): - The rope index difference between sequence length and multimodal rope. - """ + cross_attention_output, _ = self.crossattention( + self_attention_output, + None, # attention_mask + encoder_hidden_states, + encoder_attention_mask, + past_key_values=past_key_values, + **kwargs, + ) + attention_output = cross_attention_output - last_hidden_state: torch.FloatTensor | None = None - past_key_values: Cache | None = None - hidden_states: tuple[torch.FloatTensor] | None = None - attentions: tuple[torch.FloatTensor] | None = None - rope_deltas: torch.LongTensor | None = None - router_logits: tuple[torch.FloatTensor] | None = None + layer_output = apply_chunking_to_forward( + self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output + ) + return layer_output + def feed_forward_chunk(self, attention_output): + intermediate_output = self.intermediate(attention_output) + layer_output = self.output(intermediate_output, attention_output) + return layer_output -@dataclass -@auto_docstring( - custom_intro=""" - Base class for InternS2Preview causal language model (or autoregressive) outputs. - """ -) -class InternS2PreviewCausalLMOutputWithPast(ModelOutput): - r""" - loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): - Language modeling loss (for next-token prediction). - logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): - Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). - past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): - It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). - - Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see - `past_key_values` input) to speed up sequential decoding. - rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): - The rope index difference between sequence length and multimodal rope. - """ - loss: torch.FloatTensor | None = None - logits: torch.FloatTensor | None = None - past_key_values: Cache | None = None - hidden_states: tuple[torch.FloatTensor] | None = None - attentions: tuple[torch.FloatTensor] | None = None - rope_deltas: torch.LongTensor | None = None - router_logits: tuple[torch.FloatTensor] | None = None - aux_loss: torch.FloatTensor | None = None +class InternS2PreviewBertEmbeddings(nn.Module): + """Construct the embeddings from word, position and token_type embeddings.""" + def __init__(self, config): + super().__init__() + self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id) + self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size) + self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) + + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + self.register_buffer( + "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False + ) -@auto_docstring -class InternS2PreviewModel(InternS2PreviewPreTrainedModel): - base_model_prefix = "model" - _checkpoint_conversion_mapping = {} - # Reference: fix gemma3 grad acc #37208 - accepts_loss_kwargs = False + def forward( + self, + input_ids: torch.LongTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + past_key_values_length: int = 0, + ) -> torch.Tensor: + if input_ids is not None: + input_shape = input_ids.size() + else: + input_shape = inputs_embeds.size()[:-1] - config: InternS2PreviewConfig - _no_split_modules = [ - "Qwen3_5MoeTextDecoderLayer", - "Qwen3_5MoeVisionBlock", - "InternS2PreviewTimeSeriesEncoderLayer", - ] + batch_size, seq_length = input_shape - def __init__(self, config): - super().__init__(config) - self.visual = InternS2PreviewVisionModel._from_config(config.vision_config) - self.language_model = InternS2PreviewTextModel._from_config(config.text_config) - self.rope_deltas = None # cache rope_deltas here - self.time_series = InternS2PreviewTimeSeriesModel._from_config(config.ts_config) + if position_ids is None: + position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length] + + # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs + # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves + # issue #5664 + if token_type_ids is None: + if hasattr(self, "token_type_ids"): + # NOTE: We assume either pos ids to have bsz == 1 (broadcastable) or bsz == effective bsz (input_shape[0]) + buffered_token_type_ids = self.token_type_ids.expand(position_ids.shape[0], -1) + buffered_token_type_ids = torch.gather(buffered_token_type_ids, dim=1, index=position_ids) + token_type_ids = buffered_token_type_ids.expand(batch_size, seq_length) + else: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device) - # Initialize weights and apply final processing - self.post_init() + if inputs_embeds is None: + inputs_embeds = self.word_embeddings(input_ids) + token_type_embeddings = self.token_type_embeddings(token_type_ids) + embeddings = inputs_embeds + token_type_embeddings - def get_input_embeddings(self): - return self.language_model.get_input_embeddings() + position_embeddings = self.position_embeddings(position_ids) + embeddings = embeddings + position_embeddings - def set_input_embeddings(self, value): - self.language_model.set_input_embeddings(value) + embeddings = self.LayerNorm(embeddings) + embeddings = self.dropout(embeddings) + return embeddings + + +class InternS2PreviewBertEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.layer = nn.ModuleList( + [InternS2PreviewBertLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)] + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.FloatTensor | None = None, + encoder_hidden_states: torch.FloatTensor | None = None, + encoder_attention_mask: torch.FloatTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = None, + cache_position: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor] | BaseModelOutputWithPastAndCrossAttentions: + for i, layer_module in enumerate(self.layer): + hidden_states = layer_module( + hidden_states, + attention_mask, + encoder_hidden_states, # as a positional argument for gradient checkpointing + encoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + cache_position=cache_position, + **kwargs, + ) + + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + past_key_values=past_key_values if use_cache else None, + ) + + +class InternS2PreviewBertPooler(nn.Module): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.activation = nn.Tanh() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # We "pool" the model by simply taking the hidden state corresponding + # to the first token. + first_token_tensor = hidden_states[:, 0] + pooled_output = self.dense(first_token_tensor) + pooled_output = self.activation(pooled_output) + return pooled_output + + +@auto_docstring +class InternS2PreviewBertPreTrainedModel(PreTrainedModel): + config_class = InternS2PreviewBertConfig + base_model_prefix = "interns2_preview_bert" + supports_gradient_checkpointing = True + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": InternS2PreviewBertLayer, + "attentions": InternS2PreviewBertSelfAttention, + "cross_attentions": InternS2PreviewBertCrossAttention, + } + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + super()._init_weights(module) + if isinstance(module, InternS2PreviewBertEmbeddings): + init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1))) + init.zeros_(module.token_type_ids) + + +class InternS2PreviewBertModel(InternS2PreviewBertPreTrainedModel): + _no_split_modules = ["InternS2PreviewBertEmbeddings", "InternS2PreviewBertLayer"] + + def __init__(self, config, add_pooling_layer=True): + r""" + add_pooling_layer (bool, *optional*, defaults to `True`): + Whether to add a pooling layer + """ + super().__init__(config) + self.config = config + self.gradient_checkpointing = False + + self.embeddings = InternS2PreviewBertEmbeddings(config) + self.encoder = InternS2PreviewBertEncoder(config) + + self.pooler = InternS2PreviewBertPooler(config) if add_pooling_layer else None + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embeddings.word_embeddings + + def set_input_embeddings(self, value): + self.embeddings.word_embeddings = value + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + token_type_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = None, + cache_position: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor] | BaseModelOutputWithPoolingAndCrossAttentions: + if self.config.is_decoder: + use_cache = use_cache if use_cache is not None else self.config.use_cache + else: + use_cache = False + + if use_cache and past_key_values is None: + past_key_values = ( + EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config)) + if encoder_hidden_states is not None or self.config.is_encoder_decoder + else DynamicCache(config=self.config) + ) + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if input_ids is not None: + device = input_ids.device + seq_length = input_ids.shape[1] + else: + device = inputs_embeds.device + seq_length = inputs_embeds.shape[1] + + past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 + if cache_position is None: + cache_position = torch.arange(past_key_values_length, past_key_values_length + seq_length, device=device) + + embedding_output = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + past_key_values_length=past_key_values_length, + ) + + attention_mask, encoder_attention_mask = self._create_attention_masks( + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + embedding_output=embedding_output, + encoder_hidden_states=encoder_hidden_states, + cache_position=cache_position, + past_key_values=past_key_values, + ) + + encoder_outputs = self.encoder( + embedding_output, + attention_mask=attention_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + cache_position=cache_position, + position_ids=position_ids, + **kwargs, + ) + sequence_output = encoder_outputs.last_hidden_state + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None + + return BaseModelOutputWithPoolingAndCrossAttentions( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + past_key_values=encoder_outputs.past_key_values, + ) + + def _create_attention_masks( + self, + attention_mask, + encoder_attention_mask, + embedding_output, + encoder_hidden_states, + cache_position, + past_key_values, + ): + if self.config.is_decoder: + attention_mask = create_causal_mask( + config=self.config, + inputs_embeds=embedding_output, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + ) + else: + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=embedding_output, + attention_mask=attention_mask, + ) + + if encoder_attention_mask is not None: + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=embedding_output, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, + ) + + return attention_mask, encoder_attention_mask + + +class MRQFormer(nn.Module): + def __init__( + self, + num_query_list, + hidden_size: int = 1024, + num_layers: int = 6, + cross_attention_freq: int = 2, + encoder_hidden_size: int = None, + output_size: int = None, + ): + super().__init__() + + self.num_query_list = num_query_list + self.hidden_size = hidden_size + + self.query_token_list = nn.ParameterList() + for num_queries in self.num_query_list: + self.query_token_list.append(nn.Parameter(torch.randn(1, num_queries, hidden_size))) + + bert_config = InternS2PreviewBertConfig( + vocab_size=1, + hidden_size=hidden_size, + num_hidden_layers=num_layers, + num_attention_heads=hidden_size // 16, + intermediate_size=hidden_size * 4, + max_position_embeddings=max(self.num_query_list), + add_cross_attention=True, + is_decoder=True, + cross_attention_freq=cross_attention_freq, + ) + self.transformer = InternS2PreviewBertModel(bert_config) + + if encoder_hidden_size is not None and encoder_hidden_size != hidden_size: + self.encoder_proj = nn.Linear(encoder_hidden_size, hidden_size) + else: + self.encoder_proj = nn.Identity() + + if output_size == None: + output_size = hidden_size + + self.proj_out = nn.Sequential( + nn.Linear(in_features=hidden_size, out_features=4 * hidden_size), + nn.ReLU(), + nn.Linear(in_features=4 * hidden_size, out_features=output_size), + ) + + def forward(self, encoder_features: torch.Tensor, attention_mask=None, res_idx=0): + B = encoder_features.shape[0] + + encoder_features = self.encoder_proj(encoder_features) + + query_tokens = self.query_token_list[res_idx] + query_tokens = query_tokens.expand(B, -1, -1) + + encoder_attention_mask = None + if attention_mask is not None: + encoder_attention_mask = attention_mask + + outputs = self.transformer( + inputs_embeds=query_tokens, + encoder_hidden_states=encoder_features, + encoder_attention_mask=encoder_attention_mask, + return_dict=True, + ) + + outputs = outputs.last_hidden_state + + outputs = self.proj_out(outputs) + + return outputs + + +class TotalFixlenSingleResChunkQformerSubsampling(nn.Module): + def __init__(self, config: InternS2PreviewTimeSeriesConfig): + super().__init__() + + self.selfatt = config.subsampling_selfatt + self.patch_list = [config.subsampling_patch] + self.num_query_list = [config.subsampling_num_query] + + self.num_conv_layers = config.subsampling_num_conv_layers + hidden_dim = config.subsampling_hidden_dim + if self.num_conv_layers == 0: + self.conv = nn.Conv1d(in_channels=1, out_channels=hidden_dim - 2, kernel_size=5, stride=1, padding=2) + self.mask_pool = None + else: + self.conv_layers = [] + self.pool_layers = [] + current_channels = 1 + output_channels = hidden_dim - 2 + + for i in range(self.num_conv_layers): + if i == 0: + in_c = current_channels + out_c = output_channels + else: + in_c = output_channels + out_c = output_channels + + self.conv_layers.append(nn.Conv1d(in_c, out_c, kernel_size=5, stride=2, padding=2)) + self.conv_layers.append(nn.ReLU(inplace=True)) + self.pool_layers.append(nn.MaxPool1d(kernel_size=5, stride=2, padding=2)) + + self.conv = nn.Sequential(*self.conv_layers) + self.mask_pool = nn.Sequential(*self.pool_layers) + + if self.selfatt: + num_encoder_layers = 1 + encoder_layers = nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=config.subsampling_nhead) + self.transformer_encoder = nn.TransformerEncoder(encoder_layers, num_encoder_layers) + self.mrqformer = None + else: + self.mrqformer = MRQFormer( + num_query_list=self.num_query_list, + hidden_size=hidden_dim, + num_layers=6, + cross_attention_freq=2, + output_size=hidden_dim, + ) + self.transformer_encoder = None + + self.pos_encoder = InternS2PreviewTimeSeriesFixPositionalEncoding(d_model=hidden_dim) + + num_channel_encoder_layers = 1 + channel_encoder_layers = nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=32) + self.channel_encoder = nn.TransformerEncoder(channel_encoder_layers, num_channel_encoder_layers) + self.fuze_proj = nn.Linear(in_features=hidden_dim, out_features=256) + + def forward(self, inputs, mask=None, subrate=None): + if mask is None: + mask = torch.ones(inputs.shape).to(inputs.device) + features, feature_lens, learnable_lens, strides, raw_outputs = self.forward_patch(inputs, mask, subrate) + + outputs = features + output_lens = feature_lens + + return outputs, output_lens, learnable_lens, strides, raw_outputs + + def forward_patch(self, inputs, mask=None, subrate=None): + + pred_strides = None + + patched_outputs = [] + raw_outputs = [] + + with torch.set_grad_enabled(self.training): + seq = inputs # (B, L, C) + + masked_seq = seq.masked_fill(mask == 0, float("nan")) + + mean = seq.mean(dim=1, keepdim=True) + std = seq.std(dim=1, keepdim=True) + + mean_expand = mean.expand(seq.size()) + abs_mean_plus_1 = torch.abs(mean) + 1.0 # (B, 1, C), 所有值 ≥ 1 + log_abs_plus_1 = torch.log(abs_mean_plus_1) # (B, 1, C), ≥ 0 + sign_mean = torch.sign(mean) # (B, 1, C), 取值 -1, 0, +1 + mean_feat = sign_mean * log_abs_plus_1 # (B, 1, C), 保留原符号的 log 值 + + std_expand = std.expand(seq.size()) + std_feat = torch.log(std + 1e-7) + + seq = (seq - mean) / (std + 1e-7) + + P_max = torch.tensor([max(self.patch_list)]).to(seq.device) + Q_max = torch.tensor([max(self.num_query_list)]).to(seq.device) + if subrate is None: + subrate = P_max / Q_max + st = subrate * Q_max + ps = torch.ceil(st) + le = torch.tensor(inputs.shape[1]).to(seq.device) + + output_len = torch.ceil((le - ps) / st + 1) + pad_len = (torch.ceil((output_len - 1) * st + ps - le)).long().item() + + if seq.ndim == 2: + seq = seq.unsqueeze(-1) + if not isinstance(pad_len, int): + print(f"⚠️ [DEBUG] pad_len is NOT an int! value={pad_len}, type={type(pad_len)}") + seq = torch.nn.functional.pad(seq, (0, 0, 0, pad_len, 0, 0), "constant", 0) + + padded_mask = torch.nn.functional.pad(mask, (0, 0, 0, pad_len, 0, 0), "constant", 0) # (B, L, C) + + B, L, C = seq.shape + seq = seq.permute(0, 2, 1) # (B, C, L) + seq = seq.reshape(B * C, 1, L) # (B*C, 1, L) + if self.num_conv_layers == 0: + seq = nn.functional.relu(self.conv(seq)) # (B*C, D, L) + else: + seq = self.conv(seq) # (B*C, D, L) + L = seq.shape[-1] + seq = seq.reshape(B, C, -1, L) + seq = seq.permute(0, 3, 1, 2) # (B, L, C, D) + + mean_feat = mean_feat.expand([B, L, C]).unsqueeze(-1) + std_feat = std_feat.expand([B, L, C]).unsqueeze(-1) + + seq = torch.cat([seq, mean_feat, std_feat], dim=-1) # (B, L, C, D) + + D = seq.shape[-1] + + if self.mask_pool is not None: + padded_mask = padded_mask.permute(0, 2, 1) + padded_mask = self.mask_pool(padded_mask) + padded_mask = padded_mask.permute(0, 2, 1) + + idx = torch.arange(L).to(seq.device) + + output_list = [] + for res_idx, raw_patch in enumerate(self.patch_list): + num_query = self.num_query_list[res_idx] + stride = subrate * num_query + patch = torch.ceil(stride) + patch_output_len = torch.floor(output_len * st / (stride * (2**self.num_conv_layers))) + + indices = ( + torch.floor(torch.arange(0, patch_output_len.item()).to(seq.device) * stride).unsqueeze(1) + + torch.arange(torch.ceil(patch)).to(seq.device) + ).long() + patched = seq[:, indices, :, :] # (B, N, P, C, D) + _, N, P, _, _ = patched.shape + patched = patched.permute(1, 2, 0, 3, 4) + patched = patched.reshape(N, P, -1) # (N, P, B*C*D) + patched = patched.reshape(N, P, B, C, D) + + patched = patched.permute(2, 0, 1, 3, 4) # (B, N, P, C, D) + patched = patched.reshape(B * N, P, C, D) + + patched_mask = padded_mask[:, indices, :] # (B, N, P, C) + patched_mask = patched_mask.reshape(B * N, P, C) + + raw_output = self.forward_encoder(patched, patched_mask, res_idx) + if self.selfatt: + output = raw_output.mean(dim=1) # B*N, D + else: + output = raw_output.reshape(-1, D) # B*N*Q, D + output_list.append(output) + + outputs = torch.cat(output_list, dim=1) + outputs = self.fuze_proj(outputs) + _, D_out = outputs.shape + + outputs = outputs.reshape(B, -1, D_out) + raw_outputs = None + + if self.selfatt: + output_lens = torch.ceil((padded_mask[:, :, 0].sum(dim=-1) - ps) / st + 1) + else: + output_lens = torch.ceil((padded_mask[:, :, 0].sum(dim=-1) - ps) / st + 1) * Q_max + + learnable_lens = None + + return outputs, output_lens, learnable_lens, pred_strides, raw_outputs + + def forward_encoder(self, inputs, mask, res_idx=0): + + x = inputs # (B*N, P, C, D) + num_patch, patch_len, C, D = x.shape + x = x.permute(1, 0, 2, 3) # (P, B*N, C, D) + x = x.reshape(patch_len, num_patch * C, -1) # [P, B*N*C, D] + + x = self.pos_encoder(x) # [P, B*C, D] + + mask = mask.permute(0, 2, 1) # (B*N, C, P) + mask = mask.reshape(num_patch * C, patch_len) # (B*N*C, P) + bool_mask = mask == 0 # (B*N*C, P) + + chunk_size = 4096 * 8 + all_outputs = [] + + for i in range(0, x.size(1), chunk_size): + x_chunk = x[:, i : i + chunk_size, :].contiguous() # [P, chunk_len, D] + chunk_mask = mask[i : i + chunk_size] # (chunk_len, P) + chunk_bool_mask = bool_mask[i : i + chunk_size] # (chunk_len, P) + if self.selfatt: + x_chunk = self.transformer_encoder(x_chunk, src_key_padding_mask=chunk_bool_mask) # [P, chunk_len, D] + else: + x_chunk = x_chunk.permute(1, 0, 2) + x_chunk = self.mrqformer(x_chunk, attention_mask=chunk_mask, res_idx=res_idx) + x_chunk = x_chunk.permute(1, 0, 2) # [Q, chunk_len, D] + all_outputs.append(x_chunk) + + x = torch.cat(all_outputs, dim=1) # [P, B*N*C, D] / [Q, B*N*C, D] + + x = x.reshape(-1, num_patch, C, D) # [P, B*N, C, D] /# [Q, B*N, C, D] + x = x.permute(2, 1, 0, 3) # [C, B, P, D] + x = x.reshape(C, -1, D) # [C, B*N*P, D] + + all_outputs = [] + + for i in range(0, x.size(1), chunk_size): + x_chunk = x[:, i : i + chunk_size, :].contiguous() # [C, chunk_len, D] + x_chunk = self.channel_encoder(x_chunk) + x_chunk = x_chunk.mean(0) # [chunk_len, D] + all_outputs.append(x_chunk) + + x = torch.cat(all_outputs, dim=0) # [B*N*P, D] + x = x.reshape(num_patch, -1, D) # [B*N, P, D] + + return x + + +class InternS2PreviewTimeSeriesProjector(nn.Module): + def __init__(self, config: InternS2PreviewTimeSeriesConfig): + super().__init__() + self.layer_norm = nn.LayerNorm(config.ts_hidden_dim) + self.linear_1 = nn.Linear(config.ts_hidden_dim, config.out_hidden_size) + self.act = ACT2FN[config.activation_function] + self.linear_2 = nn.Linear(config.out_hidden_size, config.out_hidden_size) + + def forward(self, ts_features): + hidden_states = self.layer_norm(ts_features) + hidden_states = self.linear_1(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + + +@auto_docstring +class InternS2PreviewTimeSeriesModel(PreTrainedModel): + main_input_name = "time_series_signals" + _supports_flash_attn_2 = False + config_class = InternS2PreviewTimeSeriesConfig + _no_split_modules = ["InternS2PreviewTimeSeriesEncoderLayer"] + + def __init__(self, config: InternS2PreviewTimeSeriesConfig): + super().__init__(config) + self.config = config + self.encoder_embed = TotalFixlenSingleResChunkQformerSubsampling(config) + self.encoder = InternS2PreviewTimeSeriesEncoder(config) + self.projector = InternS2PreviewTimeSeriesProjector(config) + + def get_input_embeddings(self): + return self.encoder_embed + + def make_pad_mask(self, lengths: torch.Tensor) -> torch.Tensor: + r""" + lengths (`torch.Tensor` of shape `(batch_size,)`): + A 1-D tensor containing sentence lengths. + + Returns: + A 2-D bool tensor, where masked positions are filled with `True` and non-masked positions are filled with `False`. + Example: + ```python + >>> lengths = torch.tensor([1, 3, 2, 5]) + >>> self.make_pad_mask(lengths) + tensor([[False, True, True, True, True], + [False, False, False, True, True], + [False, False, True, True, True], + [False, False, False, False, False]]) + ``` + """ + assert lengths.ndim == 1, lengths.ndim + max_len = lengths.max() + n = lengths.size(0) + seq_range = torch.arange(0, max_len, device=lengths.device) + expaned_lengths = seq_range.unsqueeze(0).expand(n, max_len) + return expaned_lengths >= lengths.unsqueeze(-1) + + def concat_chunks_simple( + self, + chunks: torch.Tensor, # (N, T_chunk, C) + chunk_lengths: torch.Tensor = None, # (N,) 每个 chunk 的有效长度 + ) -> torch.Tensor: + """ + 简单首尾相接（无 overlap） + """ + N, T_chunk, C = chunks.shape + + if chunk_lengths is None: + signal = chunks.reshape(-1, C) # (N*T_chunk, C) + signal_len = N * T_chunk + else: + signal_list = [] + signal_len = 0 + for i in range(N): + valid_len = int(chunk_lengths[i].item()) + signal_list.append(chunks[i, :valid_len, :]) # (T_i, C) + signal_len = signal_len + valid_len + signal = torch.cat(signal_list, dim=0) # (sum(T_i), C) + + return signal, signal_len + + def chunk_tensor(self, tensor: torch.Tensor, return_indices: bool = False): + + T = tensor.shape[0] + chunk_size = self.config.chunk_size + step = self.config.chunk_step + + chunks = [] + indices = [] if return_indices else None + masks = [] + + if chunk_size > 0: + start = 0 + while start < T: + end = min(start + chunk_size, T) + chunk = tensor[start:end, :] + chunks.append(chunk) + mask = torch.zeros(chunk.shape) + mask[: end - start, :] = 1 + masks.append(mask) + if return_indices: + indices.append((start, end)) + + if end >= T: + break + + start += step + + output = nn.utils.rnn.pad_sequence(chunks, batch_first=True) + output_mask = nn.utils.rnn.pad_sequence(masks, batch_first=True).to(output.device) + else: + output = tensor.unsqueeze(0) + output_mask = torch.ones(output.shape).to(tensor.device) + + if return_indices: + return output, indices, output_mask + return output, output_mask + + def forward( + self, + time_series_signals: torch.FloatTensor = None, + ts_lens: torch.Tensor = None, + sr: torch.Tensor = None, + channels: torch.LongTensor | None = None, + output_hidden_states: bool = None, + return_dict: bool = None, + time_series_embeds: torch.FloatTensor = None, + ): + 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 time_series_signals is None and time_series_embeds is None: + raise ValueError("You have to specify time_series_signals or time_series_embeds") + if channels is None: + raise ValueError("channels must be specified") + + if isinstance(time_series_signals, list): + batch_iter = [time_series_signals[b][: ts_lens[b], : channels[b]] for b in range(len(time_series_signals))] + elif time_series_signals.ndim == 3: + batch_iter = [ + time_series_signals[b, : ts_lens[b], : channels[b]] for b in range(time_series_signals.shape[0]) + ] + else: + raise ValueError(f"wrong time_series_signals size: {time_series_signals[0].shape}") + + outputs = [] + output_lens = [] + output_chunks = [] + for b, seq in enumerate(batch_iter): + chunks, mask = self.chunk_tensor(seq) + + subrate = max(ts_lens[b] / 500, torch.tensor(1.0)) + chunk_output, chunk_lens, learnable_lens, strides, raw_feature = self.encoder_embed(chunks, mask, subrate) + + outputs.append(chunk_output) + output_lens.append(chunk_lens) + output_chunks.append(chunk_output.shape[0]) + x_lens = torch.cat(output_lens, dim=0).long() + max_len = x_lens.max().item() + + padded_list = [] + for t in outputs: + _, T, _ = t.shape # (B, T, D) + pad_T = int(max_len - T) + + padded = nn.functional.pad(t, (0, 0, 0, pad_T, 0, 0), mode="constant", value=0) + padded_list.append(padded) + x = torch.cat(padded_list, dim=0) + if x_lens.ndim == 0: + x_lens = x_lens.unsqueeze(-1) + + src_key_padding_mask = self.make_pad_mask(x_lens) + encoder_out = self.encoder(x, x_lens, src_key_padding_mask, output_hidden_states=True, return_dict=return_dict) + encoder_out = encoder_out.last_hidden_state + encoder_out_lens = (x_lens + 1) // 2 + + features = [] + feature_lens = [] + start = 0 + + for i, count in enumerate(output_chunks): + end = start + count + signal_chunks = encoder_out[start:end] # (N_i, T_chunk, C) + + feature, feature_len = self.concat_chunks_simple(signal_chunks, encoder_out_lens[start:end]) + + features.append(feature) + start = end + feature_lens.append(feature_len) + + encoder_out = nn.utils.rnn.pad_sequence(features, batch_first=True) + encoder_out_lens = torch.tensor(feature_lens).to(encoder_out.device) + + src_key_padding_mask = self.make_pad_mask(encoder_out_lens) + + ts_pad_mask = src_key_padding_mask + ts_embeds = self.projector(encoder_out) + + return ts_embeds, ts_pad_mask, encoder_out + + +class RMSNorm(nn.Module): + """RMS normalization.""" + + def __init__(self, num_features: int, *, epsilon: float = 1e-6): + super().__init__() + self.scale = nn.Parameter(torch.zeros(num_features)) + self.num_features = num_features + self.epsilon = epsilon + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + var = torch.mean(torch.square(inputs), dim=-1, keepdim=True) + normed_inputs = inputs * torch.rsqrt(var + self.epsilon) + normed_inputs = normed_inputs * self.scale + return normed_inputs + + +class ResidualBlock(nn.Module): + """Residual block with two linear layers and a linear residual connection.""" + + def __init__(self, input_dims, hidden_dims, output_dims, use_bias, activation="swish"): + super().__init__() + self.hidden_layer = nn.Linear( + in_features=input_dims, + out_features=hidden_dims, + bias=use_bias, + ) + self.output_layer = nn.Linear( + in_features=hidden_dims, + out_features=output_dims, + bias=use_bias, + ) + self.residual_layer = nn.Linear( + in_features=input_dims, + out_features=output_dims, + bias=use_bias, + ) + if activation == "relu": + self.activation = nn.ReLU() + elif activation == "swish": + self.activation = nn.SiLU() + elif activation == "none": + self.activation = nn.Identity() + else: + raise ValueError(f"Activation: {activation} not supported.") + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.output_layer(self.activation(self.hidden_layer(x))) + self.residual_layer(x) + + +@dataclass(frozen=False) +class DecodeCache: + """Cache for decoding.""" + + next_index: torch.Tensor + num_masked: torch.Tensor + key: torch.Tensor + value: torch.Tensor + kv_mask: torch.Tensor | None = None # (B, cache_size), True=padded/masked + + +class RotaryPositionalEmbedding(nn.Module): + """Rotary positional embedding.""" + + def __init__( + self, + embedding_dims: int, + min_timescale: float = 1.0, + max_timescale: float = 10000.0, + ): + super().__init__() + self.embedding_dims = embedding_dims + self.min_timescale = min_timescale + self.max_timescale = max_timescale + + def forward( + self, + inputs: torch.Tensor, + position: torch.Tensor | None = None, + ): + """Generates a JTensor of sinusoids with different frequencies.""" + if self.embedding_dims != inputs.shape[-1]: + raise ValueError( + "The embedding dims of the rotary position embeddingmust match the hidden dimension of the inputs." + ) + half_embedding_dim = self.embedding_dims // 2 + fraction = 2 * torch.arange(0, half_embedding_dim, device=inputs.device) / self.embedding_dims + timescale = (self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction).to(inputs.device) + if position is None: + seq_length = inputs.shape[1] + position = torch.arange(seq_length, dtype=torch.float32, device=inputs.device)[None, :] + + if len(inputs.shape) == 4: + position = position[..., None, None] + timescale = timescale[None, None, None, :] + elif len(inputs.shape) == 3: + position = position[..., None] + timescale = timescale[None, None, :] + else: + raise ValueError("Inputs must be of rank 3 or 4.") + + sinusoid_inp = position / timescale + sin = torch.sin(sinusoid_inp) + cos = torch.cos(sinusoid_inp) + first_half, second_half = torch.chunk(inputs, 2, dim=-1) + first_part = first_half * cos - second_half * sin + second_part = second_half * cos + first_half * sin + return torch.cat([first_part, second_part], dim=-1) + + +class PerDimScale(nn.Module): + """Per-dimension scaling.""" + + def __init__(self, num_dims: int): + super().__init__() + self.num_dims = num_dims + self.per_dim_scale = nn.Parameter(torch.zeros(num_dims)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + scale_factor = 1.442695041 / math.sqrt(self.num_dims) * nn.functional.softplus(self.per_dim_scale) + return x * scale_factor + + +def _torch_dot_product_attention(query, key, value, mask=None): + """Same (unscaled) attention as _dot_product_attention, fused kernel.""" + safe_mask = mask + fully_masked_rows = None + if mask is not None: + # F.scaled_dot_product_attention can emit NaNs/NaN-grads when a query row is + # fully masked. Keep the fused kernel by opening a single dummy key for + # those rows, then zero their outputs. + fully_masked_rows = ~mask.any(dim=-1, keepdim=True) + if fully_masked_rows.any(): + dummy_key_mask = torch.zeros_like(mask) + dummy_key_mask[..., 0] = True + safe_mask = mask | (fully_masked_rows & dummy_key_mask) + + attention_dtype = value.dtype + if query.dtype != attention_dtype: + query = query.to(attention_dtype) + if key.dtype != attention_dtype: + key = key.to(attention_dtype) + + # 1. Permute inputs from (B, L, H, D) to the expected (B, H, L, D) + query = query.permute(0, 2, 1, 3) + key = key.permute(0, 2, 1, 3) + value = value.permute(0, 2, 1, 3) + + # 2. Fused attention kernel with scale=1.0 (disable 1/sqrt(d_k) scaling). + output = nn.functional.scaled_dot_product_attention(query, key, value, attn_mask=safe_mask, scale=1.0) + + if fully_masked_rows is not None and fully_masked_rows.any(): + output = output.masked_fill(fully_masked_rows, 0.0) + + # 3. Permute back to (B, L, H, D) + output = output.permute(0, 2, 1, 3) + return output + + +class MultiHeadAttention(nn.Module): + """Multi-head attention.""" + + def __init__( + self, + num_heads: int, + in_features: int, + *, + use_per_dim_scale: bool = True, + use_rotary_position_embeddings: bool = True, + use_bias: bool = False, + attention_fn: Callable[..., torch.Tensor] = _torch_dot_product_attention, + qk_norm: str = "rms", + fuse_qkv: bool = False, + ): + super().__init__() + self.num_heads = num_heads + self.in_features = in_features + self.head_dim = in_features // num_heads + self.use_bias = use_bias + self.attention_fn = attention_fn + self.qk_norm = qk_norm + self.fuse_qkv = fuse_qkv + + if self.in_features % self.num_heads != 0: + raise ValueError( + f"Memory dimension ({self.in_features}) must be divisible by 'num_heads' heads ({self.num_heads})." + ) + + if self.fuse_qkv: + self.qkv_proj = nn.Linear(self.in_features, 3 * self.in_features, bias=use_bias) + else: + self.query = nn.Linear(self.in_features, self.in_features, bias=use_bias) + self.key = nn.Linear(self.in_features, self.in_features, bias=use_bias) + self.value = nn.Linear(self.in_features, self.in_features, bias=use_bias) + self.out = nn.Linear(self.in_features, self.in_features, bias=use_bias) + + if self.qk_norm == "rms": + self.query_ln = RMSNorm(self.head_dim) + self.key_ln = RMSNorm(self.head_dim) + else: + self.query_ln = nn.Identity() + self.key_ln = nn.Identity() + + self.use_rotary_position_embeddings = use_rotary_position_embeddings + if self.use_rotary_position_embeddings: + self.rotary_position_embedding = RotaryPositionalEmbedding( + embedding_dims=self.head_dim, + ) + + self.use_per_dim_scale = use_per_dim_scale + if use_per_dim_scale: + self.per_dim_scale = PerDimScale(num_dims=self.head_dim) + + def make_attn_mask( + self, + query_length: int, + num_all_masked_kv: torch.Tensor, + query_index_offset: torch.Tensor | None = None, + kv_length: int = 0, + ) -> torch.Tensor: + """Makes attention mask.""" + if kv_length == 0: + kv_length = query_length + + q_index = torch.arange(query_length, device=num_all_masked_kv.device)[None, None, :, None] + if query_index_offset is not None: + q_index = q_index + query_index_offset[:, None, None, None] + kv_index = torch.arange(kv_length, device=num_all_masked_kv.device)[None, None, None, :] + return torch.logical_and( + q_index >= kv_index, + kv_index >= num_all_masked_kv[:, None, None, None], + ) + + def forward( + self, + inputs_q: torch.Tensor, + *, + decode_cache: DecodeCache | None = None, + patch_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, DecodeCache | None]: + b, n_patches, _ = inputs_q.shape + if patch_mask is None: + patch_mask = torch.zeros(b, n_patches, dtype=torch.bool, device=inputs_q.device) + + if self.fuse_qkv: + qkv = self.qkv_proj(inputs_q) + query, key, value = torch.chunk(qkv, 3, dim=-1) + query = query.view(b, n_patches, self.num_heads, self.head_dim) + key = key.view(b, n_patches, self.num_heads, self.head_dim) + value = value.view(b, n_patches, self.num_heads, self.head_dim) + else: + query = self.query(inputs_q).view(b, n_patches, self.num_heads, self.head_dim) + key = self.key(inputs_q).view(b, n_patches, self.num_heads, self.head_dim) + value = self.value(inputs_q).view(b, n_patches, self.num_heads, self.head_dim) + + if decode_cache is None: + # Count only LEADING masked patches (not total). make_attn_mask assumes + # all masked positions are a contiguous left-padding block. + is_valid = ~patch_mask + has_valid = is_valid.any(dim=-1) + first_valid = torch.argmax(is_valid.to(torch.int32), dim=-1) + seq_len = torch.tensor(patch_mask.shape[-1], dtype=torch.int32, device=patch_mask.device) + num_masked = torch.where(has_valid, first_valid, seq_len) + next_index = torch.zeros_like(num_masked, dtype=torch.int32) + else: + is_valid = ~patch_mask + has_valid = is_valid.any(dim=-1) + first_valid = torch.argmax(is_valid.to(torch.int32), dim=-1) + seq_len = torch.tensor(patch_mask.shape[-1], dtype=torch.int32, device=patch_mask.device) + leading_masked = torch.where(has_valid, first_valid, seq_len) + num_masked = leading_masked + decode_cache.num_masked + next_index = decode_cache.next_index.clone() + + if self.use_rotary_position_embeddings: + position = ( + torch.arange(n_patches, device=inputs_q.device)[None, :] + next_index[:, None] - num_masked[:, None] + ) + query = self.rotary_position_embedding(query, position) + key = self.rotary_position_embedding(key, position) + + query = self.query_ln(query) + key = self.key_ln(key) + + if self.use_per_dim_scale: + query = self.per_dim_scale(query) + + if decode_cache is not None: + _, decode_cache_size, _, _ = decode_cache.value.shape + + start = decode_cache.next_index[0] + end = start + n_patches + + decode_cache.key[:, start:end] = key + decode_cache.value[:, start:end] = value + + if decode_cache.kv_mask is None: + decode_cache.kv_mask = torch.ones( + b, + decode_cache_size, + dtype=torch.bool, + device=patch_mask.device, + ) + decode_cache.kv_mask[:, start:end] = patch_mask + + key = decode_cache.key + value = decode_cache.value + decode_cache.next_index += n_patches + decode_cache.num_masked = num_masked + attn_mask = self.make_attn_mask( + query_length=n_patches, + num_all_masked_kv=num_masked, + query_index_offset=next_index, + kv_length=decode_cache_size, + ) + kv_valid = ~decode_cache.kv_mask # (B, decode_cache_size) + attn_mask = attn_mask & kv_valid[:, None, None, :] + else: + attn_mask = self.make_attn_mask(query_length=n_patches, num_all_masked_kv=num_masked) + kv_valid = ~patch_mask # (B, n_patches), True=valid + attn_mask = attn_mask & kv_valid[:, None, None, :] + + x = self.attention_fn(query, key, value, mask=attn_mask) + + x = x.reshape(b, n_patches, self.in_features) + out = self.out(x) + return out, decode_cache + + +class MultiHeadCrossAttention(nn.Module): + """Multi-head cross-attention from a query stream into a static KV context. + + Q comes from the transformer's residual stream; K/V come from an external + context tensor (e.g. Q-former output). No RoPE, no decode_cache: the KV + context is treated as a static set of tokens at every layer / decode step. + """ + + def __init__( + self, + num_heads: int, + in_features: int, + *, + kv_features: int | None = None, + use_bias: bool = False, + attention_fn: Callable[..., torch.Tensor] = _torch_dot_product_attention, + qk_norm: str = "rms", + ): + super().__init__() + self.num_heads = num_heads + self.in_features = in_features + self.kv_features = kv_features if kv_features is not None else in_features + self.head_dim = in_features // num_heads + self.attention_fn = attention_fn + + if self.in_features % self.num_heads != 0: + raise ValueError( + f"Memory dimension ({self.in_features}) must be divisible by 'num_heads' heads ({self.num_heads})." + ) + + self.q_proj = nn.Linear(self.in_features, self.in_features, bias=use_bias) + self.k_proj = nn.Linear(self.kv_features, self.in_features, bias=use_bias) + self.v_proj = nn.Linear(self.kv_features, self.in_features, bias=use_bias) + self.out = nn.Linear(self.in_features, self.in_features, bias=use_bias) + + if qk_norm == "rms": + self.query_ln = RMSNorm(self.head_dim) + self.key_ln = RMSNorm(self.head_dim) + else: + self.query_ln = nn.Identity() + self.key_ln = nn.Identity() + + self.per_dim_scale = PerDimScale(num_dims=self.head_dim) + + def forward( + self, + inputs_q: torch.Tensor, + *, + kv: torch.Tensor, + kv_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + b, n_q, _ = inputs_q.shape + n_kv = kv.shape[1] + + if kv.dtype != inputs_q.dtype: + kv = kv.to(dtype=inputs_q.dtype) + + query = self.q_proj(inputs_q).view(b, n_q, self.num_heads, self.head_dim) + key = self.k_proj(kv).view(b, n_kv, self.num_heads, self.head_dim) + value = self.v_proj(kv).view(b, n_kv, self.num_heads, self.head_dim) + + query = self.query_ln(query) + key = self.key_ln(key) + query = self.per_dim_scale(query) + + if kv_mask is None: + attn_mask = torch.ones( + b, + 1, + n_q, + n_kv, + dtype=torch.bool, + device=inputs_q.device, + ) + else: + kv_valid = ~kv_mask # True = valid + attn_mask = kv_valid[:, None, None, :].expand(b, 1, n_q, n_kv) + + x = self.attention_fn(query, key, value, mask=attn_mask) + x = x.reshape(b, n_q, self.in_features) + return self.out(x) + + +class Transformer(nn.Module): + """Classic Transformer used in Forecaster.""" + + def __init__( + self, + model_dims: int, + hidden_dims: int, + num_heads: int, + *, + attention_norm: str, + feedforward_norm: str, + qk_norm: str, + use_bias: bool, + use_rotary_position_embeddings: bool, + ff_activation: str, + fuse_qkv: bool, + use_cross_attn: bool = False, + cross_attn_kv_dim: int = 0, + use_cross_attn_gate: bool = True, + ): + super().__init__() + + if attention_norm == "rms": + self.pre_attn_ln = RMSNorm(num_features=model_dims) + self.post_attn_ln = RMSNorm(num_features=model_dims) + else: + raise ValueError(f"Layer norm: {attention_norm} not supported.") + + self.attn = MultiHeadAttention( + num_heads=num_heads, + in_features=model_dims, + use_per_dim_scale=True, + use_rotary_position_embeddings=use_rotary_position_embeddings, + qk_norm=qk_norm, + fuse_qkv=fuse_qkv, + ) + + if use_cross_attn: + _cross_kv_dim = cross_attn_kv_dim or model_dims + self.cross_attn_ln = RMSNorm(num_features=model_dims) + self.cross_attn = MultiHeadCrossAttention( + num_heads=num_heads, + in_features=model_dims, + kv_features=_cross_kv_dim, + use_bias=use_bias, + qk_norm=qk_norm, + ) + if use_cross_attn_gate: + # tanh(0)=0 → cross-attn contributes zero at init. + self.cross_attn_gate = nn.Parameter(torch.zeros(1)) + else: + self.cross_attn_gate = None + else: + self.cross_attn_ln = None + self.cross_attn = None + self.cross_attn_gate = None + + if feedforward_norm == "rms": + self.pre_ff_ln = RMSNorm(num_features=model_dims) + self.post_ff_ln = RMSNorm(num_features=model_dims) + else: + raise ValueError(f"Layer norm: {feedforward_norm} not supported.") + + self.ff0 = nn.Linear( + in_features=model_dims, + out_features=hidden_dims, + bias=use_bias, + ) + self.ff1 = nn.Linear( + in_features=hidden_dims, + out_features=model_dims, + bias=use_bias, + ) + if ff_activation == "relu": + self.activation = nn.ReLU() + elif ff_activation == "swish": + self.activation = nn.SiLU() + elif ff_activation == "none": + self.activation = nn.Identity() + else: + raise ValueError(f"Activation: {ff_activation} not supported.") + + def forward( + self, + input_embeddings: torch.Tensor, + patch_mask: torch.Tensor, + decode_cache: DecodeCache | None = None, + cross_kv: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, DecodeCache | None]: + attn_output, decode_cache = self.attn( + inputs_q=self.pre_attn_ln(input_embeddings), + decode_cache=decode_cache, + patch_mask=patch_mask, + ) + attn_output = self.post_attn_ln(attn_output) + input_embeddings + + if self.cross_attn is not None and cross_kv is not None: + cross_out = self.cross_attn( + self.cross_attn_ln(attn_output), + kv=cross_kv, + ) + if self.cross_attn_gate is not None: + cross_out = torch.tanh(self.cross_attn_gate) * cross_out + attn_output = attn_output + cross_out + + output_embeddings = ( + self.post_ff_ln(self.ff1(self.activation(self.ff0(self.pre_ff_ln(attn_output))))) + attn_output + ) + return output_embeddings, decode_cache + + +class _QFormerBlock(nn.Module): + """One Q-former block: self-attn over queries, cross-attn into source, FFN.""" + + def __init__( + self, + hidden_dim: int, + num_heads: int, + dropout: float = 0.0, + ffn_mult: float = 4.0, + ): + super().__init__() + self.self_attn_ln = nn.LayerNorm(hidden_dim) + self.self_attn = nn.MultiheadAttention( + embed_dim=hidden_dim, + num_heads=num_heads, + dropout=dropout, + batch_first=True, + ) + + self.cross_attn_ln_q = nn.LayerNorm(hidden_dim) + self.cross_attn_ln_kv = nn.LayerNorm(hidden_dim) + self.cross_attn = nn.MultiheadAttention( + embed_dim=hidden_dim, + num_heads=num_heads, + dropout=dropout, + batch_first=True, + ) + + ffn_hidden = int(round(hidden_dim * ffn_mult)) + self.ffn_ln = nn.LayerNorm(hidden_dim) + self.ffn = nn.Sequential( + nn.Linear(hidden_dim, ffn_hidden), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(ffn_hidden, hidden_dim), + nn.Dropout(dropout), + ) + + def forward( + self, + queries: torch.Tensor, + kv: torch.Tensor, + kv_key_padding_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + q_norm = self.self_attn_ln(queries) + sa_out, _ = self.self_attn(q_norm, q_norm, q_norm, need_weights=False) + queries = queries + sa_out + + q_norm = self.cross_attn_ln_q(queries) + kv_norm = self.cross_attn_ln_kv(kv) + ca_out, _ = self.cross_attn( + q_norm, + kv_norm, + kv_norm, + key_padding_mask=kv_key_padding_mask, + need_weights=False, + ) + queries = queries + ca_out + + queries = queries + self.ffn(self.ffn_ln(queries)) + return queries + + +class QFormer(nn.Module): + """BLIP-2-style Q-former: learned queries cross-attend to a source sequence. + + Args: + in_dim: Source feature dim (e.g. ts encoder hidden, or LLM hidden). + out_dim: Output (query) hidden dim. Should match the consumer module dim. + num_query_tokens: Number of learned query tokens. + num_heads: Attention heads inside the Q-former. + num_layers: Number of stacked Q-former blocks. + dropout: Attention/FFN dropout. + ffn_mult: FFN expansion ratio. + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + num_query_tokens: int, + num_heads: int = 8, + num_layers: int = 2, + dropout: float = 0.0, + ffn_mult: float = 4.0, + ): + super().__init__() + if num_query_tokens <= 0: + raise ValueError(f"num_query_tokens must be positive, got {num_query_tokens}") + if num_layers <= 0: + raise ValueError(f"num_layers must be positive, got {num_layers}") + if out_dim % num_heads != 0: + raise ValueError(f"out_dim ({out_dim}) must be divisible by num_heads ({num_heads})") + + self.in_dim = in_dim + self.out_dim = out_dim + self.num_query_tokens = num_query_tokens + + self.input_proj = nn.Linear(in_dim, out_dim) + self.input_ln = nn.LayerNorm(out_dim) + + self.query_tokens = nn.Parameter(torch.zeros(1, num_query_tokens, out_dim)) + nn.init.trunc_normal_(self.query_tokens, std=0.02) + + self.blocks = nn.ModuleList( + [ + _QFormerBlock( + hidden_dim=out_dim, + num_heads=num_heads, + dropout=dropout, + ffn_mult=ffn_mult, + ) + for _ in range(num_layers) + ] + ) + + self.output_ln = nn.LayerNorm(out_dim) + + def forward( + self, + src: torch.Tensor, + src_key_padding_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Compress `src` into `num_query_tokens` learned query vectors. + + Args: + src: (B, T, in_dim) source token embeddings. + src_key_padding_mask: (B, T) bool, True = padded position (ignored). + + Returns: + (B, num_query_tokens, out_dim) query outputs. + """ + if src.dim() != 3: + raise ValueError(f"src must be (B, T, D), got shape {tuple(src.shape)}") + if src.size(-1) != self.in_dim: + raise ValueError(f"src last dim ({src.size(-1)}) does not match in_dim ({self.in_dim})") + + kv = self.input_ln(self.input_proj(src)) + + if src_key_padding_mask is not None: + if src_key_padding_mask.shape != src.shape[:2]: + raise ValueError( + "src_key_padding_mask shape must equal src.shape[:2]:" + f" {tuple(src_key_padding_mask.shape)} != {tuple(src.shape[:2])}" + ) + kpm = src_key_padding_mask.to(dtype=torch.bool, device=kv.device) + # If a row is fully padded, MHA would NaN. Flip one position to valid. + all_padded = kpm.all(dim=1) + if all_padded.any(): + kpm = kpm.clone() + kpm[all_padded, 0] = False + else: + kpm = None + + queries = self.query_tokens.expand(src.size(0), -1, -1).to(dtype=kv.dtype) + for block in self.blocks: + queries = block(queries, kv, kv_key_padding_mask=kpm) + + return self.output_ln(queries) + + +class Aligner(nn.Module): + """Aligns precomputed LLM / TS-encoder embeddings into Forecaster's cross-attn + KV space, and predicts the forecast horizon. + + Two per-modality Q-formers compress the raw hidden-state sequences into a + fixed number of query tokens; their concatenation is the static KV stream the + Forecaster transformer cross-attends to. The forecast horizon is predicted by + a separate one-layer Transformer encoder over the original LLM hidden states, + so it does not share the LLM Q-former used for Forecaster cross-attention. + """ + + def __init__(self, config: "InternS2PreviewTimeSeriesForecasterConfig"): + super().__init__() + qformer_kwargs = dict( + out_dim=config.qformer_hidden_dim, + num_query_tokens=config.qformer_num_query_tokens, + num_heads=config.qformer_num_heads, + num_layers=config.qformer_num_layers, + dropout=config.qformer_dropout, + ) + self.ts_qformer = QFormer(in_dim=config.d_ts_encoder, **qformer_kwargs) + self.llm_qformer = QFormer(in_dim=config.d_llm, **qformer_kwargs) + + if config.use_horizon_head: + horizon_hidden_dim = config.qformer_hidden_dim + if horizon_hidden_dim % config.qformer_num_heads != 0: + raise ValueError( + f"horizon_hidden_dim ({horizon_hidden_dim}) must be divisible by qformer_num_heads " + f"({config.qformer_num_heads}) for the horizon Transformer encoder." + ) + self.horizon_input_proj = nn.Linear(config.d_llm, horizon_hidden_dim) + self.horizon_input_ln = nn.LayerNorm(horizon_hidden_dim) + self.horizon_encoder = nn.TransformerEncoderLayer( + d_model=horizon_hidden_dim, + nhead=config.qformer_num_heads, + dim_feedforward=4 * horizon_hidden_dim, + dropout=config.qformer_dropout, + activation="gelu", + batch_first=True, + norm_first=True, + ) + # Prediction-length head: LayerNorm -> Linear -> SiLU -> Linear(.,1), + # regressing horizon from pooled horizon-encoder LLM features. + self.horizon_head = nn.Sequential( + nn.LayerNorm(horizon_hidden_dim), + nn.Linear(horizon_hidden_dim, horizon_hidden_dim), + nn.SiLU(), + nn.Linear(horizon_hidden_dim, 1), + ) + else: + self.horizon_input_proj = None + self.horizon_input_ln = None + self.horizon_encoder = None + self.horizon_head = None + + def forward( + self, + llm_embedding_input: torch.Tensor, + ts_encoder_embedding_input: torch.Tensor, + llm_embedding_mask: torch.Tensor | None = None, + ts_encoder_embedding_mask: torch.Tensor | None = None, + ): + """Compress both modalities into the Forecaster cross-attention KV stream. + + Args: + llm_embedding_input: (B, T_llm, d_llm) precomputed LLM hidden states. + ts_encoder_embedding_input: (B, T_ts, d_ts_encoder) precomputed TS-encoder + hidden states. + llm_embedding_mask / ts_encoder_embedding_mask: optional (B, T) bool masks, + True = valid token. Default: all valid. + + Returns: + ctx: (B, Q_ts + Q_llm, qformer_hidden_dim) cross-attention KV stream. + """ + # --- TS-encoder branch --- + ts_param = next(self.ts_qformer.parameters()) + ts_hidden = ts_encoder_embedding_input.to(device=ts_param.device, dtype=ts_param.dtype) + if ts_encoder_embedding_mask is None: + ts_pad = torch.zeros(ts_hidden.shape[:2], dtype=torch.bool, device=ts_hidden.device) + else: + ts_pad = (~ts_encoder_embedding_mask.to(dtype=torch.bool)).to(device=ts_hidden.device) + ts_chunk = self.ts_qformer(ts_hidden, src_key_padding_mask=ts_pad) + + # --- LLM branch --- + llm_param = next(self.llm_qformer.parameters()) + llm_hidden = llm_embedding_input.to(device=llm_param.device, dtype=llm_param.dtype) + if llm_embedding_mask is None: + llm_pad = torch.zeros(llm_hidden.shape[:2], dtype=torch.bool, device=llm_hidden.device) + else: + llm_pad = (~llm_embedding_mask.to(dtype=torch.bool)).to(device=llm_hidden.device) + llm_chunk = self.llm_qformer(llm_hidden, src_key_padding_mask=llm_pad) + + # --- Concatenate: ts tokens first, then llm tokens --- + ctx = torch.cat([ts_chunk, llm_chunk], dim=1) + return ctx + + def predict_horizon( + self, + llm_embedding_input: torch.Tensor, + llm_embedding_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Predict per-sample forecast horizon in LINEAR step space from LLM hidden states. + + Args: + llm_embedding_input: (B, T_llm, d_llm) precomputed LLM hidden states. + llm_embedding_mask: optional (B, T_llm) bool mask, True = valid token. + + Returns: + (B,) float32 tensor of the predicted horizon (in steps) as float32. Linear (not log) + so a small regression error stays a small ABSOLUTE step error — log1p/expm1 amplified + errors at long horizons (e.g. ~0.05 in log ≈ 36 steps at H=720), hurting exact matching. + """ + if ( + self.horizon_input_proj is None + or self.horizon_input_ln is None + or self.horizon_encoder is None + or self.horizon_head is None + ): + raise RuntimeError("horizon_head is not enabled but predict_horizon was called") + if llm_embedding_input.dim() != 3: + raise ValueError( + f"Expected llm_embedding_input with shape (B, T, D), got {tuple(llm_embedding_input.shape)}" + ) + head_param = next(self.horizon_head.parameters()) + llm_hidden = llm_embedding_input.to(device=head_param.device, dtype=head_param.dtype) + + if llm_embedding_mask is None: + valid_mask = torch.ones(llm_hidden.shape[:2], dtype=torch.bool, device=llm_hidden.device) + else: + valid_mask = llm_embedding_mask.to(device=llm_hidden.device, dtype=torch.bool) + if valid_mask.shape != llm_hidden.shape[:2]: + raise ValueError( + "llm_embedding_mask shape must equal llm_embedding_input.shape[:2]: " + f"{tuple(valid_mask.shape)} != {tuple(llm_hidden.shape[:2])}" + ) + if not valid_mask.any(dim=1).all(): + raise ValueError("Each llm_embedding_mask row must contain at least one valid token.") + + src_key_padding_mask = ~valid_mask + + hidden = self.horizon_input_ln(self.horizon_input_proj(llm_hidden)) + encoded = self.horizon_encoder( + hidden, + src_key_padding_mask=src_key_padding_mask, + ) + pool_mask = valid_mask.to(dtype=encoded.dtype).unsqueeze(-1) + denom = pool_mask.sum(dim=1).clamp_min(1.0) + pooled = (encoded * pool_mask).sum(dim=1) / denom + return self.horizon_head(pooled).squeeze(-1).to(dtype=torch.float32) + + @staticmethod + def decode_horizon( + pred_horizon: torch.Tensor, + min_h: int = 1, + max_h: int | None = None, + ) -> torch.Tensor: + """Round the head's linear-space horizon prediction to a positive integer tensor.""" + horizon = pred_horizon.round().clamp_min(float(min_h)) + if max_h is not None: + horizon = horizon.clamp_max(float(max_h)) + return horizon.long() + + +_TOLERANCE = 1e-6 + + +def revin( + x: torch.Tensor, + mu: torch.Tensor, + sigma: torch.Tensor, + reverse: bool = False, +): + """Reversible instance normalization.""" + if len(mu.shape) == len(x.shape) - 1: + mu = mu[..., None] + sigma = sigma[..., None] + elif len(mu.shape) == len(x.shape) - 2: + mu = mu[..., None, None] + sigma = sigma[..., None, None] + + if reverse: + return x * sigma + mu + else: + return (x - mu) / torch.where(sigma < _TOLERANCE, 1.0, sigma) + + +class ForecasterBackbone(nn.Module): + """Forecaster 2.5 with 200M parameters (cross-attention capable).""" + + # Architecture constants (Forecaster 2.5 200M) + context_limit = 16384 + _quantiles = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] + _num_layers = 20 + + # Residual block configs + _tokenizer_kwargs = dict(input_dims=64, hidden_dims=1280, output_dims=1280, use_bias=True, activation="swish") + _output_point_kwargs = dict( + input_dims=1280, hidden_dims=1280, output_dims=1280, use_bias=False, activation="swish" + ) + _output_quantile_kwargs = dict( + input_dims=1280, hidden_dims=1280, output_dims=10240, use_bias=False, activation="swish" + ) + + # Transformer layer config + _xf_kwargs = dict( + model_dims=1280, + hidden_dims=1280, + num_heads=16, + attention_norm="rms", + feedforward_norm="rms", + qk_norm="rms", + use_bias=False, + use_rotary_position_embeddings=True, + ff_activation="swish", + fuse_qkv=True, + ) + + def __init__( + self, + use_cross_attn: bool = False, + cross_attn_kv_dim: int = 0, + use_cross_attn_gate: bool = True, + ): + super().__init__() + + # Named constants. + self.p = 32 + self.o = 128 + self.os = 1024 + self.m = self.o // self.p # 4 + self.x = self._num_layers # 20 + self.h = self._xf_kwargs["num_heads"] # 16 + self.md = self._xf_kwargs["model_dims"] # 1280 + self.hd = self.md // self.h # 80 + self.q = len(self._quantiles) + 1 # 10 + self.aridx = 5 + + self.use_cross_attn = bool(use_cross_attn) + + xf_kwargs = dict(self._xf_kwargs) + if self.use_cross_attn: + xf_kwargs["use_cross_attn"] = True + xf_kwargs["cross_attn_kv_dim"] = int(cross_attn_kv_dim) or self.md + xf_kwargs["use_cross_attn_gate"] = bool(use_cross_attn_gate) + + # Layers. + self.tokenizer = ResidualBlock(**self._tokenizer_kwargs) + self.stacked_xf = nn.ModuleList([Transformer(**xf_kwargs) for _ in range(self.x)]) + self.output_projection_point = ResidualBlock(**self._output_point_kwargs) + self.output_projection_quantiles = ResidualBlock(**self._output_quantile_kwargs) + + def update_running_stats( + self, + n: torch.Tensor, + mu: torch.Tensor, + sigma: torch.Tensor, + x: torch.Tensor, + mask: torch.Tensor, + ): + """Updates the running stats.""" + is_legit = torch.logical_not(mask) + inc_n = torch.sum(is_legit.to(x.dtype), dim=-1) + + inc_mu_numerator = torch.sum(x * is_legit, dim=-1) + inc_n_safe = torch.where(inc_n == 0, 1.0, inc_n) + inc_mu = inc_mu_numerator / inc_n_safe + inc_mu = torch.where(inc_n == 0, 0.0, inc_mu) + + inc_var_numerator = torch.sum(((x - inc_mu.unsqueeze(-1)) ** 2) * is_legit, dim=-1) + inc_var = inc_var_numerator / inc_n_safe + inc_var = torch.where(inc_n == 0, 0.0, inc_var) + inc_sigma = torch.sqrt(inc_var) + + new_n = n + inc_n + new_n_safe = torch.where(new_n == 0, 1.0, new_n) + + new_mu = (n * mu + inc_mu * inc_n) / new_n_safe + new_mu = torch.where(new_n == 0, 0.0, new_mu) + + term1 = n * sigma.pow(2) + term2 = inc_n * inc_sigma.pow(2) + term3 = n * (mu - new_mu).pow(2) + term4 = inc_n * (inc_mu - new_mu).pow(2) + + new_var = (term1 + term2 + term3 + term4) / new_n_safe + new_var = torch.where(new_n == 0, 0.0, new_var) + new_sigma = torch.sqrt(torch.clamp(new_var, min=0.0)) + + return (w := (new_n, new_mu, new_sigma), w) + + def forward( + self, + inputs: torch.Tensor, + masks: torch.Tensor, + decode_caches: list | None = None, + cross_kv: torch.Tensor | None = None, + ): + """Forward pass — history-only path with cross-attention injection.""" + if cross_kv is not None and not getattr(self, "use_cross_attn", False): + raise ValueError( + "cross_kv was supplied but the Forecaster module was constructed with" + " use_cross_attn=False; rebuild the module with use_cross_attn=True." + ) + + input_dtype = inputs.dtype + model_dtype = next(self.parameters()).dtype + if inputs.dtype != model_dtype: + inputs = inputs.to(dtype=model_dtype) + + tokenizer_inputs = torch.cat([inputs, masks.to(inputs.dtype)], dim=-1) + input_embeddings = self.tokenizer(tokenizer_inputs) + + if cross_kv is not None: + cross_kv = cross_kv.to(device=input_embeddings.device, dtype=input_embeddings.dtype) + + if decode_caches is None: + decode_caches = [None] * self.x + new_decode_caches = [] + + output_embeddings = input_embeddings + token_masks = masks[..., -1] + + for i, layer in enumerate(self.stacked_xf): + output_embeddings, new_cache = layer( + output_embeddings, + token_masks, + decode_caches[i], + cross_kv=cross_kv, + ) + new_decode_caches.append(new_cache) + + output_ts = self.output_projection_point(output_embeddings) + output_quantile_spread = self.output_projection_quantiles(output_embeddings) + + if output_ts.dtype != input_dtype: + input_embeddings = input_embeddings.to(dtype=input_dtype) + output_embeddings = output_embeddings.to(dtype=input_dtype) + output_ts = output_ts.to(dtype=input_dtype) + output_quantile_spread = output_quantile_spread.to(dtype=input_dtype) + + return ( + input_embeddings, + output_embeddings, + output_ts, + output_quantile_spread, + ), new_decode_caches + + def decode( + self, + horizon: int, + inputs, + masks, + enable_grad: bool = False, + cross_kv: torch.Tensor | None = None, + ): + """Decodes the time series.""" + with torch.set_grad_enabled(enable_grad): + batch_size, context = inputs.shape[0], inputs.shape[1] + num_decode_steps = (horizon - 1) // self.o + num_input_patches = context // self.p + decode_cache_size = num_input_patches + num_decode_steps * self.m + cache_dtype = next(self.parameters()).dtype + + # Prefill + patched_inputs = torch.reshape(inputs, (batch_size, -1, self.p)) + patched_masks = torch.reshape(masks, (batch_size, -1, self.p)) + + # running stats + n = torch.zeros(batch_size, device=inputs.device) + mu = torch.zeros(batch_size, device=inputs.device) + sigma = torch.zeros(batch_size, device=inputs.device) + patch_mu = [] + patch_sigma = [] + for i in range(num_input_patches): + (n, mu, sigma), _ = self.update_running_stats(n, mu, sigma, patched_inputs[:, i], patched_masks[:, i]) + patch_mu.append(mu) + patch_sigma.append(sigma) + last_n, last_mu, last_sigma = n, mu, sigma + context_mu = torch.stack(patch_mu, dim=1) + context_sigma = torch.stack(patch_sigma, dim=1) + + decode_caches = [ + DecodeCache( + next_index=torch.zeros(batch_size, dtype=torch.int32, device=inputs.device), + num_masked=torch.zeros(batch_size, dtype=torch.int32, device=inputs.device), + key=torch.zeros( + batch_size, + decode_cache_size, + self.h, + self.hd, + device=inputs.device, + dtype=cache_dtype, + ), + value=torch.zeros( + batch_size, + decode_cache_size, + self.h, + self.hd, + device=inputs.device, + dtype=cache_dtype, + ), + ) + for _ in range(self.x) + ] + + normed_inputs = revin(patched_inputs, context_mu, context_sigma, reverse=False) + normed_inputs = torch.where(patched_masks, 0.0, normed_inputs) + (_, _, normed_outputs, normed_quantile_spread), decode_caches = self( + normed_inputs, + patched_masks, + decode_caches, + cross_kv=cross_kv, + ) + renormed_outputs = torch.reshape( + revin(normed_outputs, context_mu, context_sigma, reverse=True), + (batch_size, -1, self.o, self.q), + ) + renormed_quantile_spread = torch.reshape( + revin(normed_quantile_spread, context_mu, context_sigma, reverse=True), + (batch_size, -1, self.os, self.q), + )[:, -1, ...] + + # Autoregressive decode + ar_outputs = [] + last_renormed_output = renormed_outputs[:, -1, :, self.aridx] + + for _ in range(num_decode_steps): + new_patched_input = torch.reshape(last_renormed_output, (batch_size, self.m, self.p)) + new_mask = torch.zeros_like(new_patched_input, dtype=torch.bool) + + n, mu, sigma = last_n, last_mu, last_sigma + new_mus, new_sigmas = [], [] + for i in range(self.m): + (n, mu, sigma), _ = self.update_running_stats( + n, mu, sigma, new_patched_input[:, i], new_mask[:, i] + ) + new_mus.append(mu) + new_sigmas.append(sigma) + last_n, last_mu, last_sigma = n, mu, sigma + new_mu = torch.stack(new_mus, dim=1) + new_sigma = torch.stack(new_sigmas, dim=1) + + new_normed_input = revin(new_patched_input, new_mu, new_sigma, reverse=False) + (_, _, new_normed_output, _), decode_caches = self( + new_normed_input, + new_mask, + decode_caches, + cross_kv=cross_kv, + ) + + new_renormed_output = torch.reshape( + revin(new_normed_output, new_mu, new_sigma, reverse=True), + (batch_size, self.m, self.o, self.q), + ) + ar_outputs.append(new_renormed_output[:, -1, ...]) + last_renormed_output = new_renormed_output[:, -1, :, self.aridx] + + if num_decode_steps > 0: + ar_renormed_outputs = torch.stack(ar_outputs, dim=1) + else: + ar_renormed_outputs = None + + return renormed_outputs, renormed_quantile_spread, ar_renormed_outputs + + +@dataclass +class InternS2PreviewTimeSeriesForecasterOutput(ModelOutput): + """Output of :class:`InternS2PreviewTimeSeriesForecaster`. + + Attributes: + point_forecast: list of (horizon_i, C_i) per-sample median forecasts. + quantile_forecast: list of (horizon_i, C_i, 10) per-sample quantile + forecasts (quantile order: [median, 0.1, 0.2, ..., 0.9]). + predicted_horizon: (B,) long tensor of horizons predicted by the horizon + head, or None when the head is disabled / a horizon override was given. + """ + + point_forecast: list[torch.Tensor] | None = None + quantile_forecast: list[torch.Tensor] | None = None + predicted_horizon: torch.Tensor | None = None + + +class InternS2PreviewTimeSeriesForecaster(PreTrainedModel): + """Standalone TimeOmni_v2 forecaster (cross-attention Forecaster head). + + Inputs (see :meth:`forward`): the raw multi-channel ``history``, plus the two + precomputed embedding streams ``llm_embedding_input`` and + ``ts_encoder_embedding_input``. The LLM and TS encoder themselves are NOT part + of this model. + """ + + config_class = InternS2PreviewTimeSeriesForecasterConfig + base_model_prefix = "interns2_preview_time_series_forecaster" + main_input_name = "history" + supports_gradient_checkpointing = False + + def __init__(self, config: InternS2PreviewTimeSeriesForecasterConfig): + super().__init__(config) + + self.aligner = Aligner(config) + self.forecaster = ForecasterBackbone( + use_cross_attn=True, + cross_attn_kv_dim=config.cross_attn_kv_dim, + use_cross_attn_gate=config.use_cross_attn_gate, + ) + + # Normalize max_context / max_horizon to Forecaster patch multiples. + if config.max_context % self.forecaster.p != 0: + config.max_context = math.ceil(config.max_context / self.forecaster.p) * self.forecaster.p + if config.max_horizon % self.forecaster.o != 0: + config.max_horizon = math.ceil(config.max_horizon / self.forecaster.o) * self.forecaster.o + if config.max_context + config.max_horizon > self.forecaster.context_limit: + raise ValueError( + "Context + horizon must be less than the context limit." + f" {config.max_context} + {config.max_horizon} > {self.forecaster.context_limit}." + ) + if config.use_continuous_quantile_head and (config.max_horizon > self.forecaster.os): + raise ValueError(f"Continuous quantile head is not supported for horizons > {self.forecaster.os}.") + + self._horizon_max_length = int(config.horizon_max_length) or int(config.max_horizon) + + self.post_init() + + def _init_weights(self, module): + """No-op: every submodule self-initializes in its own constructor. + + Overriding this keeps post_init() from clobbering the carefully chosen + inits (RMSNorm.scale=0, PerDimScale=0, cross_attn_gate=0, QFormer query + trunc-normal, ...).""" + return + + # ------------------------------------------------------------------ # + # Forecaster forecast orchestration (ported from Forecaster_2p5.forecast / + # Forecaster_2p5_200M_torch.compile._compiled_decode, cross-attn path). + # ------------------------------------------------------------------ # + + def _forecaster_compiled_decode( + self, + horizon, + inputs, + masks, + enable_grad: bool = False, + cross_kv=None, + ): + """Single-series decode with revin / flip-invariance / quantile post-proc.""" + fc = self.config + module = self.forecaster + if horizon > fc.max_horizon: + raise ValueError(f"Horizon must be less than the max horizon. {horizon} > {fc.max_horizon}.") + + with torch.set_grad_enabled(enable_grad): + target_device = next(module.parameters()).device + inputs = torch.as_tensor(inputs, dtype=torch.float32, device=target_device) + masks = torch.as_tensor(masks, dtype=torch.bool, device=target_device) + batch_size = inputs.shape[0] + + cross_kv_t = None + if cross_kv is not None: + cross_kv_t = torch.as_tensor(cross_kv, dtype=torch.float32, device=target_device) + if cross_kv_t.shape[0] != batch_size: + raise ValueError( + f"Batch size mismatch between cross_kv and inputs: {cross_kv_t.shape[0]} != {batch_size}" + ) + + if fc.infer_is_positive: + is_positive = torch.all(inputs >= 0, dim=-1, keepdim=True) + else: + is_positive = None + + if fc.normalize_inputs: + mu = torch.mean(inputs, dim=-1, keepdim=True) + sigma = torch.std(inputs, dim=-1, keepdim=True) + inputs = revin(inputs, mu, sigma, reverse=False) + else: + mu, sigma = None, None + + pf_outputs, quantile_spreads, ar_outputs = module.decode( + horizon, + inputs, + masks, + enable_grad=enable_grad, + cross_kv=cross_kv_t, + ) + to_cat = [pf_outputs[:, -1, ...]] + if ar_outputs is not None: + to_cat.append(ar_outputs.reshape(batch_size, -1, module.q)) + full_forecast = torch.cat(to_cat, dim=1) + + def flip_quantile_fn(x): + return torch.cat([x[..., :1], torch.flip(x[..., 1:], dims=(-1,))], dim=-1) + + if fc.force_flip_invariance: + flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = module.decode( + horizon, + -inputs, + masks, + enable_grad=enable_grad, + cross_kv=cross_kv_t, + ) + flipped_quantile_spreads = flip_quantile_fn(flipped_quantile_spreads) + flipped_pf_outputs = flip_quantile_fn(flipped_pf_outputs) + to_cat = [flipped_pf_outputs[:, -1, ...]] + if flipped_ar_outputs is not None: + to_cat.append(flipped_ar_outputs.reshape(batch_size, -1, module.q)) + flipped_full_forecast = torch.cat(to_cat, dim=1) + quantile_spreads = (quantile_spreads - flipped_quantile_spreads) / 2 + pf_outputs = (pf_outputs - flipped_pf_outputs) / 2 + full_forecast = (full_forecast - flipped_full_forecast) / 2 + + if fc.use_continuous_quantile_head: + for quantile_index in [1, 2, 3, 4, 6, 7, 8, 9]: + full_forecast[:, :horizon, quantile_index] = ( + quantile_spreads[:, :horizon, quantile_index] + - quantile_spreads[:, :horizon, 5] + + full_forecast[:, :horizon, 5] + ) + full_forecast = full_forecast[:, :horizon, :] + + if fc.return_backcast: + full_backcast = pf_outputs[:, :-1, : module.p, :].reshape(batch_size, -1, module.q) + full_forecast = torch.cat([full_backcast, full_forecast], dim=1) + + if fc.fix_quantile_crossing: + for i in [4, 3, 2, 1]: + full_forecast[:, :, i] = torch.where( + full_forecast[:, :, i] < full_forecast[:, :, i + 1], + full_forecast[:, :, i], + full_forecast[:, :, i + 1], + ) + for i in [6, 7, 8, 9]: + full_forecast[:, :, i] = torch.where( + full_forecast[:, :, i] > full_forecast[:, :, i - 1], + full_forecast[:, :, i], + full_forecast[:, :, i - 1], + ) + + if fc.normalize_inputs: + full_forecast = revin(full_forecast, mu, sigma, reverse=True) + + if is_positive is not None: + full_forecast = torch.where( + is_positive[..., None], + torch.maximum(full_forecast, torch.zeros_like(full_forecast)), + full_forecast, + ) + + return full_forecast[..., 5], full_forecast + + def strip_leading_nans(self, arr): + """Removes contiguous NaN values from the beginning of a NumPy array.""" + isnan = np.isnan(arr) + first_valid_index = np.argmax(~isnan) + return arr[first_valid_index:] + + def linear_interpolation(self, arr): + """Linear interpolation to fill NaN values in a 1D numpy array.""" + nans = np.isnan(arr) + if not np.any(nans): + return arr + + def x(z): + return z.nonzero()[0] + + nans_indices = x(nans) + non_nans_indices = x(~nans) + non_nans_values = arr[~nans] + + try: + arr[nans] = np.interp(nans_indices, non_nans_indices, non_nans_values) + except ValueError: + if non_nans_values: + mu = np.nanmean(arr) + else: + mu = 0.0 + arr = np.where(np.isfinite(arr), arr, mu) + return arr + + def _forecaster_forecast( + self, + horizon: int, + inputs: list, + enable_grad: bool = False, + cross_kv: list | None = None, + ): + """Batched forecast over a list of 1-D series (ported from Forecaster_2p5).""" + max_context = self.config.max_context + num_inputs = len(inputs) + if cross_kv is not None and len(cross_kv) != num_inputs: + raise ValueError(f"cross_kv length must match inputs length: {len(cross_kv)} != {num_inputs}") + + output_points = [] + output_quantiles = [] + for input_index, each_input in enumerate(inputs): + value = torch.as_tensor(each_input, dtype=torch.float32).reshape(-1) + if value.numel() > 0: + value_np = self.linear_interpolation(self.strip_leading_nans(value.detach().cpu().numpy())) + value = torch.from_numpy(value_np).to(dtype=torch.float32) + + if value.numel() > max_context: + value = value[-max_context:] + + context = max( + self.forecaster.p, + math.ceil(value.numel() / self.forecaster.p) * self.forecaster.p, + ) + context = min(context, max_context) + + value_len = value.numel() + if value_len >= context: + value = value[-context:] + mask = torch.zeros(context, dtype=torch.bool) + else: + pad_len = context - value_len + value = torch.cat([torch.zeros(pad_len, dtype=torch.float32), value], dim=0) + mask = torch.cat( + [torch.ones(pad_len, dtype=torch.bool), torch.zeros(value_len, dtype=torch.bool)], + dim=0, + ) + + values_t = value.unsqueeze(0) + masks_t = mask.unsqueeze(0) + cross_kv_t = None + if cross_kv is not None: + cross_kv_t = torch.as_tensor( + cross_kv[input_index], + dtype=torch.float32, + device=value.device, + ).unsqueeze(0) + + point_forecast, quantile_forecast = self._forecaster_compiled_decode( + horizon, + values_t, + masks_t, + enable_grad=enable_grad, + cross_kv=cross_kv_t, + ) + output_points.append(point_forecast) + output_quantiles.append(quantile_forecast) + + output_points = torch.cat(output_points, dim=0) + output_quantiles = torch.cat(output_quantiles, dim=0) + return output_points[:num_inputs], output_quantiles[:num_inputs] + + @staticmethod + def _flatten_channelwise(history: list[torch.Tensor], ctx: torch.Tensor): + """Split each (T, C) history into C single-channel series; replicate the + per-sample cross-attn context once per channel (Forecaster forecasts a single + univariate series at a time).""" + channel_counts = [] + flattened_inputs = [] + flattened_prefix = [] + for sample_idx, ts in enumerate(history): + if ts.dim() != 2: + raise ValueError(f"Each history series must be 2D (T, C), got shape {tuple(ts.shape)}") + channels = ts.shape[1] + channel_counts.append(channels) + for channel_idx in range(channels): + flattened_inputs.append(ts[:, channel_idx]) + flattened_prefix.append(ctx[sample_idx]) + return channel_counts, flattened_inputs, flattened_prefix + + def forward( + self, + history: list[torch.Tensor], + llm_embedding_input: torch.Tensor, + ts_encoder_embedding_input: torch.Tensor, + llm_embedding_mask: torch.Tensor | None = None, + ts_encoder_embedding_mask: torch.Tensor | None = None, + override_horizon: int | Sequence[int] | None = None, + return_dict: bool | None = None, + ) -> InternS2PreviewTimeSeriesForecasterOutput: + """Forecast multi-channel time series. + + Args: + history: list of B per-sample 2-D tensors (T_i, C_i) — the numeric + forecast context fed to Forecaster. T and C may differ across samples. + llm_embedding_input: (B, T_llm, d_llm) precomputed LLM hidden states. + ts_encoder_embedding_input: (B, T_ts, d_ts_encoder) precomputed TS-encoder + hidden states. + llm_embedding_mask / ts_encoder_embedding_mask: optional (B, T) bool masks + (True = valid). Default: all valid. + override_horizon: optional forecast horizon. An int applies to every + sample; a sequence gives one horizon per sample. When omitted, the + horizon head predicts it (if enabled), else ``config.default_pred_len``. + return_dict: kept for HF API compatibility (output is always a + :class:`TSForecasterOutput`). + + Returns: + :class:`TSForecasterOutput` with per-sample ``point_forecast`` / + ``quantile_forecast`` lists and the predicted horizon. + """ + if not isinstance(history, (list, tuple)): + raise ValueError("history must be a list/tuple of (T, C) tensors, one per sample") + batch_size = len(history) + if llm_embedding_input.size(0) != batch_size or ts_encoder_embedding_input.size(0) != batch_size: + raise ValueError( + "Batch size mismatch: len(history)=" + f"{batch_size}, llm_embedding_input={llm_embedding_input.size(0)}," + f" ts_encoder_embedding_input={ts_encoder_embedding_input.size(0)}" + ) + + # 1. Align modalities -> cross-attention KV stream. + ctx = self.aligner( + llm_embedding_input, + ts_encoder_embedding_input, + llm_embedding_mask=llm_embedding_mask, + ts_encoder_embedding_mask=ts_encoder_embedding_mask, + ) + + # 2. Resolve per-sample horizons. + predicted_horizon = None + if self.aligner.horizon_head is not None: + pred_horizon = self.aligner.predict_horizon( + llm_embedding_input, + llm_embedding_mask=llm_embedding_mask, + ) + predicted_horizon = self.aligner.decode_horizon( + pred_horizon, + min_h=1, + max_h=self._horizon_max_length, + ) + if override_horizon is not None: + if isinstance(override_horizon, int): + per_sample_horizons = [int(override_horizon)] * batch_size + else: + per_sample_horizons = [int(h) for h in override_horizon] + elif self.aligner.horizon_head is not None: + per_sample_horizons = [int(v) for v in predicted_horizon.tolist()] + else: + per_sample_horizons = [int(self.config.default_pred_len)] * batch_size + + if len(per_sample_horizons) != batch_size: + raise ValueError( + f"per-sample horizon count ({len(per_sample_horizons)}) does not match batch size ({batch_size})" + ) + # Forecaster decodes a single horizon for the whole flattened batch; use the + # max and slice each sample back to its own length below. + horizon = max(per_sample_horizons) if per_sample_horizons else 1 + + # 3. Flatten channelwise; replicate the cross-attn context per channel. + channel_counts, flattened_inputs, flattened_prefix = self._flatten_channelwise(history, ctx) + # Preprocess the KV stream in float32 (stable revin); the Forecaster module + # re-casts it to its own dtype internally. + flattened_prefix = [p.to(dtype=torch.float32) for p in flattened_prefix] + + # 4. Forecaster forecast (cross-attention injection). + point_flat, quantile_flat = self._forecaster_forecast( + horizon=horizon, + inputs=flattened_inputs, + enable_grad=self.training, + cross_kv=flattened_prefix, + ) + + # 5. Reshape the flat (Bf, horizon[, Q]) outputs back to per-sample lists. + point_outputs = [] + quantile_outputs = [] + cursor = 0 + for sample_idx, channels in enumerate(channel_counts): + sample_horizon = max(1, min(int(per_sample_horizons[sample_idx]), horizon)) + channel_point = [] + channel_quantile = [] + for _ in range(channels): + channel_point.append(point_flat[cursor, :sample_horizon]) + channel_quantile.append(quantile_flat[cursor, :sample_horizon, :]) + cursor += 1 + point_outputs.append(torch.stack(channel_point, dim=1)) + quantile_outputs.append(torch.stack(channel_quantile, dim=1)) + + return InternS2PreviewTimeSeriesForecasterOutput( + point_forecast=point_outputs, + quantile_forecast=quantile_outputs, + predicted_horizon=predicted_horizon, + ) + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Llava outputs, with hidden states and attentions. + """ +) +class InternS2PreviewModelOutputWithPast(ModelOutput): + r""" + Base class for InternS2Preview model outputs with cache and time-series + encoder states. + + Args: + last_hidden_state (`torch.FloatTensor`, *optional*): + Sequence of hidden-states at the output of the last layer. + past_key_values (`Cache`, *optional*): + Pre-computed hidden-states that can be used to speed up sequential decoding. + hidden_states (`tuple[torch.FloatTensor]`, *optional*): + Hidden-states of the model at the output of each layer. + attentions (`tuple[torch.FloatTensor]`, *optional*): + Attention weights after the attention softmax. + rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): + The rope index difference between sequence length and multimodal rope. + router_logits (`tuple[torch.FloatTensor]`, *optional*): + Router logits from MoE layers. + ts_encoder_embedding (`torch.Tensor`, *optional*): + Padded time-series encoder embeddings. + ts_encoder_embedding_mask (`torch.Tensor`, *optional*): + Boolean mask for valid time-series encoder embeddings. + ts_history (`list[torch.Tensor]`, *optional*): + Raw time-series histories used by the forecasting head. + """ + + last_hidden_state: torch.FloatTensor | None = None + past_key_values: Cache | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + rope_deltas: torch.LongTensor | None = None + router_logits: tuple[torch.FloatTensor] | None = None + + ts_encoder_embedding: torch.Tensor | None = None + ts_encoder_embedding_mask: torch.Tensor | None = None + ts_history: list[torch.Tensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for InternS2Preview causal language model (or autoregressive) outputs. + """ +) +class InternS2PreviewCausalLMOutputWithPast(ModelOutput): + r""" + Output type of InternS2Preview causal language modeling with optional + time-series forecasting outputs. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*): + Language modeling loss. + logits (`torch.FloatTensor`, *optional*): + Prediction scores of the language modeling head. + past_key_values (`Cache`, *optional*): + Pre-computed hidden-states that can be used to speed up sequential decoding. + hidden_states (`tuple[torch.FloatTensor]`, *optional*): + Hidden-states of the model at the output of each layer. + attentions (`tuple[torch.FloatTensor]`, *optional*): + Attention weights after the attention softmax. + rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): + The rope index difference between sequence length and multimodal rope. + router_logits (`tuple[torch.FloatTensor]`, *optional*): + Router logits from MoE layers. + aux_loss (`torch.FloatTensor`, *optional*): + Auxiliary router loss. + ts_point_forecast (`list[torch.Tensor]`, *optional*): + Point forecasts for each sample. + ts_quantile_forecast (`list[torch.Tensor]`, *optional*): + Quantile forecasts for each sample. + ts_predicted_horizon (`torch.Tensor`, *optional*): + Predicted forecast horizon. + """ + + loss: torch.FloatTensor | None = None + logits: torch.FloatTensor | None = None + past_key_values: Cache | None = None + hidden_states: tuple[torch.FloatTensor] | None = None + attentions: tuple[torch.FloatTensor] | None = None + rope_deltas: torch.LongTensor | None = None + router_logits: tuple[torch.FloatTensor] | None = None + aux_loss: torch.FloatTensor | None = None + ts_point_forecast: list[torch.Tensor] | None = None + ts_quantile_forecast: list[torch.Tensor] | None = None + ts_predicted_horizon: torch.Tensor | None = None + + +@auto_docstring +class InternS2PreviewModel(InternS2PreviewPreTrainedModel): + base_model_prefix = "model" + _checkpoint_conversion_mapping = {} + # Reference: fix gemma3 grad acc #37208 + accepts_loss_kwargs = False + + config: InternS2PreviewConfig + _no_split_modules = [ + "Qwen3_5MoeTextDecoderLayer", + "Qwen3_5MoeVisionBlock", + "InternS2PreviewTimeSeriesEncoderLayer", + ] + + def __init__(self, config): + super().__init__(config) + self.visual = InternS2PreviewVisionModel._from_config(config.vision_config) + self.language_model = InternS2PreviewTextModel._from_config(config.text_config) + self.rope_deltas = None # cache rope_deltas here + self.time_series = InternS2PreviewTimeSeriesModel._from_config(config.ts_config) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.language_model.set_input_embeddings(value) def get_rope_index( self, @@ -2461,12 +4967,18 @@ class InternS2PreviewModel(InternS2PreviewPreTrainedModel): **kwargs: Unpack[TransformersKwargs], ) -> tuple | InternS2PreviewModelOutputWithPast: r""" + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. ts_values (`torch.FloatTensor` of shape `(batch_size, seq_len, num_channels)`, *optional*): The tensors corresponding to the input time series signals. ts_lens (`torch.Tensor` of shape `(batch_size,)`, *optional*): The valid lengths of each time series signal in the batch. ts_sr (`torch.FloatTensor` of shape `(batch_size,)`, *optional*): The sampling rates of each time series signal in the batch. + ts_channels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + The number of channels of each time series signal in the batch. """ if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -2496,9 +5008,16 @@ class InternS2PreviewModel(InternS2PreviewPreTrainedModel): ) inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + ts_encoder_embedding = None + ts_encoder_embedding_mask = None + ts_history = None if ts_values is not None: - ts_features, ts_pad_mask = self.get_ts_feature(ts_values, ts_lens, ts_sr, ts_channels) # [B, T, C], [B, T] - ts_features = ts_features[~ts_pad_mask].to( + ts_embeds, ts_pad_mask, ts_embeds_before_project = self.get_ts_feature( + ts_values, ts_lens, ts_sr, ts_channels + ) # [B, T, C], [B, T] + ts_encoder_embedding = ts_embeds_before_project + ts_encoder_embedding_mask = ~ts_pad_mask + ts_features = ts_embeds[~ts_pad_mask].to( inputs_embeds.device, inputs_embeds.dtype ) # [num_valid_ts_tokens, C] B, N, C = inputs_embeds.shape @@ -2511,6 +5030,7 @@ class InternS2PreviewModel(InternS2PreviewPreTrainedModel): assert n_ts_placeholders == n_ts_tokens, ( f"[ERROR]: Mismatch: tokens={n_ts_placeholders}, ts_embeds_valid={n_ts_tokens}" ) + input_ids = input_ids.reshape(B, N) try: # TODO why not scatter? @@ -2523,6 +5043,8 @@ class InternS2PreviewModel(InternS2PreviewPreTrainedModel): inputs_embeds = inputs_embeds.reshape(B, N, C) + ts_history = [ts_values[i, : ts_lens[i], : ts_channels[i]] for i in range(ts_values.shape[0])] + if position_ids is None: position_ids = self.compute_3d_position_ids( input_ids=input_ids, @@ -2546,6 +5068,9 @@ class InternS2PreviewModel(InternS2PreviewPreTrainedModel): return InternS2PreviewModelOutputWithPast( **outputs, rope_deltas=self.rope_deltas, + ts_encoder_embedding=ts_encoder_embedding, + ts_encoder_embedding_mask=ts_encoder_embedding_mask, + ts_history=ts_history, ) @can_return_tuple @@ -2564,8 +5089,10 @@ class InternS2PreviewModel(InternS2PreviewPreTrainedModel): The length of the time series. sr (`torch.FloatTensor` of shape `(batch_size,)`): The sampling rate of the time series. + ts_channels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + The number of channels of the time series. """ - ts_embeds, ts_pad_mask = self.time_series( + ts_embeds, ts_pad_mask, ts_embeds_before_project = self.time_series( time_series_signals=ts_values, ts_lens=ts_lens, sr=sr, @@ -2573,7 +5100,7 @@ class InternS2PreviewModel(InternS2PreviewPreTrainedModel): output_hidden_states=False, return_dict=True, ) - return ts_embeds, ts_pad_mask + return ts_embeds, ts_pad_mask, ts_embeds_before_project # NOTE: Cannot inherit from `Qwen3_5MoeForCausalLM` here due to a converter limitation: when the modular @@ -2700,6 +5227,7 @@ class InternS2PreviewForConditionalGeneration(InternS2PreviewPreTrainedModel, Ge super().__init__(config) self.model = InternS2PreviewModel(config) self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + self.time_series_forecaster = InternS2PreviewTimeSeriesForecaster(config.ts_forecaster_config) self.post_init() @@ -2760,6 +5288,8 @@ class InternS2PreviewForConditionalGeneration(InternS2PreviewPreTrainedModel, Ge ts_channels: torch.LongTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, + ts_forecast: bool = False, + forecast_horizon: int | list[int] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | InternS2PreviewCausalLMOutputWithPast: r""" @@ -2837,6 +5367,18 @@ class InternS2PreviewForConditionalGeneration(InternS2PreviewPreTrainedModel, Ge slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(hidden_states[:, slice_indices, :]) + ts_forecast_outputs = None + if ts_forecast and ts_values is not None: + llm_embedding_mask = attention_mask.bool() if attention_mask is not None else None + ts_forecast_outputs = self.time_series_forecaster( + history=outputs.ts_history, + llm_embedding_input=hidden_states, + ts_encoder_embedding_input=outputs.ts_encoder_embedding, + llm_embedding_mask=llm_embedding_mask, + ts_encoder_embedding_mask=outputs.ts_encoder_embedding_mask, + override_horizon=forecast_horizon, + ) + loss = None if labels is not None: loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size) @@ -2863,6 +5405,9 @@ class InternS2PreviewForConditionalGeneration(InternS2PreviewPreTrainedModel, Ge attentions=outputs.attentions, rope_deltas=outputs.rope_deltas, router_logits=outputs.router_logits, + ts_point_forecast=None if ts_forecast_outputs is None else ts_forecast_outputs.point_forecast, + ts_quantile_forecast=None if ts_forecast_outputs is None else ts_forecast_outputs.quantile_forecast, + ts_predicted_horizon=None if ts_forecast_outputs is None else ts_forecast_outputs.predicted_horizon, ) def prepare_inputs_for_generation( @@ -3107,6 +5652,7 @@ class InternS2PreviewForConditionalGeneration(InternS2PreviewPreTrainedModel, Ge ts_values: torch.FloatTensor | list[torch.FloatTensor], ts_lens: torch.Tensor | list[torch.Tensor], sr: torch.FloatTensor | list[torch.FloatTensor], + ts_channels: torch.LongTensor | None = None, ): r""" ts_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, input_size)`): @@ -3115,8 +5661,63 @@ class InternS2PreviewForConditionalGeneration(InternS2PreviewPreTrainedModel, Ge The length of the time series. sr (`torch.FloatTensor` of shape `(batch_size,)`): The sampling rate of the time series. + ts_channels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + The number of channels of the time series. """ - return self.model.get_ts_feature(ts_values, ts_lens, sr) + return self.model.get_ts_feature(ts_values, ts_lens, sr, ts_channels) + + def generate_with_time_series_input( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + pixel_values: torch.Tensor | None = None, + pixel_values_videos: torch.FloatTensor | None = None, + ts_values: torch.FloatTensor | list[torch.FloatTensor] = None, + image_grid_thw: torch.LongTensor | None = None, + video_grid_thw: torch.LongTensor | None = None, + ts_lens: torch.Tensor | list[torch.Tensor] = None, + ts_sr: torch.FloatTensor | list[torch.FloatTensor] = None, + ts_channels: torch.LongTensor | None = None, + cache_position: torch.LongTensor | None = None, + forecast_horizon: int | list[int] | None = None, + enable_forecasting: bool = False, + logits_to_keep: int | torch.Tensor = 1, + return_dict: bool | None = True, + **generate_kwargs, + ): + if ts_values is None: + raise ValueError("No time series was provided.") + + model_inputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + "past_key_values": past_key_values, + "inputs_embeds": inputs_embeds, + "pixel_values": pixel_values, + "pixel_values_videos": pixel_values_videos, + "ts_values": ts_values, + "image_grid_thw": image_grid_thw, + "video_grid_thw": video_grid_thw, + "ts_lens": ts_lens, + "ts_sr": ts_sr, + "ts_channels": ts_channels, + "cache_position": cache_position, + } + + if enable_forecasting: + return self( + **model_inputs, + ts_forecast=True, + forecast_horizon=forecast_horizon, + logits_to_keep=logits_to_keep, + return_dict=return_dict, + ) + + return self.generate(**model_inputs, **generate_kwargs) __all__ = [