Merge pull request #5 from Nospoko/MIDI-126/finetuning

WojciechMat · web-flow · commit 6abb305b2e0a · 2023-12-26T17:29:32.000+01:00
Midi 126/finetuning
diff --git a/configs/T5denoise-dstart.yaml b/configs/T5denoise-dstart.yaml
@@ -0,0 +1,46 @@
+train:
+  num_epochs: 5
+  accum_iter: 5
+  batch_size: 2
+  base_lr: 3e-5
+  warmup: 4000
+  finetune: False
+
+model_name: T5
+dataset_name: 'roszcz/maestro-v1-sustain'
+target: denoise
+seed: 26
+
+overfit: False
+
+tokens_per_note: single
+time_quantization_method: dstart
+masking_probability: 0.2
+mask: tokens
+
+encoder: velocity
+time_bins: 100
+
+dataset:
+  sequence_len: 128
+  sequence_step: 42
+
+  quantization:
+    dstart: 5
+    duration: 5
+    velocity: 3
+
+device: "cuda:0"
+
+log: True
+log_frequency: 10
+run_name: midi-T5-${now:%Y-%m-%d-%H-%M}
+project: "midi-hf-transformer"
+
+pre_defined_model: null
+model:
+  d_model: 512
+  d_kv: 64
+  d_ff: 2048
+  num_layers: 6
+  num_heads: 8
diff --git a/configs/T5denoise.yaml b/configs/T5denoise.yaml
@@ -4,6 +4,7 @@ train:
   batch_size: 8
   base_lr: 3e-5
   warmup: 4000
+  finetune: False
 
 model_name: T5
 dataset_name: 'roszcz/maestro-v1-sustain'
@@ -17,14 +18,17 @@ time_quantization_method: start
 masking_probability: 0.15
 mask: notes
 
+encoder: velocity
+time_bins: 100
+
 dataset:
   sequence_duration: 5
   sequence_step: 2
 
   quantization:
-    start: 20
+    start: 50
     duration: 5
-    velocity: 5
+    velocity: 3
 
 device: "cuda:0"
 
@@ -33,6 +37,8 @@ log_frequency: 10
 run_name: midi-T5-${now:%Y-%m-%d-%H-%M}
 project: "midi-hf-transformer"
 
+pre_defined_model: null
+
 model:
   d_model: 512
   d_kv: 64
diff --git a/configs/T5start.yaml b/configs/T5start.yaml
@@ -31,6 +31,8 @@ log_frequency: 10
 run_name: midi-T5-${now:%Y-%m-%d-%H-%M}
 project: "midi-hf-transformer"
 
+pre_defined_model: null
+
 model:
   d_model: 512
   d_kv: 64
diff --git a/configs/T5velocity-dstart.yaml b/configs/T5velocity-dstart.yaml
@@ -2,13 +2,17 @@ train:
   num_epochs: 5
   accum_iter: 10
   batch_size: 8
-  base_lr: 1e-5
+  base_lr: 3e-5
+  finetune: True
   warmup: 4000
 
+
+pretrained_checkpoint: midi-T5-2023-11-15-17-18.pt
 model_name: T5
 dataset_name: 'roszcz/maestro-v1-sustain'
 target: velocity
 seed: 26
+time_bins: 100
 
 overfit: False
 
@@ -30,6 +34,8 @@ log_frequency: 10
 run_name: midi-T5-${now:%Y-%m-%d-%H-%M}
 project: "midi-hf-transformer"
 
+pre_defined_model: null
+
 model:
   d_model: 512
   d_kv: 64
diff --git a/configs/T5velocity.yaml b/configs/T5velocity.yaml
@@ -4,7 +4,9 @@ train:
   batch_size: 8
   base_lr: 3e-4
   warmup: 4000
+  finetune: True
 
+pretrained_checkpoint: midi-T5-2023-11-11-10-29.pt
 model_name: T5
 dataset_name: 'roszcz/maestro-v1-sustain'
 target: velocity
@@ -30,6 +32,8 @@ log_frequency: 10
 run_name: midi-T5-${now:%Y-%m-%d-%H-%M}
 project: "midi-hf-transformer"
 
+pre_defined_model: null
+
 model:
   d_model: 512
   d_kv: 64
diff --git a/configs/architectures/large.yaml b/configs/architectures/large.yaml
@@ -0,0 +1,5 @@
+d_model: 512
+d_kv: 64
+d_ff: 2048
+num_layers: 6
+num_heads: 8
diff --git a/configs/architectures/mid.yaml b/configs/architectures/mid.yaml
@@ -0,0 +1,5 @@
+d_model: 256
+d_kv: 32
+d_ff: 1024
+num_layers: 6
+num_heads: 8
diff --git a/configs/architectures/small.yaml b/configs/architectures/small.yaml
@@ -0,0 +1,5 @@
+d_model: 256
+d_kv: 32
+d_ff: 512
+num_layers: 4
+num_heads: 4
diff --git a/dashboard/denoise/main.py b/dashboard/denoise/main.py
@@ -15,7 +15,7 @@
 
 from utils import vocab_size
 from data.midiencoder import QuantizedMidiEncoder
-from data.multitokencoder import MultiVelocityEncoder
+from data.multitokencoder import MultiMidiEncoder
 from data.quantizer import MidiQuantizer, MidiATQuantizer
 from data.dataset import MaskedMidiDataset, load_cache_dataset
 from data.maskedmidiencoder import MaskedMidiEncoder, MaskedNoteEncoder
@@ -105,7 +105,7 @@ def model_predictions_review(
             n_dstart_bins=dataset_cfg.quantization.dstart,
         )
     if train_cfg.tokens_per_note == "multiple":
-        base_tokenizer = MultiVelocityEncoder(
+        base_tokenizer = MultiMidiEncoder(
             quantization_cfg=train_cfg.dataset.quantization,
             time_quantization_method=train_cfg.time_quantization_method,
         )
@@ -164,6 +164,9 @@ def model_predictions_review(
 
     # predict velocities and get src, tgt and model output
     print("Making predictions ...")
+
+    # widget id for streamlit_pianoroll widget
+    key = 0
     for record_id in idxs:
         # Numpy to int :(
         record: dict = dataset.get_complete_record(int(record_id))
@@ -186,21 +189,24 @@ def model_predictions_review(
             pred_piece = MidiPiece(df)
 
         except ValueError:
-            generated_df = pd.DataFrame([[23, 1, 1, 1, 1]], columns=midi_columns)
+            generated_df = pd.DataFrame([[23.0, 1.0, 1.0, 1.0, 1.0]], columns=midi_columns)
             generated_df["mask"] = [False]
             pred_piece = MidiPiece(generated_df)
 
         pred_piece.source = true_piece.source.copy()
 
         # create a dashboard
-        st.json(record_source)
+        st.json(record_source, expanded=False)
         cols = st.columns(2)
 
         source_tokens: list[str] = [dataset.encoder.vocab[idx] for idx in src_token_ids]
         tgt_tokens: list[str] = [dataset.encoder.vocab[idx] for idx in record["target_token_ids"]]
         generated_tokens: list[str] = [dataset.encoder.vocab[idx] for idx in generated_token_ids]
+
         with cols[0]:
-            from_fortepyan(true_piece)
+            fig = ff.view.draw_pianoroll_with_velocities(true_piece)
+            st.pyplot(fig)
+            from_fortepyan(true_piece, key=key)
             # Unchanged
             st.markdown("**Source tokens:**")
             st.markdown(source_tokens)
@@ -211,9 +217,10 @@ def model_predictions_review(
             # Predicted
             fig = ff.view.draw_dual_pianoroll(pred_piece)
             st.pyplot(fig)
-            from_fortepyan(pred_piece)
+            from_fortepyan(pred_piece, key=key + 1)
             st.markdown("**Predicted tokens:**")
             st.markdown(generated_tokens)
+        key += 2
 
 
 if __name__ == "__main__":
diff --git a/dashboard/download_models.py b/dashboard/download_models.py
@@ -1,15 +1,15 @@
 from huggingface_hub import hf_hub_download
 
-FILENAME_VELOCITY = "midi-T5-2023-10-20-16-03.pt"
-FILENAME_DENOISE = "midi-T5-2023-11-07-12-53.pt"
+FILENAME_VELOCITY = "velocity-T5-2023-11-11-10-29.pt"
+FILENAME_DENOISE = "midi-T5-2023-11-11-10-29.pt"
+
 
 hf_hub_download(
     repo_id="wmatejuk/midi-T5-velocity",
     filename=FILENAME_VELOCITY,
     local_dir="checkpoints/velocity",
     local_dir_use_symlinks=False,
 )
-
 hf_hub_download(
     repo_id="wmatejuk/midi-T5-denoise",
     filename=FILENAME_DENOISE,
diff --git a/dashboard/velocity/main.py b/dashboard/velocity/main.py
@@ -15,8 +15,9 @@
 
 from utils import vocab_size
 from data.midiencoder import VelocityEncoder
-from data.multitokencoder import MultiVelocityEncoder
+from data.maskedmidiencoder import MaskedMidiEncoder
 from data.dataset import MyTokenizedMidiDataset, load_cache_dataset
+from data.multitokencoder import MultiMidiEncoder, MultiVelocityEncoder
 
 # Set the layout of the Streamlit page
 st.set_page_config(layout="wide", page_title="Velocity Transformer", page_icon=":musical_keyboard")
@@ -85,7 +86,19 @@ def model_predictions_review(
         split=split,
     )
 
-    if train_cfg.tokens_per_note == "multiple":
+    if "finetune" in train_cfg.train and train_cfg.train.finetune:
+        tokenizer = MultiMidiEncoder(
+            quantization_cfg=train_cfg.dataset.quantization,
+            time_quantization_method=train_cfg.time_quantization_method,
+        )
+        pretraining_tokenizer = MaskedMidiEncoder(
+            base_encoder=tokenizer,
+        )
+        # use the same token ids as used during pre-training
+        tokenizer.vocab = pretraining_tokenizer.vocab
+        tokenizer.token_to_id = pretraining_tokenizer.token_to_id
+        tokenizer.specials = pretraining_tokenizer.specials
+    elif train_cfg.tokens_per_note == "multiple":
         tokenizer = MultiVelocityEncoder(
             quantization_cfg=train_cfg.dataset.quantization,
             time_quantization_method=train_cfg.time_quantization_method,
@@ -190,7 +203,7 @@ def model_predictions_review(
         pred_piece.source = true_piece.source.copy()
 
         # create a dashboard
-        st.json(record_source)
+        st.json(record_source, expanded=False)
         cols = st.columns(2)
         with cols[0]:
             # Unchanged
diff --git a/data/dataset.py b/data/dataset.py
@@ -12,10 +12,10 @@
 from torch.utils.data import Dataset as TorchDataset
 from datasets import Dataset, load_dataset, concatenate_datasets
 
-from data.midiencoder import MidiEncoder
 from data.multitokencoder import MultiTokEncoder
+from data.maskedmidiencoder import MaskedMidiEncoder
+from data.midiencoder import MidiEncoder, VelocityEncoder
 from data.quantizer import MidiQuantizer, MidiATQuantizer
-from data.maskedmidiencoder import MaskedMidiEncoder, MaskedNoteEncoder
 
 
 def build_AT_translation_dataset(
@@ -363,23 +363,22 @@ def load_cache_dataset(
 def main():
     dataset_name = "roszcz/maestro-v1-sustain"
     dataset_cfg = {
-        "sequence_duration": 5,
-        "sequence_step": 10,
+        "sequence_len": 128,
+        "sequence_step": 42,
         "quantization": {
             "duration": 3,
             "velocity": 3,
             # 650 start bins sound nice :)
-            "start": 20,
+            "dstart": 3,
         },
     }
     cfg = OmegaConf.create(dataset_cfg)
     dataset = load_cache_dataset(cfg, dataset_name, split="test")
 
-    quantizer = MidiATQuantizer(
+    quantizer = MidiQuantizer(
         n_duration_bins=cfg.quantization.duration,
         n_velocity_bins=cfg.quantization.velocity,
-        n_start_bins=cfg.quantization.start,
-        sequence_duration=cfg.sequence_duration,
+        n_dstart_bins=cfg.quantization.dstart,
     )
 
     lengths = [len(record["pitch"]) for record in dataset]
@@ -397,22 +396,16 @@ def main():
     print(piece.df)
     # ff.view.make_piano_roll_video(piece, "test.mp4")
 
-    from data.multitokencoder import MultiVelocityEncoder
-
     # this is for testing and debugging btw
-    base_encoder = MultiVelocityEncoder(cfg.quantization, time_quantization_method="start")
-    encoder = MaskedNoteEncoder(base_encoder=base_encoder, masking_probability=0.3)
-    test_dataset = MaskedMidiDataset(
+    encoder = VelocityEncoder(cfg.quantization, time_quantization_method="dstart")
+    test_dataset = MyTokenizedMidiDataset(
+        encoder=encoder,
         dataset=dataset,
         dataset_cfg=cfg,
-        base_encoder=base_encoder,
-        encoder=encoder,
     )
     record = test_dataset[90]
     print([encoder.vocab[token] for token in record["source_token_ids"]])
-
-    df = test_dataset.encoder.decode(record["source_token_ids"], record["target_token_ids"])
-    print(df)
+    print([encoder.vocab[token] for token in record["target_token_ids"]])
 
 
 if __name__ == "__main__":
diff --git a/data/midiencoder.py b/data/midiencoder.py
@@ -9,6 +9,7 @@ def __init__(self):
         self.token_to_id = None
         self.vocab = None
         self.time_key = None
+        self.specials = None
 
     def tokenize_src(self, record: dict) -> list[str]:
         raise NotImplementedError("Your encoder needs *tokenize* implementation")
diff --git a/data/multitokencoder.py b/data/multitokencoder.py
diff --git a/pipelines/T5/main.py b/pipelines/T5/main.py
diff --git a/requirements.txt b/requirements.txt
diff --git a/utils.py b/utils.py
