<jupyter_start><jupyter_text><jupyter_code>#@markdown Before starting please save the notebook in your drive by clicking on `File -> Save a copy in drive`
#@markdown Check GPU, should be a Tesla V100
!nvidia-smi -L
import os
print(f"We have {os.cpu_count()} CPU cores.")
#@markdown Mount google drive
from google.colab import drive
drive.mount('/content/drive')

from pathlib import Path
if not Path("/content/drive/My Drive/IRCMS_GAN_collaborative_database").exists():
    raise RuntimeError(
        "Shortcut to our shared drive folder doesn't exits.\n\n"
        "\t1. Go to the google drive web UI\n"
        "\t2. Right click shared folder IRCMS_GAN_collaborative_database and click \"Add shortcut to Drive\""
    )

def clear_on_success(msg="Ok!"):
    if _exit_code == 0:
        output.clear()
        print(msg)
#@markdown Install wandb and log in
from google.colab import output
!pip install wandb==0.10.8
clear_on_success()
!wandb login
clear_on_success("Ok! Logged into wandb!")
#@title Configuration

#@markdown Directories can be found via file explorer on the left by navigating into `drive` to the desired folders. 
#@markdown Then right-click and `Copy path`.
audio_db_dir = "/content/drive/My Drive/AUDIO DATABASE/RAW Sessions/Roberto Studio Material 22050" #@param {type:"string"}
experiment_dir = "/content/drive/My Drive/IRCMS_GAN_collaborative_database/Experiments/colab-violingan/melgan" #@param {type:"string"}
melgan_output_dir = "/content/drive/My Drive/IRCMS_GAN_collaborative_database/Experiments/colab-violingan/melgan-outputs" #@param {type:"string"}

#@markdown ### Resumption of previous runs
#@markdown Optional resumption arguments below, leaving both empty will start a new run from scratch. 
#@markdown - The ID can be found on wandb. 
#@markdown - It's 8 characters long and may contain a-z letters and digits (for example `1t212ycn`).

#@markdown Resume a previous run 
resume_run_id = "" #@param {type:"string"}
#@markdown Load initial weights from a previous run to start a new run.
load_from_run_id = "" #@param {type:"string"}

import re
from pathlib import Path

audio_db_dir = Path(audio_db_dir)
melgan_output_dir = Path(melgan_output_dir)
experiment_dir = Path(experiment_dir)

for path in [experiment_dir, melgan_output_dir]:
    path.mkdir(parents=True, exist_ok=True)

if not audio_db_dir.exists():
    raise RuntimeError(f"audio_db_dir {audio_db_dir} does not exists.")

if resume_run_id and load_from_run_id:
    raise RuntimeError("Only set `resume_run_id` or `load_from_run_id`.")

def check_wandb_id(run_id):
    if run_id and not re.match(r"^[\da-z]{8}$", run_id):
        raise RuntimeError(
            "Run ID needs to be 8 characters long and contain only letters a-z and digits.\n"
            f"Got \"{run_id}\""
        )

check_wandb_id(resume_run_id)
check_wandb_id(load_from_run_id)
!git clone https://github.com/buganart/melgan-neurips
%cd /content/melgan-neurips
%pip install -r requirements.txt
clear_on_success("Dependencies installed.")
local_wav_dir = Path("/content/wavs/")
!find "{audio_db_dir}" -maxdepth 1 -type f | xargs -t -d "\n" -I'%%' -P 10 -n 1 rsync -a '%%' "$local_wav_dir"/
clear_on_success("All files copied to this runtime.")

audio_paths = sorted(list(local_wav_dir.glob("*")))
num_files = len(audio_paths)
print(f"Database has {num_files} files in total.")
import numpy as np

np.random.seed(123)
test_fraction = 0.1
shuffled_paths = np.random.permutation(audio_paths)
split_at = int((1 - test_fraction) * num_files)
train_paths = shuffled_paths[:split_at]
test_paths = shuffled_paths[split_at:]

with open('train_files.txt', 'w') as f:
    for path in sorted(train_paths):
        print(path, file=f)

with open('test_files.txt', 'w') as f:
    for path in sorted(test_paths):
        print(path, file=f)

print("TRAIN FILES")
!head -n3 train_files.txt
print('...')
!tail -n3 train_files.txt

print()
print("TEST FILES")
!head -n3 test_files.txt
print('...')
!tail -n3 test_files.txt
# This done a bit weirdly because setting PYTHONPATH=$PWD removes variables afterwards. A colab bug, maybe.
!env PYTHONPATH="$(pwd)" python scripts/train.py --batch_size 16 --save_path "$experiment_dir" --data_path . --resume_run_id "$resume_run_id" --load_from_run_id "$load_from_run_id"
#@markdown Copy the run directory from the beginning of the cell above. 

#@markdown The line looks something like

#@markdown `wandb: Run data is saved locally in /content/drive/.../melgan/wandb/run-20201030_030400-21bhoeww`

#@markdown and we only need the part after the last slash.

run_dir_name = "run-20201030_030400-21bhoeww" #@param {type:"string"}

run_dir = experiment_dir / "wandb" / run_dir_name
export_dir = melgan_output_dir / run_dir_name
export_dir.mkdir(exist_ok=True, parents=True)

melgan_args = run_dir / "files" / "args.yml"
melgan_netG = run_dir / "files" / "best_netG.pt" 

!cp -v "$melgan_args" "$export_dir"
!cp -v "$melgan_netG" "$export_dir"
!cp -v "/content/melgan-neurips/mel2wav/modules.py" "$export_dir"

print()
print(f"Melgan files saved in {export_dir}.")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12064554221609919, 'fraction_numerical': 0.013973627238732533, 'mean_word_length': 4.76844494892168, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 6, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '515678', 'license': 'MIT', 'n_tokens_mistral': 1843, 'n_tokens_neox': 1753, 'n_words': 501, 'path': 'melgan.ipynb', 'repo_name': 'materialvision/melgan-neurips', 'size': '5081'}