03_LM_pretraining_data_preprocessing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('../../source')\n",
    "\n",
    "from utils import *\n",
    "\n",
    "import pickle\n",
    "import json\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm\n",
    "import glob\n",
    "import time\n",
    "import random\n",
    "\n",
    "import encoders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is the path to the labeled dataset\n",
    "data_path = Path(\"/home/abunn/ttmp\")\n",
    "repo_path = data_path/\"piano_bootleg_scores\"\n",
    "piano_bootleg_scores_path = repo_path/\"imslp_bootleg_dir-v1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "31834\n"
     ]
    }
   ],
   "source": [
    "# Grab all file locations\n",
    "piece_names = glob.glob(str(Path(piano_bootleg_scores_path)/\"**/*\"))\n",
    "\n",
    "# This gets only one version of each piece\n",
    "# fnames = [glob.glob(str(Path(piece_name)/\"*.pkl\"))[0] for piece_name in piece_names if len(glob.glob(str(Path(piece_name)/\"*.pkl\"))) != 0]\n",
    "\n",
    "# This gets every version of every piece\n",
    "fnames = glob.glob(str(piano_bootleg_scores_path/\"**/*.pkl\"), recursive=True)\n",
    "\n",
    "print(len(fnames))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "filler_file = '../../cfg_files/filler_imslp.txt'\n",
    "filler = {}\n",
    "with open(filler_file, 'r') as f:\n",
    "    for line in f:\n",
    "        parts = line.strip(\"\\n\").split('\\t')\n",
    "        assert len(parts) == 2\n",
    "\n",
    "        filler[parts[0]] = [int(i) for i in parts[1].split(\",\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 31834/31834 [06:31<00:00, 81.33it/s] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Seconds to complete: 391\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# CONVERTING THE DATA TO BINARY MATRICES - MIGHT TAKE A MINUTE\n",
    "\n",
    "t0 = time.time()\n",
    "\n",
    "\n",
    "# List of tuples\n",
    "# Tuples contain (binary_score, composer)\n",
    "pieces = []\n",
    "\n",
    "for fname in tqdm(fnames):\n",
    "    # Load the pages\n",
    "    pages = load_pkl(fname)\n",
    "\n",
    "    filler_key = fname.split(\"imslp_bootleg_dir-v1/\")[1].strip(\".pkl\")\n",
    "\n",
    "    filler_pages = filler[filler_key] if filler_key in filler.keys() else []\n",
    "\n",
    "    # Convert them into binary matrices\n",
    "    bscores = [ints_to_binary_matrix(page) for i, page in enumerate(pages) if i not in filler_pages]\n",
    "    bscores = [page for page in bscores if len(page.shape) == 2 and page.shape[1] == 62]\n",
    "\n",
    "    # If there were binary scores, then combine them into one and append to dataset.\n",
    "    if len(bscores) > 0:\n",
    "        piece = np.concatenate(bscores, axis=0)\n",
    "        pieces.append(piece)\n",
    "\n",
    "print(\"Seconds to complete:\", round(time.time() - t0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1738852\n"
     ]
    }
   ],
   "source": [
    "print(sum(len(piece[0]) for piece in pieces))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1807672\n"
     ]
    }
   ],
   "source": [
    "print(sum(len(piece[0]) for piece in pieces))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dataset(pieces, valid_split=.15, test_split=.15):\n",
    "    \"\"\"\n",
    "    Creates a train / valid / test split dataset of pieces.\n",
    "    pieces: The list of binary_matrices to sample from\n",
    "    valid_split: The proportion of data to use for valid\n",
    "    test_split: The proportion of data to use for valid\n",
    "    \n",
    "    returns:\n",
    "    x & y lists for train, valid, and test sets\n",
    "    \"\"\"\n",
    "    \n",
    "    # For repeatability\n",
    "    np.random.seed(42)\n",
    "\n",
    "    # shuffle pieces\n",
    "    piece_list = [piece for piece in pieces]\n",
    "    np.random.shuffle(piece_list)\n",
    "    \n",
    "    # Calculate starting places of each section - order is (test, valid, train)\n",
    "    train_start = round((valid_split+test_split)*len(piece_list))\n",
    "    valid_start = round(test_split*len(piece_list))\n",
    "    \n",
    "    # Go through and separate pieces into train, valid, test\n",
    "    train_pieces = piece_list[train_start:]\n",
    "    valid_pieces = piece_list[valid_start:train_start]\n",
    "    test_pieces = piece_list[:valid_start]\n",
    "    \n",
    "    return train_pieces, valid_pieces, test_pieces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "train, valid, test = create_dataset(pieces, valid_split=.2, test_split=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22437\n",
      "5609\n"
     ]
    }
   ],
   "source": [
    "print(len(train))\n",
    "print(len(valid))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sparse encodings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 22437/22437 [05:25<00:00, 68.89it/s] \n",
      "100%|██████████| 5609/5609 [01:22<00:00, 67.78it/s] \n",
      "100%|██████████| 22437/22437 [05:41<00:00, 65.69it/s] \n",
      "100%|██████████| 5609/5609 [01:28<00:00, 63.50it/s] \n",
      "100%|██████████| 22437/22437 [06:31<00:00, 57.30it/s]\n",
      "100%|██████████| 5609/5609 [01:37<00:00, 57.28it/s]\n"
     ]
    }
   ],
   "source": [
    "for enc_name, enc in encoders.sparse_encoders.items():\n",
    "    train_encoded = []\n",
    "    for piece in tqdm(train):\n",
    "        train_encoded.append(enc(piece))\n",
    "        \n",
    "    valid_encoded = []\n",
    "    for piece in tqdm(valid):\n",
    "        valid_encoded.append(enc(piece))\n",
    "    \n",
    "    # Data for LM pretraining\n",
    "    with open(f\"LM_pretraining_data/{enc_name}-train.txt\", \"w\") as f:\n",
    "        f.write(\"\\n\\n\".join(train_encoded))\n",
    "    with open(f\"LM_pretraining_data/{enc_name}-valid.txt\", \"w\") as f:\n",
    "        f.write(\"\\n\\n\".join(valid_encoded))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dense encodings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 22437/22437 [14:48<00:00, 25.24it/s]\n",
      "100%|██████████| 5609/5609 [03:44<00:00, 24.96it/s]\n",
      "100%|██████████| 22437/22437 [09:39<00:00, 38.71it/s]\n",
      "100%|██████████| 5609/5609 [02:26<00:00, 38.34it/s]\n",
      "100%|██████████| 22437/22437 [06:34<00:00, 56.85it/s]\n",
      "100%|██████████| 5609/5609 [01:40<00:00, 55.68it/s]\n",
      "100%|██████████| 22437/22437 [04:01<00:00, 93.07it/s] \n",
      "100%|██████████| 5609/5609 [01:01<00:00, 90.56it/s] \n"
     ]
    }
   ],
   "source": [
    "block_sizes = [\n",
    "    [1, 1],\n",
    "    [1, 2],\n",
    "    [1, 4],\n",
    "    [1, 8],\n",
    "]\n",
    "\n",
    "for block_size in block_sizes:\n",
    "    train_encoded = []\n",
    "    for piece in tqdm(train):\n",
    "        train_encoded.append(encoders.dense_encoder(piece, block_size=block_size))\n",
    "        \n",
    "    valid_encoded = []\n",
    "    for piece in tqdm(valid):\n",
    "        valid_encoded.append(encoders.dense_encoder(piece, block_size=block_size))\n",
    "    \n",
    "    # Data for LM pretraining\n",
    "    with open(f\"LM_pretraining_data/dense_{block_size[0]}_{block_size[1]}-train.txt\", \"w\") as f:\n",
    "        f.write(\"\\n\\n\".join(train_encoded))\n",
    "    with open(f\"LM_pretraining_data/dense_{block_size[0]}_{block_size[1]}-valid.txt\", \"w\") as f:\n",
    "        f.write(\"\\n\\n\".join(valid_encoded))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "EWLLMs",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}