{ "cells": [ { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "SiTIpPjArIyr" }, "source": [ "# Full example with the Hugging Face Transformers package\n", "\n", "This notebook shows how to train a model (Mistral) and generate music with it, using the Hugging Face Transformers package." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "id": "gOd93yV0sGd2" }, "source": [ "## Setup Environment" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "fX12Yquyuihc" }, "outputs": [], "source": [ "#@title Install all dependencies (run only once per session)\n", "\n", "# Run the following command only if a nvidia driver is installed\n", "%nvidia-smi\n", "\n", "%pip install miditok\n", "%pip install symusic\n", "%pip install torch\n", "%pip install transformers\n", "%pip install accelerate\n", "%pip install evaluate\n", "%pip install tensorboard\n", "%pip install scikit-learn\n", "\n", "%wget https://storage.googleapis.com/magentadata/datasets/maestro/v3.0.0/maestro-v3.0.0-midi.zip\n", "%unzip 'maestro-v3.0.0-midi.zip'\n", "%rm 'maestro-v3.0.0-midi.zip'\n", "%mv 'maestro-v3.0.0' 'Maestro'\n", "\n", "from copy import deepcopy\n", "from pathlib import Path\n", "from random import shuffle\n", "\n", "from evaluate import load as load_metric\n", "from miditok import REMI, TokenizerConfig\n", "from miditok.pytorch_data import DatasetMIDI, DataCollator\n", "from miditok.utils import split_files_for_training\n", "from miditok.data_augmentation import augment_dataset\n", "from torch import Tensor, argmax\n", "from torch.utils.data import DataLoader\n", "from torch.cuda import is_available as cuda_available, is_bf16_supported\n", "from torch.backends.mps import is_available as mps_available\n", "from transformers import AutoModelForCausalLM, MistralConfig, Trainer, TrainingArguments, GenerationConfig\n", "from transformers.trainer_utils import set_seed\n", "from tqdm import tqdm" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## Create and train the tokenizer" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Seed\n", "set_seed(777)\n", "\n", "# Our tokenizer's configuration\n", "BEAT_RES = {(0, 1): 12, (1, 2): 4, (2, 4): 2, (4, 8): 1}\n", "TOKENIZER_PARAMS = {\n", " \"pitch_range\": (21, 109),\n", " \"beat_res\": BEAT_RES,\n", " \"num_velocities\": 24,\n", " \"special_tokens\": [\"PAD\", \"BOS\", \"EOS\"],\n", " \"use_chords\": True,\n", " \"use_rests\": True,\n", " \"use_tempos\": True,\n", " \"use_time_signatures\": True,\n", " \"use_programs\": False, # no multitrack here\n", " \"num_tempos\": 32,\n", " \"tempo_range\": (50, 200), # (min_tempo, max_tempo)\n", "}\n", "config = TokenizerConfig(**TOKENIZER_PARAMS)\n", "\n", "# Creates the tokenizer\n", "tokenizer = REMI(config)\n", "\n", "# Trains the tokenizer with Byte Pair Encoding (BPE) to build the vocabulary, here 30k tokens\n", "midi_paths = list(Path(\"Maestro\").resolve().glob(\"**/*.mid\")) + list(Path(\"Maestro\").resolve().glob(\"**/*.midi\"))\n", "tokenizer.train(\n", " vocab_size=30000,\n", " files_paths=midi_paths,\n", ")\n", "tokenizer.save_params(\"tokenizer.json\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Prepare MIDIs for training\n", "\n", "Here we split the files in three subsets: train, validation and test.\n", "Then data augmentation is performed on each subset independently, and the MIDIs are split into smaller chunks that make approximately the desired token sequence length for training." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Split MIDI paths in train/valid/test sets\n", "total_num_files = len(midi_paths)\n", "num_files_valid = round(total_num_files * 0.15)\n", "num_files_test = round(total_num_files * 0.15)\n", "shuffle(midi_paths)\n", "midi_paths_valid = midi_paths[:num_files_valid]\n", "midi_paths_test = midi_paths[num_files_valid:num_files_valid + num_files_test]\n", "midi_paths_train = midi_paths[num_files_valid + num_files_test:]\n", "\n", "# Chunk MIDIs and perform data augmentation on each subset independently\n", "for files_paths, subset_name in (\n", " (midi_paths_train, \"train\"), (midi_paths_valid, \"valid\"), (midi_paths_test, \"test\")\n", "):\n", "\n", " # Split the MIDIs into chunks of sizes approximately about 1024 tokens\n", " subset_chunks_dir = Path(f\"Maestro_{subset_name}\")\n", " split_files_for_training(\n", " files_paths=files_paths,\n", " tokenizer=tokenizer,\n", " save_dir=subset_chunks_dir,\n", " max_seq_len=1024,\n", " num_overlap_bars=2,\n", " )\n", "\n", " # Perform data augmentation\n", " augment_dataset(\n", " subset_chunks_dir,\n", " pitch_offsets=[-12, 12],\n", " velocity_offsets=[-4, 4],\n", " duration_offsets=[-0.5, 0.5],\n", " )\n", "\n", "# Create Dataset and Collator for training\n", "midi_paths_train = list(Path(\"Maestro_train\").glob(\"**/*.mid\")) + list(Path(\"Maestro_train\").glob(\"**/*.midi\"))\n", "midi_paths_valid = list(Path(\"Maestro_valid\").glob(\"**/*.mid\")) + list(Path(\"Maestro_valid\").glob(\"**/*.midi\"))\n", "midi_paths_test = list(Path(\"Maestro_test\").glob(\"**/*.mid\")) + list(Path(\"Maestro_test\").glob(\"**/*.midi\"))\n", "kwargs_dataset = {\"max_seq_len\": 1024, \"tokenizer\": tokenizer, \"bos_token_id\": tokenizer[\"BOS_None\"], \"eos_token_id\": tokenizer[\"EOS_None\"]}\n", "dataset_train = DatasetMIDI(midi_paths_train, **kwargs_dataset)\n", "dataset_valid = DatasetMIDI(midi_paths_valid, **kwargs_dataset)\n", "dataset_test = DatasetMIDI(midi_paths_test, **kwargs_dataset)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## Model initialization\n", "\n", "We will use the [Mistral implementation of Hugging Face](https://huggingface.co/docs/transformers/model_doc/mistral).\n", "Feel free to explore the documentation and source code to dig deeper.\n", "\n", "**You may need to adjust the model's configuration, the training configuration and the maximum input sequence length (cell above) depending on your hardware.**" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [], "source": [ "# Creates model\n", "model_config = MistralConfig(\n", " vocab_size=len(tokenizer),\n", " hidden_size=512,\n", " intermediate_size=2048,\n", " num_hidden_layers=8,\n", " num_attention_heads=8,\n", " num_key_value_heads=4,\n", " sliding_window=256,\n", " max_position_embeddings=8192,\n", " pad_token_id=tokenizer['PAD_None'],\n", " bos_token_id=tokenizer['BOS_None'],\n", " eos_token_id=tokenizer['EOS_None'],\n", ")\n", "model = AutoModelForCausalLM.from_config(model_config)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## Model training" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "metrics = {metric: load_metric(metric) for metric in [\"accuracy\"]}\n", "\n", "def compute_metrics(eval_pred):\n", " \"\"\"\n", " Compute metrics for pretraining.\n", "\n", " Must use preprocess_logits function that converts logits to predictions (argmax or sampling).\n", "\n", " :param eval_pred: EvalPrediction containing predictions and labels\n", " :return: metrics\n", " \"\"\"\n", " predictions, labels = eval_pred\n", " not_pad_mask = labels != -100\n", " labels, predictions = labels[not_pad_mask], predictions[not_pad_mask]\n", " return metrics[\"accuracy\"].compute(predictions=predictions.flatten(), references=labels.flatten())\n", "\n", "def preprocess_logits(logits: Tensor, _: Tensor) -> Tensor:\n", " \"\"\"\n", " Preprocess the logits before accumulating them during evaluation.\n", "\n", " This allows to significantly reduce the memory usage and make the training tractable.\n", " \"\"\"\n", " pred_ids = argmax(logits, dim=-1) # long dtype\n", " return pred_ids\n", "\n", "# Create config for the Trainer\n", "USE_CUDA = cuda_available()\n", "if not cuda_available():\n", " FP16 = FP16_EVAL = BF16 = BF16_EVAL = False\n", "elif is_bf16_supported():\n", " BF16 = BF16_EVAL = True\n", " FP16 = FP16_EVAL = False\n", "else:\n", " BF16 = BF16_EVAL = False\n", " FP16 = FP16_EVAL = True\n", "USE_MPS = not USE_CUDA and mps_available()\n", "training_config = TrainingArguments(\n", " \"runs\", False, True, True, False, \"steps\",\n", " per_device_train_batch_size=16,\n", " per_device_eval_batch_size=48,\n", " gradient_accumulation_steps=3,\n", " eval_accumulation_steps=None,\n", " eval_steps=1000,\n", " learning_rate=1e-4,\n", " weight_decay=0.01,\n", " max_grad_norm=3.0,\n", " max_steps=20000,\n", " lr_scheduler_type=\"cosine_with_restarts\",\n", " warmup_ratio=0.3,\n", " log_level=\"debug\",\n", " logging_strategy=\"steps\",\n", " logging_steps=20,\n", " save_strategy=\"steps\",\n", " save_steps=1000,\n", " save_total_limit=5,\n", " no_cuda=not USE_CUDA,\n", " seed=444,\n", " fp16=FP16,\n", " fp16_full_eval=FP16_EVAL,\n", " bf16=BF16,\n", " bf16_full_eval=BF16_EVAL,\n", " load_best_model_at_end=True,\n", " label_smoothing_factor=0.,\n", " optim=\"adamw_torch\",\n", " report_to=[\"tensorboard\"],\n", " gradient_checkpointing=True,\n", ")\n", "\n", "collator = DataCollator(tokenizer[\"PAD_None\"], copy_inputs_as_labels=True)\n", "trainer = Trainer(\n", " model=model,\n", " args=training_config,\n", " data_collator=collator,\n", " train_dataset=dataset_train,\n", " eval_dataset=dataset_valid,\n", " compute_metrics=compute_metrics,\n", " callbacks=None,\n", " preprocess_logits_for_metrics=preprocess_logits,\n", ")\n", "\n", "# Training\n", "train_result = trainer.train()\n", "trainer.save_model() # Saves the tokenizer too\n", "trainer.log_metrics(\"train\", train_result.metrics)\n", "trainer.save_metrics(\"train\", train_result.metrics)\n", "trainer.save_state()" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## Generate music" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "OaNkGcFo9UP_" }, "outputs": [], "source": [ "(gen_results_path := Path('gen_res')).mkdir(parents=True, exist_ok=True)\n", "generation_config = GenerationConfig(\n", " max_new_tokens=200, # extends samples by 200 tokens\n", " num_beams=1, # no beam search\n", " do_sample=True, # but sample instead\n", " temperature=0.9,\n", " top_k=15,\n", " top_p=0.95,\n", " epsilon_cutoff=3e-4,\n", " eta_cutoff=1e-3,\n", " pad_token_id=tokenizer.pad_token_id,\n", ")\n", "\n", "# Here the sequences are padded to the left, so that the last token along the time dimension\n", "# is always the last token of each seq, allowing to efficiently generate by batch\n", "collator.pad_on_left = True\n", "collator.eos_token = None\n", "dataloader_test = DataLoader(dataset_test, batch_size=16, collate_fn=collator)\n", "model.eval()\n", "count = 0\n", "for batch in tqdm(dataloader_test, desc='Testing model / Generating results'): # (N,T)\n", " res = model.generate(\n", " inputs=batch[\"input_ids\"].to(model.device),\n", " attention_mask=batch[\"attention_mask\"].to(model.device),\n", " generation_config=generation_config) # (N,T)\n", "\n", " # Saves the generated music, as MIDI files and tokens (json)\n", " for prompt, continuation in zip(batch[\"input_ids\"], res):\n", " generated = continuation[len(prompt):]\n", " midi = tokenizer.decode([deepcopy(generated.tolist())])\n", " tokens = [generated, prompt, continuation] # list compr. as seqs of dif. lengths\n", " tokens = [seq.tolist() for seq in tokens]\n", " for tok_seq in tokens[1:]:\n", " _midi = tokenizer.decode([deepcopy(tok_seq)])\n", " midi.tracks.append(_midi.tracks[0])\n", " midi.tracks[0].name = f'Continuation of original sample ({len(generated)} tokens)'\n", " midi.tracks[1].name = f'Original sample ({len(prompt)} tokens)'\n", " midi.tracks[2].name = f'Original sample and continuation'\n", " midi.dump_midi(gen_results_path / f'{count}.mid')\n", " tokenizer.save_tokens(tokens, gen_results_path / f'{count}.json') \n", "\n", " count += 1" ] } ], "metadata": { "accelerator": "GPU", "colab": { "collapsed_sections": [], "machine_shape": "hm", "name": "Optimus_VIRTUOSO_Multi_Instrumental_RGA_Edition.ipynb", "private_outputs": true, "provenance": [] }, "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.9.6" }, "vscode": { "interpreter": { "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6" } } }, "nbformat": 4, "nbformat_minor": 4 }