public_ATM_Qwen2Audio (1).ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Iux87x6p6U3x"
      },
      "source": [
        "## 1. Install dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9p6Im3pcNx8d"
      },
      "outputs": [],
      "source": [
        "!pip install librosa pretty_midi wandb jams datasets huggingface_hub"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YzEHC7fsFj3-"
      },
      "outputs": [],
      "source": [
        "!apt-get update\n",
        "!apt-get install -y abcmidi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mSyIXpSuFkVf"
      },
      "outputs": [],
      "source": [
        "!abc2midi --version"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_lEZkuELN3cm"
      },
      "source": [
        "## Data"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**I have these datasets stored in my Google Drive**"
      ],
      "metadata": {
        "id": "kynvI_fLKX_3"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Prl3Ubd8UWji"
      },
      "source": [
        "### Download Maestro"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MYDfZB2HB3bI"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "!cp -r \"/content/drive/My Drive/maestro/\" \"/content/maestro/\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pf26424fUBaf"
      },
      "source": [
        "### Download URMP"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1P9oEidKyonL"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "!cp \"/content/drive/My Drive/automatic-music-transcription/URMP_Dataset.tar.gz\" \"/content/URMP_Dataset.tar.gz\"\n",
        "!mkdir -p /content/URMP_Dataset\n",
        "!tar -xzvf /content/URMP_Dataset.tar.gz -C /content/URMP_Dataset\n",
        "!rm /content/URMP_Dataset.tar.gz"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rb9eI2MZgknL"
      },
      "source": [
        "## MIDI to ABC notation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iEV71PEpZiZ8"
      },
      "outputs": [],
      "source": [
        "import subprocess\n",
        "from pathlib import Path\n",
        "from typing import Optional\n",
        "\n",
        "def remove_backslash_after_bar(text):\n",
        "    # Replace occurrences of \"| \\\" with just \"|\"\n",
        "    modified_text = text.replace(\"| \\\\\", \"| \")\n",
        "    return modified_text\n",
        "\n",
        "def cmd_midi_to_abc(\n",
        "    midi_path: str,\n",
        "    extract_all_tracks: bool = False,\n",
        "    verbose: bool = False,\n",
        "    min_note_length: Optional[int] = None\n",
        ") -> Optional[str]:\n",
        "    if not Path(midi_path).is_file():\n",
        "        raise FileNotFoundError(f\"MIDI file not found: {midi_path}\")\n",
        "\n",
        "    cmd = ['midi2abc', midi_path]\n",
        "    if extract_all_tracks:\n",
        "        cmd.append('-xa')\n",
        "    if verbose:\n",
        "        cmd.append('-v')\n",
        "    if min_note_length is not None:\n",
        "        cmd.extend(['-mpl', str(min_note_length)])\n",
        "\n",
        "    try:\n",
        "        result = subprocess.run(cmd, capture_output=True, text=True, check=True)\n",
        "        return result.stdout\n",
        "    except subprocess.CalledProcessError as e:\n",
        "        print(f\"Error converting MIDI to ABC: {e}\")\n",
        "        print(f\"Error output: {e.stderr}\")\n",
        "        return None\n",
        "\n",
        "def remove_key_signature(abc_string):\n",
        "    # Dictionary of key signatures and their affected notes (capitals for base notes)\n",
        "    key_signatures = {\n",
        "        # Major keys\n",
        "        'Cb': ['B-', 'E-', 'A-', 'D-', 'G-', 'C-', 'F-'],\n",
        "        'Gb': ['B-', 'E-', 'A-', 'D-', 'G-', 'C-'],\n",
        "        'Db': ['B-', 'E-', 'A-', 'D-', 'G-'],\n",
        "        'Ab': ['B-', 'E-', 'A-', 'D-'],\n",
        "        'Eb': ['B-', 'E-', 'A-'],\n",
        "        'Bb': ['B-', 'E-'],\n",
        "        'F': ['B-'],\n",
        "        'C': [],\n",
        "        'G': ['F'],\n",
        "        'D': ['F', 'C'],\n",
        "        'A': ['F', 'C', 'G'],\n",
        "        'E': ['F', 'C', 'G', 'D'],\n",
        "        'B': ['F', 'C', 'G', 'D', 'A'],\n",
        "        'F#': ['F', 'C', 'G', 'D', 'A', 'E'],\n",
        "        'C#': ['F', 'C', 'G', 'D', 'A', 'E', 'B'],\n",
        "        # Minor keys\n",
        "        'Abmin': ['B-', 'E-', 'A-', 'D-'],\n",
        "        'Ebmin': ['B-', 'E-', 'A-'],\n",
        "        'Bbmin': ['B-', 'E-'],\n",
        "        'Fmin': ['B-', 'E-', 'A-', 'D-'],\n",
        "        'Cmin': ['B-', 'E-', 'A-'],\n",
        "        'Gmin': ['B-', 'E-'],\n",
        "        'Dmin': ['B-'],\n",
        "        'Amin': [],\n",
        "        'Emin': ['F'],\n",
        "        'Bmin': ['F', 'C'],\n",
        "        'F#min': ['F', 'C', 'G'],\n",
        "        'C#min': ['F', 'C', 'G', 'D'],\n",
        "        'G#min': ['F', 'C', 'G', 'D', 'A']\n",
        "    }\n",
        "\n",
        "    # Find and store the key\n",
        "    lines = abc_string.split('\\n')\n",
        "    current_key = None\n",
        "    key_line_index = None\n",
        "\n",
        "    for i, line in enumerate(lines):\n",
        "        if line.startswith('K:'):\n",
        "            current_key = line.split(':')[1].strip().split('%')[0].strip()  # Get the key without comments\n",
        "            key_line_index = i\n",
        "            break\n",
        "\n",
        "    if not current_key or current_key == 'none':\n",
        "        return abc_string\n",
        "\n",
        "    # Get the affected notes for this key\n",
        "    affected_notes = key_signatures.get(current_key, [])\n",
        "\n",
        "    # Replace the key line with \"none\"\n",
        "    lines[key_line_index] = 'K: none'\n",
        "\n",
        "    # Process the music lines\n",
        "    for i in range(len(lines)):\n",
        "        if not any(lines[i].startswith(x) for x in ['M:', 'L:', 'K:', 'Q:', '%', 'V:', '%%']):\n",
        "            line = lines[i]\n",
        "            new_line = ''\n",
        "            j = 0\n",
        "            while j < len(line):\n",
        "                if j+1 < len(line) and line[j] == '=' and line[j+1].isalpha():\n",
        "                    # For naturals, just keep the note without the natural sign\n",
        "                    new_line += line[j+1]\n",
        "                    j += 2\n",
        "                elif line[j].isalpha():\n",
        "                    note = line[j].upper()\n",
        "                    # Check if this note is affected by key signature and doesn't already have an accidental\n",
        "                    if note in [n[0] for n in affected_notes]:\n",
        "                        # Only add accidental if there isn't one already\n",
        "                        if j == 0 or not (line[j-1] in ['^', '_', '=']):\n",
        "                            # Add sharp (^) or flat (_) based on key signature\n",
        "                            if affected_notes[0][-1] == '-':  # if it's a flat key\n",
        "                                new_line += '_'\n",
        "                            else:  # if it's a sharp key\n",
        "                                new_line += '^'\n",
        "                    new_line += line[j]\n",
        "                    j += 1\n",
        "                else:\n",
        "                    new_line += line[j]\n",
        "                    j += 1\n",
        "            lines[i] = new_line\n",
        "\n",
        "    return '\\n'.join(lines)\n",
        "\n",
        "def remove_comment_lines(abc_string):\n",
        "    # Split the input string into lines\n",
        "    lines = abc_string.splitlines()\n",
        "    # Filter out lines that start with \"%\" or \"T: \"\n",
        "    non_comment_lines = [line for line in lines if not line.strip().startswith('%') and not line.strip().startswith(\"T: \")]\n",
        "    # Join the remaining lines back into a single string\n",
        "    return \"\\n\".join(non_comment_lines)\n",
        "\n",
        "abc_questions = [\n",
        "    \"\"\"Transcribe this music clip using ABC notation. Follow this template:\n",
        "\n",
        "```\n",
        "M: 4/4\n",
        "L: 1/16\n",
        "Q:1/4=120\n",
        "K: none\n",
        "\n",
        "V:1 name=\"Instrument 1\"\n",
        "(notes here)\n",
        "\n",
        "(Add more voices if present)\n",
        "```\n",
        "\n",
        "Key points for ABC notation:\n",
        "- Notes: A-G (lowercase for higher octaves)\n",
        "- Accidentals: ^ (sharp), _ (flat)\n",
        "- Note length: Numbers after note (C2 = twice as long as C)\n",
        "- Dotted notes: . after note (C.)\n",
        "- Rests: z with optional duration (z2 = half rest)\n",
        "- Chords: [CEG]\n",
        "- Ties: -\n",
        "- Bar lines: |\n",
        "- Broken rhythms: > or < between notes (C>D = dotted-C eighth + D sixteenth)\n",
        "\n",
        "Important:\n",
        "- The key should always be set to \"none\" as this is a short clip and it's hard to identify a key.\n",
        "- Always explicitly show accidentals as the key is set to none.\n",
        "- Include a separate voice (V:) for each distinct instrument you can identify.\n",
        "- List the instrument names using one of the 128 General MIDI instrument names.\n",
        "    \"\"\"\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZCnG-tRd-TPb"
      },
      "outputs": [],
      "source": [
        "from typing import Tuple, Dict\n",
        "from pretty_midi import pretty_midi\n",
        "from mido import MidiFile, MidiTrack, Message, MetaMessage, bpm2tempo\n",
        "\n",
        "# Global cache for MIDI data\n",
        "midi_data_cache: Dict[str, pretty_midi.PrettyMIDI] = {}\n",
        "\n",
        "def get_midi_data(midi_file: str) -> pretty_midi.PrettyMIDI:\n",
        "    if midi_file in midi_data_cache:\n",
        "        return midi_data_cache[midi_file]\n",
        "    else:\n",
        "        midi_data = pretty_midi.PrettyMIDI(midi_file)\n",
        "        midi_data_cache[midi_file] = midi_data\n",
        "        return midi_data\n",
        "\n",
        "def get_tempo_and_metre(midi_file: str, midi_data: pretty_midi.PrettyMIDI, chunk_start_time: float, chunk_end_time: float) -> Tuple[float, int, int]:\n",
        "    tempo_changes_times, tempo_changes_values = midi_data.get_tempo_changes()\n",
        "\n",
        "    initial_tempo = tempo_changes_values[0]\n",
        "    for i, t in enumerate(tempo_changes_times):\n",
        "        if t <= chunk_start_time:\n",
        "            initial_tempo = tempo_changes_values[i]\n",
        "        else:\n",
        "            break\n",
        "\n",
        "    for i, t in enumerate(tempo_changes_times):\n",
        "        if chunk_start_time < t < chunk_end_time:\n",
        "            new_tempo = tempo_changes_values[i]\n",
        "            tempo_change = abs(new_tempo - initial_tempo)\n",
        "\n",
        "            if tempo_change > 10:\n",
        "                raise ValueError(f\"Significant tempo change (>{tempo_change:.2f} BPM) detected within the chunk ({chunk_start_time} - {chunk_end_time})\")\n",
        "\n",
        "    try:\n",
        "        relevant_time_signature = None\n",
        "        for ts in midi_data.time_signature_changes:\n",
        "            if ts.time <= chunk_start_time:\n",
        "                relevant_time_signature = ts\n",
        "            elif ts.time > chunk_start_time:\n",
        "                break\n",
        "\n",
        "        if relevant_time_signature is None:\n",
        "            relevant_time_signature = midi_data.time_signature_changes[0]\n",
        "\n",
        "        changes_within_chunk = [ts for ts in midi_data.time_signature_changes\n",
        "                              if chunk_start_time < ts.time < chunk_end_time]\n",
        "        if changes_within_chunk:\n",
        "            print(f\"Note: Time signature change(s) detected within the chunk ({chunk_start_time} - {chunk_end_time}). Using the first time signature for the entire chunk.\")\n",
        "\n",
        "    except IndexError:\n",
        "        raise ValueError(\"Error: The time_signatures list is empty.\")\n",
        "\n",
        "    return (initial_tempo, relevant_time_signature.numerator, relevant_time_signature.denominator)\n",
        "\n",
        "def get_tempo_and_metre_simp(midi_data: pretty_midi.PrettyMIDI) -> Tuple[float, int, int]:\n",
        "    tempo = midi_data.get_tempo_changes()[1][0]\n",
        "    time_sig = midi_data.time_signature_changes[0]\n",
        "    return (tempo, time_sig.numerator, time_sig.denominator)\n",
        "\n",
        "from mido import MidiFile, MidiTrack, Message, MetaMessage, bpm2tempo\n",
        "from typing import Tuple, Dict\n",
        "from pretty_midi import pretty_midi\n",
        "\n",
        "def extract_midi_segment(input_file: str, output_file: str, start_time: float, duration: float) -> None:\n",
        "    end_time = start_time + duration\n",
        "\n",
        "    # Get cached MIDI data\n",
        "    midi_data = get_midi_data(input_file)\n",
        "\n",
        "    tempo_bpm, numerator, denominator = get_tempo_and_metre(input_file, midi_data, start_time, end_time)\n",
        "    tempo_microseconds = bpm2tempo(tempo_bpm)\n",
        "    mid = MidiFile(input_file)\n",
        "    new_midi = MidiFile(ticks_per_beat=mid.ticks_per_beat)\n",
        "\n",
        "    # Create metadata track\n",
        "    meta_track = MidiTrack()\n",
        "    new_midi.tracks.append(meta_track)\n",
        "\n",
        "    meta_track.append(MetaMessage('time_signature',\n",
        "                                numerator=numerator,\n",
        "                                denominator=denominator,\n",
        "                                clocks_per_click=24,\n",
        "                                notated_32nd_notes_per_beat=8,\n",
        "                                time=0))\n",
        "\n",
        "    meta_track.append(MetaMessage('set_tempo',\n",
        "                                tempo=tempo_microseconds,\n",
        "                                time=0))\n",
        "\n",
        "    start_ticks = int(start_time * tempo_bpm * mid.ticks_per_beat / 60)\n",
        "    end_ticks = int(end_time * tempo_bpm * mid.ticks_per_beat / 60)\n",
        "\n",
        "    # Process each track\n",
        "    for track in mid.tracks[1:]:\n",
        "        new_track = MidiTrack()\n",
        "        current_time = 0\n",
        "        track_started = False\n",
        "        active_notes = {}  # Dictionary to track active notes: key = (note, channel), value = start_time\n",
        "        messages_added = False\n",
        "\n",
        "        # Pre-scan for notes that started before our window but should be included\n",
        "        for msg in track:\n",
        "            current_time += msg.time\n",
        "            if current_time < start_ticks:\n",
        "                if msg.type == 'note_on' and msg.velocity > 0:\n",
        "                    active_notes[(msg.note, msg.channel)] = current_time\n",
        "                elif msg.type == 'note_off' or (msg.type == 'note_on' and msg.velocity == 0):\n",
        "                    if (msg.note, msg.channel) in active_notes:\n",
        "                        del active_notes[(msg.note, msg.channel)]\n",
        "\n",
        "        # Add still-active notes from before our window\n",
        "        if active_notes:\n",
        "            for (note, channel), note_start in active_notes.items():\n",
        "                new_track.append(Message('note_on', note=note, velocity=64,\n",
        "                                      time=0 if track_started else 0,\n",
        "                                      channel=channel))\n",
        "                track_started = True\n",
        "                messages_added = True\n",
        "\n",
        "        # Reset for main processing\n",
        "        current_time = 0\n",
        "\n",
        "        # Process messages within our window\n",
        "        for msg in track:\n",
        "            current_time += msg.time\n",
        "\n",
        "            if current_time <= end_ticks:\n",
        "                if start_ticks <= current_time:\n",
        "                    # Handle note_on messages\n",
        "                    if msg.type == 'note_on' and msg.velocity > 0:\n",
        "                        if not track_started:\n",
        "                            new_msg = msg.copy(time=0)\n",
        "                            track_started = True\n",
        "                        else:\n",
        "                            new_msg = msg.copy(time=msg.time)\n",
        "                        new_track.append(new_msg)\n",
        "                        messages_added = True\n",
        "                        active_notes[(msg.note, msg.channel)] = current_time\n",
        "\n",
        "                    # Handle note_off messages (including note_on with velocity 0)\n",
        "                    elif msg.type == 'note_off' or (msg.type == 'note_on' and msg.velocity == 0):\n",
        "                        if (msg.note, msg.channel) in active_notes:\n",
        "                            if not track_started:\n",
        "                                new_msg = msg.copy(time=0)\n",
        "                                track_started = True\n",
        "                            else:\n",
        "                                new_msg = msg.copy(time=msg.time)\n",
        "                            new_track.append(new_msg)\n",
        "                            messages_added = True\n",
        "                            del active_notes[(msg.note, msg.channel)]\n",
        "\n",
        "                    # Handle other control messages (pitch bend, control change, etc.)\n",
        "                    elif msg.type in ['pitch_bend', 'control_change', 'program_change']:\n",
        "                        if not track_started:\n",
        "                            new_msg = msg.copy(time=0)\n",
        "                            track_started = True\n",
        "                        else:\n",
        "                            new_msg = msg.copy(time=msg.time)\n",
        "                        new_track.append(new_msg)\n",
        "                        messages_added = True\n",
        "\n",
        "        # For any notes still active at the end of the chunk,\n",
        "        # add note_off messages at the chunk boundary\n",
        "        if active_notes:\n",
        "            last_time = current_time if current_time <= end_ticks else end_ticks\n",
        "            time_to_end = end_ticks - last_time\n",
        "\n",
        "            # Add note_off for first note with the time delta to chunk end\n",
        "            first_note = True\n",
        "            for (note, channel) in list(active_notes.keys()):\n",
        "                if first_note:\n",
        "                    new_track.append(Message('note_off', note=note, velocity=0,\n",
        "                                          time=time_to_end,\n",
        "                                          channel=channel))\n",
        "                    first_note = False\n",
        "                else:\n",
        "                    # Subsequent note_offs happen simultaneously\n",
        "                    new_track.append(Message('note_off', note=note, velocity=0,\n",
        "                                          time=0,\n",
        "                                          channel=channel))\n",
        "                messages_added = True\n",
        "                del active_notes[(note, channel)]\n",
        "\n",
        "        # Only add tracks that had actual content\n",
        "        if messages_added:\n",
        "            new_midi.tracks.append(new_track)\n",
        "            new_track.append(MetaMessage('end_of_track', time=0))\n",
        "\n",
        "    new_midi.save(output_file)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5t6k6KsKrPCg"
      },
      "source": [
        "## Datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iqrrqAiqIJBU"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from torch.utils.data import Dataset\n",
        "from collections import OrderedDict\n",
        "import random\n",
        "from typing import List, Tuple\n",
        "import numpy as np\n",
        "import librosa\n",
        "\n",
        "def generate_short_qa_data(tempo: float, instruments: List[str]) -> Tuple[str, str]:\n",
        "    tempo = int(tempo)\n",
        "    question_type = random.choices([\"tempo\", \"instrument\"], weights=[20, 80], k=1)[0]\n",
        "\n",
        "    if question_type == \"tempo\":\n",
        "        # Variations of tempo-related questions\n",
        "        tempo_questions = [\n",
        "            \"What is the tempo of this audio clip?\",\n",
        "            \"Can you tell me the tempo of the track?\",\n",
        "            \"How fast is the tempo in this music?\",\n",
        "            \"What beats per minute (BPM) is this audio clip playing at?\",\n",
        "            \"Identify the tempo in this audio track.\"\n",
        "        ]\n",
        "        question = random.choice(tempo_questions)\n",
        "\n",
        "        # Variations of tempo-related answers\n",
        "        tempo_answers = [\n",
        "            f\"The tempo of this audio clip is {tempo} BPM.\",\n",
        "            f\"This track has a tempo of {tempo} beats per minute.\",\n",
        "            f\"The BPM for this audio is {tempo}.\",\n",
        "            f\"The music plays at a tempo of {tempo} BPM.\",\n",
        "            f\"The speed of this track is {tempo} beats per minute.\"\n",
        "        ]\n",
        "        answer = random.choice(tempo_answers)\n",
        "\n",
        "    elif question_type == \"instrument\":\n",
        "        instrument_questions = [\n",
        "            \"What instrument is playing in this audio clip?\",\n",
        "            \"What instruments are playing in this audio clip?\",\n",
        "            \"Can you identify the instrument in this track?\",\n",
        "            \"Can you name the instruments playing in this track?\",\n",
        "            \"Which instruments are present in this audio?\",\n",
        "            \"List the instruments that are featured in this recording.\"\n",
        "        ]\n",
        "\n",
        "        question = random.choice(instrument_questions)\n",
        "\n",
        "        # Variations of instrument-related answers for one instrument\n",
        "        if len(instruments) == 1:\n",
        "            instrument_answers_single = [\n",
        "                f\"In this audio clip, a {instruments[0]} is playing.\",\n",
        "                f\"You can hear a {instruments[0]} playing in this audio.\",\n",
        "                f\"This track features a {instruments[0]}.\",\n",
        "                f\"The sound in this audio is produced by a {instruments[0]}.\",\n",
        "                f\"In this recording, a {instruments[0]} is the main instrument.\"\n",
        "            ]\n",
        "            answer = random.choice(instrument_answers_single)\n",
        "\n",
        "        # Variations of instrument-related answers for multiple instruments\n",
        "        else:\n",
        "            instrument_answers_multiple = [\n",
        "                f\"In this audio clip, the following instruments are playing: {', '.join(instruments[:-1])}, and {instruments[-1]}.\",\n",
        "                f\"This track features these instruments: {', '.join(instruments[:-1])}, and {instruments[-1]}.\",\n",
        "                f\"The instruments heard in this recording are: {', '.join(instruments[:-1])}, and {instruments[-1]}.\",\n",
        "                f\"You can hear the following instruments in this audio: {', '.join(instruments[:-1])}, and {instruments[-1]}.\",\n",
        "                f\"In this audio, you can hear a combination of: {', '.join(instruments[:-1])}, and {instruments[-1]}.\"\n",
        "            ]\n",
        "            answer = random.choice(instrument_answers_multiple)\n",
        "\n",
        "    return question, answer\n",
        "\n",
        "class BaseAudioDataset(Dataset):\n",
        "    audio_cache = OrderedDict()\n",
        "    max_cache_size = 1000\n",
        "\n",
        "    def __init__(self, root_dir, dataset_name, duration=25, target_sr=16000, return_audio=True):\n",
        "        self.root_dir = root_dir\n",
        "        self.dataset_name = dataset_name\n",
        "        self.duration = duration\n",
        "        self.target_sr = target_sr\n",
        "        self.audio_cache = OrderedDict()\n",
        "        self.max_cache_size = 1000\n",
        "        self.samples = []\n",
        "        self.return_audio = return_audio\n",
        "\n",
        "        BaseAudioDataset.max_cache_size = 1000\n",
        "\n",
        "    def load_audio(self, audio_path):\n",
        "        if not self.return_audio:\n",
        "            return None\n",
        "\n",
        "        cache = BaseAudioDataset.audio_cache\n",
        "        if audio_path in cache:\n",
        "            cache.move_to_end(audio_path)\n",
        "            return cache[audio_path]\n",
        "        else:\n",
        "            audio, sr = librosa.load(audio_path, sr=None, mono=True)\n",
        "            if sr != self.target_sr:\n",
        "                audio = librosa.resample(audio, orig_sr=sr, target_sr=self.target_sr)\n",
        "            cache[audio_path] = audio\n",
        "            if len(cache) > BaseAudioDataset.max_cache_size:\n",
        "                cache.popitem(last=False)\n",
        "            return audio\n",
        "\n",
        "    def get_audio_chunk(self, audio, start_time, duration):\n",
        "        start_sample = int(start_time * self.target_sr)\n",
        "        end_sample = int((start_time + duration) * self.target_sr)\n",
        "        audio_chunk = audio[start_sample:end_sample]\n",
        "\n",
        "        return normalize_audio(audio_chunk)\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.samples)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        if idx >= len(self.samples):\n",
        "            raise IndexError(f\"Index {idx} is out of range for dataset with {len(self.samples)} samples\")\n",
        "\n",
        "        print(f\"Processing sample {idx}\")\n",
        "        audio_path, _, start_time, chunk_duration, folder_path, _ = self.samples[idx]\n",
        "\n",
        "        audio_chunk = None\n",
        "        if self.return_audio:\n",
        "            audio = self.load_audio(audio_path)\n",
        "            audio_chunk = self.get_audio_chunk(audio, start_time, chunk_duration)\n",
        "\n",
        "        result = self.process_item(idx, audio_chunk)\n",
        "        if result is not None:\n",
        "            return result\n",
        "\n",
        "    def process_item(self, idx, audio_chunk):\n",
        "        # Implemented by child classes\n",
        "        raise NotImplementedError\n",
        "\n",
        "def normalize_audio(audio_chunk):\n",
        "    mean_audio = np.mean(audio_chunk)\n",
        "    std_audio = np.std(audio_chunk)\n",
        "    if std_audio > 0:\n",
        "        return (audio_chunk - mean_audio) / std_audio\n",
        "    else:\n",
        "        return audio_chunk - mean_audio"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4xgE9ARnsDGT"
      },
      "source": [
        "### Maestro"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZLeKlmvhDcti"
      },
      "outputs": [],
      "source": [
        "def process_maestro_chunk(midi_file, synthetic_data: bool) -> Tuple[str, str]:\n",
        "    midi_data = pretty_midi.PrettyMIDI(midi_file)\n",
        "\n",
        "    tempo, metre_numerator, metre_denominator = get_tempo_and_metre_simp(midi_data)\n",
        "\n",
        "    if synthetic_data:\n",
        "      instruments = [\n",
        "          pretty_midi.program_to_instrument_name(instrument.program)\n",
        "          for instrument in midi_data.instruments\n",
        "      ]\n",
        "      return generate_short_qa_data(tempo, instruments)\n",
        "\n",
        "    abc_notation = cmd_midi_to_abc(midi_file, extract_all_tracks=True)\n",
        "    abc_notation = remove_key_signature(abc_notation)\n",
        "    abc_notation = remove_comment_lines(abc_notation)\n",
        "    abc_notation = remove_backslash_after_bar(abc_notation)\n",
        "    abc_notation = abc_notation.replace('V:1', 'V:1 name=\"Acoustic Grand Piano\"')\n",
        "\n",
        "    return random.choice(abc_questions), abc_notation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o2KLroMAA8Cs"
      },
      "outputs": [],
      "source": [
        "!mkdir /content/maestro/temp_midi/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tUWDRbKI7zGk"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import yaml\n",
        "import librosa\n",
        "import numpy as np\n",
        "import pretty_midi\n",
        "import torch\n",
        "from torch.utils.data import Dataset\n",
        "import random\n",
        "import pickle\n",
        "import os\n",
        "import torch.nn.functional as F\n",
        "import glob\n",
        "import mido\n",
        "\n",
        "class MaestroDataset(BaseAudioDataset):\n",
        "    def __init__(self, root_dir, dataset_name, duration=25, target_sr=16000, return_audio=True):\n",
        "        super().__init__(root_dir, dataset_name, duration, target_sr, return_audio)\n",
        "        self.initialize_samples()\n",
        "\n",
        "    def initialize_samples(self):\n",
        "        for year_folder in os.listdir(self.root_dir):\n",
        "            year_path = os.path.join(self.root_dir, year_folder)\n",
        "            if os.path.isdir(year_path):\n",
        "\n",
        "                wav_files = glob.glob(os.path.join(year_path, '*.wav'))\n",
        "\n",
        "                for wav_file in wav_files:\n",
        "                    # Get the base name without extension\n",
        "                    base_name = os.path.splitext(os.path.basename(wav_file))[0]\n",
        "                    midi_path = os.path.join(year_path, base_name + '.midi')\n",
        "\n",
        "                    total_duration = librosa.get_duration(path=wav_file)\n",
        "\n",
        "                    num_chunks = int(total_duration // self.duration)  # Full chunks\n",
        "                    for i in range(num_chunks):\n",
        "                        start_time = i * self.duration\n",
        "                        midi_file_name = os.path.basename(midi_path)\n",
        "                        output_midi_path = f'/content/maestro/temp_midi/{midi_file_name}_{str(start_time)}.mid'\n",
        "\n",
        "                        try:\n",
        "                            extract_midi_segment(midi_path, output_midi_path, float(start_time), float(self.duration))\n",
        "                            self.samples.append((wav_file, output_midi_path, start_time, self.duration, \"\", False))\n",
        "                        except Exception as e:\n",
        "                            print(f\"Skipping sample due to MIDI extraction error: {e}\")\n",
        "                            continue\n",
        "\n",
        "        '''\n",
        "        # Take 25% of the samples\n",
        "        sample_size = len(self.samples) // 4\n",
        "        self.samples = random.sample(self.samples, sample_size)\n",
        "        '''\n",
        "\n",
        "    def process_item(self, idx, audio_chunk):\n",
        "        audio_path, midi_path, start_time, chunk_duration, _, synthetic_flag = self.samples[idx]\n",
        "\n",
        "        query, answer = process_maestro_chunk(midi_path, synthetic_flag)\n",
        "\n",
        "        return audio_chunk, query, answer, self.dataset_name, audio_path, int(start_time)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kBx_v58IZjxz"
      },
      "source": [
        "### URMP"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kZn5JN3Q4Ljz"
      },
      "outputs": [],
      "source": [
        "from pretty_midi import Instrument\n",
        "from typing import Optional, List, Tuple\n",
        "import re\n",
        "import random\n",
        "\n",
        "def insert_instrument_names(abc_notation, instrument_names):\n",
        "    lines = abc_notation.split('\\n')\n",
        "    instrument_index = 0\n",
        "\n",
        "    # Process each line\n",
        "    for i in range(len(lines)):\n",
        "        # Look for voice headers\n",
        "        if lines[i].startswith('V:'):\n",
        "            # Extract voice number\n",
        "            voice_num = lines[i].split(':')[1].strip()\n",
        "\n",
        "            # Replace the line with new format including instrument name\n",
        "            if instrument_index < len(instrument_names):\n",
        "                lines[i] = f'V:{voice_num} name=\"{instrument_names[instrument_index]}\"'\n",
        "                instrument_index += 1\n",
        "\n",
        "    return '\\n'.join(lines)\n",
        "\n",
        "def process_urmp_chunk(midi_track_number: int, instruments: List[str], midi_file: str,\n",
        "                      chunk_start_time: int, sample_duration: int, synthetic_data: bool) -> Tuple[str, str]:\n",
        "    def get_instrument_name(instrument: Instrument, instruments: List[str], index: int) -> str:\n",
        "        if not instrument.name:\n",
        "            print(f\"Warning: instrument names not available: {midi_data.instruments}\")\n",
        "            return instruments[index]\n",
        "        return instrument.name\n",
        "\n",
        "    midi_data = pretty_midi.PrettyMIDI(midi_file)\n",
        "    tempo, metre_numerator, metre_denominator = get_tempo_and_metre_simp(midi_data)\n",
        "\n",
        "    # Handle instrument names\n",
        "    if midi_track_number is not None:\n",
        "        instrument_names = [instruments[0]] if synthetic_data else instruments\n",
        "    else:\n",
        "        instrument_names = [\n",
        "            get_instrument_name(instrument, instruments, i)\n",
        "            for i, instrument in enumerate(midi_data.instruments)\n",
        "        ]\n",
        "\n",
        "    if synthetic_data:\n",
        "        return generate_short_qa_data(tempo, instrument_names)\n",
        "\n",
        "    abc_notation = cmd_midi_to_abc(midi_file, extract_all_tracks=True)\n",
        "    abc_notation = remove_key_signature(abc_notation)\n",
        "    abc_notation = remove_comment_lines(abc_notation)\n",
        "    abc_notation = remove_backslash_after_bar(abc_notation)\n",
        "    abc_notation = insert_instrument_names(abc_notation, instrument_names)\n",
        "\n",
        "    return random.choice(abc_questions), abc_notation\n",
        "\n",
        "# Map of instrument abbreviations to full names\n",
        "instrument_map = {\n",
        "    \"Vn\": \"Violin\",\n",
        "    \"Va\": \"Viola\",\n",
        "    \"Vc\": \"Cello\",\n",
        "    \"Db\": \"Double Bass\",\n",
        "    \"Fl\": \"Flute\",\n",
        "    \"Ob\": \"Oboe\",\n",
        "    \"Cl\": \"Clarinet\",\n",
        "    \"Sax\": \"Saxophone\",\n",
        "    \"Bn\": \"Bassoon\",\n",
        "    \"Tpt\": \"Trumpet\",\n",
        "    \"Hn\": \"Horn\",\n",
        "    \"Tbn\": \"Trombone\",\n",
        "    \"Tba\": \"Tuba\"\n",
        "}\n",
        "\n",
        "def extract_instruments(string):\n",
        "    instruments = []\n",
        "    for abbr, name in instrument_map.items():\n",
        "        # Use re.findall to find all occurrences of the abbreviation, case insensitive\n",
        "        matches = re.findall(rf'(?<![a-zA-Z]){abbr}(?![a-zA-Z])', string, re.IGNORECASE)\n",
        "        instruments.extend([name] * len(matches))\n",
        "    return instruments"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VsGkB-j_AJfi"
      },
      "outputs": [],
      "source": [
        "!mkdir /content/URMP_Dataset/temp_midi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z0JNJx_fU9yf"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import librosa\n",
        "\n",
        "class URMPDataset(BaseAudioDataset):\n",
        "    def __init__(self, root_dir, dataset_name, duration=25, target_sr=16000, return_audio=True):\n",
        "        super().__init__(root_dir, dataset_name, duration, target_sr, return_audio)\n",
        "        self.initialize_samples()\n",
        "\n",
        "    def initialize_samples(self):\n",
        "        for sub_folder in os.listdir(self.root_dir):\n",
        "            sub_folder_path = os.path.join(self.root_dir, sub_folder)\n",
        "\n",
        "            if os.path.isdir(sub_folder_path):\n",
        "                sub_folder_name = os.path.basename(sub_folder_path)\n",
        "                midi_path = os.path.join(sub_folder_path, f\"Sco_{sub_folder_name}.mid\")\n",
        "\n",
        "                for file in os.listdir(sub_folder_path):\n",
        "                    # OS specific files\n",
        "                    if file.startswith('._') or not file.endswith('.wav'):\n",
        "                        continue\n",
        "\n",
        "                    audio_path = os.path.join(sub_folder_path, file)\n",
        "\n",
        "                    if os.path.exists(midi_path):\n",
        "                        total_duration = librosa.get_duration(path=audio_path)\n",
        "                        num_chunks = math.ceil(total_duration / self.duration)\n",
        "\n",
        "                        for i in range(num_chunks):\n",
        "                            start_time = i * self.duration\n",
        "                            chunk_duration = min(self.duration, total_duration - start_time)\n",
        "                            midi_file_name = os.path.basename(midi_path)\n",
        "                            output_midi_path = f'/content/URMP_Dataset/temp_midi/{midi_file_name}_{str(start_time)}.mid'\n",
        "\n",
        "                            try:\n",
        "                                extract_midi_segment(midi_path, output_midi_path, start_time, chunk_duration)\n",
        "                                self.samples.append((audio_path, output_midi_path, start_time, chunk_duration, \"\", False))\n",
        "\n",
        "                                # Generate short QA sample\n",
        "                                if random.randint(1, 1) == 1:\n",
        "                                    self.samples.append((audio_path, output_midi_path, start_time, chunk_duration, \"\", True))\n",
        "                            except Exception as e:\n",
        "                                print(f\"Skipping sample due to MIDI extraction error: {e}\")\n",
        "                                continue\n",
        "\n",
        "\n",
        "    def process_item(self, idx, audio_chunk):\n",
        "        audio_path, midi_path, start_time, chunk_duration, _, synthetic_flag = self.samples[idx]\n",
        "\n",
        "        audio_path_basename = os.path.basename(audio_path)\n",
        "        instruments = extract_instruments(audio_path_basename)\n",
        "\n",
        "        if \"AuSep\" in audio_path_basename:\n",
        "            midi_track_number = int(audio_path_basename[6])\n",
        "        else:\n",
        "            midi_track_number = None\n",
        "\n",
        "        query, answer = process_urmp_chunk(midi_track_number, instruments, midi_path, start_time, chunk_duration, synthetic_flag)\n",
        "\n",
        "        id = os.path.basename(audio_path)\n",
        "\n",
        "        sample_type = \"qa\" if synthetic_flag else \"abc\"\n",
        "\n",
        "        return audio_chunk, query, answer, self.dataset_name, id, start_time, sample_type"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_cF-Q94KBG6O"
      },
      "outputs": [],
      "source": [
        "import re\n",
        "import math\n",
        "\n",
        "urmp_dataset = URMPDataset(root_dir=\"/content/URMP_Dataset/Dataset\", dataset_name=\"urmp\")\n",
        "maestro_dataset = MaestroDataset(root_dir=\"/content/maestro\", dataset_name=\"maestro\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NFobVzJI653T"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torch.utils.data import Dataset\n",
        "import random\n",
        "\n",
        "class CombinedDataset(Dataset):\n",
        "    def __init__(self, datasets):\n",
        "        self.datasets = datasets\n",
        "        self.lengths = [len(dataset) for dataset in datasets]\n",
        "        self.cumulative_lengths = torch.cumsum(torch.tensor(self.lengths), dim=0)\n",
        "\n",
        "        # Create a list of indices that covers all datasets\n",
        "        self.indices = [(i, j) for i, dataset in enumerate(datasets) for j in range(len(dataset))]\n",
        "\n",
        "        random.shuffle(self.indices)\n",
        "\n",
        "    def __len__(self):\n",
        "        return sum(self.lengths)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        dataset_idx, sample_idx = self.indices[idx]\n",
        "        return self.datasets[dataset_idx][sample_idx]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IfX1z6iw68L_"
      },
      "outputs": [],
      "source": [
        "combined_dataset = CombinedDataset([urmp_dataset, maestro_dataset])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eN7KwECklnFL"
      },
      "source": [
        "## Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yEpT3GfzVm2Z"
      },
      "outputs": [],
      "source": [
        "# Import the login function\n",
        "from huggingface_hub import login\n",
        "\n",
        "# Log in to Hugging Face\n",
        "login(token='hf_token')"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q -U bitsandbytes\n",
        "!pip install git+https://github.com/huggingface/transformers\n",
        "!pip install -q -U git+https://github.com/huggingface/peft.git\n",
        "!pip install -q -U git+https://github.com/huggingface/accelerate.git\n",
        "!pip install datasets\n",
        "!pip install trl==0.11.4"
      ],
      "metadata": {
        "id": "ZEfFJgCz7fdT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DBHLPEvPYTd_"
      },
      "source": [
        "**RESTART THE SESSION OTHERWISE THE TRAINER WON'T WORK!**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IhLOXq5Beii7"
      },
      "outputs": [],
      "source": [
        "from transformers import TrainingArguments, Qwen2AudioForConditionalGeneration\n",
        "from peft import get_peft_model, LoraConfig, TaskType\n",
        "from trl import SFTTrainer\n",
        "import torch\n",
        "\n",
        "model = Qwen2AudioForConditionalGeneration.from_pretrained(\"Qwen/Qwen2-Audio-7B-Instruct\", torch_dtype=torch.bfloat16, device_map=\"auto\")\n",
        "# model = Qwen2AudioForConditionalGeneration.from_pretrained(\"Qwen/Qwen2-Audio-7B-Instruct\", load_in_4bit=True, device_map=\"auto\")\n",
        "\n",
        "peft_config = LoraConfig(task_type=TaskType.CAUSAL_LM, inference_mode=False, r=128, lora_alpha=256, lora_dropout=0, target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"])\n",
        "model.enable_input_require_grads()\n",
        "model = get_peft_model(model, peft_config)\n",
        "#model.gradient_checkpointing_enable()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NtgVobWTViL-"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RF5g4AS6GtVN"
      },
      "source": [
        "### Load wandb"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kMs2GDzA1fWe"
      },
      "outputs": [],
      "source": [
        "import wandb, os\n",
        "wandb.login()\n",
        "\n",
        "wandb_project = \"autotab\"\n",
        "if len(wandb_project) > 0:\n",
        "    os.environ[\"WANDB_PROJECT\"] = wandb_project"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "26AZA7qId07Z"
      },
      "outputs": [],
      "source": [
        "project = \"Qwen2Audio-1\"\n",
        "run_name = \"run\"\n",
        "project_and_run_name = project + \"-\" + run_name\n",
        "output_dir = \"./\" + project_and_run_name"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tCMWxBCd9DIL"
      },
      "source": [
        "### Train with SFT"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LgMJBHulADCN"
      },
      "outputs": [],
      "source": [
        "from transformers import Trainer, TrainingArguments\n",
        "from datetime import datetime\n",
        "from trl import SFTTrainer\n",
        "\n",
        "wandbname = project + \"-\" + run_name\n",
        "\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "\n",
        "training_args = TrainingArguments(\n",
        "    output_dir=output_dir,\n",
        "    num_train_epochs=3,\n",
        "    #max_steps=1000,\n",
        "    # per_device_eval_batch_size = 0,\n",
        "    per_device_train_batch_size = 1,\n",
        "    gradient_accumulation_steps = 1,\n",
        "    warmup_ratio = 0.1,\n",
        "    logging_dir='./logs',\n",
        "    learning_rate = 1e-5,\n",
        "    logging_steps = 1,\n",
        "    #evaluation_strategy=\"steps\",\n",
        "    #eval_steps=75,\n",
        "    #save_steps=100,\n",
        "    max_grad_norm=10.0,\n",
        "    fp16 = not torch.cuda.is_bf16_supported(),\n",
        "    gradient_checkpointing=True,\n",
        "    bf16 = torch.cuda.is_bf16_supported(),\n",
        "    optim = \"adamw_8bit\",\n",
        "    weight_decay = 0.001,\n",
        "    #seed = 3407,\n",
        "    save_strategy=\"no\",\n",
        "    #save_strategy=\"epoch\",\n",
        "    lr_scheduler_type = \"cosine\",\n",
        "    #load_best_model_at_end=True,\n",
        "    report_to=\"wandb\",\n",
        "    run_name=f\"{wandbname}-{datetime.now().strftime('%m-%d-%H-%M')}\"\n",
        ")\n",
        "\n",
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    args=training_args,\n",
        "    train_dataset=train_dataset,\n",
        "    # eval_dataset=test_dataset,\n",
        "    data_collator=data_collator\n",
        ")\n",
        "\n",
        "trainer.train()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YnRW5cM6q5ZR"
      },
      "source": [
        "### Train with PPO"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SPDWPSg_Q5Uk"
      },
      "outputs": [],
      "source": [
        "project = \"Qwen2Audio-1\"\n",
        "run_name = \"run\"\n",
        "project_and_run_name = project + \"-\" + run_name\n",
        "output_dir = \"./\" + project_and_run_name"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NJYDTAu-UfoK"
      },
      "outputs": [],
      "source": [
        "import wandb, os\n",
        "wandb.login()\n",
        "\n",
        "wandb_project = \"autotab\"\n",
        "if len(wandb_project) > 0:\n",
        "    os.environ[\"WANDB_PROJECT\"] = wandb_project"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mOs_gWcjrBgv"
      },
      "source": [
        "#### Custom loss/reward model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sKS0iihrq6wD"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import re\n",
        "import math\n",
        "import pretty_midi\n",
        "import numpy as np\n",
        "import difflib\n",
        "\n",
        "class ABCLoss(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(ABCLoss, self).__init__()\n",
        "\n",
        "        # Mapping keys to their positions on the circle of fifths, accounting for both major and minor keys\n",
        "        self.key_to_circle_position = {\n",
        "            'C': 0, 'G': 1, 'D': 2, 'A': 3, 'E': 4, 'B': 5, 'F#': 6, 'C#': 7,\n",
        "            'F': -1, 'Bb': -2, 'Eb': -3, 'Ab': -4, 'Db': -5, 'Gb': -6, 'Cb': -7,\n",
        "            'Am': 0, 'Em': 1, 'Bm': 2, 'F#m': 3, 'C#m': 4, 'G#m': 5, 'D#m': 6, 'A#m': 7,\n",
        "            'Dm': -1, 'Gm': -2, 'Cm': -3, 'Fm': -4, 'Bbm': -5, 'Ebm': -6, 'Abm': -7\n",
        "        }\n",
        "\n",
        "    def extract_metre(self, abc_string):\n",
        "        match = re.search(r'M:\\s*\\(?\\s*(\\d+/\\d+)\\s*\\)?', abc_string)\n",
        "        return match.group(1) if match else None\n",
        "\n",
        "    def extract_length(self, abc_string):\n",
        "        match = re.search(r'L:\\s*\\(?\\s*(\\d+/\\d+)\\s*\\)?', abc_string)\n",
        "        return match.group(1) if match else None\n",
        "\n",
        "    def extract_tempo(self, abc_string):\n",
        "        match = re.search(r'Q:\\s*\\(?\\s*(\\d+)\\s*\\)?', abc_string)\n",
        "        return match.group(1) if match else None\n",
        "\n",
        "    def extract_key(self, abc_string):\n",
        "        match = re.search(r'K:\\s*\\(?\\s*(\\S+)\\s*\\)?', abc_string)\n",
        "        return match.group(1) if match else None\n",
        "\n",
        "    def extract_voice_notes(self, abc_string):\n",
        "        # Extracts the content of each voice, including the 'V:...' line\n",
        "        voices = re.findall(r'(V:\\d+.*?\\n(?:.*?\\n)*?)(?=V:|$)', abc_string)\n",
        "        voice_contents = []\n",
        "        for voice in voices:\n",
        "            # Remove the voice header line\n",
        "            content = '\\n'.join(voice.split('\\n')[1:])\n",
        "            voice_contents.append(content)\n",
        "        return voice_contents\n",
        "\n",
        "    def parse_fraction(self, fraction_string):\n",
        "        \"\"\"Convert a fraction string like '4/4' into separate numerator and denominator.\"\"\"\n",
        "        num, denom = fraction_string.split('/')\n",
        "        return int(num), int(denom)\n",
        "\n",
        "    def compute_metre_loss(self, pred_metre, gt_metre):\n",
        "        pred_num, pred_denom = self.parse_fraction(pred_metre)\n",
        "        gt_num, gt_denom = self.parse_fraction(gt_metre)\n",
        "\n",
        "        # Compute loss for numerator and denominator separately\n",
        "        num_diff = abs(pred_num - gt_num)\n",
        "        denom_diff = abs(pred_denom - gt_denom)\n",
        "\n",
        "        # Compute a ratio difference\n",
        "        ratio_pred = pred_num / pred_denom\n",
        "        ratio_gt = gt_num / gt_denom\n",
        "        ratio_diff = abs(ratio_pred - ratio_gt)\n",
        "\n",
        "        # Define thresholds\n",
        "        linear_threshold_num = 1\n",
        "        linear_threshold_denom = 2\n",
        "        linear_threshold_ratio = 0.25\n",
        "\n",
        "        # Compute individual losses\n",
        "        def compute_component_loss(diff, threshold):\n",
        "            if diff <= threshold:\n",
        "                return diff / threshold * 0.5\n",
        "            else:\n",
        "                base_loss = 0.5\n",
        "                exp_scale = 1 - math.exp(-(diff - threshold) / threshold)\n",
        "                return base_loss + 0.5 * exp_scale\n",
        "\n",
        "        num_loss = compute_component_loss(num_diff, linear_threshold_num)\n",
        "        denom_loss = compute_component_loss(denom_diff, linear_threshold_denom)\n",
        "        ratio_loss = compute_component_loss(ratio_diff, linear_threshold_ratio)\n",
        "\n",
        "        # Combine losses with emphasis on the ratio difference\n",
        "        total_loss = 0.3 * num_loss + 0.3 * denom_loss + 0.4 * ratio_loss\n",
        "\n",
        "        # Ensure loss is between 0 and 1\n",
        "        total_loss = min(total_loss, 1.0)\n",
        "\n",
        "        return torch.tensor(total_loss, dtype=torch.float32)\n",
        "\n",
        "    def compute_tempo_loss(self, pred_tempo, gt_tempo):\n",
        "        if not pred_tempo or not gt_tempo:\n",
        "            return torch.tensor(0.0)\n",
        "\n",
        "        pred_value = float(pred_tempo)\n",
        "        target_value = float(gt_tempo)\n",
        "\n",
        "        # Calculate the absolute difference\n",
        "        diff = abs(pred_value - target_value)\n",
        "\n",
        "        # Define thresholds\n",
        "        linear_threshold = 10  # BPM\n",
        "        max_diff = 200  # Maximum expected difference in BPM\n",
        "\n",
        "        # Compute loss\n",
        "        if diff <= linear_threshold:\n",
        "            # Linear scaling for small differences\n",
        "            loss = diff / linear_threshold * 0.25  # Max 0.5 for linear part\n",
        "        else:\n",
        "            # Exponential scaling for larger differences\n",
        "            base_loss = 0.5  # Starting point after linear threshold\n",
        "            exp_scale = 1 - math.exp(-(diff - linear_threshold) / 50)  # Adjust '50' to control steepness\n",
        "            loss = base_loss + 0.5 * exp_scale  # Additional 0.5 for exponential part\n",
        "\n",
        "        # Ensure loss is between 0 and 1\n",
        "        loss = min(loss, 1.0)\n",
        "\n",
        "        return torch.tensor(loss, dtype=torch.float32)\n",
        "\n",
        "    def extract_instruments(self, abc_string):\n",
        "        # Handle instrument names with or without quotes\n",
        "        voice_headers = re.findall(r'V:\\d+[^\\n]*name=\"?([^\"\\n]*)\"?', abc_string)\n",
        "        return [self.instrument_name_to_program(name) for name in voice_headers]\n",
        "\n",
        "    def instrument_name_to_program(self, name):\n",
        "        # Clean up the name\n",
        "        name = name.strip().lower()\n",
        "        # Get list of all MIDI instrument names\n",
        "        instrument_names = [pretty_midi.program_to_instrument_name(p).lower() for p in range(128)]\n",
        "        # Use difflib to get the closest match\n",
        "        matches = difflib.get_close_matches(name, instrument_names, n=1, cutoff=0.6)\n",
        "        if matches:\n",
        "            closest_name = matches[0]\n",
        "            # Get the program number for the closest name\n",
        "            for program in range(128):\n",
        "                if pretty_midi.program_to_instrument_name(program).lower() == closest_name:\n",
        "                    return program\n",
        "        else:\n",
        "            # No close match found\n",
        "            return -1  # Return -1 if no matching instrument is found\n",
        "\n",
        "    def program_to_family(self, program):\n",
        "        # Returns the family index for a given program number\n",
        "        return program // 8\n",
        "\n",
        "    def compute_instrument_loss(self, pred_abc, gt_abc):\n",
        "        pred_instruments = self.extract_instruments(pred_abc)\n",
        "        gt_instruments = self.extract_instruments(gt_abc)\n",
        "\n",
        "        total_instruments = max(len(pred_instruments), len(gt_instruments))\n",
        "        if total_instruments == 0:\n",
        "            return torch.tensor(0.0)  # No instruments in either prediction or ground truth\n",
        "\n",
        "        total_loss = 0.0\n",
        "        for pred_inst, gt_inst in zip(pred_instruments, gt_instruments):\n",
        "            if pred_inst == gt_inst:\n",
        "                loss = 0.0  # Exact match\n",
        "            elif pred_inst == -1 or gt_inst == -1:\n",
        "                # One of the instruments is unknown\n",
        "                loss = 1.0\n",
        "            elif self.program_to_family(pred_inst) == self.program_to_family(gt_inst):\n",
        "                # Same family\n",
        "                loss = 0.5\n",
        "            else:\n",
        "                # Compute similarity between instrument names\n",
        "                pred_name = pretty_midi.program_to_instrument_name(pred_inst).lower()\n",
        "                gt_name = pretty_midi.program_to_instrument_name(gt_inst).lower()\n",
        "                similarity = difflib.SequenceMatcher(None, pred_name, gt_name).ratio()\n",
        "                loss = 1 - similarity  # Higher similarity, lower loss\n",
        "            total_loss += loss\n",
        "\n",
        "        avg_loss = total_loss / total_instruments\n",
        "        return torch.tensor(avg_loss, dtype=torch.float32)\n",
        "\n",
        "    def note_to_number(self, note):\n",
        "        note_map = {'C': 0, 'D': 2, 'E': 4, 'F': 5, 'G': 7, 'A': 9, 'B': 11}\n",
        "        base_note = note[0].upper()\n",
        "        base = note_map.get(base_note, 0)\n",
        "        modifiers = note[1:]\n",
        "        octave = 0\n",
        "        accidental = 0\n",
        "        idx = 1\n",
        "        while idx < len(note):\n",
        "            char = note[idx]\n",
        "            if char == ',':\n",
        "                octave -= 1\n",
        "            elif char == \"'\":\n",
        "                octave += 1\n",
        "            elif char == '^':\n",
        "                accidental += 1\n",
        "            elif char == '_':\n",
        "                accidental -= 1\n",
        "            else:\n",
        "                break\n",
        "            idx += 1\n",
        "        # Apply accidentals\n",
        "        base += accidental\n",
        "        # Shift by octave\n",
        "        base += (octave + 5) * 12  # shift to avoid negative numbers\n",
        "        return base\n",
        "\n",
        "    def extract_notes_from_voice(self, voice_string):\n",
        "        # Remove whitespace and newlines\n",
        "        voice_string = re.sub(r'\\s+', '', voice_string)\n",
        "        # Extract notes\n",
        "        notes = re.findall(r'[\\^_]*[A-Ga-g][,\\']*', voice_string)\n",
        "        note_numbers = [self.note_to_number(note) for note in notes]\n",
        "        return note_numbers\n",
        "\n",
        "    def normalized_levenshtein_distance(self, seq1, seq2):\n",
        "        len_seq1 = len(seq1)\n",
        "        len_seq2 = len(seq2)\n",
        "        max_len = max(len_seq1, len_seq2)\n",
        "        if max_len == 0:\n",
        "            return 0.0\n",
        "        # Compute the Levenshtein distance\n",
        "        d = np.zeros((len_seq1 + 1, len_seq2 + 1), dtype=int)\n",
        "        for i in range(len_seq1 + 1):\n",
        "            d[i][0] = i\n",
        "        for j in range(len_seq2 + 1):\n",
        "            d[0][j] = j\n",
        "        for i in range(1, len_seq1 + 1):\n",
        "            for j in range(1, len_seq2 + 1):\n",
        "                if seq1[i - 1] == seq2[j - 1]:\n",
        "                    cost = 0\n",
        "                else:\n",
        "                    cost = 1\n",
        "                d[i][j] = min(\n",
        "                    d[i - 1][j] + 1,      # deletion\n",
        "                    d[i][j - 1] + 1,      # insertion\n",
        "                    d[i - 1][j - 1] + cost  # substitution\n",
        "                )\n",
        "        distance = d[len_seq1][len_seq2]\n",
        "        normalized_distance = distance / max_len\n",
        "        return normalized_distance\n",
        "\n",
        "    def forward(self, predicted_abc, ground_truth_abc):\n",
        "        # Extract components\n",
        "        pred_metre = self.extract_metre(predicted_abc)\n",
        "        gt_metre = self.extract_metre(ground_truth_abc)\n",
        "\n",
        "        pred_tempo = self.extract_tempo(predicted_abc)\n",
        "        gt_tempo = self.extract_tempo(ground_truth_abc)\n",
        "\n",
        "        pred_voice_notes = self.extract_voice_notes(predicted_abc)\n",
        "        gt_voice_notes = self.extract_voice_notes(ground_truth_abc)\n",
        "\n",
        "        metre_loss = self.compute_metre_loss(pred_metre, gt_metre) if pred_metre and gt_metre else torch.tensor(0.0)\n",
        "        tempo_loss = self.compute_tempo_loss(pred_tempo, gt_tempo) if pred_tempo and gt_tempo else torch.tensor(0.0)\n",
        "\n",
        "        # Compute pitch loss\n",
        "        pred_voice_notes_list = [self.extract_notes_from_voice(voice) for voice in pred_voice_notes]\n",
        "        gt_voice_notes_list = [self.extract_notes_from_voice(voice) for voice in gt_voice_notes]\n",
        "\n",
        "        pitch_losses = []\n",
        "        for pred_notes in pred_voice_notes_list:\n",
        "            # For this predicted voice, compute distances to all ground truth voices\n",
        "            distances = []\n",
        "            for gt_notes in gt_voice_notes_list:\n",
        "                distance = self.normalized_levenshtein_distance(pred_notes, gt_notes)\n",
        "                distances.append(distance)\n",
        "            if distances:\n",
        "                # Take the smallest distance\n",
        "                min_distance = min(distances)\n",
        "                pitch_losses.append(min_distance)\n",
        "            else:\n",
        "                # No ground truth voices, maximum loss\n",
        "                pitch_losses.append(1.0)\n",
        "        # Now, compute the average pitch loss\n",
        "        if pitch_losses:\n",
        "            pitch_loss = torch.tensor(sum(pitch_losses) / len(pitch_losses), dtype=torch.float32)\n",
        "        else:\n",
        "            pitch_loss = torch.tensor(1.0)\n",
        "\n",
        "        instrument_loss = self.compute_instrument_loss(predicted_abc, ground_truth_abc)\n",
        "\n",
        "        # Total loss\n",
        "        total_loss = (0.5 * pitch_loss +\n",
        "                      0.15 * metre_loss +\n",
        "                      0.15 * tempo_loss +\n",
        "                      0.2 * instrument_loss)\n",
        "\n",
        "        loss_dict = {\n",
        "            'Total Loss': total_loss.item(),\n",
        "            'Pitch Loss': pitch_loss.item(),\n",
        "            'Metre Loss': metre_loss.item(),\n",
        "            'Tempo Loss': tempo_loss.item(),\n",
        "            'Instrument Loss': instrument_loss.item()\n",
        "        }\n",
        "\n",
        "        return total_loss, loss_dict\n",
        "\n",
        "# Example usage\n",
        "predicted_abc = \"\"\"\n",
        "M: 1/16 L: 1/16 Q: 120 K: none V:1 name=\"Piano\"\n",
        "\n",
        "c2 | c4 e3 d3 c'3 | G3 A3 B3 c'3 d3 E3 F3 |\n",
        "G3 A3 B3 c'3 d3 E3 F3 | c2 d2 c2 e2 f2 |\n",
        "g2 f2 e2 d2 c2 B2 A2 | G2 F2 E2 D2 C2 B2 A2 |\n",
        "c2 d2 c2 e2 f2 | g2 f2 e2 d2 c2 B2 A2 |\n",
        "c2 d2 c2 e2 f2 | g2 f2 e2 d2 c2 B2 A2 |\n",
        "G2 F2 E2 D2 C2 B2 A2 | c2 d2 c2 e2 f2 |\n",
        "g2 f2 e2 d2 c2 B2 A2 | G2 F2 E2 D2 C2 B2 A2 |\n",
        "c2 d2 c2 e2 f2 | g2 f2 e2 d2 c2 B2 A2 |\n",
        "G2 F2 E2 D2 C2 B2 A2 | c2 d2 c2 e2 f2 |\n",
        "g2 f2 e2 d2 c2 B2 A2 | G2 F2 E2 D2 C2 B2 A2 |\n",
        "c2 d2 c2 e2 f2 | g2 f2 e2 d2 c2 B2 A2 |\n",
        "G2 F2 E2 D2 C2 B2 A2 | c2 d2 c2 e2 f2 |\n",
        "g2 f2 e2 d2 c2 B2 A2 | G2 F2 E2 D2 C2 B2 A2 |\n",
        "c2 d2 c2 e2 f2 | g2 f2 e2 d2 c2 B2 A2 |\n",
        "G2 F2 E2 D2 C2 B2 A2 | c2 d2 c2 e2 f2 |\n",
        "g2 f2 e2 d2 c2 B2 A2 | G2 F2 E2 D2 C2 B2 A2 |\n",
        "c2 d2 c2 e2 f2 | g2 f2 e2 d2 c2 B2 A2 |\n",
        "G2 F2 E2 D2 C2 B2 A2 | c2 d2 c2 e2 f2 |\n",
        "g2 f2 e2 d2 c2 B2 A2 | G2 F2 E2 D2 C2 B2 A2 |\n",
        "c2 d2 c2 e2 f2 | g2 f2 e2 d2 c2 B2 A2 |\n",
        "G2 F2 E2 D2 C2 B2 A2 | c2 d2 c2 e2 f2 |\n",
        "g2 f2 e2 d2 c2 B2 A2 | G2 F2 E2 D2 C2 B2 A2 |\n",
        "c2 d2 c2 e2 f2 | g2 f2 e2 d2 c2 B2 A2 |\n",
        "G2 F2 E2 D2 C2 B2 A2 | c2 d2 c2 e2 f2 |\n",
        "g2 f2 e2 d2 c2 B2 A2 | G2 F2 E2 D2 C2 B2 A2 |\n",
        "c2 d2 c2 e2 f2 | g2 f2 e2 d2 c2 B2 A2 |\n",
        "G2 F2 E2 D2 C2 B2 A2 | c2 d2 c2 e2 f2 |\n",
        "g2 f2 e2 d2 c2 B2 A2 | G2 F2 E2 D2 C2 B2 A2 |\n",
        "c2 d2 c2 e2 f2 | g2 f2 e2 d2 c2 B2 A2 |\n",
        "G2 F2 E2 D2 C2 B2 A2 | c2 d2 c2 e2 f2 |\n",
        "g2 f2 e2 d2 c2 B2 A2 | G2 F2 E2 D2 C2 B2 A2\n",
        "\"\"\"\n",
        "\n",
        "ground_truth_abc = \"\"\"\n",
        "X:1\n",
        "M:3/4\n",
        "L:1/16\n",
        "K:none\n",
        "Q:220\n",
        "\n",
        "V:1 name=\"Trumpet in Bb\"\n",
        "%%octave-default C6\n",
        "B,,B,,2 | C,C,2B,, B,,G2 D^2DC | DDD^F4G2D^2D |\n",
        "CDDCA^,4 G2D^2 | DCDD D^ F4G2D^2 |\n",
        "DCD DCA^,4G2 | D2A^,G, A,A,A^,G, A,A,A^,G,G2 |\n",
        "D2B, G,A,A,B, A,2<G,2 | G2D2 A^,G,A, A, A^,G,A,A, |\n",
        "\n",
        "V:2 name=\"Trumpet in Bb\"\n",
        "%%octave-default C5\n",
        "DDD | EF^G2 G,D^'2 A2AA | A^2>D'2 D'D^'F'D^'2A2A |\n",
        "A2<A^2 D'D'C'A^ D^'2A2 | AAA^2>D'2D'D^' F'D^'2A2 |\n",
        "AA2<A^2D' D' C'A^A^2 | A^2GG F^2D F^2DB2 |\n",
        "B2G GF^2 F^2<D2 | A^2A^2 GGF^2DF^ |\n",
        "\n",
        "V:3 name=\"Horn in F\"\n",
        "%%octave-default C5\n",
        "A,B,C | B,A,G,D G,G2 G2FD^ | D2>A^,2 A^,CDG2G2F |\n",
        "D^2<D2 A^,FD^D G2G2 | FD^D2>A^,2A^,C DG2G2 |\n",
        "FD^2<D2A^, F D^DD2 | D2DA^, C2A^, C2A^,D2 |\n",
        "D2D B,C2 C2<B,2 | D2D2 DA^,C2A^,C |\n",
        "\n",
        "V:4 name=\"Trombone\"\n",
        "%%octave-default C3\n",
        "D'3 | G3C'2 F2F2 | A^A^C'D'4C'2F2F2 |\n",
        "A^F'D^'D'4 C'2F2 | F2A^A^ C' D'4C'2F2 |\n",
        "F2A^ F'D^'D'4GA | A^C'D'D^' D'F^G D'F^GG |\n",
        "GABBC'D' D^'D'E'F^' D'GGAB | GAA^C' D'D^'D' F^GD'F^ |\n",
        "\"\"\"\n",
        "\n",
        "abc_loss_fn = ABCLoss()\n",
        "loss, loss_dict = abc_loss_fn(predicted_abc, ground_truth_abc)\n",
        "\n",
        "print(f\"loss: {loss}\")\n",
        "print(f\"Detailed Losses: {loss_dict}\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hlpJ7NCa2_RJ"
      },
      "source": [
        "#### Custom loss tests"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mnOcivpC27-I"
      },
      "outputs": [],
      "source": [
        "import unittest\n",
        "import torch\n",
        "\n",
        "class TestABCLoss(unittest.TestCase):\n",
        "    def setUp(self):\n",
        "        self.loss_fn = ABCLoss()\n",
        "\n",
        "    def test_extract_metre(self):\n",
        "        abc_string = \"M:4/4\\nL:1/16\\nK:C\"\n",
        "        self.assertEqual(self.loss_fn.extract_metre(abc_string), \"4/4\")\n",
        "\n",
        "        abc_string = \"M:3/4\\nL:1/16\\nK:C\"\n",
        "        self.assertEqual(self.loss_fn.extract_metre(abc_string), \"3/4\")\n",
        "\n",
        "    def test_extract_tempo(self):\n",
        "        abc_string = \"Q:120\\nM:4/4\\nL:1/16\\nK:C\"\n",
        "        self.assertEqual(self.loss_fn.extract_tempo(abc_string), \"120\")\n",
        "\n",
        "        # Test when tempo is missing\n",
        "        abc_string = \"M:4/4\\nL:1/16\\nK:C\"\n",
        "        self.assertIsNone(self.loss_fn.extract_tempo(abc_string))\n",
        "\n",
        "    def test_compute_metre_loss(self):\n",
        "        # Test exact match\n",
        "        loss = self.loss_fn.compute_metre_loss(\"4/4\", \"4/4\")\n",
        "        self.assertEqual(loss.item(), 0.0)\n",
        "\n",
        "        # Test similar metres\n",
        "        loss = self.loss_fn.compute_metre_loss(\"3/4\", \"4/4\")\n",
        "        self.assertLess(loss.item(), 1.0)\n",
        "        self.assertGreater(loss.item(), 0.0)\n",
        "\n",
        "        # Test very different metres\n",
        "        loss = self.loss_fn.compute_metre_loss(\"2/4\", \"6/8\")\n",
        "        self.assertLess(loss.item(), 1.0)\n",
        "        self.assertGreater(loss.item(), 0.0)\n",
        "\n",
        "    def test_compute_tempo_loss(self):\n",
        "        # Test exact match\n",
        "        loss = self.loss_fn.compute_tempo_loss(\"120\", \"120\")\n",
        "        self.assertEqual(loss.item(), 0.0)\n",
        "\n",
        "        # Test small difference\n",
        "        loss = self.loss_fn.compute_tempo_loss(\"125\", \"120\")\n",
        "        self.assertLess(loss.item(), 0.5)\n",
        "\n",
        "        # Test large difference\n",
        "        loss = self.loss_fn.compute_tempo_loss(\"200\", \"120\")\n",
        "        self.assertGreater(loss.item(), 0.5)\n",
        "\n",
        "    def test_compute_key_loss(self):\n",
        "        # Test exact match\n",
        "        loss = self.loss_fn.compute_key_loss(\"C\", \"C\")\n",
        "        self.assertEqual(loss.item(), 0.0)\n",
        "\n",
        "        # Test adjacent keys on circle of fifths\n",
        "        loss = self.loss_fn.compute_key_loss(\"C\", \"G\")\n",
        "        self.assertLess(loss.item(), 0.5)\n",
        "\n",
        "        # Test opposite keys on circle of fifths\n",
        "        loss = self.loss_fn.compute_key_loss(\"C\", \"F#\")\n",
        "        self.assertGreater(loss.item(), 0.5)\n",
        "\n",
        "    def test_instrument_name_to_program(self):\n",
        "        # Test common instruments\n",
        "        self.assertGreaterEqual(self.loss_fn.instrument_name_to_program(\"Piano\"), 0)\n",
        "        self.assertGreaterEqual(self.loss_fn.instrument_name_to_program(\"Guitar\"), 0)\n",
        "\n",
        "        # Test invalid instrument\n",
        "        self.assertEqual(self.loss_fn.instrument_name_to_program(\"InvalidInstrument\"), -1)\n",
        "\n",
        "    def test_complete_loss_computation(self):\n",
        "        simple_pred = \"\"\"X:1\n",
        "M:4/4\n",
        "L:1/16\n",
        "K:C\n",
        "Q:120\n",
        "V:1 name=\"Piano\"\n",
        "CDEF|\"\"\"\n",
        "\n",
        "        simple_gt = \"\"\"X:1\n",
        "M:4/4\n",
        "L:1/16\n",
        "K:C\n",
        "Q:120\n",
        "V:1 name=\"Piano\"\n",
        "CDEF|\"\"\"\n",
        "\n",
        "        total_loss, loss_dict = self.loss_fn(simple_pred, simple_gt)\n",
        "\n",
        "        # Check if loss is there\n",
        "        self.assertIsInstance(total_loss, torch.Tensor)\n",
        "        self.assertGreaterEqual(total_loss.item(), 0.0)\n",
        "        self.assertLessEqual(total_loss.item(), 1.0)\n",
        "\n",
        "        # Check if loss dictionary contains all components\n",
        "        expected_keys = ['Total Loss', 'Pitch Loss', 'Metre Loss',\n",
        "                        'Tempo Loss', 'Instrument Loss']\n",
        "        for key in expected_keys:\n",
        "            self.assertIn(key, loss_dict)\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    unittest.main()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1aSB-4zmfF2f"
      },
      "source": [
        "#### Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hyhrHBgAMZNF"
      },
      "source": [
        "**Ensure trl is on version 0.11.4 as the PPOTrainer (v1) is deprecated from 0.12.0 onward!**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XtbULJFrF1C8"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoTokenizer, AutoProcessor\n",
        "from transformers import Qwen2AudioForConditionalGeneration\n",
        "from trl import AutoModelForCausalLMWithValueHead\n",
        "import torch.nn as nn\n",
        "\n",
        "class Qwen2AudioForPPO(AutoModelForCausalLMWithValueHead):\n",
        "    def __init__(self, pretrained_model):\n",
        "        # We create a wrapper that exposes the lm_head and includes necessary attributes\n",
        "        class LanguageModelWrapper(nn.Module):\n",
        "            def __init__(self, language_model):\n",
        "                super().__init__()\n",
        "                self.language_model = language_model\n",
        "                self.lm_head = language_model.lm_head  # Expose the lm_head\n",
        "                self.config = language_model.config  # Include the config\n",
        "                # Include other necessary attributes\n",
        "                self.prepare_inputs_for_generation = language_model.prepare_inputs_for_generation\n",
        "                self.get_output_embeddings = language_model.get_output_embeddings\n",
        "                self.get_input_embeddings = language_model.get_input_embeddings\n",
        "\n",
        "            def forward(self, *args, **kwargs):\n",
        "                return self.language_model(*args, **kwargs)\n",
        "\n",
        "        # Pass the wrapped language model to the parent class\n",
        "        wrapped_lm = LanguageModelWrapper(pretrained_model.language_model)\n",
        "        super().__init__(wrapped_lm)\n",
        "        self.pretrained_model = pretrained_model\n",
        "\n",
        "    def forward(self, input_features=None, feature_attention_mask=None, input_ids=None, attention_mask=None, **kwargs):\n",
        "        outputs = self.pretrained_model(\n",
        "            input_features=input_features,\n",
        "            feature_attention_mask=feature_attention_mask,\n",
        "            input_ids=input_ids,\n",
        "            attention_mask=attention_mask,\n",
        "            output_hidden_states=True,\n",
        "            **kwargs\n",
        "        )\n",
        "\n",
        "        last_hidden_state = outputs.hidden_states[-1]\n",
        "        value = self.v_head(last_hidden_state).squeeze(-1)\n",
        "\n",
        "        return outputs.logits, outputs.hidden_states, value\n",
        "\n",
        "    def generate(self, *args, **kwargs):\n",
        "        return self.pretrained_model.generate(*args, **kwargs)\n",
        "\n",
        "processor = AutoProcessor.from_pretrained(\"Qwen/Qwen2-Audio-7B-Instruct\")\n",
        "\n",
        "model_for_ppo = Qwen2AudioForPPO(model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3pMwm_jIiwxl"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torch import nn\n",
        "from transformers import (\n",
        "    Qwen2AudioForConditionalGeneration,\n",
        "    AutoProcessor,\n",
        "    PreTrainedModel,\n",
        ")\n",
        "from trl import PPOConfig, PPOTrainer, AutoModelForCausalLMWithValueHead\n",
        "from torch.utils.data import DataLoader, Dataset\n",
        "from transformers import AdamW\n",
        "import bitsandbytes as bnb\n",
        "\n",
        "ppo_config = PPOConfig(\n",
        "    model_name=\"Qwen/Qwen2-Audio-7B-Instruct\",\n",
        "    learning_rate=5e-5,          # Other options: [5e-5]\n",
        "    batch_size=8,                # Other options: [16, 32]\n",
        "    mini_batch_size=4,           # Should divide batch_size; other options: [2, 8]\n",
        "    gradient_accumulation_steps=1,  # Increase if batch_size is too large\n",
        "    ppo_epochs=4,                # Other options: [5]\n",
        "    max_grad_norm=10.0,           # Other options: [0.5]\n",
        "    cliprange=0.2,               # Other options: [0.1]\n",
        "    # gradient_checkpointing=True,  # Uncomment to save memory\n",
        "    # seed=3407,\n",
        "    log_with='wandb',\n",
        ")\n",
        "\n",
        "optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=ppo_config.learning_rate, weight_decay=0.001)\n",
        "\n",
        "def your_data_collator(batch):\n",
        "    queries = [item[1] for item in batch]\n",
        "    audios = [item[0] for item in batch]\n",
        "    answers = [item[2] for item in batch]\n",
        "    return {\"queries\": queries, \"audios\": audios, \"answers\": answers}\n",
        "\n",
        "ppo_trainer = PPOTrainer(\n",
        "    config=ppo_config,\n",
        "    model=model_for_ppo,\n",
        "    tokenizer=processor.tokenizer,\n",
        "    dataset=urmp_dataset,\n",
        "    data_collator=your_data_collator,\n",
        ")\n",
        "\n",
        "abc_loss_fn = ABCLoss()\n",
        "\n",
        "def compute_rewards(answers, responses):\n",
        "    rewards = []\n",
        "    for i in range(0, len(responses)):\n",
        "        loss, loss_dict = abc_loss_fn(answers[i], responses[i])\n",
        "        reward = torch.tensor(1 - loss, device=inputs['input_ids'].device, dtype=torch.float32)\n",
        "        rewards.append(reward)\n",
        "\n",
        "    print(f\"rewards {rewards}\")\n",
        "    return rewards\n",
        "\n",
        "import json\n",
        "\n",
        "def log_to_file(predicted_responses, target_responses, rewards, file_path='ppo_training_log.txt'):\n",
        "    with open(file_path, 'a') as f:\n",
        "        for pred, target, reward in zip(predicted_responses, target_responses, rewards):\n",
        "            log_entry = {\n",
        "                'predicted_response': pred,\n",
        "                'target_response': target,\n",
        "                'reward': reward.item() if isinstance(reward, torch.Tensor) else reward\n",
        "            }\n",
        "            f.write(json.dumps(log_entry) + '\\n')\n",
        "\n",
        "for epoch in range(ppo_config.ppo_epochs):\n",
        "    for batch in ppo_trainer.dataloader:\n",
        "        queries = batch[\"queries\"]\n",
        "        audios = batch[\"audios\"]\n",
        "        answers = batch[\"answers\"]\n",
        "\n",
        "        texts = []\n",
        "        for query in queries:\n",
        "            conversation = [\n",
        "                {'role': 'system', 'content': 'You are a helpful assistant.'},\n",
        "                {\"role\": \"user\", \"content\": [\n",
        "                    {\"type\": \"audio\", \"audio_url\": \"https://example.com/audio.mp3\"},\n",
        "                    {\"type\": \"text\", \"text\": query},\n",
        "                ]},\n",
        "            ]\n",
        "            text = processor.apply_chat_template(\n",
        "                conversation, tokenize=False, add_generation_prompt=True\n",
        "            )\n",
        "            texts.append(text)\n",
        "\n",
        "        inputs = processor(\n",
        "            text=texts,\n",
        "            audios=audios,\n",
        "            return_tensors=\"pt\",\n",
        "            padding=True,\n",
        "            sampling_rate=16000,\n",
        "        )\n",
        "\n",
        "        inputs = {k: v.to('cuda') for k, v in inputs.items()}\n",
        "\n",
        "        with torch.no_grad():\n",
        "            response_ids = model_for_ppo.generate(\n",
        "                **inputs,\n",
        "                do_sample=True, top_p=0.8, temperature=1, max_length=2048, use_cache=True\n",
        "            )\n",
        "\n",
        "        generated_responses = response_ids[:, inputs['input_ids'].size(1):]\n",
        "\n",
        "        # Decode responses\n",
        "        responses = processor.tokenizer.batch_decode(generated_responses, skip_special_tokens=True)\n",
        "\n",
        "        print(f\"responses: {responses}\")\n",
        "\n",
        "        # Prepare query and response tensors\n",
        "        # Tokenize queries separately to get their token lengths\n",
        "        tokenized_queries = processor.tokenizer(queries, return_tensors='pt', padding=True)\n",
        "        query_lengths = [len(q) for q in tokenized_queries['input_ids']]\n",
        "\n",
        "        query_tensors = [input_ids[:ql] for input_ids, ql in zip(inputs['input_ids'], query_lengths)]\n",
        "        response_tensors = [ids[ql:] for ids, ql in zip(response_ids, query_lengths)]\n",
        "\n",
        "        rewards = compute_rewards(answers, responses)\n",
        "\n",
        "        log_to_file(responses, answers, rewards)\n",
        "\n",
        "        with torch.amp.autocast(device_type='cuda'):\n",
        "            stats = ppo_trainer.step(\n",
        "                queries=query_tensors,\n",
        "                responses=response_tensors,\n",
        "                scores=rewards\n",
        "            )\n",
        "\n",
        "            ppo_trainer.log_stats(stats, batch, rewards, columns_to_log=['queries', 'answers'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EY3lc3LaMPcH"
      },
      "outputs": [],
      "source": [
        "model.save_pretrained(\"/content/qwen_audio_finetune_3\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4_8ZoqtDGeYf"
      },
      "outputs": [],
      "source": [
        "!cp -r \"/content/qwen_audio_finetune_3\" \"/content/drive/My Drive/automatic-music-transcription/saved_models/\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4EbkJr4Z9EOg"
      },
      "source": [
        "## Try model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "edmrLWd0PqH5"
      },
      "source": [
        "### Test model manually"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Rx_Gjg-XXf8l"
      },
      "outputs": [],
      "source": [
        "!cp -r \"/content/drive/My Drive/automatic-music-transcription/saved_models/qwen_audio_finetune_2\" \"/content/\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VAuZBmt8aROK"
      },
      "outputs": [],
      "source": [
        "from transformers import TrainingArguments, Qwen2AudioForConditionalGeneration\n",
        "from peft import get_peft_model, LoraConfig, TaskType\n",
        "from trl import SFTTrainer\n",
        "import torch\n",
        "\n",
        "model = Qwen2AudioForConditionalGeneration.from_pretrained(\"Qwen/Qwen2-Audio-7B-Instruct\", torch_dtype=torch.bfloat16, device_map=\"auto\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UFpYW-obwqhr"
      },
      "outputs": [],
      "source": [
        "peft_config = LoraConfig(task_type=TaskType.CAUSAL_LM, inference_mode=True, r=128, lora_alpha=256, lora_dropout=0, target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"])\n",
        "model = get_peft_model(model, peft_config)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GHMCyZRmXgB7"
      },
      "outputs": [],
      "source": [
        "model.load_adapter(\"/content/qwen_audio_finetune_2\", adapter_name=\"test\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "34f6laoUcD06"
      },
      "outputs": [],
      "source": [
        "urmp_dataset = URMPDataset(root_dir=\"/content/URMP_Dataset/Dataset\", dataset_name=\"urmp\", return_audio=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8b_dZIgouThM"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoProcessor\n",
        "\n",
        "processor = AutoProcessor.from_pretrained(\"Qwen/Qwen2-Audio-7B-Instruct\" ,trust_remote_code=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jOc-ZuuVcUS_"
      },
      "outputs": [],
      "source": [
        "audio = urmp_dataset[22][0]\n",
        "query = urmp_dataset[22][1]\n",
        "answer = urmp_dataset[22][2]\n",
        "query, answer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RmHQxgK1arbY"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoProcessor\n",
        "\n",
        "processor = AutoProcessor.from_pretrained(\"Qwen/Qwen2-Audio-7B-Instruct\" ,trust_remote_code=True)\n",
        "\n",
        "conversation = [\n",
        "    {'role': 'system', 'content': 'You are a helpful assistant.'},\n",
        "    {\"role\": \"user\", \"content\": [\n",
        "        # audio_url is required despite irrelevant!\n",
        "        {\"type\": \"audio\", \"audio_url\": \"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-Audio/audio/glass-breaking-151256.mp3\"},\n",
        "        {\"type\": \"text\", \"text\": query},\n",
        "    ]},\n",
        "]\n",
        "\n",
        "text = processor.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)\n",
        "texts = [text]\n",
        "audios = [audio]\n",
        "\n",
        "inputs = processor(text=text, audios=audios, return_tensors=\"pt\", padding=True, sampling_rate=16000)\n",
        "inputs.to(\"cuda\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ItU29mHcc3xU"
      },
      "outputs": [],
      "source": [
        "generate_ids = model.generate(**inputs, do_sample=True, top_p=0.8, temperature=1.3, max_length=2048, use_cache=True)\n",
        "generate_ids = generate_ids[:, inputs.input_ids.size(1):]\n",
        "\n",
        "response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]\n",
        "response"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "erebhZzjPwuC"
      },
      "source": [
        "### Run model over evaluation dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OfPlGsARpXsE"
      },
      "outputs": [],
      "source": [
        "import locale\n",
        "locale.getpreferredencoding = lambda: \"UTF-8\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1CHdjvRPQJQc"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "import torch\n",
        "import datetime\n",
        "\n",
        "# Set random seed for reproducibility\n",
        "random.seed(42)\n",
        "\n",
        "# Generate 30 random unique indices\n",
        "dataset_size = len(urmp_dataset)\n",
        "random_indices = random.sample(range(dataset_size), 30)\n",
        "\n",
        "# Create a timestamp for the log file\n",
        "timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
        "log_file = f\"/content/qwen2audio_original_predictions_{timestamp}.txt\"\n",
        "\n",
        "def run_inference(audio, query):\n",
        "    try:\n",
        "        conversation = [\n",
        "            {'role': 'system', 'content': 'You are a helpful assistant.'},\n",
        "            {\"role\": \"user\", \"content\": [\n",
        "                {\"type\": \"audio\", \"audio_url\": \"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-Audio/audio/glass-breaking-151256.mp3\"},\n",
        "                {\"type\": \"text\", \"text\": query},\n",
        "            ]},\n",
        "        ]\n",
        "\n",
        "        text = processor.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)\n",
        "        inputs = processor(text=text, audios=[audio], return_tensors=\"pt\", padding=True, sampling_rate=16000)\n",
        "        inputs = inputs.to(\"cuda\")\n",
        "\n",
        "        generate_ids = model.generate(**inputs, do_sample=True, top_p=0.7, temperature=1.3, max_length=2048, use_cache=True)\n",
        "        generate_ids = generate_ids[:, inputs.input_ids.size(1):]\n",
        "\n",
        "        prediction = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]\n",
        "        return prediction, None\n",
        "    except Exception as e:\n",
        "        return None, str(e)\n",
        "\n",
        "# Process each sample and write to log file\n",
        "with open(log_file, 'w', encoding='utf-8') as f:\n",
        "    for sample_num, idx in enumerate(random_indices, 1):\n",
        "        print(f\"Processing sample {idx} ({sample_num}/20)\")\n",
        "\n",
        "        try:\n",
        "            # Get sample data\n",
        "            audio, query, answer, _, _, _ = urmp_dataset[idx]\n",
        "\n",
        "            # Write sample information\n",
        "            f.write(f\"Sample {sample_num} (Index: {idx})\\n\")\n",
        "            f.write(f\"Question: {query}\\n\")\n",
        "            f.write(f\"Ground Truth: {answer}\\n\")\n",
        "\n",
        "            # Run inference 3 times\n",
        "            f.write(\"Model Predictions:\\n\")\n",
        "            for i in range(3):\n",
        "                print(f\"  Running prediction {i+1}/3 for sample {idx}\")\n",
        "                prediction, error = run_inference(audio, query)\n",
        "\n",
        "                if prediction is not None:\n",
        "                    f.write(f\"Prediction {i+1}: {prediction}\\n\")\n",
        "                else:\n",
        "                    f.write(f\"Prediction {i+1} failed: {error}\\n\")\n",
        "\n",
        "                # Flush the file after each prediction\n",
        "                f.flush()\n",
        "\n",
        "                # Clear CUDA cache\n",
        "                torch.cuda.empty_cache()\n",
        "\n",
        "            torch.cuda.empty_cache()\n",
        "\n",
        "            f.write(\"\\n\" + \"=\"*30 + \"\\n\\n\")\n",
        "            # Flush the file after each sample\n",
        "            f.flush()\n",
        "\n",
        "        except Exception as e:\n",
        "            print(f\"Error processing sample {idx}: {str(e)}\")\n",
        "            f.write(f\"Error processing sample {idx}: {str(e)}\\n\")\n",
        "            f.write(\"\\n\" + \"=\"*30 + \"\\n\\n\")\n",
        "            f.flush()\n",
        "            continue\n",
        "\n",
        "print(f\"Results have been logged to: {log_file}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JpH7EUr0zogQ"
      },
      "outputs": [],
      "source": [
        "!cp \"{log_file}\" \"/content/drive/My Drive/automatic-music-transcription/logs/\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "DZJiYK0Kzz3q"
      },
      "outputs": [],
      "source": [
        "from google.colab import runtime\n",
        "runtime.unassign()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "machine_shape": "hm",
      "provenance": [],
      "toc_visible": true,
      "gpuType": "L4"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}