<reponame>payoto/graphcore_examples
# Copyright (c) 2021 Graphcore Ltd. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import logging
import os
import sys

from tensorflow.python import ipu
from fastspeech2 import FastSpeech2Config


def _str_to_bool(value):
    if isinstance(value, bool) or value is None:
        return value
    if value.lower() in {'false', 'f', '0', 'no', 'n'}:
        return False
    elif value.lower() in {'true', 't', '1', 'yes', 'y'}:
        return True
    raise argparse.ArgumentTypeError(f'{value} is not a valid boolean value')


def add_common_arguments(parser, required=True):
    group = parser.add_argument_group('Common options')

    # Training options
    group.add_argument('--train', action='store_true',
                       help="Whether to run training loop or not. Default is True.")
    group.add_argument('--eval', action='store_true',
                       help="""
                       Whether to run evaluation loop or not. Default is False.
                       If both `--train` and `--eval` are set, then the evaluation will be done after trainning.
                       """)
    group.add_argument('--batch-size', type=int,
                       help="Set batch-size for training graph")
    group.add_argument('--base-learning-rate', type=float, default=2e-5,
                       help="Base learning rate exponent (2**N). blr = lr /  bs")
    group.add_argument('--loss-scaling', type=float, default=1,
                       help="Loss scaling factor.")
    group.add_argument('--optimizer', type=str, default="adamw", choices=['adam', 'adamw'],
                       help="Optimizer")
    group.add_argument('--beta1', type=float, default=0.9,
                       help="Adam/AdamW beta1 coefficient.")
    group.add_argument('--beta2', type=float, default=0.999,
                       help="Adam/AdamW beta2 coefficient.")
    group.add_argument('--weight-decay-rate', type=float, default=0.0,
                       help="Weight decay to use during optimisation.")
    group.add_argument('--decay-steps', type=int, default=1000,
                       help="Learning rate schedule decay steps.")
    group.add_argument('--decay-date', type=float, default=0.9,
                       help="Learning rate schedule decay rate. Only for Exponential.")
    group.add_argument('--epsilon', type=float, default=1e-4,
                       help="Optimiser epsilon value.")
    group.add_argument('--lr-schedule', default='exponential', choices=["exponential", "cosine", "polynomial"],
                       help="Learning rate schedule function.")
    group.add_argument('--warmup', default=0.1, type=float,
                       help="Learning rate schedule warm-up period, in epochs (float) or number of steps (integer).")
    group.add_argument('--seed', default=None,
                       help="Seed for randomizing training")
    group.add_argument('--wandb', type=_str_to_bool, default=False,
                       help="Enable logging and experiment tracking with Weights & Biases.")
    group.add_argument('--wandb-name', type=str, default=None,
                       help="Override Weights&Biases run name.")
    group.add_argument('--init-checkpoint', type=str, default=None,
                       help='Initialise a new training session from this checkpoint.')
    group.add_argument('--epochs', default=None, type=int,
                       help='Number of epochs we want to let the training last.')
    group.add_argument('--epochs-per-save', default=5, type=int,
                       help='Number of epochs to save model.')
    group.add_argument('--steps-per-epoch', default=None, type=int,
                       help='Specifies the total number of steps to be performed per epoch.')

    # FastSpeech2 options
    group.add_argument('--vocab-size', type=int,
                       help="""Vocabulary size of `FastSpeech2`.""")
    # Encoder
    group.add_argument('--encoder-hidden-size', type=int,
                       help="""Size ofthe encoder layers and the pooler layer.""")
    group.add_argument('--encoder-num-hidden-layers', type=int,
                       help="""Number of hidden layers in the encoder.""")
    group.add_argument('--encoder-num-attention-heads', type=int,
                       help="""Number of attention heads for each attention layer in the encoder.""")
    group.add_argument('--encoder-intermediate-size', type=int,
                       help="""The size of the "intermediate" (i.e., feed-forward) layer in the encoder.""")
    group.add_argument('--encoder-intermediate-kernel-size', type=int,
                       help="""The convolution kernel size of the "intermediate" (i.e., feed-forward) layer in the encoder.""")
    group.add_argument('--encoder-hidden-act', type=int,
                       help="""The non-linear activation function (function or string) in the encoder and pooler.""")
    group.add_argument('--output-attentions', type=_str_to_bool, default=False,
                       help="Whether to output the attentions.")
    group.add_argument('--output-hidden-states', type=_str_to_bool, default=False,
                       help="Whether to output the hidden states.")
    group.add_argument('--layer-norm-eps', type=float, default=1e-4,
                       help="Layer Normalization epsilon value.")
    # Decoder
    group.add_argument('--decoder-hidden-size', type=int,
                       help="""Size ofthe decoder layers and the pooler layer.""")
    group.add_argument('--decoder-num-hidden-layers', type=int,
                       help="""Number of hidden layers in the decoder.""")
    group.add_argument('--decoder-num-attention-heads', type=int,
                       help="""Number of attention heads for each attention layer in the decoder.""")
    group.add_argument('--decoder-intermediate-size', type=int,
                       help="""The size of the "intermediate" (i.e., feed-forward) layer in the decoder.""")
    group.add_argument('--decoder-intermediate-kernel-size', type=int,
                       help="""The convolution kernel size of the "intermediate" (i.e., feed-forward) layer in the decoder.""")
    group.add_argument('--decoder-hidden-act', type=int,
                       help="""The non-linear activation function (function or string) in the decoder and pooler.""")

    # Duration predictor(Only for FastSpeech)
    group.add_argument('--duration-predictor-num-conv-layers', type=int, default=2,
                       help="""Number of convolution layers in the duration predictor.""")
    group.add_argument('--duration-predictor-kernel-size', type=int, default=3,
                       help="""The convolution kernel size of the duration predictor.""")
    group.add_argument('--duration-predictor-dropout-probs', type=int, default=0.5,
                       help="""The dropout probability for duration predictor.""")
    group.add_argument('--duration-predictor-filters', type=int, default=256,
                       help="""Number of filters in the duration predictor.""")
    # Variance predictor and Postnet
    group.add_argument('--variant-predictor-num-conv-layers', type=int, default=2,
                       help="""Number of convolution layers in the variant predictor.""")
    group.add_argument('--variant-predictor-kernel-size', type=int, default=3,
                       help="""The convolution kernel size of the variant predictor.""")
    group.add_argument('--variant-predictor-dropout-rate', type=int, default=0.5,
                       help="""The dropout probability for variant predictor.""")
    group.add_argument('--variant-predictor-filter', type=int, default=256,
                       help="""Number of filters in the variant predictor.""")
    group.add_argument('--use-postnet', type=_str_to_bool, default=True,
                       help="Whether to use postnet after decoder.")
    group.add_argument('--postnet-num-conv-layers', type=int,
                       help="""Number of convolution layers in the postnet.""")
    group.add_argument('--postnet-conv-kernel-size', type=int,
                       help="""The convolution kernel size of the postnet.""")
    group.add_argument('--postnet-dropout-rate', type=int,
                       help="""The dropout probability for postnet.""")
    group.add_argument('--postnet-conv-filters', type=int,
                       help="""Number of convolution filters in the postnet.""")

    group.add_argument('--hidden-dropout-prob', type=float, default=0.2,
                       help="""The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.""")
    group.add_argument('--attention-probs-dropout-prob', type=float, default=0.1,
                       help="""The dropout ratio for the attention probabilities.""")
    group.add_argument('--max-position-embeddings', type=int,
                       help="""The maximum sequence length that this model might ever be used with. Typically set this to something large just in case (e.g., 512 or 1024 or 2048).""")
    group.add_argument('--initializer-range', type=float, default=0.02,
                       help="""The stdev of the truncated-normal-initializer for initializing all weight matrices.""")
    group.add_argument('--max-seq-length', type=int, default=135,
                       help='The maximuim sequence length.')
    group.add_argument('--max-wave-length', type=int, default=1000,
                       help='The maximuim wave length.')
    group.add_argument('--num-mels', type=int, default=80,
                       help='The number of mel-spectrograms dimensions.')

    # IPU options
    pipeline_schedule_options = [
        _.name for _ in ipu.pipelining_ops.PipelineSchedule]
    schedulers_available = ['Clustering',
                            'PostOrder', 'LookAhead', 'ShortestPath']
    recomputation_modes_available = [
        p.name for p in ipu.pipelining_ops.RecomputationMode]

    group.add_argument('--gradient-accumulation-count', type=int, default=None,
                       help="Number of gradients to accumulate in the pipeline. Must also set --shards > 1.")
    group.add_argument('--pipeline-schedule', type=str, default="Interleaved", choices=pipeline_schedule_options,
                       help="Pipelining scheduler.")
    group.add_argument('--replicas', type=int, default=1,
                       help="Replicate graph over N workers to increase batch to batch-size*N")
    group.add_argument('--precision', type=str, default="16", choices=["16", "32"],
                       help="Precision of Ops(weights/activations/gradients) data types: 16, 32.")
    group.add_argument('--available-memory-proportion', type=str, default=0.23,
                       help="Proportion of IPU memory available to matmul operations. A list can be used to specify the value for each IPU.")
    group.add_argument('--variable-offloading', type=_str_to_bool, default=True,
                       help="Enable offloading of training variables into remote memory.")
    group.add_argument('--stochastic-rounding', type=_str_to_bool, default=True,
                       help="Enable stochastic rounding. Set to False when run evaluation.")
    group.add_argument("--xla-recompute", default=True, action="store_true",
                       help="Recompute activations during backward pass")
    group.add_argument('--fp-exceptions', default=False, action="store_true",
                       help="Enable floating-point exeptions.")
    group.add_argument('--partials-type', type=str, default="half", choices=["half", "float"],
                       help="Floating-point precision of data in matmul and convolution operations..")
    group.add_argument('--scheduler', type=str, default='Clustering', choices=schedulers_available,
                       help="Forces the compiler to use a specific scheduler when ordering the instructions.")
    group.add_argument('--recomputation-mode', type=str,
                       default="RecomputeAndBackpropagateInterleaved", choices=recomputation_modes_available)
    group.add_argument('--reduction-type', type=str, choices=['sum', 'mean'], default='mean',
                       help='The reduction type applied to the pipeline, the choice is between summation and mean.')
    group.add_argument('--weight-norm-clip', type=float, default=0.,
                       help='The value from which we want to clip the w_norm value, value of 0 is no weight clipping.')

    # Dataset options
    group.add_argument('--data-path', type=str, required=False,
                       help="path to training/validation dataset tfrecord file.")
    group.add_argument('--parallell-io-threads', type=int, default=4,
                       help="Number of cpu threads used to do data prefetch.")
    group.add_argument('--generated-data', action="store_true", default=False,
                       help="Generates synthetic-data on the host and then use it for training.")

    # Add logging-specific arguments
    group_log = parser.add_argument_group('Logging')
    group_log.add_argument('--log-dir', type=str,
                           help="Log and weights save directory")
    group_log.add_argument('--steps-per-logs', type=int, default=1,
                           help="Logs per epoch (if number of epochs specified)")
    return parser


def create_command_line_parser():
    parser = argparse.ArgumentParser(description="""FastSpeech2 on Graphcore's IPU""",
                                     add_help=False,
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    group = parser.add_argument_group("""Main options""")
    group.add_argument('--help', action='store_true',
                       default=False, help="Display help.")
    group.add_argument('--config', type=str,
                       help='FastSpeech2 configuration file in JSON format.')
    return parser


def create_all_options_parser():
    parser = create_command_line_parser()
    parser = add_common_arguments(parser)
    return parser


def make_global_options(task_specific_parsers=[]):
    # Parse command-line arguments
    command_line_parser = create_command_line_parser()
    all_options_parser = create_all_options_parser()

    for task_parser in task_specific_parsers:
        all_options_parser = task_parser(all_options_parser)

    known_command_line_args, unknown_command_line_args = command_line_parser.parse_known_args()

    if known_command_line_args.help or known_command_line_args.config is None:
        all_options_parser.print_help()
        sys.exit(os.EX_OK)

    # Parse options specified in the configuration file into
    config_file_path = known_command_line_args.config
    opts_from_config_file = FastSpeech2Config.from_json_file(config_file_path)

    # Build the global options structure from the default options
    current_options = vars(all_options_parser.parse_args())

    unknown_options = [opt for opt in opts_from_config_file.keys(
    ) if opt not in current_options.keys()]

    if unknown_options:
        logging.error(f"Unonwn options: {unknown_options}")
        sys.exit(os.EX_USAGE)

    # Overwrite global options by those specified in the config file.
    current_options.update(opts_from_config_file)
    options_namespace = argparse.Namespace(**current_options)

    # Overwrite with command-line arguments
    all_options_namespace = all_options_parser.parse_args(
        unknown_command_line_args, options_namespace)
    logging.info(
        f"Overwrite configuration parameters: {', '.join(unknown_command_line_args)}")

    # argparse.Namespace -> dict()
    opts = vars(all_options_namespace)

    return opts

# Copyright (c) 2021 Graphcore Ltd. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import logging
import os
import sys

from tensorflow.python import ipu
from fastspeech2 import FastSpeech2Config


def _str_to_bool(value):
    if isinstance(value, bool) or value is None:
        return value
    if value.lower() in {'false', 'f', '0', 'no', 'n'}:
        return False
    elif value.lower() in {'true', 't', '1', 'yes', 'y'}:
        return True
    raise argparse.ArgumentTypeError(f'{value} is not a valid boolean value')


def add_common_arguments(parser, required=True):
    group = parser.add_argument_group('Common options')

    # Training options
    group.add_argument('--train', action='store_true',
                       help="Whether to run training loop or not. Default is True.")
    group.add_argument('--eval', action='store_true',
                       help="""
                       Whether to run evaluation loop or not. Default is False.
                       If both `--train` and `--eval` are set, then the evaluation will be done after trainning.
                       """)
    group.add_argument('--batch-size', type=int,
                       help="Set batch-size for training graph")
    group.add_argument('--base-learning-rate', type=float, default=2e-5,
                       help="Base learning rate exponent (2**N). blr = lr /  bs")
    group.add_argument('--loss-scaling', type=float, default=1,
                       help="Loss scaling factor.")
    group.add_argument('--optimizer', type=str, default="adamw", choices=['adam', 'adamw'],
                       help="Optimizer")
    group.add_argument('--beta1', type=float, default=0.9,
                       help="Adam/AdamW beta1 coefficient.")
    group.add_argument('--beta2', type=float, default=0.999,
                       help="Adam/AdamW beta2 coefficient.")
    group.add_argument('--weight-decay-rate', type=float, default=0.0,
                       help="Weight decay to use during optimisation.")
    group.add_argument('--decay-steps', type=int, default=1000,
                       help="Learning rate schedule decay steps.")
    group.add_argument('--decay-date', type=float, default=0.9,
                       help="Learning rate schedule decay rate. Only for Exponential.")
    group.add_argument('--epsilon', type=float, default=1e-4,
                       help="Optimiser epsilon value.")
    group.add_argument('--lr-schedule', default='exponential', choices=["exponential", "cosine", "polynomial"],
                       help="Learning rate schedule function.")
    group.add_argument('--warmup', default=0.1, type=float,
                       help="Learning rate schedule warm-up period, in epochs (float) or number of steps (integer).")
    group.add_argument('--seed', default=None,
                       help="Seed for randomizing training")
    group.add_argument('--wandb', type=_str_to_bool, default=False,
                       help="Enable logging and experiment tracking with Weights & Biases.")
    group.add_argument('--wandb-name', type=str, default=None,
                       help="Override Weights&Biases run name.")
    group.add_argument('--init-checkpoint', type=str, default=None,
                       help='Initialise a new training session from this checkpoint.')
    group.add_argument('--epochs', default=None, type=int,
                       help='Number of epochs we want to let the training last.')
    group.add_argument('--epochs-per-save', default=5, type=int,
                       help='Number of epochs to save model.')
    group.add_argument('--steps-per-epoch', default=None, type=int,
                       help='Specifies the total number of steps to be performed per epoch.')

    # FastSpeech2 options
    group.add_argument('--vocab-size', type=int,
                       help="""Vocabulary size of `FastSpeech2`.""")
    # Encoder
    group.add_argument('--encoder-hidden-size', type=int,
                       help="""Size ofthe encoder layers and the pooler layer.""")
    group.add_argument('--encoder-num-hidden-layers', type=int,
                       help="""Number of hidden layers in the encoder.""")
    group.add_argument('--encoder-num-attention-heads', type=int,
                       help="""Number of attention heads for each attention layer in the encoder.""")
    group.add_argument('--encoder-intermediate-size', type=int,
                       help="""The size of the "intermediate" (i.e., feed-forward) layer in the encoder.""")
    group.add_argument('--encoder-intermediate-kernel-size', type=int,
                       help="""The convolution kernel size of the "intermediate" (i.e., feed-forward) layer in the encoder.""")
    group.add_argument('--encoder-hidden-act', type=int,
                       help="""The non-linear activation function (function or string) in the encoder and pooler.""")
    group.add_argument('--output-attentions', type=_str_to_bool, default=False,
                       help="Whether to output the attentions.")
    group.add_argument('--output-hidden-states', type=_str_to_bool, default=False,
                       help="Whether to output the hidden states.")
    group.add_argument('--layer-norm-eps', type=float, default=1e-4,
                       help="Layer Normalization epsilon value.")
    # Decoder
    group.add_argument('--decoder-hidden-size', type=int,
                       help="""Size ofthe decoder layers and the pooler layer.""")
    group.add_argument('--decoder-num-hidden-layers', type=int,
                       help="""Number of hidden layers in the decoder.""")
    group.add_argument('--decoder-num-attention-heads', type=int,
                       help="""Number of attention heads for each attention layer in the decoder.""")
    group.add_argument('--decoder-intermediate-size', type=int,
                       help="""The size of the "intermediate" (i.e., feed-forward) layer in the decoder.""")
    group.add_argument('--decoder-intermediate-kernel-size', type=int,
                       help="""The convolution kernel size of the "intermediate" (i.e., feed-forward) layer in the decoder.""")
    group.add_argument('--decoder-hidden-act', type=int,
                       help="""The non-linear activation function (function or string) in the decoder and pooler.""")

    # Duration predictor(Only for FastSpeech)
    group.add_argument('--duration-predictor-num-conv-layers', type=int, default=2,
                       help="""Number of convolution layers in the duration predictor.""")
    group.add_argument('--duration-predictor-kernel-size', type=int, default=3,
                       help="""The convolution kernel size of the duration predictor.""")
    group.add_argument('--duration-predictor-dropout-probs', type=int, default=0.5,
                       help="""The dropout probability for duration predictor.""")
    group.add_argument('--duration-predictor-filters', type=int, default=256,
                       help="""Number of filters in the duration predictor.""")
    # Variance predictor and Postnet
    group.add_argument('--variant-predictor-num-conv-layers', type=int, default=2,
                       help="""Number of convolution layers in the variant predictor.""")
    group.add_argument('--variant-predictor-kernel-size', type=int, default=3,
                       help="""The convolution kernel size of the variant predictor.""")
    group.add_argument('--variant-predictor-dropout-rate', type=int, default=0.5,
                       help="""The dropout probability for variant predictor.""")
    group.add_argument('--variant-predictor-filter', type=int, default=256,
                       help="""Number of filters in the variant predictor.""")
    group.add_argument('--use-postnet', type=_str_to_bool, default=True,
                       help="Whether to use postnet after decoder.")
    group.add_argument('--postnet-num-conv-layers', type=int,
                       help="""Number of convolution layers in the postnet.""")
    group.add_argument('--postnet-conv-kernel-size', type=int,
                       help="""The convolution kernel size of the postnet.""")
    group.add_argument('--postnet-dropout-rate', type=int,
                       help="""The dropout probability for postnet.""")
    group.add_argument('--postnet-conv-filters', type=int,
                       help="""Number of convolution filters in the postnet.""")

    group.add_argument('--hidden-dropout-prob', type=float, default=0.2,
                       help="""The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.""")
    group.add_argument('--attention-probs-dropout-prob', type=float, default=0.1,
                       help="""The dropout ratio for the attention probabilities.""")
    group.add_argument('--max-position-embeddings', type=int,
                       help="""The maximum sequence length that this model might ever be used with. Typically set this to something large just in case (e.g., 512 or 1024 or 2048).""")
    group.add_argument('--initializer-range', type=float, default=0.02,
                       help="""The stdev of the truncated-normal-initializer for initializing all weight matrices.""")
    group.add_argument('--max-seq-length', type=int, default=135,
                       help='The maximuim sequence length.')
    group.add_argument('--max-wave-length', type=int, default=1000,
                       help='The maximuim wave length.')
    group.add_argument('--num-mels', type=int, default=80,
                       help='The number of mel-spectrograms dimensions.')

    # IPU options
    pipeline_schedule_options = [
        _.name for _ in ipu.pipelining_ops.PipelineSchedule]
    schedulers_available = ['Clustering',
                            'PostOrder', 'LookAhead', 'ShortestPath']
    recomputation_modes_available = [
        p.name for p in ipu.pipelining_ops.RecomputationMode]

    group.add_argument('--gradient-accumulation-count', type=int, default=None,
                       help="Number of gradients to accumulate in the pipeline. Must also set --shards > 1.")
    group.add_argument('--pipeline-schedule', type=str, default="Interleaved", choices=pipeline_schedule_options,
                       help="Pipelining scheduler.")
    group.add_argument('--replicas', type=int, default=1,
                       help="Replicate graph over N workers to increase batch to batch-size*N")
    group.add_argument('--precision', type=str, default="16", choices=["16", "32"],
                       help="Precision of Ops(weights/activations/gradients) data types: 16, 32.")
    group.add_argument('--available-memory-proportion', type=str, default=0.23,
                       help="Proportion of IPU memory available to matmul operations. A list can be used to specify the value for each IPU.")
    group.add_argument('--variable-offloading', type=_str_to_bool, default=True,
                       help="Enable offloading of training variables into remote memory.")
    group.add_argument('--stochastic-rounding', type=_str_to_bool, default=True,
                       help="Enable stochastic rounding. Set to False when run evaluation.")
    group.add_argument("--xla-recompute", default=True, action="store_true",
                       help="Recompute activations during backward pass")
    group.add_argument('--fp-exceptions', default=False, action="store_true",
                       help="Enable floating-point exeptions.")
    group.add_argument('--partials-type', type=str, default="half", choices=["half", "float"],
                       help="Floating-point precision of data in matmul and convolution operations..")
    group.add_argument('--scheduler', type=str, default='Clustering', choices=schedulers_available,
                       help="Forces the compiler to use a specific scheduler when ordering the instructions.")
    group.add_argument('--recomputation-mode', type=str,
                       default="RecomputeAndBackpropagateInterleaved", choices=recomputation_modes_available)
    group.add_argument('--reduction-type', type=str, choices=['sum', 'mean'], default='mean',
                       help='The reduction type applied to the pipeline, the choice is between summation and mean.')
    group.add_argument('--weight-norm-clip', type=float, default=0.,
                       help='The value from which we want to clip the w_norm value, value of 0 is no weight clipping.')

    # Dataset options
    group.add_argument('--data-path', type=str, required=False,
                       help="path to training/validation dataset tfrecord file.")
    group.add_argument('--parallell-io-threads', type=int, default=4,
                       help="Number of cpu threads used to do data prefetch.")
    group.add_argument('--generated-data', action="store_true", default=False,
                       help="Generates synthetic-data on the host and then use it for training.")

    # Add logging-specific arguments
    group_log = parser.add_argument_group('Logging')
    group_log.add_argument('--log-dir', type=str,
                           help="Log and weights save directory")
    group_log.add_argument('--steps-per-logs', type=int, default=1,
                           help="Logs per epoch (if number of epochs specified)")
    return parser


def create_command_line_parser():
    parser = argparse.ArgumentParser(description="""FastSpeech2 on Graphcore's IPU""",
                                     add_help=False,
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    group = parser.add_argument_group("""Main options""")
    group.add_argument('--help', action='store_true',
                       default=False, help="Display help.")
    group.add_argument('--config', type=str,
                       help='FastSpeech2 configuration file in JSON format.')
    return parser


def create_all_options_parser():
    parser = create_command_line_parser()
    parser = add_common_arguments(parser)
    return parser


def make_global_options(task_specific_parsers=[]):
    # Parse command-line arguments
    command_line_parser = create_command_line_parser()
    all_options_parser = create_all_options_parser()

    for task_parser in task_specific_parsers:
        all_options_parser = task_parser(all_options_parser)

    known_command_line_args, unknown_command_line_args = command_line_parser.parse_known_args()

    if known_command_line_args.help or known_command_line_args.config is None:
        all_options_parser.print_help()
        sys.exit(os.EX_OK)

    # Parse options specified in the configuration file into
    config_file_path = known_command_line_args.config
    opts_from_config_file = FastSpeech2Config.from_json_file(config_file_path)

    # Build the global options structure from the default options
    current_options = vars(all_options_parser.parse_args())

    unknown_options = [opt for opt in opts_from_config_file.keys(
    ) if opt not in current_options.keys()]

    if unknown_options:
        logging.error(f"Unonwn options: {unknown_options}")
        sys.exit(os.EX_USAGE)

    # Overwrite global options by those specified in the config file.
    current_options.update(opts_from_config_file)
    options_namespace = argparse.Namespace(**current_options)

    # Overwrite with command-line arguments
    all_options_namespace = all_options_parser.parse_args(
        unknown_command_line_args, options_namespace)
    logging.info(
        f"Overwrite configuration parameters: {', '.join(unknown_command_line_args)}")

    # argparse.Namespace -> dict()
    opts = vars(all_options_namespace)

    return opts
# Copyright (c) 2021 Graphcore Ltd. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # #    http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # Training options Whether to run evaluation loop or not. Default is False.                        If both `--train` and `--eval` are set, then the evaluation will be done after trainning. # FastSpeech2 options Vocabulary size of `FastSpeech2`. # Encoder Size ofthe encoder layers and the pooler layer. Number of hidden layers in the encoder. Number of attention heads for each attention layer in the encoder. The size of the "intermediate" (i.e., feed-forward) layer in the encoder. The convolution kernel size of the "intermediate" (i.e., feed-forward) layer in the encoder. The non-linear activation function (function or string) in the encoder and pooler. # Decoder Size ofthe decoder layers and the pooler layer. Number of hidden layers in the decoder. Number of attention heads for each attention layer in the decoder. The size of the "intermediate" (i.e., feed-forward) layer in the decoder. The convolution kernel size of the "intermediate" (i.e., feed-forward) layer in the decoder. The non-linear activation function (function or string) in the decoder and pooler. # Duration predictor(Only for FastSpeech) Number of convolution layers in the duration predictor. The convolution kernel size of the duration predictor. The dropout probability for duration predictor. Number of filters in the duration predictor. # Variance predictor and Postnet Number of convolution layers in the variant predictor. The convolution kernel size of the variant predictor. The dropout probability for variant predictor. Number of filters in the variant predictor. Number of convolution layers in the postnet. The convolution kernel size of the postnet. The dropout probability for postnet. Number of convolution filters in the postnet. The dropout probability for all fully connected layers in the embeddings, encoder, and pooler. The dropout ratio for the attention probabilities. The maximum sequence length that this model might ever be used with. Typically set this to something large just in case (e.g., 512 or 1024 or 2048). The stdev of the truncated-normal-initializer for initializing all weight matrices. # IPU options # Dataset options # Add logging-specific arguments FastSpeech2 on Graphcore's IPU Main options # Parse command-line arguments # Parse options specified in the configuration file into # Build the global options structure from the default options # Overwrite global options by those specified in the config file. # Overwrite with command-line arguments # argparse.Namespace -> dict()