Added new traing data

.ipynb_checkpoints/HuggingFace_Mistral_Transformer_Single_Instrument-checkpoint.ipynb

@@ -0,0 +1,803 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "SiTIpPjArIyr"
+   },
+   "source": [
+    "# Using Midi traning data and MidiTok Remi to generate music with Mistral model \n",
+    "# split music into Single Instrument and split into 1024\n"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "gOd93yV0sGd2"
+   },
+   "source": [
+    "## Setup Environment"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "To compile Symusic \n",
+    "\n",
+    "Get g++11 or higher\n",
+    "\n",
+    "git clone --recursive https://github.com/Yikai-Liao/symusic\n",
+    "CXX=/usr/bin/g++-11 pip install ./symusic\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install torch==2.6.0\n",
+    "%pip install evaluate  transformers[torch]==4.55.4 tqdm miditok  accelerate tensorboardX scikit-learn\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {
+    "cellView": "form",
+    "id": "fX12Yquyuihc"
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2025-09-12 09:17:37.410013: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
+      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
+      "2025-09-12 09:17:38.509451: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "\n",
+    "from copy import deepcopy\n",
+    "from pathlib import Path\n",
+    "from random import shuffle, sample\n",
+    "\n",
+    "from evaluate import load as load_metric\n",
+    "from miditok import REMI, TokenizerConfig, TokTrainingIterator\n",
+    "from miditok.pytorch_data import DatasetMIDI, DataCollator\n",
+    "from miditok.utils import split_files_for_training\n",
+    "\n",
+    "from miditok.data_augmentation import augment_dataset\n",
+    "from torch import Tensor, argmax, torch\n",
+    "from torch.utils.data import DataLoader\n",
+    "from torch.cuda import is_available as cuda_available, is_bf16_supported\n",
+    "from torch.backends.mps import is_available as mps_available\n",
+    "from transformers import AutoModelForCausalLM, MistralConfig, Trainer, TrainingArguments, GenerationConfig, AutoConfig\n",
+    "from transformers.trainer_utils import set_seed\n",
+    "from tqdm import tqdm"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Setup Tokenizer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Seed\n",
+    "set_seed(777)\n",
+    "\n",
+    "# Our tokenizer's configuration\n",
+    "BEAT_RES = {(0, 1): 12, (1, 2): 4, (2, 4): 2, (4, 8): 1}\n",
+    "TOKENIZER_PARAMS = {\n",
+    "    \"pitch_range\": (21, 108),\n",
+    "    \"beat_res\": BEAT_RES,\n",
+    "    \"num_velocities\": 32,\n",
+    "    \"special_tokens\": [\"PAD\", \"BOS\", \"EOS\"],\n",
+    "    \"use_chords\": True,\n",
+    "    \"use_rests\": True,\n",
+    "    \"use_tempos\": True,\n",
+    "    \"use_time_signatures\": True,\n",
+    "    \"use_programs\": False,  # We want single track \n",
+    "    \"one_token_stream_for_programs\": False, # We want single track\n",
+    "    \"programs\": list(range(0, 128)), #-1 drums, skip drums\n",
+    "    \"num_tempos\": 32,\n",
+    "    \"tempo_range\": (40, 250),  # (min_tempo, max_tempo)\n",
+    "}\n",
+    "config = TokenizerConfig(**TOKENIZER_PARAMS)\n",
+    "\n",
+    "# Creates the tokenizer REMI PLUS\n",
+    "tokenizer = REMI(config)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Load Midi filed and train the the tokenizer on the midi files"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "root_data_dir = Path('/home/wombat/Documents/projects/music/midiTok/data/')\n",
+    "root_save = Path(root_data_dir / 'HuggingFace_Mistral_Transformer_Single_Instrument')\n",
+    "\n",
+    "tokenizer_name = \"HuggingFace_Mistral_Transformer_Single_Instrument_v4_single_track.json\"\n",
+    "dataset_dir = root_save / \"data\"\n",
+    "dataset_dir.mkdir(parents=True, exist_ok=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "# Trains the tokenizer with Byte Pair Encoding (BPE) to build the vocabulary, here 30k tokens\n",
+    "#data_dirs = [\"adl-piano-midi\", \"maestro-v3.0.0\", \"musicnet_midis\" ] # for single \n",
+    "data_dirs = [\"MIDIs\"]\n",
+    "midi_paths = []\n",
+    "for data_dir in data_dirs:\n",
+    "    path = Path(root_data_dir / 'Traning Data' / data_dir)\n",
+    "    midi_paths.extend(list(path.resolve().glob(\"**/*.mid\")) + list(path.resolve().glob(\"**/*.midi\")))\n",
+    "\n",
+    "print(f\"Found {len(midi_paths)} MIDI files\")\n",
+    "\n",
+    "shuffle(midi_paths)\n",
+    "\n",
+    "# We need a subset of files otherwise training tokenizer takes too long\n",
+    "percentage_to_select = 0.15\n",
+    "num_files_to_select = int(len(midi_paths) * percentage_to_select)\n",
+    "\n",
+    "subset_midi_paths = sample(midi_paths, num_files_to_select)\n",
+    "print(f\"Found {len(subset_midi_paths)} MIDI files\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#Note the size of the dataset is quite large, so it requires a huge amount of memory to train the tokenizer for 61749 files it took 64gb of memory\n",
+    "tokenizer.train(\n",
+    "    vocab_size=24000,\n",
+    "    files_paths=subset_midi_paths,\n",
+    ")\n",
+    "tokenizer.save(root_save / tokenizer_name)\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tokenizer = REMI(params=Path(root_save / tokenizer_name))\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Prepare MIDIs for training\n",
+    "\n",
+    "Here we split the files in three subsets: train, validation and test.\n",
+    "Then data augmentation is performed on each subset independently, and the MIDIs are split into smaller chunks that make approximately the desired token sequence length for training."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sequence_length = 1024  # The maximum sequence length for data samples.\n",
+    "kwargs_dataset = {\"max_seq_len\": sequence_length, \"tokenizer\": tokenizer, \"bos_token_id\": tokenizer[\"BOS_None\"], \"eos_token_id\": tokenizer[\"EOS_None\"], \"pre_tokenize\": True, \"pre_tokenize_thread_count\": 7}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Test splitting files for training and testing purposes"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "# Split will need to add the BPM to the files its split\n",
+    "# \n",
+    "file_paths_test = [\n",
+    "    Path('/media/wombat/c6928dc9-ba03-411d-9483-8e28df5973b9/Music Data/Traning Data/clean_midi/Fatboy Slim/Right Here, Right Now.mid'),\n",
+    "    Path('/media/wombat/c6928dc9-ba03-411d-9483-8e28df5973b9/Music Data/Traning Data/clean_midi/Fatboy Slim/Praise You.mid'),\n",
+    "    Path('/media/wombat/c6928dc9-ba03-411d-9483-8e28df5973b9/Music Data/Traning Data/clean_midi/Goo Goo Dolls/Iris.mid'),\n",
+    "    Path('/media/wombat/c6928dc9-ba03-411d-9483-8e28df5973b9/Music Data/Traning Data/clean_midi/Goo Goo Dolls/Slide.mid'),\n",
+    "    Path('/media/wombat/c6928dc9-ba03-411d-9483-8e28df5973b9/Music Data/Traning Data/clean_midi/James Brown/Sex Machine (Get Up I Feel Like Being A).mid'),\n",
+    "    Path('/media/wombat/c6928dc9-ba03-411d-9483-8e28df5973b9/Music Data/Traning Data/clean_midi/Jamiroquai/Virtual Insanity.1.mid'),\n",
+    "    Path('/media/wombat/c6928dc9-ba03-411d-9483-8e28df5973b9/Music Data/Traning Data/clean_midi/Jamiroquai/Virtual Insanity.mid')\n",
+    "]\n",
+    "\n",
+    "split_files_for_training(\n",
+    "    files_paths=file_paths_test,\n",
+    "    tokenizer=tokenizer,\n",
+    "    save_dir=Path('/home/wombat/Documents/projects/music/midiTok/data/HuggingFace_Mistral_Transformer_Single_Instrument/test'),\n",
+    "    max_seq_len=sequence_length,\n",
+    "    num_overlap_bars=2,\n",
+    "    skip_drums=True\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Split MIDI paths in train/valid/test sets\n",
+    "total_num_files = len(midi_paths)\n",
+    "\n",
+    "num_files_valid = round(total_num_files * 0.15)\n",
+    "num_files_test = round(total_num_files * 0.15)\n",
+    "shuffle(midi_paths)\n",
+    "midi_paths_valid = midi_paths[:num_files_valid]\n",
+    "midi_paths_test = midi_paths[num_files_valid:num_files_valid + num_files_test]\n",
+    "midi_paths_train = midi_paths[num_files_valid + num_files_test:]\n",
+    "\n",
+    "\n",
+    "\n",
+    "# Chunk MIDIs and perform data augmentation on each subset independently\n",
+    "for files_paths, subset_name in (\n",
+    "    (midi_paths_train, \"train\"), (midi_paths_valid, \"valid\"), (midi_paths_test, \"test\")\n",
+    "):\n",
+    "\n",
+    "    # Split the MIDIs into chunks of sizes approximately about 1024 tokens\n",
+    "    subset_chunks_dir = root_save / f\"Maestro_{subset_name}\"\n",
+    "    print(subset_chunks_dir)\n",
+    "    split_files_for_training(\n",
+    "        files_paths=files_paths,\n",
+    "        tokenizer=tokenizer,\n",
+    "        save_dir=subset_chunks_dir,\n",
+    "        max_seq_len=sequence_length,\n",
+    "        num_overlap_bars=2,\n",
+    "        skip_drums=True\n",
+    "    )\n",
+    "\n",
+    "    if subset_name == 'train':\n",
+    "        print(\"Augmentation\")\n",
+    "        # Perform data augmentation\n",
+    "        augment_dataset(\n",
+    "            subset_chunks_dir,\n",
+    "            pitch_offsets=[-12, 12],\n",
+    "            velocity_offsets=[-4, 4],\n",
+    "            duration_offsets=[-0.5, 0.5],\n",
+    "        )\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#Since the datasets are too large after splitting we only want 50% of the split data to train against\n",
+    "sample_subset_per = .5"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "TypeError",
+     "evalue": "slice indices must be integers or None or have an __index__ method",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[8], line 3\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# Create Dataset and Collator for training\u001b[39;00m\n\u001b[1;32m      2\u001b[0m midi_paths_train \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m(root_save\u001b[38;5;241m.\u001b[39mjoinpath(Path(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMaestro_train\u001b[39m\u001b[38;5;124m\"\u001b[39m))\u001b[38;5;241m.\u001b[39mglob(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m**/*.mid\u001b[39m\u001b[38;5;124m\"\u001b[39m)) \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mlist\u001b[39m(root_save\u001b[38;5;241m.\u001b[39mjoinpath(Path(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMaestro_train\u001b[39m\u001b[38;5;124m\"\u001b[39m))\u001b[38;5;241m.\u001b[39mglob(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m**/*.midi\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n\u001b[0;32m----> 3\u001b[0m midi_paths_train \u001b[38;5;241m=\u001b[39m \u001b[43mmidi_paths_train\u001b[49m\u001b[43m[\u001b[49m\u001b[43m:\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mlen\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mmidi_paths_train\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43msample_subset_per\u001b[49m\u001b[43m)\u001b[49m\u001b[43m]\u001b[49m\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;28mlen\u001b[39m(midi_paths_train))\n\u001b[1;32m      5\u001b[0m dataset_train \u001b[38;5;241m=\u001b[39m DatasetMIDI(midi_paths_train, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs_dataset)\n",
+      "\u001b[0;31mTypeError\u001b[0m: slice indices must be integers or None or have an __index__ method"
+     ]
+    }
+   ],
+   "source": [
+    "# Create Dataset and Collator for training\n",
+    "midi_paths_train = list(root_save.joinpath(Path(\"Maestro_train\")).glob(\"**/*.mid\")) + list(root_save.joinpath(Path(\"Maestro_train\")).glob(\"**/*.midi\"))\n",
+    "sample_count = (len(midi_paths_train)*sample_subset_per)\n",
+    "midi_paths_train = midi_paths_train[:]\n",
+    "print(len(midi_paths_train))\n",
+    "dataset_train = DatasetMIDI(midi_paths_train, **kwargs_dataset)\n",
+    "torch.save(dataset_train, Path(dataset_dir / \"dataset_train.pt\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "midi_paths_valid = list(root_save.joinpath(Path(\"Maestro_valid\")).glob(\"**/*.mid\")) + list(root_save.joinpath(Path(\"Maestro_valid\")).glob(\"**/*.midi\")) \n",
+    "midi_paths_valid = midi_paths_valid[:(len(midi_paths_valid)*sample_subset_per]\n",
+    "print(len(midi_paths_valid))\n",
+    "dataset_valid = DatasetMIDI(midi_paths_valid, **kwargs_dataset)\n",
+    "torch.save(dataset_valid, Path(dataset_dir / \"dataset_valid.pt\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "midi_paths_test = list(root_save.joinpath(Path(\"Maestro_test\")).glob(\"**/*.mid\")) + list(root_save.joinpath(Path(\"Maestro_test\")).glob(\"**/*.midi\"))\n",
+    "midi_paths_test = midi_paths_test[:(len(midi_paths_test)*sample_subset_per]\n",
+    "print(len(midi_paths_test))\n",
+    "dataset_test = DatasetMIDI(midi_paths_test, **kwargs_dataset)\n",
+    "torch.save(dataset_test, Path(dataset_dir / \"dataset_test.pt\"))\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print (len(midi_paths_train), len(midi_paths_valid), len(midi_paths_test))\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Save and Load datasets"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "dataset_train = torch.load(Path(dataset_dir / \"dataset_train.pt\"), weights_only=False)\n",
+    "dataset_valid = torch.load(Path(dataset_dir / \"dataset_valid.pt\"), weights_only=False)\n",
+    "dataset_test = torch.load(Path(dataset_dir / \"dataset_test.pt\"), weights_only=False)\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pickle\n",
+    "\n",
+    "test_file = open(Path(dataset_dir / \"dataset_test.pickle\"), 'ab')\n",
+    "pickle.dump(dataset_test, test_file)\n",
+    "test_file.close()\n",
+    "\n",
+    "print(dataset_test[0])\n",
+    "\n",
+    "test_file = open(Path(dataset_dir / \"dataset_test.pickle\"), 'rb')\n",
+    "test_pickle = pickle.load(test_file)\n",
+    "print(test_pickle)\n",
+    "print(test_pickle[0])\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Preview files data load and split"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "#testing_files = \n",
+    "preview_files_path = []\n",
+    "for testing_file in testing_files:\n",
+    "    preview_files_path.append(Path(testing_file))\n",
+    "\n",
+    "preview_dir = Path(root_save / \"preview\")\n",
+    "split_files_for_training(\n",
+    "        files_paths=preview_files_path,\n",
+    "        tokenizer=tokenizer,\n",
+    "        save_dir=preview_dir,\n",
+    "        max_seq_len=sequence_length,\n",
+    "        num_overlap_bars=2,\n",
+    "    )\n",
+    "\n",
+    "valid_midi_path = root_save / \"Maestro_valid\"\n",
+    "midi_split_preview = list(valid_midi_path.resolve().glob(\"**/*.mid\")) + list(valid_midi_path.resolve().glob(\"**/*.midi\"))\n",
+    "\n",
+    "print(len(midi_split_preview))\n",
+    "file_name_lookup = []\n",
+    "def func_to_get_labels(p1, p2, p3):\n",
+    "    if p3.name not in file_name_lookup:\n",
+    "        file_name_lookup.append(p3.name)\n",
+    "    return file_name_lookup.index(p3.name)\n",
+    "    \n",
+    "kwargs_dataset = {\"max_seq_len\": sequence_length, \"tokenizer\": tokenizer, \"bos_token_id\": tokenizer[\"BOS_None\"], \"eos_token_id\": tokenizer[\"EOS_None\"], \"func_to_get_labels\" : func_to_get_labels}\n",
+    "dataset_preview = DatasetMIDI(midi_split_preview, **kwargs_dataset)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Model initialization\n",
+    "\n",
+    "We will use the [Mistral implementation of Hugging Face](https://huggingface.co/docs/transformers/model_doc/mistral).\n",
+    "Feel free to explore the documentation and source code to dig deeper.\n",
+    "\n",
+    "**You may need to adjust the model's configuration, the training configuration and the maximum input sequence length (cell above) depending on your hardware.**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Creates model\n",
+    "model_config = MistralConfig(\n",
+    "    vocab_size=len(tokenizer), #from miditok output default 32K\n",
+    "    hidden_size=512, # default 4096\n",
+    "    intermediate_size=2048, # default  14336\n",
+    "    num_hidden_layers=8, # default  32\n",
+    "    num_attention_heads=8, # default  32\n",
+    "    num_key_value_heads=4, # default 8\n",
+    "    sliding_window=256, # default  4096\n",
+    "    max_position_embeddings=8192, #has no effect on the parms count or training just limits the input length  # default 4096*32\n",
+    "    pad_token_id=tokenizer['PAD_None'],\n",
+    "    bos_token_id=tokenizer['BOS_None'],\n",
+    "    eos_token_id=tokenizer['EOS_None'],\n",
+    ")\n",
+    "model = AutoModelForCausalLM.from_config(model_config)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Model training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_dir = root_save / 'run'\n",
+    "model_dir_str = str(model_dir)\n",
+    "print(model_dir)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "metrics = {metric: load_metric(metric) for metric in [\"accuracy\"]}\n",
+    "\n",
+    "def compute_metrics(eval_pred):\n",
+    "    \"\"\"\n",
+    "    Compute metrics for pretraining.\n",
+    "\n",
+    "    Must use preprocess_logits function that converts logits to predictions (argmax or sampling).\n",
+    "\n",
+    "    :param eval_pred: EvalPrediction containing predictions and labels\n",
+    "    :return: metrics\n",
+    "    \"\"\"\n",
+    "    predictions, labels = eval_pred\n",
+    "    not_pad_mask = labels != -100\n",
+    "    labels, predictions = labels[not_pad_mask], predictions[not_pad_mask]\n",
+    "    return metrics[\"accuracy\"].compute(predictions=predictions.flatten(), references=labels.flatten())\n",
+    "\n",
+    "def preprocess_logits(logits: Tensor, _: Tensor) -> Tensor:\n",
+    "    \"\"\"\n",
+    "    Preprocess the logits before accumulating them during evaluation.\n",
+    "\n",
+    "    This allows to significantly reduce the memory usage and make the training tractable.\n",
+    "    \"\"\"\n",
+    "    pred_ids = argmax(logits, dim=-1)  # long dtype\n",
+    "    return pred_ids\n",
+    "\n",
+    "# Create config for the Trainer\n",
+    "USE_CUDA = cuda_available()\n",
+    "print(USE_CUDA)\n",
+    "if not cuda_available():\n",
+    "    FP16 = FP16_EVAL = BF16 = BF16_EVAL = False\n",
+    "elif is_bf16_supported():\n",
+    "    BF16 = BF16_EVAL = True\n",
+    "    FP16 = FP16_EVAL = False\n",
+    "else:\n",
+    "    BF16 = BF16_EVAL = False\n",
+    "    FP16 = FP16_EVAL = True\n",
+    "USE_MPS = not USE_CUDA and mps_available()\n",
+    "training_config = TrainingArguments(\n",
+    "    model_dir_str, False, True, True, False, \"steps\",\n",
+    "    per_device_train_batch_size=24, #76% @ 24 batch size #76% @ 32 batch size try 64 batch size next time \n",
+    "    per_device_eval_batch_size=24, #was 24 now 32\n",
+    "    gradient_accumulation_steps=3, #change this to 4\n",
+    "    eval_accumulation_steps=None,\n",
+    "    eval_steps=1000,\n",
+    "    learning_rate=1e-4,\n",
+    "    weight_decay=0.01,\n",
+    "    max_grad_norm=3.0,\n",
+    "    max_steps=40000,\n",
+    "    lr_scheduler_type=\"cosine_with_restarts\",\n",
+    "    warmup_ratio=0.3,\n",
+    "    log_level=\"debug\",\n",
+    "    logging_strategy=\"steps\",\n",
+    "    logging_steps=20,\n",
+    "    save_strategy=\"steps\",\n",
+    "    save_steps=1000,\n",
+    "    save_total_limit=5,\n",
+    "    no_cuda=not USE_CUDA,\n",
+    "    seed=444,\n",
+    "    fp16=FP16,\n",
+    "    fp16_full_eval=FP16_EVAL,\n",
+    "    bf16=BF16,\n",
+    "    bf16_full_eval=BF16_EVAL,\n",
+    "    load_best_model_at_end=True,\n",
+    "    label_smoothing_factor=0.,\n",
+    "    optim=\"adamw_torch\",\n",
+    "    report_to=[\"tensorboard\"],\n",
+    "    gradient_checkpointing=True,\n",
+    "    dataloader_num_workers=8, #added to fix trashing isssue with the gpu not having enough data to process\n",
+    "    dataloader_pin_memory=True, #we want the dataset in memory\n",
+    "    torch_compile=True #added to speed up \n",
+    "    \n",
+    ")\n",
+    "\n",
+    "collator = DataCollator(tokenizer[\"PAD_None\"], copy_inputs_as_labels=True, pad_on_left=True) #not sure about the pad_on_left, it might get better results\n",
+    "trainer = Trainer(\n",
+    "    model=model,\n",
+    "    args=training_config,\n",
+    "    data_collator=collator,\n",
+    "    train_dataset=dataset_train,\n",
+    "    eval_dataset=dataset_valid,\n",
+    "    compute_metrics=compute_metrics,\n",
+    "    callbacks=None,\n",
+    "    preprocess_logits_for_metrics=preprocess_logits,\n",
+    "    \n",
+    ")\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "del model\n",
+    "torch.cuda.empty_cache()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(model)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%env PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Training\n",
+    "train_result = trainer.train()\n",
+    "trainer.save_model()  # Saves the tokenizer too\n",
+    "trainer.log_metrics(\"train\", train_result.metrics)\n",
+    "trainer.save_metrics(\"train\", train_result.metrics)\n",
+    "trainer.save_state()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.create_model_card(tags=[\"mistral\", \"midi\", \"miditok\", \"music\", \"instrument\"],\n",
+    "                          model_name=\"Mistral_MidiTok_Transformer_Single_Instrument_Small\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "model.hub_model_id = \"adricl/midi_single_instrument_mistral_transformer\"\n",
+    "\n",
+    "model.push_to_hub(commit_message=\"Training Basic Model for Mistral MidiTok Transformer Single Instrument Small\", repo_id=\"adricl/midi_single_instrument_mistral_transformer\",\n",
+    "                 token=\"\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# For Tensorboard tensorboard --logdir runs/"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "config = AutoConfig.from_pretrained(str(model_dir / \"config.json\"))\n",
+    "model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=str(model_dir / \"model.safetensors\"), from_tf=False, config=config)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Generate music"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "cellView": "form",
+    "id": "OaNkGcFo9UP_"
+   },
+   "outputs": [],
+   "source": [
+    "# for single track midi files splits \n",
+    "\n",
+    "gen_results_path = root_save / 'gen_res'\n",
+    "gen_results_path.mkdir(parents=True, exist_ok=True)\n",
+    "generation_config = GenerationConfig(\n",
+    "    max_new_tokens=200,  # extends samples by 200 tokens\n",
+    "    num_beams=1,         # no beam search\n",
+    "    do_sample=True,      # but sample instead\n",
+    "    temperature=0.9,\n",
+    "    top_k=15,\n",
+    "    top_p=0.95,\n",
+    "    epsilon_cutoff=3e-4,\n",
+    "    eta_cutoff=1e-3,\n",
+    "    pad_token_id=tokenizer.pad_token_id,\n",
+    ")\n",
+    "\n",
+    "# Here the sequences are padded to the left, so that the last token along the time dimension\n",
+    "# is always the last token of each seq, allowing to efficiently generate by batch\n",
+    "collator.pad_on_left = True\n",
+    "collator.eos_token = None\n",
+    "dataloader_test = DataLoader(dataset_preview, batch_size=24, collate_fn=collator)\n",
+    "model.eval()\n",
+    "count = 0\n",
+    "for batch in tqdm(dataloader_test, desc='Testing model / Generating results'):  # (N,T)\n",
+    "    print(batch)\n",
+    "    res = model.generate(\n",
+    "        inputs=batch[\"input_ids\"].to(model.device),\n",
+    "        attention_mask=batch[\"attention_mask\"].to(model.device),\n",
+    "        generation_config=generation_config)  # (N,T)\n",
+    "\n",
+    "\n",
+    "    # Saves the generated music, as MIDI files and tokens (json)\n",
+    "    for prompt, continuation in zip(batch[\"input_ids\"], res):\n",
+    "        generated = continuation[len(prompt):]\n",
+    "        midi = tokenizer.decode([deepcopy(generated.tolist())])\n",
+    "        tokens = [generated, prompt, continuation]  # list compr. as seqs of dif. lengths\n",
+    "        tokens = [seq.tolist() for seq in tokens]\n",
+    "        for tok_seq in tokens[1:]:\n",
+    "            _midi = tokenizer.decode([deepcopy(tok_seq)])\n",
+    "            midi.tracks.append(_midi.tracks[0])\n",
+    "            \n",
+    "        file_name = file_name_lookup[count]\n",
+    "        print(file_name)\n",
+    "        midi.tracks[0].name = f'Continuation of original sample ({len(generated)} tokens) Original file {file_name}'\n",
+    "        midi.tracks[1].name = f'Original sample ({len(prompt)} tokens)'\n",
+    "        if (len(midi.tracks) > 2):\n",
+    "            midi.tracks[2].name = f'Original sample and continuation'\n",
+    "        midi.dump_midi(gen_results_path / f'{count}_{file_name}.mid')\n",
+    "        tokenizer.save_tokens(tokens, gen_results_path / f'{count}_{file_name}.json') \n",
+    "\n",
+    "        count += 1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(file_name_lookup)"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "collapsed_sections": [],
+   "machine_shape": "hm",
+   "name": "Optimus_VIRTUOSO_Multi_Instrumental_RGA_Edition.ipynb",
+   "private_outputs": true,
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.5"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

HuggingFace_Mistral_Transformer_Single_Instrument.ipynb

@@ -37,32 +37,181 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Requirement already satisfied: pip in /usr/local/lib/python3.11/dist-packages (25.2)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning.\u001b[0m\u001b[33m\n",
+      "\u001b[0m"
+     ]
+    }
+   ],
+   "source": [
+    "%%python -m pip install --upgrade pip\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Requirement already satisfied: evaluate in /usr/local/lib/python3.11/dist-packages (0.4.6)\n",
+      "Requirement already satisfied: transformers in /usr/local/lib/python3.11/dist-packages (4.56.2)\n",
+      "Requirement already satisfied: tqdm in /usr/local/lib/python3.11/dist-packages (4.67.1)\n",
+      "Requirement already satisfied: miditok in /usr/local/lib/python3.11/dist-packages (3.0.6.post1)\n",
+      "Requirement already satisfied: accelerate in /usr/local/lib/python3.11/dist-packages (1.10.1)\n",
+      "Requirement already satisfied: tensorboardX in /usr/local/lib/python3.11/dist-packages (2.6.4)\n",
+      "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.11/dist-packages (1.7.2)\n",
+      "Requirement already satisfied: wandb in /usr/local/lib/python3.11/dist-packages (0.22.0)\n",
+      "Requirement already satisfied: datasets>=2.0.0 in /usr/local/lib/python3.11/dist-packages (from evaluate) (4.1.1)\n",
+      "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.11/dist-packages (from evaluate) (2.1.2)\n",
+      "Requirement already satisfied: dill in /usr/local/lib/python3.11/dist-packages (from evaluate) (0.4.0)\n",
+      "Requirement already satisfied: pandas in /usr/local/lib/python3.11/dist-packages (from evaluate) (2.3.2)\n",
+      "Requirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.11/dist-packages (from evaluate) (2.32.3)\n",
+      "Requirement already satisfied: xxhash in /usr/local/lib/python3.11/dist-packages (from evaluate) (3.5.0)\n",
+      "Requirement already satisfied: multiprocess in /usr/local/lib/python3.11/dist-packages (from evaluate) (0.70.16)\n",
+      "Requirement already satisfied: fsspec>=2021.05.0 in /usr/local/lib/python3.11/dist-packages (from fsspec[http]>=2021.05.0->evaluate) (2024.10.0)\n",
+      "Requirement already satisfied: huggingface-hub>=0.7.0 in /usr/local/lib/python3.11/dist-packages (from evaluate) (0.35.1)\n",
+      "Requirement already satisfied: packaging in /usr/local/lib/python3.11/dist-packages (from evaluate) (24.2)\n",
+      "Requirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from transformers) (3.16.1)\n",
+      "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.11/dist-packages (from transformers) (6.0.2)\n",
+      "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.11/dist-packages (from transformers) (2025.9.18)\n",
+      "Requirement already satisfied: tokenizers<=0.23.0,>=0.22.0 in /usr/local/lib/python3.11/dist-packages (from transformers) (0.22.1)\n",
+      "Requirement already satisfied: safetensors>=0.4.3 in /usr/local/lib/python3.11/dist-packages (from transformers) (0.6.2)\n",
+      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.11/dist-packages (from huggingface-hub>=0.7.0->evaluate) (4.12.2)\n",
+      "Requirement already satisfied: hf-xet<2.0.0,>=1.1.3 in /usr/local/lib/python3.11/dist-packages (from huggingface-hub>=0.7.0->evaluate) (1.1.10)\n",
+      "Requirement already satisfied: symusic>=0.5.0 in /usr/local/lib/python3.11/dist-packages (from miditok) (0.5.8)\n",
+      "Requirement already satisfied: psutil in /usr/local/lib/python3.11/dist-packages (from accelerate) (7.0.0)\n",
+      "Requirement already satisfied: torch>=2.0.0 in /usr/local/lib/python3.11/dist-packages (from accelerate) (2.8.0.dev20250319+cu128)\n",
+      "Requirement already satisfied: protobuf>=3.20 in /usr/local/lib/python3.11/dist-packages (from tensorboardX) (6.32.1)\n",
+      "Requirement already satisfied: scipy>=1.8.0 in /usr/local/lib/python3.11/dist-packages (from scikit-learn) (1.16.2)\n",
+      "Requirement already satisfied: joblib>=1.2.0 in /usr/local/lib/python3.11/dist-packages (from scikit-learn) (1.5.2)\n",
+      "Requirement already satisfied: threadpoolctl>=3.1.0 in /usr/local/lib/python3.11/dist-packages (from scikit-learn) (3.6.0)\n",
+      "Requirement already satisfied: click>=8.0.1 in /usr/local/lib/python3.11/dist-packages (from wandb) (8.3.0)\n",
+      "Requirement already satisfied: gitpython!=3.1.29,>=1.0.0 in /usr/local/lib/python3.11/dist-packages (from wandb) (3.1.45)\n",
+      "Requirement already satisfied: platformdirs in /usr/local/lib/python3.11/dist-packages (from wandb) (4.3.7)\n",
+      "Requirement already satisfied: pydantic<3 in /usr/local/lib/python3.11/dist-packages (from wandb) (2.11.9)\n",
+      "Requirement already satisfied: sentry-sdk>=2.0.0 in /usr/local/lib/python3.11/dist-packages (from wandb) (2.39.0)\n",
+      "Requirement already satisfied: annotated-types>=0.6.0 in /usr/local/lib/python3.11/dist-packages (from pydantic<3->wandb) (0.7.0)\n",
+      "Requirement already satisfied: pydantic-core==2.33.2 in /usr/local/lib/python3.11/dist-packages (from pydantic<3->wandb) (2.33.2)\n",
+      "Requirement already satisfied: typing-inspection>=0.4.0 in /usr/local/lib/python3.11/dist-packages (from pydantic<3->wandb) (0.4.1)\n",
+      "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.11/dist-packages (from requests>=2.19.0->evaluate) (3.4.1)\n",
+      "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.11/dist-packages (from requests>=2.19.0->evaluate) (3.10)\n",
+      "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.11/dist-packages (from requests>=2.19.0->evaluate) (2.3.0)\n",
+      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.11/dist-packages (from requests>=2.19.0->evaluate) (2025.1.31)\n",
+      "Requirement already satisfied: pyarrow>=21.0.0 in /usr/local/lib/python3.11/dist-packages (from datasets>=2.0.0->evaluate) (21.0.0)\n",
+      "Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /usr/local/lib/python3.11/dist-packages (from fsspec[http]>=2021.05.0->evaluate) (3.12.15)\n",
+      "Requirement already satisfied: aiohappyeyeballs>=2.5.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2021.05.0->evaluate) (2.6.1)\n",
+      "Requirement already satisfied: aiosignal>=1.4.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2021.05.0->evaluate) (1.4.0)\n",
+      "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2021.05.0->evaluate) (25.3.0)\n",
+      "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.11/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2021.05.0->evaluate) (1.7.0)\n",
+      "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.11/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2021.05.0->evaluate) (6.6.4)\n",
+      "Requirement already satisfied: propcache>=0.2.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2021.05.0->evaluate) (0.3.2)\n",
+      "Requirement already satisfied: yarl<2.0,>=1.17.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2021.05.0->evaluate) (1.20.1)\n",
+      "Requirement already satisfied: gitdb<5,>=4.0.1 in /usr/local/lib/python3.11/dist-packages (from gitpython!=3.1.29,>=1.0.0->wandb) (4.0.12)\n",
+      "Requirement already satisfied: smmap<6,>=3.0.1 in /usr/local/lib/python3.11/dist-packages (from gitdb<5,>=4.0.1->gitpython!=3.1.29,>=1.0.0->wandb) (5.0.2)\n",
+      "Requirement already satisfied: pySmartDL in /usr/local/lib/python3.11/dist-packages (from symusic>=0.5.0->miditok) (1.3.4)\n",
+      "Requirement already satisfied: sympy>=1.13.3 in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (1.13.3)\n",
+      "Requirement already satisfied: networkx in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (3.4.2)\n",
+      "Requirement already satisfied: jinja2 in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (3.1.4)\n",
+      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.8.61 in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (12.8.61)\n",
+      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.8.57 in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (12.8.57)\n",
+      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.8.57 in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (12.8.57)\n",
+      "Requirement already satisfied: nvidia-cudnn-cu12==9.8.0.87 in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (9.8.0.87)\n",
+      "Requirement already satisfied: nvidia-cublas-cu12==12.8.3.14 in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (12.8.3.14)\n",
+      "Requirement already satisfied: nvidia-cufft-cu12==11.3.3.41 in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (11.3.3.41)\n",
+      "Requirement already satisfied: nvidia-curand-cu12==10.3.9.55 in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (10.3.9.55)\n",
+      "Requirement already satisfied: nvidia-cusolver-cu12==11.7.2.55 in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (11.7.2.55)\n",
+      "Requirement already satisfied: nvidia-cusparse-cu12==12.5.7.53 in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (12.5.7.53)\n",
+      "Requirement already satisfied: nvidia-cusparselt-cu12==0.6.3 in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (0.6.3)\n",
+      "Requirement already satisfied: nvidia-nccl-cu12==2.25.1 in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (2.25.1)\n",
+      "Requirement already satisfied: nvidia-nvtx-cu12==12.8.55 in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (12.8.55)\n",
+      "Requirement already satisfied: nvidia-nvjitlink-cu12==12.8.61 in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (12.8.61)\n",
+      "Requirement already satisfied: nvidia-cufile-cu12==1.13.0.11 in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (1.13.0.11)\n",
+      "Requirement already satisfied: pytorch-triton==3.3.0+git96316ce5 in /usr/local/lib/python3.11/dist-packages (from torch>=2.0.0->accelerate) (3.3.0+git96316ce5)\n",
+      "Requirement already satisfied: setuptools>=40.8.0 in /usr/local/lib/python3.11/dist-packages (from pytorch-triton==3.3.0+git96316ce5->torch>=2.0.0->accelerate) (77.0.1)\n",
+      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.11/dist-packages (from sympy>=1.13.3->torch>=2.0.0->accelerate) (1.3.0)\n",
+      "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.11/dist-packages (from jinja2->torch>=2.0.0->accelerate) (2.1.5)\n",
+      "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.11/dist-packages (from pandas->evaluate) (2.9.0.post0)\n",
+      "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.11/dist-packages (from pandas->evaluate) (2025.2)\n",
+      "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.11/dist-packages (from pandas->evaluate) (2025.2)\n",
+      "Requirement already satisfied: six>=1.5 in /usr/lib/python3/dist-packages (from python-dateutil>=2.8.2->pandas->evaluate) (1.16.0)\n",
+      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning.\u001b[0m\u001b[33m\n",
+      "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n"
+     ]
+    }
+   ],
+   "source": [
+    "%pip install evaluate transformers tqdm miditok  accelerate tensorboardX scikit-learn wandb"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
    "metadata": {
     "cellView": "form",
     "id": "fX12Yquyuihc"
    },
    "outputs": [],
    "source": [
-    "\n",
-    "\n",
     "from copy import deepcopy\n",
     "from pathlib import Path\n",
-    "from random import shuffle\n",
+    "from random import shuffle, sample\n",
     "\n",
     "from evaluate import load as load_metric\n",
-    "from miditok import REMI, TokenizerConfig, TokTrainingIterator\n",
+    "from miditok import REMI, TokenizerConfig\n",
     "from miditok.pytorch_data import DatasetMIDI, DataCollator\n",
     "from miditok.utils import split_files_for_training\n",
     "\n",
     "from miditok.data_augmentation import augment_dataset\n",
-    "from torch import Tensor, argmax\n",
+    "from torch import Tensor, argmax, torch\n",
     "from torch.utils.data import DataLoader\n",
     "from torch.cuda import is_available as cuda_available, is_bf16_supported\n",
     "from torch.backends.mps import is_available as mps_available\n",
     "from transformers import AutoModelForCausalLM, MistralConfig, Trainer, TrainingArguments, GenerationConfig, AutoConfig\n",
     "from transformers.trainer_utils import set_seed\n",
-    "from tqdm import tqdm"
+    "from tqdm import tqdm\n",
+    "\n",
+    "#Seed\n",
+    "set_seed(777)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33madric-landman\u001b[0m (\u001b[33madric-landman-hobby\u001b[0m) to \u001b[32mhttps://api.wandb.ai\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Wand db\n",
+    "import wandb\n",
+    "import os\n",
+    "wandb.login() #6bfb401d368c508cb01a291a8ae84c0ecce2310d \n",
+    "\n",
+    "os.environ[\"WANDB_PROJECT\"]=\"midi_music_maker\"\n",
+    "\n"
    ]
   },
   {
@@ -75,13 +224,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Seed\n",
-    "set_seed(777)\n",
-    "\n",
     "# Our tokenizer's configuration\n",
     "BEAT_RES = {(0, 1): 12, (1, 2): 4, (2, 4): 2, (4, 8): 1}\n",
     "TOKENIZER_PARAMS = {\n",
@@ -114,14 +260,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
-    "root_data_dir = Path('/home/wombat/Documents/projects/music/midiTok/data/')\n",
+    "root_data_dir = Path('/workspace/traindata/data')\n",
     "root_save = Path(root_data_dir / 'HuggingFace_Mistral_Transformer_Single_Instrument')\n",
     "\n",
-    "tokenizer_name = \"HuggingFace_Mistral_Transformer_Single_Instrument.json\""
+    "tokenizer_name = \"HuggingFace_Mistral_Transformer_Single_Instrument_v4_single_track.json\"\n",
+    "dataset_dir = root_save / \"data\"\n",
+    "dataset_dir.mkdir(parents=True, exist_ok=True)"
    ]
   },
   {
@@ -132,13 +280,23 @@
    "source": [
     "\n",
     "# Trains the tokenizer with Byte Pair Encoding (BPE) to build the vocabulary, here 30k tokens\n",
-    "data_dirs = [\"adl-piano-midi\", \"maestro-v3.0.0\", \"musicnet_midis\" ] # for single \n",
+    "#data_dirs = [\"adl-piano-midi\", \"maestro-v3.0.0\", \"musicnet_midis\" ] # for single \n",
+    "data_dirs = [\"MIDIs\"]\n",
     "midi_paths = []\n",
     "for data_dir in data_dirs:\n",
     "    path = Path(root_data_dir / 'Traning Data' / data_dir)\n",
     "    midi_paths.extend(list(path.resolve().glob(\"**/*.mid\")) + list(path.resolve().glob(\"**/*.midi\")))\n",
     "\n",
-    "print(f\"Found {len(midi_paths)} MIDI files\")"
+    "print(f\"Found {len(midi_paths)} MIDI files\")\n",
+    "\n",
+    "shuffle(midi_paths)\n",
+    "\n",
+    "# We need a subset of files otherwise training tokenizer takes too long\n",
+    "percentage_to_select = 0.15\n",
+    "num_files_to_select = int(len(midi_paths) * percentage_to_select)\n",
+    "\n",
+    "subset_midi_paths = sample(midi_paths, num_files_to_select)\n",
+    "print(f\"Found {len(subset_midi_paths)} MIDI files\")"
    ]
   },
   {
@@ -149,8 +307,8 @@
    "source": [
     "#Note the size of the dataset is quite large, so it requires a huge amount of memory to train the tokenizer for 61749 files it took 64gb of memory\n",
     "tokenizer.train(\n",
-    "    vocab_size=32000,\n",
-    "    files_paths=midi_paths,\n",
+    "    vocab_size=24000,\n",
+    "    files_paths=subset_midi_paths,\n",
     ")\n",
     "tokenizer.save(root_save / tokenizer_name)\n",
     "\n"
@@ -158,11 +316,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [],
    "source": [
-    "tokenizer = REMI(params=Path(root_save / tokenizer_name))"
+    "tokenizer = REMI(params=Path(root_save / tokenizer_name))\n"
    ]
   },
   {
@@ -177,12 +335,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [],
    "source": [
     "sequence_length = 1024  # The maximum sequence length for data samples.\n",
-    "kwargs_dataset = {\"max_seq_len\": sequence_length, \"tokenizer\": tokenizer, \"bos_token_id\": tokenizer[\"BOS_None\"], \"eos_token_id\": tokenizer[\"EOS_None\"]}"
+    "kwargs_dataset = {\"max_seq_len\": sequence_length, \"tokenizer\": tokenizer, \"bos_token_id\": tokenizer[\"BOS_None\"], \"eos_token_id\": tokenizer[\"EOS_None\"], \"pre_tokenize\": True, \"pre_tokenize_thread_count\": 7}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Test splitting files for training and testing purposes"
    ]
   },
   {
@@ -191,30 +356,27 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def remove_files_with_boring_data( file_paths: list[Path], rms_threshold: float = 0.01) -> list[Path]:\n",
-    "    \"\"\"\n",
-    "    Remove files with boring data, i.e. files with low RMS.\n",
-    "    \"\"\"\n",
-    "    from symusic import Score\n",
-    "    from tqdm import tqdm\n",
-    "    import numpy as np\n",
-    "\n",
-    "    rms = lambda data: (sum(x * x for x in data) / len(data)) ** 0.5\n",
-    "\n",
-    "    filtered_files = []\n",
-    "    for file_path in tqdm(file_paths, desc=\"Filtering boring files\"):\n",
-    "        try:\n",
-    "            scores = [Score(file_path)]\n",
-    "        except SCORE_LOADING_EXCEPTION:\n",
-    "            continue\n",
-    "\n",
-    "        for track in scores[0].tracks:\n",
-    "            values = track.notes['pitch']\n",
-    "            result = rms(values)\n",
-    "\n",
+    "from pathlib import Path\n",
+    "# Split will need to add the BPM to the files its split\n",
+    "# \n",
+    "file_paths_test = [\n",
+    "    Path('/media/wombat/c6928dc9-ba03-411d-9483-8e28df5973b9/Music Data/Traning Data/clean_midi/Fatboy Slim/Right Here, Right Now.mid'),\n",
+    "    Path('/media/wombat/c6928dc9-ba03-411d-9483-8e28df5973b9/Music Data/Traning Data/clean_midi/Fatboy Slim/Praise You.mid'),\n",
+    "    Path('/media/wombat/c6928dc9-ba03-411d-9483-8e28df5973b9/Music Data/Traning Data/clean_midi/Goo Goo Dolls/Iris.mid'),\n",
+    "    Path('/media/wombat/c6928dc9-ba03-411d-9483-8e28df5973b9/Music Data/Traning Data/clean_midi/Goo Goo Dolls/Slide.mid'),\n",
+    "    Path('/media/wombat/c6928dc9-ba03-411d-9483-8e28df5973b9/Music Data/Traning Data/clean_midi/James Brown/Sex Machine (Get Up I Feel Like Being A).mid'),\n",
+    "    Path('/media/wombat/c6928dc9-ba03-411d-9483-8e28df5973b9/Music Data/Traning Data/clean_midi/Jamiroquai/Virtual Insanity.1.mid'),\n",
+    "    Path('/media/wombat/c6928dc9-ba03-411d-9483-8e28df5973b9/Music Data/Traning Data/clean_midi/Jamiroquai/Virtual Insanity.mid')\n",
+    "]\n",
     "\n",
-    "            filtered_files.append(file_path)\n",
-    "    return filtered_files"
+    "split_files_for_training(\n",
+    "    files_paths=file_paths_test,\n",
+    "    tokenizer=tokenizer,\n",
+    "    save_dir=Path('/home/wombat/Documents/projects/music/midiTok/data/HuggingFace_Mistral_Transformer_Single_Instrument/test'),\n",
+    "    max_seq_len=sequence_length,\n",
+    "    num_overlap_bars=2,\n",
+    "    skip_drums=True\n",
+    ")"
    ]
   },
   {
@@ -225,6 +387,7 @@
    "source": [
     "# Split MIDI paths in train/valid/test sets\n",
     "total_num_files = len(midi_paths)\n",
+    "\n",
     "num_files_valid = round(total_num_files * 0.15)\n",
     "num_files_test = round(total_num_files * 0.15)\n",
     "shuffle(midi_paths)\n",
@@ -248,6 +411,7 @@
     "        save_dir=subset_chunks_dir,\n",
     "        max_seq_len=sequence_length,\n",
     "        num_overlap_bars=2,\n",
+    "        skip_drums=True\n",
     "    )\n",
     "\n",
     "    if subset_name == 'train':\n",
@@ -261,6 +425,16 @@
     "        )\n"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#Since the datasets are too large after splitting we only want 50% of the split data to train against\n",
+    "sample_subset_per = .25"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -269,47 +443,39 @@
    "source": [
     "# Create Dataset and Collator for training\n",
     "midi_paths_train = list(root_save.joinpath(Path(\"Maestro_train\")).glob(\"**/*.mid\")) + list(root_save.joinpath(Path(\"Maestro_train\")).glob(\"**/*.midi\"))\n",
-    "midi_paths_valid = list(root_save.joinpath(Path(\"Maestro_valid\")).glob(\"**/*.mid\")) + list(root_save.joinpath(Path(\"Maestro_valid\")).glob(\"**/*.midi\")) \n",
-    "midi_paths_test = list(root_save.joinpath(Path(\"Maestro_test\")).glob(\"**/*.mid\")) + list(root_save.joinpath(Path(\"Maestro_test\")).glob(\"**/*.midi\"))\n",
+    "sample_count = round(len(midi_paths_train)*sample_subset_per)\n",
+    "print(f\"sample count length: {sample_count} total count: {len(midi_paths_train)}\")\n",
     "\n",
-    "\n",
-    "\n",
-    "dataset_train = DatasetMIDI(midi_paths_train, **kwargs_dataset)\n",
-    "dataset_valid = DatasetMIDI(midi_paths_valid, **kwargs_dataset)\n",
-    "dataset_test = DatasetMIDI(midi_paths_test, **kwargs_dataset)\n",
-    "print (len(midi_paths_train), len(midi_paths_valid), len(midi_paths_test))"
+    "midi_paths_train_sample = midi_paths_train[0:sample_count]\n",
+    "print(len(midi_paths_train_sample))\n",
+    "dataset_train = DatasetMIDI(midi_paths_train_sample, **kwargs_dataset)\n",
+    "torch.save(dataset_train, Path(dataset_dir / \"dataset_train.pt\"))"
    ]
   },
   {
-   "cell_type": "markdown",
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
-    "# Preview files data load and split"
+    "midi_paths_valid = list(root_save.joinpath(Path(\"Maestro_valid\")).glob(\"**/*.mid\")) + list(root_save.joinpath(Path(\"Maestro_valid\")).glob(\"**/*.midi\")) \n",
+    "midi_paths_valid = midi_paths_valid[:(len(midi_paths_valid)*sample_subset_per)]\n",
+    "print(len(midi_paths_valid))\n",
+    "dataset_valid = DatasetMIDI(midi_paths_valid, **kwargs_dataset)\n",
+    "torch.save(dataset_valid, Path(dataset_dir / \"dataset_valid.pt\"))"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "tags": [
-     "Generate Preview Files"
-    ]
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "#testing_files = \n",
-    "preview_files_path = []\n",
-    "for testing_file in testing_files:\n",
-    "    preview_files_path.append(Path(testing_file))\n",
-    "\n",
-    "preview_dir = Path(root_save / \"preview\")\n",
-    "split_files_for_training(\n",
-    "        files_paths=preview_files_path,\n",
-    "        tokenizer=tokenizer,\n",
-    "        save_dir=preview_dir,\n",
-    "        max_seq_len=sequence_length,\n",
-    "        num_overlap_bars=2,\n",
-    "    )\n"
+    "midi_paths_test = list(root_save.joinpath(Path(\"Maestro_test\")).glob(\"**/*.mid\")) + list(root_save.joinpath(Path(\"Maestro_test\")).glob(\"**/*.midi\"))\n",
+    "midi_paths_test = midi_paths_test[:(len(midi_paths_test)*sample_subset_per)]\n",
+    "print(len(midi_paths_test))\n",
+    "dataset_test = DatasetMIDI(midi_paths_test, **kwargs_dataset)\n",
+    "torch.save(dataset_test, Path(dataset_dir / \"dataset_test.pt\"))\n"
    ]
   },
   {
@@ -318,18 +484,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "valid_midi_path = root_save / \"Maestro_valid\"\n",
-    "midi_split_preview = list(valid_midi_path.resolve().glob(\"**/*.mid\")) + list(valid_midi_path.resolve().glob(\"**/*.midi\"))\n",
-    "\n",
-    "print(len(midi_split_preview))\n",
-    "file_name_lookup = []\n",
-    "def func_to_get_labels(p1, p2, p3):\n",
-    "    if p3.name not in file_name_lookup:\n",
-    "        file_name_lookup.append(p3.name)\n",
-    "    return file_name_lookup.index(p3.name)\n",
-    "    \n",
-    "kwargs_dataset = {\"max_seq_len\": sequence_length, \"tokenizer\": tokenizer, \"bos_token_id\": tokenizer[\"BOS_None\"], \"eos_token_id\": tokenizer[\"EOS_None\"], \"func_to_get_labels\" : func_to_get_labels}\n",
-    "dataset_preview = DatasetMIDI(midi_split_preview, **kwargs_dataset)"
+    "print (len(midi_paths_train), len(midi_paths_valid), len(midi_paths_test))\n"
    ]
   },
   {
@@ -341,12 +496,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [],
    "source": [
-    "dataset_dir = root_save / \"data\"\n",
-    "dataset_dir.mkdir(parents=True, exist_ok=True)"
+    "\n",
+    "dataset_train = torch.load(Path(dataset_dir / \"dataset_train.pt\"), weights_only=False)\n",
+    "dataset_valid = torch.load(Path(dataset_dir / \"dataset_valid.pt\"), weights_only=False)\n",
+    "dataset_test = torch.load(Path(dataset_dir / \"dataset_test.pt\"), weights_only=False)\n",
+    "\n"
    ]
   },
   {
@@ -355,23 +513,27 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import torch\n",
-    "torch.save(dataset_train, Path(dataset_dir / \"dataset_train.pt\"))\n",
-    "torch.save(dataset_valid, Path(dataset_dir / \"dataset_valid.pt\"))\n",
-    "torch.save(dataset_test, Path(dataset_dir / \"dataset_test.pt\"))\n"
+    "import pickle\n",
+    "\n",
+    "test_file = open(Path(dataset_dir / \"dataset_test.pickle\"), 'ab')\n",
+    "pickle.dump(dataset_test, test_file)\n",
+    "test_file.close()\n",
+    "\n",
+    "print(dataset_test[0])\n",
+    "\n",
+    "test_file = open(Path(dataset_dir / \"dataset_test.pickle\"), 'rb')\n",
+    "test_pickle = pickle.load(test_file)\n",
+    "print(test_pickle)\n",
+    "print(test_pickle[0])\n",
+    "\n",
+    "\n"
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
+   "cell_type": "markdown",
    "metadata": {},
-   "outputs": [],
    "source": [
-    "import torch\n",
-    "dataset_train = torch.load(Path(dataset_dir / \"dataset_train.pt\"))\n",
-    "dataset_valid = torch.load(Path(dataset_dir / \"dataset_valid.pt\"))\n",
-    "dataset_test = torch.load(Path(dataset_dir / \"dataset_test.pt\"))\n",
-    "\n"
+    "# Preview files data load and split"
    ]
   },
   {
@@ -380,7 +542,33 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "print(dataset_train[0])\n"
+    "\n",
+    "#testing_files = \n",
+    "preview_files_path = []\n",
+    "for testing_file in testing_files:\n",
+    "    preview_files_path.append(Path(testing_file))\n",
+    "\n",
+    "preview_dir = Path(root_save / \"preview\")\n",
+    "split_files_for_training(\n",
+    "        files_paths=preview_files_path,\n",
+    "        tokenizer=tokenizer,\n",
+    "        save_dir=preview_dir,\n",
+    "        max_seq_len=sequence_length,\n",
+    "        num_overlap_bars=2,\n",
+    "    )\n",
+    "\n",
+    "valid_midi_path = root_save / \"Maestro_valid\"\n",
+    "midi_split_preview = list(valid_midi_path.resolve().glob(\"**/*.mid\")) + list(valid_midi_path.resolve().glob(\"**/*.midi\"))\n",
+    "\n",
+    "print(len(midi_split_preview))\n",
+    "file_name_lookup = []\n",
+    "def func_to_get_labels(p1, p2, p3):\n",
+    "    if p3.name not in file_name_lookup:\n",
+    "        file_name_lookup.append(p3.name)\n",
+    "    return file_name_lookup.index(p3.name)\n",
+    "    \n",
+    "kwargs_dataset = {\"max_seq_len\": sequence_length, \"tokenizer\": tokenizer, \"bos_token_id\": tokenizer[\"BOS_None\"], \"eos_token_id\": tokenizer[\"EOS_None\"], \"func_to_get_labels\" : func_to_get_labels}\n",
+    "dataset_preview = DatasetMIDI(midi_split_preview, **kwargs_dataset)"
    ]
   },
   {
@@ -398,9 +586,22 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 19,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Generate config GenerationConfig {\n",
+      "  \"bos_token_id\": 1,\n",
+      "  \"eos_token_id\": 2,\n",
+      "  \"pad_token_id\": 0\n",
+      "}\n",
+      "\n"
+     ]
+    }
+   ],
    "source": [
     "# Creates model\n",
     "model_config = MistralConfig(\n",
@@ -411,12 +612,80 @@
     "    num_attention_heads=8, # default  32\n",
     "    num_key_value_heads=4, # default 8\n",
     "    sliding_window=256, # default  4096\n",
-    "    max_position_embeddings=sequence_length + 256, # 8192 this was before  # default 4096*32\n",
+    "    max_position_embeddings=8192, #has no effect on the parms count or training just limits the input length  # default 4096*32\n",
     "    pad_token_id=tokenizer['PAD_None'],\n",
     "    bos_token_id=tokenizer['BOS_None'],\n",
     "    eos_token_id=tokenizer['EOS_None'],\n",
     ")\n",
-    "model = AutoModelForCausalLM.from_config(model_config)"
+    "\n",
+    "model = AutoModelForCausalLM.from_config(model_config)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "loading configuration file /workspace/traindata/train/checkpoint-22000/config.json\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Model config MistralConfig {\n",
+      "  \"architectures\": [\n",
+      "    \"MistralForCausalLM\"\n",
+      "  ],\n",
+      "  \"attention_dropout\": 0.0,\n",
+      "  \"bos_token_id\": 1,\n",
+      "  \"dtype\": \"float32\",\n",
+      "  \"eos_token_id\": 2,\n",
+      "  \"head_dim\": null,\n",
+      "  \"hidden_act\": \"silu\",\n",
+      "  \"hidden_size\": 512,\n",
+      "  \"initializer_range\": 0.02,\n",
+      "  \"intermediate_size\": 2048,\n",
+      "  \"max_position_embeddings\": 8192,\n",
+      "  \"model_type\": \"mistral\",\n",
+      "  \"num_attention_heads\": 8,\n",
+      "  \"num_hidden_layers\": 8,\n",
+      "  \"num_key_value_heads\": 4,\n",
+      "  \"pad_token_id\": 0,\n",
+      "  \"rms_norm_eps\": 1e-06,\n",
+      "  \"rope_theta\": 10000.0,\n",
+      "  \"sliding_window\": 256,\n",
+      "  \"tie_word_embeddings\": false,\n",
+      "  \"transformers_version\": \"4.56.2\",\n",
+      "  \"use_cache\": true,\n",
+      "  \"vocab_size\": 24000\n",
+      "}\n",
+      "\n",
+      "loading weights file /workspace/traindata/train/checkpoint-22000/model.safetensors\n",
+      "Generate config GenerationConfig {\n",
+      "  \"bos_token_id\": 1,\n",
+      "  \"eos_token_id\": 2,\n",
+      "  \"pad_token_id\": 0\n",
+      "}\n",
+      "\n",
+      "All model checkpoint weights were used when initializing MistralForCausalLM.\n",
+      "\n",
+      "All the weights of MistralForCausalLM were initialized from the model checkpoint at /workspace/traindata/train/checkpoint-22000/model.safetensors.\n",
+      "If your task is similar to the task the model of the checkpoint was trained on, you can already use MistralForCausalLM for predictions without further training.\n",
+      "Generation config file not found, using a generation config created from the model config.\n"
+     ]
+    }
+   ],
+   "source": [
+    "# This is only for training existing models not new ones\n",
+    "model_dir = Path(\"/workspace/traindata/train/checkpoint-22000\")\n",
+    "\n",
+    "config = AutoConfig.from_pretrained(str(model_dir / \"config.json\"))\n",
+    "model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=str(model_dir / \"model.safetensors\"), from_tf=False, config=config)"
    ]
   },
   {
@@ -429,9 +698,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 12,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run\n"
+     ]
+    }
+   ],
    "source": [
     "model_dir = root_save / 'run'\n",
     "model_dir_str = str(model_dir)\n",
@@ -442,7 +719,25 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "PyTorch: setting up devices\n",
+      "average_tokens_across_devices is True but world size is 1. Setting it to False automatically.\n",
+      "max_steps is given, it will override any value given in num_train_epochs\n",
+      "Using auto half precision backend\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "True\n"
+     ]
+    }
+   ],
    "source": [
     "metrics = {metric: load_metric(metric) for metric in [\"accuracy\"]}\n",
     "\n",
@@ -483,22 +778,23 @@
     "USE_MPS = not USE_CUDA and mps_available()\n",
     "training_config = TrainingArguments(\n",
     "    model_dir_str, False, True, True, False, \"steps\",\n",
-    "    per_device_train_batch_size=30, #76% @ 24 batch size #76% @ 32 batch size try 64 batch size next time \n",
-    "    per_device_eval_batch_size=30, #was 24 now 32\n",
+    "    per_device_train_batch_size=32, #76% @ 24 batch size #76% @ 32 batch size try 64 batch size next time \n",
+    "    per_device_eval_batch_size=32, #was 24 now 32\n",
     "    gradient_accumulation_steps=3, #change this to 4\n",
     "    eval_accumulation_steps=None,\n",
-    "    eval_steps=1000,\n",
+    "    eval_steps=3000,\n",
+    "    eval_delay=6000,\n",
     "    learning_rate=1e-4,\n",
     "    weight_decay=0.01,\n",
-    "    max_grad_norm=3.0,\n",
-    "    max_steps=40000,\n",
-    "    lr_scheduler_type=\"cosine_with_restarts\",\n",
-    "    warmup_ratio=0.3,\n",
+    "    max_grad_norm=1.0,\n",
+    "    max_steps=30000,\n",
+    "    lr_scheduler_type=\"cosine\",\n",
+    "    warmup_ratio=0.08,\n",
     "    log_level=\"debug\",\n",
     "    logging_strategy=\"steps\",\n",
-    "    logging_steps=20,\n",
+    "    logging_steps=100,\n",
     "    save_strategy=\"steps\",\n",
-    "    save_steps=1000,\n",
+    "    save_steps=3000,\n",
     "    save_total_limit=5,\n",
     "    no_cuda=not USE_CUDA,\n",
     "    seed=444,\n",
@@ -509,11 +805,11 @@
     "    load_best_model_at_end=True,\n",
     "    label_smoothing_factor=0.,\n",
     "    optim=\"adamw_torch\",\n",
-    "    report_to=[\"tensorboard\"],\n",
-    "    gradient_checkpointing=True,\n",
+    "    report_to=[\"tensorboard\", \"wandb\"],\n",
+    "    gradient_checkpointing=False,\n",
     "    dataloader_num_workers=8, #added to fix trashing isssue with the gpu not having enough data to process\n",
     "    dataloader_pin_memory=True, #we want the dataset in memory\n",
-    "    torch_compile=True #added to speed up \n",
+    "    torch_compile=False #added to speed up \n",
     "    \n",
     ")\n",
     "\n",
@@ -538,7 +834,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "print(model)"
+    "torch.cuda.empty_cache()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(model)\n",
+    "os.environ['CUDA_LAUNCH_BLOCKING']=\"1\"\n",
+    "os.environ['TORCH_USE_CUDA_DSA'] = \"1\""
    ]
   },
   {
@@ -546,6 +853,159 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
+   "source": [
+    "%env PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "seed=1\n",
+    "# # print(max(dataset_train[\"input_ids\"].max().item(), 0))\n",
+    "\n",
+    "torch.manual_seed(seed)\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Currently training with a batch size of: 32\n",
+      "***** Running training *****\n",
+      "  Num examples = 5,570,752\n",
+      "  Num Epochs = 1\n",
+      "  Instantaneous batch size per device = 32\n",
+      "  Total train batch size (w. parallel, distributed & accumulation) = 96\n",
+      "  Gradient Accumulation steps = 3\n",
+      "  Total optimization steps = 30,000\n",
+      "  Number of trainable parameters = 56,041,984\n",
+      "Automatic Weights & Biases logging enabled, to disable set os.environ[\"WANDB_DISABLED\"] = \"true\"\n",
+      "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "\n",
+       "    <div>\n",
+       "      \n",
+       "      <progress value='15166' max='30000' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
+       "      [15166/30000 3:11:26 < 3:07:16, 1.32 it/s, Epoch 0.26/1]\n",
+       "    </div>\n",
+       "    <table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       " <tr style=\"text-align: left;\">\n",
+       "      <th>Step</th>\n",
+       "      <th>Training Loss</th>\n",
+       "      <th>Validation Loss</th>\n",
+       "      <th>Accuracy</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <td>6000</td>\n",
+       "      <td>1.602200</td>\n",
+       "      <td>1.751858</td>\n",
+       "      <td>0.010508</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>9000</td>\n",
+       "      <td>1.584300</td>\n",
+       "      <td>1.732275</td>\n",
+       "      <td>0.010426</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>12000</td>\n",
+       "      <td>1.547300</td>\n",
+       "      <td>1.712772</td>\n",
+       "      <td>0.010505</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>15000</td>\n",
+       "      <td>1.540700</td>\n",
+       "      <td>1.694235</td>\n",
+       "      <td>0.010407</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table><p>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Saving model checkpoint to /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-3000\n",
+      "Configuration saved in /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-3000/config.json\n",
+      "Configuration saved in /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-3000/generation_config.json\n",
+      "Model weights saved in /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-3000/model.safetensors\n",
+      "\n",
+      "***** Running Evaluation *****\n",
+      "  Num examples = 849907\n",
+      "  Batch size = 32\n",
+      "Saving model checkpoint to /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-6000\n",
+      "Configuration saved in /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-6000/config.json\n",
+      "Configuration saved in /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-6000/generation_config.json\n",
+      "Model weights saved in /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-6000/model.safetensors\n",
+      "\n",
+      "***** Running Evaluation *****\n",
+      "  Num examples = 849907\n",
+      "  Batch size = 32\n",
+      "Saving model checkpoint to /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-9000\n",
+      "Configuration saved in /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-9000/config.json\n",
+      "Configuration saved in /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-9000/generation_config.json\n",
+      "Model weights saved in /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-9000/model.safetensors\n",
+      "\n",
+      "***** Running Evaluation *****\n",
+      "  Num examples = 849907\n",
+      "  Batch size = 32\n",
+      "Saving model checkpoint to /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-12000\n",
+      "Configuration saved in /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-12000/config.json\n",
+      "Configuration saved in /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-12000/generation_config.json\n",
+      "Model weights saved in /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-12000/model.safetensors\n",
+      "\n",
+      "***** Running Evaluation *****\n",
+      "  Num examples = 849907\n",
+      "  Batch size = 32\n",
+      "Saving model checkpoint to /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-15000\n",
+      "Configuration saved in /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-15000/config.json\n",
+      "Configuration saved in /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-15000/generation_config.json\n",
+      "Model weights saved in /workspace/traindata/data/HuggingFace_Mistral_Transformer_Single_Instrument/run/checkpoint-15000/model.safetensors\n"
+     ]
+    },
+    {
+     "ename": "KeyboardInterrupt",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
+      "\u001b[31mKeyboardInterrupt\u001b[39m                         Traceback (most recent call last)",
+      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[24]\u001b[39m\u001b[32m, line 2\u001b[39m\n\u001b[32m      1\u001b[39m \u001b[38;5;66;03m# Training\u001b[39;00m\n\u001b[32m----> \u001b[39m\u001b[32m2\u001b[39m train_result = \u001b[43mtrainer\u001b[49m\u001b[43m.\u001b[49m\u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m      3\u001b[39m trainer.save_model()  \u001b[38;5;66;03m# Saves the tokenizer too\u001b[39;00m\n\u001b[32m      4\u001b[39m trainer.log_metrics(\u001b[33m\"\u001b[39m\u001b[33mtrain\u001b[39m\u001b[33m\"\u001b[39m, train_result.metrics)\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/transformers/trainer.py:2328\u001b[39m, in \u001b[36mTrainer.train\u001b[39m\u001b[34m(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)\u001b[39m\n\u001b[32m   2326\u001b[39m         hf_hub_utils.enable_progress_bars()\n\u001b[32m   2327\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m-> \u001b[39m\u001b[32m2328\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43minner_training_loop\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m   2329\u001b[39m \u001b[43m        \u001b[49m\u001b[43margs\u001b[49m\u001b[43m=\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   2330\u001b[39m \u001b[43m        \u001b[49m\u001b[43mresume_from_checkpoint\u001b[49m\u001b[43m=\u001b[49m\u001b[43mresume_from_checkpoint\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   2331\u001b[39m \u001b[43m        \u001b[49m\u001b[43mtrial\u001b[49m\u001b[43m=\u001b[49m\u001b[43mtrial\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   2332\u001b[39m \u001b[43m        \u001b[49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[43m=\u001b[49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   2333\u001b[39m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/transformers/trainer.py:2672\u001b[39m, in \u001b[36mTrainer._inner_training_loop\u001b[39m\u001b[34m(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)\u001b[39m\n\u001b[32m   2665\u001b[39m context = (\n\u001b[32m   2666\u001b[39m     functools.partial(\u001b[38;5;28mself\u001b[39m.accelerator.no_sync, model=model)\n\u001b[32m   2667\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m i != \u001b[38;5;28mlen\u001b[39m(batch_samples) - \u001b[32m1\u001b[39m\n\u001b[32m   2668\u001b[39m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m.accelerator.distributed_type != DistributedType.DEEPSPEED\n\u001b[32m   2669\u001b[39m     \u001b[38;5;28;01melse\u001b[39;00m contextlib.nullcontext\n\u001b[32m   2670\u001b[39m )\n\u001b[32m   2671\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m context():\n\u001b[32m-> \u001b[39m\u001b[32m2672\u001b[39m     tr_loss_step = \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mtraining_step\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnum_items_in_batch\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m   2674\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m (\n\u001b[32m   2675\u001b[39m     args.logging_nan_inf_filter\n\u001b[32m   2676\u001b[39m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m is_torch_xla_available()\n\u001b[32m   2677\u001b[39m     \u001b[38;5;129;01mand\u001b[39;00m (torch.isnan(tr_loss_step) \u001b[38;5;129;01mor\u001b[39;00m torch.isinf(tr_loss_step))\n\u001b[32m   2678\u001b[39m ):\n\u001b[32m   2679\u001b[39m     \u001b[38;5;66;03m# if loss is nan or inf simply add the average of previous logged losses\u001b[39;00m\n\u001b[32m   2680\u001b[39m     tr_loss = tr_loss + tr_loss / (\u001b[32m1\u001b[39m + \u001b[38;5;28mself\u001b[39m.state.global_step - \u001b[38;5;28mself\u001b[39m._globalstep_last_logged)\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/transformers/trainer.py:4060\u001b[39m, in \u001b[36mTrainer.training_step\u001b[39m\u001b[34m(***failed resolving arguments***)\u001b[39m\n\u001b[32m   4057\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m.accelerator.distributed_type == DistributedType.DEEPSPEED:\n\u001b[32m   4058\u001b[39m         kwargs[\u001b[33m\"\u001b[39m\u001b[33mscale_wrt_gas\u001b[39m\u001b[33m\"\u001b[39m] = \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[32m-> \u001b[39m\u001b[32m4060\u001b[39m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43maccelerator\u001b[49m\u001b[43m.\u001b[49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[43mloss\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m   4062\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m loss.detach()\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/accelerate/accelerator.py:2734\u001b[39m, in \u001b[36mAccelerator.backward\u001b[39m\u001b[34m(self, loss, **kwargs)\u001b[39m\n\u001b[32m   2732\u001b[39m     \u001b[38;5;28mself\u001b[39m.lomo_backward(loss, learning_rate)\n\u001b[32m   2733\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m-> \u001b[39m\u001b[32m2734\u001b[39m     \u001b[43mloss\u001b[49m\u001b[43m.\u001b[49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/torch/_tensor.py:648\u001b[39m, in \u001b[36mTensor.backward\u001b[39m\u001b[34m(self, gradient, retain_graph, create_graph, inputs)\u001b[39m\n\u001b[32m    638\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[32m    639\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[32m    640\u001b[39m         Tensor.backward,\n\u001b[32m    641\u001b[39m         (\u001b[38;5;28mself\u001b[39m,),\n\u001b[32m   (...)\u001b[39m\u001b[32m    646\u001b[39m         inputs=inputs,\n\u001b[32m    647\u001b[39m     )\n\u001b[32m--> \u001b[39m\u001b[32m648\u001b[39m \u001b[43mtorch\u001b[49m\u001b[43m.\u001b[49m\u001b[43mautograd\u001b[49m\u001b[43m.\u001b[49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m    649\u001b[39m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgradient\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m=\u001b[49m\u001b[43minputs\u001b[49m\n\u001b[32m    650\u001b[39m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/torch/autograd/__init__.py:353\u001b[39m, in \u001b[36mbackward\u001b[39m\u001b[34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[39m\n\u001b[32m    348\u001b[39m     retain_graph = create_graph\n\u001b[32m    350\u001b[39m \u001b[38;5;66;03m# The reason we repeat the same comment below is that\u001b[39;00m\n\u001b[32m    351\u001b[39m \u001b[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[32m    352\u001b[39m \u001b[38;5;66;03m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m353\u001b[39m \u001b[43m_engine_run_backward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m    354\u001b[39m \u001b[43m    \u001b[49m\u001b[43mtensors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    355\u001b[39m \u001b[43m    \u001b[49m\u001b[43mgrad_tensors_\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    356\u001b[39m \u001b[43m    \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    357\u001b[39m \u001b[43m    \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    358\u001b[39m \u001b[43m    \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    359\u001b[39m \u001b[43m    \u001b[49m\u001b[43mallow_unreachable\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[32m    360\u001b[39m \u001b[43m    \u001b[49m\u001b[43maccumulate_grad\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[32m    361\u001b[39m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/torch/autograd/graph.py:824\u001b[39m, in \u001b[36m_engine_run_backward\u001b[39m\u001b[34m(t_outputs, *args, **kwargs)\u001b[39m\n\u001b[32m    822\u001b[39m     unregister_hooks = _register_logging_hooks_on_whole_graph(t_outputs)\n\u001b[32m    823\u001b[39m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[32m--> \u001b[39m\u001b[32m824\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mVariable\u001b[49m\u001b[43m.\u001b[49m\u001b[43m_execution_engine\u001b[49m\u001b[43m.\u001b[49m\u001b[43mrun_backward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# Calls into the C++ engine to run the backward pass\u001b[39;49;00m\n\u001b[32m    825\u001b[39m \u001b[43m        \u001b[49m\u001b[43mt_outputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\n\u001b[32m    826\u001b[39m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# Calls into the C++ engine to run the backward pass\u001b[39;00m\n\u001b[32m    827\u001b[39m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[32m    828\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m attach_logging_hooks:\n",
+      "\u001b[31mKeyboardInterrupt\u001b[39m: "
+     ]
+    }
+   ],
    "source": [
     "# Training\n",
     "train_result = trainer.train()\n",
@@ -555,6 +1015,62 @@
     "trainer.save_state()"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<br>    <style><br>        .wandb-row {<br>            display: flex;<br>            flex-direction: row;<br>            flex-wrap: wrap;<br>            justify-content: flex-start;<br>            width: 100%;<br>        }<br>        .wandb-col {<br>            display: flex;<br>            flex-direction: column;<br>            flex-basis: 100%;<br>            flex: 1;<br>            padding: 10px;<br>        }<br>    </style><br><div class=\"wandb-row\"><div class=\"wandb-col\"><h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>eval/accuracy</td><td>█▂█▁</td></tr><tr><td>eval/loss</td><td>█▆▃▁</td></tr><tr><td>eval/runtime</td><td>█▁▃▆</td></tr><tr><td>eval/samples_per_second</td><td>▁█▆▃</td></tr><tr><td>eval/steps_per_second</td><td>▁█▆▃</td></tr><tr><td>train/epoch</td><td>▁▁▁▁▁▂▂▂▂▁▂▂▂▂▃▃▃▃▃▄▄▄▄▅▅▅▅▅▅▅▆▆▆▆▆▇▇███</td></tr><tr><td>train/global_step</td><td>▁▁▁▁▂▂▂▂▂▂▁▁▁▁▁▂▂▂▃▃▃▄▄▄▄▅▅▅▅▅▆▆▆▆▇▇████</td></tr><tr><td>train/grad_norm</td><td>▅▃▄▅▆█▇▆▁▁▂▂▂▂▁▂▁▂▂▁▁▁▁▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>train/learning_rate</td><td>▃▅▅▆▆█▁▂▄▄▆▇█████████▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▅</td></tr><tr><td>train/loss</td><td>█▇▆▆▆▅▄▄▃▃▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr></table><br/></div><div class=\"wandb-col\"><h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>eval/accuracy</td><td>0.01041</td></tr><tr><td>eval/loss</td><td>1.69424</td></tr><tr><td>eval/runtime</td><td>1748.5708</td></tr><tr><td>eval/samples_per_second</td><td>486.058</td></tr><tr><td>eval/steps_per_second</td><td>15.19</td></tr><tr><td>train/epoch</td><td>0.26022</td></tr><tr><td>train/global_step</td><td>15100</td></tr><tr><td>train/grad_norm</td><td>0.56171</td></tr><tr><td>train/learning_rate</td><td>6e-05</td></tr><tr><td>train/loss</td><td>1.5401</td></tr></table><br/></div></div>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View run <strong style=\"color:#cdcd00\">fast-yogurt-6</strong> at: <a href='https://wandb.ai/adric-landman-hobby/midi_music_maker/runs/g1fn393k' target=\"_blank\">https://wandb.ai/adric-landman-hobby/midi_music_maker/runs/g1fn393k</a><br> View project at: <a href='https://wandb.ai/adric-landman-hobby/midi_music_maker' target=\"_blank\">https://wandb.ai/adric-landman-hobby/midi_music_maker</a><br>Synced 5 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Find logs at: <code></code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "wandb.finish()"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -567,9 +1083,32 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 26,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "ename": "HfHubHTTPError",
+     "evalue": "401 Client Error: Unauthorized for url: https://huggingface.co/api/repos/create (Request ID: Root=1-68d628b1-575691d056937c56340182c7;4f0ea033-8e43-4260-938f-d74b744942be)\n\nInvalid credentials in Authorization header",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
+      "\u001b[31mHTTPError\u001b[39m                                 Traceback (most recent call last)",
+      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_http.py:407\u001b[39m, in \u001b[36mhf_raise_for_status\u001b[39m\u001b[34m(response, endpoint_name)\u001b[39m\n\u001b[32m    406\u001b[39m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[32m--> \u001b[39m\u001b[32m407\u001b[39m     \u001b[43mresponse\u001b[49m\u001b[43m.\u001b[49m\u001b[43mraise_for_status\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    408\u001b[39m \u001b[38;5;28;01mexcept\u001b[39;00m HTTPError \u001b[38;5;28;01mas\u001b[39;00m e:\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/requests/models.py:1024\u001b[39m, in \u001b[36mResponse.raise_for_status\u001b[39m\u001b[34m(self)\u001b[39m\n\u001b[32m   1023\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m http_error_msg:\n\u001b[32m-> \u001b[39m\u001b[32m1024\u001b[39m     \u001b[38;5;28;01mraise\u001b[39;00m HTTPError(http_error_msg, response=\u001b[38;5;28mself\u001b[39m)\n",
+      "\u001b[31mHTTPError\u001b[39m: 401 Client Error: Unauthorized for url: https://huggingface.co/api/repos/create",
+      "\nThe above exception was the direct cause of the following exception:\n",
+      "\u001b[31mHfHubHTTPError\u001b[39m                            Traceback (most recent call last)",
+      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[26]\u001b[39m\u001b[32m, line 3\u001b[39m\n\u001b[32m      1\u001b[39m model.hub_model_id = \u001b[33m\"\u001b[39m\u001b[33madricl/midi_single_instrument_mistral_transformer\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m3\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m.\u001b[49m\u001b[43mpush_to_hub\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcommit_message\u001b[49m\u001b[43m=\u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43mTraining Basic Model for Mistral MidiTok Transformer Single Instrument Small\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mrepo_id\u001b[49m\u001b[43m=\u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43madricl/midi_single_instrument_mistral_transformer\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[32m      4\u001b[39m \u001b[43m                 \u001b[49m\u001b[43mtoken\u001b[49m\u001b[43m=\u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/transformers/modeling_utils.py:4346\u001b[39m, in \u001b[36mPreTrainedModel.push_to_hub\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m   4344\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m tags:\n\u001b[32m   4345\u001b[39m     kwargs[\u001b[33m\"\u001b[39m\u001b[33mtags\u001b[39m\u001b[33m\"\u001b[39m] = tags\n\u001b[32m-> \u001b[39m\u001b[32m4346\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m.\u001b[49m\u001b[43mpush_to_hub\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/transformers/utils/hub.py:955\u001b[39m, in \u001b[36mPushToHubMixin.push_to_hub\u001b[39m\u001b[34m(self, repo_id, use_temp_dir, commit_message, private, token, max_shard_size, create_pr, safe_serialization, revision, commit_description, tags, **deprecated_kwargs)\u001b[39m\n\u001b[32m    952\u001b[39m repo_url = deprecated_kwargs.pop(\u001b[33m\"\u001b[39m\u001b[33mrepo_url\u001b[39m\u001b[33m\"\u001b[39m, \u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[32m    953\u001b[39m organization = deprecated_kwargs.pop(\u001b[33m\"\u001b[39m\u001b[33morganization\u001b[39m\u001b[33m\"\u001b[39m, \u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[32m--> \u001b[39m\u001b[32m955\u001b[39m repo_id = \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_create_repo\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m    956\u001b[39m \u001b[43m    \u001b[49m\u001b[43mrepo_id\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mprivate\u001b[49m\u001b[43m=\u001b[49m\u001b[43mprivate\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtoken\u001b[49m\u001b[43m=\u001b[49m\u001b[43mtoken\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mrepo_url\u001b[49m\u001b[43m=\u001b[49m\u001b[43mrepo_url\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43morganization\u001b[49m\u001b[43m=\u001b[49m\u001b[43morganization\u001b[49m\n\u001b[32m    957\u001b[39m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    959\u001b[39m \u001b[38;5;66;03m# Create a new empty model card and eventually tag it\u001b[39;00m\n\u001b[32m    960\u001b[39m model_card = create_and_tag_model_card(\n\u001b[32m    961\u001b[39m     repo_id, tags, token=token, ignore_metadata_errors=ignore_metadata_errors\n\u001b[32m    962\u001b[39m )\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/transformers/utils/hub.py:759\u001b[39m, in \u001b[36mPushToHubMixin._create_repo\u001b[39m\u001b[34m(self, repo_id, private, token, repo_url, organization)\u001b[39m\n\u001b[32m    756\u001b[39m             repo_id = repo_id.split(\u001b[33m\"\u001b[39m\u001b[33m/\u001b[39m\u001b[33m\"\u001b[39m)[-\u001b[32m1\u001b[39m]\n\u001b[32m    757\u001b[39m         repo_id = \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00morganization\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m/\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mrepo_id\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m\"\u001b[39m\n\u001b[32m--> \u001b[39m\u001b[32m759\u001b[39m url = \u001b[43mcreate_repo\u001b[49m\u001b[43m(\u001b[49m\u001b[43mrepo_id\u001b[49m\u001b[43m=\u001b[49m\u001b[43mrepo_id\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtoken\u001b[49m\u001b[43m=\u001b[49m\u001b[43mtoken\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mprivate\u001b[49m\u001b[43m=\u001b[49m\u001b[43mprivate\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexist_ok\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[32m    760\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m url.repo_id\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_validators.py:114\u001b[39m, in \u001b[36mvalidate_hf_hub_args.<locals>._inner_fn\u001b[39m\u001b[34m(*args, **kwargs)\u001b[39m\n\u001b[32m    111\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m check_use_auth_token:\n\u001b[32m    112\u001b[39m     kwargs = smoothly_deprecate_use_auth_token(fn_name=fn.\u001b[34m__name__\u001b[39m, has_token=has_token, kwargs=kwargs)\n\u001b[32m--> \u001b[39m\u001b[32m114\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/huggingface_hub/hf_api.py:3766\u001b[39m, in \u001b[36mHfApi.create_repo\u001b[39m\u001b[34m(self, repo_id, token, private, repo_type, exist_ok, resource_group_id, space_sdk, space_hardware, space_storage, space_sleep_time, space_secrets, space_variables)\u001b[39m\n\u001b[32m   3763\u001b[39m     \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[32m   3765\u001b[39m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[32m-> \u001b[39m\u001b[32m3766\u001b[39m     \u001b[43mhf_raise_for_status\u001b[49m\u001b[43m(\u001b[49m\u001b[43mr\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m   3767\u001b[39m \u001b[38;5;28;01mexcept\u001b[39;00m HTTPError \u001b[38;5;28;01mas\u001b[39;00m err:\n\u001b[32m   3768\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m exist_ok \u001b[38;5;129;01mand\u001b[39;00m err.response.status_code == \u001b[32m409\u001b[39m:\n\u001b[32m   3769\u001b[39m         \u001b[38;5;66;03m# Repo already exists and `exist_ok=True`\u001b[39;00m\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_http.py:480\u001b[39m, in \u001b[36mhf_raise_for_status\u001b[39m\u001b[34m(response, endpoint_name)\u001b[39m\n\u001b[32m    476\u001b[39m     \u001b[38;5;28;01mraise\u001b[39;00m _format(HfHubHTTPError, message, response) \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01me\u001b[39;00m\n\u001b[32m    478\u001b[39m \u001b[38;5;66;03m# Convert `HTTPError` into a `HfHubHTTPError` to display request information\u001b[39;00m\n\u001b[32m    479\u001b[39m \u001b[38;5;66;03m# as well (request id and/or server error message)\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m480\u001b[39m \u001b[38;5;28;01mraise\u001b[39;00m _format(HfHubHTTPError, \u001b[38;5;28mstr\u001b[39m(e), response) \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01me\u001b[39;00m\n",
+      "\u001b[31mHfHubHTTPError\u001b[39m: 401 Client Error: Unauthorized for url: https://huggingface.co/api/repos/create (Request ID: Root=1-68d628b1-575691d056937c56340182c7;4f0ea033-8e43-4260-938f-d74b744942be)\n\nInvalid credentials in Authorization header"
+     ]
+    }
+   ],
    "source": [
     "\n",
     "model.hub_model_id = \"adricl/midi_single_instrument_mistral_transformer\"\n",
@@ -699,7 +1238,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.5"
+   "version": "3.11.11"
   },
   "vscode": {
    "interpreter": {

train_tokenizer.py

@@ -0,0 +1,56 @@
+from miditok import REMI, TokenizerConfig
+from random import shuffle, sample
+from pathlib import Path
+
+# Our tokenizer's configuration
+BEAT_RES = {(0, 1): 12, (1, 2): 4, (2, 4): 2, (4, 8): 1}
+TOKENIZER_PARAMS = {
+    "pitch_range": (21, 108),
+    "beat_res": BEAT_RES,
+    "num_velocities": 32,
+    "special_tokens": ["PAD", "BOS", "EOS"],
+    "use_chords": True,
+    "use_rests": True,
+    "use_tempos": True,
+    "use_time_signatures": True,
+    "use_programs": False,  # We want single track 
+    "one_token_stream_for_programs": False, # We want single track
+    "programs": list(range(0, 128)), #-1 drums, skip drums
+    "num_tempos": 32,
+    "tempo_range": (40, 250),  # (min_tempo, max_tempo)
+}
+config = TokenizerConfig(**TOKENIZER_PARAMS)
+
+# Creates the tokenizer REMI PLUS
+tokenizer = REMI(config)
+
+root_data_dir = Path('/root')
+root_save = Path(root_data_dir / 'HuggingFace_Mistral_Transformer_Single_Instrument')
+
+tokenizer_name = "HuggingFace_Mistral_Transformer_Single_Instrument_v4_single_track.json"
+
+
+
+data_dirs = ["MIDIs"]
+midi_paths = []
+for data_dir in data_dirs:
+    path = Path(root_data_dir / data_dir)
+    midi_paths.extend(list(path.resolve().glob("**/*.mid")) + list(path.resolve().glob("**/*.midi")))
+
+print(f"Found {len(midi_paths)} MIDI files")
+
+midi_paths = midi_paths.shuffle()
+
+# We need a subset of files otherwise training tokenizer takes too long
+percentage_to_select = 0.20
+num_files_to_select = int(len(midi_paths) * percentage_to_select)
+
+subset_midi_paths = sample(midi_paths, num_files_to_select)
+print(f"Found {len(subset_midi_paths)} MIDI files")
+
+
+tokenizer.train(
+    vocab_size=24000,
+    files_paths=subset_midi_paths,
+)
+tokenizer.save(root_save / tokenizer_name)