Prep for ALICE training

train.py

@@ -182,47 +182,4 @@ trainer.save_model()  # Saves the tokenizer too
 trainer.log_metrics("train", train_result.metrics)
 trainer.save_metrics("train", train_result.metrics)
 trainer.save_state()
-
-
-(gen_results_path := Path('gen_res')).mkdir(parents=True, exist_ok=True)
-generation_config = GenerationConfig(
-    max_new_tokens=512,  # extends samples by 512 tokens
-    num_beams=1,        # no beam search
-    do_sample=True,     # but sample instead
-    temperature=0.9,
-    top_k=15,
-    top_p=0.95,
-    epsilon_cutoff=3e-4,
-    eta_cutoff=1e-3,
-    pad_token_id=config.padding_token_id,
-)
-
-# Here the sequences are padded to the left, so that the last token along the time dimension
-# is always the last token of each seq, allowing to efficiently generate by batch
-collator.pad_on_left = True
-collator.eos_token = None
-dataloader_test = DataLoader(dataset_test, batch_size=16, collate_fn=collator)
-model.eval()
-count = 0
-for batch in tqdm(dataloader_test, desc='Testing model / Generating results'):  # (N,T)
-    res = model.generate(
-        inputs=batch["input_ids"].to(model.device),
-        attention_mask=batch["attention_mask"].to(model.device),
-        generation_config=generation_config)  # (N,T)
-
-    # Saves the generated music, as MIDI files and tokens (json)
-    for prompt, continuation in zip(batch["input_ids"], res):
-        generated = continuation[len(prompt):]
-        midi = tokenizer.tokens_to_midi([deepcopy(generated.tolist())])
-        tokens = [generated, prompt, continuation]  # list compr. as seqs of dif. lengths
-        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.dump_midi(gen_results_path / f'{count}.mid')
-        tokenizer.save_tokens(tokens, gen_results_path / f'{count}.json') 
-
-        count += 1
+trainer.push_to_hub()