{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# AI4Code PyTorch Training Starter + AMP + 🤗 + W&B 📉 - Code BERT Large\n\nThis training kernel is a fork of my original training kernel [[TRAIN] AI4Code PyTorch - 🤗 BERT Large + W&B 📉](https://www.kaggle.com/code/heyytanay/train-ai4code-pytorch-bert-large-w-b) (check it out and upvote if you found it helpful!)\n\nThis is training script written in Vanilla PyTorch with a Custom Trainer Class. I hope it can help you in developing more sophisticated models.\n\nIn this notebook, I am using CodeBERT model by Microsoft (large variant of the model to be exact). For more details on it, check out it's [HuggingFace docs](https://huggingface.co/microsoft/codebert-base) and [Github repo](https://github.com/microsoft/CodeBERT).\n\nThink of this notebook has a skeleton for all similar Models (in-fact any PyTorch Hugginface model in reality). You can change chunks of code to suit your needs and it will work efficiently in most cases.\n\nI've borrowed chunks of code from Ahmet Erdem's notebook [here](https://www.kaggle.com/code/aerdem4/ai4code-pytorch-distilbert-baseline). Please check that out, it's very informative to get started!\n\n**Feel free to fork and change the models and do some preprocessing, but if you do please leave an upvote :)**","metadata":{}},{"cell_type":"code","source":"%%sh\npip install -q --upgrade transformers","metadata":{"execution":{"iopub.status.busy":"2022-05-18T14:12:15.55611Z","iopub.execute_input":"2022-05-18T14:12:15.556743Z","iopub.status.idle":"2022-05-18T14:12:24.415398Z","shell.execute_reply.started":"2022-05-18T14:12:15.556708Z","shell.execute_reply":"2022-05-18T14:12:24.411938Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import platform\nimport numpy as np\nimport pandas as pd\nfrom tqdm.notebook import tqdm\nimport matplotlib.pyplot as plt\n\nimport gc\nimport os\nimport wandb\nimport json\nimport glob\nfrom scipy import sparse\nfrom pathlib import Path\n\nimport torch\nimport transformers\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.cuda.amp import GradScaler, autocast\nfrom torch.utils.data import Dataset, DataLoader\n\nfrom sklearn.model_selection import GroupShuffleSplit\nfrom sklearn.metrics import mean_squared_error\nimport warnings\nwarnings.simplefilter('ignore')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-05-18T14:12:24.434925Z","iopub.execute_input":"2022-05-18T14:12:24.436354Z","iopub.status.idle":"2022-05-18T14:12:31.990338Z","shell.execute_reply.started":"2022-05-18T14:12:24.436314Z","shell.execute_reply":"2022-05-18T14:12:31.989527Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We define a Config class to store variables and functions that are to be used globally inside our training script.\nThis makes the code more modular and easy to approach at the same time.","metadata":{}},{"cell_type":"markdown","source":"Keeping the number of epochs to 2 since a single epoch on 15K samples take ~1 hour and 13 minutes. You can change it to however much you need.\n\nThe notebook will not give an OOM error at any point, should you change the epoch size since I have written the code to be heavily optimized.","metadata":{}},{"cell_type":"code","source":"class Config:\n NB_EPOCHS = 2\n LR = 3e-4\n T_0 = 20\n η_min = 1e-4\n MAX_LEN = 120\n TRAIN_BS = 16\n VALID_BS = 16\n MODEL_NAME = 'microsoft/codebert-base'\n data_dir = Path('../input/AI4Code')\n TOKENIZER = transformers.RobertaTokenizer.from_pretrained(MODEL_NAME)\n scaler = GradScaler()","metadata":{"execution":{"iopub.status.busy":"2022-05-18T14:13:21.286039Z","iopub.execute_input":"2022-05-18T14:13:21.286761Z","iopub.status.idle":"2022-05-18T14:13:26.364779Z","shell.execute_reply.started":"2022-05-18T14:13:21.286721Z","shell.execute_reply":"2022-05-18T14:13:26.364001Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## About W&B:\n
\"Weights

\n

WandB is a developer tool for companies turn deep learning research projects into deployed software by helping teams track their models, visualize model performance and easily automate training and improving models.\nWe will use their tools to log hyperparameters and output metrics from your runs, then visualize and compare results and quickly share findings with your colleagues.

","metadata":{}},{"cell_type":"markdown","source":"To login to W&B, you can use below snippet.\n\n```python\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nwb_key = user_secrets.get_secret(\"WANDB_API_KEY\")\n\nwandb.login(key=wb_key)\n```\nMake sure you have your W&B key stored as `WANDB_API_KEY` under Add-ons -> Secrets\n\nYou can view [this](https://www.kaggle.com/ayuraj/experiment-tracking-with-weights-and-biases) notebook to learn more about W&B tracking.\n\nIf you don't want to login to W&B, the kernel will still work and log everything to W&B in anonymous mode.","metadata":{}},{"cell_type":"code","source":"WANDB_CONFIG = {\n 'TRAIN_BS': Config.TRAIN_BS,\n 'VALID_BS': Config.VALID_BS,\n 'N_EPOCHS': Config.NB_EPOCHS,\n 'ARCH': Config.MODEL_NAME,\n 'MAX_LEN': Config.MAX_LEN,\n 'LR': Config.LR,\n 'NUM_WORKERS': 8,\n 'OPTIM': \"AdamW\",\n 'LOSS': \"MSELoss\",\n 'DEVICE': \"cuda\",\n 'T_0': 20,\n 'η_min': 1e-4,\n 'infra': \"Kaggle\",\n 'competition': 'ai4code',\n '_wandb_kernel': 'tanaym'\n}\n\n# Start W&B logging\n# W&B Login\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nwb_key = user_secrets.get_secret(\"WANDB_API_KEY\")\n\nwandb.login(key=wb_key)\n\nrun = wandb.init(\n project='pytorch',\n config=WANDB_CONFIG,\n group='nlp',\n job_type='train',\n)","metadata":{"execution":{"iopub.status.busy":"2022-05-18T14:13:45.253915Z","iopub.execute_input":"2022-05-18T14:13:45.254649Z","iopub.status.idle":"2022-05-18T14:13:51.742364Z","shell.execute_reply.started":"2022-05-18T14:13:45.254601Z","shell.execute_reply":"2022-05-18T14:13:51.741575Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Below are some utility functions that we will be using.","metadata":{}},{"cell_type":"code","source":"from bisect import bisect\n\ndef wandb_log(**kwargs):\n \"\"\"\n Logs a key-value pair to W&B\n \"\"\"\n for k, v in kwargs.items():\n wandb.log({k: v})\n\ndef count_inversions(a):\n inversions = 0\n sorted_so_far = []\n for i, u in enumerate(a):\n j = bisect(sorted_so_far, u)\n inversions += i - j\n sorted_so_far.insert(j, u)\n return inversions\n\ndef kendall_tau(ground_truth, predictions):\n total_inversions = 0\n total_2max = 0 # twice the maximum possible inversions across all instances\n for gt, pred in zip(ground_truth, predictions):\n ranks = [gt.index(x) for x in pred] # rank predicted order in terms of ground truth\n total_inversions += count_inversions(ranks)\n n = len(gt)\n total_2max += n * (n - 1)\n return 1 - 4 * total_inversions / total_2max\n\ndef read_notebook(path):\n return (\n pd.read_json(\n path,\n dtype={'cell_type': 'category', 'source': 'str'})\n .assign(id=path.stem)\n .rename_axis('cell_id')\n )\n\ndef get_ranks(base, derived):\n return [base.index(d) for d in derived]","metadata":{"execution":{"iopub.status.busy":"2022-05-18T14:13:54.105307Z","iopub.execute_input":"2022-05-18T14:13:54.106009Z","iopub.status.idle":"2022-05-18T14:13:54.126888Z","shell.execute_reply.started":"2022-05-18T14:13:54.105965Z","shell.execute_reply":"2022-05-18T14:13:54.125968Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Data preprocessing for further training process taken from the Starter notebook for this competition.","metadata":{}},{"cell_type":"code","source":"NUM_TRAIN = 15000\n\npaths_train = list((Config.data_dir / 'train').glob('*.json'))[:NUM_TRAIN]\nnotebooks_train = [\n read_notebook(path) for path in tqdm(paths_train, desc='Train NBs')\n]\ndf = (\n pd.concat(notebooks_train)\n .set_index('id', append=True)\n .swaplevel()\n .sort_index(level='id', sort_remaining=False)\n)\n\ndf_orders = pd.read_csv(\n Config.data_dir / 'train_orders.csv',\n index_col='id',\n squeeze=True,\n).str.split()\n\ndf_orders_ = df_orders.to_frame().join(\n df.reset_index('cell_id').groupby('id')['cell_id'].apply(list),\n how='right',\n)\n\nranks = {}\nfor id_, cell_order, cell_id in df_orders_.itertuples():\n ranks[id_] = {'cell_id': cell_id, 'rank': get_ranks(cell_order, cell_id)}\n\ndf_ranks = (\n pd.DataFrame\n .from_dict(ranks, orient='index')\n .rename_axis('id')\n .apply(pd.Series.explode)\n .set_index('cell_id', append=True)\n)\n\ndf_ancestors = pd.read_csv(Config.data_dir / 'train_ancestors.csv', index_col='id')\ndf = df.reset_index().merge(df_ranks, on=[\"id\", \"cell_id\"]).merge(df_ancestors, on=[\"id\"])\ndf[\"pct_rank\"] = df[\"rank\"] / df.groupby(\"id\")[\"cell_id\"].transform(\"count\")","metadata":{"execution":{"iopub.status.busy":"2022-05-18T14:21:15.719176Z","iopub.execute_input":"2022-05-18T14:21:15.719852Z","iopub.status.idle":"2022-05-18T14:21:19.728013Z","shell.execute_reply.started":"2022-05-18T14:21:15.719812Z","shell.execute_reply":"2022-05-18T14:21:19.727243Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"NVALID = 0.1 # size of validation set\n\nsplitter = GroupShuffleSplit(n_splits=1, test_size=NVALID, random_state=0)\n\ntrain_ind, val_ind = next(splitter.split(df, groups=df[\"ancestor_id\"]))\n\ntrain_df = df.loc[train_ind].reset_index(drop=True)\nval_df = df.loc[val_ind].reset_index(drop=True)\n\ntrain_df_mark = train_df[train_df[\"cell_type\"] == \"code\"].reset_index(drop=True)\nval_df_mark = val_df[val_df[\"cell_type\"] == \"code\"].reset_index(drop=True)","metadata":{"execution":{"iopub.status.busy":"2022-05-18T14:21:19.731267Z","iopub.execute_input":"2022-05-18T14:21:19.731741Z","iopub.status.idle":"2022-05-18T14:21:19.805261Z","shell.execute_reply.started":"2022-05-18T14:21:19.731701Z","shell.execute_reply":"2022-05-18T14:21:19.804459Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Model class for CodeBert Small model**","metadata":{}},{"cell_type":"code","source":"class CodeBertModel(nn.Module):\n def __init__(self):\n super(CodeBertModel, self).__init__()\n self.bert = transformers.RobertaModel.from_pretrained(Config.MODEL_NAME)\n self.drop = nn.Dropout(0.3)\n self.fc = nn.Linear(768, 1)\n \n def forward(self, ids, mask, token_type_ids):\n _, output = self.bert(ids, attention_mask=mask, token_type_ids=token_type_ids, return_dict=False)\n output = self.drop(output)\n output = self.fc(output)\n return output","metadata":{"execution":{"iopub.status.busy":"2022-05-18T14:21:19.806784Z","iopub.execute_input":"2022-05-18T14:21:19.807061Z","iopub.status.idle":"2022-05-18T14:21:19.815962Z","shell.execute_reply.started":"2022-05-18T14:21:19.807013Z","shell.execute_reply":"2022-05-18T14:21:19.815296Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Custom dataset for the Markdown cells","metadata":{}},{"cell_type":"code","source":"class AI4CodeDataset(Dataset):\n def __init__(self, df, is_test=False):\n self.df = df.reset_index(drop=True)\n self.is_test = is_test\n\n def __getitem__(self, idx):\n sample = self.df.iloc[idx]\n \n inputs = Config.TOKENIZER.encode_plus(\n sample['source'],\n None,\n add_special_tokens=True,\n max_length=Config.MAX_LEN,\n padding=\"max_length\",\n return_token_type_ids=True,\n truncation=True\n )\n ids = torch.tensor(inputs['input_ids'], dtype=torch.long)\n mask = torch.tensor(inputs['attention_mask'], dtype=torch.long)\n token_type_ids = torch.tensor(inputs['token_type_ids'], dtype=torch.long)\n\n if self.is_test:\n return (ids, mask, token_type_ids)\n else: \n targets = torch.tensor([sample.pct_rank], dtype=torch.float)\n return (ids, mask, token_type_ids, targets)\n\n def __len__(self):\n return len(self.df)","metadata":{"execution":{"iopub.status.busy":"2022-05-18T14:21:19.818399Z","iopub.execute_input":"2022-05-18T14:21:19.8191Z","iopub.status.idle":"2022-05-18T14:21:19.832805Z","shell.execute_reply.started":"2022-05-18T14:21:19.81906Z","shell.execute_reply":"2022-05-18T14:21:19.831892Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Below is a custom `Trainer` class that I wrote from scratch to facilitate my training and validation sub-routines.","metadata":{}},{"cell_type":"code","source":"class Trainer:\n def __init__(self, config, dataloaders, optimizer, model, loss_fns, scheduler, device=\"cuda:0\"):\n self.train_loader, self.valid_loader = dataloaders\n self.train_loss_fn, self.valid_loss_fn = loss_fns\n self.scheduler = scheduler\n self.optimizer = optimizer\n self.model = model\n self.device = torch.device(device)\n self.config = config\n\n def train_one_epoch(self):\n \"\"\"\n Trains the model for 1 epoch\n \"\"\"\n self.model.train()\n train_pbar = tqdm(enumerate(self.train_loader), total=len(self.train_loader))\n train_preds, train_targets = [], []\n\n for bnum, cache in train_pbar:\n ids = self._convert_if_not_tensor(cache[0], dtype=torch.long)\n mask = self._convert_if_not_tensor(cache[1], dtype=torch.long)\n ttis = self._convert_if_not_tensor(cache[2], dtype=torch.long)\n targets = self._convert_if_not_tensor(cache[3], dtype=torch.float)\n \n with autocast(enabled=True):\n outputs = self.model(ids=ids, mask=mask, token_type_ids=ttis).view(-1)\n \n loss = self.train_loss_fn(outputs, targets)\n loss_itm = loss.item()\n \n wandb_log(\n train_batch_loss = loss_itm\n )\n \n train_pbar.set_description('loss: {:.2f}'.format(loss_itm))\n\n Config.scaler.scale(loss).backward()\n Config.scaler.step(self.optimizer)\n Config.scaler.update()\n self.optimizer.zero_grad()\n self.scheduler.step()\n \n train_targets.extend(targets.cpu().detach().numpy().tolist())\n train_preds.extend(outputs.cpu().detach().numpy().tolist())\n \n # Tidy\n del outputs, targets, ids, mask, ttis, loss_itm, loss\n gc.collect()\n torch.cuda.empty_cache()\n \n return train_preds, train_targets\n\n @torch.no_grad()\n def valid_one_epoch(self):\n \"\"\"\n Validates the model for 1 epoch\n \"\"\"\n self.model.eval()\n valid_pbar = tqdm(enumerate(self.valid_loader), total=len(self.valid_loader))\n valid_preds, valid_targets = [], []\n\n for idx, cache in valid_pbar:\n ids = self._convert_if_not_tensor(cache[0], dtype=torch.long)\n mask = self._convert_if_not_tensor(cache[1], dtype=torch.long)\n ttis = self._convert_if_not_tensor(cache[2], dtype=torch.long)\n targets = self._convert_if_not_tensor(cache[3], dtype=torch.float)\n\n outputs = self.model(ids=ids, mask=mask, token_type_ids=ttis).view(-1)\n valid_loss = self.valid_loss_fn(outputs, targets)\n \n wandb_log(\n valid_batch_loss = valid_loss.item()\n )\n \n valid_pbar.set_description(desc=f\"val_loss: {valid_loss.item():.4f}\")\n\n valid_targets.extend(targets.cpu().detach().numpy().tolist())\n valid_preds.extend(outputs.cpu().detach().numpy().tolist())\n\n # Tidy\n del outputs, targets, ids, mask, ttis, valid_loss\n gc.collect()\n torch.cuda.empty_cache()\n \n return valid_preds, valid_targets\n\n def fit(self, epochs: int = 10, output_dir: str = \"/kaggle/working/\", custom_name: str = 'model.pth'):\n \"\"\"\n Low-effort alternative for doing the complete training and validation process\n \"\"\"\n best_loss = int(1e+7)\n best_preds = None\n for epx in range(epochs):\n print(f\"{'='*20} Epoch: {epx+1} / {epochs} {'='*20}\")\n\n train_preds, train_targets = self.train_one_epoch()\n train_mse = mean_squared_error(train_targets, train_preds)\n print(f\"Training loss: {train_mse:.4f}\")\n\n valid_preds, valid_targets = self.valid_one_epoch()\n valid_mse = mean_squared_error(valid_targets, valid_preds)\n print(f\"Validation loss: {valid_mse:.4f}\")\n \n wandb_log(\n train_mse = train_mse,\n valid_mse = valid_mse\n )\n \n if valid_mse < best_loss:\n best_loss = valid_mse\n self.save_model(output_dir, custom_name)\n print(f\"Saved model with val_loss: {best_loss:.4f}\")\n \n def save_model(self, path, name, verbose=False):\n \"\"\"\n Saves the model at the provided destination\n \"\"\"\n try:\n if not os.path.exists(path):\n os.makedirs(path)\n except:\n print(\"Errors encountered while making the output directory\")\n\n torch.save(self.model.state_dict(), os.path.join(path, name))\n if verbose:\n print(f\"Model Saved at: {os.path.join(path, name)}\")\n\n def _convert_if_not_tensor(self, x, dtype):\n if self._tensor_check(x):\n return x.to(self.device, dtype=dtype)\n else:\n return torch.tensor(x, dtype=dtype, device=self.device)\n\n def _tensor_check(self, x):\n return isinstance(x, torch.Tensor)","metadata":{"execution":{"iopub.status.busy":"2022-05-18T14:21:20.459682Z","iopub.execute_input":"2022-05-18T14:21:20.46013Z","iopub.status.idle":"2022-05-18T14:21:20.499219Z","shell.execute_reply.started":"2022-05-18T14:21:20.460097Z","shell.execute_reply":"2022-05-18T14:21:20.498227Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Optimizer only for certain parameters in the model","metadata":{}},{"cell_type":"code","source":"def yield_optimizer(model):\n \"\"\"\n Returns optimizer for specific parameters\n \"\"\"\n param_optimizer = list(model.named_parameters())\n no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n optimizer_parameters = [\n {\n \"params\": [\n p for n, p in param_optimizer if not any(nd in n for nd in no_decay)\n ],\n \"weight_decay\": 0.003,\n },\n {\n \"params\": [\n p for n, p in param_optimizer if any(nd in n for nd in no_decay)\n ],\n \"weight_decay\": 0.0,\n },\n ]\n return transformers.AdamW(optimizer_parameters, lr=Config.LR)","metadata":{"execution":{"iopub.status.busy":"2022-05-18T14:21:21.019545Z","iopub.execute_input":"2022-05-18T14:21:21.019971Z","iopub.status.idle":"2022-05-18T14:21:21.03245Z","shell.execute_reply.started":"2022-05-18T14:21:21.019932Z","shell.execute_reply":"2022-05-18T14:21:21.031744Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Main training code. I will be adding KFolds support soon!","metadata":{}},{"cell_type":"code","source":"# Training Code\nif __name__ == '__main__':\n if torch.cuda.is_available():\n print(\"[INFO] Using GPU: {}\\n\".format(torch.cuda.get_device_name()))\n DEVICE = torch.device('cuda:0')\n else:\n print(\"\\n[INFO] GPU not found. Using CPU: {}\\n\".format(platform.processor()))\n DEVICE = torch.device('cpu')\n\n train_set = AI4CodeDataset(train_df_mark)\n valid_set = AI4CodeDataset(val_df_mark)\n\n train_loader = DataLoader(\n train_set,\n batch_size = Config.TRAIN_BS,\n shuffle = True,\n num_workers = 8\n )\n\n valid_loader = DataLoader(\n valid_set,\n batch_size = Config.VALID_BS,\n shuffle = False,\n num_workers = 8\n )\n\n model = CodeBertModel().to(DEVICE)\n nb_train_steps = int(len(train_df_mark) / Config.TRAIN_BS * Config.NB_EPOCHS)\n optimizer = yield_optimizer(model)\n scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n optimizer, \n T_0=Config.T_0, \n eta_min=Config.η_min\n )\n train_loss_fn, valid_loss_fn = nn.MSELoss(), nn.MSELoss()\n \n wandb.watch(model, criterion=train_loss_fn)\n \n trainer = Trainer(\n config = Config,\n dataloaders = (train_loader, valid_loader),\n loss_fns = (train_loss_fn, valid_loss_fn),\n optimizer = optimizer,\n model = model,\n scheduler = scheduler,\n )\n\n best_pred = trainer.fit(\n epochs = Config.NB_EPOCHS,\n custom_name = f\"ai4code_codebert_small.bin\"\n )","metadata":{"execution":{"iopub.status.busy":"2022-05-18T14:21:21.699845Z","iopub.execute_input":"2022-05-18T14:21:21.700311Z","iopub.status.idle":"2022-05-18T14:22:07.932468Z","shell.execute_reply.started":"2022-05-18T14:21:21.700265Z","shell.execute_reply":"2022-05-18T14:22:07.931021Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Finish the logging run\nrun.finish()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n \n
","metadata":{}}]}