## Tasks

A task is the PyTorch lightning module that defines your training, validation and test steps as well as optimizer and learning rate scheduler configuration.
There are two key python modules in this subdirectory that are needed for foundation model training:

- **`mae_reconstruction.py`**: Implements a pretraining task using a Masked Autoencoder (MAE) for reconstruction purposes. The model processes input data (such as spectrograms), masks parts of it, and uses an encoder-decoder structure to predict the missing patches. It supports various optimizers and logging metrics for training and validation.

- **`finetune_pretrained_model.py`**: Fine-tunes a pretrained MAE model for classification tasks. It incorporates augmentations like SpecAugment and white noise addition. Metrics such as accuracy, AUROC, and balanced accuracy are tracked during training, validation, and testing. The module supports freezing the pretrained encoder while training the classification head.

**How to add a new task**
1. Add the code for the task to this subdirectory.
2. Add the configuration file of the model to [`./config/task`](https://github.com/ofsoundof/TimeFM/tree/split_attention_foundation_model/config/task).


