dsipts 1.1.10__tar.gz → 1.1.11__tar.gz

{dsipts-1.1.10 → dsipts-1.1.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dsipts
-Version: 1.1.10
+Version: 1.1.11
 Summary: Unified library for timeseries modelling
 Author-email: Andrea Gobbi <agobbi@fbk.eu>
 Project-URL: Homepage, https://github.com/DSIP-FBK/DSIPTS

{dsipts-1.1.10 → dsipts-1.1.11}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "dsipts"
-version = "1.1.10"
+version = "1.1.11"
 description = "Unified library for timeseries modelling"
 readme = "README.md"
 requires-python = "==3.11.13"

{dsipts-1.1.10 → dsipts-1.1.11}/src/dsipts/data_structure/data_structure.py

@@ -35,7 +35,18 @@ from .modifiers import *
 from aim.pytorch_lightning import AimLogger
 import time
 
-
+class DummyScaler():
+    def __init__(self):
+        pass
+    def fit(self,x):
+        pass
+    def transform(self,x):
+        return x
+    def inverse_transform(self,x):
+        return x
+    def fit_transform(self,x):
+        return x
+    
 
 pd.options.mode.chained_assignment = None 
 log = logging.getLogger(__name__)
@@ -210,20 +221,23 @@ class TimeSeries():
         self.future_variables = []
         self.target_variables = ['signal']
         self.num_var = list(set(self.past_variables).union(set(self.future_variables)).union(set(self.target_variables)))
-        
+        self.num_var = list(np.sort(self.num_var))
         
     def enrich(self,dataset,columns):
-        if columns =='hour':
-            dataset[columns] = dataset.time.dt.hour
-        elif columns=='dow':
-            dataset[columns] = dataset.time.dt.weekday
-        elif columns=='month':
-            dataset[columns] = dataset.time.dt.month
-        elif columns=='minute':
-            dataset[columns] = dataset.time.dt.minute
-        else:
-            if columns  not in dataset.columns:
-                beauty_string(f'I can not automatically enrich column {columns}. Please contact the developers or add it manually to your dataset.','section',True)
+        try:
+            if columns =='hour':
+                dataset[columns] = dataset.time.dt.hour
+            elif columns=='dow':
+                dataset[columns] = dataset.time.dt.weekday
+            elif columns=='month':
+                dataset[columns] = dataset.time.dt.month
+            elif columns=='minute':
+                dataset[columns] = dataset.time.dt.minute
+            else:
+                if columns  not in dataset.columns:
+                    beauty_string(f'I can not automatically enrich column {columns}. Please contact the developers or add it manually to your dataset.','section',True)
+        except:
+            beauty_string(f'I can not automatically enrich column {columns}. Probably not a temporal index.','section',True)
 
     def load_signal(self,data:pd.DataFrame,
                     enrich_cat:List[str] = [],
@@ -300,7 +314,7 @@ class TimeSeries():
             if check_past:
                 beauty_string('I will update past column adding all target columns, if you want to avoid this beahviour please use check_pass as false','info',self.verbose)
                 past_variables = list(set(past_variables).union(set(target_variables)))
-
+                past_variables = list(np.sort(past_variables))
         self.cat_past_var = cat_past_var
         self.cat_fut_var = cat_fut_var
 
@@ -321,14 +335,18 @@ class TimeSeries():
                 beauty_string('Categorical {c} already present, it will be added to categorical variable but not call the enriching function','info',self.verbose) 
             else:
                 self.enrich(dataset,c)
+        self.cat_past_var = list(np.sort(self.cat_past_var))
+        self.cat_fut_var = list(np.sort(self.cat_fut_var))
+        
         self.cat_var = list(set(self.cat_past_var+self.cat_fut_var)) ## all categorical data
-
+        self.cat_var = list(np.sort(self.cat_var))
         self.dataset = dataset
         self.past_variables = past_variables
         self.future_variables = future_variables
         self.target_variables = target_variables
         self.out_vars = len(target_variables)
         self.num_var = list(set(self.past_variables).union(set(self.future_variables)).union(set(self.target_variables)))
+        self.num_var = list(np.sort(self.num_var))
         if silly_model:
             beauty_string('YOU ARE TRAINING A SILLY MODEL WITH THE TARGETS IN THE INPUTS','section',self.verbose) 
             self.future_variables+=self.target_variables
@@ -665,7 +683,11 @@ class TimeSeries():
         #self.model.apply(weight_init_zeros)
 
         self.config = config
-        
+        try:
+            self.model = torch.compile(self.model)
+        except:
+            beauty_string('Can not compile the model','block',self.verbose)
+
         beauty_string('Setting the model','block',self.verbose)
         beauty_string(model,'',self.verbose)
         

{dsipts-1.1.10 → dsipts-1.1.11}/src/dsipts/models/ITransformer.py

@@ -8,6 +8,8 @@ import numpy as np
 from .itransformer.Transformer_EncDec import Encoder, EncoderLayer
 from .itransformer.SelfAttention_Family import FullAttention, AttentionLayer
 from .itransformer.Embed import DataEmbedding_inverted
+from ..data_structure.utils import beauty_string
+from .utils import  get_scope,get_activation,Embedding_cat_variables
 
 try:
     import lightning.pytorch as pl
@@ -17,12 +19,6 @@ except:
     import pytorch_lightning as pl
     OLD_PL = True
     from .base import Base
-from .utils import QuantileLossMO,Permute, get_activation
-
-from typing import List, Union
-from ..data_structure.utils import beauty_string
-from .utils import  get_scope
-from .utils import Embedding_cat_variables
 
 
 
@@ -34,8 +30,6 @@ class ITransformer(Base):
     description = get_scope(handle_multivariate,handle_future_covariates,handle_categorical_variables,handle_quantile_loss)
     
     def __init__(self, 
-                
-
                  # specific params
                  hidden_size:int,
                  d_model: int,

dsipts-1.1.11/src/dsipts/models/TTM.py

@@ -0,0 +1,145 @@
+import torch
+import numpy as np
+from torch import  nn
+
+try:
+    import lightning.pytorch as pl
+    from .base_v2 import Base
+    OLD_PL = False
+except:
+    import pytorch_lightning as pl
+    OLD_PL = True
+    from .base import Base
+
+
+from .ttm.utils import get_model, get_frequency_token, count_parameters
+from ..data_structure.utils import beauty_string
+from .utils import  get_scope
+
+class TTM(Base):
+    handle_multivariate = True
+    handle_future_covariates = True
+    handle_categorical_variables = True
+    handle_quantile_loss = True
+    description = get_scope(handle_multivariate,handle_future_covariates,handle_categorical_variables,handle_quantile_loss)
+    
+    def __init__(self, 
+                model_path:str,
+                prefer_l1_loss:bool,  # exog: set true to use l1 loss
+                prefer_longer_context:bool,
+                prediction_channel_indices,
+                exogenous_channel_indices_cont,
+                exogenous_channel_indices_cat,
+                decoder_mode,
+                freq,
+                freq_prefix_tuning,
+                fcm_context_length,
+                fcm_use_mixer,
+                fcm_mix_layers,
+                fcm_prepend_past,
+                enable_forecast_channel_mixing,
+                **kwargs)->None:
+   
+        super().__init__(**kwargs)
+        self.save_hyperparameters(logger=False)
+
+
+        self.index_fut = list(exogenous_channel_indices_cont)
+
+        if len(exogenous_channel_indices_cat)>0:
+            self.index_fut_cat = (self.past_channels+len(self.embs_past))+list(exogenous_channel_indices_cat)
+        else:
+            self.index_fut_cat = []
+        self.freq = freq
+
+        self.model = get_model(
+            model_path=model_path,
+            context_length=self.past_steps,
+            prediction_length=self.future_steps,
+            prefer_l1_loss=prefer_l1_loss,
+            prefer_longer_context=prefer_longer_context,
+            num_input_channels=self.past_channels+len(self.embs_past), #giusto
+            decoder_mode=decoder_mode,
+            prediction_channel_indices=list(prediction_channel_indices),
+            exogenous_channel_indices=self.index_fut + self.index_fut_cat,
+            fcm_context_length=fcm_context_length,
+            fcm_use_mixer=fcm_use_mixer,
+            fcm_mix_layers=fcm_mix_layers,
+            freq=freq,
+            freq_prefix_tuning=freq_prefix_tuning,
+            fcm_prepend_past=fcm_prepend_past,
+            enable_forecast_channel_mixing=enable_forecast_channel_mixing,
+            
+        )
+        hidden_size =  self.model.config.hidden_size
+        self.model.prediction_head = torch.nn.Linear(hidden_size, self.out_channels*self.mul)
+        self._freeze_backbone()
+
+    def _freeze_backbone(self):
+        """
+        Freeze the backbone of the model.
+        This is useful when you want to fine-tune only the head of the model.
+        """
+        beauty_string(f"Number of params before freezing backbone:{count_parameters(self.model)}",'info',self.verbose)
+        
+        # Freeze the backbone of the model
+        for param in self.model.backbone.parameters():
+            param.requires_grad = False
+        # Count params
+        beauty_string(f"Number of params after freezing the backbone: {count_parameters(self.model)}",'info',self.verbose)
+            
+    
+    def _scaler_past(self, input):
+        for i, e in enumerate(self.embs_past):
+            input[:,:,i] = input[:, :, i] / (e-1)
+        return input
+    def _scaler_fut(self, input):
+        for i, e in enumerate(self.embs_fut):
+            input[:,:,i] = input[:, :, i] / (e-1)
+        return input
+
+
+    def forward(self, batch):
+        x_enc = batch['x_num_past']
+        original_indexes = batch['idx_target'][0].tolist()
+
+
+
+
+        if 'x_cat_past' in batch.keys():
+            x_mark_enc = batch['x_cat_past'].to(torch.float32).to(self.device)
+            x_mark_enc = self._scaler_past(x_mark_enc)
+            past_values = torch.cat((x_enc,x_mark_enc), axis=-1).type(torch.float32)
+        else:
+            past_values = x_enc
+        
+        future_values = torch.zeros_like(past_values)
+        future_values = future_values[:,:self.future_steps,:]
+
+        if 'x_num_future' in batch.keys(): 
+            future_values[:,:,self.index_fut] = batch['x_num_future'].to(self.device)
+        if 'x_cat_future' in batch.keys():
+            x_mark_dec = batch['x_cat_future'].to(torch.float32).to(self.device)
+            x_mark_dec = self._scaler_fut(x_mark_dec)
+            future_values[:,:,self.index_cat_fut] = x_mark_dec
+        
+
+        #investigating!!
+        freq_token = get_frequency_token(self.freq).repeat(past_values.shape[0])
+
+        res = self.model(
+            past_values= past_values,
+            future_values= future_values,# future_values if future_values.shape[0]>0 else None,
+            past_observed_mask = None,
+            future_observed_mask = None,
+            output_hidden_states =  False,
+            return_dict = False,
+            freq_token= freq_token, ##investigating
+            static_categorical_values = None
+        )
+
+
+        BS = res.shape[0]
+        return res.reshape(BS,self.future_steps,-1,self.mul)
+        
+    

{dsipts-1.1.10 → dsipts-1.1.11}/src/dsipts/models/base.py

@@ -111,8 +111,11 @@ class Base(pl.LightningModule):
         self.train_loss_epoch = -100.0
         self.verbose = verbose
         self.name = self.__class__.__name__
-        self.train_epoch_metrics = []
-        self.validation_epoch_metrics = []
+        self.register_buffer("train_epoch_metrics", torch.tensor(0.0))
+        self.register_buffer("validation_epoch_metrics", torch.tensor(0.0))
+        self.register_buffer("train_epoch_count", torch.tensor(0))
+        self.register_buffer("validation_epoch_count", torch.tensor(0))
+        
         
         self.use_quantiles = True if len(quantiles)>0 else False
         self.quantiles =  quantiles
@@ -295,7 +298,8 @@ class Base(pl.LightningModule):
                 y_hat = self(batch)
                 loss = self.compute_loss(batch,y_hat)
             
-        self.train_epoch_metrics.append(loss.item())
+        self.train_epoch_metrics+=loss.detach()
+        self.train_epoch_count +=1
         return loss
 
     
@@ -311,27 +315,20 @@ class Base(pl.LightningModule):
             y_hat = self(batch)
             score = 0
         if batch_idx==0:
-            if self.use_quantiles:
-                idx = 1
-            else:
-                idx = 0
-            #track the predictions! We can do better than this but maybe it is better to firstly update pytorch-lightening 
-
+           
             if self.count_epoch%int(max(self.trainer.max_epochs/100,1))==1:
-
-                for i in range(batch['y'].shape[2]):
-                    real =  batch['y'][0,:,i].cpu().detach().numpy()
-                    pred =  y_hat[0,:,i,idx].cpu().detach().numpy()
-                    fig, ax = plt.subplots(figsize=(7,5))  
-                    ax.plot(real,'o-',label='real')
-                    ax.plot(pred,'o-',label='pred')
-                    ax.legend()
-                    ax.set_title(f'Channel {i} first element first batch validation {int(100*self.count_epoch/self.trainer.max_epochs)}%')
-                    self.logger.experiment.track(Image(fig), name='cm_training_end')
-                    #self.log(f"example_{i}", np.stack([real, pred]).T,sync_dist=True)
-
-        return self.compute_loss(batch,y_hat)+score
-
+                self._val_outputs.append({
+                    "y": batch['y'].detach().cpu(),
+                    "y_hat": y_hat.detach().cpu()
+                })                
+        self.validation_epoch_metrics = (self.compute_loss(batch,y_hat)+score).detach()
+        self.validation_epoch_count+=1
+
+        return None #self.compute_loss(batch,y_hat)+score
+    
+    def on_validation_start(self):
+        # reset buffer each epoch
+        self._val_outputs = []
 
     def validation_epoch_end(self, outs):
         """
@@ -339,14 +336,30 @@ class Base(pl.LightningModule):
         
         :meta private:
         """
-        if len(outs)==0:
-            loss = 10000
-            beauty_string(f'THIS IS A BUG, It should be polulated','info',self.verbose)
-        else:
-            loss = torch.stack(outs).mean()
-        
-        self.log("val_loss", loss.item(),sync_dist=True)
-        beauty_string(f'Epoch: {self.count_epoch} train error: {self.train_loss_epoch:.4f} validation loss: {loss.item():.4f}','info',self.verbose)
+        if len(self._val_outputs)>0:
+            ys = torch.cat([o["y"] for o in self._val_outputs])
+            y_hats = torch.cat([o["y_hat"] for o in self._val_outputs])
+            if self.use_quantiles:
+                idx = 1
+            else:
+                idx = 0
+            for i in range(ys.shape[2]):
+                real =  ys[0,:,i].cpu().detach().numpy()
+                pred =  y_hats[0,:,i,idx].cpu().detach().numpy()
+                fig, ax = plt.subplots(figsize=(7,5))  
+                ax.plot(real,'o-',label='real')
+                ax.plot(pred,'o-',label='pred')
+                ax.legend()
+                ax.set_title(f'Channel {i} first element first batch validation {int(100*self.count_epoch/self.trainer.max_epochs)}%')
+                self.logger.experiment.track(Image(fig), name='cm_training_end')
+                #self.log(f"example_{i}", np.stack([real, pred]).T,sync_dist=True)
+                plt.close(fig) 
+        avg = self.validation_epoch_metrics/self.validation_epoch_count
+
+        self.validation_epoch_metrics.zero_()
+        self.validation_epoch_count.zero_()
+        self.log("val_loss", avg,sync_dist=True)
+        beauty_string(f'Epoch: {self.count_epoch} train error: {self.train_loss_epoch:.4f} validation loss: {avg:.4f}','info',self.verbose)
 
     def training_epoch_end(self, outs):
         """
@@ -355,12 +368,11 @@ class Base(pl.LightningModule):
         :meta private:
         """
 
-        loss = sum(outs['loss'] for outs in outs) / len(outs)
-        self.log("train_loss", loss.item(),sync_dist=True)
+        loss = self.train_epoch_metrics/self.global_step
+        self.log("train_loss", loss,sync_dist=True)
         self.count_epoch+=1    
 
-        self.train_loss_epoch = loss.item()
-
+        self.train_loss_epoch = loss
     def compute_loss(self,batch,y_hat):
         """
         custom loss calculation

{dsipts-1.1.10 → dsipts-1.1.11}/src/dsipts/models/base_v2.py

@@ -15,7 +15,6 @@ from typing import List, Union
 from .utils import QuantileLossMO, CPRS
 import torch.nn as nn
 
-
 def standardize_momentum(x,order):
     mean = torch.mean(x,1).unsqueeze(1).repeat(1,x.shape[1],1)
     num = torch.pow(x-mean,order).mean(axis=1)
@@ -113,8 +112,13 @@ class Base(pl.LightningModule):
         self.train_loss_epoch = -100.0
         self.verbose = verbose
         self.name = self.__class__.__name__
-        self.train_epoch_metrics = []
-        self.validation_epoch_metrics = []
+        #self.train_epoch_metrics = 0
+        #self.validation_epoch_metrics = 0
+        
+        self.register_buffer("train_epoch_metrics", torch.tensor(0.0))
+        self.register_buffer("validation_epoch_metrics", torch.tensor(0.0))
+        self.register_buffer("train_epoch_count", torch.tensor(0))
+        self.register_buffer("validation_epoch_count", torch.tensor(0))
         
         self.use_quantiles = True if len(quantiles)>0 else False
         self.quantiles =  quantiles
@@ -299,7 +303,8 @@ class Base(pl.LightningModule):
                 y_hat = self(batch)
                 loss = self.compute_loss(batch,y_hat)
             
-        self.train_epoch_metrics.append(loss.item())
+        self.train_epoch_metrics+=loss.detach()
+        self.train_epoch_count +=1
         return loss
 
     
@@ -316,41 +321,54 @@ class Base(pl.LightningModule):
             y_hat = self(batch)
             score = 0
         if batch_idx==0:
-            if self.use_quantiles:
-                idx = 1
-            else:
-                idx = 0
+
             #track the predictions! We can do better than this but maybe it is better to firstly update pytorch-lightening 
 
             if self.count_epoch%int(max(self.trainer.max_epochs/100,1))==1:
+                self._val_outputs.append({
+                    "y": batch['y'].detach().cpu(),
+                    "y_hat": y_hat.detach().cpu()
+                })                
+        self.validation_epoch_metrics+= (self.compute_loss(batch,y_hat)+score).detach()
+        self.validation_epoch_count+=1
+        return None
 
-                for i in range(batch['y'].shape[2]):
-                    real =  batch['y'][0,:,i].cpu().detach().numpy()
-                    pred =  y_hat[0,:,i,idx].cpu().detach().numpy()
-                    fig, ax = plt.subplots(figsize=(7,5))  
-                    ax.plot(real,'o-',label='real')
-                    ax.plot(pred,'o-',label='pred')
-                    ax.legend()
-                    ax.set_title(f'Channel {i} first element first batch validation {int(100*self.count_epoch/self.trainer.max_epochs)}%')
-                    self.logger.experiment.track(Image(fig), name='cm_training_end')
-                    #self.log(f"example_{i}", np.stack([real, pred]).T,sync_dist=True)
-        self.validation_epoch_metrics.append(self.compute_loss(batch,y_hat)+score)
-        return 
-
+    def on_validation_start(self):
+        # reset buffer each epoch
+        self._val_outputs = []
+    
 
     def on_validation_epoch_end(self):
         """
         pythotrch lightening stuff
         
         :meta private:
-        """
+        """   
 
-        if len(self.validation_epoch_metrics)==0:
-            avg = 10000
-            beauty_string(f'THIS IS A BUG, It should be polulated','info',self.verbose)
-        else:
-            avg = torch.stack(self.validation_epoch_metrics).mean()
-        self.validation_epoch_metrics = []
+        if len(self._val_outputs)>0:
+            ys = torch.cat([o["y"] for o in self._val_outputs])
+            y_hats = torch.cat([o["y_hat"] for o in self._val_outputs])
+            if self.use_quantiles:
+                idx = 1
+            else:
+                idx = 0
+            for i in range(ys.shape[2]):
+                real =  ys[0,:,i].cpu().detach().numpy()
+                pred =  y_hats[0,:,i,idx].cpu().detach().numpy()
+                fig, ax = plt.subplots(figsize=(7,5))  
+                ax.plot(real,'o-',label='real')
+                ax.plot(pred,'o-',label='pred')
+                ax.legend()
+                ax.set_title(f'Channel {i} first element first batch validation {int(100*self.count_epoch/self.trainer.max_epochs)}%')
+                self.logger.experiment.track(Image(fig), name='cm_training_end')
+                #self.log(f"example_{i}", np.stack([real, pred]).T,sync_dist=True)
+                plt.close(fig) 
+            
+         
+        avg = self.validation_epoch_metrics/self.validation_epoch_count
+
+        self.validation_epoch_metrics.zero_()
+        self.validation_epoch_count.zero_()
         self.log("val_loss", avg,sync_dist=True)
         beauty_string(f'Epoch: {self.count_epoch} train error: {self.train_loss_epoch:.4f} validation loss: {avg:.4f}','info',self.verbose)
 
@@ -361,14 +379,12 @@ class Base(pl.LightningModule):
         
         :meta private:
         """
-        if len(self.train_epoch_metrics)==0:
-            avg = 0
-            beauty_string(f'THIS IS A BUG, It should be polulated','info',self.verbose)
-        else:
-            avg = np.stack(self.train_epoch_metrics).mean()
+
+        avg = self.train_epoch_metrics/self.train_epoch_count
         self.log("train_loss", avg,sync_dist=True)
         self.count_epoch+=1    
-        self.train_epoch_metrics = []
+        self.train_epoch_metrics.zero_()
+        self.train_epoch_count.zero_()
         self.train_loss_epoch = avg
 
     def compute_loss(self,batch,y_hat):

{dsipts-1.1.10 → dsipts-1.1.11}/src/dsipts.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dsipts
-Version: 1.1.10
+Version: 1.1.11
 Summary: Unified library for timeseries modelling
 Author-email: Andrea Gobbi <agobbi@fbk.eu>
 Project-URL: Homepage, https://github.com/DSIP-FBK/DSIPTS

dsipts-1.1.10/src/dsipts/models/TTM.py

@@ -1,252 +0,0 @@
-import torch
-import numpy as np
-from torch import  nn
-
-try:
-    import lightning.pytorch as pl
-    from .base_v2 import Base
-    OLD_PL = False
-except:
-    import pytorch_lightning as pl
-    OLD_PL = True
-    from .base import Base
-
-
-
-from typing import List,Union
-
-from .utils import  QuantileLossMO
-from ..data_structure.utils import beauty_string
-from .ttm.utils import get_model, get_frequency_token, count_parameters, RMSELoss
-
-
-class TTM(Base):
-    def __init__(self, 
-                model_path:str,
-                past_steps:int,
-                future_steps:int,
-                freq_prefix_tuning:bool,
-                freq:str,
-                prefer_l1_loss:bool,  # exog: set true to use l1 loss
-                prefer_longer_context:bool,
-                loss_type:str,
-                num_input_channels,
-                prediction_channel_indices,
-                exogenous_channel_indices,
-                decoder_mode,
-                fcm_context_length,
-                fcm_use_mixer,
-                fcm_mix_layers,
-                fcm_prepend_past,
-                enable_forecast_channel_mixing,
-                out_channels:int,
-                embs:List[int],
-                remove_last = False,
-                optim:Union[str,None]=None,
-                optim_config:dict=None,
-                scheduler_config:dict=None,
-                verbose = False,
-                use_quantiles=False,
-                persistence_weight:float=0.0,
-                quantiles:List[int]=[],
-                **kwargs)->None:
-        """TODO and FIX for future and past categorical variables
-        
-        Args:
-            model_path (str): _description_
-            past_steps (int): _description_
-            future_steps (int): _description_
-            freq_prefix_tuning (bool): _description_
-            freq (str): _description_
-            prefer_l1_loss (bool): _description_
-            loss_type (str): _description_
-            num_input_channels (_type_): _description_
-            prediction_channel_indices (_type_): _description_
-            exogenous_channel_indices (_type_): _description_
-            decoder_mode (_type_): _description_
-            fcm_context_length (_type_): _description_
-            fcm_use_mixer (_type_): _description_
-            fcm_mix_layers (_type_): _description_
-            fcm_prepend_past (_type_): _description_
-            enable_forecast_channel_mixing (_type_): _description_
-            out_channels (int): _description_
-            embs (List[int]): _description_
-            remove_last (bool, optional): _description_. Defaults to False.
-            optim (Union[str,None], optional): _description_. Defaults to None.
-            optim_config (dict, optional): _description_. Defaults to None.
-            scheduler_config (dict, optional): _description_. Defaults to None.
-            verbose (bool, optional): _description_. Defaults to False.
-            use_quantiles (bool, optional): _description_. Defaults to False.
-            persistence_weight (float, optional): _description_. Defaults to 0.0.
-            quantiles (List[int], optional): _description_. Defaults to [].
-        """
-        super(TTM, self).__init__(verbose)
-        self.save_hyperparameters(logger=False)
-        self.future_steps = future_steps
-        self.use_quantiles = use_quantiles
-        self.optim = optim
-        self.optim_config = optim_config
-        self.scheduler_config = scheduler_config
-        self.persistence_weight = persistence_weight 
-        self.loss_type = loss_type
-        self.remove_last = remove_last
-        self.embs = embs
-        self.freq = freq
-        self.extend_variables = False
-
-        # NOTE: For Hydra
-        prediction_channel_indices = list(prediction_channel_indices)
-        exogenous_channel_indices = list(exogenous_channel_indices)
-
-        if len(quantiles)>0:
-            assert len(quantiles)==3, beauty_string('ONLY 3 quantiles premitted','info',True)
-            self.use_quantiles = True
-            self.mul = len(quantiles)
-            self.loss = QuantileLossMO(quantiles)
-            self.extend_variables = True
-            if out_channels * 3 != len(prediction_channel_indices):
-                prediction_channel_indices, exogenous_channel_indices, num_input_channels = self.__add_quantile_features(prediction_channel_indices, 
-                                                                                                                         exogenous_channel_indices, 
-                                                                                                                         out_channels)
-        else:
-            self.mul = 1
-            if self.loss_type == 'mse':
-                self.loss = nn.MSELoss(reduction="mean")
-            elif self.loss_type == 'rmse':
-                self.loss = RMSELoss()
-            else:
-                self.loss = nn.L1Loss()
-        
-        self.model = get_model(
-            model_path=model_path,
-            context_length=past_steps,
-            prediction_length=future_steps,
-            freq_prefix_tuning=freq_prefix_tuning,
-            freq=freq,
-            prefer_l1_loss=prefer_l1_loss,
-            prefer_longer_context=prefer_longer_context,
-            num_input_channels=num_input_channels,
-            decoder_mode=decoder_mode,
-            prediction_channel_indices=list(prediction_channel_indices),
-            exogenous_channel_indices=list(exogenous_channel_indices),
-            fcm_context_length=fcm_context_length,
-            fcm_use_mixer=fcm_use_mixer,
-            fcm_mix_layers=fcm_mix_layers,
-            fcm_prepend_past=fcm_prepend_past,
-            #loss='mse',
-            enable_forecast_channel_mixing=enable_forecast_channel_mixing,
-        )
-        self.__freeze_backbone()
-
-    def __add_quantile_features(self, prediction_channel_indices, exogenous_channel_indices, out_channels):
-        prediction_channel_indices = list(range(out_channels * 3))
-        exogenous_channel_indices = [prediction_channel_indices[-1] + i for i in range(1, len(exogenous_channel_indices)+1)]
-        num_input_channels = len(prediction_channel_indices) + len(exogenous_channel_indices)
-        return prediction_channel_indices, exogenous_channel_indices, num_input_channels
-
-    def __freeze_backbone(self):
-        """
-        Freeze the backbone of the model.
-        This is useful when you want to fine-tune only the head of the model.
-        """
-        print(
-            "Number of params before freezing backbone",
-            count_parameters(self.model),
-        )
-        # Freeze the backbone of the model
-        for param in self.model.backbone.parameters():
-            param.requires_grad = False
-        # Count params
-        print(
-            "Number of params after freezing the backbone",
-            count_parameters(self.model),
-        )
-    
-    def __scaler(self, input):
-        #new_data = torch.tensor([MinMaxScaler().fit_transform(step_data) for step_data in data])
-        for i, e in enumerate(self.embs):
-            input[:,:,i] = input[:, :, i] / (e-1)
-        return input
-    
-    def __build_tupla_indexes(self, size, target_idx, current_idx):
-        permute = list(range(size))
-        history = dict()
-        for j, i in enumerate(target_idx):
-            c = history.get(current_idx[j], current_idx[j])
-            permute[i], permute[c] = current_idx[j], i
-            history[i] = current_idx[j]
-
-
-    def __permute_indexes(self, values, target_idx, current_idx):
-        if current_idx is None or target_idx is None:
-            raise ValueError("Indexes cannot be None")
-        if sorted(current_idx) != sorted(target_idx):
-            return values[..., self.__build_tupla_indexes(values.shape[-1], target_idx, current_idx)]
-        return values
-    
-    def __extend_with_quantile_variables(self, x, original_indexes):
-        covariate_indexes = [i for i in range(x.shape[-1]) if i not in original_indexes]
-        covariate_tensors = x[..., covariate_indexes]
-
-        new_tensors = [x[..., target_index] for target_index in original_indexes for _ in range(3)]
-
-        new_original_indexes = list(range(len(original_indexes) * 3))
-        return torch.cat([torch.stack(new_tensors, dim=-1), covariate_tensors], dim=-1), new_original_indexes
-    
-    def forward(self, batch):
-        x_enc = batch['x_num_past']
-        original_indexes = batch['idx_target'][0].tolist()
-        original_indexes_future = batch['idx_target_future'][0].tolist()
-
-
-        if self.extend_variables:
-            x_enc, original_indexes = self.__extend_with_quantile_variables(x_enc, original_indexes)
-
-        if 'x_cat_past' in batch.keys():
-            x_mark_enc = batch['x_cat_past'].to(torch.float32).to(self.device)
-            x_mark_enc = self.__scaler(x_mark_enc)
-            past_values = torch.cat((x_enc,x_mark_enc), axis=-1).type(torch.float32)
-        else:
-            past_values = x_enc
-        
-        x_dec = torch.tensor([]).to(self.device)
-        if 'x_num_future' in batch.keys(): 
-            x_dec = batch['x_num_future'].to(self.device)
-            if self.extend_variables:
-                x_dec, original_indexes_future = self.__extend_with_quantile_variables(x_dec, original_indexes_future)
-        if 'x_cat_future' in batch.keys():
-            x_mark_dec = batch['x_cat_future'].to(torch.float32).to(self.device)
-            x_mark_dec = self.__scaler(x_mark_dec)
-            future_values = torch.cat((x_dec, x_mark_dec), axis=-1).type(torch.float32)
-        else:
-            future_values = x_dec
-
-        if self.remove_last:
-            idx_target = batch['idx_target'][0]
-            x_start = x_enc[:,-1,idx_target].unsqueeze(1)
-            x_enc[:,:,idx_target]-=x_start 
-
-
-        past_values = self.__permute_indexes(past_values, self.model.prediction_channel_indices, original_indexes)
-
-        
-        future_values = self.__permute_indexes(future_values, self.model.prediction_channel_indices, original_indexes_future)
-
-        freq_token = get_frequency_token(self.freq).repeat(x_enc.shape[0])
-
-        res = self.model(
-            past_values= past_values,
-            future_values= future_values,
-            past_observed_mask = None,
-            future_observed_mask = None,
-            output_hidden_states =  False,
-            return_dict = False,
-            freq_token= freq_token,
-            static_categorical_values = None
-        )
-        #args = None
-        #res = self.model(**args)
-        BS = res.shape[0]
-        return res.reshape(BS,self.future_steps,-1,self.mul)
-        
-