<jupyter_start><jupyter_text>Table of Contents1&nbsp;&nbsp;Main Table of Contents1&nbsp;&nbsp;Main Table of Contents1&nbsp;&nbsp;New heading<jupyter_code># default_exp exec.train_tacotron2
# export
import argparse
import json
import librosa  # NOTE(zach): importing torch before librosa causes LLVM issues for some unknown reason.
import sys

import torch
from torch import multiprocessing as mp

from uberduck_ml_dev.trainer.base import TTSTrainer, MellotronTrainer
from uberduck_ml_dev.vendor.tfcompat.hparam import HParams
from uberduck_ml_dev.models.mellotron import DEFAULTS as MELLOTRON_DEFAULTS


def parse_args(args):
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", help="Path to JSON config")
    args = parser.parse_args(args)
    return args
assert isinstance(MELLOTRON_DEFAULTS, HParams)
assert parse_args(["--config", "foo.json"]).config == "foo.json"
assert parse_args([]).config is None
assert HParams(**MELLOTRON_DEFAULTS.values()).values() == MELLOTRON_DEFAULTS.values()
# export
def run(rank, device_count, hparams):
    trainer = MellotronTrainer(hparams, rank=rank, world_size=device_count)
    try:
        trainer.train()
    except Exception as e:
        print(f"Exception raised while training: {e}")
        # TODO: save state.
        raise e<jupyter_output><empty_output><jupyter_text>Main<jupyter_code># export
try:
    from nbdev.imports import IN_NOTEBOOK
except:
    IN_NOTEBOOK = False
if __name__ == "__main__" and not IN_NOTEBOOK:
    args = parse_args(sys.argv[1:])
    config = MELLOTRON_DEFAULTS.values()
    if args.config:
        with open(args.config) as f:
            config.update(json.load(f))
    hparams = HParams(**config)
    if hparams.distributed_run:
        device_count = torch.cuda.device_count()
        mp.spawn(run, (device_count, hparams), device_count)
    else:
        run(None, None, hparams)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10193413486670151, 'fraction_numerical': 0.0026136957658128594, 'mean_word_length': 4.629411764705883, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, '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': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '151531', 'license': 'Apache-2.0', 'n_tokens_mistral': 651, 'n_tokens_neox': 639, 'n_words': 148, 'path': 'nbs/exec.train_tacotron2.ipynb', 'repo_name': 'A913440/uberduck-ml-dev', 'size': '1913'}