Debuged generate.py

.gitignore

@@ -0,0 +1 @@
+.DS_Store

generate.py

@@ -8,13 +8,40 @@ from torch import Tensor, argmax
 from torch.utils.data import DataLoader
 from torch.cuda import is_available as cuda_available, is_bf16_supported
 from torch.backends.mps import is_available as mps_available
-from transformers import AutoModelForCausalLM, MistralConfig, Trainer, TrainingArguments, GenerationConfig, AutoTokenizer, AutoModel
+from transformers import AutoModelForCausalLM, MistralConfig, Trainer, TrainingArguments, GenerationConfig, AutoTokenizer, MistralForCausalLM
 from transformers.trainer_utils import set_seed
 from evaluate import load as load_metric
 from miditok import REMI, TokenizerConfig
 from miditok.pytorch_data import DatasetTok, DataCollator
 from tqdm import tqdm
 
+# Our tokenizer's configuration
+PITCH_RANGE = (21, 109)
+BEAT_RES = {(0, 1): 8, (1, 2): 4, (2, 4): 2, (4, 8): 1}
+NUM_VELOCITIES = 24
+SPECIAL_TOKENS = ["PAD", "MASK", "BOS", "EOS"]
+USE_CHORDS = False
+USE_RESTS = False
+USE_TEMPOS = True
+USE_TIME_SIGNATURE = False
+USE_PROGRAMS = False
+NUM_TEMPOS = 32
+TEMPO_RANGE = (50, 200)  # (min_tempo, max_tempo)
+TOKENIZER_PARAMS = {
+    "pitch_range": PITCH_RANGE,
+    "beat_res": BEAT_RES,
+    "num_velocities": NUM_VELOCITIES,
+    "special_tokens": SPECIAL_TOKENS,
+    "use_chords": USE_CHORDS,
+    "use_rests": USE_RESTS,
+    "use_tempos": USE_TEMPOS,
+    "use_time_signatures": USE_TIME_SIGNATURE,
+    "use_programs": USE_PROGRAMS,
+    "num_tempos": NUM_TEMPOS,
+    "tempo_range": TEMPO_RANGE,
+}
+config = TokenizerConfig(**TOKENIZER_PARAMS)
+
 # Seed
 set_seed(777)
 
@@ -39,7 +66,7 @@ collator = DataCollator(
 )
 
 # Creates model using the correct configuration
-model = AutoModelForCausalLM.from_pretrained("sunsetsobserver/MIDI/runs")
+model = MistralForCausalLM.from_pretrained("./runs")
 
 collator = DataCollator(tokenizer["PAD_None"], tokenizer["BOS_None"], tokenizer["EOS_None"], copy_inputs_as_labels=True)
 
@@ -53,7 +80,6 @@ generation_config = GenerationConfig(
     top_p=0.95,
     epsilon_cutoff=3e-4,
     eta_cutoff=1e-3,
-    pad_token_id=tokenizer.padding_token_id,
 )
 
 # Here the sequences are padded to the left, so that the last token along the time dimension
@@ -77,10 +103,10 @@ for batch in tqdm(dataloader_test, desc='Testing model / Generating results'):
         tokens = [seq.tolist() for seq in tokens]
         for tok_seq in tokens[1:]:
             _midi = tokenizer.tokens_to_midi([deepcopy(tok_seq)])
-            midi.instruments.append(_midi.instruments[0])
-        midi.instruments[0].name = f'Continuation of original sample ({len(generated)} tokens)'
-        midi.instruments[1].name = f'Original sample ({len(prompt)} tokens)'
-        midi.instruments[2].name = f'Original sample and continuation'
+            midi.tracks.append(_midi.tracks[0])
+        midi.tracks[0].name = f'Continuation of original sample ({len(generated)} tokens)'
+        midi.tracks[1].name = f'Original sample ({len(prompt)} tokens)'
+        midi.tracks[2].name = f'Original sample and continuation'
         midi.dump_midi(gen_results_path / f'{count}.mid')
         tokenizer.save_tokens(tokens, gen_results_path / f'{count}.json')