lab1/solutions/PT_Part2_Music_Generation_Solution.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/solutions/PT_Part2_Music_Generation_Solution.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/solutions/PT_Part2_Music_Generation_Solution.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 PyTorch 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 using PyTorch. 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 PyTorch and other relevant libraries\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\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 torch.cuda.is_available(), \"Please enable GPU from runtime settings\"\n",
        "assert COMET_API_KEY != \"\", \"Please insert your Comet API Key\"\n",
        "\n"
      ]
    },
    {
      "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 an 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.\n"
      ]
    },
    {
      "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",
        "  vectorized_output = np.array([char2idx[char] for char in string])\n",
        "  return vectorized_output\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:\n"
      ]
    },
    {
      "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.\n"
      ]
    },
    {
      "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 = [vectorized_songs[i: i + seq_length] for i in idx]\n",
        "\n",
        "    '''TODO: construct a list of output sequences for the training batch'''\n",
        "    output_batch = [vectorized_songs[i + 1: i + seq_length + 1] for i in idx]\n",
        "\n",
        "    # Convert the input and output batches to tensors\n",
        "    x_batch = torch.tensor(input_batch, dtype=torch.long)\n",
        "    y_batch = torch.tensor(output_batch, dtype=torch.long)\n",
        "\n",
        "    return x_batch, y_batch\n",
        "\n",
        "# Perform some simple tests to make sure your batch function is working properly!\n",
        "test_args = (vectorized_songs, 10, 2)\n",
        "x_batch, y_batch = get_batch(*test_args)\n",
        "assert x_batch.shape == (2, 10), \"x_batch shape is incorrect\"\n",
        "assert y_batch.shape == (2, 10), \"y_batch shape is incorrect\"\n",
        "print(\"Batch function works correctly!\")\n"
      ]
    },
    {
      "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(x_batch[0], y_batch[0])):\n",
        "    print(\"Step {:3d}\".format(i))\n",
        "    print(\"  input: {} ({:s})\".format(input_idx, repr(idx2char[input_idx.item()])))\n",
        "    print(\"  expected output: {} ({:s})\".format(target_idx, repr(idx2char[target_idx.item()])))\n"
      ]
    },
    {
      "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 an 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 linear [`nn.Linear`](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html) 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 PyTorch's [`nn.Module`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html) to define the model. Three components are used to define the model:\n",
        "\n",
        "* [`nn.Embedding`](https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html): 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",
        "* [`nn.LSTM`](https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html): Our LSTM network, with size `hidden_size`.\n",
        "* [`nn.Linear`](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html): The output layer, with `vocab_size` outputs.\n",
        "\n",
        "<img src=\"https://raw.githubusercontent.com/MITDeepLearning/introtodeeplearning/2019/lab1/img/lstm_unrolled-01-01.png\" alt=\"Drawing\"/>\n",
        "\n",
        "\n",
        "\n",
        "<!--\n",
        "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",
        "Let's define our model as an `nn.Module`. Fill in the `TODOs` to define the RNN model.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8DsWzojvkbc7"
      },
      "outputs": [],
      "source": [
        "### Defining the RNN Model ###\n",
        "\n",
        "'''TODO: Add LSTM and Linear layers to define the RNN model using nn.Module'''\n",
        "class LSTMModel(nn.Module):\n",
        "    def __init__(self, vocab_size, embedding_dim, hidden_size):\n",
        "        super(LSTMModel, self).__init__()\n",
        "        self.hidden_size = hidden_size\n",
        "\n",
        "        # Define each of the network layers\n",
        "        # Layer 1: Embedding layer to transform indices into dense vectors\n",
        "        #   of a fixed embedding size\n",
        "        self.embedding = nn.Embedding(vocab_size, embedding_dim)\n",
        "\n",
        "        # Layer 2: LSTM with hidden_size `hidden_size`. note: number of layers defaults to 1.\n",
        "        # TODO: Use the nn.LSTM() module from pytorch.\n",
        "        self.lstm = nn.LSTM(embedding_dim, hidden_size, batch_first=True)\n",
        "        # self.lstm = nn.LSTM('''TODO''')\n",
        "\n",
        "        # Layer 3: Linear (fully-connected) layer that transforms the LSTM output\n",
        "        #   into the vocabulary size.\n",
        "        # TODO: Add the Linear layer.\n",
        "        self.fc = nn.Linear(hidden_size, vocab_size)\n",
        "        # self.fc = nn.Linear('''TODO''')\n",
        "\n",
        "    def init_hidden(self, batch_size, device):\n",
        "        # Initialize hidden state and cell state with zeros\n",
        "        return (torch.zeros(1, batch_size, self.hidden_size).to(device),\n",
        "                torch.zeros(1, batch_size, self.hidden_size).to(device))\n",
        "\n",
        "    def forward(self, x, state=None, return_state=False):\n",
        "        x = self.embedding(x)\n",
        "\n",
        "        if state is None:\n",
        "            state = self.init_hidden(x.size(0), x.device)\n",
        "        out, state = self.lstm(x, state)\n",
        "\n",
        "        out = self.fc(out)\n",
        "        return out if not return_state else (out, state)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IbWU4dMJmMvq"
      },
      "source": [
        "The time has come! Let's instantiate the model!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MtCrdfzEI2N0"
      },
      "outputs": [],
      "source": [
        "# Instantiate the model! Build a simple model with default hyperparameters. You\n",
        "#     will get the chance to change these later.\n",
        "vocab_size = len(vocab)\n",
        "embedding_dim = 256\n",
        "hidden_size = 1024\n",
        "batch_size = 8\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "model = LSTMModel(vocab_size, embedding_dim, hidden_size).to(device)\n",
        "\n",
        "# print out a summary of the model\n",
        "print(model)"
      ]
    },
    {
      "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",
        "We can quickly check the layers in the model, the shape of the output of each of the layers, the batch size, and the dimensionality of the output. Note that the model can be run on inputs of any length."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C-_70kKAPrPU"
      },
      "outputs": [],
      "source": [
        "# Test the model with some sample data\n",
        "x, y = get_batch(vectorized_songs, seq_length=100, batch_size=32)\n",
        "x = x.to(device)\n",
        "y = y.to(device)\n",
        "\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 torch.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. [`torch.multinomial`](https://pytorch.org/docs/stable/generated/torch.multinomial.html#torch.multinomial) samples over a categorical distribution to generate predictions.\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 repetitive 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 = torch.multinomial(torch.softmax(pred[0], dim=-1), num_samples=1)\n",
        "sampled_indices = sampled_indices.squeeze(-1).cpu().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].cpu()])))\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? Well, 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 (i.e., negative log likelihood loss). Specifically, we will use PyTorch's [`CrossEntropyLoss`](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html), as it combines the application of a log-softmax ([`LogSoftmax`](https://pytorch.org/docs/stable/generated/torch.nn.LogSoftmax.html#torch.nn.LogSoftmax)) and negative log-likelihood ([`NLLLoss`](https://pytorch.org/docs/stable/generated/torch.nn.NLLLoss.html#torch.nn.NLLLoss) in a single class and accepts 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 define a function to compute the loss, and then use that function to 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 compute_loss function to compute and return the loss between\n",
        "#     the true labels and predictions (logits). '''\n",
        "cross_entropy = nn.CrossEntropyLoss() # instantiates the function\n",
        "def compute_loss(labels, logits):\n",
        "    \"\"\"\n",
        "    Inputs:\n",
        "      labels: (batch_size, sequence_length)\n",
        "      logits: (batch_size, sequence_length, vocab_size)\n",
        "\n",
        "    Output:\n",
        "      loss: scalar cross entropy loss over the batch and sequence length\n",
        "    \"\"\"\n",
        "\n",
        "    # Batch the labels so that the shape of the labels should be (B * L,)\n",
        "    batched_labels = labels.view(-1)\n",
        "\n",
        "    ''' TODO: Batch the logits so that the shape of the logits should be (B * L, V) '''\n",
        "    batched_logits = logits.view(-1, logits.size(-1))\n",
        "    # batched_logits = \"\"\" TODO \"\"\" # TODO\n",
        "\n",
        "    '''TODO: Compute the cross-entropy loss using the batched  next characters and predictions'''\n",
        "    loss = cross_entropy(batched_logits, batched_labels)\n",
        "    # loss = \"\"\" TODO \"\"\" # TODO\n",
        "    return loss"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "### compute the loss on the predictions from the untrained model from earlier. ###\n",
        "y.shape  # (batch_size, sequence_length)\n",
        "pred.shape  # (batch_size, sequence_length, vocab_size)\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(y, pred)\n",
        "# example_batch_loss = compute_loss('''TODO''', '''TODO''') # TODO\n",
        "\n",
        "print(f\"Prediction shape: {pred.shape} # (batch_size, sequence_length, vocab_size)\")\n",
        "print(f\"scalar_loss:      {example_batch_loss.mean().item()}\")"
      ],
      "metadata": {
        "id": "GuGUJB0ZT_Uo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "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",
        "  hidden_size = 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\")\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://pytorch.org/docs/stable/generated/torch.optim.Adam.html) and [`Adagrad`](https://pytorch.org/docs/stable/generated/torch.optim.Adagrad.html).\n",
        "\n",
        "First, we will instantiate a new model and an optimizer, and ready them for training. Then, we will use [`loss.backward()`](https://pytorch.org/docs/stable/generated/torch.Tensor.backward.html), enabled by PyTorch's [autograd](https://pytorch.org/docs/stable/generated/torch.autograd.grad.html) method, to perform the backpropagation. Finally, to update the model's parameters based on the computed gradients, we will utake a step with the optimizer, using [`optimizer.step()`](https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.step.html).\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 LSTMModel model for training using the hyperparameters\n",
        "    created above.'''\n",
        "model = LSTMModel(vocab_size, params[\"embedding_dim\"], params[\"hidden_size\"])\n",
        "# model = LSTMModel('''TODO: arguments''')\n",
        "\n",
        "# Move the model to the GPU\n",
        "model.to(device)\n",
        "\n",
        "'''TODO: instantiate an optimizer with its learning rate.\n",
        "  Checkout the PyTorch website for a list of supported optimizers.\n",
        "  https://pytorch.org/docs/stable/optim.html\n",
        "  Try using the Adam optimizer to start.'''\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=params[\"learning_rate\"])\n",
        "# optimizer = # TODO\n",
        "\n",
        "def train_step(x, y):\n",
        "  # Set the model's mode to train\n",
        "  model.train()\n",
        "\n",
        "  # Zero gradients for every step\n",
        "  optimizer.zero_grad()\n",
        "\n",
        "  # Forward pass\n",
        "  '''TODO: feed the current input into the model and generate predictions'''\n",
        "  y_hat = model(x) # TODO\n",
        "  # y_hat = model('''TODO''')\n",
        "\n",
        "  # Compute the loss\n",
        "  '''TODO: compute the loss!'''\n",
        "  loss = compute_loss(y, y_hat) # TODO\n",
        "  # loss = compute_loss('''TODO''', '''TODO''')\n",
        "\n",
        "  # Backward pass\n",
        "  '''TODO: complete the gradient computation and update step.\n",
        "    Remember that in PyTorch there are two steps to the training loop:\n",
        "    1. Backpropagate the loss\n",
        "    2. Update the model parameters using the optimizer\n",
        "  '''\n",
        "  loss.backward() # TODO\n",
        "  optimizer.step() # TODO\n",
        "\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",
        "\n",
        "    # Convert numpy arrays to PyTorch tensors\n",
        "    x_batch = torch.tensor(x_batch, dtype=torch.long).to(device)\n",
        "    y_batch = torch.tensor(y_batch, dtype=torch.long).to(device)\n",
        "\n",
        "    # Take a train step\n",
        "    loss = train_step(x_batch, y_batch)\n",
        "\n",
        "    # Log the loss to the Comet interface\n",
        "    experiment.log_metric(\"loss\", loss.item(), step=iter)\n",
        "\n",
        "    # Update the progress bar and visualize within notebook\n",
        "    history.append(loss.item())\n",
        "    plotter.plot(history)\n",
        "\n",
        "    # Save model checkpoint\n",
        "    if iter % 100 == 0:\n",
        "        torch.save(model.state_dict(), checkpoint_prefix)\n",
        "\n",
        "# Save the final trained model\n",
        "torch.save(model.state_dict(), 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": "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_idx = [char2idx[s] for s in start_string] # TODO\n",
        "  # input_idx = ['''TODO''']\n",
        "  input_idx = torch.tensor([input_idx], dtype=torch.long).to(device)\n",
        "\n",
        "  # Initialize the hidden state\n",
        "  state = model.init_hidden(input_idx.size(0), device)\n",
        "\n",
        "  # Empty string to store our results\n",
        "  text_generated = []\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, state = model(input_idx, state, return_state=True)\n",
        "    # predictions, hidden_state = model('''TODO''', '''TODO''', return_state=True)\n",
        "\n",
        "    # Remove the batch dimension\n",
        "    predictions = predictions.squeeze(0)\n",
        "\n",
        "    '''TODO: use a multinomial distribution to sample over the probabilities'''\n",
        "    input_idx = torch.multinomial(torch.softmax(predictions, dim=-1), num_samples=1)\n",
        "    # input_idx = torch.multinomial('''TODO''', dim=-1), num_samples=1)\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(idx2char[input_idx].item()) # TODO\n",
        "    # text_generated.append('''TODO''')\n",
        "\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(model, start_string=\"X\", generation_length=1000) # TODO\n",
        "# generated_text = generate_text('''TODO''', '''TODO''', '''TODO''')"
      ]
    },
    {
      "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 & Artifacts` 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": "PT_Part2_Music_Generation_Solution.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
}