def main(cfg: DictConfig) -> None:
torch.manual_seed(42)
cfg = OmegaConf.merge(
OmegaConf.structured(PunctuationCapitalizationConfig()), cfg)
trainer = pl.Trainer(**cfg.trainer)
exp_manager(trainer, cfg.get("exp_manager", None))
if not cfg.do_training and not cfg.do_testing:
raise ValueError(
"At least one of config parameters `do_training` and `do_testing` has to `true`."
)
if cfg.do_training:
if cfg.model.get('train_ds') is None:
raise ValueError(
'`model.train_ds` config section is required if `do_training` config item is `True`.'
)
if cfg.do_testing:
if cfg.model.get('test_ds') is None:
raise ValueError(
'`model.test_ds` config section is required if `do_testing` config item is `True`.'
)
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 '
f'{PunctuationCapitalizationModel.list_available_models()}')
model.update_config_after_restoring_from_checkpoint(
class_labels=cfg.model.class_labels,
common_dataset_parameters=cfg.model.common_dataset_parameters,
train_ds=cfg.model.get('train_ds') if cfg.do_training else None,
validation_ds=cfg.model.get('validation_ds')
if cfg.do_training else None,
test_ds=cfg.model.get('test_ds') if cfg.do_testing else None,
optim=cfg.model.get('optim') if cfg.do_training else None,
)
model.set_trainer(trainer)
if cfg.do_training:
model.setup_training_data()
model.setup_validation_data()
model.setup_optimization()
else:
model.setup_test_data()
if cfg.do_training:
trainer.fit(model)
if cfg.do_testing:
trainer.test(model)
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')