Upload modeling_moment.py

modeling_moment.py

@@ -432,6 +432,7 @@ class MomentEmbeddingModel(MomentPreTrainedModel):
         x_enc = self.normalizer(x=x_enc, mask=input_mask, mode="norm")
         x_enc = torch.nan_to_num(x_enc, nan=0, posinf=0, neginf=0)
 
+        # [batch_size x n_patches]
         input_mask_patch_view = Masking.convert_seq_to_patch_view(
             input_mask, self.patch_len
         )
@@ -453,11 +454,6 @@ class MomentEmbeddingModel(MomentPreTrainedModel):
         enc_out = enc_out.reshape((-1, n_channels, n_patches, self.config.d_model))
         # [batch_size x n_channels x n_patches x d_model]
 
-        # For Mists model
-        # [batch_size, n_channels x n_patches, d_model]
-        # Ensure hidden_states are consistent for both short and long inputs with input_mask specified
-        hidden_states = hidden_states.reshape(batch_size, n_channels, n_patches, self.config.d_model).transpose(1, 2).reshape(batch_size, -1, self.config.d_model)
-
         if reduction == "mean":
             enc_out = enc_out.mean(dim=1, keepdim=False)  # Mean across channels
             # [batch_size x n_patches x d_model]
@@ -469,6 +465,20 @@ class MomentEmbeddingModel(MomentPreTrainedModel):
             ) / input_mask_patch_view.sum(dim=1)
         else:
             raise NotImplementedError(f"Reduction method {reduction} not implemented.")
+        
+        # For Mists model
+        # [batch_size, n_channels x n_patches, d_model]
+        # Ensure hidden_states are consistent for both short and long inputs with input_mask specified
+        # hidden_states = hidden_states.reshape(batch_size, n_channels, n_patches, self.config.d_model).transpose(1, 2).reshape(batch_size, -1, self.config.d_model)
+        # [batch_size x n_channels x n_patches x d_model]
+        hidden_states = hidden_states.reshape(batch_size, n_channels, n_patches, self.config.d_model)
+        # [batch_size x n_patches]
+        input_mask_patch_view_for_hidden_states = Masking.convert_seq_to_patch_view(input_mask, self.patch_len)
+        # [batch_size　x n_channels x n_patches x d_model]
+        input_mask_patch_view_for_hidden_states = input_mask_patch_view_for_hidden_states.unsqueeze(-1).repeat(
+            1, n_channels, 1, self.config.d_model
+        )
+        hidden_states = input_mask_patch_view_for_hidden_states * hidden_states
 
         return TimeseriesOutputs(
             embeddings=enc_out, input_mask=input_mask, metadata=reduction, hidden_states=hidden_states