amt/src/model/ymt3.py

# Copyright 2024 The YourMT3 Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Please see the details in the LICENSE file.
"""ymt3.py"""
import os
from typing import Union, Optional, Tuple, Dict, List, Any
from collections import Counter

import torch
import torch.nn as nn
from torch.nn import CrossEntropyLoss
import torchaudio  # for debugging audio
import pytorch_lightning as pl
import numpy as np
import wandb
from einops import rearrange

from transformers import T5Config
from model.t5mod import T5EncoderYMT3, T5DecoderYMT3, MultiChannelT5Decoder
from model.t5mod_helper import task_cond_dec_generate
from model.perceiver_mod import PerceiverTFEncoder
from model.perceiver_helper import PerceiverTFConfig
from model.conformer_mod import ConformerYMT3Encoder
from model.conformer_helper import ConformerYMT3Config
from model.lm_head import LMHead
from model.pitchshift_layer import PitchShiftLayer
from model.spectrogram import get_spectrogram_layer_from_audio_cfg
from model.conv_block import PreEncoderBlockRes3B
from model.conv_block import PreEncoderBlockHFTT, PreEncoderBlockRes3BHFTT  # added for hFTT-like pre-encoder
from model.projection_layer import get_projection_layer, get_multi_channel_projection_layer
from model.optimizers import get_optimizer
from model.lr_scheduler import get_lr_scheduler

from utils.note_event_dataclasses import Note
from utils.note2event import mix_notes
from utils.event2note import merge_zipped_note_events_and_ties_to_notes, DECODING_ERR_TYPES
from utils.metrics import compute_track_metrics
from utils.metrics import AMTMetrics
# from utils.utils import write_model_output_as_npy
from utils.utils import write_model_output_as_midi, create_inverse_vocab, write_err_cnt_as_json
from utils.utils import Timer
from utils.task_manager import TaskManager

from config.config import audio_cfg as default_audio_cfg
from config.config import model_cfg as default_model_cfg
from config.config import shared_cfg as default_shared_cfg
from config.config import T5_BASE_CFG


class YourMT3(pl.LightningModule):
    """YourMT3:
    
    Lightning wrapper for multi-task music transcription Transformer.
    
    """

    def __init__(
            self,
            audio_cfg: Optional[Dict] = None,
            model_cfg: Optional[Dict] = None,
            shared_cfg: Optional[Dict] = None,
            pretrained: bool = False,
            optimizer_name: str = 'adamwscale',
            scheduler_name: str = 'cosine',
            base_lr: float = None,  # None: 'auto' for AdaFactor, 1e-3 for constant, 1e-2 for cosine
            max_steps: Optional[int] = None,
            weight_decay: float = 0.0,
            init_factor: Optional[Union[str, float]] = None,
            task_manager: TaskManager = TaskManager(),
            eval_subtask_key: Optional[str] = "default",
            eval_vocab: Optional[Dict] = None,
            eval_drum_vocab: Optional[Dict] = None,
            write_output_dir: Optional[str] = None,
            write_output_vocab: Optional[Dict] = None,
            onset_tolerance: float = 0.05,
            add_pitch_class_metric: Optional[List[str]] = None,
            add_melody_metric_to_singing: bool = True,
            test_optimal_octave_shift: bool = False,
            test_pitch_shift_layer: Optional[str] = None,
            **kwargs: Any) -> None:
        super().__init__()
        if pretrained is True:
            raise NotImplementedError("Pretrained model is not supported in this version.")
        self.test_pitch_shift_layer = test_pitch_shift_layer  # debug only

        # Config
        if model_cfg is None:
            model_cfg = default_model_cfg  # default config, not overwritten by args of trainer
        if audio_cfg is None:
            audio_cfg = default_audio_cfg  # default config, not overwritten by args of trainer
        if shared_cfg is None:
            shared_cfg = default_shared_cfg  # default config, not overwritten by args of trainer

        # Spec Layer (need to define here to infer max token length)
        self.spectrogram, spec_output_shape = get_spectrogram_layer_from_audio_cfg(
            audio_cfg)  # can be spec or melspec; output_shape is (T, F)
        model_cfg["feat_length"] = spec_output_shape[0]  # T of (T, F)

        # Task manger and Tokens
        self.task_manager = task_manager
        self.max_total_token_length = self.task_manager.max_total_token_length

        # Task Conditioning
        self.use_task_cond_encoder = bool(model_cfg["use_task_conditional_encoder"])
        self.use_task_cond_decoder = bool(model_cfg["use_task_conditional_decoder"])

        # Select Encoder type, Model-specific Config
        assert model_cfg["encoder_type"] in ["t5", "perceiver-tf", "conformer"]
        assert model_cfg["decoder_type"] in ["t5", "multi-t5"]
        self.encoder_type = model_cfg["encoder_type"]  # {"t5", "perceiver-tf", "conformer"}
        self.decoder_type = model_cfg["decoder_type"]  # {"t5", "multi-t5"}
        encoder_config = model_cfg["encoder"][self.encoder_type]  # mutable
        decoder_config = model_cfg["decoder"][self.decoder_type]  # mutable

        # Positional Encoding
        if isinstance(model_cfg["num_max_positions"], str) and model_cfg["num_max_positions"] == 'auto':
            encoder_config["num_max_positions"] = int(model_cfg["feat_length"] +
                                                      self.task_manager.max_task_token_length + 10)
            decoder_config["num_max_positions"] = int(self.max_total_token_length + 10)
        else:
            assert isinstance(model_cfg["num_max_positions"], int)
            encoder_config["num_max_positions"] = model_cfg["num_max_positions"]
            decoder_config["num_max_positions"] = model_cfg["num_max_positions"]

        # Select Pre-Encoder and Pre-Decoder type
        if model_cfg["pre_encoder_type"] == "default":
            model_cfg["pre_encoder_type"] = model_cfg["pre_encoder_type_default"].get(model_cfg["encoder_type"], None)
        elif model_cfg["pre_encoder_type"] in [None, "none", "None", "0"]:
            model_cfg["pre_encoder_type"] = None
        if model_cfg["pre_decoder_type"] == "default":
            model_cfg["pre_decoder_type"] = model_cfg["pre_decoder_type_default"].get(model_cfg["encoder_type"]).get(
                model_cfg["decoder_type"], None)
        elif model_cfg["pre_decoder_type"] in [None, "none", "None", "0"]:
            model_cfg["pre_decoder_type"] = None
        self.pre_encoder_type = model_cfg["pre_encoder_type"]
        self.pre_decoder_type = model_cfg["pre_decoder_type"]

        # Pre-encoder
        self.pre_encoder = nn.Sequential()
        if self.pre_encoder_type in ["conv", "conv1d_t", "conv1d_f"]:
            kernel_size = (3, 3)
            avp_kernel_size = (1, 2)
            if self.pre_encoder_type == "conv1d_t":
                kernel_size = (3, 1)
            elif self.pre_encoder_type == "conv1d_f":
                kernel_size = (1, 3)
            self.pre_encoder.append(
                PreEncoderBlockRes3B(1,
                                     model_cfg["conv_out_channels"],
                                     kernel_size=kernel_size,
                                     avp_kernerl_size=avp_kernel_size,
                                     activation="relu"))
            pre_enc_output_shape = (spec_output_shape[0], spec_output_shape[1] // 2**3, model_cfg["conv_out_channels"]
                                   )  # (T, F, C) excluding batch dim
        elif self.pre_encoder_type == "hftt":
            self.pre_encoder.append(PreEncoderBlockHFTT())
            pre_enc_output_shape = (spec_output_shape[0], spec_output_shape[1], 128)  # (T, F, C) excluding batch dim
        elif self.pre_encoder_type == "res3b_hftt":
            self.pre_encoder.append(PreEncoderBlockRes3BHFTT())
            pre_enc_output_shape = (spec_output_shape[0], spec_output_shape[1] // 2**3, 128)
        else:
            pre_enc_output_shape = spec_output_shape  # (T, F) excluding batch dim

        # Auto-infer `d_feat` and `d_model`, `vocab_size`, and `num_max_positions`
        if isinstance(model_cfg["d_feat"], str) and model_cfg["d_feat"] == 'auto':
            if self.encoder_type == "perceiver-tf" and encoder_config["attention_to_channel"] is True:
                model_cfg["d_feat"] = pre_enc_output_shape[-2]  # TODO: better readablity
            else:
                model_cfg["d_feat"] = pre_enc_output_shape[-1]  # C of (T, F, C) or F or (T, F)

        if self.encoder_type == "perceiver-tf" and isinstance(encoder_config["d_model"], str):
            if encoder_config["d_model"] == 'q':
                encoder_config["d_model"] = encoder_config["d_latent"]
            elif encoder_config["d_model"] == 'kv':
                encoder_config["d_model"] = model_cfg["d_feat"]
            else:
                raise ValueError(f"Unknown d_model: {encoder_config['d_model']}")

        # # required for PerceiverTF with attention_to_channel option
        # if self.encoder_type == "perceiver-tf":
        #     if encoder_config["attention_to_channel"] is True:
        #         encoder_config["kv_dim"] = model_cfg["d_feat"]  # TODO: better readablity
        #     else:
        #         encoder_config["kv_dim"] = model_cfg["conv_out_channels"]

        if isinstance(model_cfg["vocab_size"], str) and model_cfg["vocab_size"] == 'auto':
            model_cfg["vocab_size"] = task_manager.num_tokens

        if isinstance(model_cfg["num_max_positions"], str) and model_cfg["num_max_positions"] == 'auto':
            model_cfg["num_max_positions"] = int(
                max(model_cfg["feat_length"], model_cfg["event_length"]) + self.task_manager.max_task_token_length + 10)

        # Pre-decoder
        self.pre_decoder = nn.Sequential()
        if self.encoder_type == "perceiver-tf" and self.decoder_type == "t5":
            t, f, c = pre_enc_output_shape  # perceiver-tf: (110, 128, 128) for 2s
            encoder_output_shape = (t, encoder_config["num_latents"], encoder_config["d_latent"])  # (T, K, D_source)
            decoder_input_shape = (t, decoder_config["d_model"])  # (T, D_target)
            proj_layer = get_projection_layer(input_shape=encoder_output_shape,
                                              output_shape=decoder_input_shape,
                                              proj_type=self.pre_decoder_type)
            self.pre_encoder_output_shape = pre_enc_output_shape
            self.encoder_output_shape = encoder_output_shape
            self.decoder_input_shape = decoder_input_shape
            self.pre_decoder.append(proj_layer)
        elif self.encoder_type in ["t5", "conformer"] and self.decoder_type == "t5":
            pass
        elif self.encoder_type == "perceiver-tf" and self.decoder_type == "multi-t5":
            # NOTE: this is experiemental, only for multi-channel decoding with 13 classes
            assert encoder_config["num_latents"] % decoder_config["num_channels"] == 0
            encoder_output_shape = (encoder_config["num_latents"], encoder_config["d_model"])
            decoder_input_shape = (decoder_config["num_channels"], decoder_config["d_model"])
            proj_layer = get_multi_channel_projection_layer(input_shape=encoder_output_shape,
                                                            output_shape=decoder_input_shape,
                                                            proj_type=self.pre_decoder_type)
            self.pre_decoder.append(proj_layer)
        else:
            raise NotImplementedError(
                f"Encoder type {self.encoder_type} and decoder type {self.decoder_type} is not implemented yet.")

        # Positional Encoding, Vocab, etc.
        if self.encoder_type in ["t5", "conformer"]:
            encoder_config["num_max_positions"] = decoder_config["num_max_positions"] = model_cfg["num_max_positions"]
        else:  # perceiver-tf uses separate positional encoding
            encoder_config["num_max_positions"] = model_cfg["feat_length"]
            decoder_config["num_max_positions"] = model_cfg["num_max_positions"]
        encoder_config["vocab_size"] = decoder_config["vocab_size"] = model_cfg["vocab_size"]

        # Print and save updated configs
        self.audio_cfg = audio_cfg
        self.model_cfg = model_cfg
        self.shared_cfg = shared_cfg
        self.save_hyperparameters()
        if self.global_rank == 0:
            print(self.hparams)

        # Encoder and Decoder and LM-head
        self.encoder = None
        self.decoder = None
        self.lm_head = LMHead(decoder_config, 1.0, model_cfg["tie_word_embeddings"])
        self.embed_tokens = nn.Embedding(decoder_config["vocab_size"], decoder_config["d_model"])
        self.embed_tokens.weight.data.normal_(mean=0.0, std=1.0)
        self.shift_right_fn = None
        self.set_encoder_decoder()  # shift_right_fn is also set here

        # Model as ModuleDict
        # self.model = nn.ModuleDict({
        #     "pitchshift": self.pitchshift,   # no grad; created in setup() only for training,
        #                                        and called by training_step()
        #     "spectrogram": self.spectrogram,  # no grad
        #     "pre_encoder": self.pre_encoder,
        #     "encoder": self.encoder,
        #     "pre_decoder": self.pre_decoder,
        #     "decoder": self.decoder,
        #     "embed_tokens": self.embed_tokens,
        #     "lm_head": self.lm_head,
        # })

        # Tables (for logging)
        columns = ['Ep', 'Track ID', 'Pred Events', 'Actual Events', 'Pred Notes', 'Actual Notes']
        self.sample_table = wandb.Table(columns=columns)

        # Output MIDI
        if write_output_dir is not None:
            if write_output_vocab is None:
                from config.vocabulary import program_vocab_presets
                self.midi_output_vocab = program_vocab_presets["gm_ext_plus"]
            else:
                self.midi_output_vocab = write_output_vocab
            self.midi_output_inverse_vocab = create_inverse_vocab(self.midi_output_vocab)

    def set_encoder_decoder(self) -> None:
        """Set encoder, decoder, lm_head and emb_tokens from self.model_cfg"""

        # Generate and update T5Config
        t5_basename = self.model_cfg["t5_basename"]
        if t5_basename in T5_BASE_CFG.keys():
            # Load from pre-defined config in config.py
            t5_config = T5Config(**T5_BASE_CFG[t5_basename])
        else:
            # Load from HuggingFace hub
            t5_config = T5Config.from_pretrained(t5_basename)

        # Create encoder, decoder, lm_head and embed_tokens
        if self.encoder_type == "t5":
            self.encoder = T5EncoderYMT3(self.model_cfg["encoder"]["t5"], t5_config)
        elif self.encoder_type == "perceiver-tf":
            perceivertf_config = PerceiverTFConfig()
            perceivertf_config.update(self.model_cfg["encoder"]["perceiver-tf"])
            self.encoder = PerceiverTFEncoder(perceivertf_config)
        elif self.encoder_type == "conformer":
            conformer_config = ConformerYMT3Config()
            conformer_config.update(self.model_cfg["encoder"]["conformer"])
            self.encoder = ConformerYMT3Encoder(conformer_config)

        if self.decoder_type == "t5":
            self.decoder = T5DecoderYMT3(self.model_cfg["decoder"]["t5"], t5_config)
        elif self.decoder_type == "multi-t5":
            self.decoder = MultiChannelT5Decoder(self.model_cfg["decoder"]["multi-t5"], t5_config)

        # `shift_right` function for decoding
        self.shift_right_fn = self.decoder._shift_right

    def setup(self, stage: str) -> None:
        # Defining metrics
        if self.hparams.eval_vocab is None:
            extra_classes_per_dataset = [None]
        else:
            extra_classes_per_dataset = [
                list(v.keys()) if v is not None else None for v in self.hparams.eval_vocab
            ]  # e.g. [['Piano'], ['Guitar'], ['Piano'], ['Piano', 'Strings', 'Winds'], None]

        # For direct addition of extra metrics using full metric name
        extra_metrics = None
        if self.hparams.add_melody_metric_to_singing is True:
            extra_metrics = ["melody_rpa_Singing Voice", "melody_rca_Singing Voice", "melody_oa_Singing Voice"]

        # Add pitch class metric
        if self.hparams.add_pitch_class_metric is not None:
            for sublist in extra_classes_per_dataset:
                for name in self.hparams.add_pitch_class_metric:
                    if sublist is not None and name in sublist:
                        sublist += [name + "_pc"]

        extra_classes_unique = list(
            set(item for sublist in extra_classes_per_dataset if sublist is not None
                for item in sublist))  # e.g. ['Strings', 'Winds', 'Guitar', 'Piano']
        dm = self.trainer.datamodule

        # Train/Vaidation-only
        if stage == "fit":
            self.val_metrics_macro = AMTMetrics(prefix=f'validation/macro_', extra_classes=extra_classes_unique)
            self.val_metrics = nn.ModuleList()  # val_metric is a list of AMTMetrics objects
            for i in range(dm.num_val_dataloaders):
                self.val_metrics.append(
                    AMTMetrics(prefix=f'validation/({dm.get_val_dataset_name(i)})',
                               extra_classes=extra_classes_per_dataset[i],
                               error_types=DECODING_ERR_TYPES))

            # Add pitchshift layer
            if self.shared_cfg["AUGMENTATION"]["train_pitch_shift_range"] in [None, [0, 0]]:
                self.pitchshift = None
            else:
                # torchaudio pitchshifter requires a dummy input for initialization in DDP
                input_shape = (self.shared_cfg["BSZ"]["train_local"], 1, self.audio_cfg["input_frames"])
                self.pitchshift = PitchShiftLayer(
                    pshift_range=self.shared_cfg["AUGMENTATION"]["train_pitch_shift_range"],
                    expected_input_shape=input_shape,
                    device=self.device)

        # Test-only
        elif stage == "test":
            # self.test_metrics_macro = AMTMetrics(
            #     prefix=f'test/macro_', extra_classes=extra_classes_unique)
            self.test_metrics = nn.ModuleList()
            for i in range(dm.num_test_dataloaders):
                self.test_metrics.append(
                    AMTMetrics(prefix=f'test/({dm.get_test_dataset_name(i)})',
                               extra_classes=extra_classes_per_dataset[i],
                               extra_metrics=extra_metrics,
                               error_types=DECODING_ERR_TYPES))

            # Test pitch shift layer: debug only
            if self.test_pitch_shift_layer is not None:
                self.test_pitch_shift_semitone = int(self.test_pitch_shift_layer)
                self.pitchshift = PitchShiftLayer(
                    pshift_range=[self.test_pitch_shift_semitone, self.test_pitch_shift_semitone])

    def configure_optimizers(self) -> None:
        """Configure optimizer and scheduler"""
        optimizer, base_lr = get_optimizer(models_dict=self.named_parameters(),
                                           optimizer_name=self.hparams.optimizer_name,
                                           base_lr=self.hparams.base_lr,
                                           weight_decay=self.hparams.weight_decay)

        if self.hparams.optimizer_name.lower() == 'adafactor' and self.hparams.base_lr == None:
            print("Using AdaFactor with auto learning rate and no scheduler")
            return [optimizer]
        if self.hparams.optimizer_name.lower() == 'dadaptadam':
            print("Using dAdaptAdam with auto learning rate and no scheduler")
            return [optimizer]
        elif self.hparams.base_lr == None:
            print(f"Using default learning rate {base_lr} of {self.hparams.optimizer_name} as base learning rate.")
            self.hparams.base_lr = base_lr

        scheduler_cfg = self.shared_cfg["LR_SCHEDULE"]
        if self.hparams.max_steps != -1:
            # overwrite total_steps
            scheduler_cfg["total_steps"] = self.hparams.max_steps
        _lr_scheduler = get_lr_scheduler(optimizer,
                                         scheduler_name=self.hparams.scheduler_name,
                                         base_lr=base_lr,
                                         scheduler_cfg=scheduler_cfg)

        lr_scheduler = {'scheduler': _lr_scheduler, 'interval': 'step', 'frequency': 1}
        return [optimizer], [lr_scheduler]

    def forward(
            self,
            x: torch.FloatTensor,
            target_tokens: torch.LongTensor,
            # task_tokens: Optional[torch.LongTensor] = None,
            **kwargs) -> Dict:
        """ Forward pass with teacher-forcing for training and validation.
        Args:
            x: (B, 1, T) waveform with default T=32767
            target_tokens: (B, C, N) tokenized sequence of length N=event_length
            task_tokens: (B, C, task_len) tokenized task

        Returns:
            {
                'logits': (B, N + task_len + 1, vocab_size)
                'loss': (1, )
            }

        NOTE: all the commented shapes are in the case of original MT3 setup.
        """
        x = self.spectrogram(x)  # mel-/spectrogram: (b, 256, 512) or (B, T, F)
        x = self.pre_encoder(x)  # projection to d_model: (B, 256, 512)

        # TODO: task_cond_encoder would not work properly because of 3-d task_tokens
        # if task_tokens is not None and task_tokens.numel() > 0 and self.use_task_cond_encoder is True:
        #     # append task embedding to encoder input
        #     task_embed = self.embed_tokens(task_tokens)  # (B, task_len, 512)
        #     x = torch.cat([task_embed, x], dim=1)  # (B, task_len + 256, 512)
        enc_hs = self.encoder(inputs_embeds=x)["last_hidden_state"]  # (B, T', D)
        enc_hs = self.pre_decoder(enc_hs)  # (B, T', D) or (B, K, T, D)

        # if task_tokens is not None and task_tokens.numel() > 0 and self.use_task_cond_decoder is True:
        #     # append task token to decoder input and output label
        #     labels = torch.cat([task_tokens, target_tokens], dim=2)  # (B, C, task_len + N)
        # else:
        #     labels = target_tokens  # (B, C, N)
        labels = target_tokens  # (B, C, N)
        if labels.shape[1] == 1:  # for single-channel decoders, e.g. t5.
            labels = labels.squeeze(1)  # (B, N)

        dec_input_ids = self.shift_right_fn(labels)  # t5:(B, N), multi-t5:(B, C, N)
        dec_inputs_embeds = self.embed_tokens(dec_input_ids)  # t5:(B, N, D), multi-t5:(B, C, N, D)
        dec_hs, _ = self.decoder(inputs_embeds=dec_inputs_embeds, encoder_hidden_states=enc_hs, return_dict=False)

        if self.model_cfg["tie_word_embeddings"] is True:
            dec_hs = dec_hs * (self.model_cfg["decoder"][self.decoder_type]["d_model"]**-0.5)

        logits = self.lm_head(dec_hs)

        loss = None
        labels = labels.masked_fill(labels == 0, value=-100)  # ignore pad tokens for loss
        loss_fct = CrossEntropyLoss(ignore_index=-100)
        loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
        return {"logits": logits, "loss": loss}

    def inference(self,
                  x: torch.FloatTensor,
                  task_tokens: Optional[torch.LongTensor] = None,
                  max_token_length: Optional[int] = None,
                  **kwargs: Any) -> torch.Tensor:
        """ Inference from audio batch by cached autoregressive decoding.
        Args:
            x: (b, 1, t) waveform with t=32767
            task_token: (b, c, task_len) tokenized task. If None, will not append task embeddings (from task_tokens) to input.
            max_length: Maximum length of generated sequence. If None, self.max_total_token_length.
            **kwargs: https://huggingface.co/docs/transformers/v4.27.2/en/main_classes/text_generation#transformers.GenerationMixin.generate
            
        Returns:
            res_tokens: (b, n) resulting tokenized sequence of variable length < max_length
        """
        if self.test_pitch_shift_layer is not None:
            x_ps = self.pitchshift(x, self.test_pitch_shift_semitone)
            x = x_ps

        # From spectrogram to pre-decoder is the same pipeline as in forward()
        x = self.spectrogram(x)  # mel-/spectrogram: (b, 256, 512) or (B, T, F)
        x = self.pre_encoder(x)  # projection to d_model: (B, 256, 512)
        if task_tokens is not None and task_tokens.numel() > 0 and self.use_task_cond_encoder is True:
            # append task embedding to encoder input
            task_embed = self.embed_tokens(task_tokens)  # (B, task_len, 512)
            x = torch.cat([task_embed, x], dim=1)  # (B, task_len + 256, 512)
        enc_hs = self.encoder(inputs_embeds=x)["last_hidden_state"]  # (B, task_len + 256, 512)
        enc_hs = self.pre_decoder(enc_hs)  # (B, task_len + 256, 512)

        # Cached-autoregressive decoding with task token (can be None) as prefix
        if max_token_length is None:
            max_token_length = self.max_total_token_length

        pred_ids = task_cond_dec_generate(decoder=self.decoder,
                                          decoder_type=self.decoder_type,
                                          embed_tokens=self.embed_tokens,
                                          lm_head=self.lm_head,
                                          encoder_hidden_states=enc_hs,
                                          shift_right_fn=self.shift_right_fn,
                                          prefix_ids=task_tokens,
                                          max_length=max_token_length)  # (B, task_len + N) or (B, C, task_len + N)
        if pred_ids.dim() == 2:
            pred_ids = pred_ids.unsqueeze(1)  # (B, 1, task_len + N)

        if self.test_pitch_shift_layer is None:
            return pred_ids
        else:
            return pred_ids, x_ps

    def inference_file(
        self,
        bsz: int,
        audio_segments: torch.FloatTensor,  # (n_items, 1, segment_len): from a single file
        note_token_array: Optional[torch.LongTensor] = None,
        task_token_array: Optional[torch.LongTensor] = None,
        # subtask_key: Optional[str] = "default"
    ) -> Tuple[List[np.ndarray], Optional[torch.Tensor]]:
        """ Inference from audio batch by autoregressive decoding:
        Args:
            bsz: batch size
            audio_segments: (n_items, 1, segment_len): segmented audio from a single file
            note_token_array: (n_items, max_token_len): Optional. If token_array is None, will not return loss.
            subtask_key: (str): If None, not using subtask prefix. By default, using "default" defined in task manager.
        """
        # if subtask_key is not None:
        #     _subtask_token = torch.LongTensor(
        #         self.task_manager.get_eval_subtask_prefix_dict()[subtask_key]).to(self.device)

        n_items = audio_segments.shape[0]
        loss = 0.
        pred_token_array_file = []  # each element is (B, C, L) np.ndarray
        x_ps_concat = []

        for i in range(0, n_items, bsz):
            if i + bsz > n_items:  # last batch can be smaller
                x = audio_segments[i:n_items].to(self.device)
                # if subtask_key is not None:
                #     b = n_items - i  # bsz for the last batch
                #     task_tokens = _subtask_token.expand((b, -1))  # (b, task_len)
                if note_token_array is not None:
                    target_tokens = note_token_array[i:n_items].to(self.device)
                if task_token_array is not None and task_token_array.numel() > 0:
                    task_tokens = task_token_array[i:n_items].to(self.device)
                else:
                    task_tokens = None
            else:
                x = audio_segments[i:i + bsz].to(self.device)  # (bsz, 1, segment_len)
                # if subtask_key is not None:
                #     task_tokens = _subtask_token.expand((bsz, -1))  # (bsz, task_len)
                if note_token_array is not None:
                    target_tokens = note_token_array[i:i + bsz].to(self.device)  # (bsz, token_len)
                if task_token_array is not None and task_token_array.numel() > 0:
                    task_tokens = task_token_array[i:i + bsz].to(self.device)
                else:
                    task_tokens = None

            # token prediction (fast-autoregressive decoding)
            # if subtask_key is not None:
            #     preds = self.inference(x, task_tokens).detach().cpu().numpy()
            # else:
            #     preds = self.inference(x).detach().cpu().numpy()

            if self.test_pitch_shift_layer is not None:  # debug only
                preds, x_ps = self.inference(x, task_tokens)
                preds = preds.detach().cpu().numpy()
                x_ps_concat.append(x_ps.detach().cpu())
            else:
                preds = self.inference(x, task_tokens).detach().cpu().numpy()
            if len(preds) != len(x):
                raise ValueError(f'preds: {len(preds)}, x: {len(x)}')
            pred_token_array_file.append(preds)

            # validation loss (by teacher forcing)
            if note_token_array is not None:
                loss_weight = x.shape[0] / n_items
                loss += self(x, target_tokens)['loss'] * loss_weight
                # loss += self(x, target_tokens, task_tokens)['loss'] * loss_weight
            else:
                loss = None

        if self.test_pitch_shift_layer is not None:  # debug only
            if self.hparams.write_output_dir is not None:
                x_ps_concat = torch.cat(x_ps_concat, dim=0)
                return pred_token_array_file, loss, x_ps_concat.flatten().unsqueeze(0)
        else:
            return pred_token_array_file, loss

    def training_step(self, batch, batch_idx) -> torch.Tensor:
        # batch: {
        # 'dataset1': [Tuple[audio_segments(b, 1, t), tokens(b, max_token_len), ...]]
        # 'dataset2': [Tuple[audio_segments(b, 1, t), tokens(b, max_token_len), ...]]
        # 'dataset3': ...
        # }
        audio_segments, note_tokens, pshift_steps = [torch.cat(t, dim=0) for t in zip(*batch.values())]

        if self.pitchshift is not None:
            # Pitch shift
            n_groups = len(batch)
            audio_segments = torch.chunk(audio_segments, n_groups, dim=0)
            pshift_steps = torch.chunk(pshift_steps, n_groups, dim=0)
            for p in pshift_steps:
                assert p.eq(p[0]).all().item()

            audio_segments = torch.cat([self.pitchshift(a, p[0].item()) for a, p in zip(audio_segments, pshift_steps)],
                                       dim=0)

        loss = self(audio_segments, note_tokens)['loss']
        self.log('train_loss',
                 loss,
                 on_step=True,
                 on_epoch=True,
                 prog_bar=True,
                 batch_size=note_tokens.shape[0],
                 sync_dist=True)
        # print('lr', self.trainer.optimizers[0].param_groups[0]['lr'])
        return loss

    def validation_step(self, batch, batch_idx, dataloader_idx=0) -> Dict:
        # File-wise validation
        if self.task_manager.num_decoding_channels == 1:
            bsz = self.shared_cfg["BSZ"]["validation"]
        else:
            bsz = self.shared_cfg["BSZ"]["validation"] // self.task_manager.num_decoding_channels * 3
        # audio_segments, notes_dict, note_token_array, task_token_array = batch
        audio_segments, notes_dict, note_token_array = batch
        task_token_array = None

        # Loop through the tensor in chunks of bsz (=subbsz actually)
        n_items = audio_segments.shape[0]
        start_secs_file = [32767 * i / 16000 for i in range(n_items)]
        with Timer() as t:
            pred_token_array_file, loss = self.inference_file(bsz, audio_segments, note_token_array, task_token_array)
            """
            notes_dict: # Ground truth notes
                {
                    'mtrack_id': int,
                    'program': List[int],
                    'is_drum': bool,
                    'duration_sec': float,
                    'notes': List[Note],
                }
            """
            # Process a list of channel-wise token arrays for a file
            num_channels = self.task_manager.num_decoding_channels
            pred_notes_in_file = []
            n_err_cnt = Counter()
            for ch in range(num_channels):
                pred_token_array_ch = [arr[:, ch, :] for arr in pred_token_array_file]  # (B, L)
                zipped_note_events_and_tie, list_events, ne_err_cnt = self.task_manager.detokenize_list_batches(
                    pred_token_array_ch, start_secs_file, return_events=True)
                pred_notes_ch, n_err_cnt_ch = merge_zipped_note_events_and_ties_to_notes(zipped_note_events_and_tie)
                pred_notes_in_file.append(pred_notes_ch)
                n_err_cnt += n_err_cnt_ch
            pred_notes = mix_notes(pred_notes_in_file)  # This is the mixed notes from all channels

            if self.hparams.write_output_dir is not None:
                track_info = [notes_dict[k] for k in notes_dict.keys() if k.endswith("_id")][0]
                dataset_info = [k for k in notes_dict.keys() if k.endswith('_id')][0][:-3]
                # write_model_output_as_npy(zipped_note_events_and_tie, self.hparams.write_output_dir,
                #                           track_info)
                write_model_output_as_midi(pred_notes,
                                           self.hparams.write_output_dir,
                                           track_info,
                                           self.midi_output_inverse_vocab,
                                           output_dir_suffix=str(dataset_info) + '_' +
                                           str(self.hparams.eval_subtask_key))
            # generate sample text to display in log table
            # pred_events_text = [str([list_events[0][:200]])]
            # pred_notes_text = [str([pred_notes[:200]])]

            # this is local GPU metric per file, not global metric in DDP
            drum_metric, non_drum_metric, instr_metric = compute_track_metrics(
                pred_notes,
                notes_dict['notes'],
                eval_vocab=self.hparams.eval_vocab[dataloader_idx],
                eval_drum_vocab=self.hparams.eval_drum_vocab,
                onset_tolerance=self.hparams.onset_tolerance,
                add_pitch_class_metric=self.hparams.add_pitch_class_metric)
            self.val_metrics[dataloader_idx].bulk_update(drum_metric)
            self.val_metrics[dataloader_idx].bulk_update(non_drum_metric)
            self.val_metrics[dataloader_idx].bulk_update(instr_metric)
            self.val_metrics_macro.bulk_update(drum_metric)
            self.val_metrics_macro.bulk_update(non_drum_metric)
            self.val_metrics_macro.bulk_update(instr_metric)

        # Log sample table: predicted notes and ground truth notes
        # if batch_idx in (0, 1) and self.global_rank == 0:
        #     actual_notes_text = [str([notes_dict['notes'][:200]])]
        #     actual_tokens = token_array[0, :200].detach().cpu().numpy().tolist()
        #     actual_events_text = [str(self.tokenizer._decode(actual_tokens))]
        #     track_info = [notes_dict[k] for k in notes_dict.keys() if k.endswith("_id")]
        #     self.sample_table.add_data(self.current_epoch, track_info, pred_events_text,
        #                                actual_events_text, pred_notes_text, actual_notes_text)
        #     self.logger.log_table('Samples', self.sample_table.columns, self.sample_table.data)

        decoding_time_sec = t.elapsed_time()
        self.log('val_loss', loss, prog_bar=True, batch_size=n_items, sync_dist=True)
        # self.val_metrics[dataloader_idx].bulk_update_errors({'decoding_time': decoding_time_sec})

    def on_validation_epoch_end(self) -> None:
        for val_metrics in self.val_metrics:
            self.log_dict(val_metrics.bulk_compute(), sync_dist=True)
            val_metrics.bulk_reset()
        self.log_dict(self.val_metrics_macro.bulk_compute(), sync_dist=True)
        self.val_metrics_macro.bulk_reset()

    def test_step(self, batch, batch_idx, dataloader_idx=0) -> Dict:
        # File-wise evaluation
        if self.task_manager.num_decoding_channels == 1:
            bsz = self.shared_cfg["BSZ"]["validation"]
        else:
            bsz = self.shared_cfg["BSZ"]["validation"] // self.task_manager.num_decoding_channels * 3
        # audio_segments, notes_dict, note_token_array, task_token_array = batch
        audio_segments, notes_dict, note_token_array = batch
        task_token_array = None

        # Test pitch shift layer: debug only
        if self.test_pitch_shift_layer is not None and self.test_pitch_shift_semitone != 0:
            for n in notes_dict['notes']:
                if n.is_drum == False:
                    n.pitch = n.pitch + self.test_pitch_shift_semitone

        # Loop through the tensor in chunks of bsz (=subbsz actually)
        n_items = audio_segments.shape[0]
        start_secs_file = [32767 * i / 16000 for i in range(n_items)]

        if self.test_pitch_shift_layer is not None and self.hparams.write_output_dir is not None:
            pred_token_array_file, loss, x_ps = self.inference_file(bsz, audio_segments, None, None)
        else:
            pred_token_array_file, loss = self.inference_file(bsz, audio_segments, None, None)
        if len(pred_token_array_file) > 0:

            # Process a list of channel-wise token arrays for a file
            num_channels = self.task_manager.num_decoding_channels
            pred_notes_in_file = []
            n_err_cnt = Counter()
            for ch in range(num_channels):
                pred_token_array_ch = [arr[:, ch, :] for arr in pred_token_array_file]  # (B, L)
                zipped_note_events_and_tie, list_events, ne_err_cnt = self.task_manager.detokenize_list_batches(
                    pred_token_array_ch, start_secs_file, return_events=True)
                pred_notes_ch, n_err_cnt_ch = merge_zipped_note_events_and_ties_to_notes(zipped_note_events_and_tie)
                pred_notes_in_file.append(pred_notes_ch)
                n_err_cnt += n_err_cnt_ch
            pred_notes = mix_notes(pred_notes_in_file)  # This is the mixed notes from all channels

            if self.test_pitch_shift_layer is not None and self.hparams.write_output_dir is not None:
                # debug only
                wav_output_dir = os.path.join(self.hparams.write_output_dir, f"model_output_{dataset_info}")
                os.makedirs(wav_output_dir, exist_ok=True)
                wav_output_file = os.path.join(wav_output_dir, f"{track_info}_ps_{self.test_pitch_shift_semitone}.wav")
                torchaudio.save(wav_output_file, x_ps.squeeze(1), 16000, bits_per_sample=16)

            drum_metric, non_drum_metric, instr_metric = compute_track_metrics(
                pred_notes,
                notes_dict['notes'],
                eval_vocab=self.hparams.eval_vocab[dataloader_idx],
                eval_drum_vocab=self.hparams.eval_drum_vocab,
                onset_tolerance=self.hparams.onset_tolerance,
                add_pitch_class_metric=self.hparams.add_pitch_class_metric,
                add_melody_metric=['Singing Voice'] if self.hparams.add_melody_metric_to_singing else None,
                add_frame_metric=True,
                add_micro_metric=True,
                add_multi_f_metric=True)

            if self.hparams.write_output_dir is not None and self.global_rank == 0:
                # write model output to file
                track_info = [notes_dict[k] for k in notes_dict.keys() if k.endswith("_id")][0]
                dataset_info = [k for k in notes_dict.keys() if k.endswith('_id')][0][:-3]
                f_score = f"OnF{non_drum_metric['onset_f']:.2f}_MulF{instr_metric['multi_f']:.2f}"
                write_model_output_as_midi(pred_notes,
                                           self.hparams.write_output_dir,
                                           track_info,
                                           self.midi_output_inverse_vocab,
                                           output_dir_suffix=str(dataset_info) + '_' +
                                           str(self.hparams.eval_subtask_key) + '_' + f_score)
                write_err_cnt_as_json(track_info, self.hparams.write_output_dir,
                                      str(dataset_info) + '_' + str(self.hparams.eval_subtask_key) + '_' + f_score,
                                      n_err_cnt, ne_err_cnt)

            # Test with optimal octave shift
            if self.hparams.test_optimal_octave_shift:
                track_info = [notes_dict[k] for k in notes_dict.keys() if k.endswith("_id")][0]
                dataset_info = [k for k in notes_dict.keys() if k.endswith('_id')][0][:-3]
                score = [instr_metric['onset_f_Bass']]
                ref_notes_plus = []
                ref_notes_minus = []
                for note in notes_dict['notes']:
                    if note.is_drum == True:
                        ref_notes_plus.append(note)
                        ref_notes_minus.append(note)
                    else:
                        ref_notes_plus.append(
                            Note(is_drum=note.is_drum,
                                 program=note.program,
                                 onset=note.onset,
                                 offset=note.offset,
                                 pitch=note.pitch + 12,
                                 velocity=note.velocity))
                        ref_notes_minus.append(
                            Note(is_drum=note.is_drum,
                                 program=note.program,
                                 onset=note.onset,
                                 offset=note.offset,
                                 pitch=note.pitch - 12,
                                 velocity=note.velocity))

                drum_metric_plus, non_drum_metric_plus, instr_metric_plus = compute_track_metrics(
                    pred_notes,
                    ref_notes_plus,
                    eval_vocab=self.hparams.eval_vocab[dataloader_idx],
                    eval_drum_vocab=self.hparams.eval_drum_vocab,
                    onset_tolerance=self.hparams.onset_tolerance,
                    add_pitch_class_metric=self.hparams.add_pitch_class_metric)
                drum_metric_minus, non_drum_metric_minus, instr_metric_minus = compute_track_metrics(
                    ref_notes_minus,
                    notes_dict['notes'],
                    eval_vocab=self.hparams.eval_vocab[dataloader_idx],
                    eval_drum_vocab=self.hparams.eval_drum_vocab,
                    onset_tolerance=self.hparams.onset_tolerance,
                    add_pitch_class_metric=self.hparams.add_pitch_class_metric)

                score.append(instr_metric_plus['onset_f_Bass'])
                score.append(instr_metric_minus['onset_f_Bass'])
                max_index = score.index(max(score))
                if max_index == 0:
                    print(f"ZERO: {track_info}, z/p/m: {score[0]:.2f}/{score[1]:.2f}/{score[2]:.2f}")
                elif max_index == 1:
                    # plus
                    instr_metric['onset_f_Bass'] = instr_metric_plus['onset_f_Bass']
                    print(f"PLUS: {track_info}, z/p/m: {score[0]:.2f}/{score[1]:.2f}/{score[2]:.2f}")
                    write_model_output_as_midi(ref_notes_plus,
                                               self.hparams.write_output_dir,
                                               track_info + '_ref_octave_plus',
                                               self.midi_output_inverse_vocab,
                                               output_dir_suffix=str(dataset_info) + '_' +
                                               str(self.hparams.eval_subtask_key))
                else:
                    # minus
                    instr_metric['onset_f_Bass'] = instr_metric_minus['onset_f_Bass']
                    print(f"MINUS: {track_info}, z/p/m: {score[0]:.2f}/{score[1]:.2f}/{score[2]:.2f}")
                    write_model_output_as_midi(ref_notes_minus,
                                               self.hparams.write_output_dir,
                                               track_info + '_ref_octave_minus',
                                               self.midi_output_,
                                               output_dir_suffix=str(dataset_info) + '_' +
                                               str(self.hparams.eval_subtask_key))

            self.test_metrics[dataloader_idx].bulk_update(drum_metric)
            self.test_metrics[dataloader_idx].bulk_update(non_drum_metric)
            self.test_metrics[dataloader_idx].bulk_update(instr_metric)
            # self.test_metrics_macro.bulk_update(drum_metric)
            # self.test_metrics_macro.bulk_update(non_drum_metric)
            # self.test_metrics_macro.bulk_update(instr_metric)

    def on_test_epoch_end(self) -> None:
        # all_gather is done seeminglesly by torchmetrics
        for test_metrics in self.test_metrics:
            self.log_dict(test_metrics.bulk_compute(), sync_dist=True)
            test_metrics.bulk_reset()
        # self.log_dict(self.test_metrics_macro.bulk_compute(), sync_dist=True)
        # self.test_metrics_macro.bulk_reset()


def test_case_forward_mt3():
    import torch
    from config.config import audio_cfg, model_cfg, shared_cfg
    from model.ymt3 import YourMT3
    model = YourMT3()
    model.eval()
    x = torch.randn(2, 1, 32767)
    labels = torch.randint(0, 596, (2, 1, 1024), requires_grad=False)  # (B, C=1, T)
    task_tokens = torch.LongTensor([])
    output = model.forward(x, labels, task_tokens)
    logits, loss = output['logits'], output['loss']
    assert logits.shape == (2, 1024, 596)  # (B, N, vocab_size)


def test_case_inference_mt3():
    import torch
    from config.config import audio_cfg, model_cfg, shared_cfg
    from model.ymt3 import YourMT3
    model_cfg["num_max_positions"] = 1024 + 3 + 1
    model = YourMT3(model_cfg=model_cfg)
    model.eval()
    x = torch.randn(2, 1, 32767)
    task_tokens = torch.randint(0, 596, (2, 3), requires_grad=False)
    pred_ids = model.inference(x, task_tokens, max_token_length=10)  # (2, 3, 9) (B, C, L-task_len)
    # TODO: need to check the length of pred_ids when task_tokens is not None


def test_case_forward_enc_perceiver_tf_dec_t5():
    import torch
    from model.ymt3 import YourMT3
    from config.config import audio_cfg, model_cfg, shared_cfg
    model_cfg["encoder_type"] = "perceiver-tf"
    audio_cfg["codec"] = "spec"
    audio_cfg["hop_length"] = 300

    model = YourMT3(audio_cfg=audio_cfg, model_cfg=model_cfg)
    model.eval()

    x = torch.randn(2, 1, 32767)
    labels = torch.randint(0, 596, (2, 1, 1024), requires_grad=False)

    # forward
    output = model.forward(x, labels)
    logits, loss = output['logits'], output['loss']  # logits: (2, 1024, 596) (B, N, vocab_size)

    # inference
    pred_ids = model.inference(x, None, max_token_length=3)  # (2, 1, 3) (B, C, L)


def test_case_forward_enc_conformer_dec_t5():
    import torch
    from model.ymt3 import YourMT3
    from config.config import audio_cfg, model_cfg, shared_cfg
    model_cfg["encoder_type"] = "conformer"
    audio_cfg["codec"] = "melspec"
    audio_cfg["hop_length"] = 128
    model = YourMT3(audio_cfg=audio_cfg, model_cfg=model_cfg)
    model.eval()

    x = torch.randn(2, 1, 32767)
    labels = torch.randint(0, 596, (2, 1024), requires_grad=False)

    # forward
    output = model.forward(x, labels)
    logits, loss = output['logits'], output['loss']  # logits: (2, 1024, 596) (B, N, vocab_size)

    # inference
    pred_ids = model.inference(x, None, 20)  # (2, 1, 20) (B, C, L)


def test_case_enc_perceiver_tf_dec_multi_t5():
    import torch
    from model.ymt3 import YourMT3
    from config.config import audio_cfg, model_cfg, shared_cfg
    model_cfg["encoder_type"] = "perceiver-tf"
    model_cfg["decoder_type"] = "multi-t5"
    model_cfg["encoder"]["perceiver-tf"]["attention_to_channel"] = True
    model_cfg["encoder"]["perceiver-tf"]["num_latents"] = 26
    audio_cfg["codec"] = "spec"
    audio_cfg["hop_length"] = 300
    model = YourMT3(audio_cfg=audio_cfg, model_cfg=model_cfg)
    model.eval()

    x = torch.randn(2, 1, 32767)
    labels = torch.randint(0, 596, (2, 13, 200), requires_grad=False)  # (B, C, T)

    # x = model.spectrogram(x)
    # x = model.pre_encoder(x)  # (2, 110, 128, 128) (B, T, C, D)
    # enc_hs = model.encoder(inputs_embeds=x)["last_hidden_state"]  # (2, 110, 128, 128) (B, T, C, D)
    # enc_hs = model.pre_decoder(enc_hs)  # (2, 13, 110, 512) (B, C, T, D)

    # dec_input_ids = model.shift_right_fn(labels)  # (2, 13, 200) (B, C, T)
    # dec_inputs_embeds = model.embed_tokens(dec_input_ids)  # (2, 13, 200, 512) (B, C, T, D)
    # dec_hs, _ = model.decoder(
    #     inputs_embeds=dec_inputs_embeds, encoder_hidden_states=enc_hs, return_dict=False)
    # logits = model.lm_head(dec_hs)  # (2, 13, 200, 596) (B, C, T, vocab_size)

    # forward
    x = torch.randn(2, 1, 32767)
    labels = torch.randint(0, 596, (2, 13, 200), requires_grad=False)  # (B, C, T)
    output = model.forward(x, labels)
    logits, loss = output['logits'], output['loss']  # (2, 13, 200, 596) (B, C, T, vocab_size)

    # inference
    model.max_total_token_length = 123  # to save time..
    pred_ids = model.inference(x, None)  # (2, 13, 123) (B, C, L)