RNN_TestMIDI.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qzulj1okkrSq"
      },
      "outputs": [],
      "source": [
        "#@title Get a smaller version of the Lakh MIDI Dataset v0.1\n",
        "%%capture\n",
        "!wget http://hog.ee.columbia.edu/craffel/lmd/clean_midi.tar.gz\n",
        "!tar xvf clean_midi.tar.gz\n",
        "!rm clean_midi.tar.gz\n",
        "!rm -rf sample_data\n",
        "\n",
        "dataset_path = \"/content/clean_midi\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-E_cQreQk-8u"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install pretty_midi\n",
        "!pip install miditok"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Installing libraries to hear a MIDI\n",
        "%%capture\n",
        "!apt-get update -qq && apt-get install -y fluidsynth\n",
        "!pip install pretty_midi midi-clip\n",
        "\n",
        "# GS2\n",
        "!gdown 1wlpTIS70nQHMrYBjDT0M6nyg07kUejUv\n",
        "!unzip GeneralUser_GS_v2.0.0--doc_r2.zip\n",
        "!rm -rf GeneralUser_GS_v2.0.0--doc_r2.zip support documentation demo\\ MIDIs\n",
        "!mv GeneralUser\\ GS\\ v2.0.0.sf2 guGS.sf2\n",
        "\n",
        "# PICONICA\n",
        "!gdown 1uk51T9Gvo1n2JRl3_CHCg2FVGWiNI4qJ\n",
        "\n",
        "# Utility library\n",
        "!wget https://raw.githubusercontent.com/roostico/NesGen/refs/heads/main/utility.py\n",
        "!wget https://raw.githubusercontent.com/roostico/NesGen/refs/heads/main/transformer.py\n",
        "\n",
        "from utility import *"
      ],
      "metadata": {
        "id": "NHKsN90VswB-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Import libraries"
      ],
      "metadata": {
        "id": "Ln0a94DmCede"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import random\n",
        "import shutil\n",
        "from tqdm import tqdm\n",
        "from pathlib import Path\n",
        "import pretty_midi\n",
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "from miditok import REMI, TokenizerConfig\n",
        "from utility import playMidi, show_midi_info\n",
        "import json"
      ],
      "metadata": {
        "id": "nmXk2h-HCf42"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Utility functions"
      ],
      "metadata": {
        "id": "rvcrINLqacaF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def play_tokens(tokens: np.ndarray, tokenizer: REMI, delete_after: bool = True, show_info: bool = False):\n",
        "  \"\"\"\n",
        "  Plays the given tokens, decoded using the given tokenizer\n",
        "  \"\"\"\n",
        "  dumped_midi = \"decoded.mid\"\n",
        "  tokenizer.decode([tokens]).dump_midi(dumped_midi)\n",
        "  to_play = playMidi(dumped_midi)\n",
        "  if show_info:\n",
        "    show_midi_info(dumped_midi)\n",
        "  if delete_after:\n",
        "    os.remove(dumped_midi)\n",
        "  return to_play\n",
        "\n",
        "\n",
        "def random_filtered(collection: list, predicate):\n",
        "  \"\"\"\n",
        "  Returns a random element from a collection that satisfies the given predicate.\n",
        "  If no element satisfies the filter, returns None.\n",
        "  \"\"\"\n",
        "  for elem in random.sample(collection, len(collection)):\n",
        "    if predicate(elem):\n",
        "      return elem\n",
        "\n",
        "def lasts_less_than(midi_path: str, time_seconds: int) -> bool:\n",
        "  \"\"\"\n",
        "  Returns true if the last note of the MIDI file is less than the given time in seconds.\n",
        "  \"\"\"\n",
        "  return pretty_midi.PrettyMIDI(midi_path).get_end_time() <= time_seconds"
      ],
      "metadata": {
        "id": "ExvDUjp7aeXD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Move files and rename them"
      ],
      "metadata": {
        "id": "AEgwypV4k_-t"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a6elj9pecSnQ"
      },
      "outputs": [],
      "source": [
        "from pathlib import Path\n",
        "import os\n",
        "import shutil\n",
        "import random\n",
        "\n",
        "# Paths to the files of the dataset\n",
        "midi_paths = list(Path(dataset_path).resolve().glob(\"**/*.mid\"))\n",
        "\n",
        "midis_dir = \"midis\"\n",
        "os.makedirs(midis_dir, exist_ok=True)\n",
        "\n",
        "for i, midi_path in enumerate(midi_paths):\n",
        "  new_midi_path = os.path.join(midis_dir, f\"{i}.mid\")\n",
        "  shutil.move(str(midi_path), new_midi_path)\n",
        "\n",
        "midis = list(Path(\"midis\").resolve().glob(\"**/*.mid\"))\n",
        "\n",
        "def sample():\n",
        "  return str(random.choice(midis))"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Select a sample of these files"
      ],
      "metadata": {
        "id": "h2Y-PIyzlHmM"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Define a filtering function"
      ],
      "metadata": {
        "id": "Mas7DxYj991I"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def is_valid(file: str) -> bool:\n",
        "  \"\"\"Checks if a MIDI file is valid. If any of its instruments has no name,\n",
        "  it is invalid.\n",
        "\n",
        "  Args:\n",
        "      file: The path to the MIDI file.\n",
        "\n",
        "  Returns:\n",
        "      True if the MIDI file is valid, False otherwise.\n",
        "  \"\"\"\n",
        "  try:\n",
        "    midi = pretty_midi.PrettyMIDI(file)\n",
        "    if any([len(instrument.name) == 0 for instrument in midi.instruments]):\n",
        "      return False\n",
        "    return True\n",
        "  except Exception as e:\n",
        "    print(e)\n",
        "    return False"
      ],
      "metadata": {
        "id": "sSMgrOb49FPA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Move files"
      ],
      "metadata": {
        "id": "9vJmyA4g-CkV"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LxkoLRvYcmXj"
      },
      "outputs": [],
      "source": [
        "def copy_random_files(source_dir: str, dest_dir: str, num_files: int, is_file_valid) -> list:\n",
        "  \"\"\"Copies a specified number of random files from a source directory to a destination directory.\n",
        "\n",
        "  Args:\n",
        "      source_dir: The path to the source directory.\n",
        "      dest_dir: The path to the destination directory.\n",
        "      num_files: The number of files to move.\n",
        "  \"\"\"\n",
        "  if not os.path.exists(source_dir):\n",
        "    print(f\"Error: Source directory '{source_dir}' not found.\")\n",
        "    return\n",
        "\n",
        "  if os.path.exists(dest_dir):\n",
        "    shutil.rmtree(dest_dir)\n",
        "\n",
        "  os.makedirs(dest_dir, exist_ok=True)\n",
        "  files = [f for f in os.listdir(source_dir) if os.path.isfile(os.path.join(source_dir, f))]\n",
        "\n",
        "  if len(files) < num_files:\n",
        "    print(f\"Warning: Only {len(files)} files found in '{source_dir}'. Moving all of them.\")\n",
        "    num_files = len(files)\n",
        "\n",
        "  files_to_move = []\n",
        "  i = 0\n",
        "  with tqdm(total=num_files, position=0, leave=True) as pbar:\n",
        "    while i < num_files:\n",
        "      random_file = random.choice(files)\n",
        "      if is_file_valid(os.path.join(source_dir, random_file)):\n",
        "        files_to_move.append(random_file)\n",
        "        i = i + 1\n",
        "        pbar.update()\n",
        "      else:\n",
        "        continue\n",
        "\n",
        "  result = []\n",
        "  for file in files_to_move:\n",
        "    source_path = os.path.join(source_dir, file)\n",
        "    dest_path = os.path.join(dest_dir, file)\n",
        "    shutil.copy(source_path, dest_path)\n",
        "    result.append(dest_path)\n",
        "  return result\n",
        "\n",
        "source_directory = \"midis\"\n",
        "destination_directory = \"selected\"\n",
        "number_of_files_to_move = 500\n",
        "\n",
        "sample = copy_random_files(source_directory, destination_directory, number_of_files_to_move, is_valid)\n",
        "assert len(sample) == 500"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Listen one random MIDI from sample"
      ],
      "metadata": {
        "id": "fmvUMw52BdxU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def valid_example(file: str) -> float:\n",
        "  \"\"\"\n",
        "  Returns true if the midi file lasts less that 120 seconds and has more than 1 instrument\n",
        "  \"\"\"\n",
        "  return lasts_less_than(file, 120) and len(pretty_midi.PrettyMIDI(file).instruments) > 1\n",
        "\n",
        "example = random_filtered(sample, valid_example)\n",
        "\n",
        "print(f\"Showing file {example}\")\n",
        "show_midi_info(example)\n",
        "playMidi(example)"
      ],
      "metadata": {
        "id": "B7sfptC6Bhan"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Pre-processing"
      ],
      "metadata": {
        "id": "RA677BGpNMNC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "TRACK_MIN_DENSITY_PERC = 0.20\n",
        "\n",
        "def is_excluded(track_name):\n",
        "  \"\"\"\n",
        "  Exclusion criteria for MIDI tracks\n",
        "  \"\"\"\n",
        "  exclude_keywords = [\"drum\", \"effect\", \"percussion\"]\n",
        "  return any(keyword.lower() in track_name.lower() for keyword in exclude_keywords)\n",
        "\n",
        "def merge_tracks_to_single_instrument(input_file, output_file, target_channel=0):\n",
        "    midi_data = pretty_midi.PrettyMIDI(input_file)\n",
        "    merged_midi = pretty_midi.PrettyMIDI()\n",
        "    merged_instrument = pretty_midi.Instrument(program=pretty_midi.instrument_name_to_program('Acoustic Grand Piano'), is_drum=False)\n",
        "\n",
        "    mean_notes = map(lambda x: len(x.notes), midi_data.instruments)\n",
        "    mean_notes = sum(mean_notes) / len(midi_data.instruments)\n",
        "\n",
        "    for instrument in midi_data.instruments:\n",
        "        track_name = instrument.name\n",
        "\n",
        "        # Exclude drum instruments or effect instruments\n",
        "        if instrument.is_drum or is_excluded(track_name):\n",
        "            continue\n",
        "\n",
        "        # Exclude instruments that have a low number of notes\n",
        "        if len(instrument.notes) / mean_notes < TRACK_MIN_DENSITY_PERC:\n",
        "            continue\n",
        "\n",
        "        for note in instrument.notes:\n",
        "            note.velocity = max(1, note.velocity)  # Ensure velocity is within MIDI range\n",
        "            merged_instrument.notes.append(note)\n",
        "\n",
        "    merged_instrument.notes.sort(key=lambda note: note.start)\n",
        "\n",
        "    tempo_times, tempi = midi_data.get_tempo_changes()\n",
        "    if len(tempi) > 0:\n",
        "        merged_midi._tick_scales = midi_data._tick_scales  # Copy tempo-related timing\n",
        "\n",
        "    merged_midi.instruments.append(merged_instrument)\n",
        "    merged_midi.time_signature_changes = midi_data.time_signature_changes\n",
        "    merged_midi.key_signature_changes = midi_data.key_signature_changes\n",
        "\n",
        "    merged_midi.write(output_file)\n",
        "\n",
        "if os.path.exists(\"pre-processed\"):\n",
        "  shutil.rmtree(\"pre-processed\")\n",
        "\n",
        "os.makedirs(\"pre-processed\")\n",
        "\n",
        "for file in tqdm(sample):\n",
        "  try:\n",
        "    merge_tracks_to_single_instrument(file, f\"pre-processed/{os.path.basename(file)}\")\n",
        "  except Exception as e:\n",
        "    print(f\"There was an error: {e}\")\n",
        "    continue\n",
        "\n",
        "processed = list(Path(\"pre-processed\").resolve().glob(\"**/*.mid\"))"
      ],
      "metadata": {
        "id": "8bRjAI8HNOA5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Listen the same file from before, but with merged MIDI tracks"
      ],
      "metadata": {
        "id": "bqUCP3U-HEa2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "processed_example = os.path.join(\"pre-processed\", os.path.basename(example))\n",
        "print(f\"Showing file {processed_example}\")\n",
        "show_midi_info(processed_example)\n",
        "playMidi(processed_example)"
      ],
      "metadata": {
        "id": "_zFSNP-zHQOX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Tokenizer\n",
        "\n"
      ],
      "metadata": {
        "id": "X_qYrXTErrsH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tok_config = {\n",
        "    \"use_pitchdrum_tokens\": False\n",
        "}\n",
        "\n",
        "tok_config = TokenizerConfig(**tok_config)\n",
        "tokenizer = REMI(tok_config)"
      ],
      "metadata": {
        "id": "fM44SXMEHNmo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### (Optional): train the tokenizer"
      ],
      "metadata": {
        "id": "VqhBEeOGHUZB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer.train(vocab_size=1000, files_paths=processed)"
      ],
      "metadata": {
        "id": "ELI4Z2GhHZFB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Tokenizer the dataset"
      ],
      "metadata": {
        "id": "CGr1Z52HHaaI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def midi_valid(midi) -> bool:\n",
        "    if any(ts.numerator != 4 for ts in midi.time_signature_changes):\n",
        "        return False  # time signature different from 4/*, 4 beats per bar\n",
        "    return True\n",
        "\n",
        "if os.path.exists(\"tokenized\"):\n",
        "  shutil.rmtree(\"tokenized\")\n",
        "\n",
        "tokenizer.tokenize_dataset(        # 2 velocity and 1 duration values\n",
        "    Path(\"/content\", \"pre-processed\"),\n",
        "    Path(\"/content\", \"tokenized\"),\n",
        "    midi_valid,\n",
        ")"
      ],
      "metadata": {
        "id": "uB8MAhdvHGEB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Utility function to read a JSON tokenized file"
      ],
      "metadata": {
        "id": "_txg7WlqOoXf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def read_json(path: str) -> dict:\n",
        "  with open(path, \"r\") as f:\n",
        "    return json.load(f)"
      ],
      "metadata": {
        "id": "bZzHxEPeOrnU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### See the tokenized version of the previous file"
      ],
      "metadata": {
        "id": "K8sNToqWOcs2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tokenized_example = os.path.join(\"tokenized\", Path(example).stem + \".json\")\n",
        "example_ids = read_json(tokenized_example)[\"ids\"][0]\n",
        "print(f\"Showing IDS of {tokenized_example}\")\n",
        "print(np.array(example_ids))"
      ],
      "metadata": {
        "id": "ZVpb2XnBOgSC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Read the tokenized version of files from the JSON"
      ],
      "metadata": {
        "id": "ceLrlILzQg7v"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def read_json_files(json_file_paths):\n",
        "    \"\"\"Reads a list of JSON files and returns a list of objects.\n",
        "\n",
        "    Args:\n",
        "        json_file_paths: A list of file paths to JSON files.\n",
        "\n",
        "    Returns:\n",
        "        A list of objects, where each object represents the data from a JSON file.\n",
        "        Returns an empty list if any error occurs during file processing.\n",
        "    \"\"\"\n",
        "    objects = []\n",
        "    for file_path in tqdm(json_file_paths):\n",
        "        try:\n",
        "            objects.append(read_json(file_path))\n",
        "        except FileNotFoundError:\n",
        "            print(f\"Error: File not found - {file_path}\")\n",
        "            return [] # Return empty list on error\n",
        "        except json.JSONDecodeError:\n",
        "            print(f\"Error decoding JSON in file: {file_path}\")\n",
        "            return [] # Return empty list on error\n",
        "    return objects\n",
        "\n",
        "# Example usage (assuming 'tokenized' directory contains JSON files):\n",
        "tokenized_files = list(Path(\"/content\", \"tokenized\").resolve().glob(\"**/*.json\"))\n",
        "data_objects = read_json_files(tokenized_files)\n",
        "\n",
        "if data_objects:\n",
        "    print(f\"\\nSuccessfully read {len(data_objects)} JSON files.\")\n",
        "    # Now you can work with the 'data_objects' list\n",
        "    # For example, print the first object:\n",
        "    # print(data_objects[0])\n",
        "else:\n",
        "    print(\"Error reading JSON files.\")"
      ],
      "metadata": {
        "id": "iR0bPrk-H1O1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Create the list of tokenized songs, taking the IDs of each one"
      ],
      "metadata": {
        "id": "P9nbkZgfQllX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "encoded = [np.array(song[\"ids\"][0]) for song in data_objects]"
      ],
      "metadata": {
        "id": "_B81lpcTKL6m"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Listen the same example from before, decoded from its tokenization"
      ],
      "metadata": {
        "id": "Cg9YYbdlSpzI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(f\"Showing decoded IDS of {tokenized_example}\")\n",
        "\n",
        "to_play = play_tokens(example_ids, tokenizer, show_info = True)\n",
        "to_play"
      ],
      "metadata": {
        "id": "EZU2w5GtSxF0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Trim of initial and ending silence in each song"
      ],
      "metadata": {
        "id": "BsY64PWJQ7Uh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def trim(ids: np.ndarray, token_to_remove: int) -> np.ndarray:\n",
        "  \"\"\"\n",
        "  Returns a new numpy array with initial and ending `token_to_remove` removed.\n",
        "  \"\"\"\n",
        "  start_idx = 0\n",
        "  end_idx = len(ids)\n",
        "\n",
        "  while start_idx < len(ids) and ids[start_idx] == token_to_remove:\n",
        "      start_idx += 1\n",
        "  while end_idx > start_idx and ids[end_idx - 1] == token_to_remove:\n",
        "      end_idx -= 1\n",
        "  return ids[start_idx:end_idx]\n",
        "\n",
        "bar_token = tokenizer.vocab[\"Bar_None\"]\n",
        "encoded = [trim(ids, bar_token) for ids in encoded]"
      ],
      "metadata": {
        "id": "wNAB9oeoRBwE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Remove long rests in the song"
      ],
      "metadata": {
        "id": "E_KK0LtlZoMJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def shorten_sequences(arr: np.ndarray, target_value: int, max_length: int) -> np.ndarray:\n",
        "  result = []\n",
        "  current_sequence = []\n",
        "\n",
        "  for value in arr:\n",
        "    if value == target_value:\n",
        "      current_sequence.append(value)\n",
        "    else:\n",
        "      if len(current_sequence) > max_length:\n",
        "        result.extend([target_value] * max_length)\n",
        "      else:\n",
        "        result.extend(current_sequence)\n",
        "\n",
        "        current_sequence = []\n",
        "        result.append(value)\n",
        "\n",
        "  if len(current_sequence) > max_length:\n",
        "    result.extend([target_value] * max_length)\n",
        "  else:\n",
        "    result.extend(current_sequence)\n",
        "\n",
        "  return np.array(result)\n",
        "\n",
        "bar_token = tokenizer.vocab[\"Bar_None\"]\n",
        "max_rest_length = 20\n",
        "encoded = [shorten_sequences(ids, bar_token, max_rest_length) for ids in encoded]"
      ],
      "metadata": {
        "id": "D_iZoc74Zr_Y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Show an example from the resulting IDs"
      ],
      "metadata": {
        "id": "W2YwfFIKZdg0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "random_ids = random.choice(encoded)[:1000]\n",
        "\n",
        "print(\"Decoded:\")\n",
        "play_tokens(random_ids, tokenizer, show_info = True)"
      ],
      "metadata": {
        "id": "0PTftO6OZidw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "End of pre-processing, proceding with data and model preparation with Tensorflow\n",
        "\n",
        "---\n",
        "\n"
      ],
      "metadata": {
        "id": "VmJjxVpaQr13"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Tensorflow data and model setup"
      ],
      "metadata": {
        "id": "MFobYqPQQ2mL"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Creating a Tensorflow dataset with all IDs"
      ],
      "metadata": {
        "id": "PZYhXakurnaB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "\n",
        "all_ids = np.concatenate(encoded)\n",
        "ids_dataset = tf.data.Dataset.from_tensor_slices(all_ids)"
      ],
      "metadata": {
        "id": "0yIWzFmmowXC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Convert into sequences"
      ],
      "metadata": {
        "id": "gKA8lKB5srKw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "seq_length = 400 #@param {type: 'slider', max: 500, min: 50, step: 50}\n",
        "sequences = ids_dataset.batch(seq_length+1, drop_remainder=True)"
      ],
      "metadata": {
        "id": "X33uqO3KrWNA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Preparing labels"
      ],
      "metadata": {
        "id": "ZnJUygt7uUFi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def split_input_target(sequence):\n",
        "    input_seq = sequence[:-1]\n",
        "    input_mask = np.ones(seq_length)\n",
        "    target_seq = sequence[1:]\n",
        "    return (input_seq, input_mask), target_seq\n",
        "\n",
        "dataset = sequences.map(split_input_target)"
      ],
      "metadata": {
        "id": "Xlsk46y3sVRS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Creating training batches"
      ],
      "metadata": {
        "id": "o6Ssr4f_um6f"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Batch size\n",
        "BATCH_SIZE = 128 #@param {type: 'slider', max: 256, min: 32, step: 32}\n",
        "\n",
        "# Buffer size to shuffle the dataset\n",
        "# (TF data is designed to work with possibly infinite sequences,\n",
        "# so it doesn't attempt to shuffle the entire sequence in memory. Instead,\n",
        "# it maintains a buffer in which it shuffles elements).\n",
        "BUFFER_SIZE = 10000\n",
        "\n",
        "dataset = (\n",
        "    dataset\n",
        "    .shuffle(BUFFER_SIZE)\n",
        "    .batch(BATCH_SIZE, drop_remainder=True)\n",
        "    .prefetch(tf.data.experimental.AUTOTUNE))\n",
        "\n",
        "dataset"
      ],
      "metadata": {
        "id": "G1nUYxujuorC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Splitting in Train, Validation and Test"
      ],
      "metadata": {
        "id": "sryXDHvk6ElG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "ds_size = dataset.cardinality().numpy()\n",
        "train_size = int(0.8 * ds_size)\n",
        "val_size = int(0.1 * ds_size)\n",
        "test_size = int(0.1 * ds_size)"
      ],
      "metadata": {
        "id": "pRUgrcIV6HPI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_ds = dataset.take(train_size)\n",
        "remaining = dataset.skip(train_size)\n",
        "valid_ds = remaining.take(val_size)\n",
        "test_ds = remaining.skip(val_size)"
      ],
      "metadata": {
        "id": "a6u4odFU6dhT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Building the model"
      ],
      "metadata": {
        "id": "13TNnkIRuwKa"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Transformer"
      ],
      "metadata": {
        "id": "Hj1_6V3FCVXX"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Hyper-parameters"
      ],
      "metadata": {
        "id": "Wr6h6007A7K-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "num_layers = 4\n",
        "d_model = 128\n",
        "dff = 256\n",
        "num_heads = 2\n",
        "dropout_rate = 0.1"
      ],
      "metadata": {
        "id": "INYXNRtkA5zd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from transformer import Transformer\n",
        "\n",
        "model = Transformer(\n",
        "    num_layers=num_layers,\n",
        "    d_model=d_model,\n",
        "    num_heads=num_heads,\n",
        "    dff=dff,\n",
        "    input_vocab_size=tokenizer.vocab_size,\n",
        "    target_vocab_size=tokenizer.vocab_size,\n",
        "    dropout_rate=dropout_rate\n",
        "    )"
      ],
      "metadata": {
        "id": "EKXVOp2v7EyL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "loss = tf.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
        "model.compile(optimizer='adam',loss=loss, metrics=['accuracy'])"
      ],
      "metadata": {
        "id": "UJlS7YtFBM0a"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "---"
      ],
      "metadata": {
        "id": "YW9ZJXh4Buwh"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Try generating with un-trained model"
      ],
      "metadata": {
        "id": "Cv7uUEpv3LUn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from utility import playMidi\n",
        "\n",
        "# Use the model\n",
        "for input_example_batch, target_example_batch in dataset.take(1):\n",
        "  example_batch_predictions = model(input_example_batch)\n",
        "  print(example_batch_predictions.shape, \"# (batch_size, sequence_length, vocab_size)\")\n",
        "\n",
        "# Sample indices from predictions\n",
        "sampled_indices = tf.random.categorical(example_batch_predictions[0], num_samples=1)\n",
        "sampled_indices = tf.squeeze(sampled_indices, axis=-1).numpy()\n",
        "\n",
        "# Decode\n",
        "midi = tokenizer.decode([sampled_indices])\n",
        "midi.dump_midi(\"example.mid\")\n",
        "playMidi(\"example.mid\")"
      ],
      "metadata": {
        "id": "2L95vK8M0grD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Calculate the loss\n",
        "#loss = tf.keras.losses.sparse_categorical_crossentropy(target_example_batch, example_batch_predictions, from_logits=True)\n",
        "# Reduce mean to get a single scalar loss value\n",
        "#loss = tf.reduce_mean(loss)\n",
        "\n",
        "#print(\"Loss:\", loss.numpy())"
      ],
      "metadata": {
        "id": "Bgq7J1sTy-Jx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Taken from tensorflow tutorial:\n",
        "\n",
        "A newly initialized model shouldn't be too sure of itself, the output logits should all have similar magnitudes. To confirm this you can check that the exponential of the mean loss is approximately equal to the vocabulary size. A much higher loss means the model is sure of its wrong answers, and is badly initialized:"
      ],
      "metadata": {
        "id": "od6NsuqV8y2f"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#print(\"Checking if it is near to vocabulary size\")\n",
        "#print(tf.exp(example_batch_mean_loss).numpy())\n",
        "#print(\"Vocab size: \", encoding.vocab_size)"
      ],
      "metadata": {
        "id": "MXnmou978tLu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training"
      ],
      "metadata": {
        "id": "CdAbc4nvy_Ub"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "loss = tf.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
        "model.compile(optimizer='adam', loss=loss, metrics=['accuracy'])"
      ],
      "metadata": {
        "id": "e9hwzFK6zDYw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Directory where the checkpoints will be saved\n",
        "checkpoint_dir = './training_checkpoints'\n",
        "# Name of the checkpoint files\n",
        "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt_{epoch}.weights.h5\")\n",
        "early_stopping = tf.keras.callbacks.EarlyStopping(\n",
        "    monitor='val_loss',\n",
        "    patience=3,\n",
        "    restore_best_weights=True\n",
        ")\n",
        "\n",
        "checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(\n",
        "    filepath=checkpoint_prefix,\n",
        "    save_weights_only=True\n",
        ")"
      ],
      "metadata": {
        "id": "lADTb-9GzVsO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "EPOCHS = 5\n",
        "\n",
        "history = model.fit(train_ds,\n",
        "                    epochs=EPOCHS,\n",
        "                    validation_data=valid_ds,\n",
        "                    callbacks=[checkpoint_callback, early_stopping]\n",
        "                    )"
      ],
      "metadata": {
        "id": "4EWiw3RYzo_5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Generation"
      ],
      "metadata": {
        "id": "wQp_3lVX-wLK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class OneStep(tf.keras.Model):\n",
        "  def __init__(self, model, decoding, encoding, vocab_size, temperature=1.0):\n",
        "    super().__init__()\n",
        "    self.temperature = temperature\n",
        "    self.model = model\n",
        "    self.decode = decoding\n",
        "    self.encode = encoding\n",
        "\n",
        "    # Taken from tensorflow tutorial: useful to skip ids\n",
        "\n",
        "    #skip_ids = self.ids_from_chars(['[UNK]'])[:, None]\n",
        "    #sparse_mask = tf.SparseTensor(\n",
        "        # Put a -inf at each bad index.\n",
        "    #    values=[-float('inf')]*len(skip_ids),\n",
        "    #    indices=skip_ids,\n",
        "        # Match the shape to the vocabulary\n",
        "    #    dense_shape=[vocab_size])\n",
        "    #self.prediction_mask = tf.sparse.to_dense(sparse_mask)\n",
        "\n",
        "  @tf.function\n",
        "  def generate_one_step(self, input_ids):\n",
        "    # Ensure input_ids is 2D: (batch_size, sequence_length)\n",
        "    if len(tf.shape(input_ids)) == 1:\n",
        "        input_ids = tf.expand_dims(input_ids, axis=0)\n",
        "\n",
        "    # Create a context tensor with the same shape as input_ids\n",
        "    context = tf.ones_like(input_ids)\n",
        "\n",
        "    # Package the context and input_ids as a tuple\n",
        "    inputs = (context, input_ids)\n",
        "\n",
        "    # Run the model with the tuple (context, input_ids)\n",
        "    predicted_logits = self.model(inputs=inputs)\n",
        "\n",
        "    # Take only the last time step’s predictions\n",
        "    predicted_logits = predicted_logits[:, -1, :]\n",
        "    predicted_logits = predicted_logits / self.temperature\n",
        "\n",
        "    # Sample from the output logits to generate token IDs\n",
        "    predicted_ids = tf.random.categorical(predicted_logits, num_samples=1)\n",
        "    predicted_ids = tf.squeeze(predicted_ids, axis=-1)\n",
        "\n",
        "    return predicted_ids"
      ],
      "metadata": {
        "id": "_2z0FLUK-yYf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "one_step_model = OneStep(model, tokenizer.decode, tokenizer.encode, tokenizer.vocab_size)"
      ],
      "metadata": {
        "id": "bgXr-8m9AWVB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_random_file() -> str:\n",
        "  return random.choice(midis)\n",
        "\n",
        "seed_ids = np.array(tokenizer.encode(get_random_file())[0].ids)\n",
        "seed_ids = seed_ids[:seq_length]"
      ],
      "metadata": {
        "id": "yN7j8IoUCQxT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tf.config.run_functions_eagerly(True)\n",
        "one_step_model.generate_one_step(seed_ids)"
      ],
      "metadata": {
        "id": "lg1mzRUWAMDp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import time\n",
        "\n",
        "start = time.time()\n",
        "next_ids = seed_ids\n",
        "result = [next_ids]\n",
        "\n",
        "for n in range(1000):\n",
        "  next_ids = one_step_model.generate_one_step(next_ids)\n",
        "  result.append(next_ids)\n",
        "\n",
        "end = time.time()\n",
        "print('\\nRun time:', end - start, \"\\n\", '_'*80, \"\\n\")\n",
        "result = np.concatenate(result[1:])\n",
        "print(\"Shape of result: \", result.shape)"
      ],
      "metadata": {
        "id": "ND5aYRs-Agfi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Save the generator"
      ],
      "metadata": {
        "id": "hC3N64vLT8gm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tf.saved_model.save(one_step_model, 'one_step')"
      ],
      "metadata": {
        "id": "89POLl6ETDLu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Reload the generator"
      ],
      "metadata": {
        "id": "PBcd9qc1VmiR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "one_step_reloaded = tf.saved_model.load('one_step')"
      ],
      "metadata": {
        "id": "MFxLMn6TVolF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def tokens_from_ids(ids, tokenizer):\n",
        "  tokens = []\n",
        "  for id in ids:\n",
        "    for key, value in tokenizer.vocab.items():\n",
        "      if value == id:\n",
        "        tokens.append(key)\n",
        "  return np.array(tokens)"
      ],
      "metadata": {
        "id": "sgnRbYJ228eq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "example = train_ds.take(1).as_numpy_iterator().next()[0]"
      ],
      "metadata": {
        "id": "oGMKVG8H40xj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokens_from_ids(example[0], tokenizer)"
      ],
      "metadata": {
        "id": "wWETYlrn4sl4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokens_from_ids(result, tokenizer)"
      ],
      "metadata": {
        "id": "URAWsmhN5Zbp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Hear the result"
      ],
      "metadata": {
        "id": "6BuVmTnKT_xl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from utility import playMidi\n",
        "\n",
        "midi = tokenizer.decode([result])\n",
        "midi.dump_midi(\"result.mid\")\n",
        "playMidi(\"result.mid\")"
      ],
      "metadata": {
        "id": "w8tRExb_T-f4"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}