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