AI4Code_PyTorch.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "AI4Code-PyTorch.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "mount_file_id": "1gzFRQ0T91t8vYx-qzuHvYBW-IZRTQbyE",
      "authorship_tag": "ABX9TyNUph+LtndyTpsXP1minyyf",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "2a0858a074804aaa89c2b7484513529d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_174c839c1281482a94736e45876990b6",
              "IPY_MODEL_c3d113798e4847eeb2c0e8af6905ce4a",
              "IPY_MODEL_e9d2f7d74f8d485485f6cff8f992e87c"
            ],
            "layout": "IPY_MODEL_f85b6eec2fed44a984e11ab233340b51"
          }
        },
        "174c839c1281482a94736e45876990b6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_4a66b1d2b77048228d9cb96279e5ab18",
            "placeholder": "​",
            "style": "IPY_MODEL_e2b74c76fb0c4437a8511669b9ebb34d",
            "value": "Downloading: 100%"
          }
        },
        "c3d113798e4847eeb2c0e8af6905ce4a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_17da8b41943448669f1376949e708f75",
            "max": 1344997306,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_29a766a2de144665ae1e85901d76a867",
            "value": 1344997306
          }
        },
        "e9d2f7d74f8d485485f6cff8f992e87c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_34e11419597e4538ab938911b2ca5e5d",
            "placeholder": "​",
            "style": "IPY_MODEL_df1164bd238a499ea6f08461381ef164",
            "value": " 1.25G/1.25G [00:20&lt;00:00, 67.9MB/s]"
          }
        },
        "f85b6eec2fed44a984e11ab233340b51": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "4a66b1d2b77048228d9cb96279e5ab18": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "e2b74c76fb0c4437a8511669b9ebb34d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "17da8b41943448669f1376949e708f75": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "29a766a2de144665ae1e85901d76a867": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "34e11419597e4538ab938911b2ca5e5d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "df1164bd238a499ea6f08461381ef164": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "4d2a1f9598224094b5bfe4ca424b8b95": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_13d840bb407c4784a1d14f29327a5159",
              "IPY_MODEL_126d3e07731b4fed931a1293f3995ae4",
              "IPY_MODEL_460274a114a84c11a86375986fd6cd5f"
            ],
            "layout": "IPY_MODEL_07c5ffc9611d4ec485000d3800c0ca4f"
          }
        },
        "13d840bb407c4784a1d14f29327a5159": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_af624fa94e124d9382ad436eda96efd7",
            "placeholder": "​",
            "style": "IPY_MODEL_171ef91b77f14df78c4ad319c8b565af",
            "value": "loss: 0.08:  29%"
          }
        },
        "126d3e07731b4fed931a1293f3995ae4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_58bc23b3eb974b06b94a12df6c68661c",
            "max": 8836,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_a6267a0a8bfd45b288b2d1c4e155bed7",
            "value": 2525
          }
        },
        "460274a114a84c11a86375986fd6cd5f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_01d38445a3f840cdab15e20f3dc2f45d",
            "placeholder": "​",
            "style": "IPY_MODEL_a9d965747ead4859a480d7f38324a866",
            "value": " 2524/8836 [27:27&lt;1:08:23,  1.54it/s]"
          }
        },
        "07c5ffc9611d4ec485000d3800c0ca4f": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "af624fa94e124d9382ad436eda96efd7": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "171ef91b77f14df78c4ad319c8b565af": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "58bc23b3eb974b06b94a12df6c68661c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a6267a0a8bfd45b288b2d1c4e155bed7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "01d38445a3f840cdab15e20f3dc2f45d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a9d965747ead4859a480d7f38324a866": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jasonsgraham/kaggle_notebooks/blob/main/AI4Code_PyTorch.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# AI4Code PyTorch Training Starter + AMP + 🤗 + W&B 📉 - BERT Large\n",
        "\n",
        "This is training script written in Vanilla PyTorch with a Custom Trainer Class. I hope it can help you in developing more sophisticated models.\n",
        "\n",
        "Along with BERT Large uncased. If you want to use that, you'll have to make some small changes in the code. Refer to official HuggingFace documentation for more.\n",
        "\n",
        "Think of this notebook has a skeleton for all BERT based 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",
        "\n",
        "I'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": {
        "id": "Im3UJpSQ4Ksc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Default title text\n",
        "!wget -q https://gist.githubusercontent.com/jasonsgraham/f63e1737121e2154ee3ad228398137e2/raw/setup_colab.py -O colab_setup.py\n",
        "%run colab_setup.py"
      ],
      "metadata": {
        "id": "jcL9Qj9m4tg5",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "864b2ee0-8c8c-4881-a99f-61d24efd639e"
      },
      "execution_count": 38,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loading WANDB api key.\n",
            "[Tesla P100-PCIE-16GB] RAM: 0.00/0.00 GB used\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%sh\n",
        "pip install -q --upgrade transformers\n",
        "pip install -q --upgrade wandb"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-14T04:11:00.07767Z",
          "iopub.execute_input": "2022-05-14T04:11:00.077996Z",
          "iopub.status.idle": "2022-05-14T04:11:16.098168Z",
          "shell.execute_reply.started": "2022-05-14T04:11:00.077901Z",
          "shell.execute_reply": "2022-05-14T04:11:16.097333Z"
        },
        "trusted": true,
        "id": "XmCqytz54Ksv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import platform\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from tqdm.notebook import tqdm\n",
        "import matplotlib.pyplot as plt\n",
        "import os\n",
        "import gc\n",
        "import os\n",
        "import wandb\n",
        "import json\n",
        "import glob\n",
        "from scipy import sparse\n",
        "from pathlib import Path\n",
        "\n",
        "import torch\n",
        "import transformers\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.cuda.amp import GradScaler, autocast\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "\n",
        "from sklearn.model_selection import GroupShuffleSplit\n",
        "from sklearn.metrics import mean_squared_error\n",
        "import warnings\n",
        "warnings.simplefilter('ignore')"
      ],
      "metadata": {
        "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
        "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
        "execution": {
          "iopub.status.busy": "2022-05-14T04:11:16.101315Z",
          "iopub.execute_input": "2022-05-14T04:11:16.101764Z",
          "iopub.status.idle": "2022-05-14T04:11:23.590434Z",
          "shell.execute_reply.started": "2022-05-14T04:11:16.101729Z",
          "shell.execute_reply": "2022-05-14T04:11:23.589678Z"
        },
        "trusted": true,
        "id": "uUmNQ_Ln4Ks1"
      },
      "execution_count": 40,
      "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.\n",
        "This makes the code more modular and easy to approach at the same time."
      ],
      "metadata": {
        "id": "B-kOIGo-4Ks3"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Keeping the number of epochs to 2 since a single epoch on 15K samples take ~2 hours and 13 minutes. You can change it to however you need.\n",
        "\n",
        "The notebook will not give an OOM error at any point, should you change the epoch since I have written the code to be heavily optimized."
      ],
      "metadata": {
        "id": "TlUhebnp4Ks4"
      }
    },
    {
      "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 = 'bert-large-uncased'\n",
        "    data_dir = Path(DRIVE / '../input/AI4Code')\n",
        "    TOKENIZER = transformers.BertTokenizer.from_pretrained(MODEL_NAME, do_lower_case=True)\n",
        "    scaler = GradScaler()"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-14T04:11:23.591732Z",
          "iopub.execute_input": "2022-05-14T04:11:23.592008Z",
          "iopub.status.idle": "2022-05-14T04:11:26.941421Z",
          "shell.execute_reply.started": "2022-05-14T04:11:23.591963Z",
          "shell.execute_reply": "2022-05-14T04:11:26.940677Z"
        },
        "trusted": true,
        "id": "ci3dDwHM4Ks6"
      },
      "execution_count": 41,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## About W&B:\n",
        "<center><img src=\"https://i.imgur.com/gb6B4ig.png\" width=\"400\" alt=\"Weights & Biases\"/></center><br>\n",
        "<p style=\"text-align:center\">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.\n",
        "We 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.<br><br></p>"
      ],
      "metadata": {
        "id": "lV03rvIm4Ks9"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "To login to W&B, you can use below snippet.\n",
        "\n",
        "```python\n",
        "from kaggle_secrets import UserSecretsClient\n",
        "user_secrets = UserSecretsClient()\n",
        "wb_key = user_secrets.get_secret(\"WANDB_API_KEY\")\n",
        "\n",
        "wandb.login(key=wb_key)\n",
        "```\n",
        "Make sure you have your W&B key stored as `WANDB_API_KEY` under Add-ons -> Secrets\n",
        "\n",
        "You can view [this](https://www.kaggle.com/ayuraj/experiment-tracking-with-weights-and-biases) notebook to learn more about W&B tracking.\n",
        "\n",
        "If you don't want to login to W&B, the kernel will still work and log everything to W&B in anonymous mode."
      ],
      "metadata": {
        "id": "mEDvpy5q4Ks-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(os.environ.keys())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xYq5nFqlZI-S",
        "outputId": "5b8b7dab-1c64-4707-9721-19a80bbdd2ad"
      },
      "execution_count": 42,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "KeysView(environ({'CUDNN_VERSION': '8.0.5.39', '__EGL_VENDOR_LIBRARY_DIRS': '/usr/lib64-nvidia:/usr/share/glvnd/egl_vendor.d/', 'LD_LIBRARY_PATH': '/usr/lib64-nvidia', 'CLOUDSDK_PYTHON': 'python3', 'LANG': 'en_US.UTF-8', 'HOSTNAME': 'd953166fb88c', 'OLDPWD': '/', 'CLOUDSDK_CONFIG': '/content/.config', 'NVIDIA_VISIBLE_DEVICES': 'all', 'DATALAB_SETTINGS_OVERRIDES': '{\"kernelManagerProxyPort\":6000,\"kernelManagerProxyHost\":\"172.28.0.3\",\"jupyterArgs\":[\"--ip=172.28.0.2\"],\"debugAdapterMultiplexerPath\":\"/usr/local/bin/dap_multiplexer\",\"enableLsp\":true}', 'ENV': '/root/.bashrc', 'NCCL_VERSION': '2.7.8', 'TF_FORCE_GPU_ALLOW_GROWTH': 'true', 'NO_GCE_CHECK': 'False', 'PWD': '/', 'HOME': '/root', 'LAST_FORCED_REBUILD': '20220614', 'DEBIAN_FRONTEND': 'noninteractive', 'LIBRARY_PATH': '/usr/local/cuda/lib64/stubs', 'GCE_METADATA_TIMEOUT': '3', 'GLIBCPP_FORCE_NEW': '1', 'TBE_CREDS_ADDR': '172.28.0.1:8008', 'SHELL': '/bin/bash', 'GCS_READ_CACHE_BLOCK_SIZE_MB': '16', 'PYTHONWARNINGS': 'ignore:::pip._internal.cli.base_command', 'CUDA_VERSION': '11.1.1', 'NVIDIA_DRIVER_CAPABILITIES': 'compute,utility', 'SHLVL': '0', 'PYTHONPATH': '/env/python', 'NVIDIA_REQUIRE_CUDA': 'cuda>=11.1 brand=tesla,driver>=418,driver<419 brand=tesla,driver>=440,driver<441 brand=tesla,driver>=450,driver<451', 'TBE_EPHEM_CREDS_ADDR': '172.28.0.1:8009', 'COLAB_GPU': '1', 'GLIBCXX_FORCE_NEW': '1', 'PATH': '/opt/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/tools/node/bin:/tools/google-cloud-sdk/bin', 'LD_PRELOAD': '/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4', 'JPY_PARENT_PID': '56', 'TERM': 'xterm-color', 'CLICOLOR': '1', 'PAGER': 'cat', 'GIT_PAGER': 'cat', 'MPLBACKEND': 'module://ipykernel.pylab.backend_inline', 'ENABLE_DIRECTORYPREFETCHER': '1', 'USE_AUTH_EPHEM': '1', 'PYDEVD_USE_FRAME_EVAL': 'NO', 'GDRIVE': '/content/drive/MyDrive', 'COLAB': '/content/drive/MyDrive/colab', 'GH_USER': 'jasonsgraham', 'GH_PAT': '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', 'GITHUB_TOKEN': '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', 'MLFLOW_TRACKING_URI': 'http://18.214.164.39:5000/', 'MLFLOW_TRACKING_USERNAME': 'mlflow', 'MLFLOW_TRACKING_PASSWORD': 'HZ^a483Lu%ZCCQ', 'KMP_DUPLICATE_LIB_OK': 'True', 'KMP_INIT_AT_FORK': 'FALSE', 'WANDB_API_KEY': 'fbb8605231534383480c88a5740acf47d1adbb6c'}))\n"
          ]
        }
      ]
    },
    {
      "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\n",
        "# from kaggle_secrets import UserSecretsClient\n",
        "# user_secrets = UserSecretsClient()\n",
        "wb_key = os.environ['WANDB_API_KEY']\n",
        "\n",
        "wandb.login(key=wb_key)\n",
        "\n",
        "run = wandb.init(\n",
        "    project='pytorch',\n",
        "    config=WANDB_CONFIG,\n",
        "    group='nlp',\n",
        "    job_type='train',\n",
        ")"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-14T04:11:26.944268Z",
          "iopub.execute_input": "2022-05-14T04:11:26.944776Z",
          "iopub.status.idle": "2022-05-14T04:11:40.306001Z",
          "shell.execute_reply.started": "2022-05-14T04:11:26.944738Z",
          "shell.execute_reply": "2022-05-14T04:11:40.305126Z"
        },
        "trusted": true,
        "id": "E3sup5YW4KtA",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 139
        },
        "outputId": "02e4533e-2aec-4110-d776-384b3d471794"
      },
      "execution_count": 43,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mjgraham20\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m If you're specifying your api key in code, ensure this code is not shared publicly.\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Consider setting the WANDB_API_KEY environment variable, or running `wandb login` from the command line.\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: Appending key for api.wandb.ai to your netrc file: /root/.netrc\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Tracking run with wandb version 0.12.19"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Run data is saved locally in <code>/content/wandb/run-20220629_031252-28336ijo</code>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Syncing run <strong><a href=\"https://wandb.ai/jgraham20/pytorch/runs/28336ijo\" target=\"_blank\">quiet-dream-1</a></strong> to <a href=\"https://wandb.ai/jgraham20/pytorch\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://wandb.me/run\" target=\"_blank\">docs</a>)<br/>"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Below are some utility functions that we will be using."
      ],
      "metadata": {
        "id": "idY4jNQu4KtD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from bisect import bisect\n",
        "\n",
        "def 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",
        "\n",
        "def 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",
        "\n",
        "def 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",
        "\n",
        "def 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",
        "\n",
        "def get_ranks(base, derived):\n",
        "    return [base.index(d) for d in derived]"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-14T04:11:40.307592Z",
          "iopub.execute_input": "2022-05-14T04:11:40.308009Z",
          "iopub.status.idle": "2022-05-14T04:11:40.319771Z",
          "shell.execute_reply.started": "2022-05-14T04:11:40.307959Z",
          "shell.execute_reply": "2022-05-14T04:11:40.319002Z"
        },
        "trusted": true,
        "id": "7AVNsWzK4KtE"
      },
      "execution_count": 44,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import sys\n",
        "GOOGLE_COLAB = 'google.colab' in sys.modules\n",
        "\n",
        "if GOOGLE_COLAB:\n",
        "  data_dir = Path('/content/drive/MyDrive/Colab Notebooks/input/AI4Code')\n",
        "  output_dir = Path('/content/drive/MyDrive/Colab Notebooks/output/AI4Code')\n",
        "  train_parquet_file = data_dir / 'train.parquet'\n",
        "else:\n",
        "  data_dir = data_dir = Path('../input/AI4Code')\n",
        "  output_dir = Path('./')"
      ],
      "metadata": {
        "id": "oylX7FiUdSsR"
      },
      "execution_count": 46,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Data preprocessing for further training process taken from the Starter notebook for this competition."
      ],
      "metadata": {
        "id": "o1RQQ6Zl4KtF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "NUM_TRAIN = 10000\n",
        "\n",
        "\n",
        "def 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",
        "\n",
        "if train_parquet_file.exists():\n",
        "  df = pd.read_parquet(train_parquet_file)\n",
        "else:\n",
        "  paths_train = list((data_dir / 'train').glob('*.json'))[:NUM_TRAIN]\n",
        "  notebooks_train = [\n",
        "      read_notebook(path) for path in tqdm(paths_train, desc='Train NBs')\n",
        "  ]\n",
        "  df = (\n",
        "      pd.concat(notebooks_train)\n",
        "      .set_index('id', append=True)\n",
        "      .swaplevel()\n",
        "      .sort_index(level='id', sort_remaining=False)\n",
        "  )\n",
        "  df.to_parquet(train_parquet_file)\n",
        "\n",
        "df_orders = pd.read_csv(\n",
        "    data_dir / 'train_orders.csv',\n",
        "    index_col='id',\n",
        "    squeeze=True,\n",
        ").str.split()\n",
        "\n",
        "df_orders_ = df_orders.to_frame().join(\n",
        "    df.reset_index('cell_id').groupby('id')['cell_id'].apply(list),\n",
        "    how='right',\n",
        ")\n",
        "\n",
        "ranks = {}\n",
        "for id_, cell_order, cell_id in df_orders_.itertuples():\n",
        "    ranks[id_] = {'cell_id': cell_id, 'rank': get_ranks(cell_order, cell_id)}\n",
        "\n",
        "df_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",
        "\n",
        "df_ancestors = pd.read_csv(data_dir / 'train_ancestors.csv', index_col='id')\n",
        "df = df.reset_index().merge(df_ranks, on=[\"id\", \"cell_id\"]).merge(df_ancestors, on=[\"id\"])\n",
        "df[\"pct_rank\"] = df[\"rank\"] / df.groupby(\"id\")[\"cell_id\"].transform(\"count\")"
      ],
      "metadata": {
        "jupyter": {
          "outputs_hidden": false
        },
        "execution": {
          "iopub.status.busy": "2022-06-22T20:12:00.441182Z",
          "iopub.execute_input": "2022-06-22T20:12:00.441447Z",
          "iopub.status.idle": "2022-06-22T20:13:53.598925Z",
          "shell.execute_reply.started": "2022-06-22T20:12:00.441420Z",
          "shell.execute_reply": "2022-06-22T20:13:53.598409Z"
        },
        "trusted": true,
        "id": "itkesgMJImYN"
      },
      "execution_count": 52,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "NVALID = 0.1  # size of validation set\n",
        "\n",
        "splitter = GroupShuffleSplit(n_splits=1, test_size=NVALID, random_state=0)\n",
        "\n",
        "train_ind, val_ind = next(splitter.split(df, groups=df[\"ancestor_id\"]))\n",
        "\n",
        "train_df = df.loc[train_ind].reset_index(drop=True)\n",
        "val_df = df.loc[val_ind].reset_index(drop=True)\n",
        "\n",
        "train_df_mark = train_df[train_df[\"cell_type\"] == \"markdown\"].reset_index(drop=True)\n",
        "val_df_mark = val_df[val_df[\"cell_type\"] == \"markdown\"].reset_index(drop=True)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-14T04:13:50.45171Z",
          "iopub.execute_input": "2022-05-14T04:13:50.451956Z",
          "iopub.status.idle": "2022-05-14T04:13:52.056939Z",
          "shell.execute_reply.started": "2022-05-14T04:13:50.451923Z",
          "shell.execute_reply": "2022-05-14T04:13:52.0562Z"
        },
        "trusted": true,
        "id": "VnruVJyn4KtH"
      },
      "execution_count": 53,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Model class for BERT Large model.\n",
        "\n",
        "**Note: Keep in mind to run this model carefully as it's pretty big and a little large parameter somewhere can give you can OOM (Out Of Memory) error for your GPU**"
      ],
      "metadata": {
        "id": "88tUNEbu4KtJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class BERTLargeModel(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(BERTLargeModel, self).__init__()\n",
        "        self.bert = transformers.BertModel.from_pretrained(Config.MODEL_NAME)\n",
        "        self.drop = nn.Dropout(0.3)\n",
        "        self.fc = nn.Linear(1024, 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-14T04:13:52.058115Z",
          "iopub.execute_input": "2022-05-14T04:13:52.058555Z",
          "iopub.status.idle": "2022-05-14T04:13:52.066832Z",
          "shell.execute_reply.started": "2022-05-14T04:13:52.058518Z",
          "shell.execute_reply": "2022-05-14T04:13:52.066164Z"
        },
        "trusted": true,
        "id": "h3FLmlyp4KtK"
      },
      "execution_count": 54,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Custom dataset for the Markdown cells"
      ],
      "metadata": {
        "id": "wYEQH-Tf4KtL"
      }
    },
    {
      "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-14T04:13:52.069904Z",
          "iopub.execute_input": "2022-05-14T04:13:52.072114Z",
          "iopub.status.idle": "2022-05-14T04:13:52.088924Z",
          "shell.execute_reply.started": "2022-05-14T04:13:52.072072Z",
          "shell.execute_reply": "2022-05-14T04:13:52.087769Z"
        },
        "trusted": true,
        "id": "RDxoBJpk4KtL"
      },
      "execution_count": 55,
      "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": {
        "id": "Jiz638-w4KtM"
      }
    },
    {
      "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-14T04:13:52.091978Z",
          "iopub.execute_input": "2022-05-14T04:13:52.092286Z",
          "iopub.status.idle": "2022-05-14T04:13:52.123871Z",
          "shell.execute_reply.started": "2022-05-14T04:13:52.092206Z",
          "shell.execute_reply": "2022-05-14T04:13:52.123194Z"
        },
        "trusted": true,
        "id": "DazwsCCq4KtN"
      },
      "execution_count": 56,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Optimizer only for certain parameters in the model"
      ],
      "metadata": {
        "id": "PLNAbKQy4KtP"
      }
    },
    {
      "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-14T04:13:52.125664Z",
          "iopub.execute_input": "2022-05-14T04:13:52.126192Z",
          "iopub.status.idle": "2022-05-14T04:13:52.138376Z",
          "shell.execute_reply.started": "2022-05-14T04:13:52.126156Z",
          "shell.execute_reply": "2022-05-14T04:13:52.137704Z"
        },
        "trusted": true,
        "id": "5Riw8u_44KtR"
      },
      "execution_count": 57,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Main training code. I will be adding KFolds support soon!"
      ],
      "metadata": {
        "id": "hR_LBSvN4KtS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Training Code\n",
        "if __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 = BERTLargeModel().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_bert_large.bin\"\n",
        "    )"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-14T04:13:52.139816Z",
          "iopub.execute_input": "2022-05-14T04:13:52.140297Z"
        },
        "trusted": true,
        "id": "Gd56pe7L4KtT",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 205,
          "referenced_widgets": [
            "2a0858a074804aaa89c2b7484513529d",
            "174c839c1281482a94736e45876990b6",
            "c3d113798e4847eeb2c0e8af6905ce4a",
            "e9d2f7d74f8d485485f6cff8f992e87c",
            "f85b6eec2fed44a984e11ab233340b51",
            "4a66b1d2b77048228d9cb96279e5ab18",
            "e2b74c76fb0c4437a8511669b9ebb34d",
            "17da8b41943448669f1376949e708f75",
            "29a766a2de144665ae1e85901d76a867",
            "34e11419597e4538ab938911b2ca5e5d",
            "df1164bd238a499ea6f08461381ef164",
            "4d2a1f9598224094b5bfe4ca424b8b95",
            "13d840bb407c4784a1d14f29327a5159",
            "126d3e07731b4fed931a1293f3995ae4",
            "460274a114a84c11a86375986fd6cd5f",
            "07c5ffc9611d4ec485000d3800c0ca4f",
            "af624fa94e124d9382ad436eda96efd7",
            "171ef91b77f14df78c4ad319c8b565af",
            "58bc23b3eb974b06b94a12df6c68661c",
            "a6267a0a8bfd45b288b2d1c4e155bed7",
            "01d38445a3f840cdab15e20f3dc2f45d",
            "a9d965747ead4859a480d7f38324a866"
          ]
        },
        "outputId": "757e6fcd-fbb3-4251-d693-93b6ccd50e5c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[INFO] Using GPU: Tesla P100-PCIE-16GB\n",
            "\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading:   0%|          | 0.00/1.25G [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "2a0858a074804aaa89c2b7484513529d"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Some weights of the model checkpoint at bert-large-uncased were not used when initializing BertModel: ['cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.weight', 'cls.seq_relationship.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.bias', 'cls.predictions.transform.LayerNorm.weight']\n",
            "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
            "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "==================== Epoch: 1 / 2 ====================\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/8836 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "4d2a1f9598224094b5bfe4ca424b8b95"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Finish the logging run\n",
        "run.finish()"
      ],
      "metadata": {
        "trusted": true,
        "id": "oFqngeFY4KtU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<center>\n",
        "    <img src=\"https://img.shields.io/badge/Upvote-If%20you%20like%20my%20work-07b3c8?style=for-the-badge&logo=kaggle\">\n",
        "</center>"
      ],
      "metadata": {
        "id": "i616dn3h4KtU"
      }
    }
  ]
}