{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Introduction\n\nThis notebook is a PyTorch starter that uses TPUs.\n\nIt is based on HuggingFace's library. I use the already processed data with multilingual Bert.\n\n\n\nThe dataset  used for the pretrained Bert is here : https://www.kaggle.com/theoviel/bertconfigs"},{"metadata":{},"cell_type":"markdown","source":"thanks to @youhanlee > https://www.kaggle.com/youhanlee/bert-pytorch-huggingface-tpu-version-xla\n- Hey, Pietro I hope you will be able to adapt this code for your own project, this code is for binary classification of text data, as part of competition on Kaggle, I collected pieces of code and wrote some myself to make it simple for you and to enable to play around with this code."},{"metadata":{},"cell_type":"markdown","source":"# Intialization"},{"metadata":{},"cell_type":"markdown","source":"### Switching to PyTorch Lightning"},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n!python pytorch-xla-env-setup.py --apt-packages libomp5 libopenblas-dev","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Imports"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"import gc\nimport os\nimport time\nimport math\nimport random\nimport warnings\nimport numpy as np\nimport pandas as pd\nimport seaborn as sns\nimport matplotlib.pyplot as plt\n\nfrom datetime import date\nfrom transformers import *\nfrom sklearn.metrics import *\nfrom tqdm.notebook import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.utils.data\nimport torch.nn.functional as F\n\nfrom torch import Tensor\nfrom torch.optim import *\nfrom torch.nn.modules.loss import *\nfrom torch.optim.lr_scheduler import * \nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.utils.data.sampler import RandomSampler\n\nimport torch_xla\nimport torch_xla.core.xla_model as xm","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Seeding"},{"metadata":{"trusted":true},"cell_type":"code","source":"def seed_everything(seed):\n    \"\"\"\n    Seeds basic parameters for reproductibility of results\n    \n    Arguments:\n        seed {int} -- Number of the seed\n    \"\"\"\n    random.seed(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"seed = 2020\nseed_everything(seed)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data"},{"metadata":{"trusted":true},"cell_type":"code","source":"MODEL_PATHS = {\n    'bert-multi-cased': '../input/bertconfigs/multi_cased_L-12_H-768_A-12/multi_cased_L-12_H-768_A-12/',\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DATA_PATH = '../input/jigsaw-multilingual-toxic-comment-classification/'\n\ndf_val = pd.read_csv(DATA_PATH + 'validation-processed-seqlen128.csv')\ndf_test =  pd.read_csv(DATA_PATH + 'test-processed-seqlen128.csv')\ndf_train = pd.read_csv(DATA_PATH + 'jigsaw-toxic-comment-train-processed-seqlen128.csv')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure(figsize=(10, 5))\nplt.subplot(1, 2, 1)\nsns.countplot(df_train['toxic'])\nplt.title('Target repartition on training data')\n\nplt.subplot(1, 2, 2)\nsns.countplot(df_val['toxic'])\nplt.title('Target repartition on validation data')\n\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class JigsawDataset(Dataset):\n    \"\"\"\n    Torch dataset for the competition.\n    \"\"\"\n    def __init__(self, df):\n        \"\"\"\n        Constructor\n        \n        Arguments:\n            df {pandas dataframe} -- Dataframe where the data is. Expects to be one of the []-processed-seqlen128.csv files\n        \"\"\"\n            \n        super().__init__()\n        self.df = df \n        self.word_ids = np.array([word_ids[1:-1].split(', ') for word_ids in df['input_word_ids']]).astype(int)\n        \n        try:\n            self.y = df['toxic'].values\n        except KeyError: # test data\n            self.y = np.zeros(len(df))\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        return torch.tensor(self.word_ids[idx]), torch.tensor(self.y[idx])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"TRANSFORMERS = {\n    \"bert-multi-cased\": (BertModel, BertTokenizer, \"bert-multi-cased\"),\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# you can play around with this very easily, by defining your own neck structure for the network this doesn't include dropout but it is quite\n# to add in the next cell i found an example that will help you to do so, again there are a ton of ways to do this, and you can play around with \n# as much as you want, and this whole code can run on TPU making it fast, that's why we are using here torch_xla, it is similar to what JAX is to\n#tensorflow and enables us to use the great processing power of TPU's\n\nclass Transformer(nn.Module):\n    def __init__(self, model, num_classes=1):\n        \"\"\"\n        Constructor\n        \n        Arguments:\n            model {string} -- Transformer to build the model on. Expects \"camembert-base\".\n            num_classes {int} -- Number of classes (default: {1})\n        \"\"\"\n        super().__init__()\n        self.name = model\n\n        model_class, tokenizer_class, pretrained_weights = TRANSFORMERS[model]\n\n        bert_config = BertConfig.from_json_file(MODEL_PATHS[model] + 'bert_config.json')\n        bert_config.output_hidden_states = True\n        \n        self.transformer = BertModel(bert_config)\n\n        self.nb_features = self.transformer.pooler.dense.out_features\n\n        self.pooler = nn.Sequential(\n            nn.Linear(self.nb_features, self.nb_features), \n            nn.Tanh(),\n        )\n\n        self.logit = nn.Linear(self.nb_features, num_classes)\n\n    def forward(self, tokens):\n        \"\"\"\n        Usual torch forward function\n        \n        Arguments:\n            tokens {torch tensor} -- Sentence tokens\n        \n        Returns:\n            torch tensor -- Class logits\n        \"\"\"\n        _, _, hidden_states = self.transformer(\n            tokens, attention_mask=(tokens > 0).long()\n        )\n\n        hidden_states = hidden_states[-1][:, 0] # Use the representation of the first token of the last layer\n\n        ft = self.pooler(hidden_states)\n         x = self.logit(ft)\n\n        return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"## an example bert model for customRobert\nclass CustomRoberta(nn.Module):\n    def __init__(self):\n        super(CustomRoberta, self).__init__()\n        self.num_labels = 1\n        self.roberta = transformers.XLMRobertaModel.from_pretrained(\"xlm-roberta-large\", output_hidden_states=False, num_labels=1)\n        self.dropout = nn.Dropout(p=0.2)\n        self.ln = nn.LayerNorm(1024)\n        self.classifier = nn.Linear(1024, self.num_labels)\n\n    def forward(self,\n                input_ids=None,\n                attention_mask=None,\n                position_ids=None,\n                head_mask=None,\n                inputs_embeds=None):\n\n        o1, o2 = self.roberta(input_ids,\n                               attention_mask=attention_mask,\n                               position_ids=position_ids,\n                               head_mask=head_mask,\n                               inputs_embeds=inputs_embeds)\n        \n        x1 = torch.mean(o1, 1)\n        \n        x = x1\n        \n        x = self.ln(x)\n        x = self.dropout(x)\n\n        logits = self.classifier(x)       \n        \n        return logits","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"def fit(model, train_dataset, val_dataset, epochs=1, batch_size=32, warmup_prop=0, lr=5e-5):\n    device = xm.xla_device()\n    model.to(device)\n    \n    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)\n\n    optimizer = AdamW(model.parameters(), lr=lr)\n    \n    num_warmup_steps = int(warmup_prop * epochs * len(train_loader))\n    num_training_steps = epochs * len(train_loader)\n    \n    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps)\n\n    loss_fct = nn.BCEWithLogitsLoss(reduction='mean').to(device)\n    \n    for epoch in range(epochs):\n        model.train()\n        start_time = time.time()\n        \n        optimizer.zero_grad()\n        avg_loss = 0\n        \n        for step, (x, y_batch) in tqdm(enumerate(train_loader), total=len(train_loader)): \n            y_pred = model(x.to(device))\n            \n            loss = loss_fct(y_pred.view(-1).float(), y_batch.float().to(device))\n            loss.backward()\n            avg_loss += loss.item() / len(train_loader)\n\n            xm.optimizer_step(optimizer, barrier=True)\n            scheduler.step()\n            model.zero_grad()\n            optimizer.zero_grad()\n                \n        model.eval()\n        preds = []\n        truths = []\n        avg_val_loss = 0.\n\n        with torch.no_grad():\n            for x, y_batch in val_loader:                \n                y_pred = model(x.to(device))\n                loss = loss_fct(y_pred.detach().view(-1).float(), y_batch.float().to(device))\n                avg_val_loss += loss.item() / len(val_loader)\n                \n                probs = torch.sigmoid(y_pred).detach().cpu().numpy()\n                preds += list(probs.flatten())\n                truths += list(y_batch.numpy().flatten())\n            score = roc_auc_score(truths, preds)\n            \n        \n        dt = time.time() - start_time\n        lr = scheduler.get_last_lr()[0]\n        print(f'Epoch {epoch + 1}/{epochs} \\t lr={lr:.1e} \\t t={dt:.0f}s \\t loss={avg_loss:.4f} \\t val_loss={avg_val_loss:.4f} \\t val_auc={score:.4f}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = Transformer(\"bert-multi-cased\")\ndevice = xm.xla_device()\nmodel = model.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"epochs = 3\nbatch_size = 32\nwarmup_prop = 0.1\nlr = 2e-5  # Important parameter to tweak","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"n = 10000  # I do not train on the entier data as it will take too long (for now)\ntrain_dataset = JigsawDataset(df_train.sample(n))\nval_dataset = JigsawDataset(df_val)\ntest_dataset = JigsawDataset(df_test)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fit(model, train_dataset, val_dataset, epochs=epochs, batch_size=batch_size, warmup_prop=warmup_prop, lr=lr)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Predicting"},{"metadata":{"trusted":true},"cell_type":"code","source":"def predict(model, dataset, batch_size=64):\n    \"\"\"\n    Usual predict torch function\n    \n    Arguments:\n        model {torch model} -- Model to predict with\n        dataset {torch dataset} -- Dataset to get predictions from\n    \n    Keyword Arguments:\n        batch_size {int} -- Batch size (default: {32})\n    \n    Returns:\n        numpy array -- Predictions\n    \"\"\"\n    device = xm.xla_device()\n    model.eval().to(device)\n    preds = np.empty((0, 1))\n    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n\n    with torch.no_grad():\n        for x, _ in tqdm(loader):\n            probs = torch.sigmoid(model(x.to(device))).detach().cpu().numpy()\n            preds = np.concatenate([preds, probs])\n            \n    return preds","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pred_val = predict(model, val_dataset)\ndf_val['pred'] = pred_val","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for language in df_val['lang'].unique():\n    lang_score = roc_auc_score(\n        df_val[df_val['lang'] == language]['toxic'], \n        df_val[df_val['lang']  == language]['pred']\n    )\n    print(f'AUC for language {language}: {lang_score:.4f}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"score = roc_auc_score(df_val['toxic'], pred_val)\nprint(f'Scored {score:.4f} on validation data')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pred_test = predict(model, test_dataset)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sub = pd.read_csv(DATA_PATH + 'sample_submission.csv')\nsub['toxic'] = pred_test\nsub.to_csv('submission.csv', index=False)\nsub.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"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":4,"nbformat_minor":4}