colab-notebooks/Example_HuggingFace_Mistral_Transformer.ipynb

{
 "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
}