train_vgcn_bert.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5f325a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import numpy as np\n",
    "import pickle as pkl\n",
    "import argparse\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from tqdm import tqdm, trange\n",
    "from sklearn.metrics import classification_report\n",
    "from sklearn.metrics import f1_score\n",
    "\n",
    "from pytorch_pretrained_bert.modeling import BertModel, BertConfig, WEIGHTS_NAME, CONFIG_NAME\n",
    "from pytorch_pretrained_bert.tokenization import BertTokenizer\n",
    "from pytorch_pretrained_bert.optimization import BertAdam  # , warmup_linear\n",
    "\n",
    "\n",
    "from utils import *\n",
    "from model_vgcn_bert import VGCN_Bert\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87a9c6f6",
   "metadata": {},
   "source": [
    "#### Step 1:   Configurations for Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57d47f8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# parser = argparse.ArgumentParser()\n",
    "# # ds = dataset, sw = stopwords, lr = learning rate, l2 = L2 regularization\n",
    "# parser.add_argument('--ds', type=str, default='pheme')\n",
    "# parser.add_argument('--load', type=int, default='0')\n",
    "# parser.add_argument('--sw', type=int, default='1')\n",
    "# parser.add_argument('--dim', type=int, default='16')\n",
    "# parser.add_argument('--lr', type=float, default=1e-5)  # 2e-5\n",
    "# parser.add_argument('--l2', type=float, default=0.01)  # 0.001\n",
    "# parser.add_argument('--model', type=str, default='VGCN_BERT')\n",
    "# args = parser.parse_args()\n",
    "args = {\"ds\": \"pheme\", \"load\": 0, \"sw\": 1, \"dim\": 16,\n",
    "        \"lr\": 1e-5, \"l2\": 0.01, \"model\": \"VGCN_BERT\"}\n",
    "\n",
    "# config_dataset = args.ds\n",
    "# config_load_model_from_checkpoint = True if args.load == 1 else False\n",
    "# config_use_stopwords = True if args.sw == 1 else False\n",
    "# config_gcn_embedding_dim = args.dim\n",
    "# config_learning_rate0 = args.lr\n",
    "# config_l2_decay = args.l2\n",
    "# config_model_type = args.model\n",
    "\n",
    "config_dataset = args[\"ds\"]\n",
    "config_load_model_from_checkpoint = True if args[\"load\"] == 1 else False\n",
    "config_use_stopwords = True if args[\"sw\"] == 1 else False\n",
    "config_gcn_embedding_dim = args[\"dim\"]\n",
    "config_learning_rate0 = args[\"lr\"]\n",
    "config_l2_decay = args[\"l2\"]\n",
    "config_model_type = args[\"model\"]\n",
    "config_data_type = \"text_comments\" # text_comments / text_only / comments_only\n",
    "\n",
    "\n",
    "config_warmup_proportion = 0.1\n",
    "config_vocab_adj = 'all'  # pmi / tf / all\n",
    "config_adj_npmi_threshold = 0.2\n",
    "config_adj_tf_threshold = 0\n",
    "config_loss_criterion = 'cross_entropy'\n",
    "\n",
    "MAX_SEQ_LENGTH = 200 + config_gcn_embedding_dim\n",
    "total_train_epochs = 9\n",
    "batch_size = 16  # 12\n",
    "gradient_accumulation_steps = 1\n",
    "\n",
    "do_lower_case = True\n",
    "perform_metrics_str = ['weighted avg', 'f1-score']\n",
    "do_softmax_before_mse = True\n",
    "\n",
    "if config_dataset == 'pheme':\n",
    "    bert_model_scale = 'bert-base-uncased'\n",
    "elif config_dataset == 'weibo':\n",
    "    bert_model_scale = 'bert-base-chinese'\n",
    "\n",
    "data_dir = './prepared_data/' + config_dataset + '_' + config_data_type\n",
    "output_dir = './model_output/'\n",
    "if not os.path.exists(output_dir):\n",
    "    os.mkdir(output_dir)\n",
    "\n",
    "\n",
    "model_save_file = 'Model_' + config_model_type + '_' \\\n",
    "                + config_dataset + '_' + config_data_type + '.pt'\n",
    "\n",
    "print('----------STEP 1: CONFIGURATIONS FOR TRAINING--------')\n",
    "print('Dataset: ', config_dataset)\n",
    "print('Will Load Model from Checkpoint: ', config_load_model_from_checkpoint)\n",
    "print('Will Delete Stop Words: ', config_use_stopwords)\n",
    "print('Vocab GCN Hidden Dim: vocab_size -> 128 -> ' + str(config_gcn_embedding_dim))\n",
    "print('Learning Rate0: ', config_learning_rate0)\n",
    "print('Weight Decay: ', config_l2_decay)\n",
    "print('Loss Criterion: ', config_loss_criterion)\n",
    "print('Will Perform Softmax before MSE: ', do_softmax_before_mse)\n",
    "print('Vocab Adjcent: ', config_vocab_adj)\n",
    "print('MAX_SEQ_LENGTH: ', MAX_SEQ_LENGTH)\n",
    "print('Perform Metrics: ', perform_metrics_str)\n",
    "print('Saved Model File Name: ', model_save_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49fab5ac",
   "metadata": {},
   "source": [
    "#### Step 2.1: Prepare Dataset & Load Vocabulary Adjacent Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbc1a634",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('----------STEP 2: PREPARE DATASET & LOAD VOCABULARY ADJACENT MATRIX----------')\n",
    "print(' Load and seperate', config_dataset, 'dataset, with vocabulary graph adjacent matrix')\n",
    "\n",
    "objects = []\n",
    "names = ['index_label', 'train_label', 'train_label_prob', 'test_label',\n",
    "         'test_label_prob', 'clean_docs', 'vocab_adj_tf', 'vocab_adj_pmi', 'vocab_map']\n",
    "\n",
    "for i in range(len(names)):\n",
    "    datafile = data_dir + \"/data_%s.%s\" % (config_dataset, names[i])\n",
    "    with open(datafile, 'rb') as f:\n",
    "        objects.append(pkl.load(f, encoding='latin1'))\n",
    "\n",
    "index_labels_list, train_label, train_label_prob, test_label, test_label_prob, shuffled_clean_docs, gcn_vocab_adj_tf, gcn_vocab_adj_pmi, gcn_vocab_map = tuple(objects)\n",
    "\n",
    "label2idx = index_labels_list[0]\n",
    "idx2label = index_labels_list[1]\n",
    "\n",
    "all_labels = np.hstack((train_label, test_label))\n",
    "all_labels_prob = np.vstack((train_label_prob, test_label_prob))\n",
    "\n",
    "examples = []\n",
    "for i, text in enumerate(shuffled_clean_docs):\n",
    "    example = InputExample(i, text.strip(), confidence=all_labels_prob[i], label=all_labels[i])\n",
    "    examples.append(example)\n",
    "\n",
    "num_classes = len(label2idx)\n",
    "gcn_vocab_size = len(gcn_vocab_map)\n",
    "train_size = len(train_label)\n",
    "test_size = len(test_label)\n",
    "\n",
    "indexs = np.arange(0, len(examples))\n",
    "train_examples = [examples[i] for i in indexs[:train_size]]\n",
    "test_examples = [examples[i] for i in indexs[train_size:train_size + test_size]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "363fd7e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "if config_adj_tf_threshold > 0:\n",
    "    gcn_vocab_adj_tf.data *= (gcn_vocab_adj_tf.data > config_adj_tf_threshold)\n",
    "    gcn_vocab_adj_tf.eliminate_zeros()\n",
    "if config_adj_npmi_threshold > 0:\n",
    "    gcn_vocab_adj_pmi.data *= (gcn_vocab_adj_pmi.data > config_adj_npmi_threshold)\n",
    "    gcn_vocab_adj_pmi.eliminate_zeros()\n",
    "\n",
    "if config_vocab_adj == 'pmi':\n",
    "    gcn_vocab_adj_list = [gcn_vocab_adj_pmi]\n",
    "elif config_vocab_adj == 'tf':\n",
    "    gcn_vocab_adj_list = [gcn_vocab_adj_tf]\n",
    "elif config_vocab_adj == 'all':\n",
    "    gcn_vocab_adj_list = [gcn_vocab_adj_tf, gcn_vocab_adj_pmi]\n",
    "\n",
    "norm_gcn_vocab_adj_list = []\n",
    "for i in range(len(gcn_vocab_adj_list)):\n",
    "    adj = gcn_vocab_adj_list[i]\n",
    "\n",
    "    print('Zero ratio for vocab adj %dth: %.8f' %\n",
    "          (i, 100 * (1 - adj.count_nonzero() / (adj.shape[0] * adj.shape[1]))))\n",
    "\n",
    "    adj = normalize_adj(adj)\n",
    "    norm_gcn_vocab_adj_list.append(sparse_scipy2torch(adj.tocoo()).to(device))\n",
    "\n",
    "gcn_adj_list = norm_gcn_vocab_adj_list\n",
    "\n",
    "\n",
    "train_classes_num, train_classes_weight = get_class_count_and_weight(train_label, len(label2idx))\n",
    "loss_weight = torch.tensor(train_classes_weight).to(device)\n",
    "loss_weight = torch.tensor(loss_weight, dtype=torch.float32).to(device)\n",
    "\n",
    "tokenizer = BertTokenizer.from_pretrained(bert_model_scale, do_lower_case=do_lower_case)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e606fcc",
   "metadata": {},
   "source": [
    "#### Step 2.2:   Prepare PyTorch DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ed4db8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_pytorch_dataloader(examples, tokenizer, batch_size):\n",
    "    dataset = CorpusDataset(examples, tokenizer, gcn_vocab_map, MAX_SEQ_LENGTH, config_gcn_embedding_dim)\n",
    "    return DataLoader(dataset=dataset, batch_size=batch_size, shuffle=False, num_workers=4, collate_fn=dataset.pad)\n",
    "\n",
    "\n",
    "train_dataloader = get_pytorch_dataloader(train_examples, tokenizer, batch_size)\n",
    "test_dataloader = get_pytorch_dataloader(test_examples, tokenizer, batch_size)\n",
    "\n",
    "total_train_steps = int(len(train_dataloader) / gradient_accumulation_steps * total_train_epochs)\n",
    "\n",
    "print('Train Classes Count: ', train_classes_num)\n",
    "print('Batch size: ', batch_size)\n",
    "print('Num steps: ', total_train_steps)\n",
    "print('Number of Examples for Training: ', len(train_examples))\n",
    "print('Number of Examples for Training After Dataloader: ', len(train_dataloader) * batch_size)\n",
    "print('Number of Examples for Test: ', len(test_examples))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5d32200",
   "metadata": {},
   "source": [
    "#### Step 3.1:  Define Evaluating Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc8f26c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(model, examples, tokenizer, batch_size):\n",
    "    dataloader = get_pytorch_dataloader(examples, tokenizer, batch_size)\n",
    "    predict_out = []\n",
    "    confidence_out = []\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for batch in dataloader:\n",
    "            batch = tuple(t.to(device) for t in batch)\n",
    "            input_ids, input_mask, segment_ids, _, _, gcn_swop_eye = batch\n",
    "\n",
    "            _, score_out = model(gcn_adj_list, gcn_swop_eye,\n",
    "                                 input_ids, segment_ids, input_mask)\n",
    "            if config_loss_criterion == 'mse' and do_softmax_before_mse:\n",
    "                score_out = torch.nn.functional.softmax(score_out, dim=-1)\n",
    "            predict_out.extend(score_out.max(1)[1].tolist())\n",
    "            confidence_out.extend(score_out.max(1)[0].tolist())\n",
    "\n",
    "    return np.array(predict_out).reshape(-1), np.array(confidence_out).reshape(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "865ad863",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model, gcn_adj_list, predict_dataloader, epoch_th, dataset_name):\n",
    "    model.eval()\n",
    "    predict_out = []\n",
    "    all_label_ids = []\n",
    "    ev_loss = 0\n",
    "    total = 0\n",
    "    correct = 0\n",
    "    start = time.time()\n",
    "    with torch.no_grad():\n",
    "        for batch in predict_dataloader:\n",
    "            batch = tuple(t.to(device) for t in batch)\n",
    "            input_ids, input_mask, segment_ids, y_prob, label_ids, gcn_swop_eye = batch\n",
    "            _, logits = model(gcn_adj_list, gcn_swop_eye,input_ids,  segment_ids, input_mask)\n",
    "\n",
    "            if config_loss_criterion == 'mse':\n",
    "                if do_softmax_before_mse:\n",
    "                    logits = F.softmax(logits, -1)\n",
    "                loss = F.mse_loss(logits, y_prob)\n",
    "            else:\n",
    "                loss = F.cross_entropy(logits.view(-1, num_classes), label_ids)\n",
    "            ev_loss += loss.item()\n",
    "\n",
    "            _, predicted = torch.max(logits, -1)\n",
    "            predict_out.extend(predicted.tolist())\n",
    "            all_label_ids.extend(label_ids.tolist())\n",
    "            eval_accuracy = predicted.eq(label_ids).sum().item()\n",
    "            total += len(label_ids)\n",
    "            correct += eval_accuracy\n",
    "\n",
    "        f1_metrics = f1_score(np.array(all_label_ids).reshape(-1), np.array(predict_out).reshape(-1), average='weighted')\n",
    "        print(\"Report:\\n\" + classification_report(np.array(all_label_ids).reshape(-1), np.array(predict_out).reshape(-1), digits=4))\n",
    "\n",
    "    ev_acc = correct/total\n",
    "    end = time.time()\n",
    "    print('Epoch : %d, %s: %.3f Acc : %.3f on %s, Spend:%.3f minutes for evaluation'\n",
    "          % (epoch_th, ' '.join(perform_metrics_str), 100 * f1_metrics, 100. * ev_acc, dataset_name, (end - start) / 60.0))\n",
    "    print('*' * 50)\n",
    "    return ev_loss, ev_acc, f1_metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fa0f581",
   "metadata": {},
   "source": [
    "#### Step 3.2:   Load / Initialize VGCN_BERT Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77bd8c81",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('----------STEP 3: START TRAINING VGCN_BERT MODEL----------')\n",
    "\n",
    "if config_load_model_from_checkpoint and os.path.exists(os.path.join(output_dir, model_save_file)):\n",
    "    checkpoint = torch.load(os.path.join(output_dir, model_save_file), map_location='cpu')\n",
    "    if 'step' in checkpoint:\n",
    "        prev_save_step = checkpoint['step']\n",
    "        start_epoch = checkpoint['epoch']\n",
    "    else:\n",
    "        prev_save_step = -1\n",
    "        start_epoch = checkpoint['epoch'] + 1\n",
    "\n",
    "    test_acc_prev = checkpoint['test_acc']\n",
    "    perform_metrics_prev = checkpoint['perform_metrics']\n",
    "    model = VGCN_Bert.from_pretrained(bert_model_scale, state_dict=checkpoint['model_state'], gcn_adj_dim=gcn_vocab_size, \n",
    "        gcn_adj_num=len(gcn_adj_list), gcn_embedding_dim=config_gcn_embedding_dim, num_labels=len(label2idx))\n",
    "\n",
    "    pretrained_dict = checkpoint['model_state']\n",
    "    net_state_dict = model.state_dict()\n",
    "    pretrained_dict_selected = {\n",
    "        k: v for k, v in pretrained_dict.items() if k in net_state_dict}\n",
    "    net_state_dict.update(pretrained_dict_selected)\n",
    "    model.load_state_dict(net_state_dict)\n",
    "\n",
    "    print('Loaded the pre-trained model:', model_save_file, ', epoch:', checkpoint['epoch'], 'step:', prev_save_step, 'test acc:',\n",
    "          checkpoint['test_acc'], ' '.join(perform_metrics_str) + '_test:', checkpoint['perform_metrics'])\n",
    "\n",
    "else:\n",
    "    start_epoch = 0\n",
    "    test_acc_prev = 0\n",
    "    perform_metrics_prev = 0\n",
    "    model = VGCN_Bert.from_pretrained(bert_model_scale, gcn_adj_dim=gcn_vocab_size, gcn_adj_num=len(\n",
    "        gcn_adj_list), gcn_embedding_dim=config_gcn_embedding_dim, num_labels=len(label2idx))\n",
    "    prev_save_step = -1\n",
    "\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af5f99d6",
   "metadata": {},
   "source": [
    "#### Step 3.3:   Start Training VGCN_BERT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "040c2068",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = BertAdam(model.parameters(), lr=config_learning_rate0,\n",
    "                     warmup=config_warmup_proportion, t_total=total_train_steps, weight_decay=config_l2_decay)\n",
    "\n",
    "train_start = time.time()\n",
    "global_step_th = int(len(train_examples) / batch_size /\n",
    "                     gradient_accumulation_steps * start_epoch)\n",
    "\n",
    "all_loss_list = {'train': [], 'test': []}\n",
    "all_f1_list = {'train': [], 'test': []}\n",
    "for epoch in range(start_epoch, total_train_epochs):\n",
    "    train_loss = 0\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    for step, batch in enumerate(train_dataloader):\n",
    "        if prev_save_step > -1:\n",
    "            if step <= prev_save_step:\n",
    "                continue\n",
    "        if prev_save_step > -1:\n",
    "            prev_save_step = -1\n",
    "        batch = tuple(t.to(device) for t in batch)\n",
    "        input_ids, input_mask, segment_ids, y_prob, label_ids, gcn_swop_eye = batch\n",
    "\n",
    "        _, logits = model(gcn_adj_list, gcn_swop_eye,\n",
    "                          input_ids, segment_ids, input_mask)\n",
    "\n",
    "        if config_loss_criterion == 'mse':\n",
    "            if do_softmax_before_mse:\n",
    "                logits = F.softmax(logits, -1)\n",
    "            loss = F.mse_loss(logits, y_prob)\n",
    "        else:\n",
    "            if loss_weight is None:\n",
    "                loss = F.cross_entropy(logits, label_ids)\n",
    "            else:\n",
    "                loss = F.cross_entropy(\n",
    "                    logits.view(-1, num_classes), label_ids, loss_weight)\n",
    "\n",
    "        if gradient_accumulation_steps > 1:\n",
    "            loss = loss / gradient_accumulation_steps\n",
    "        loss.backward()\n",
    "\n",
    "        train_loss += loss.item()\n",
    "        if (step + 1) % gradient_accumulation_steps == 0:\n",
    "            optimizer.step()\n",
    "            optimizer.zero_grad()\n",
    "            global_step_th += 1\n",
    "        if step % 40 == 0:\n",
    "            print(\"Epoch:{}-{}/{}, Train {} Loss: {}, Cumulated Time: {}m \".format(epoch, step,\n",
    "                  len(train_dataloader), config_loss_criterion, loss.item(), (time.time() - train_start)/60.0))\n",
    "\n",
    "    print('*' * 50)\n",
    "    test_loss, test_acc, perform_metrics = evaluate(model, gcn_adj_list, test_dataloader, epoch, 'Test_set')\n",
    "    all_loss_list['train'].append(train_loss)\n",
    "    all_loss_list['test'].append(test_loss)\n",
    "    all_f1_list['test'].append(perform_metrics)\n",
    "    print(\"Epoch:{} Completed, Total Train Loss:{}, Test Loss:{}, Spend {}m \".format(\n",
    "        epoch, train_loss, test_loss, (time.time() - train_start) / 60.0))\n",
    "\n",
    "    if perform_metrics > perform_metrics_prev:\n",
    "        to_save = {'epoch': epoch, 'model_state': model.state_dict(),\n",
    "                   'test_acc': test_acc, 'lower_case': do_lower_case,\n",
    "                   'perform_metrics': perform_metrics}\n",
    "        torch.save(to_save, os.path.join(output_dir, model_save_file))\n",
    "\n",
    "        perform_metrics_prev = perform_metrics\n",
    "\n",
    "        test_f1_best_epoch = epoch\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ed9b61c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Optimization Finished! Total Spend Time:', (time.time() - train_start)/60.0)\n",
    "print('Test Weighted F1: %.3f at %d Epoch.' % (100 * perform_metrics_prev, test_f1_best_epoch))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6eec3f2a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "409.6px"
   },
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}