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}')