models/MOMENT.py

"""
This function is adapted from [moment] by [mononitogoswami]
Original source: [https://github.com/moment-timeseries-foundation-model/moment]
"""

from momentfm import MOMENTPipeline
from momentfm.utils.masking import Masking
from momentfm.utils.utils import control_randomness
from sklearn.preprocessing import MinMaxScaler
import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
from torch import nn
import math

from .base import BaseDetector
from ..utils.dataset import ReconstructDataset_Moment
from ..utils.torch_utility import EarlyStoppingTorch, get_gpu

class MOMENT(BaseDetector):
    def __init__(self, 
                 win_size=256, 
                 input_c=1, 
                 batch_size=128,
                 epochs=2,
                 validation_size=0,
                 lr=1e-4):

        self.model_name = 'MOMENT'
        self.win_size = win_size
        self.input_c = input_c
        self.batch_size = batch_size
        self.anomaly_criterion = nn.MSELoss(reduce=False)
        self.epochs = epochs
        self.validation_size = validation_size
        self.lr = lr

        cuda = True        
        self.cuda = cuda
        self.device = get_gpu(self.cuda)

        # Control randomness for reproducibility
        control_randomness(seed=42)

        # Load the model properly with config
        try:
            self.model = MOMENTPipeline.from_pretrained(
                "AutonLab/MOMENT-1-base",
                model_kwargs={
                    "task_name": "reconstruction",
                    "n_channels": self.input_c,
                    "max_seq_len": self.win_size
                }
            )
            self.model.init()
        except Exception as e:
            # Fallback: try alternative initialization
            print(f"Failed to load MOMENT model with from_pretrained: {e}")
            print("Attempting alternative initialization...")
            from transformers import AutoConfig
            config = AutoConfig.from_pretrained("AutonLab/MOMENT-1-base")
            config.task_name = "reconstruction"
            config.n_channels = self.input_c
            config.max_seq_len = self.win_size
            self.model = MOMENTPipeline(config)
            self.model.init()

        self.model = self.model.to(self.device).float()
        # Optimize Mean Squarred Error using your favourite optimizer
        self.criterion = torch.nn.MSELoss() 
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)
        self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=5, gamma=0.75)
        self.save_path = None
        self.early_stopping = EarlyStoppingTorch(save_path=self.save_path, patience=3)

    def zero_shot(self, data):

        test_loader = DataLoader(
            dataset=ReconstructDataset_Moment(data, window_size=self.win_size, normalize=True),
            batch_size=self.batch_size,
            shuffle=False)

        trues, preds = [], []
        self.score_list = []
        with torch.no_grad():
            for batch_x, batch_masks in tqdm(test_loader, total=len(test_loader)):
                batch_x = batch_x.to("cuda").float()
                batch_masks = batch_masks.to("cuda")
                batch_x = batch_x.permute(0,2,1)

                # print('batch_x: ', batch_x.shape)             # [batch_size, n_channels, window_size]
                # print('batch_masks: ', batch_masks.shape)     # [batch_size, window_size]

                output = self.model(x_enc=batch_x, input_mask=batch_masks) # [batch_size, n_channels, window_size]
                score = torch.mean(self.anomaly_criterion(batch_x, output.reconstruction), dim=-1).detach().cpu().numpy()[:, -1]
                self.score_list.append(score)

        self.__anomaly_score = np.concatenate(self.score_list, axis=0).reshape(-1)

        if self.__anomaly_score.shape[0] < len(data):
            self.__anomaly_score = np.array([self.__anomaly_score[0]]*math.ceil((self.win_size-1)/2) + 
                        list(self.__anomaly_score) + [self.__anomaly_score[-1]]*((self.win_size-1)//2))
        self.decision_scores_ = self.__anomaly_score


    def fit(self, data):
        tsTrain = data[:int((1-self.validation_size)*len(data))]
        tsValid = data[int((1-self.validation_size)*len(data)):]

        train_loader = DataLoader(
            dataset=ReconstructDataset_Moment(tsTrain, window_size=self.win_size),
            batch_size=self.batch_size,
            shuffle=True
        )
        
        valid_loader = DataLoader(
            dataset=ReconstructDataset_Moment(tsValid, window_size=self.win_size),
            batch_size=self.batch_size,
            shuffle=False
        )

        mask_generator = Masking(mask_ratio=0.3) # Mask 30% of patches randomly 


        for epoch in range(1, self.epochs + 1):
            self.model.train()
            for batch_x, batch_masks in tqdm(train_loader, total=len(train_loader)):
                batch_x = batch_x.to(self.device).float()
                batch_x = batch_x.permute(0,2,1)
                # print('batch_x: ', batch_x.shape)

                original = batch_x
                n_channels = batch_x.shape[1]
                
                # Reshape to [batch_size * n_channels, 1, window_size]
                batch_x = batch_x.reshape((-1, 1, self.win_size)) 
                
                batch_masks = batch_masks.to(self.device).long()
                batch_masks = batch_masks.repeat_interleave(n_channels, axis=0)
                
                # Randomly mask some patches of data
                mask = mask_generator.generate_mask(
                    x=batch_x, input_mask=batch_masks).to(self.device).long()
                
                mask = torch.nn.functional.pad(mask, (0, batch_masks.size(1) - mask.size(1)), mode='constant', value=1)

                # Forward
                model_output = self.model(batch_x, input_mask=batch_masks, mask=mask).reconstruction
                model_output = torch.nn.functional.pad(model_output, (0, original.size(2)-model_output.size(2)), mode='replicate')

                output = model_output.reshape(original.size(0), n_channels, self.win_size)

                # Compute loss
                loss = self.criterion(output, original)
                    
                # print(f"loss: {loss.item()}")
                
                # Backward
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()

            # self.model.eval()
            # avg_loss = 0
            # with torch.no_grad():
            #     for batch_x, batch_masks in tqdm(valid_loader, total=len(valid_loader)):
            #         batch_x = batch_x.to("cuda").float()
            #         batch_masks = batch_masks.to("cuda")
            #         batch_x = batch_x.permute(0,2,1)

            #         print('batch_x: ', batch_x.shape)
            #         print('batch_masks: ', batch_masks.shape)

            #         output = self.model(batch_x, input_mask=batch_masks) 

            #         loss = self.criterion(output.reconstruction.reshape(-1, n_channels, self.win_size), batch_x)
            #         print(f"loss: {loss.item()}")
            #         avg_loss += loss.cpu().item()

            # valid_loss = avg_loss/max(len(valid_loader), 1)
            # self.scheduler.step()
            # self.early_stopping(valid_loss, self.model)
            # if self.early_stopping.early_stop:
            #     print("   Early stopping<<<")
            #     break
        
    def decision_function(self, data):
        """
        Not used, present for API consistency by convention.
        """

        test_loader = DataLoader(
            dataset=ReconstructDataset_Moment(data, window_size=self.win_size),
            batch_size=self.batch_size,
            shuffle=False)

        trues, preds = [], []
        self.score_list = []
        with torch.no_grad():
            for batch_x, batch_masks in tqdm(test_loader, total=len(test_loader)):
                batch_x = batch_x.to("cuda").float()
                batch_masks = batch_masks.to("cuda")
                batch_x = batch_x.permute(0,2,1)

                # print('batch_x: ', batch_x.shape)             # [batch_size, n_channels, window_size]
                # print('batch_masks: ', batch_masks.shape)     # [batch_size, window_size]

                output = self.model(batch_x, input_mask=batch_masks) 
                score = torch.mean(self.anomaly_criterion(batch_x, output.reconstruction), dim=-1).detach().cpu().numpy()[:, -1]
                self.score_list.append(score)

        self.__anomaly_score = np.concatenate(self.score_list, axis=0).reshape(-1)

        if self.__anomaly_score.shape[0] < len(data):
            self.__anomaly_score = np.array([self.__anomaly_score[0]]*math.ceil((self.win_size-1)/2) + 
                        list(self.__anomaly_score) + [self.__anomaly_score[-1]]*((self.win_size-1)//2))

        return self.__anomaly_score