Peterz98/DiffspeechResearch

1import shutil2import os3import time4from montreal_forced_aligner import __version__5from montreal_forced_aligner.corpus.align_corpus import AlignableCorpus6from montreal_forced_aligner.dictionary import Dictionary, MultispeakerDictionary7from montreal_forced_aligner.aligner import TrainableAligner, PretrainedAligner8from montreal_forced_aligner.models import AcousticModel9from montreal_forced_aligner.config import TEMP_DIR, align_yaml_to_config, load_basic_align, load_command_configuration, \10    train_yaml_to_config11from montreal_forced_aligner.utils import get_available_acoustic_languages, get_pretrained_acoustic_path, \12    get_available_dict_languages, validate_dictionary_arg13from montreal_forced_aligner.helper import setup_logger, log_config14from montreal_forced_aligner.exceptions import ArgumentError15 16 17def load_adapt_config():18    training_config, align_config = train_yaml_to_config('mfa_usr/adapt_config.yaml', require_mono=False)19    training_config.training_configs[0].fmllr_iterations = list(20        range(0, training_config.training_configs[0].num_iterations))21    training_config.training_configs[0].realignment_iterations = list(range(0, training_config.training_configs[22        0].num_iterations))23    return training_config, align_config24 25 26class AcousticModel2(AcousticModel):27    def adaptation_config(self):28        train, align = load_adapt_config()29        return train30 31 32def adapt_model(args, unknown_args=None):33    command = 'align'34    all_begin = time.time()35    if not args.temp_directory:36        temp_dir = TEMP_DIR37    else:38        temp_dir = os.path.expanduser(args.temp_directory)39    corpus_name = os.path.basename(args.corpus_directory)40    if corpus_name == '':41        args.corpus_directory = os.path.dirname(args.corpus_directory)42        corpus_name = os.path.basename(args.corpus_directory)43    data_directory = os.path.join(temp_dir, corpus_name)44    if args.config_path:45        align_config = align_yaml_to_config(args.config_path)46    else:47        align_config = load_basic_align()48    align_config.use_mp = not args.disable_mp49    align_config.debug = args.debug50    align_config.overwrite = args.overwrite51    align_config.cleanup_textgrids = not args.disable_textgrid_cleanup52 53    if unknown_args:54        align_config.update_from_args(unknown_args)55    conf_path = os.path.join(data_directory, 'config.yml')56    if getattr(args, 'clean', False) and os.path.exists(data_directory):57        print('Cleaning old directory!')58        shutil.rmtree(data_directory, ignore_errors=True)59    if getattr(args, 'verbose', False):60        log_level = 'debug'61    else:62        log_level = 'info'63    logger = setup_logger(command, data_directory, console_level=log_level)64    logger.debug('ALIGN CONFIG:')65    log_config(logger, align_config)66    conf = load_command_configuration(conf_path, {'dirty': False,67                                                  'begin': all_begin,68                                                  'version': __version__,69                                                  'type': command,70                                                  'corpus_directory': args.corpus_directory,71                                                  'dictionary_path': args.dictionary_path,72                                                  'acoustic_model_path': args.acoustic_model_path})73    if conf['dirty'] or conf['type'] != command \74            or conf['corpus_directory'] != args.corpus_directory \75            or conf['version'] != __version__ \76            or conf['dictionary_path'] != args.dictionary_path:77        logger.warning(78            'WARNING: Using old temp directory, this might not be ideal for you, use the --clean flag to ensure no '79            'weird behavior for previous versions of the temporary directory.')80        if conf['dirty']:81            logger.debug('Previous run ended in an error (maybe ctrl-c?)')82        if conf['type'] != command:83            logger.debug('Previous run was a different subcommand than {} (was {})'.format(command, conf['type']))84        if conf['corpus_directory'] != args.corpus_directory:85            logger.debug('Previous run used source directory '86                         'path {} (new run: {})'.format(conf['corpus_directory'], args.corpus_directory))87        if conf['version'] != __version__:88            logger.debug('Previous run was on {} version (new run: {})'.format(conf['version'], __version__))89        if conf['dictionary_path'] != args.dictionary_path:90            logger.debug('Previous run used dictionary path {} '91                         '(new run: {})'.format(conf['dictionary_path'], args.dictionary_path))92        if conf['acoustic_model_path'] != args.acoustic_model_path:93            logger.debug('Previous run used acoustic model path {} '94                         '(new run: {})'.format(conf['acoustic_model_path'], args.acoustic_model_path))95 96    os.makedirs(data_directory, exist_ok=True)97    model_directory = os.path.join(data_directory, 'acoustic_models')98    os.makedirs(model_directory, exist_ok=True)99    acoustic_model = AcousticModel2(args.acoustic_model_path, root_directory=model_directory)100    print("| acoustic_model.meta", acoustic_model.meta)101    acoustic_model.log_details(logger)102    training_config = acoustic_model.adaptation_config()103    training_config.training_configs[0].update({'beam': align_config.beam, 'retry_beam': align_config.retry_beam})104    training_config.update_from_align(align_config)105    logger.debug('ADAPT TRAINING CONFIG:')106    log_config(logger, training_config)107    audio_dir = None108    if args.audio_directory:109        audio_dir = args.audio_directory110    try:111        corpus = AlignableCorpus(args.corpus_directory, data_directory,112                                 speaker_characters=args.speaker_characters,113                                 num_jobs=args.num_jobs, sample_rate=align_config.feature_config.sample_frequency,114                                 logger=logger, use_mp=align_config.use_mp, punctuation=align_config.punctuation,115                                 clitic_markers=align_config.clitic_markers, audio_directory=audio_dir)116        if corpus.issues_check:117            logger.warning('Some issues parsing the corpus were detected. '118                           'Please run the validator to get more information.')119        logger.info(corpus.speaker_utterance_info())120        if args.dictionary_path.lower().endswith('.yaml'):121            dictionary = MultispeakerDictionary(args.dictionary_path, data_directory, logger=logger,122                                                punctuation=align_config.punctuation,123                                                clitic_markers=align_config.clitic_markers,124                                                compound_markers=align_config.compound_markers,125                                                multilingual_ipa=acoustic_model.meta['multilingual_ipa'],126                                                strip_diacritics=acoustic_model.meta.get('strip_diacritics', None),127                                                digraphs=acoustic_model.meta.get('digraphs', None))128        else:129            dictionary = Dictionary(args.dictionary_path, data_directory, logger=logger,130                                    punctuation=align_config.punctuation,131                                    clitic_markers=align_config.clitic_markers,132                                    compound_markers=align_config.compound_markers,133                                    multilingual_ipa=acoustic_model.meta['multilingual_ipa'],134                                    strip_diacritics=acoustic_model.meta.get('strip_diacritics', None),135                                    digraphs=acoustic_model.meta.get('digraphs', None))136        acoustic_model.validate(dictionary)137 138        begin = time.time()139        previous = PretrainedAligner(corpus, dictionary, acoustic_model, align_config,140                                     temp_directory=data_directory,141                                     debug=getattr(args, 'debug', False), logger=logger)142        a = TrainableAligner(corpus, dictionary, training_config, align_config,143                             temp_directory=data_directory,144                             debug=getattr(args, 'debug', False), logger=logger, pretrained_aligner=previous)145        logger.debug('Setup adapter in {} seconds'.format(time.time() - begin))146        a.verbose = args.verbose147 148        begin = time.time()149        a.train()150        logger.debug('Performed adaptation in {} seconds'.format(time.time() - begin))151 152        begin = time.time()153        a.save(args.output_model_path, root_directory=model_directory)154        a.export_textgrids(args.output_directory)155        logger.debug('Exported TextGrids in {} seconds'.format(time.time() - begin))156        logger.info('All done!')157 158    except Exception as _:159        conf['dirty'] = True160        raise161    finally:162        handlers = logger.handlers[:]163        for handler in handlers:164            handler.close()165            logger.removeHandler(handler)166        conf.save(conf_path)167 168 169def validate_args(args, downloaded_acoustic_models, download_dictionaries):170    if not os.path.exists(args.corpus_directory):171        raise ArgumentError('Could not find the corpus directory {}.'.format(args.corpus_directory))172    if not os.path.isdir(args.corpus_directory):173        raise ArgumentError('The specified corpus directory ({}) is not a directory.'.format(args.corpus_directory))174 175    args.dictionary_path = validate_dictionary_arg(args.dictionary_path, download_dictionaries)176 177    if args.acoustic_model_path.lower() in downloaded_acoustic_models:178        args.acoustic_model_path = get_pretrained_acoustic_path(args.acoustic_model_path.lower())179    elif args.acoustic_model_path.lower().endswith(AcousticModel.extension):180        if not os.path.exists(args.acoustic_model_path):181            raise ArgumentError('The specified model path does not exist: ' + args.acoustic_model_path)182    else:183        raise ArgumentError(184            'The language \'{}\' is not currently included in the distribution, '185            'please align via training or specify one of the following language names: {}.'.format(186                args.acoustic_model_path.lower(), ', '.join(downloaded_acoustic_models)))187 188 189def run_adapt_model(args, unknown_args=None, downloaded_acoustic_models=None, download_dictionaries=None):190    if downloaded_acoustic_models is None:191        downloaded_acoustic_models = get_available_acoustic_languages()192    if download_dictionaries is None:193        download_dictionaries = get_available_dict_languages()194    try:195        args.speaker_characters = int(args.speaker_characters)196    except ValueError:197        pass198    args.corpus_directory = args.corpus_directory.rstrip('/').rstrip('\\')199 200    validate_args(args, downloaded_acoustic_models, download_dictionaries)201    adapt_model(args, unknown_args)202