Upload 5 files

.gitattributes

@@ -35,3 +35,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 inference/input/test.mp3 filter=lfs diff=lfs merge=lfs -text
 m2e.png filter=lfs diff=lfs merge=lfs -text
+framework.png filter=lfs diff=lfs merge=lfs -text

data_loader.py

@@ -0,0 +1,78 @@
+import os
+import numpy as np
+import pickle
+from torch.utils import data
+import torchaudio.transforms as T
+import torchaudio
+import torch
+import csv
+import pytorch_lightning as pl
+from music2latent import EncoderDecoder
+import json
+import math
+from sklearn.preprocessing import StandardScaler
+
+from dataset_loaders.jamendo import JamendoDataset
+from dataset_loaders.pmemo import PMEmoDataset
+from dataset_loaders.deam import DEAMDataset
+from dataset_loaders.emomusic import EmoMusicDataset
+
+from omegaconf import DictConfig
+
+DATASET_REGISTRY = {
+    "jamendo": JamendoDataset,
+    "pmemo": PMEmoDataset,
+    "deam": DEAMDataset,
+    "emomusic": EmoMusicDataset
+}
+
+class DataModule(pl.LightningDataModule):
+    def __init__(self, cfg: DictConfig):
+        super().__init__()
+        self.cfg = cfg
+
+        self.train_datasets = []
+        self.val_datasets = []
+        self.test_datasets = []
+        
+    def setup(self, stage=None):
+        # Clear previous dataset lists
+        self.train_datasets = []
+        self.val_datasets = []
+        self.test_datasets = []
+
+        # Register the datasets and load them
+        for dataset_name in self.cfg.datasets:
+            dataset_cfg = self.cfg.dataset[dataset_name]
+
+            if dataset_name in DATASET_REGISTRY:
+                train_dataset = DATASET_REGISTRY[dataset_name](**dataset_cfg, cfg=self.cfg, tr_val='train')
+                val_dataset = DATASET_REGISTRY[dataset_name](**dataset_cfg, cfg=self.cfg, tr_val='validation')
+                test_dataset = DATASET_REGISTRY[dataset_name](**dataset_cfg, cfg=self.cfg, tr_val='test')
+
+                self.train_datasets.append(train_dataset)
+                self.val_datasets.append(val_dataset)
+                self.test_datasets.append(test_dataset)
+            else:
+                raise ValueError(f"Dataset {dataset_name} not found in registry")
+
+    def train_dataloader(self):
+        return [data.DataLoader(ds, batch_size=self.cfg.dataset[ds_name].batch_size, 
+                                shuffle=True, num_workers=self.cfg.dataset[ds_name].num_workers, 
+                                persistent_workers=True)
+                for ds, ds_name in zip(self.train_datasets, self.cfg.datasets)]
+
+    def val_dataloader(self):
+        return [data.DataLoader(ds, batch_size=self.cfg.dataset[ds_name].batch_size, 
+                                shuffle=False, num_workers=self.cfg.dataset[ds_name].num_workers, 
+                                persistent_workers=True)
+                for ds, ds_name in zip(self.val_datasets, self.cfg.datasets)]
+
+    def test_dataloader(self):
+        return [data.DataLoader(ds, batch_size=self.cfg.dataset[ds_name].batch_size, 
+                                shuffle=False, num_workers=self.cfg.dataset[ds_name].num_workers, 
+                                persistent_workers=True)
+                for ds, ds_name in zip(self.test_datasets, self.cfg.datasets)]
+
+
+

test.py

@@ -0,0 +1,98 @@
+import os
+import logging
+
+import torch
+torch.set_float32_matmul_precision("medium")
+
+import pytorch_lightning as pl
+from pytorch_lightning.loggers import TensorBoardLogger
+from pytorch_lightning.callbacks import ModelCheckpoint
+from data_loader import DataModule
+from trainer import MusicClassifier
+import yaml
+from omegaconf import DictConfig
+import hydra
+from hydra.utils import to_absolute_path
+from hydra.core.hydra_config import HydraConfig
+from pytorch_lightning.utilities.combined_loader import CombinedLoader
+
+
+log = logging.getLogger(__name__)
+
+def get_latest_version(log_dir):
+    version_dirs = [d for d in os.listdir(log_dir) if d.startswith('version_')]
+    version_dirs.sort(key=lambda x: int(x.split('_')[-1]))  # Sort by version number
+    return version_dirs[-1] if version_dirs else None
+
+def save_metrics_and_checkpoint(metrics, checkpoint, output_file):
+    data = {
+        'checkpoint': checkpoint,
+        'metrics': metrics
+    }
+    with open(output_file, 'w') as f:
+        yaml.dump(data, f)
+
+def read_best_checkpoint_info(file_path, dataset_type=None):
+    """Read the best checkpoint file."""
+    if not os.path.exists(file_path):
+        raise FileNotFoundError(f"Checkpoint info file not found: {file_path}")
+    
+    with open(file_path, 'r') as f:
+        lines = f.readlines()
+    
+    if dataset_type == "mood":
+        checkpoint_line = next((line for line in lines if line.startswith("Best checkpoint (mood):")), None)
+    elif dataset_type == "va":
+        checkpoint_line = next((line for line in lines if line.startswith("Best checkpoint (va):")), None)
+    else:
+        checkpoint_line = next((line for line in lines if line.startswith("Best checkpoint:")), None)
+
+    if not checkpoint_line:
+        raise ValueError(f"No checkpoint found for dataset type '{dataset_type}' in the file.")
+    
+    return checkpoint_line.split(": ")[-1].strip()
+
+@hydra.main(version_base=None, config_path="config", config_name="test_config")
+def main(config: DictConfig):
+    log.info("Testing starts")
+    log_base_dir = 'tb_logs/train_audio_classification'
+    # log_base_dir = to_absolute_path('tb_logs/train_audio_classification')
+
+    latest_version = get_latest_version(log_base_dir)
+    if not latest_version:
+        raise FileNotFoundError("No version directories found in log base directory.")
+    version_log_dir = os.path.join(log_base_dir, latest_version)
+    output_file = os.path.join(version_log_dir, 'test_metrics.txt')
+
+    if config.checkpoint_latest:
+        if config.multitask:
+            dataset_type = config.dataset_type  # Expecting 'mood' or 'va'
+            best_checkpoint_file = os.path.join(version_log_dir, 'best_checkpoint.txt')
+            ckpt = read_best_checkpoint_info(best_checkpoint_file, dataset_type)
+        else:
+            best_checkpoint_file = os.path.join(version_log_dir, 'best_checkpoint.txt')
+            ckpt = read_best_checkpoint_info(best_checkpoint_file)
+    else:
+        ckpt = config.checkpoint
+    if not os.path.exists(ckpt):
+        raise FileNotFoundError(f"Checkpoint file not found: {ckpt}")
+    
+    log.info(f"Using checkpoint: {ckpt}")
+    data_module = DataModule( config )
+    data_module.setup()
+
+    testloaders = {dataset_name: loader for dataset_name, loader in zip(config.datasets, data_module.test_dataloader())}
+    combined_test_loader = CombinedLoader(testloaders, mode="max_size")
+
+    model = MusicClassifier.load_from_checkpoint(ckpt, cfg=config, output_file=output_file)
+    logger = TensorBoardLogger(save_dir=log_base_dir,  
+                                name="",  
+                                version=latest_version)
+    trainer = pl.Trainer(**config.trainer,
+                        logger=logger)
+    
+    
+    trainer.test(model, combined_test_loader)
+
+if __name__ == '__main__':
+    main()

train.py

@@ -0,0 +1,107 @@
+import os
+import logging
+
+import torch
+torch.set_float32_matmul_precision("medium") 
+
+
+import pytorch_lightning as pl
+from pytorch_lightning.loggers import TensorBoardLogger
+from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
+from data_loader import DataModule
+from trainer import MusicClassifier
+from omegaconf import DictConfig
+import hydra
+from hydra.utils import to_absolute_path
+from hydra.core.hydra_config import HydraConfig
+from pytorch_lightning.callbacks import EarlyStopping
+from pytorch_lightning.utilities.combined_loader import CombinedLoader
+from pytorch_lightning.strategies import DDPStrategy
+
+#from utilities.custom_early_stopping import MultiMetricEarlyStopping
+def get_latest_version(log_dir):
+    version_dirs = [d for d in os.listdir(log_dir) if d.startswith('version_')]
+    version_dirs.sort(key=lambda x: int(x.split('_')[-1]))  # Sort by version number
+    return version_dirs[-1] if version_dirs else None
+
+log = logging.getLogger(__name__)
+@hydra.main(version_base=None, config_path="config", config_name="train_config")
+def main(config: DictConfig):
+
+    log_base_dir = 'tb_logs/train_audio_classification'
+    # log_base_dir = to_absolute_path('tb_logs/train_audio_classification')
+    is_mt = False
+    if "mt" in config.model.classifier:
+        is_mt = True
+
+    logger = TensorBoardLogger("tb_logs", name="train_audio_classification")
+    logger.log_hyperparams(config)
+    train_log_dir = logger.log_dir
+    print(f"Logging to {train_log_dir}")
+    log.info("Training starts")
+
+    data_module = DataModule( config )
+    data_module.setup()
+
+    # Get the list of dataloaders for both train and validation, with dataset names
+    trainloaders = {dataset_name: loader for dataset_name, loader in zip(config.datasets, data_module.train_dataloader())}
+    vallowers = {dataset_name: loader for dataset_name, loader in zip(config.datasets, data_module.val_dataloader())}
+    
+    # Combine multiple loaders using CombinedLoader, now with dataset names
+    combined_train_loader = CombinedLoader(trainloaders, mode="max_size")
+    combined_val_loader = CombinedLoader(vallowers, mode="max_size")
+
+    latest_version = get_latest_version(log_base_dir)
+    next_version = int(latest_version.split('_')[-1]) + 1 if latest_version else 0
+    next_version = f"version_{next_version}"
+
+    val_epoch_file = os.path.join(log_base_dir,  latest_version, 'val_epoch.txt')
+
+    model = MusicClassifier( config, output_file = val_epoch_file)
+
+    if is_mt:
+        checkpoint_callback_mood = ModelCheckpoint(**config.checkpoint_mood)
+        checkpoint_callback_va = ModelCheckpoint(**config.checkpoint_va)
+        early_stop_callback = EarlyStopping(**config.earlystopping)
+
+        if config.model.kd == True:
+            trainer = pl.Trainer(
+                **config.trainer,
+                strategy=DDPStrategy(find_unused_parameters=True),
+                callbacks=[checkpoint_callback_mood, checkpoint_callback_va, early_stop_callback],
+                logger=logger,
+                num_sanity_val_steps=0
+            )
+        else:
+            trainer = pl.Trainer(
+                **config.trainer,
+                strategy=DDPStrategy(find_unused_parameters=False),
+                callbacks=[checkpoint_callback_mood, checkpoint_callback_va, early_stop_callback],
+                logger=logger,
+                num_sanity_val_steps=0
+            )
+
+    else:
+        checkpoint_callback = ModelCheckpoint(**config.checkpoint)
+        # early_stop_callback = EarlyStopping(**config.earlystopping)
+        trainer = pl.Trainer(
+            **config.trainer,
+            callbacks=[checkpoint_callback, early_stop_callback],
+            logger=logger,
+            num_sanity_val_steps = 0
+        )
+    
+    trainer.fit(model, combined_train_loader, combined_val_loader)
+
+    if trainer.global_rank == 0:
+        best_checkpoint_file = os.path.join(train_log_dir, 'best_checkpoint.txt')
+        with open(best_checkpoint_file, 'w') as f:
+            if is_mt:
+                f.write(f"Best checkpoint (mood): {checkpoint_callback_mood.best_model_path}\n")
+                f.write(f"Best checkpoint (va): {checkpoint_callback_va.best_model_path}\n")
+            else:
+                f.write(f"Best checkpoint: {checkpoint_callback.best_model_path}\n")
+            f.write(f"Version: {logger.version}\n")
+
+if __name__ == '__main__':
+    main()

trainer.py

@@ -0,0 +1,478 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pytorch_lightning as pl
+from sklearn import metrics
+from transformers import AutoModelForAudioClassification
+import numpy as np
+from collections import OrderedDict
+from torchmetrics import MeanMetric, MaxMetric, Accuracy
+import torchmetrics.functional as tmf
+
+from model.linear import FeedforwardModel
+from model.linear_small import FeedforwardModelSmall
+from model.linear_attn_ck import FeedforwardModelAttnCK
+from model.linear_mt import FeedforwardModelMT
+from model.linear_mt_attn_ck import FeedforwardModelMTAttnCK
+
+import logging
+import yaml
+from omegaconf import DictConfig
+
+import torch
+from torch.distributed import all_gather, get_world_size
+# from lion_pytorch import Lion
+from torch_optimizer import RAdam
+
+def gather_all_results(tensor):
+    """
+    Gather tensors from all GPUs in distributed training.
+    """
+    gathered_tensors = [torch.zeros_like(tensor) for _ in range(get_world_size())]
+    all_gather(gathered_tensors, tensor)
+    return torch.cat(gathered_tensors, dim=0)
+
+# torch.set_float32_matmul_precision('medium')
+
+log = logging.getLogger(__name__)
+class MusicClassifier(pl.LightningModule):
+    def __init__(self, cfg: DictConfig, output_file = None):
+        super(MusicClassifier, self).__init__()
+        self.cfg = cfg
+        self.encoder = cfg.model.encoder
+        self.classifier = cfg.model.classifier
+        self.lr = cfg.model.lr
+        self.output_file = output_file
+        self.kd = cfg.model.kd
+        self.kd_weight = cfg.model.kd_weight
+        self.kd_temperature = self.cfg.model.kd_temperature
+
+        layer_size = len(self.cfg.model.layers)
+        mert_dim = 768 * layer_size
+
+        self.feature_dim_dict = {
+            "MERT": mert_dim
+        }
+
+        encoders = self.encoder.split("-")
+        self.input_size = sum(self.feature_dim_dict[encoder] for encoder in encoders)
+        self.num_datasets = len(self.cfg.datasets)
+
+        if "mt" in self.classifier:
+            if self.num_datasets < 2:
+                raise Exception("Error: Dataset size >= 2 needed for MT classifier")
+            classifiers = {
+                "linear-mt-attn-ck": FeedforwardModelMTAttnCK,
+            }
+            if self.classifier in classifiers:
+                self.model = classifiers[self.classifier](
+                    input_size=self.input_size,
+                    output_size_classification=56,
+                    output_size_regression=2
+                )
+            else:
+                raise Exception(f"Unknown classifier: {self.classifier}")          
+        else:
+            if self.num_datasets  >= 2:
+                raise Exception(f"Error: Dataset size == 1 needed for classifier") 
+            dataset_name = self.cfg.datasets[0]
+            self.output_size = self.cfg.dataset[dataset_name].output_size
+            classifiers = {
+                "linear": FeedforwardModel,
+                "linear-attn-ck": FeedforwardModelAttnCK
+            }
+
+            if self.classifier in classifiers:
+                self.model = classifiers[self.classifier](input_size=self.input_size, output_size=self.output_size)
+            else:
+                raise Exception(f"Unknown classifier: {self.classifier}")
+        
+                
+        if self.kd:
+            self.teacher_models = {}
+            
+            for dataset in self.cfg.datasets:
+                self.output_size = self.cfg.dataset[dataset].output_size
+                teacher_model_path = getattr(self.cfg, f"checkpoint_{dataset}", None)
+                
+                if teacher_model_path:
+                    # Create a new teacher model instance
+                    teacher_model = FeedforwardModelAttnCK(
+                        input_size=self.input_size,
+                        output_size=self.output_size,
+                    )
+                    
+                    # Load the checkpoint
+                    checkpoint = torch.load(teacher_model_path, map_location=self.device, weights_only=False)
+                    state_dict = checkpoint["state_dict"]
+                    
+                    # Adjust the keys in the state_dict
+                    state_dict = {key.replace("model.", ""): value for key, value in state_dict.items()}
+                    
+                    # Filter state_dict to match model's keys
+                    model_keys = set(teacher_model.state_dict().keys())
+                    filtered_state_dict = {key: value for key, value in state_dict.items() if key in model_keys}
+                    
+                    # Load the filtered state_dict and set the model to evaluation mode
+                    teacher_model.load_state_dict(filtered_state_dict)
+                    teacher_model.to(self.device)
+
+                    teacher_model.eval()
+                    
+                    # Store the teacher model in the dictionary with the dataset name as the key
+                    self.teacher_models[dataset] = teacher_model
+
+        probas = torch.from_numpy(np.load("dataset/jamendo/meta/probas_train.npy"))
+        pos_weight = torch.tensor(1.) / probas
+        weight = torch.tensor(2.) / (torch.tensor(1.) + pos_weight)
+        
+        self.loss_fn_classification = nn.BCEWithLogitsLoss(
+            pos_weight=pos_weight,reduction="mean",weight=weight
+        )
+        self.loss_fn_classification_eval = nn.BCEWithLogitsLoss(
+            pos_weight=pos_weight,reduction="none",weight=weight
+        )
+        
+        self.loss_fn_regression = nn.MSELoss()  
+
+        self.loss_kd = nn.KLDivLoss(reduction="batchmean")
+
+        self.prd_array = []
+        self.gt_array = []
+        self.song_array = []
+
+        self.prd_array_va = []
+        self.gt_array_va = []
+        self.song_array_va = []
+
+        self.validation_predictions = []
+        self.validation_targets = []
+        self.validation_results = {'preds': [], 'gt': []}
+
+        self.trn_loss = MeanMetric()
+        self.val_loss = MeanMetric()
+
+    def forward(self, model_input_dic, output_idx = 0):
+        if "mt" in self.classifier:
+            classification_output, regression_output = self.model(model_input_dic)
+            if output_idx == 0:
+                return classification_output
+            elif output_idx == 1:
+                return regression_output
+            elif output_idx == 2:
+                return classification_output, regression_output
+        else:
+            output = self.model(model_input_dic)
+            return output
+        
+    def compute_classification_loss(self, model_input_dic, y_mood):
+        classification_logits = self(model_input_dic, 0)
+        loss= self.loss_fn_classification(classification_logits, y_mood)
+        return loss
+
+    def compute_regression_loss(self, model_input_dic, y_va):
+        regression_output = self(model_input_dic, 1)
+        loss = self.loss_fn_regression(regression_output, y_va)
+        return loss
+    
+    def compute_mt_loss(self, model_input_dic, y_mood, y_va):
+        classification_logits, regression_output = self(model_input_dic, 2)
+        loss_classification = self.loss_fn_classification(classification_logits, y_mood)
+        loss_regression = self.loss_fn_regression(regression_output, y_va)
+        return loss_classification, loss_regression
+    
+
+    def compute_kd_loss(self, model_input_dic, y_mood, y_va, dataset_name):
+        """
+        Compute knowledge distillation loss for a given dataset.
+        """
+        # Forward pass through student model
+        s_logits_mood, s_logits_va = self(model_input_dic, 2)
+
+        # Compute student losses
+        s_loss_mood = self.loss_fn_classification(s_logits_mood, y_mood)
+        s_loss_va = self.loss_fn_regression(s_logits_va, y_va)
+
+        # Get the corresponding teacher model for the dataset
+        teacher_model = self.teacher_models.get(dataset_name)
+        teacher_model.to(self.device)
+
+        # Ensure teacher model exists
+        if teacher_model is None:
+            raise ValueError(f"No teacher model found for dataset: {dataset_name}")
+
+        with torch.no_grad():
+            # Forward pass through teacher model
+            t_logits = teacher_model(model_input_dic)
+
+        # Compute knowledge distillation losses
+        t_probs = torch.softmax(t_logits / self.kd_temperature, dim=-1)
+        if dataset_name == "jamendo":
+            s_probs_mood = torch.log_softmax(s_logits_mood / self.kd_temperature, dim=-1)
+            kd_loss = self.loss_kd(s_probs_mood, t_probs)
+        else:
+            s_probs_va = torch.log_softmax(s_logits_va / self.kd_temperature, dim=-1)
+            kd_loss = self.loss_kd(s_probs_va, t_probs)
+
+        return kd_loss, s_loss_mood, s_loss_va
+    
+    def handle_dataset(self, dataset_name, batch, losses, total_loss, stage):
+        dataset_batch = batch[dataset_name]
+
+        model_input_dic = {}
+        model_input_dic["x_mert"] = dataset_batch["x_mert"]
+        model_input_dic["x_chord"] = dataset_batch["x_chord"]
+        model_input_dic["x_chord_root"] = dataset_batch["x_chord_root"]
+        model_input_dic["x_chord_attr"] = dataset_batch["x_chord_attr"]
+        model_input_dic["x_key"] = dataset_batch["x_key"]
+
+        if "mt" in self.classifier:            
+            if dataset_name == "jamendo":
+                y_mood = dataset_batch["y_mood"]
+                y_va = dataset_batch["y_va"]
+                if self.kd:
+                    kd_loss, s_loss_mood, s_loss_va = self.compute_kd_loss(model_input_dic, y_mood, y_va, dataset_name)
+                    if stage == "train":
+                        losses['loss_mood'] = s_loss_mood
+
+                        total_loss += self.kd_weight * kd_loss + (1 - self.kd_weight) * s_loss_mood
+                    else:
+                        losses['loss_mood'] = s_loss_mood
+                        total_loss += s_loss_mood
+                else:
+                    s_loss_mood, s_loss_va = self.compute_mt_loss(model_input_dic, y_mood, y_va)
+                    if stage == "train":
+                        losses['loss_mood'] = s_loss_mood
+                        total_loss += s_loss_mood 
+                    else:
+                        losses['loss_mood'] = s_loss_mood
+                        total_loss += s_loss_mood
+            else:
+                y_mood = dataset_batch["y_mood"]
+                y_va = dataset_batch["y_va"]
+
+                if self.kd:
+                    kd_loss, s_loss_mood, s_loss_va = self.compute_kd_loss(model_input_dic, y_mood, y_va, dataset_name)
+                    if stage == "train":
+                        losses['loss_va'] = s_loss_va
+                        total_loss += self.kd_weight * kd_loss + (1 - self.kd_weight) * s_loss_va
+                    else:
+                        losses['loss_va'] = s_loss_va
+                        total_loss += s_loss_va
+                else:
+                    s_loss_mood, s_loss_va = self.compute_mt_loss(model_input_dic, y_mood, y_va)
+                    if stage == "train":
+                        losses['loss_va'] = s_loss_va
+                        total_loss += s_loss_va
+                    else:
+                        losses['loss_va'] = s_loss_va
+                        total_loss += s_loss_va
+        else:
+            if dataset_name == "jamendo":
+                y_mood = dataset_batch["y_mood"]
+                loss_classification = self.compute_classification_loss(model_input_dic, y_mood)
+                losses['loss_mood'] = loss_classification
+                total_loss += loss_classification
+            else:
+                y_va = dataset_batch["y_va"]
+                loss_regression = self.compute_regression_loss(model_input_dic, y_va)
+                losses['loss_va'] = loss_regression
+                total_loss += loss_regression
+
+        return total_loss
+
+    def training_step(self, batch, batch_idx):
+        total_loss = 0
+        losses = {}
+        datasets = ["jamendo",  "deam", "emomusic", "pmemo"]
+        
+        for dataset in datasets:
+            if dataset in batch and batch[dataset] is not None:
+                total_loss = self.handle_dataset(dataset, batch, losses, total_loss, "train")
+        
+        batch_size = batch[next(iter(batch))]["x_mert"].size(0)
+
+        self.log('train_loss_mood', losses.get('loss_mood', 0), on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)
+        self.log('train_loss_va', losses.get('loss_va', 0), on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)
+        self.log('train_loss', total_loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)
+
+        return total_loss
+    
+    def validation_step(self, batch, batch_idx):
+
+        total_loss = 0
+        losses = {}
+        datasets = ["jamendo", "deam", "emomusic", "pmemo"]
+
+        for dataset in datasets:
+            if dataset in batch and batch[dataset] is not None:
+                total_loss = self.handle_dataset(dataset, batch, losses, total_loss, "val")
+        
+        batch_size = batch[next(iter(batch))]["x_mert"].size(0)
+
+        self.log('val_loss_mood', losses.get('loss_mood', 0), on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)
+        self.log('val_loss_va', losses.get('loss_va', 0), on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)
+        self.log('val_loss', total_loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)
+        return total_loss
+
+    def test_step(self, batch, batch_idx):
+        total_loss = 0
+        losses = {}
+        datasets = ["jamendo",  "deam", "emomusic", "pmemo"]
+
+        for dataset in datasets:
+            if dataset in batch and batch[dataset] is not None:
+                dataset_batch = batch[dataset]
+
+                model_input_dic = {}
+                model_input_dic["x_mert"] = dataset_batch["x_mert"]
+                model_input_dic["x_chord"] = dataset_batch["x_chord"]
+                model_input_dic["x_chord_root"] = dataset_batch["x_chord_root"]
+                model_input_dic["x_chord_attr"] = dataset_batch["x_chord_attr"]
+                
+                model_input_dic["x_key"] = dataset_batch["x_key"]
+
+                if dataset == "jamendo":
+                    y_mood = dataset_batch["y_mood"]
+                    classification_logits = self(model_input_dic, 0)
+                    
+                    loss_classification = self.loss_fn_classification(classification_logits, y_mood)
+                    total_loss += loss_classification
+
+                    probs = torch.sigmoid(classification_logits)
+                    if not hasattr(self, 'jamendo_results'):
+                        self.jamendo_results = {'preds': [], 'gt': [], 'paths': []}
+                    
+                    self.jamendo_results['preds'].extend(probs.detach().cpu().numpy())
+                    self.jamendo_results['gt'].extend(y_mood.detach().cpu().numpy())
+                    self.jamendo_results['paths'].extend(dataset_batch["path"])
+
+                    losses['test_loss_mood'] = loss_classification
+
+                else:  # Handle regression for all other datasets
+                    if batch[dataset] is not None:
+                        y_va = dataset_batch["y_va"]
+                        regression_output = self(model_input_dic, 1)
+
+                        loss_regression = self.loss_fn_regression(regression_output, y_va)
+                        total_loss += loss_regression
+
+                        # Track results separately for each dataset
+                        if not hasattr(self, f'{dataset}_results'):
+                            setattr(self, f'{dataset}_results', {'preds': [], 'gt': [], 'paths': []})
+                        
+                        dataset_results = getattr(self, f'{dataset}_results')
+                        dataset_results['preds'].extend(regression_output.detach().cpu().numpy())
+                        dataset_results['gt'].extend(y_va.detach().cpu().numpy())
+                        dataset_results['paths'].extend(dataset_batch["path"])
+
+                        losses['test_loss_va'] = loss_regression
+
+        batch_size = batch[next(iter(batch))]["x_mert"].size(0)
+
+        # Log the classification and regression losses
+        self.log('test_loss_mood', losses.get('test_loss_mood', 0), on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)
+        self.log('test_loss_va', losses.get('test_loss_va', 0), on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)
+        self.log('test_loss', total_loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size)
+
+        return total_loss
+
+    def on_test_end(self):
+        output_dic = {}
+
+        # Jamendo classification metrics (AUC and PR AUC)
+        if hasattr(self, 'jamendo_results') and self.jamendo_results['preds']:
+            roc_auc, pr_auc = self.get_auc(self.jamendo_results['preds'], self.jamendo_results['gt'])
+
+            roc_auc = roc_auc.item()
+            pr_auc = pr_auc.item()
+
+            log.info('*** Display ROC_AUC_MACRO scores (Jamendo) ***')
+            log.info(f"ROC_AUC_MACRO: {round(roc_auc, 4)}")
+            log.info(f"PR_AUC_MACRO: {round(pr_auc, 4)}")
+
+            if self.output_file is not None:
+                with open(self.output_file, 'a') as f:
+                    f.write(f"ROC_AUC_MACRO (Jamendo): {round(roc_auc, 4)}\n")
+                    f.write(f"PR_AUC_MACRO (Jamendo): {round(pr_auc, 4)}\n")
+
+            output_dic["test_roc_auc_jamendo"] = round(roc_auc, 4)
+            output_dic["test_pr_auc_jamendo"] = round(pr_auc, 4)
+
+        # Metrics for each regression dataset (DMDD, DEAM, EmoMusic, PMEmo)
+        for dataset in ["deam", "emomusic", "pmemo"]:
+            dataset_results = getattr(self, f'{dataset}_results', None)
+
+            if dataset_results and dataset_results['preds']:
+                preds = torch.tensor(np.array(dataset_results['preds']))
+                gts = torch.tensor(np.array(dataset_results['gt']))
+
+                # Assuming valence is the first column and arousal is the second
+                preds_valence = preds[:, 0]
+                preds_arousal = preds[:, 1]
+                gts_valence = gts[:, 0]
+                gts_arousal = gts[:, 1]
+
+                rmse = torch.sqrt(tmf.mean_squared_error(preds, gts))
+                r2 = tmf.r2_score(preds, gts)
+
+                # Calculate metrics for valence
+                rmse_valence = torch.sqrt(tmf.mean_squared_error(preds_valence, gts_valence))
+                r2_valence = tmf.r2_score(preds_valence, gts_valence)
+
+                # Calculate metrics for arousal
+                rmse_arousal = torch.sqrt(tmf.mean_squared_error(preds_arousal, gts_arousal))
+                r2_arousal = tmf.r2_score(preds_arousal, gts_arousal)
+
+                log.info(f'*** Display RMSE and R² scores ({dataset.upper()}) ***')
+                log.info(f"RMSE: {round(rmse.item(), 4)}")
+                log.info(f"R²: {round(r2.item(), 4)}")
+                log.info(f"Valence - RMSE: {round(rmse_valence.item(), 4)}, R²: {round(r2_valence.item(), 4)}")
+                log.info(f"Arousal - RMSE: {round(rmse_arousal.item(), 4)}, R²: {round(r2_arousal.item(), 4)}")
+
+                if self.output_file is not None:
+                    with open(self.output_file, 'a') as f:
+                        f.write(f"RMSE ({dataset.upper()}): {round(rmse.item(), 4)}\n")
+                        f.write(f"R² ({dataset.upper()}): {round(r2.item(), 4)}\n")
+                        f.write(f"Valence - RMSE ({dataset.upper()}): {round(rmse_valence.item(), 4)}\n")
+                        f.write(f"Valence - R² ({dataset.upper()}): {round(r2_valence.item(), 4)}\n")
+                        f.write(f"Arousal - RMSE ({dataset.upper()}): {round(rmse_arousal.item(), 4)}\n")
+                        f.write(f"Arousal - R² ({dataset.upper()}): {round(r2_arousal.item(), 4)}\n")
+
+                output_dic[f"test_rmse_{dataset}"] = round(rmse.item(), 4)
+                output_dic[f"test_r2_{dataset}"] = round(r2.item(), 4)
+                output_dic[f"test_rmse_valence_{dataset}"] = round(rmse_valence.item(), 4)
+                output_dic[f"test_r2_valence_{dataset}"] = round(r2_valence.item(), 4)
+                output_dic[f"test_rmse_arousal_{dataset}"] = round(rmse_arousal.item(), 4)
+                output_dic[f"test_r2_arousal_{dataset}"] = round(r2_arousal.item(), 4)
+
+        # Clear results for each dataset
+        for dataset in ["jamendo",  "deam", "emomusic", "pmemo"]:
+            if hasattr(self, f'{dataset}_results'):
+                getattr(self, f'{dataset}_results')['preds'].clear()
+                getattr(self, f'{dataset}_results')['gt'].clear()
+                getattr(self, f'{dataset}_results')['paths'].clear()
+
+        return output_dic
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=self.lr)
+    
+    def get_auc(self, prd_array, gt_array):
+        prd_array = np.array(prd_array)
+        gt_array = np.array(gt_array)
+
+        prd_tensor = torch.tensor(prd_array)
+        gt_tensor = torch.tensor(gt_array)
+        
+        try:
+            roc_auc = tmf.auroc(prd_tensor, gt_tensor.int(), task='multilabel', num_labels = 56 , average='macro', num_classes=gt_tensor.size(1))
+            pr_auc = tmf.average_precision(prd_tensor, gt_tensor.int(), task='multilabel', num_labels = 56, average='macro', num_classes=gt_tensor.size(1))            
+        except ValueError as e:
+            print(f"Error computing metrics: {e}")
+            roc_auc = None
+            pr_auc = None
+        return roc_auc, pr_auc
+    
+