{"cells":[{"metadata":{},"cell_type":"markdown","source":"# <center> Text classification with DistillBert and Catalyst\n    \n<img src='https://habrastorage.org/webt/ne/n_/ow/nen_ow49hxu8zrkgolq1rv3xkhi.png'>\n    \nHere we classify Amazon product reviews with the `DistillBert` model from [transformers](https://github.com/huggingface/transformers). We use [catalyst](https://github.com/catalyst-team/catalyst) to run the experiment. We reuse many powerful techniques for faster model training, almost all of them are handled by Catalyst:\n\n1. **Gradient accumulation.** Doing one optimization step for several bachward steps. Well explained in [this post](https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255) by HuggingFace\n1. **Mixed-precision training.** Handled by [Nvidia Apex](https://github.com/NVIDIA/apex) and reused by Catalyst\n1. **Learning rate schedule.** Standard thing when training deep neural networks, Catalysts handles lot of them\n1. **Sequence bucketing (soon).** The main idea is that you can group long sentences with long ones, short ones with short ones and thus do less padding. Three approaches are described in [this Kernel](https://www.kaggle.com/bminixhofer/speed-up-your-rnn-with-sequence-bucketing)"},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install -U catalyst==19.11.1 transformers > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Python \nimport os\nimport warnings\nimport logging\nfrom typing import Mapping, List\nfrom pprint import pprint\n\n# Numpy and Pandas \nimport numpy as np\nimport pandas as pd\n\n# PyTorch \nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\n\n# Transformers \nfrom transformers import AutoConfig, AutoModel, AutoTokenizer\n\n# Catalyst\nfrom catalyst.dl import SupervisedRunner\nfrom catalyst.dl.callbacks import AccuracyCallback, F1ScoreCallback, OptimizerCallback\nfrom catalyst.dl.callbacks import CheckpointCallback, InferCallback\nfrom catalyst.utils import set_global_seed, prepare_cudnn","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Setup**"},{"metadata":{"trusted":true},"cell_type":"code","source":"MODEL_NAME = 'distilbert-base-uncased' # pretrained model from Transformers\nLOG_DIR = \"./logdir\"                   # for training logs and tensorboard visualizations\nNUM_EPOCHS = 3                         # smth around 2-6 epochs is typically fine when finetuning transformers\nBATCH_SIZE = 8                        # depends on your available GPU memory (in combination with max seq length)\nMAX_SEQ_LENGTH = 512                   # depends on your available GPU memory (in combination with batch size)\nNUM_CLASSES = 3                        # solving 3-class classification problem\nLEARN_RATE = 5e-5                      # learning rate is typically ~1e-5 for transformers\nACCUM_STEPS = 4                        # one optimization step for that many backward passes\nSEED = 17                              # random seed for reproducibility","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Additionaly, we install [Nvidia Apex](https://github.com/NVIDIA/apex) to reuse AMP - automatic mixed-precision training.**\n\nThe idea is that we can use float16 format for faster training, only switching tio float32 when necessary. \nHere we'll only need to tell Catalyst to use fp16."},{"metadata":{"trusted":true},"cell_type":"code","source":"FP16_PARAMS = None","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"%%capture\n# if Your machine doesn't support FP16, comment this 3 lines below\n!git clone https://github.com/NVIDIA/apex \n!pip install -v --no-cache-dir --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" ./apex \n!rm -rf ./apex\nFP16_PARAMS = dict(opt_level=\"O1\") ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Dataset**\n\nClickbait news detection - [Kaggle Inclass competition](https://www.kaggle.com/c/clickbait-news-detection). Given:\n\n - title - title of a news\n - text - text of a news\nNeed to predict label - either news, clickbait or other."},{"metadata":{"trusted":true},"cell_type":"code","source":"# to reproduce, download the data and customize this path\nPATH_TO_DATA = '../input/clickbait-news-detection//'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(PATH_TO_DATA + 'train.csv', index_col='id').fillna('')\nvalid_df = pd.read_csv(PATH_TO_DATA + 'valid.csv', index_col='id').fillna('')\ntest_df = pd.read_csv(PATH_TO_DATA + 'test.csv', index_col='id').fillna('')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# target distribution\ntrain_df['label'].value_counts(normalize=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# statistics of text length (in words)\ntrain_df['text'].apply(lambda s: len(s.split())).describe()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We'll combine title and text and create a new field. This will be passed to the transformer."},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df['title_n_text'] = train_df['title'] + '_' + train_df['text']\nvalid_df['title_n_text'] = valid_df['title'] + '_' + valid_df['text']\ntest_df['title_n_text'] = test_df['title'] + '_' + test_df['text']","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Torch Dataset\n\nThis is left for user to be defined. Catalyst will take care of the rest. "},{"metadata":{"trusted":true},"cell_type":"code","source":"class TextClassificationDataset(Dataset):\n    def __init__(self, \n                 texts: List, \n                 labels: List=None, \n                 label_dict: Mapping[str, int]=None,\n                 max_seq_length: int=512,\n                 model_name: str='distilbert-base-uncased'):\n        '''\n        :param texts: a list with texts to classify or to train the classifier on\n        :param labels: a list with classification labels, can be strings as well (optional)\n        :param label_dict: a dictionary mapping class names to class ids, to be passed to the validation data\n        :param max_seq_length: maximal sequence length, texts will be stripped \n        :model_name: transformer model name, we need here to perform appropriate tokenization\n        '''\n        self.texts = texts\n        self.labels = labels\n        self.label_dict = label_dict\n        \n        if self.label_dict is None and labels is not None:\n            # {'class1': 0, 'class2': 1, 'class3': 2, ...}\n            # using this instead of `sklearn.preprocessing.LabelEncoder`\n            # no easily handle unknown target values\n            self.label_dict = dict(zip(sorted(set(labels)), range(len(set(labels)))))\n        \n        self.max_seq_length = max_seq_length\n\n        self.tokenizer = AutoTokenizer.from_pretrained(model_name)\n        # suppresses tokenizer warnings\n        logging.getLogger(\"transformers.tokenization_utils\").setLevel(logging.FATAL)\n        \n        # special tokens for transformers \n        # in the simplest case a [CLS] token is added in the beginning\n        # and [SEP] token is added in the end of a piece of text\n        self.sep_vid = self.tokenizer.vocab[\"[SEP]\"]\n        self.cls_vid = self.tokenizer.vocab[\"[CLS]\"]\n        self.pad_vid = self.tokenizer.vocab[\"[PAD]\"]\n\n    def __len__(self):\n        return len(self.texts)\n\n    def __getitem__(self, index) -> Mapping[str, torch.Tensor]:\n\n        x = self.texts[index]\n\n        # encoding the text\n        x_encoded = self.tokenizer.encode(\n            x,\n            add_special_tokens=True,\n            max_length=self.max_seq_length,\n            return_tensors=\"pt\",\n        ).squeeze(0)\n        \n        # padding short texts\n        true_seq_length = x_encoded.size(0)\n        pad_size = self.max_seq_length - true_seq_length\n        pad_ids = torch.Tensor([self.pad_vid] * pad_size).long()\n        x_tensor = torch.cat((x_encoded, pad_ids))\n        \n        # encoding target\n        if self.labels is not None:\n            y = self.labels[index]\n            y_encoded = torch.Tensor([self.label_dict.get(y, -1)]).long().squeeze(0)\n        \n        # dealing with masks\n        mask = torch.ones_like(x_encoded, dtype=torch.int8)\n        mask_pad = torch.zeros_like(pad_ids, dtype=torch.int8)\n        mask = torch.cat((mask, mask_pad))\n        \n        output_dict = {\n            \"features\": x_tensor,\n            'attention_mask': mask\n        }\n        \n        if self.labels is not None:\n            output_dict[\"targets\"] = y_encoded\n\n        return output_dict","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Create Torch Datasets with train, validation, and test data.**"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = TextClassificationDataset(\n    texts=train_df['text'],\n    labels=train_df['label'],\n    label_dict=None,\n    max_seq_length=MAX_SEQ_LENGTH)\n\nvalid_dataset = TextClassificationDataset(\n    texts=valid_df['text'],\n    labels=valid_df['label'],\n    label_dict=train_dataset.label_dict,\n    max_seq_length=MAX_SEQ_LENGTH)\n\ntest_dataset = TextClassificationDataset(\n    texts=test_df['text'],\n    labels=None,\n    label_dict=None,\n    max_seq_length=MAX_SEQ_LENGTH)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"One of the training dataset instances:"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df.loc[1]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pprint(train_dataset[1])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Finally, we define standard PyTorch loaders. This dictionary will be fed to Catalyst.**"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_val_loaders = {\n    \"train\": DataLoader(dataset=train_dataset,\n                        batch_size=BATCH_SIZE, \n                        shuffle=True),\n    \"valid\": DataLoader(dataset=valid_dataset,\n                        batch_size=BATCH_SIZE, \n                        shuffle=False)    \n}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# The model\n\nIt's going to be a slightly simplified version of [`DistilBertForSequenceClassification`](https://github.com/huggingface/transformers/blob/master/transformers/modeling_distilbert.py#L547) by HuggingFace.\nWe need only predicted probabilities as output, nothing more - we don't need neither loss to be output nor hidden states or attentions (as in the original implementation)."},{"metadata":{"trusted":true},"cell_type":"code","source":"class DistilBertForSequenceClassification(nn.Module):\n    def __init__(self, model_name, num_classes=None):\n        super().__init__()\n        \n        config = AutoConfig.from_pretrained(\n            model_name, num_labels=num_classes)\n        \n        self.distilbert = AutoModel.from_pretrained(model_name, \n                                                    config=config)\n        self.pre_classifier = nn.Linear(config.dim, config.dim)\n        self.classifier = nn.Linear(config.dim, num_classes)\n        self.dropout = nn.Dropout(config.seq_classif_dropout)\n\n    def forward(self, features, attention_mask=None, head_mask=None):\n        assert attention_mask is not None, \"attention mask is none\"\n        distilbert_output = self.distilbert(input_ids=features,\n                                            attention_mask=attention_mask,\n                                            head_mask=head_mask)\n        hidden_state = distilbert_output[0]                  # (bs, seq_len, dim)\n        pooled_output = hidden_state[:, 0]                   # (bs, dim)\n        pooled_output = self.pre_classifier(pooled_output)   # (bs, dim)\n        pooled_output = nn.ReLU()(pooled_output)             # (bs, dim)\n        pooled_output = self.dropout(pooled_output)          # (bs, dim)\n        logits = self.classifier(pooled_output)              # (bs, dim)\n\n        return logits","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = DistilBertForSequenceClassification(model_name=MODEL_NAME, num_classes=NUM_CLASSES)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model training\n\nFirst we specify criterion, optimizer and scheduler (pure PyTorch). Then Catalyst stuff."},{"metadata":{"trusted":true},"cell_type":"code","source":"criterion = torch.nn.CrossEntropyLoss()\noptimizer = torch.optim.Adam(model.parameters(), lr=LEARN_RATE)\nscheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"To run Deep Learning experiments, Catalyst resorts to the [`Runner`](https://catalyst-team.github.io/catalyst/api/dl.html#catalyst.dl.core.runner.Runner) abstraction, in particular, to [`SupervisedRunner`](https://catalyst-team.github.io/catalyst/api/dl.html#module-catalyst.dl.runner.supervised).\n\n`SupervisedRunner` implements the following methods:\n - `train` - starts the training process of the model\n - `predict_loader` - makes a prediction on the whole loader with the specified model\n - `infer` - makes the inference on the model\n \nTo train the model within this interface you pass the following to the `train` method:\n - model (`torch.nn.Module`) – PyTorch model to train\n - criterion (`nn.Module`) – PyTorch criterion function for training\n - optimizer (`optim.Optimizer`) – PyTorch optimizer for training\n - loaders (dict) – dictionary containing one or several `torch.utils.data.DataLoader` for training and validation\n - logdir (str) – path to output directory. There Catalyst will write logs, will dump the best model and the actual code to train the model\n - callbacks – list of Catalyst callbacks\n - scheduler (`optim.lr_scheduler._LRScheduler`) – PyTorch scheduler for training\n - ...\n \nIn our case we'll pass the created `DistilBertForSequenceClassification` model, cross-entropy criterion, Adam optimizer, scheduler and data loaders that we created earlier. Also, we'll be tracking accuracy and thus will need `AccuracyCallback`. To perform batch accumulation, we'll be using `OptimizationCallback`.\n\nThere are many more useful [callbacks](https://catalyst-team.github.io/catalyst/api/dl.html#module-catalyst.dl.callbacks.checkpoint) implemented, also check out [Catalyst examples](https://github.com/catalyst-team/catalyst/tree/master/examples/notebooks)."},{"metadata":{"trusted":true},"cell_type":"code","source":"os.environ['CUDA_VISIBLE_DEVICES'] = \"0\"    # can be changed in case of multiple GPUs onboard\nset_global_seed(SEED)                       # reproducibility\nprepare_cudnn(deterministic=True)           # reproducibility","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# we need a small wrapper around Catalyst's runner to be able to pass masks to it\nclass BertSupervisedRunner(SupervisedRunner):\n    def __init__(self, *args, **kwargs):\n        super().__init__(*args, input_key=(\n            \"features\",\n            \"attention_mask\",\n        ), **kwargs)","execution_count":null,"outputs":[]},{"metadata":{"scrolled":true,"trusted":true},"cell_type":"code","source":"# model runner\nrunner = BertSupervisedRunner()\n\n# model training\nrunner.train(\n    model=model,\n    criterion=criterion,\n    optimizer=optimizer,\n    scheduler=scheduler,\n    loaders=train_val_loaders,\n    callbacks=[\n        AccuracyCallback(num_classes=NUM_CLASSES),\n#         F1ScoreCallback(activation='Softmax'), # throws a tensor shape mismatch error\n        OptimizerCallback(accumulation_steps=ACCUM_STEPS)\n    ],\n    fp16=FP16_PARAMS,\n    logdir=LOG_DIR,\n    num_epochs=NUM_EPOCHS,\n    verbose=True\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"!nvidia-smi","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"torch.cuda.empty_cache()","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"!nvidia-smi","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Plot metrics\n\n<img src=\"https://habrastorage.org/webt/ki/ib/hy/kiibhyp373r65zriwruroiqitky.jpeg\" width=30% />\n\nThere are at least 4 ways to monitor training:\n\n### 1. Good old tqdm\nThere above it's set with a flag `verbose` in `runner.train`. Actually, it's not that bad :)\n\n<img src='https://habrastorage.org/webt/ta/1s/98/ta1s988ghabz412weaq0lgs_cke.png'> \n\n\n### 2. Weights & Biases\n\nBefore launching training, you can run [Weighs & Biases](https://app.wandb.ai/) inititialization for this project. Execute `wandb init` in a separate terminal window (from the same directory where this notebook is running). `wandb` will ask your API key from https://app.wandb.ai/authorize and project name. The rest will be picked up by Catalyst's `SupervisedWandbRunner` (so you'll need to import this instead of `SupervisedRunner`). \nFollowing the links printed above (smth. like  https://app.wandb.ai/yorko/catalyst-nlp-bert) we can keep track of loss and metrics.\n\n### 3. Tensorboard\nDuring training, logs are written to `LOG_DIR` specified above. \nSimiltaneously with training, you can run `tensorboard --logdir $LOG_DIR` (in another terminal tab, in case of training on a server, I also had to add a `--bin_all` flag),\nand you'll get a nice dashboard. Here we see how accuracy and loss change during training.\n\n<img src=\"https://habrastorage.org/webt/2a/sx/mo/2asxmoizgcpf2fnhjjkfhvf70aw.png\" width=50% />\n\n### 4. Offline metric plotting\n\nIf your training is pretty fast and/or you're not interested in tracking training progress, you can just plot losses and metrics once the training is done. Looks like it won't work in Kernels though but try it locally."},{"metadata":{"trusted":false},"cell_type":"code","source":"from catalyst.dl.utils import plot_metrics","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"plot_metrics(\n    logdir=LOG_DIR,\n    step='batch',\n    metrics=['loss', 'accuracy01']\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Inference for the test set"},{"metadata":{},"cell_type":"markdown","source":"Let's create a Torch loader for the test set and launch `infer` to actually make predictions fot the test set. First, we load the best model checkpoint, then make inference with this model."},{"metadata":{"trusted":false},"cell_type":"code","source":"test_loaders = {\n    \"test\": DataLoader(dataset=test_dataset,\n                        batch_size=BATCH_SIZE, \n                        shuffle=False) \n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"runner.infer(\n    model=model,\n    loaders=test_loaders,\n    callbacks=[\n        CheckpointCallback(\n            resume=f\"{LOG_DIR}/checkpoints/best.pth\"\n        ),\n        InferCallback(),\n    ],   \n    verbose=True\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"predicted_probs = runner.callbacks[0].predictions['logits']","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Now that we have predicted probabilities, let's finally create a submission file."},{"metadata":{"trusted":false},"cell_type":"code","source":"sample_sub_df = pd.read_csv(PATH_TO_DATA + 'sample_submission.csv',\n                           index_col='id')","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"train_dataset.label_dict","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"sample_sub_df['label'] = predicted_probs.argmax(axis=1)\nsample_sub_df['label'] = sample_sub_df['label'].map({v:k for k, v in train_dataset.label_dict.items()})","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"sample_sub_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"sample_sub_df.to_csv('distillbert_submission.csv')","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.6.6"}},"nbformat":4,"nbformat_minor":1}