def main(cfg: DictConfig) -> None: try: plugin = NLPDDPPlugin() except (ImportError, ModuleNotFoundError): plugin = None trainer = pl.Trainer(plugins=plugin, **cfg.trainer) exp_manager(trainer, cfg.get("exp_manager", None)) if not cfg.pretrained_model: logging.info(f'Config: {OmegaConf.to_yaml(cfg)}') model = TokenClassificationModel(cfg.model, trainer=trainer) else: if os.path.exists(cfg.pretrained_model): # TODO: can we drop strict=False? model = TokenClassificationModel.restore_from(cfg.pretrained_model, trainer=trainer, strict=False) elif cfg.pretrained_model in TokenClassificationModel.get_available_model_names( ): model = TokenClassificationModel.from_pretrained( cfg.pretrained_model) else: raise ValueError( f'Provide path to the pre-trained .nemo file or choose from {TokenClassificationModel.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 it will require # setup the data dir to get class weights statistics model.update_data_dir(data_dir=data_dir) # finally, setup train and validation Pytorch DataLoaders model.setup_training_data() model.setup_validation_data() # then we're setting up loss, use model.dataset.class_balancing, # if you want to add class weights to the CrossEntropyLoss model.setup_loss(class_balancing=cfg.model.dataset.class_balancing) 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)
def main(cfg: DictConfig) -> None: logging.info( 'During evaluation/testing, it is currently advisable to construct a new Trainer with single GPU and \ no DDP to obtain accurate results') if not hasattr(cfg.model, 'test_ds'): raise ValueError( f'model.test_ds was not found in the config, skipping evaluation') else: gpu = 1 if cfg.trainer.gpus != 0 else 0 trainer = pl.Trainer( gpus=gpu, precision=cfg.trainer.precision, amp_level=cfg.trainer.amp_level, logger=False, checkpoint_callback=False, ) exp_dir = exp_manager(trainer, cfg.exp_manager) if not cfg.pretrained_model: raise ValueError( 'To run evaluation and inference script a pre-trained model or .nemo file must be provided.' f'Choose from {TokenClassificationModel.list_available_models()} or "pretrained_model"="your_model.nemo"' ) if os.path.exists(cfg.pretrained_model): model = TokenClassificationModel.restore_from(cfg.pretrained_model) elif cfg.pretrained_model in TokenClassificationModel.get_available_model_names( ): model = TokenClassificationModel.from_pretrained(cfg.pretrained_model) else: raise ValueError( f'Provide path to the pre-trained .nemo checkpoint or choose from {TokenClassificationModel.list_available_models()}' ) data_dir = cfg.model.dataset.get('data_dir', None) if data_dir is None: logging.error( 'No dataset directory provided. Skipping evaluation. ' 'To run evaluation on a file, specify path to the directory that contains test_ds.text_file and test_ds.labels_file with "model.dataset.data_dir" argument.' ) elif not os.path.exists(data_dir): logging.error( f'{data_dir} is not found, skipping evaluation on the test set.') else: model.update_data_dir(data_dir=data_dir) model._cfg.dataset = cfg.model.dataset if not hasattr(cfg.model, 'test_ds'): logging.error( f'model.test_ds was not found in the config, skipping evaluation' ) elif model.prepare_test(trainer): model.setup_test_data(cfg.model.test_ds) trainer.test(model) else: logging.error( 'Skipping the evaluation. The trainer is not setup properly.') model.evaluate_from_file( text_file=os.path.join(data_dir, cfg.model.test_ds.text_file), labels_file=os.path.join(data_dir, cfg.model.test_ds.labels_file), output_dir=exp_dir, add_confusion_matrix=True, normalize_confusion_matrix=True, ) # run an inference on a few examples queries = [ 'we bought four shirts from the nvidia gear store in santa clara.', 'Nvidia is a company.' ] results = model.add_predictions(queries, output_file='predictions.txt') for query, result in zip(queries, results): logging.info(f'Query : {query}') logging.info(f'Result: {result.strip()}\n') logging.info(f'Results are saved at {exp_dir}')