def main(cfg: DictConfig) -> None: trainer = pl.Trainer(**cfg.trainer) exp_manager(trainer, cfg.get("exp_manager", None)) do_training = True if not cfg.pretrained_model: logging.info(f'Config: {OmegaConf.to_yaml(cfg)}') model = PunctuationCapitalizationModel(cfg.model, trainer=trainer) else: logging.info(f'Loading pretrained model {cfg.pretrained_model}') # TODO: Remove strict, when lightning has persistent parameter support for add_state() model = PunctuationCapitalizationModel.from_pretrained( cfg.pretrained_model, strict=False) data_dir = cfg.model.dataset.get('data_dir', None) if data_dir: # we can also do finetunining of the pretrained model but it will require # setting up train and validation Pytorch DataLoaders model.setup_training_data(data_dir=data_dir) # evaluation could be done on multiple files, use model.validation_ds.ds_items to specify multiple # data directories if needed model.setup_validation_data(data_dirs=data_dir) logging.info(f'Using config file of the pretrained model') else: do_training = False logging.info( f'Data dir should be specified for training/finetuning. ' f'Using pretrained {cfg.pretrained_model} model weights and skipping finetuning.' ) if do_training: trainer.fit(model) if cfg.model.nemo_path: model.save_to(cfg.model.nemo_path) logging.info( 'During evaluation/testing, it is currently advisable to construct a new Trainer with single GPU ' 'and no DDP to obtain accurate results') gpu = 1 if cfg.trainer.gpus != 0 else 0 trainer = pl.Trainer(gpus=gpu) model.set_trainer(trainer) # run an inference on a few examples queries = [ 'we bought four shirts one pen and a mug from the nvidia gear store in santa clara', 'what can i do for you today', 'how are you', ] inference_results = model.add_punctuation_capitalization(queries) for query, result in zip(queries, inference_results): logging.info(f'Query : {query}') logging.info(f'Result: {result.strip()}\n')
def main(cfg: DictConfig) -> None: trainer = pl.Trainer(**cfg.trainer) exp_manager(trainer, cfg.get("exp_manager", None)) do_training = True if not cfg.pretrained_model: logging.info(f'Config: {OmegaConf.to_yaml(cfg)}') model = PunctuationCapitalizationModel(cfg.model, trainer=trainer) else: logging.info(f'Loading pretrained model {cfg.pretrained_model}') model = PunctuationCapitalizationModel.from_pretrained(cfg.pretrained_model) data_dir = cfg.model.dataset.get('data_dir', None) if data_dir: model.update_data_dir(data_dir) model.setup_training_data() model.setup_validation_data() logging.info(f'Using config file of the pretrained model') else: do_training = False logging.info( f'Data dir should be specified for training/finetuning. ' f'Using pretrained {cfg.pretrained_model} model weights and skipping finetuning.' ) if do_training: trainer.fit(model) if cfg.model.nemo_path: model.save_to(cfg.model.nemo_path) logging.info( 'During evaluation/testing, it is currently advisable to construct a new Trainer with single GPU ' 'and no DDP to obtain accurate results' ) gpu = 1 if cfg.trainer.gpus != 0 else 0 trainer = pl.Trainer(gpus=gpu) model.set_trainer(trainer) # run an inference on a few examples queries = [ 'we bought four shirts one pen and a mug from the nvidia gear store in santa clara', 'what can i do for you today', 'how are you', ] inference_results = model.add_punctuation_capitalization(queries) for query, result in zip(queries, inference_results): logging.info(f'Query : {query}') logging.info(f'Result: {result.strip()}\n')
def main(cfg: DictConfig) -> None: trainer = pl.Trainer(**cfg.trainer) exp_manager(trainer, cfg.get("exp_manager", None)) if not cfg.pretrained_model: logging.info(f'Config: {OmegaConf.to_yaml(cfg)}') model = PunctuationCapitalizationModel(cfg.model, trainer=trainer) else: if os.path.exists(cfg.pretrained_model): model = PunctuationCapitalizationModel.restore_from( cfg.pretrained_model) elif cfg.pretrained_model in PunctuationCapitalizationModel.get_available_model_names( ): model = PunctuationCapitalizationModel.from_pretrained( cfg.pretrained_model) else: raise ValueError( f'Provide path to the pre-trained .nemo file or choose from {PunctuationCapitalizationModel.list_available_models()}' ) data_dir = cfg.model.dataset.get('data_dir', None) if data_dir: if not os.path.exists(data_dir): raise ValueError(f'{data_dir} is not found at') # we can also do finetuning of the pretrained model but we would need to update the data dir model.update_data_dir(data_dir) # setup train and validation Pytorch DataLoaders model.setup_training_data() model.setup_validation_data() logging.info(f'Using config file of the pretrained model') else: raise ValueError( 'Specify a valid dataset directory that contains test_ds.text_file and test_ds.labels_file \ with "model.dataset.data_dir" argument') trainer.fit(model) if cfg.model.nemo_path: model.save_to(cfg.model.nemo_path) logging.info(f'The model was saved to {cfg.model.nemo_path}')