Introduce code from Rust monorepo. This is a mess

2026-10-28 09:08:54 +00:00 · 2022-02-28 05:43:54 -05:00
parent 921fb49f9a
commit ba326bfcfd
1058 changed files with 192192 additions and 0 deletions
@@ -0,0 +1,16 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.collections.nlp.models.language_modeling.bert_lm_model import BERTLMModel
+from nemo.collections.nlp.models.language_modeling.transformer_lm_model import TransformerLMModel
@@ -0,0 +1,272 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from typing import Dict, Optional
+
+import torch
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+from torch.utils.data import DataLoader
+
+from nemo.collections.common.losses import AggregatorLoss, CrossEntropyLoss, SmoothedCrossEntropyLoss
+from nemo.collections.common.metrics import Perplexity
+from nemo.collections.nlp.data.language_modeling.lm_bert_dataset import (
+    BertPretrainingDataset,
+    BertPretrainingPreprocessedDataloader,
+)
+from nemo.collections.nlp.modules.common import BertPretrainingTokenClassifier, SequenceClassifier
+from nemo.collections.nlp.modules.common.lm_utils import get_lm_model
+from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer
+from nemo.core.classes.common import PretrainedModelInfo, typecheck
+from nemo.core.classes.modelPT import ModelPT
+from nemo.core.neural_types import NeuralType
+from nemo.utils import logging
+
+__all__ = ["BERTLMModel"]
+
+
+class BERTLMModel(ModelPT):
+    """
+    BERT language model pretraining.
+    """
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        return self.bert_model.input_types
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        output_types_dict = {"mlm_log_probs": self.mlm_classifier.output_types["log_probs"]}
+        if not self.only_mlm_loss:
+            output_types_dict["nsp_logits"] = self.nsp_classifier.output_types["logits"]
+        return output_types_dict
+
+    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=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)
+            if cfg.tokenizer
+            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)
+
+    @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,
+        )
+        mlm_log_probs = self.mlm_classifier(hidden_states=hidden_states)
+        if self.only_mlm_loss:
+            return (mlm_log_probs,)
+        nsp_logits = self.nsp_classifier(hidden_states=hidden_states)
+        return mlm_log_probs, nsp_logits
+
+    def _compute_losses(self, mlm_log_probs, nsp_logits, output_ids, output_mask, labels):
+        mlm_loss = self.mlm_loss(log_probs=mlm_log_probs, labels=output_ids, output_mask=output_mask)
+        if self.only_mlm_loss:
+            loss, nsp_loss = mlm_loss, None
+        else:
+            nsp_loss = self.nsp_loss(logits=nsp_logits, labels=labels)
+            loss = self.agg_loss(loss_1=mlm_loss, loss_2=nsp_loss)
+        return mlm_loss, nsp_loss, loss
+
+    def _parse_forward_outputs(self, forward_outputs):
+        if self.only_mlm_loss:
+            return forward_outputs[0], None
+        else:
+            return forward_outputs
+
+    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`.
+        """
+        input_ids, input_type_ids, input_mask, output_ids, output_mask, labels = batch
+        forward_outputs = self.forward(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)
+        mlm_log_probs, nsp_logits = self._parse_forward_outputs(forward_outputs)
+        _, _, loss = self._compute_losses(mlm_log_probs, nsp_logits, output_ids, output_mask, labels)
+        lr = self._optimizer.param_groups[0]['lr']
+        self.log('train_loss', loss)
+        self.log('lr', lr, prog_bar=True)
+        return {"loss": 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`.
+        """
+        input_ids, input_type_ids, input_mask, output_ids, output_mask, labels = batch
+        forward_outputs = self.forward(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask)
+        mlm_log_probs, nsp_logits = self._parse_forward_outputs(forward_outputs)
+        _, _, loss = self._compute_losses(mlm_log_probs, nsp_logits, output_ids, output_mask, labels)
+        self.validation_perplexity(logits=mlm_log_probs)
+        return {'val_loss': loss}
+
+    def validation_epoch_end(self, outputs):
+        """Called at the end of validation to aggregate outputs.
+
+        Args:
+            outputs (list): The individual outputs of each validation step.
+
+        Returns:
+            dict: Validation loss and tensorboard logs.
+        """
+        if outputs:
+            avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
+            perplexity = self.validation_perplexity.compute()
+            logging.info(f"evaluation perplexity {perplexity.cpu().item()}")
+            self.log(f'val_loss', avg_loss)
+
+    def setup_training_data(self, train_data_config: Optional[DictConfig]):
+        self._train_dl = (
+            self._setup_preprocessed_dataloader(train_data_config)
+            if self.tokenizer is None
+            else self._setup_dataloader(train_data_config)
+        )
+
+    def setup_validation_data(self, val_data_config: Optional[DictConfig]):
+        self._validation_dl = (
+            self._setup_preprocessed_dataloader(val_data_config)
+            if self.tokenizer is None
+            else self._setup_dataloader(val_data_config)
+        )
+
+    def setup_test_data(self, test_data_config: Optional[DictConfig]):
+        pass
+
+    def _setup_preprocessed_dataloader(self, cfg: Optional[DictConfig]):
+        dataset = cfg.data_file
+        max_predictions_per_seq = cfg.max_predictions_per_seq
+        batch_size = cfg.batch_size
+
+        if os.path.isdir(dataset):
+            files = [os.path.join(dataset, f) for f in os.listdir(dataset) if os.path.isfile(os.path.join(dataset, f))]
+        else:
+            files = [dataset]
+        files.sort()
+        dl = BertPretrainingPreprocessedDataloader(
+            data_files=files, max_predictions_per_seq=max_predictions_per_seq, batch_size=batch_size,
+        )
+        return dl
+
+    def _setup_tokenizer(self, cfg: DictConfig):
+        tokenizer = get_tokenizer(
+            tokenizer_name=cfg.tokenizer_name,
+            tokenizer_model=cfg.tokenizer_model,
+            special_tokens=OmegaConf.to_container(cfg.special_tokens) if cfg.special_tokens else None,
+            vocab_file=cfg.vocab_file,
+        )
+        self.tokenizer = tokenizer
+
+    def _setup_dataloader(self, cfg: DictConfig):
+        dataset = BertPretrainingDataset(
+            tokenizer=self.tokenizer,
+            data_file=cfg.data_file,
+            max_seq_length=cfg.max_seq_length,
+            mask_prob=cfg.mask_prob,
+            short_seq_prob=cfg.short_seq_prob,
+        )
+        dl = torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=cfg.batch_size,
+            collate_fn=dataset.collate_fn,
+            drop_last=cfg.get("drop_last", False),
+            shuffle=cfg.shuffle,
+            num_workers=cfg.get("num_workers", 0),
+        )
+        return dl
+
+    @classmethod
+    def list_available_models(cls) -> Optional[PretrainedModelInfo]:
+        """
+        This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
+
+        Returns:
+            List of available pre-trained models.
+        """
+        result = []
+
+        result.append(
+            PretrainedModelInfo(
+                pretrained_model_name="bertbaseuncased",
+                location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/bertbaseuncased/versions/1.0.0rc1/files/bertbaseuncased.nemo",
+                description="The model was trained EN Wikipedia and BookCorpus on a sequence length of 512.",
+            )
+        )
+
+        result.append(
+            PretrainedModelInfo(
+                pretrained_model_name="bertlargeuncased",
+                location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/bertlargeuncased/versions/1.0.0rc1/files/bertlargeuncased.nemo",
+                description="The model was trained EN Wikipedia and BookCorpus on a sequence length of 512.",
+            )
+        )
+        return result
@@ -0,0 +1,276 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import math
+from typing import Dict, Optional
+
+import numpy as np
+import torch
+import torch.utils.data as pt_data
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning import Trainer
+
+from nemo.collections.common.losses import SmoothedCrossEntropyLoss
+from nemo.collections.common.metrics import GlobalAverageLossMetric
+from nemo.collections.common.parts import transformer_weights_init
+from nemo.collections.nlp.data import SentenceDataset, TarredSentenceDataset
+from nemo.collections.nlp.metrics import SequencePerplexity
+from nemo.collections.nlp.modules.common import TokenClassifier
+from nemo.collections.nlp.modules.common.lm_utils import get_transformer
+from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer
+from nemo.core.classes.common import typecheck
+from nemo.core.classes.modelPT import ModelPT
+from nemo.utils import logging, model_utils
+
+__all__ = ["TransformerLMModel"]
+
+
+class TransformerLMModel(ModelPT):
+    """
+    Left-to-right Transformer language model.
+    """
+
+    def __init__(self, cfg: DictConfig, trainer: Trainer = None):
+
+        # Get global rank and total number of GPU workers for IterableDataset partitioning, if applicable
+        self.world_size = 1
+        if trainer is not None:
+            self.world_size = trainer.num_nodes * trainer.num_gpus
+
+        cfg = model_utils.convert_model_config_to_dict_config(cfg)
+        cfg = model_utils.maybe_update_config_version(cfg)
+
+        # Instantiates tokenizer and register to be saved with NeMo Model archive
+        # After this call, ther will be self.tokenizer which can convert between tokens and token_ids.
+        self.setup_tokenizer(
+            tokenizer_name=cfg.tokenizer.get("tokenizer_name", "yttm"),
+            tokenizer_model=cfg.tokenizer.get("tokenizer_model", None),
+            vocab_file=cfg.tokenizer.get("vocab_file", None),
+            bpe_dropout=cfg.tokenizer.get("bpe_dropout", 0.0),
+        )
+
+        # init superclass
+        super().__init__(cfg=cfg, trainer=trainer)
+
+        # make vocabulary size divisible by 8 for fast fp16 training
+        vocab_size = 8 * math.ceil(self.tokenizer.vocab_size / 8)
+
+        # encoder from NeMo, Megatron-LM, or HuggingFace
+        encoder_cfg_dict = OmegaConf.to_container(cfg.get('encoder'))
+        encoder_cfg_dict['vocab_size'] = vocab_size
+        library = encoder_cfg_dict.pop('library', 'nemo')
+        model_name = encoder_cfg_dict.pop('model_name', None)
+        pretrained = encoder_cfg_dict.pop('pretrained', False)
+        self.encoder = get_transformer(
+            library=library, model_name=model_name, pretrained=pretrained, config_dict=encoder_cfg_dict, encoder=True,
+        )
+
+        self.log_softmax = TokenClassifier(
+            hidden_size=self.encoder.hidden_size,
+            num_classes=vocab_size,
+            activation=cfg.head.activation,
+            log_softmax=cfg.head.log_softmax,
+            dropout=cfg.head.dropout,
+            use_transformer_init=cfg.head.use_transformer_init,
+        )
+
+        # tie weights of embedding and softmax matrices
+        self.log_softmax.mlp.layer0.weight = self.encoder.embedding.token_embedding.weight
+
+        std_init_range = 1 / self.encoder.hidden_size ** 0.5
+
+        # initialize weights if not using pretrained encoder
+        if not self._cfg.encoder.get('pretrained', False):
+            self.encoder.apply(lambda module: transformer_weights_init(module, std_init_range))
+
+        self.log_softmax.apply(lambda module: transformer_weights_init(module, std_init_range))
+
+        self.loss_fn = SmoothedCrossEntropyLoss(pad_id=self.tokenizer.pad_id, label_smoothing=cfg.label_smoothing)
+        self.eval_loss_fn = SmoothedCrossEntropyLoss(pad_id=self.tokenizer.pad_id)
+        self.eval_loss = GlobalAverageLossMetric(dist_sync_on_step=False, take_avg_loss=True)
+        self.eval_ppl = SequencePerplexity()
+
+    @typecheck()
+    def forward(self, input_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.encoder(input_ids=input_ids, encoder_mask=attention_mask)
+        log_probs = self.log_softmax(hidden_states=hidden_states)
+
+        return log_probs
+
+    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
+        for i in range(len(batch)):
+            if batch[i].ndim == 3:
+                # Dataset returns already batched data and the first dimension of size 1
+                # added by DataLoader is excess.
+                batch[i] = batch[i].squeeze(dim=0)
+        ids, mask = batch
+        input_ids, labels = ids[:, :-1], ids[:, 1:]
+        input_mask = mask[:, :-1]
+        log_probs = self(input_ids=input_ids, attention_mask=input_mask)
+
+        train_loss = self.loss_fn(log_probs=log_probs, labels=labels)
+
+        tensorboard_logs = {
+            "train_loss": train_loss,
+            "lr": self._optimizer.param_groups[0]["lr"],
+        }
+        return {"loss": train_loss, "log": tensorboard_logs}
+
+    def eval_step(self, batch, batch_idx):
+        for i in range(len(batch)):
+            if batch[i].ndim == 3:
+                # Dataset returns already batched data and the first dimension of size 1
+                # added by DataLoader is excess.
+                batch[i] = batch[i].squeeze(dim=0)
+        ids, mask = batch
+        input_ids, labels = ids[:, :-1], ids[:, 1:]
+        input_mask, output_mask = mask[:, :-1], mask[:, 1:]
+        log_probs = self(input_ids=input_ids, attention_mask=input_mask)
+        eval_loss = self.eval_loss_fn(log_probs=log_probs, labels=labels)
+
+        self.eval_loss(loss=eval_loss, num_measurements=log_probs.shape[0] * log_probs.shape[1])
+        self.eval_ppl(log_probs=log_probs, labels=labels, mask=output_mask)
+        return {}
+
+    def test_step(self, batch, batch_idx):
+        return self.eval_step(batch, batch_idx)
+
+    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`.
+        """
+        return self.eval_step(batch, batch_idx)
+
+    def eval_epoch_end(self, outputs, mode):
+        eval_loss = self.eval_loss.compute()
+        eval_ppl = self.eval_ppl.compute()
+        ans = {f"{mode}_loss": eval_loss, f"{mode}_ppl": eval_ppl}
+        ans['log'] = dict(ans)
+        dataset_name = "Validation" if mode == 'val' else "Test"
+        logging.info(f"\n\n\n\n{dataset_name} PPL: {np.round(eval_ppl.item(), 2)}")
+        return ans
+
+    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.
+        """
+        self.log_dict(self.eval_epoch_end(outputs, 'val'), sync_dist=True)
+        self.eval_loss.reset()
+        self.eval_ppl.reset()
+
+    def test_epoch_end(self, outputs):
+        return self.eval_epoch_end(outputs, 'test')
+
+    def setup_tokenizer(
+        self, tokenizer_name=None, tokenizer_model=None, vocab_file=None, bpe_dropout=0.0,
+    ):
+
+        supported_tokenizers = ['yttm', 'huggingface', 'sentencepiece', 'word']
+        if tokenizer_name not in supported_tokenizers:
+            raise NotImplementedError(f"Currently we only support tokenizers in {supported_tokenizers}.")
+
+        self.tokenizer = get_tokenizer(
+            tokenizer_name=tokenizer_name,
+            tokenizer_model=self.register_artifact("cfg.tokenizer.tokenizer_model", tokenizer_model),
+            vocab_file=vocab_file,
+            bpe_dropout=bpe_dropout,
+            special_tokens=None,
+            use_fast=False,
+        )
+
+    def setup_training_data(self, train_data_config: Optional[DictConfig]):
+        self._train_dl = self._setup_dataloader_from_config(cfg=train_data_config)
+
+    def setup_validation_data(self, val_data_config: Optional[DictConfig]):
+        self._validation_dl = self._setup_dataloader_from_config(cfg=val_data_config)
+
+    def setup_test_data(self, test_data_config: Optional[DictConfig]):
+        self._test_dl = self._setup_dataloader_from_config(cfg=test_data_config)
+
+    def _setup_dataloader_from_config(self, cfg: DictConfig, predict_last_k=0):
+
+        if cfg.get("use_tarred_dataset", False):
+            if cfg.get("metadata_file") is None:
+                raise FileNotFoundError("Trying to use tarred data set but could not find metadata path in config.")
+            else:
+                metadata_file = cfg.get('metadata_file')
+                with open(metadata_file) as metadata_reader:
+                    metadata = json.load(metadata_reader)
+                if cfg.get('tar_files') is None:
+                    tar_files = metadata.get('tar_files')
+                    if tar_files is not None:
+                        logging.info(f'Loading from tarred dataset {tar_files}')
+                    else:
+                        raise FileNotFoundError("Could not find tarred dataset in config or metadata.")
+                else:
+                    tar_files = cfg.get('tar_files')
+                    if metadata.get('tar_files') is not None:
+                        raise ValueError(
+                            'Tar files specified in config and in metadata file. Tar files should only be specified once.'
+                        )
+            dataset = TarredSentenceDataset(
+                text_tar_filepaths=tar_files,
+                metadata_path=metadata_file,
+                tokenizer=self.tokenizer,
+                shuffle_n=cfg.get("tar_shuffle_n", 100),
+                shard_strategy=cfg.get("shard_strategy", "scatter"),
+                global_rank=self.global_rank,
+                world_size=self.world_size,
+            )
+            return torch.utils.data.DataLoader(
+                dataset=dataset,
+                batch_size=1,
+                num_workers=cfg.get("num_workers", 2),
+                pin_memory=cfg.get("pin_memory", False),
+                drop_last=cfg.get("drop_last", False),
+            )
+        else:
+            dataset = SentenceDataset(
+                tokenizer=self.tokenizer,
+                dataset=cfg.file_name,
+                tokens_in_batch=cfg.tokens_in_batch,
+                clean=cfg.get("clean", False),
+                max_seq_length=cfg.get("max_seq_length", 512),
+                min_seq_length=cfg.get("min_seq_length", 1),
+                cache_ids=cfg.get("cache_ids", False),
+            )
+        if cfg.shuffle:
+            sampler = pt_data.RandomSampler(dataset)
+        else:
+            sampler = pt_data.SequentialSampler(dataset)
+        return torch.utils.data.DataLoader(
+            dataset=dataset,
+            batch_size=1,
+            sampler=sampler,
+            num_workers=cfg.get("num_workers", 2),
+            pin_memory=cfg.get("pin_memory", False),
+            drop_last=cfg.get("drop_last", False),
+        )
+
+    @classmethod
+    def list_available_models(cls) -> Optional[Dict[str, str]]:
+        pass