def __init__(self, cfg: DictConfig, trainer: Trainer = None):
"""Initializes the BERTTextClassifier model."""
# shared params for dataset and data loaders
self.dataset_cfg = cfg.dataset
# tokenizer needs to get initialized before the super.__init__()
# as dataloaders and datasets need it to process the data
self.setup_tokenizer(cfg.tokenizer)
super().__init__(cfg=cfg, trainer=trainer)
self.bert_model = get_lm_model(
pretrained_model_name=cfg.language_model.pretrained_model_name,
config_file=cfg.language_model.config_file,
config_dict=cfg.language_model.config,
checkpoint_file=cfg.language_model.lm_checkpoint,
)
self.classifier = SequenceClassifier(
hidden_size=self.bert_model.config.hidden_size,
num_classes=cfg.dataset.num_classes,
num_layers=cfg.classifier_head.num_output_layers,
activation='relu',
log_softmax=False,
dropout=cfg.classifier_head.fc_dropout,
use_transformer_init=True,
idx_conditioned_on=0,
)
class_weights = None
if cfg.dataset.class_balancing == 'weighted_loss':
if cfg.train_ds.file_path:
class_weights = calc_class_weights(cfg.train_ds.file_path,
cfg.dataset.num_classes)
else:
logging.info(
'Class_balancing feature is enabled but no train file is given. Calculating the class weights is skipped.'
)
if class_weights:
# You may need to increase the number of epochs for convergence when using weighted_loss
self.loss = CrossEntropyLoss(weight=class_weights)
else:
self.loss = CrossEntropyLoss()
# setup to track metrics
self.classification_report = ClassificationReport(
num_classes=cfg.dataset.num_classes,
mode='micro',
dist_sync_on_step=True)
# register the file containing the labels into the artifacts to get stored in the '.nemo' file later
if 'class_labels' in cfg and 'class_labels_file' in cfg.class_labels and cfg.class_labels.class_labels_file:
self.register_artifact('class_labels',
cfg.class_labels.class_labels_file)
def setup_training_data(self, train_data_config: Optional[DictConfig]):
if not train_data_config or not train_data_config.file_path:
logging.info(
f"Dataloader config or file_path for the train is missing, so no data loader for test is created!"
)
self._test_dl = None
return
self._train_dl = self._setup_dataloader_from_config(cfg=train_data_config)
# calculate the class weights to be used in the loss function
if self.cfg.dataset.class_balancing == 'weighted_loss':
self.class_weights = calc_class_weights(train_data_config.file_path, self.cfg.dataset.num_classes)
else:
self.class_weights = None
# we need to create/update the loss module by using the weights calculated from the training data
self.create_loss_module()
