lab1/TF_Part2_Music_Generation.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uoJsVjtCMunI"
      },
      "source": [
        "<table align=\"center\">\n",
        "  <td align=\"center\"><a target=\"_blank\" href=\"http://introtodeeplearning.com\">\n",
        "        <img src=\"https://i.ibb.co/Jr88sn2/mit.png\" style=\"padding-bottom:5px;\" />\n",
        "      Visit MIT Deep Learning</a></td>\n",
        "  <td align=\"center\"><a target=\"_blank\" href=\"https://colab.research.google.com/github/MITDeepLearning/introtodeeplearning/blob/master/lab1/TF_Part2_Music_Generation.ipynb\">\n",
        "        <img src=\"https://i.ibb.co/2P3SLwK/colab.png\"  style=\"padding-bottom:5px;\" />Run in Google Colab</a></td>\n",
        "  <td align=\"center\"><a target=\"_blank\" href=\"https://github.com/MITDeepLearning/introtodeeplearning/blob/master/lab1/TF_Part2_Music_Generation.ipynb\">\n",
        "        <img src=\"https://i.ibb.co/xfJbPmL/github.png\"  height=\"70px\" style=\"padding-bottom:5px;\"  />View Source on GitHub</a></td>\n",
        "</table>\n",
        "\n",
        "# Copyright Information"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bUik05YqMyCH"
      },
      "outputs": [],
      "source": [
        "# Copyright 2025 MIT Introduction to Deep Learning. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the MIT License. You may not use this file except in compliance\n",
        "# with the License. Use and/or modification of this code outside of MIT Introduction\n",
        "# to Deep Learning must reference:\n",
        "#\n",
        "# © MIT Introduction to Deep Learning\n",
        "# http://introtodeeplearning.com\n",
        "#"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O-97SDET3JG-"
      },
      "source": [
        "# Lab 1: Intro to TensorFlow and Music Generation with RNNs\n",
        "\n",
        "# Part 2: Music Generation with RNNs\n",
        "\n",
        "In this portion of the lab, we will explore building a Recurrent Neural Network (RNN) for music generation. We will train a model to learn the patterns in raw sheet music in [ABC notation](https://en.wikipedia.org/wiki/ABC_notation) and then use this model to generate new music."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rsvlBQYCrE4I"
      },
      "source": [
        "## 2.1 Dependencies\n",
        "First, let's download the course repository, install dependencies, and import the relevant packages we'll need for this lab.\n",
        "\n",
        "We will be using [Comet ML](https://www.comet.com/docs/v2/) to track our model development and training runs. First, sign up for a Comet account [at this link](https://www.comet.com/signup?utm_source=mit_dl&utm_medium=partner&utm_content=github\n",
        ") (you can use your Google or Github account). This will generate a personal API Key, which you can find either in the first 'Get Started with Comet' page, under your account settings, or by pressing the '?' in the top right corner and then 'Quickstart Guide'. Enter this API key as the global variable `COMET_API_KEY`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "riVZCVK65QTH"
      },
      "outputs": [],
      "source": [
        "!pip install comet_ml > /dev/null 2>&1\n",
        "import comet_ml\n",
        "# TODO: ENTER YOUR API KEY HERE!! instructions above\n",
        "COMET_API_KEY = \"\"\n",
        "\n",
        "# Import Tensorflow 2.0\n",
        "import tensorflow as tf\n",
        "\n",
        "# Download and import the MIT Introduction to Deep Learning package\n",
        "!pip install mitdeeplearning --quiet\n",
        "import mitdeeplearning as mdl\n",
        "\n",
        "# Import all remaining packages\n",
        "import numpy as np\n",
        "import os\n",
        "import time\n",
        "import functools\n",
        "from IPython import display as ipythondisplay\n",
        "from tqdm import tqdm\n",
        "from scipy.io.wavfile import write\n",
        "!apt-get install abcmidi timidity > /dev/null 2>&1\n",
        "\n",
        "\n",
        "# Check that we are using a GPU, if not switch runtimes\n",
        "#   using Runtime > Change Runtime Type > GPU\n",
        "# assert len(tf.config.list_physical_devices('GPU')) > 0\n",
        "# assert COMET_API_KEY != \"\", \"Please insert your Comet API Key\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_ajvp0No4qDm"
      },
      "source": [
        "## 2.2 Dataset\n",
        "\n",
        "![Let's Dance!](http://33.media.tumblr.com/3d223954ad0a77f4e98a7b87136aa395/tumblr_nlct5lFVbF1qhu7oio1_500.gif)\n",
        "\n",
        "We've gathered a dataset of thousands of Irish folk songs, represented in the ABC notation. Let's download the dataset and inspect it:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P7dFnP5q3Jve"
      },
      "outputs": [],
      "source": [
        "# Download the dataset\n",
        "songs = mdl.lab1.load_training_data()\n",
        "\n",
        "# Print one of the songs to inspect it in greater detail!\n",
        "example_song = songs[0]\n",
        "print(\"\\nExample song: \")\n",
        "print(example_song)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hKF3EHJlCAj2"
      },
      "source": [
        "We can easily convert a song in ABC notation to an audio waveform and play it back. Be patient for this conversion to run, it can take some time."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "11toYzhEEKDz"
      },
      "outputs": [],
      "source": [
        "# Convert the ABC notation to audio file and listen to it\n",
        "mdl.lab1.play_song(example_song)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7vH24yyquwKQ"
      },
      "source": [
        "One important thing to think about is that this notation of music does not simply contain information on the notes being played, but additionally there is meta information such as the song title, key, and tempo. How does the number of different characters that are present in the text file impact the complexity of the learning problem? This will become important soon, when we generate a numerical representation for the text data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IlCgQBRVymwR"
      },
      "outputs": [],
      "source": [
        "# Join our list of song strings into a single string containing all songs\n",
        "songs_joined = \"\\n\\n\".join(songs)\n",
        "\n",
        "# Find all unique characters in the joined string\n",
        "vocab = sorted(set(songs_joined))\n",
        "print(\"There are\", len(vocab), \"unique characters in the dataset\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rNnrKn_lL-IJ"
      },
      "source": [
        "## 2.3 Process the dataset for the learning task\n",
        "\n",
        "Let's take a step back and consider our prediction task. We're trying to train a RNN model to learn patterns in ABC music, and then use this model to generate (i.e., predict) a new piece of music based on this learned information.\n",
        "\n",
        "Breaking this down, what we're really asking the model is: given a character, or a sequence of characters, what is the most probable next character? We'll train the model to perform this task.\n",
        "\n",
        "To achieve this, we will input a sequence of characters to the model, and train the model to predict the output, that is, the following character at each time step. RNNs maintain an internal state that depends on previously seen elements, so information about all characters seen up until a given moment will be taken into account in generating the prediction."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LFjSVAlWzf-N"
      },
      "source": [
        "### Vectorize the text\n",
        "\n",
        "Before we begin training our RNN model, we'll need to create a numerical representation of our text-based dataset. To do this, we'll generate two lookup tables: one that maps characters to numbers, and a second that maps numbers back to characters. Recall that we just identified the unique characters present in the text."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IalZLbvOzf-F"
      },
      "outputs": [],
      "source": [
        "### Define numerical representation of text ###\n",
        "\n",
        "# Create a mapping from character to unique index.\n",
        "# For example, to get the index of the character \"d\",\n",
        "#   we can evaluate `char2idx[\"d\"]`.\n",
        "char2idx = {u:i for i, u in enumerate(vocab)}\n",
        "\n",
        "# Create a mapping from indices to characters. This is\n",
        "#   the inverse of char2idx and allows us to convert back\n",
        "#   from unique index to the character in our vocabulary.\n",
        "idx2char = np.array(vocab)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tZfqhkYCymwX"
      },
      "source": [
        "This gives us an integer representation for each character. Observe that the unique characters (i.e., our vocabulary) in the text are mapped as indices from 0 to `len(unique)`. Let's take a peek at this numerical representation of our dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FYyNlCNXymwY"
      },
      "outputs": [],
      "source": [
        "print('{')\n",
        "for char,_ in zip(char2idx, range(20)):\n",
        "    print('  {:4s}: {:3d},'.format(repr(char), char2idx[char]))\n",
        "print('  ...\\n}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g-LnKyu4dczc"
      },
      "outputs": [],
      "source": [
        "### Vectorize the songs string ###\n",
        "\n",
        "'''TODO: Write a function to convert the all songs string to a vectorized\n",
        "    (i.e., numeric) representation. Use the appropriate mapping\n",
        "    above to convert from vocab characters to the corresponding indices.\n",
        "\n",
        "  NOTE: the output of the `vectorize_string` function\n",
        "  should be a np.array with `N` elements, where `N` is\n",
        "  the number of characters in the input string\n",
        "'''\n",
        "def vectorize_string(string):\n",
        "  '''TODO'''\n",
        "\n",
        "vectorized_songs = vectorize_string(songs_joined)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IqxpSuZ1w-ub"
      },
      "source": [
        "We can also look at how the first part of the text is mapped to an integer representation:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l1VKcQHcymwb"
      },
      "outputs": [],
      "source": [
        "print ('{} ---- characters mapped to int ----> {}'.format(repr(songs_joined[:10]), vectorized_songs[:10]))\n",
        "# check that vectorized_songs is a numpy array\n",
        "assert isinstance(vectorized_songs, np.ndarray), \"returned result should be a numpy array\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hgsVvVxnymwf"
      },
      "source": [
        "### Create training examples and targets\n",
        "\n",
        "Our next step is to actually divide the text into example sequences that we'll use during training. Each input sequence that we feed into our RNN will contain `seq_length` characters from the text. We'll also need to define a target sequence for each input sequence, which will be used in training the RNN to predict the next character. For each input, the corresponding target will contain the same length of text, except shifted one character to the right.\n",
        "\n",
        "To do this, we'll break the text into chunks of `seq_length+1`. Suppose `seq_length` is 4 and our text is \"Hello\". Then, our input sequence is \"Hell\" and the target sequence is \"ello\".\n",
        "\n",
        "The batch method will then let us convert this stream of character indices to sequences of the desired size."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LF-N8F7BoDRi"
      },
      "outputs": [],
      "source": [
        "### Batch definition to create training examples ###\n",
        "\n",
        "def get_batch(vectorized_songs, seq_length, batch_size):\n",
        "  # the length of the vectorized songs string\n",
        "  n = vectorized_songs.shape[0] - 1\n",
        "  # randomly choose the starting indices for the examples in the training batch\n",
        "  idx = np.random.choice(n-seq_length, batch_size)\n",
        "\n",
        "  '''TODO: construct a list of input sequences for the training batch'''\n",
        "  input_batch = # TODO\n",
        "\n",
        "  '''TODO: construct a list of output sequences for the training batch'''\n",
        "  output_batch = # TODO\n",
        "\n",
        "  # x_batch, y_batch provide the true inputs and targets for network training\n",
        "  x_batch = np.reshape(input_batch, [batch_size, seq_length])\n",
        "  y_batch = np.reshape(output_batch, [batch_size, seq_length])\n",
        "  return x_batch, y_batch\n",
        "\n",
        "\n",
        "# Perform some simple tests to make sure your batch function is working properly!\n",
        "test_args = (vectorized_songs, 10, 2)\n",
        "if not mdl.lab1.test_batch_func_types(get_batch, test_args) or \\\n",
        "   not mdl.lab1.test_batch_func_shapes(get_batch, test_args) or \\\n",
        "   not mdl.lab1.test_batch_func_next_step(get_batch, test_args):\n",
        "   print(\"======\\n[FAIL] could not pass tests\")\n",
        "else:\n",
        "   print(\"======\\n[PASS] passed all tests!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_33OHL3b84i0"
      },
      "source": [
        "For each of these vectors, each index is processed at a single time step. So, for the input at time step 0, the model receives the index for the first character in the sequence, and tries to predict the index of the next character. At the next timestep, it does the same thing, but the RNN considers the information from the previous step, i.e., its updated state, in addition to the current input.\n",
        "\n",
        "We can make this concrete by taking a look at how this works over the first several characters in our text:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0eBu9WZG84i0"
      },
      "outputs": [],
      "source": [
        "x_batch, y_batch = get_batch(vectorized_songs, seq_length=5, batch_size=1)\n",
        "\n",
        "for i, (input_idx, target_idx) in enumerate(zip(np.squeeze(x_batch), np.squeeze(y_batch))):\n",
        "    print(\"Step {:3d}\".format(i))\n",
        "    print(\"  input: {} ({:s})\".format(input_idx, repr(idx2char[input_idx])))\n",
        "    print(\"  expected output: {} ({:s})\".format(target_idx, repr(idx2char[target_idx])))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r6oUuElIMgVx"
      },
      "source": [
        "## 2.4 The Recurrent Neural Network (RNN) model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m8gPwEjRzf-Z"
      },
      "source": [
        "Now we're ready to define and train a RNN model on our ABC music dataset, and then use that trained model to generate a new song. We'll train our RNN using batches of song snippets from our dataset, which we generated in the previous section.\n",
        "\n",
        "The model is based off the LSTM architecture, where we use a state vector to maintain information about the temporal relationships between consecutive characters. The final output of the LSTM is then fed into a fully connected [`Dense`](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Dense) layer where we'll output a softmax over each character in the vocabulary, and then sample from this distribution to predict the next character.\n",
        "\n",
        "As we introduced in the first portion of this lab, we'll be using the Keras API, specifically, [`tf.keras.Sequential`](https://www.tensorflow.org/api_docs/python/tf/keras/models/Sequential), to define the model. Three layers are used to define the model:\n",
        "\n",
        "* [`tf.keras.layers.Embedding`](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Embedding): This is the input layer, consisting of a trainable lookup table that maps the numbers of each character to a vector with `embedding_dim` dimensions.\n",
        "* [`tf.keras.layers.LSTM`](https://www.tensorflow.org/api_docs/python/tf/keras/layers/LSTM): Our LSTM network, with size `units=rnn_units`.\n",
        "* [`tf.keras.layers.Dense`](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Dense): The output layer, with `vocab_size` outputs.\n",
        "\n",
        "\n",
        "<img src=\"https://raw.githubusercontent.com/MITDeepLearning/introtodeeplearning/2019/lab1/img/lstm_unrolled-01-01.png\" alt=\"Drawing\"/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rlaOqndqBmJo"
      },
      "source": [
        "### Define the RNN model\n",
        "\n",
        "Now, we will define a function that we will use to actually build the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8DsWzojvkbc7"
      },
      "outputs": [],
      "source": [
        "def LSTM(rnn_units):\n",
        "  return tf.keras.layers.LSTM(\n",
        "    rnn_units,\n",
        "    return_sequences=True,\n",
        "    recurrent_initializer='glorot_uniform',\n",
        "    recurrent_activation='sigmoid',\n",
        "    stateful=True,\n",
        "  )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IbWU4dMJmMvq"
      },
      "source": [
        "The time has come! Fill in the `TODOs` to define the RNN model within the `build_model` function, and then call the function you just defined to instantiate the model!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MtCrdfzEI2N0"
      },
      "outputs": [],
      "source": [
        "### Defining the RNN Model ###\n",
        "\n",
        "'''TODO: Add LSTM and Dense layers to define the RNN model using the Sequential API.'''\n",
        "def build_model(vocab_size, embedding_dim, rnn_units, batch_size):\n",
        "  model = tf.keras.Sequential([\n",
        "    # Layer 1: Embedding layer to transform indices into dense vectors\n",
        "    #   of a fixed embedding size\n",
        "    tf.keras.layers.Embedding(vocab_size, embedding_dim),\n",
        "\n",
        "    # Layer 2: LSTM with `rnn_units` number of units.\n",
        "    # TODO: Call the LSTM function defined above to add this layer.\n",
        "    LSTM('''TODO'''),\n",
        "\n",
        "    # Layer 3: Dense (fully-connected) layer that transforms the LSTM output\n",
        "    #   into the vocabulary size.\n",
        "    # TODO: Add the Dense layer.\n",
        "    '''TODO: DENSE LAYER HERE'''\n",
        "  ])\n",
        "\n",
        "  return model\n",
        "\n",
        "# Build a simple model with default hyperparameters. You will get the\n",
        "#   chance to change these later.\n",
        "model = build_model(len(vocab), embedding_dim=256, rnn_units=1024, batch_size=32)\n",
        "model.build(tf.TensorShape([32, 100]))  # [batch_size, sequence_length]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-ubPo0_9Prjb"
      },
      "source": [
        "### Test out the RNN model\n",
        "\n",
        "It's always a good idea to run a few simple checks on our model to see that it behaves as expected.  \n",
        "\n",
        "First, we can use the `Model.summary` function to print out a summary of our model's internal workings. Here we can check the layers in the model, the shape of the output of each of the layers, the batch size, etc."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RwG1DD6rDrRM"
      },
      "outputs": [],
      "source": [
        "model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8xeDn5nZD0LX"
      },
      "source": [
        "We can also quickly check the dimensionality of our output, using a sequence length of 100. Note that the model can be run on inputs of any length."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C-_70kKAPrPU"
      },
      "outputs": [],
      "source": [
        "x, y = get_batch(vectorized_songs, seq_length=100, batch_size=32)\n",
        "pred = model(x)\n",
        "print(\"Input shape:      \", x.shape, \" # (batch_size, sequence_length)\")\n",
        "print(\"Prediction shape: \", pred.shape, \"# (batch_size, sequence_length, vocab_size)\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mT1HvFVUGpoE"
      },
      "source": [
        "### Predictions from the untrained model\n",
        "\n",
        "Let's take a look at what our untrained model is predicting.\n",
        "\n",
        "To get actual predictions from the model, we sample from the output distribution, which is defined by a `softmax` over our character vocabulary. This will give us actual character indices. This means we are using a [categorical distribution](https://en.wikipedia.org/wiki/Categorical_distribution) to sample over the example prediction. This gives a prediction of the next character (specifically its index) at each timestep.\n",
        "\n",
        "Note here that we sample from this probability distribution, as opposed to simply taking the `argmax`, which can cause the model to get stuck in a loop.\n",
        "\n",
        "Let's try this sampling out for the first example in the batch."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4V4MfFg0RQJg"
      },
      "outputs": [],
      "source": [
        "sampled_indices = tf.random.categorical(pred[0], num_samples=1)\n",
        "sampled_indices = tf.squeeze(sampled_indices,axis=-1).numpy()\n",
        "sampled_indices"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LfLtsP3mUhCG"
      },
      "source": [
        "We can now decode these to see the text predicted by the untrained model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xWcFwPwLSo05"
      },
      "outputs": [],
      "source": [
        "print(\"Input: \\n\", repr(\"\".join(idx2char[x[0]])))\n",
        "print()\n",
        "print(\"Next Char Predictions: \\n\", repr(\"\".join(idx2char[sampled_indices])))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HEHHcRasIDm9"
      },
      "source": [
        "As you can see, the text predicted by the untrained model is pretty nonsensical! How can we do better? We can train the network!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LJL0Q0YPY6Ee"
      },
      "source": [
        "## 2.5 Training the model: loss and training operations\n",
        "\n",
        "Now it's time to train the model!\n",
        "\n",
        "At this point, we can think of our next character prediction problem as a standard classification problem. Given the previous state of the RNN, as well as the input at a given time step, we want to predict the class of the next character -- that is, to actually predict the next character.\n",
        "\n",
        "To train our model on this classification task, we can use a form of the `crossentropy` loss (negative log likelihood loss). Specifically, we will use the [`sparse_categorical_crossentropy`](https://www.tensorflow.org/api_docs/python/tf/keras/losses/sparse_categorical_crossentropy) loss, as it utilizes integer targets for categorical classification tasks. We will want to compute the loss using the true targets -- the `labels` -- and the predicted targets -- the `logits`.\n",
        "\n",
        "Let's first compute the loss using our example predictions from the untrained model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4HrXTACTdzY-"
      },
      "outputs": [],
      "source": [
        "### Defining the loss function ###\n",
        "\n",
        "'''TODO: define the loss function to compute and return the loss between\n",
        "    the true labels and predictions (logits). Set the argument from_logits=True.'''\n",
        "def compute_loss(labels, logits):\n",
        "  loss = tf.keras.losses.sparse_categorical_crossentropy('''TODO''', '''TODO''', from_logits=True) # TODO\n",
        "  return loss\n",
        "\n",
        "'''TODO: compute the loss using the true next characters from the example batch\n",
        "    and the predictions from the untrained model several cells above'''\n",
        "example_batch_loss = compute_loss('''TODO''', '''TODO''') # TODO\n",
        "\n",
        "print(\"Prediction shape: \", pred.shape, \" # (batch_size, sequence_length, vocab_size)\")\n",
        "print(\"scalar_loss:      \", example_batch_loss.numpy().mean())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Seh7e6eRqd7"
      },
      "source": [
        "Let's start by defining some hyperparameters for training the model. To start, we have provided some reasonable values for some of the parameters. It is up to you to use what we've learned in class to help optimize the parameter selection here!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JQWUUhKotkAY"
      },
      "outputs": [],
      "source": [
        "### Hyperparameter setting and optimization ###\n",
        "\n",
        "vocab_size = len(vocab)\n",
        "\n",
        "# Model parameters:\n",
        "params = dict(\n",
        "  num_training_iterations = 3000,  # Increase this to train longer\n",
        "  batch_size = 8,  # Experiment between 1 and 64\n",
        "  seq_length = 100,  # Experiment between 50 and 500\n",
        "  learning_rate = 5e-3,  # Experiment between 1e-5 and 1e-1\n",
        "  embedding_dim = 256,\n",
        "  rnn_units = 1024,  # Experiment between 1 and 2048\n",
        ")\n",
        "\n",
        "# Checkpoint location:\n",
        "checkpoint_dir = './training_checkpoints'\n",
        "checkpoint_prefix = os.path.join(checkpoint_dir, \"my_ckpt.weights.h5\")\n",
        "os.makedirs(checkpoint_dir, exist_ok=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AyLzIPeAIqfg"
      },
      "source": [
        "Having defined our hyperparameters we can set up for experiment tracking with Comet. [`Experiment`](https://www.comet.com/docs/v2/api-and-sdk/python-sdk/reference/Experiment/) are the core objects in Comet and will allow us to track training and model development. Here we have written a short function to create a new comet experiment. Note that in this setup, when hyperparameters change, you can run the `create_experiment()` function to initiate a new experiment. All experiments defined with the same `project_name` will live under that project in your Comet interface.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MBsN1vvxInmN"
      },
      "outputs": [],
      "source": [
        "### Create a Comet experiment to track our training run ###\n",
        "\n",
        "def create_experiment():\n",
        "  # end any prior experiments\n",
        "  if 'experiment' in locals():\n",
        "    experiment.end()\n",
        "\n",
        "  # initiate the comet experiment for tracking\n",
        "  experiment = comet_ml.Experiment(\n",
        "                  api_key=COMET_API_KEY,\n",
        "                  project_name=\"6S191_Lab1_Part2\")\n",
        "  # log our hyperparameters, defined above, to the experiment\n",
        "  for param, value in params.items():\n",
        "    experiment.log_parameter(param, value)\n",
        "  experiment.flush()\n",
        "\n",
        "  return experiment"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5cu11p1MKYZd"
      },
      "source": [
        "Now, we are ready to define our training operation -- the optimizer and duration of training -- and use this function to train the model. You will experiment with the choice of optimizer and the duration for which you train your models, and see how these changes affect the network's output. Some optimizers you may like to try are [`Adam`](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adam?version=stable) and [`Adagrad`](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adagrad?version=stable).\n",
        "\n",
        "First, we will instantiate a new model and an optimizer. Then, we will use the [`tf.GradientTape`](https://www.tensorflow.org/api_docs/python/tf/GradientTape) method to perform the backpropagation operations.\n",
        "\n",
        "We will also generate a print-out of the model's progress through training, which will help us easily visualize whether or not we are minimizing the loss."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F31vzJ_u66cb"
      },
      "outputs": [],
      "source": [
        "### Define optimizer and training operation ###\n",
        "\n",
        "'''TODO: instantiate a new model for training using the `build_model`\n",
        "  function and the hyperparameters created above.'''\n",
        "model = build_model('''TODO: arguments''')\n",
        "\n",
        "'''TODO: instantiate an optimizer with its learning rate.\n",
        "  Checkout the tensorflow website for a list of supported optimizers.\n",
        "  https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/\n",
        "  Try using the Adam optimizer to start.'''\n",
        "optimizer = # TODO\n",
        "\n",
        "@tf.function\n",
        "def train_step(x, y):\n",
        "  # Use tf.GradientTape()\n",
        "  with tf.GradientTape() as tape:\n",
        "\n",
        "    '''TODO: feed the current input into the model and generate predictions'''\n",
        "    y_hat = model('''TODO''')\n",
        "\n",
        "    '''TODO: compute the loss!'''\n",
        "    loss = compute_loss('''TODO''', '''TODO''')\n",
        "\n",
        "  # Now, compute the gradients\n",
        "  '''TODO: complete the function call for gradient computation.\n",
        "      Remember that we want the gradient of the loss with respect all\n",
        "      of the model parameters.\n",
        "      HINT: use `model.trainable_variables` to get a list of all model\n",
        "      parameters.'''\n",
        "  grads = tape.gradient('''TODO''', '''TODO''')\n",
        "\n",
        "  # Apply the gradients to the optimizer so it can update the model accordingly\n",
        "  optimizer.apply_gradients(zip(grads, model.trainable_variables))\n",
        "  return loss\n",
        "\n",
        "##################\n",
        "# Begin training!#\n",
        "##################\n",
        "\n",
        "history = []\n",
        "plotter = mdl.util.PeriodicPlotter(sec=2, xlabel='Iterations', ylabel='Loss')\n",
        "experiment = create_experiment()\n",
        "\n",
        "if hasattr(tqdm, '_instances'): tqdm._instances.clear() # clear if it exists\n",
        "for iter in tqdm(range(params[\"num_training_iterations\"])):\n",
        "\n",
        "  # Grab a batch and propagate it through the network\n",
        "  x_batch, y_batch = get_batch(vectorized_songs, params[\"seq_length\"], params[\"batch_size\"])\n",
        "  loss = train_step(x_batch, y_batch)\n",
        "\n",
        "  # log the loss to the Comet interface! we will be able to track it there.\n",
        "  experiment.log_metric(\"loss\", loss.numpy().mean(), step=iter)\n",
        "  # Update the progress bar and also visualize within notebook\n",
        "  history.append(loss.numpy().mean())\n",
        "  plotter.plot(history)\n",
        "\n",
        "  # Update the model with the changed weights!\n",
        "  if iter % 100 == 0:\n",
        "    model.save_weights(checkpoint_prefix)\n",
        "\n",
        "# Save the trained model and the weights\n",
        "model.save_weights(checkpoint_prefix)\n",
        "experiment.flush()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kKkD5M6eoSiN"
      },
      "source": [
        "## 2.6 Generate music using the RNN model\n",
        "\n",
        "Now, we can use our trained RNN model to generate some music! When generating music, we'll have to feed the model some sort of seed to get it started (because it can't predict anything without something to start with!).\n",
        "\n",
        "Once we have a generated seed, we can then iteratively predict each successive character (remember, we are using the ABC representation for our music) using our trained RNN. More specifically, recall that our RNN outputs a `softmax` over possible successive characters. For inference, we iteratively sample from these distributions, and then use our samples to encode a generated song in the ABC format.\n",
        "\n",
        "Then, all we have to do is write it to a file and listen!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JIPcXllKjkdr"
      },
      "source": [
        "\n",
        "### Restore the latest checkpoint\n",
        "\n",
        "To keep this inference step simple, we will use a batch size of 1. Because of how the RNN state is passed from timestep to timestep, the model will only be able to accept a fixed batch size once it is built.\n",
        "\n",
        "To run the model with a different `batch_size`, we'll need to rebuild the model and restore the weights from the latest checkpoint, i.e., the weights after the last checkpoint during training:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LycQ-ot_jjyu"
      },
      "outputs": [],
      "source": [
        "'''TODO: Rebuild the model using a batch_size=1'''\n",
        "model = build_model('''TODO''', '''TODO''', '''TODO''', batch_size=1)\n",
        "\n",
        "# Restore the model weights for the last checkpoint after training\n",
        "model.build(tf.TensorShape([1, None]))\n",
        "model.load_weights(checkpoint_prefix)\n",
        "\n",
        "model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I9b4V2C8N62l"
      },
      "source": [
        "Notice that we have fed in a fixed `batch_size` of 1 for inference."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DjGz1tDkzf-u"
      },
      "source": [
        "### The prediction procedure\n",
        "\n",
        "Now, we're ready to write the code to generate text in the ABC music format:\n",
        "\n",
        "* Initialize a \"seed\" start string and the RNN state, and set the number of characters we want to generate.\n",
        "\n",
        "* Use the start string and the RNN state to obtain the probability distribution over the next predicted character.\n",
        "\n",
        "* Sample from multinomial distribution to calculate the index of the predicted character. This predicted character is then used as the next input to the model.\n",
        "\n",
        "* At each time step, the updated RNN state is fed back into the model, so that it now has more context in making the next prediction. After predicting the next character, the updated RNN states are again fed back into the model, which is how it learns sequence dependencies in the data, as it gets more information from the previous predictions.\n",
        "\n",
        "![LSTM inference](https://raw.githubusercontent.com/MITDeepLearning/introtodeeplearning/2019/lab1/img/lstm_inference.png)\n",
        "\n",
        "Complete and experiment with this code block (as well as some of the aspects of network definition and training!), and see how the model performs. How do songs generated after training with a small number of epochs compare to those generated after a longer duration of training?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WvuwZBX5Ogfd"
      },
      "outputs": [],
      "source": [
        "### Prediction of a generated song ###\n",
        "\n",
        "def generate_text(model, start_string, generation_length=1000):\n",
        "  # Evaluation step (generating ABC text using the learned RNN model)\n",
        "\n",
        "  '''TODO: convert the start string to numbers (vectorize)'''\n",
        "  input_eval = ['''TODO''']\n",
        "  input_eval = tf.expand_dims(input_eval, 0)\n",
        "\n",
        "  # Empty string to store our results\n",
        "  text_generated = []\n",
        "\n",
        "  # Here batch size == 1\n",
        "  model.reset_states()\n",
        "  tqdm._instances.clear()\n",
        "\n",
        "  for i in tqdm(range(generation_length)):\n",
        "      '''TODO: evaluate the inputs and generate the next character predictions'''\n",
        "      predictions = model('''TODO''')\n",
        "\n",
        "      # Remove the batch dimension\n",
        "      predictions = tf.squeeze(predictions, 0)\n",
        "\n",
        "      '''TODO: use a multinomial distribution to sample'''\n",
        "      predicted_id = tf.random.categorical('''TODO''', num_samples=1)[-1,0].numpy()\n",
        "\n",
        "      # Pass the prediction along with the previous hidden state\n",
        "      #   as the next inputs to the model\n",
        "      input_eval = tf.expand_dims([predicted_id], 0)\n",
        "\n",
        "      '''TODO: add the predicted character to the generated text!'''\n",
        "      # Hint: consider what format the prediction is in vs. the output\n",
        "      text_generated.append('''TODO''')\n",
        "\n",
        "  return (start_string + ''.join(text_generated))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ktovv0RFhrkn"
      },
      "outputs": [],
      "source": [
        "'''TODO: Use the model and the function defined above to generate ABC format text of length 1000!\n",
        "    As you may notice, ABC files start with \"X\" - this may be a good start string.'''\n",
        "generated_text = generate_text('''TODO''', start_string=\"X\", generation_length=1000)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AM2Uma_-yVIq"
      },
      "source": [
        "### Play back the generated music!\n",
        "\n",
        "We can now call a function to convert the ABC format text to an audio file, and then play that back to check out our generated music! Try training longer if the resulting song is not long enough, or re-generating the song!\n",
        "\n",
        "We will save the song to Comet -- you will be able to find your songs under the `Audio` and `Assets & Artificats` pages in your Comet interface for the project. Note the [`log_asset()`](https://www.comet.com/docs/v2/api-and-sdk/python-sdk/reference/Experiment/#experimentlog_asset) documentation, where you will see how to specify file names and other parameters for saving your assets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LrOtG64bfLto"
      },
      "outputs": [],
      "source": [
        "### Play back generated songs ###\n",
        "\n",
        "generated_songs = mdl.lab1.extract_song_snippet(generated_text)\n",
        "\n",
        "for i, song in enumerate(generated_songs):\n",
        "  # Synthesize the waveform from a song\n",
        "  waveform = mdl.lab1.play_song(song)\n",
        "\n",
        "  # If its a valid song (correct syntax), lets play it!\n",
        "  if waveform:\n",
        "    print(\"Generated song\", i)\n",
        "    ipythondisplay.display(waveform)\n",
        "\n",
        "    numeric_data = np.frombuffer(waveform.data, dtype=np.int16)\n",
        "    wav_file_path = f\"output_{i}.wav\"\n",
        "    write(wav_file_path, 88200, numeric_data)\n",
        "\n",
        "    # save your song to the Comet interface -- you can access it there\n",
        "    experiment.log_asset(wav_file_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4353qSV76gnJ"
      },
      "outputs": [],
      "source": [
        "# when done, end the comet experiment\n",
        "experiment.end()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HgVvcrYmSKGG"
      },
      "source": [
        "## 2.7 Experiment and **get awarded for the best songs**!\n",
        "\n",
        "Congrats on making your first sequence model in TensorFlow! It's a pretty big accomplishment, and hopefully you have some sweet tunes to show for it.\n",
        "\n",
        "Consider how you may improve your model and what seems to be most important in terms of performance. Here are some ideas to get you started:\n",
        "\n",
        "*  How does the number of training epochs affect the performance?\n",
        "*  What if you alter or augment the dataset?\n",
        "*  Does the choice of start string significantly affect the result?\n",
        "\n",
        "Try to optimize your model and submit your best song! **Participants will be eligible for prizes during the January 2025 offering. To enter the competition, you must upload the following to [this submission link](https://www.dropbox.com/request/U8nND6enGjirujVZKX1n):**\n",
        "\n",
        "* a recording of your song;\n",
        "* iPython notebook with the code you used to generate the song;\n",
        "* a description and/or diagram of the architecture and hyperparameters you used -- if there are any additional or interesting modifications you made to the template code, please include these in your description.\n",
        "\n",
        "**Name your file in the following format: ``[FirstName]_[LastName]_RNNMusic``, followed by the file format (.zip, .mp4, .ipynb, .pdf, etc). ZIP files of all three components are preferred over individual files. If you submit individual files, you must name the individual files according to the above nomenclature.**\n",
        "\n",
        "You can also tweet us at [@MITDeepLearning](https://twitter.com/MITDeepLearning) a copy of the song (but this will not enter you into the competition)! See this example song generated by a previous student (credit Ana Heart): <a href=\"https://twitter.com/AnaWhatever16/status/1263092914680410112?s=20\">song from May 20, 2020.</a>\n",
        "<script async src=\"https://platform.twitter.com/widgets.js\" charset=\"utf-8\"></script>\n",
        "\n",
        "Have fun and happy listening!\n",
        "\n",
        "![Let's Dance!](http://33.media.tumblr.com/3d223954ad0a77f4e98a7b87136aa395/tumblr_nlct5lFVbF1qhu7oio1_500.gif)\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "uoJsVjtCMunI"
      ],
      "name": "TF_Part2_Music_Generation.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.11.11"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}