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 __init__(self, cfg: DictConfig, trainer: Trainer = None): """Initializes the BERTTextClassifier model.""" # shared params for dataset and data loaders self.dataset_cfg = cfg.dataset self.class_weights = None super().__init__(cfg=cfg, trainer=trainer) self.classifier = SequenceClassifier( hidden_size=self.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, ) self.create_loss_module() # 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.class_labels_file', cfg.class_labels.class_labels_file)
def __init__(self, cfg: DictConfig, trainer: Trainer = None): if cfg.tokenizer is not None: self._setup_tokenizer(cfg.tokenizer) else: self.tokenizer = None 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=OmegaConf.to_container(cfg.language_model.config) if cfg.language_model.config else None, checkpoint_file=cfg.language_model.lm_checkpoint, vocab_file=cfg.tokenizer.get('vocab_file') if cfg.tokenizer is not None else None, ) self.hidden_size = self.bert_model.config.hidden_size self.vocab_size = self.bert_model.config.vocab_size self.only_mlm_loss = cfg.only_mlm_loss self.mlm_classifier = BertPretrainingTokenClassifier( hidden_size=self.hidden_size, num_classes=self.vocab_size, num_layers=cfg.num_tok_classification_layers, activation="gelu", log_softmax=True, use_transformer_init=True, ) self.mlm_loss = SmoothedCrossEntropyLoss() if not self.only_mlm_loss: self.nsp_classifier = SequenceClassifier( hidden_size=self.hidden_size, num_classes=2, num_layers=cfg.num_seq_classification_layers, log_softmax=False, activation="tanh", use_transformer_init=True, ) self.nsp_loss = CrossEntropyLoss() self.agg_loss = AggregatorLoss(num_inputs=2) # # tie weights of MLM softmax layer and embedding layer of the encoder if (self.mlm_classifier.mlp.last_linear_layer.weight.shape != self.bert_model.embeddings.word_embeddings.weight.shape): raise ValueError( "Final classification layer does not match embedding layer.") self.mlm_classifier.mlp.last_linear_layer.weight = self.bert_model.embeddings.word_embeddings.weight # create extra bias # setup to track metrics self.validation_perplexity = Perplexity(compute_on_step=False) self.setup_optimization(cfg.optim)
def __init__(self, cfg: DictConfig, trainer: Trainer = None): """ Initializes model to use BERT model for GLUE tasks. """ self.data_dir = cfg.dataset.data_dir if not os.path.exists(self.data_dir): raise FileNotFoundError( "GLUE datasets not found. For more details on how to get the data, see: " "https://gist.github.com/W4ngatang/60c2bdb54d156a41194446737ce03e2e" ) if cfg.task_name not in cfg.supported_tasks: raise ValueError( f'{cfg.task_name} not in supported task. Choose from {cfg.supported_tasks}' ) self.task_name = cfg.task_name # MNLI task has two separate dev sets: matched and mismatched cfg.train_ds.file_name = os.path.join(self.data_dir, cfg.train_ds.file_name) if self.task_name == "mnli": cfg.validation_ds.file_name = [ os.path.join(self.data_dir, 'dev_matched.tsv'), os.path.join(self.data_dir, 'dev_mismatched.tsv'), ] else: cfg.validation_ds.file_name = os.path.join( self.data_dir, cfg.validation_ds.file_name) logging.info( f'Using {cfg.validation_ds.file_name} for model evaluation.') self._setup_tokenizer(cfg.tokenizer) super().__init__(cfg=cfg, trainer=trainer) num_labels = GLUE_TASKS_NUM_LABELS[self.task_name] self.bert_model = get_lm_model( pretrained_model_name=cfg.language_model.pretrained_model_name, config_file=cfg.language_model.config_file, config_dict=OmegaConf.to_container(cfg.language_model.config) if cfg.language_model.config else None, checkpoint_file=cfg.language_model.lm_checkpoint, ) # uses [CLS] token for classification (the first token) if self.task_name == "sts-b": self.pooler = SequenceRegression( hidden_size=self.bert_model.config.hidden_size) self.loss = MSELoss() else: self.pooler = SequenceClassifier( hidden_size=self.bert_model.config.hidden_size, num_classes=num_labels, log_softmax=False) self.loss = CrossEntropyLoss() # Optimizer setup needs to happen after all model weights are ready self.setup_optimization(cfg.optim)
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) self.class_weights = None super().__init__(cfg=cfg, trainer=trainer) self.bert_model = get_lm_model( pretrained_model_name=cfg.language_model.pretrained_model_name, config_file=self.register_artifact('language_model.config_file', cfg.language_model.config_file), config_dict=cfg.language_model.config, checkpoint_file=cfg.language_model.lm_checkpoint, nemo_file=self.register_artifact( 'language_model.nemo_file', cfg.language_model.get('nemo_file', None)), vocab_file=self.register_artifact('tokenizer.vocab_file', cfg.tokenizer.vocab_file), trainer=trainer, ) if cfg.language_model.get('nemo_file', None) is not None: hidden_size = self.bert_model.cfg.hidden_size else: hidden_size = self.bert_model.config.hidden_size self.classifier = SequenceClassifier( hidden_size=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, ) self.create_loss_module() # 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.class_labels_file', cfg.class_labels.class_labels_file)
def __init__(self, cfg: DictConfig, trainer: Trainer = None): """ Initializes model to use BERT model for GLUE tasks. """ if cfg.task_name not in cfg.supported_tasks: raise ValueError( f'{cfg.task_name} not in supported task. Choose from {cfg.supported_tasks}' ) self.task_name = cfg.task_name # needed to setup validation on multiple datasets # MNLI task has two separate dev sets: matched and mismatched if not self._is_model_being_restored(): if self.task_name == "mnli": cfg.validation_ds.ds_item = [ os.path.join(cfg.dataset.data_dir, 'dev_matched.tsv'), os.path.join(cfg.dataset.data_dir, 'dev_mismatched.tsv'), ] else: cfg.validation_ds.ds_item = os.path.join( cfg.dataset.data_dir, cfg.validation_ds.ds_item) cfg.train_ds.ds_item = os.path.join(cfg.dataset.data_dir, cfg.train_ds.ds_item) logging.info( f'Using {cfg.validation_ds.ds_item} for model evaluation.') self.setup_tokenizer(cfg.tokenizer) super().__init__(cfg=cfg, trainer=trainer) num_labels = GLUE_TASKS_NUM_LABELS[self.task_name] self.bert_model = get_lm_model( pretrained_model_name=cfg.language_model.pretrained_model_name, config_file=self.register_artifact('language_model.config_file', cfg.language_model.config_file), config_dict=OmegaConf.to_container(cfg.language_model.config) if cfg.language_model.config else None, checkpoint_file=cfg.language_model.lm_checkpoint, vocab_file=self.register_artifact('tokenizer.vocab_file', cfg.tokenizer.vocab_file), ) # uses [CLS] token for classification (the first token) if self.task_name == "sts-b": self.pooler = SequenceRegression( hidden_size=self.bert_model.config.hidden_size) self.loss = MSELoss() else: self.pooler = SequenceClassifier( hidden_size=self.bert_model.config.hidden_size, num_classes=num_labels, log_softmax=False) self.loss = CrossEntropyLoss()
def test_sequence_classifier_export_to_onnx(self): for num_layers in [1, 2, 4]: classifier_export(SequenceClassifier(hidden_size=256, num_classes=16, num_layers=num_layers))
class TextClassificationModel(NLPModel, Exportable): @property def input_types(self) -> Optional[Dict[str, NeuralType]]: return self.bert_model.input_types @property def output_types(self) -> Optional[Dict[str, NeuralType]]: return self.classifier.output_types 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) # init superclass 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) def _setup_tokenizer(self, cfg: DictConfig): tokenizer = get_tokenizer( tokenizer_name=cfg.tokenizer_name, vocab_file=self.register_artifact( config_path='tokenizer.vocab_file', src=cfg.vocab_file), special_tokens=OmegaConf.to_container(cfg.special_tokens) if cfg.special_tokens else None, tokenizer_model=self.register_artifact( config_path='tokenizer.tokenizer_model', src=cfg.tokenizer_model), ) self.tokenizer = tokenizer @typecheck() def forward(self, input_ids, token_type_ids, attention_mask): """ No special modification required for Lightning, define it as you normally would in the `nn.Module` in vanilla PyTorch. """ hidden_states = self.bert_model(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask) logits = self.classifier(hidden_states=hidden_states) return logits def training_step(self, batch, batch_idx): """ Lightning calls this inside the training loop with the data from the training dataloader passed in as `batch`. """ # forward pass input_ids, input_type_ids, input_mask, labels = batch logits = self.forward(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask) train_loss = self.loss(logits=logits, labels=labels) lr = self._optimizer.param_groups[0]['lr'] self.log('train_loss', train_loss) self.log('lr', lr, prog_bar=True) return { 'loss': train_loss, 'lr': lr, } def validation_step(self, batch, batch_idx): """ Lightning calls this inside the validation loop with the data from the validation dataloader passed in as `batch`. """ if self.testing: prefix = 'test' else: prefix = 'val' input_ids, input_type_ids, input_mask, labels = batch logits = self.forward(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask) val_loss = self.loss(logits=logits, labels=labels) preds = torch.argmax(logits, axis=-1) tp, fn, fp, _ = self.classification_report(preds, labels) return {'val_loss': val_loss, 'tp': tp, 'fn': fn, 'fp': fp} def validation_epoch_end(self, outputs): """ Called at the end of validation to aggregate outputs. :param outputs: list of individual outputs of each validation step. """ if not outputs: return {} if self.testing: prefix = 'test' else: prefix = 'val' avg_loss = torch.stack([x[f'val_loss'] for x in outputs]).mean() # calculate metrics and classification report precision, recall, f1, report = self.classification_report.compute() logging.info(f'{prefix}_report: {report}') self.log(f'{prefix}_loss', avg_loss, prog_bar=True) self.log(f'{prefix}_precision', precision) self.log(f'{prefix}_f1', f1) self.log(f'{prefix}_recall', recall) def test_step(self, batch, batch_idx): """ Lightning calls this inside the test loop with the data from the test dataloader passed in as `batch`. """ return self.validation_step(batch, batch_idx) def test_epoch_end(self, outputs): """ Called at the end of test to aggregate outputs. :param outputs: list of individual outputs of each test step. """ return self.validation_epoch_end(outputs) 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) def setup_validation_data(self, val_data_config: Optional[DictConfig]): if not val_data_config or not val_data_config.file_path: logging.info( f"Dataloader config or file_path for the validation is missing, so no data loader for test is created!" ) self._test_dl = None return self._validation_dl = self._setup_dataloader_from_config( cfg=val_data_config) def setup_test_data(self, test_data_config: Optional[DictConfig]): if not test_data_config or not test_data_config.file_path: logging.info( f"Dataloader config or file_path for the test is missing, so no data loader for test is created!" ) self._test_dl = None return self._test_dl = self._setup_dataloader_from_config( cfg=test_data_config) def _setup_dataloader_from_config( self, cfg: Dict) -> 'torch.utils.data.DataLoader': input_file = cfg.file_path if not os.path.exists(input_file): raise FileNotFoundError( f'{input_file} not found! The data should be be stored in TAB-separated files \n\ "validation_ds.file_path" and "train_ds.file_path" for train and evaluation respectively. \n\ Each line of the files contains text sequences, where words are separated with spaces. \n\ The label of the example is separated with TAB at the end of each line. \n\ Each line of the files should follow the format: \n\ [WORD][SPACE][WORD][SPACE][WORD][...][TAB][LABEL]') dataset = TextClassificationDataset( tokenizer=self.tokenizer, input_file=input_file, max_seq_length=self.dataset_cfg.max_seq_length, num_samples=cfg.get("num_samples", -1), shuffle=cfg.shuffle, use_cache=self.dataset_cfg.use_cache, ) return torch.utils.data.DataLoader( dataset=dataset, batch_size=cfg.batch_size, shuffle=cfg.shuffle, num_workers=cfg.get("num_workers", 0), pin_memory=cfg.get("pin_memory", False), drop_last=cfg.get("drop_last", False), collate_fn=dataset.collate_fn, ) @torch.no_grad() def classifytext(self, queries: List[str], batch_size: int = 1, max_seq_length: int = -1) -> List[int]: """ Get prediction for the queries Args: queries: text sequences batch_size: batch size to use during inference max_seq_length: sequences longer than max_seq_length will get truncated. default -1 disables truncation. Returns: all_preds: model predictions """ # store predictions for all queries in a single list all_preds = [] mode = self.training device = next(self.parameters()).device try: # Switch model to evaluation mode self.eval() logging_level = logging.get_verbosity() logging.set_verbosity(logging.WARNING) dataloader_cfg = { "batch_size": batch_size, "num_workers": 3, "pin_memory": False } infer_datalayer = self._setup_infer_dataloader( dataloader_cfg, queries, max_seq_length) for i, batch in enumerate(infer_datalayer): input_ids, input_type_ids, input_mask, subtokens_mask = batch logits = self.forward( input_ids=input_ids.to(device), token_type_ids=input_type_ids.to(device), attention_mask=input_mask.to(device), ) preds = tensor2list(torch.argmax(logits, axis=-1)) all_preds.extend(preds) finally: # set mode back to its original value self.train(mode=mode) logging.set_verbosity(logging_level) return all_preds def _setup_infer_dataloader( self, cfg: Dict, queries: List[str], max_seq_length: int = -1) -> 'torch.utils.data.DataLoader': """ Setup function for a infer data loader. Args: cfg: config dictionary containing data loader params like batch_size, num_workers and pin_memory queries: text max_seq_length: maximum length of queries, default is -1 for no limit Returns: A pytorch DataLoader. """ dataset = TextClassificationDataset(tokenizer=self.tokenizer, queries=queries, max_seq_length=max_seq_length) return torch.utils.data.DataLoader( dataset=dataset, batch_size=cfg["batch_size"], shuffle=False, num_workers=cfg.get("num_workers", 0), pin_memory=cfg.get("pin_memory", False), drop_last=False, collate_fn=dataset.collate_fn, ) @classmethod def list_available_models(cls) -> Optional[Dict[str, str]]: pass @classmethod def from_pretrained(cls, name: str): pass def _prepare_for_export(self): return self.bert_model._prepare_for_export() def export( self, output: str, input_example=None, output_example=None, verbose=False, export_params=True, do_constant_folding=True, keep_initializers_as_inputs=False, onnx_opset_version: int = 12, try_script: bool = False, set_eval: bool = True, check_trace: bool = True, use_dynamic_axes: bool = True, ): if input_example is not None or output_example is not None: logging.warning( "Passed input and output examples will be ignored and recomputed since" " TextClassificationModel consists of two separate models with different" " inputs and outputs.") qual_name = self.__module__ + '.' + self.__class__.__qualname__ output1 = os.path.join(os.path.dirname(output), 'bert_' + os.path.basename(output)) output1_descr = qual_name + ' BERT exported to ONNX' bert_model_onnx = self.bert_model.export( output1, None, # computed by input_example() None, verbose, export_params, do_constant_folding, keep_initializers_as_inputs, onnx_opset_version, try_script, set_eval, check_trace, use_dynamic_axes, ) output2 = os.path.join(os.path.dirname(output), 'classifier_' + os.path.basename(output)) output2_descr = qual_name + ' Classifier exported to ONNX' classifier_onnx = self.classifier.export( output2, None, # computed by input_example() None, verbose, export_params, do_constant_folding, keep_initializers_as_inputs, onnx_opset_version, try_script, set_eval, check_trace, use_dynamic_axes, ) output_model = attach_onnx_to_onnx(bert_model_onnx, classifier_onnx, "CL") output_descr = qual_name + ' BERT+Classifier exported to ONNX' onnx.save(output_model, output) return ([output, output1, output2], [output_descr, output1_descr, output2_descr])