Output Selection.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a9301121",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mimikit as mmk\n",
    "import h5mapper as h5m\n",
    "import torch.nn as nn\n",
    "import torch\n",
    "import os\n",
    "import json\n",
    "from random import randint\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from models.wavenets import WaveNetFFT, WaveNetQx\n",
    "from models.srnns import SampleRNN\n",
    "from models.s2s import Seq2SeqLSTM\n",
    "from ensemble import Ensemble\n",
    "from mains import train, generate\n",
    "\n",
    "from checkpoints import Checkpoint\n",
    "from datasets import Trainset\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from mimikit.extract.from_neighbors import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c60d87b0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Checkpoint(id='dc63b032229eafd03f52f56d14aefbb4274bea5ca5fcdd55c6bdfa5cc43d1e8e', epoch=50, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='dc63b032229eafd03f52f56d14aefbb4274bea5ca5fcdd55c6bdfa5cc43d1e8e', epoch=30, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='dc63b032229eafd03f52f56d14aefbb4274bea5ca5fcdd55c6bdfa5cc43d1e8e', epoch=10, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='dc63b032229eafd03f52f56d14aefbb4274bea5ca5fcdd55c6bdfa5cc43d1e8e', epoch=40, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='dc63b032229eafd03f52f56d14aefbb4274bea5ca5fcdd55c6bdfa5cc43d1e8e', epoch=20, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='d9746c79b8e54a1ef67374b3ab711f2f440416b5a9adbda2b6c5e1f16434c505', epoch=50, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='d9746c79b8e54a1ef67374b3ab711f2f440416b5a9adbda2b6c5e1f16434c505', epoch=30, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='d9746c79b8e54a1ef67374b3ab711f2f440416b5a9adbda2b6c5e1f16434c505', epoch=10, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='d9746c79b8e54a1ef67374b3ab711f2f440416b5a9adbda2b6c5e1f16434c505', epoch=40, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='d9746c79b8e54a1ef67374b3ab711f2f440416b5a9adbda2b6c5e1f16434c505', epoch=20, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='2265d9b3469c7494f68b7ffd54258653d7fd6a73470fac427b942ef848f0d334', epoch=50, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='2265d9b3469c7494f68b7ffd54258653d7fd6a73470fac427b942ef848f0d334', epoch=30, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='2265d9b3469c7494f68b7ffd54258653d7fd6a73470fac427b942ef848f0d334', epoch=10, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='2265d9b3469c7494f68b7ffd54258653d7fd6a73470fac427b942ef848f0d334', epoch=40, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='2265d9b3469c7494f68b7ffd54258653d7fd6a73470fac427b942ef848f0d334', epoch=20, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='867974df7116e02659e3f01c01527f3c366ff754e2787c4b14fd30e20240e15c', epoch=50, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='867974df7116e02659e3f01c01527f3c366ff754e2787c4b14fd30e20240e15c', epoch=30, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='867974df7116e02659e3f01c01527f3c366ff754e2787c4b14fd30e20240e15c', epoch=10, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='867974df7116e02659e3f01c01527f3c366ff754e2787c4b14fd30e20240e15c', epoch=40, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='867974df7116e02659e3f01c01527f3c366ff754e2787c4b14fd30e20240e15c', epoch=20, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='7e2133e202f951398d3c41ce75fb15d961c4d171ae0ae7dba10a0c22cf6286ec', epoch=50, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='7e2133e202f951398d3c41ce75fb15d961c4d171ae0ae7dba10a0c22cf6286ec', epoch=30, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='7e2133e202f951398d3c41ce75fb15d961c4d171ae0ae7dba10a0c22cf6286ec', epoch=10, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='7e2133e202f951398d3c41ce75fb15d961c4d171ae0ae7dba10a0c22cf6286ec', epoch=40, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='7e2133e202f951398d3c41ce75fb15d961c4d171ae0ae7dba10a0c22cf6286ec', epoch=20, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='6fbf57f19288695ac02d8d3146d0a17836329f7445a6333d83aa568380415f01', epoch=30, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='6fbf57f19288695ac02d8d3146d0a17836329f7445a6333d83aa568380415f01', epoch=10, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='6fbf57f19288695ac02d8d3146d0a17836329f7445a6333d83aa568380415f01', epoch=20, root_dir='trainings/s2s-verdi-x-44k'),\n",
       " Checkpoint(id='86ef6d33c9b53d6f8da13c1529551cfaa4443fed7d2c43becdc7f72aadbefbc4', epoch=15, root_dir='trainings/verdi-x-3-bis'),\n",
       " Checkpoint(id='86ef6d33c9b53d6f8da13c1529551cfaa4443fed7d2c43becdc7f72aadbefbc4', epoch=5, root_dir='trainings/verdi-x-3-bis'),\n",
       " Checkpoint(id='86ef6d33c9b53d6f8da13c1529551cfaa4443fed7d2c43becdc7f72aadbefbc4', epoch=25, root_dir='trainings/verdi-x-3-bis'),\n",
       " Checkpoint(id='86ef6d33c9b53d6f8da13c1529551cfaa4443fed7d2c43becdc7f72aadbefbc4', epoch=10, root_dir='trainings/verdi-x-3-bis'),\n",
       " Checkpoint(id='86ef6d33c9b53d6f8da13c1529551cfaa4443fed7d2c43becdc7f72aadbefbc4', epoch=20, root_dir='trainings/verdi-x-3-bis'),\n",
       " Checkpoint(id='cf0cd64e586978c81267c138e9a4e12a21054b1a22e30131477b7fd86e987826', epoch=50, root_dir='trainings/verdi-x-3-bis'),\n",
       " Checkpoint(id='cf0cd64e586978c81267c138e9a4e12a21054b1a22e30131477b7fd86e987826', epoch=30, root_dir='trainings/verdi-x-3-bis'),\n",
       " Checkpoint(id='cf0cd64e586978c81267c138e9a4e12a21054b1a22e30131477b7fd86e987826', epoch=10, root_dir='trainings/verdi-x-3-bis'),\n",
       " Checkpoint(id='cf0cd64e586978c81267c138e9a4e12a21054b1a22e30131477b7fd86e987826', epoch=40, root_dir='trainings/verdi-x-3-bis'),\n",
       " Checkpoint(id='cf0cd64e586978c81267c138e9a4e12a21054b1a22e30131477b7fd86e987826', epoch=20, root_dir='trainings/verdi-x-3-bis'),\n",
       " Checkpoint(id='f7efba79172a8c08960fa5092e4fb2d8bc4cc5cbc8595e97c5f6d5e1a318ca84', epoch=60, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='f7efba79172a8c08960fa5092e4fb2d8bc4cc5cbc8595e97c5f6d5e1a318ca84', epoch=100, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='f7efba79172a8c08960fa5092e4fb2d8bc4cc5cbc8595e97c5f6d5e1a318ca84', epoch=40, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='f7efba79172a8c08960fa5092e4fb2d8bc4cc5cbc8595e97c5f6d5e1a318ca84', epoch=20, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='f7efba79172a8c08960fa5092e4fb2d8bc4cc5cbc8595e97c5f6d5e1a318ca84', epoch=80, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='c041af0248ad10f3da08c1e3744aeae55a859a6d0339f25cf39136e330d38cdf', epoch=60, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='c041af0248ad10f3da08c1e3744aeae55a859a6d0339f25cf39136e330d38cdf', epoch=40, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='c041af0248ad10f3da08c1e3744aeae55a859a6d0339f25cf39136e330d38cdf', epoch=20, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='85af6593cfc21ec1bacfc6d16e4f6af01cb3c7c22cebae4af764eabe4aa8089d', epoch=60, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='85af6593cfc21ec1bacfc6d16e4f6af01cb3c7c22cebae4af764eabe4aa8089d', epoch=100, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='85af6593cfc21ec1bacfc6d16e4f6af01cb3c7c22cebae4af764eabe4aa8089d', epoch=40, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='85af6593cfc21ec1bacfc6d16e4f6af01cb3c7c22cebae4af764eabe4aa8089d', epoch=20, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='85af6593cfc21ec1bacfc6d16e4f6af01cb3c7c22cebae4af764eabe4aa8089d', epoch=80, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='f9aa425541545a41a1eb0f8af9c9eef11e0c1f625411712a6bd480757317891a', epoch=60, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='f9aa425541545a41a1eb0f8af9c9eef11e0c1f625411712a6bd480757317891a', epoch=100, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='f9aa425541545a41a1eb0f8af9c9eef11e0c1f625411712a6bd480757317891a', epoch=40, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='f9aa425541545a41a1eb0f8af9c9eef11e0c1f625411712a6bd480757317891a', epoch=20, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='f9aa425541545a41a1eb0f8af9c9eef11e0c1f625411712a6bd480757317891a', epoch=80, root_dir='trainings/srnn-verdi-x'),\n",
       " Checkpoint(id='7f07011c0a621a6a02edafd7132e6874f017f81403e28bc727f218c4c366b7a6', epoch=50, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7f07011c0a621a6a02edafd7132e6874f017f81403e28bc727f218c4c366b7a6', epoch=30, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7f07011c0a621a6a02edafd7132e6874f017f81403e28bc727f218c4c366b7a6', epoch=10, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7f07011c0a621a6a02edafd7132e6874f017f81403e28bc727f218c4c366b7a6', epoch=40, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7f07011c0a621a6a02edafd7132e6874f017f81403e28bc727f218c4c366b7a6', epoch=20, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='c48be95495cd518fa962332604fb9c571103d07a5c63d292fb8d7f7d20469743', epoch=50, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='c48be95495cd518fa962332604fb9c571103d07a5c63d292fb8d7f7d20469743', epoch=30, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='c48be95495cd518fa962332604fb9c571103d07a5c63d292fb8d7f7d20469743', epoch=10, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='c48be95495cd518fa962332604fb9c571103d07a5c63d292fb8d7f7d20469743', epoch=40, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='c48be95495cd518fa962332604fb9c571103d07a5c63d292fb8d7f7d20469743', epoch=20, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7b7cc1b88850bb4dc05b38fcfa1b786101830386d6998a946d57e84b10384fef', epoch=50, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7b7cc1b88850bb4dc05b38fcfa1b786101830386d6998a946d57e84b10384fef', epoch=30, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7b7cc1b88850bb4dc05b38fcfa1b786101830386d6998a946d57e84b10384fef', epoch=10, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7b7cc1b88850bb4dc05b38fcfa1b786101830386d6998a946d57e84b10384fef', epoch=40, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7b7cc1b88850bb4dc05b38fcfa1b786101830386d6998a946d57e84b10384fef', epoch=20, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='dbe5b4c5f2b915f647cba66510a9570c4640ca712c0dd434368fed18392a1567', epoch=50, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='dbe5b4c5f2b915f647cba66510a9570c4640ca712c0dd434368fed18392a1567', epoch=30, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='dbe5b4c5f2b915f647cba66510a9570c4640ca712c0dd434368fed18392a1567', epoch=10, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='dbe5b4c5f2b915f647cba66510a9570c4640ca712c0dd434368fed18392a1567', epoch=40, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='dbe5b4c5f2b915f647cba66510a9570c4640ca712c0dd434368fed18392a1567', epoch=20, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7faaa866d9e8360a29776a270737343341a25d646811fc117cb9ff604137b6b1', epoch=50, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7faaa866d9e8360a29776a270737343341a25d646811fc117cb9ff604137b6b1', epoch=30, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7faaa866d9e8360a29776a270737343341a25d646811fc117cb9ff604137b6b1', epoch=10, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7faaa866d9e8360a29776a270737343341a25d646811fc117cb9ff604137b6b1', epoch=40, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7faaa866d9e8360a29776a270737343341a25d646811fc117cb9ff604137b6b1', epoch=20, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7c690428ffe4190d65f0dae6655ad4468633ac7b305ca84b392d3a033c1e52d5', epoch=50, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7c690428ffe4190d65f0dae6655ad4468633ac7b305ca84b392d3a033c1e52d5', epoch=30, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7c690428ffe4190d65f0dae6655ad4468633ac7b305ca84b392d3a033c1e52d5', epoch=10, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7c690428ffe4190d65f0dae6655ad4468633ac7b305ca84b392d3a033c1e52d5', epoch=40, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='7c690428ffe4190d65f0dae6655ad4468633ac7b305ca84b392d3a033c1e52d5', epoch=20, root_dir='trainings/s2s-verdi-x'),\n",
       " Checkpoint(id='8deaf18e14e3a86f4039c6c3a80ab511e31a987fcc3daccfc6c0b7d3aae621f7', epoch=15, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='8deaf18e14e3a86f4039c6c3a80ab511e31a987fcc3daccfc6c0b7d3aae621f7', epoch=5, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='8deaf18e14e3a86f4039c6c3a80ab511e31a987fcc3daccfc6c0b7d3aae621f7', epoch=25, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='8deaf18e14e3a86f4039c6c3a80ab511e31a987fcc3daccfc6c0b7d3aae621f7', epoch=10, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='8deaf18e14e3a86f4039c6c3a80ab511e31a987fcc3daccfc6c0b7d3aae621f7', epoch=20, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='8cec89f4539ed86e57f9cf5e1193dc2586867828288ee67800c318dca95aedd3', epoch=15, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='8cec89f4539ed86e57f9cf5e1193dc2586867828288ee67800c318dca95aedd3', epoch=5, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='8cec89f4539ed86e57f9cf5e1193dc2586867828288ee67800c318dca95aedd3', epoch=25, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='8cec89f4539ed86e57f9cf5e1193dc2586867828288ee67800c318dca95aedd3', epoch=10, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='8cec89f4539ed86e57f9cf5e1193dc2586867828288ee67800c318dca95aedd3', epoch=20, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='821c83e498c30e71fb5456dde192fcdf31d22e88bf34bb54a4c887606bde6ac7', epoch=15, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='821c83e498c30e71fb5456dde192fcdf31d22e88bf34bb54a4c887606bde6ac7', epoch=5, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='821c83e498c30e71fb5456dde192fcdf31d22e88bf34bb54a4c887606bde6ac7', epoch=25, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='821c83e498c30e71fb5456dde192fcdf31d22e88bf34bb54a4c887606bde6ac7', epoch=10, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='821c83e498c30e71fb5456dde192fcdf31d22e88bf34bb54a4c887606bde6ac7', epoch=20, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='319461474819f3973f01b92d7647f376ced667cbf23f6f2fd111ae173097efe8', epoch=15, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='319461474819f3973f01b92d7647f376ced667cbf23f6f2fd111ae173097efe8', epoch=5, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='319461474819f3973f01b92d7647f376ced667cbf23f6f2fd111ae173097efe8', epoch=25, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='319461474819f3973f01b92d7647f376ced667cbf23f6f2fd111ae173097efe8', epoch=10, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='319461474819f3973f01b92d7647f376ced667cbf23f6f2fd111ae173097efe8', epoch=20, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='7bde91869b4011558313299d156d821f412aa9d8eb26e92cb5690d8870d7572d', epoch=15, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='7bde91869b4011558313299d156d821f412aa9d8eb26e92cb5690d8870d7572d', epoch=5, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='7bde91869b4011558313299d156d821f412aa9d8eb26e92cb5690d8870d7572d', epoch=25, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='7bde91869b4011558313299d156d821f412aa9d8eb26e92cb5690d8870d7572d', epoch=10, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='7bde91869b4011558313299d156d821f412aa9d8eb26e92cb5690d8870d7572d', epoch=20, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='d1044de5f29bce635d65c46883a13df408da1468b1debe0568129506d8588ab5', epoch=15, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='d1044de5f29bce635d65c46883a13df408da1468b1debe0568129506d8588ab5', epoch=5, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='d1044de5f29bce635d65c46883a13df408da1468b1debe0568129506d8588ab5', epoch=25, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='d1044de5f29bce635d65c46883a13df408da1468b1debe0568129506d8588ab5', epoch=10, root_dir='trainings/wn-verdi-x'),\n",
       " Checkpoint(id='d1044de5f29bce635d65c46883a13df408da1468b1debe0568129506d8588ab5', epoch=20, root_dir='trainings/wn-verdi-x')]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "root = \"trainings/\"\n",
    "\n",
    "ckpts = [Checkpoint.from_path(f) for f in h5m.FileWalker(r\".*verdi.*epoch=.*\\.h5\", root)\n",
    "        if \"qx\" not in f and '88k' not in f]\n",
    "ckpts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "359bfc2d",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "\n",
      "-----------------------------------------\n",
      "\n",
      "\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Generate', layout=Layout(flex='2'), max=4306.0, style=Pro…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Generate', layout=Layout(flex='2'), max=4306.0, style=Pro…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Generate', layout=Layout(flex='2'), max=4306.0, style=Pro…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Generate', layout=Layout(flex='2'), max=4306.0, style=Pro…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Generate', layout=Layout(flex='2'), max=4306.0, style=Pro…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Generate', layout=Layout(flex='2'), max=4306.0, style=Pro…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    }
   ],
   "source": [
    "import gc\n",
    "\n",
    "for ck in ckpts[:]:\n",
    "    net = ck.network\n",
    "    feature = ck.feature\n",
    "    bank = Trainset(keyword=ck.train_hp[\"trainset\"], sr=feature.sr)\n",
    "    saved = {}\n",
    "    \n",
    "    \n",
    "    def process_outputs(outputs, bidx):\n",
    "        outputs = outputs[0]\n",
    "        y = feature.transform(bank.bank.snd[:])\n",
    "        y = torch.from_numpy(y).to(outputs)\n",
    "        nn = torch.stack([nearest_neighbor(out, y)[1] for out in outputs])\n",
    "        hx = torch.stack([cum_entropy(n, neg_diff=False) for n in nn]).detach().cpu().numpy()\n",
    "        idx = np.argsort(hx)\n",
    "        for i in idx:\n",
    "            saved[hx[i]] = outputs[i].detach().cpu().numpy().T\n",
    "        del y\n",
    "        del nn\n",
    "        gc.collect()\n",
    "        torch.cuda.empty_cache()\n",
    "        \n",
    "\n",
    "    prompt_files = bank.bank\n",
    "    batch_item = feature.batch_item(shift=0, length=net.rf, training=False)\n",
    "    indices = mmk.IndicesSampler(N=500,\n",
    "                                  indices=torch.arange(0,\n",
    "                                                       prompt_files.snd.shape[0]-batch_item.getter.length,\n",
    "                                                       (prompt_files.snd.shape[0]-batch_item.getter.length)//500))\n",
    "    dl = prompt_files.serve(\n",
    "        (batch_item, ),\n",
    "        sampler=indices,\n",
    "        shuffle=False,\n",
    "        batch_size=64,\n",
    "    )\n",
    "    \n",
    "    \n",
    "    loop = mmk.GenerateLoop(\n",
    "        network=net,\n",
    "        dataloader=dl,\n",
    "        inputs=(h5m.Input(None, \n",
    "                          getter=h5m.AsSlice(dim=1, shift=-net.rf, length=net.rf),\n",
    "                          setter=h5m.Setter(dim=1)),),\n",
    "        n_steps=feature.sr*25//feature.hop_length,\n",
    "        add_blank=True,\n",
    "        time_hop=net.hp.get(\"hop\", 1),\n",
    "        process_outputs=process_outputs\n",
    "    )\n",
    "#     print(\"\\n\")\n",
    "    print(\"\\n\")\n",
    "    print(\"\\n\")\n",
    "    print(\"-----------------------------------------\")\n",
    "    print(\"\\n\")\n",
    "    print(\"\\n\")  \n",
    "    loop.run()\n",
    "    print(\"\\n\")\n",
    "    print(\"\\n\")\n",
    "    print(ck)\n",
    "\n",
    "    for k in list(sorted(saved))[-8:]:\n",
    "        print(\"SCORE = \", k)\n",
    "        mmk.audio(saved[k], hop_length=feature.hop_length, sr=feature.sr)\n",
    "        \n",
    "    del net\n",
    "    gc.collect()\n",
    "    torch.cuda.empty_cache()\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83d34e03",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b4036e6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "950428e3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b83bf5c9",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "CKPTS = group_ckpts_by_trainset(\"trainings\")\n",
    "[*CKPTS.keys()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0d23efa",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "k = [*CKPTS.keys()][7]\n",
    "CKPTS[k]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fff4132",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import h5mapper as h5m\n",
    "import mimikit as mmk\n",
    "from pbind import *\n",
    "\n",
    "from models.ensemble import Ensemble\n",
    "from datasets import COUGH\n",
    "from checkpoints import load_files\n",
    "\n",
    "stream = Pseq([\n",
    "    Pbind(\n",
    "        \"id\", \"80cb7d5b4ff7af169e74b3617c43580a41d5de5bd6c25e3251db2d11213755cd\",\n",
    "        \"epoch\", Prand([40, 50], inf),\n",
    "        \"seconds\", Pwhite(lo=1., hi=8., repeats=1)\n",
    "        ),\n",
    "    Pbind(\n",
    "        \"id\", \"80cb7d5b4ff7af169e74b3617c43580a41d5de5bd6c25e3251db2d11213755cd\",\n",
    "        \"epoch\", Prand([40, 50], inf),\n",
    "        \"seconds\", Pwhite(lo=0.5, hi=1.5, repeats=1)\n",
    "        ),\n",
    "], inf).asStream()\n",
    "    \n",
    "ensemble = Ensemble(60., 22050, stream)\n",
    "\n",
    "def process_outputs(outputs, bidx):\n",
    "    mmk.audio(outputs[0][0].cpu().numpy(), sr=ensemble.base_sr)\n",
    "\n",
    "prompt_files = load_files(COUGH[\"Cough\"], ensemble.base_sr)\n",
    "prompt = prompt_files.snd[0:44100]\n",
    "prompt = torch.as_tensor(prompt).unsqueeze(0)\n",
    "\n",
    "loop = mmk.GenerateLoop(\n",
    "    network=ensemble,\n",
    "    dataloader=[(prompt,)],\n",
    "    inputs=(h5m.Input(None, \n",
    "                      getter=h5m.AsSlice(dim=1, shift=-ensemble.base_sr, length=ensemble.base_sr),\n",
    "                      setter=h5m.Setter(dim=1)),),\n",
    "    n_steps=int(ensemble.base_sr * ensemble.max_seconds),\n",
    "    add_blank=True,\n",
    "    process_outputs=process_outputs\n",
    ")\n",
    "loop.run()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eef9a65c",
   "metadata": {},
   "outputs": [],
   "source": [
    "feat = mmk.Spectrogram(sr=44100, normalize=True, emphasis=0.0, n_fft=1024, hop_length=256, coordinate='mag', center=True)\n",
    "feat.inverse_transform_(torch.randn(1, 20, 513)).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe2c4f77",
   "metadata": {},
   "outputs": [],
   "source": [
    "ensemble.max_seconds * ensemble.base_sr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ed95656",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08894ee5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecaa8544",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d515bde",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "8aa6a601",
   "metadata": {},
   "source": [
    "# Stream Declaration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09ac43b2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import h5mapper as h5m\n",
    "\n",
    "from checkpoints import group_ckpts_by_trainset\n",
    "\n",
    "group_ckpts_by_trainset()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9625ed16",
   "metadata": {},
   "source": [
    "# Split Checkpoint Banks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bc860d4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import h5mapper as h5m\n",
    "from google.cloud import storage\n",
    "\n",
    "from checkpoints import CkptBank, load_trainings_hp, load_network_cls, Checkpoint\n",
    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
    "\n",
    "def split_by_src(ckpt_path):\n",
    "    print(ckpt_path)\n",
    "    bank = CkptBank(ckpt_path)\n",
    "    hp = bank.ckpt.load_hp()\n",
    "    to_upload = []\n",
    "    \n",
    "    dirname = os.path.dirname(ckpt_path)\n",
    "    train_hp = load_trainings_hp(dirname)\n",
    "    net_cls = load_network_cls(train_hp[\"network_class\"])\n",
    "    hp[\"cls\"] = net_cls\n",
    "    for ep_id in bank.index.keys():\n",
    "        new_path = os.path.join(\n",
    "            dirname, ep_id.split(\"-\")[0] + \".h5\"\n",
    "        )\n",
    "        if os.path.isfile(new_path):\n",
    "            to_upload += [new_path]\n",
    "            continue\n",
    "        new = CkptBank(new_path, mode=\"w\")\n",
    "        new.ckpt.save_hp(hp)\n",
    "        new.flush()\n",
    "        new.ckpt.add(\"state_dict\", new.ckpt.format(bank.get(ep_id)['ckpt']))\n",
    "        new.flush()\n",
    "        new.close()\n",
    "        to_upload += [new_path]\n",
    "    return to_upload\n",
    "\n",
    "def upload_to_gcp(ckpt_path):\n",
    "#     raise ValueError\n",
    "    ck = Checkpoint(*Checkpoint.get_id_and_epoch(ckpt_path))\n",
    "    if not ck.blob.exists():\n",
    "        print(\"uploading\", ck)\n",
    "        ck.blob.upload_from_filename(ckpt_path, timeout=None)\n",
    "    print(ck.blob, ck.blob.exists())\n",
    "    return\n",
    "\n",
    "\n",
    "\n",
    "# to_split = h5m.FileWalker(r\"checkpoints\\.h5\", \"/home/antoine/ktonal/ax6/trainings/s2s-grid-cough\")\n",
    "to_upload = h5m.FileWalker(r\"epoch=.*\\.h5\", \"/home/antoine/ktonal/ax6/trainings/s2s-grid-lungs\")\n",
    "\n",
    "executor = ThreadPoolExecutor(max_workers=4)\n",
    "# to_upload = [path for x in to_split for path in split_by_src(x)]\n",
    "\n",
    "as_completed([*executor.map(upload_to_gcp, to_upload)])\n",
    "executor.shutdown()\n",
    "1, 2, 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfbf69a2",
   "metadata": {},
   "source": [
    "# Download Checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6700a58b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from google.cloud import storage\n",
    "import dataclasses as dtc\n",
    "import os\n",
    "\n",
    "from checkpoints import CkptBank, load_feature\n",
    "\n",
    "\n",
    "client = storage.Client(\"ax6-Project\")\n",
    "\n",
    "@dtc.dataclass\n",
    "class Checkpoint:\n",
    "    id: str\n",
    "    epoch: int\n",
    "    bucket = \"ax6-outputs\"\n",
    "    root_dir = \"./\"\n",
    "    \n",
    "    @staticmethod\n",
    "    def get_id_and_epoch(path):\n",
    "        id_, epoch = path.split(\"/\")[-2:]\n",
    "        return id_.strip(\"/\"), int(epoch.split(\".h5\")[0].split(\"=\")[-1])\n",
    "    \n",
    "    @staticmethod\n",
    "    def from_blob(blob):\n",
    "        path = blob.name\n",
    "        id_, epoch = Checkpoint.get_id_and_epoch(path)\n",
    "        ckpt = Checkpoint(id_, epoch)\n",
    "        ckpt.bucket = blob.bucket.name\n",
    "        return ckpt\n",
    "    \n",
    "    @property\n",
    "    def gcp_path(self):\n",
    "        return f\"gs://{self.bucket}/checkpoints/{self.id}/epoch={self.epoch}.h5\"\n",
    "    \n",
    "    @property\n",
    "    def os_path(self):\n",
    "        return os.path.join(self.root_dir, f\"{self.id}_epoch={self.epoch}.h5\")\n",
    "    \n",
    "    @property\n",
    "    def blob(self):\n",
    "        return client.bucket(self.bucket).blob(f\"checkpoints/{self.id}/epoch={self.epoch}.h5\")\n",
    "    \n",
    "    def download(self):\n",
    "        os.makedirs(self.root_dir, exist_ok=True)\n",
    "        client.download_blob_to_file(self.gcp_path, open(self.os_path, \"wb\"))\n",
    "        return self\n",
    "    \n",
    "    @property\n",
    "    def network(self):\n",
    "        if not os.path.isfile(self.os_path):\n",
    "            self.download()\n",
    "        bank = CkptBank(self.os_path, 'r')\n",
    "        hp = bank.ckpt.load_hp()\n",
    "        return bank.ckpt.load_checkpoint(hp[\"cls\"], \"state_dict\")\n",
    "    \n",
    "    @property\n",
    "    def feature(self):\n",
    "        if not os.path.isfile(self.os_path):\n",
    "            self.download()\n",
    "        bank = CkptBank(self.os_path, 'r')\n",
    "        hp = bank.ckpt.load_hp()\n",
    "        return hp['feature']\n",
    "    \n",
    "\n",
    "[Checkpoint.from_blob(blob) \n",
    " for blob in client.list_blobs(Checkpoint.bucket, prefix='checkpoints')\n",
    " if \"epoch=\" in blob.name and \".h5\" == blob.name[-3:]\n",
    "][-1].feature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51379730",
   "metadata": {},
   "outputs": [],
   "source": [
    "list(client.list_blobs(\"ax6-outputs\", prefix='checkpoints'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76592396",
   "metadata": {},
   "outputs": [],
   "source": [
    "from models.nnn import *\n",
    "from checkpoints import *\n",
    "from datasets import *\n",
    "import mimikit as mmk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1610101f",
   "metadata": {},
   "outputs": [],
   "source": [
    "bank = Trainset(\"Cough\").bank\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97584351",
   "metadata": {},
   "outputs": [],
   "source": [
    "fft = mmk.Spectrogram(n_fft=2048, hop_length=512, coordinate=\"mag\").transform\n",
    "optimal_path(fft(bank.snd[:8000]), fft(bank.snd[5000:15000]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b985715d",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "nnn = NearestNextNeighbor(\n",
    "    mmk.Spectrogram(n_fft=2048, hop_length=512, coordinate=\"mag\"),\n",
    "    bank.snd\n",
    ")\n",
    "prompt = fft(bank.snd[3000:8000])\n",
    "\n",
    "mmk.GenerateLoop(\n",
    "    nnn,\n",
    "    [(torch.from_numpy(prompt).unsqueeze(0), )],\n",
    "    inputs=(h5m.Input(None,\n",
    "        getter=h5m.AsSlice(dim=1, shift=-22050, length=22050),\n",
    "        setter=h5m.Setter(dim=1)),),\n",
    "    n_steps=32,\n",
    "    device=\"cpu\"\n",
    ").run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d494e23b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47b43af8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}