Pytorch_BERT_Huggingface_w_SAM_+_EMA.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Pytorch BERT Huggingface w SAM + EMA.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "309e1e4e33c04bc791a65a71f0bd7b2c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_1520628276e0420ba8b0e836548dff4f",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_8f04b2225493458bba4a9a59eb001e37",
              "IPY_MODEL_15675fed27e343c49011be18b4516cd2",
              "IPY_MODEL_ac71a9c9737943b1b59490cdc8a09472"
            ]
          }
        },
        "1520628276e0420ba8b0e836548dff4f": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "8f04b2225493458bba4a9a59eb001e37": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_a44389f762844940968d63127601526b",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": "Downloading: 100%",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_85932f6b4bca47359e5e6a30407e472c"
          }
        },
        "15675fed27e343c49011be18b4516cd2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_96e6d4ca69854ff48e2abeca7a44c5e9",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 480,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 480,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_4334531aedc8492face55cb8b3b8e185"
          }
        },
        "ac71a9c9737943b1b59490cdc8a09472": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_62477d74cd9a499aadc77b50df43d6e2",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 480/480 [00:00&lt;00:00, 6.70kB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_efc744b41f114eef9f4c5b42eceed6a7"
          }
        },
        "a44389f762844940968d63127601526b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "85932f6b4bca47359e5e6a30407e472c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "96e6d4ca69854ff48e2abeca7a44c5e9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "4334531aedc8492face55cb8b3b8e185": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "62477d74cd9a499aadc77b50df43d6e2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "efc744b41f114eef9f4c5b42eceed6a7": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "3e9c721872984948b270861f559ccf84": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_f1d9323151e54c7496f39c08883f4797",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_b2e4b1dfb4d44b0393d4ae64a560ccf7",
              "IPY_MODEL_d95e77b7a3304d6ba7d689b2fe3ada6a",
              "IPY_MODEL_a1972ab2ebab498ebb25c01308cd2bf9"
            ]
          }
        },
        "f1d9323151e54c7496f39c08883f4797": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "b2e4b1dfb4d44b0393d4ae64a560ccf7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_9013670cc49b4f75be80890541c07083",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": "Downloading: 100%",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_600107e7c9bc48fba13b92c8783f76ff"
          }
        },
        "d95e77b7a3304d6ba7d689b2fe3ada6a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_56d280e8dc594efd8b08ccec371e0488",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 898823,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 898823,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_0f8c9df34ed14281b0fd6ea9506c6c05"
          }
        },
        "a1972ab2ebab498ebb25c01308cd2bf9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_d1671866b6264ef8938be1fb7b497a01",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 878k/878k [00:00&lt;00:00, 3.47MB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_9012b7e1ffee465f885042cd94e56807"
          }
        },
        "9013670cc49b4f75be80890541c07083": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "600107e7c9bc48fba13b92c8783f76ff": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "56d280e8dc594efd8b08ccec371e0488": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "0f8c9df34ed14281b0fd6ea9506c6c05": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "d1671866b6264ef8938be1fb7b497a01": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "9012b7e1ffee465f885042cd94e56807": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "76468abe92464762aba93dabca915bb9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_bb94070c24b44c18a4d4b439f2d97518",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_c56d3e50931b4228a8f1be38534fcdd9",
              "IPY_MODEL_1b318c65e9d54ae4b4d16df358f3f414",
              "IPY_MODEL_325dfe6964c0432ea85b2ee0e5453bde"
            ]
          }
        },
        "bb94070c24b44c18a4d4b439f2d97518": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "c56d3e50931b4228a8f1be38534fcdd9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_a7ef59870ab345278d6588493db171b7",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": "Downloading: 100%",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_79440f16c0b3405286c32c57fc78bc4d"
          }
        },
        "1b318c65e9d54ae4b4d16df358f3f414": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_4ff90b4f306241bcacaff31c5e4ea3f9",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 456318,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 456318,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_682920010b854eca82ec0b4ae68ad520"
          }
        },
        "325dfe6964c0432ea85b2ee0e5453bde": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_08d502762b2448c586f8b473dfa29419",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 446k/446k [00:00&lt;00:00, 998kB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_e6ff457a77734929a48f5a9bc06cc00c"
          }
        },
        "a7ef59870ab345278d6588493db171b7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "79440f16c0b3405286c32c57fc78bc4d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "4ff90b4f306241bcacaff31c5e4ea3f9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "682920010b854eca82ec0b4ae68ad520": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "08d502762b2448c586f8b473dfa29419": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "e6ff457a77734929a48f5a9bc06cc00c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "2f3d456ccb0e4b58a9dfd440928e48dd": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_7cb60e89e93c4abf951a713829ea37a8",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_d0af7aa5354a4a79a431c07aacacf6f4",
              "IPY_MODEL_6b7574be6a2f4a45b1bc6f63ab654c52",
              "IPY_MODEL_bf76fee0ebae422eb8c574042da28740"
            ]
          }
        },
        "7cb60e89e93c4abf951a713829ea37a8": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "d0af7aa5354a4a79a431c07aacacf6f4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_142a1930fab74fd3b8884d62eca3f0b5",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": "Downloading: 100%",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_a5de16ceca0b483a847647f0c751fd80"
          }
        },
        "6b7574be6a2f4a45b1bc6f63ab654c52": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_f45611c167ba4b7e8cde71862e1ac096",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 1355863,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 1355863,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_4acfff4c7346452f95aa44584e6f8b76"
          }
        },
        "bf76fee0ebae422eb8c574042da28740": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_02a0d24699fb4c22b0266ba2ece578f4",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 1.29M/1.29M [00:00&lt;00:00, 5.38MB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_f883056c3c224a5698bcda4980ca6ea9"
          }
        },
        "142a1930fab74fd3b8884d62eca3f0b5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "a5de16ceca0b483a847647f0c751fd80": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "f45611c167ba4b7e8cde71862e1ac096": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "4acfff4c7346452f95aa44584e6f8b76": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "02a0d24699fb4c22b0266ba2ece578f4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "f883056c3c224a5698bcda4980ca6ea9": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "3422030378b749b39d7a4ea1b3198c89": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_642830d29d114b54ba9172be8f973347",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_029975b31b78474b8d889904b082fcab",
              "IPY_MODEL_8efdf37ce118451f834d04ef7600500e",
              "IPY_MODEL_f9e7745a931c4dac9b72ab8e06851177"
            ]
          }
        },
        "642830d29d114b54ba9172be8f973347": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "029975b31b78474b8d889904b082fcab": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_d982752ade25453d98dada6191618519",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": "Downloading: 100%",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_7d0e3ad75ed14a3090f690669d568248"
          }
        },
        "8efdf37ce118451f834d04ef7600500e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_ff7f3893495f43eb8f97a0bd4c003db2",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 331070498,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 331070498,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_70904926804144af99781021b4552800"
          }
        },
        "f9e7745a931c4dac9b72ab8e06851177": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_b6ac2679d0d2476eb1984354b6d287f5",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 316M/316M [00:21&lt;00:00, 33.2MB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_31b42bb5759746369648137d357d1c6b"
          }
        },
        "d982752ade25453d98dada6191618519": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "7d0e3ad75ed14a3090f690669d568248": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "ff7f3893495f43eb8f97a0bd4c003db2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "70904926804144af99781021b4552800": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "b6ac2679d0d2476eb1984354b6d287f5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "31b42bb5759746369648137d357d1c6b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/maxmatical/ml-cheatsheet/blob/master/Pytorch_BERT_Huggingface_w_SAM_%2B_EMA.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nzBFZ3cJG7AV"
      },
      "source": [
        "# Inspirations\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1o9pMmzzoq7e",
        "outputId": "d3aeee39-c34e-4fa4-811c-d1d2e84de12d"
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Sat Jan  1 15:51:45 2022       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 495.44       Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla K80           Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   30C    P8    27W / 149W |      0MiB / 11441MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5QkBmYey_BMu"
      },
      "source": [
        "%%capture\n",
        "!pip install transformers\n",
        "!pip install torchmetrics\n",
        "!pip install torch-ema\n",
        "!pip install koila"
      ],
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zDygGNnC_VGv"
      },
      "source": [
        "import math\n",
        "import time\n",
        "from typing import Optional, Tuple\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import functional as F\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torch.optim.lr_scheduler import OneCycleLR, LambdaLR, CosineAnnealingLR, ReduceLROnPlateau\n",
        "from torch.optim import AdamW\n",
        "\n",
        "from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModel\n",
        "  \n",
        "import torchmetrics\n",
        "from torchmetrics.classification import F1, Accuracy\n",
        "from torchmetrics.functional import accuracy, f1, auroc\n",
        "\n",
        "from sklearn.model_selection import train_test_split\n",
        "from sklearn.metrics import classification_report, multilabel_confusion_matrix\n",
        "\n",
        "from torch_ema import ExponentialMovingAverage"
      ],
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 252,
          "referenced_widgets": [
            "309e1e4e33c04bc791a65a71f0bd7b2c",
            "1520628276e0420ba8b0e836548dff4f",
            "8f04b2225493458bba4a9a59eb001e37",
            "15675fed27e343c49011be18b4516cd2",
            "ac71a9c9737943b1b59490cdc8a09472",
            "a44389f762844940968d63127601526b",
            "85932f6b4bca47359e5e6a30407e472c",
            "96e6d4ca69854ff48e2abeca7a44c5e9",
            "4334531aedc8492face55cb8b3b8e185",
            "62477d74cd9a499aadc77b50df43d6e2",
            "efc744b41f114eef9f4c5b42eceed6a7",
            "3e9c721872984948b270861f559ccf84",
            "f1d9323151e54c7496f39c08883f4797",
            "b2e4b1dfb4d44b0393d4ae64a560ccf7",
            "d95e77b7a3304d6ba7d689b2fe3ada6a",
            "a1972ab2ebab498ebb25c01308cd2bf9",
            "9013670cc49b4f75be80890541c07083",
            "600107e7c9bc48fba13b92c8783f76ff",
            "56d280e8dc594efd8b08ccec371e0488",
            "0f8c9df34ed14281b0fd6ea9506c6c05",
            "d1671866b6264ef8938be1fb7b497a01",
            "9012b7e1ffee465f885042cd94e56807",
            "76468abe92464762aba93dabca915bb9",
            "bb94070c24b44c18a4d4b439f2d97518",
            "c56d3e50931b4228a8f1be38534fcdd9",
            "1b318c65e9d54ae4b4d16df358f3f414",
            "325dfe6964c0432ea85b2ee0e5453bde",
            "a7ef59870ab345278d6588493db171b7",
            "79440f16c0b3405286c32c57fc78bc4d",
            "4ff90b4f306241bcacaff31c5e4ea3f9",
            "682920010b854eca82ec0b4ae68ad520",
            "08d502762b2448c586f8b473dfa29419",
            "e6ff457a77734929a48f5a9bc06cc00c",
            "2f3d456ccb0e4b58a9dfd440928e48dd",
            "7cb60e89e93c4abf951a713829ea37a8",
            "d0af7aa5354a4a79a431c07aacacf6f4",
            "6b7574be6a2f4a45b1bc6f63ab654c52",
            "bf76fee0ebae422eb8c574042da28740",
            "142a1930fab74fd3b8884d62eca3f0b5",
            "a5de16ceca0b483a847647f0c751fd80",
            "f45611c167ba4b7e8cde71862e1ac096",
            "4acfff4c7346452f95aa44584e6f8b76",
            "02a0d24699fb4c22b0266ba2ece578f4",
            "f883056c3c224a5698bcda4980ca6ea9",
            "3422030378b749b39d7a4ea1b3198c89",
            "642830d29d114b54ba9172be8f973347",
            "029975b31b78474b8d889904b082fcab",
            "8efdf37ce118451f834d04ef7600500e",
            "f9e7745a931c4dac9b72ab8e06851177",
            "d982752ade25453d98dada6191618519",
            "7d0e3ad75ed14a3090f690669d568248",
            "ff7f3893495f43eb8f97a0bd4c003db2",
            "70904926804144af99781021b4552800",
            "b6ac2679d0d2476eb1984354b6d287f5",
            "31b42bb5759746369648137d357d1c6b"
          ]
        },
        "id": "SWso8EH2IDcU",
        "outputId": "de6540f5-296c-401d-e777-8f70450450ed"
      },
      "source": [
        "# pretrained model and tokenizer\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"distilroberta-base\")\n",
        "model = AutoModel.from_pretrained(\"distilroberta-base\")"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "309e1e4e33c04bc791a65a71f0bd7b2c",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "Downloading:   0%|          | 0.00/480 [00:00<?, ?B/s]"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "3e9c721872984948b270861f559ccf84",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "Downloading:   0%|          | 0.00/878k [00:00<?, ?B/s]"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "76468abe92464762aba93dabca915bb9",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "Downloading:   0%|          | 0.00/446k [00:00<?, ?B/s]"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "2f3d456ccb0e4b58a9dfd440928e48dd",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "Downloading:   0%|          | 0.00/1.29M [00:00<?, ?B/s]"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "3422030378b749b39d7a4ea1b3198c89",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "Downloading:   0%|          | 0.00/316M [00:00<?, ?B/s]"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Some weights of the model checkpoint at distilroberta-base were not used when initializing RobertaModel: ['lm_head.layer_norm.bias', 'lm_head.dense.bias', 'lm_head.dense.weight', 'lm_head.layer_norm.weight', 'lm_head.decoder.weight', 'lm_head.bias']\n",
            "- This IS expected if you are initializing RobertaModel 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 RobertaModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4lqOT-HXGYIN"
      },
      "source": [
        "# Getting Data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7WLrlT3NGZGT",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "359b2fe7-c00d-4c1f-b403-f4ab19f2dbde"
      },
      "source": [
        "!gdown --id 1VuQ-U7TtggShMeuRSA_hzC8qGDl2LRkr"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=1VuQ-U7TtggShMeuRSA_hzC8qGDl2LRkr\n",
            "To: /content/toxic_comments.csv\n",
            "100% 68.8M/68.8M [00:00<00:00, 148MB/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xuywc9Y5Gbf1",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 206
        },
        "outputId": "9d3fb408-ede3-41a5-90af-7e3630b4d952"
      },
      "source": [
        "df = pd.read_csv(\"toxic_comments.csv\")\n",
        "\n",
        "df.head()"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/html": [
              "\n",
              "  <div id=\"df-089d73a4-dd00-4d57-9070-e2e01846958d\">\n",
              "    <div class=\"colab-df-container\">\n",
              "      <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>id</th>\n",
              "      <th>comment_text</th>\n",
              "      <th>toxic</th>\n",
              "      <th>severe_toxic</th>\n",
              "      <th>obscene</th>\n",
              "      <th>threat</th>\n",
              "      <th>insult</th>\n",
              "      <th>identity_hate</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>0000997932d777bf</td>\n",
              "      <td>Explanation\\nWhy the edits made under my usern...</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>000103f0d9cfb60f</td>\n",
              "      <td>D'aww! He matches this background colour I'm s...</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>000113f07ec002fd</td>\n",
              "      <td>Hey man, I'm really not trying to edit war. It...</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>0001b41b1c6bb37e</td>\n",
              "      <td>\"\\nMore\\nI can't make any real suggestions on ...</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>0001d958c54c6e35</td>\n",
              "      <td>You, sir, are my hero. Any chance you remember...</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "      <td>0</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-089d73a4-dd00-4d57-9070-e2e01846958d')\"\n",
              "              title=\"Convert this dataframe to an interactive table.\"\n",
              "              style=\"display:none;\">\n",
              "        \n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "       width=\"24px\">\n",
              "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
              "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
              "  </svg>\n",
              "      </button>\n",
              "      \n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      flex-wrap:wrap;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "      <script>\n",
              "        const buttonEl =\n",
              "          document.querySelector('#df-089d73a4-dd00-4d57-9070-e2e01846958d button.colab-df-convert');\n",
              "        buttonEl.style.display =\n",
              "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "        async function convertToInteractive(key) {\n",
              "          const element = document.querySelector('#df-089d73a4-dd00-4d57-9070-e2e01846958d');\n",
              "          const dataTable =\n",
              "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                     [key], {});\n",
              "          if (!dataTable) return;\n",
              "\n",
              "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "            + ' to learn more about interactive tables.';\n",
              "          element.innerHTML = '';\n",
              "          dataTable['output_type'] = 'display_data';\n",
              "          await google.colab.output.renderOutput(dataTable, element);\n",
              "          const docLink = document.createElement('div');\n",
              "          docLink.innerHTML = docLinkHtml;\n",
              "          element.appendChild(docLink);\n",
              "        }\n",
              "      </script>\n",
              "    </div>\n",
              "  </div>\n",
              "  "
            ],
            "text/plain": [
              "                 id  ... identity_hate\n",
              "0  0000997932d777bf  ...             0\n",
              "1  000103f0d9cfb60f  ...             0\n",
              "2  000113f07ec002fd  ...             0\n",
              "3  0001b41b1c6bb37e  ...             0\n",
              "4  0001d958c54c6e35  ...             0\n",
              "\n",
              "[5 rows x 8 columns]"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bDVf_wTaGgVP"
      },
      "source": [
        "train_df, val_df = train_test_split(df, test_size=0.15)"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3ro0J6WPGv69",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "970b2554-6d24-464a-a518-9603c5b022e5"
      },
      "source": [
        "# subsample clean comments\n",
        "LABEL_COLUMNS = df.columns.tolist()[2:]\n",
        "\n",
        "train_toxic = train_df[train_df[LABEL_COLUMNS].sum(axis=1) > 0]\n",
        "train_clean = train_df[train_df[LABEL_COLUMNS].sum(axis=1) == 0]\n",
        "\n",
        "train_df = pd.concat([\n",
        "  train_toxic,\n",
        "  train_clean.sample(15_000)\n",
        "])\n",
        "\n",
        "train_df.shape, val_df.shape"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((28866, 8), (23936, 8))"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EF3tWEZ74Vww",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "fc7f68ff-4f36-4fec-809c-91ebc5b38915"
      },
      "source": [
        "# take only a subsample of each train_df and val_df for faster iterations\n",
        "train_df = train_df.sample(500)\n",
        "val_df = val_df.sample(500)\n",
        "\n",
        "train_df.shape, val_df.shape"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((500, 8), (500, 8))"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dn-Q0BjUoB9Y"
      },
      "source": [
        "# Creating Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QV-deDgoo1Rx"
      },
      "source": [
        "# set batch size max seq_len\n",
        "bs = 12\n",
        "seq_len = 256"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RL6vWIO4HoO2"
      },
      "source": [
        "class ToxicCommentsDataset(Dataset):\n",
        "\n",
        "  def __init__(\n",
        "    self,\n",
        "    data: pd.DataFrame,\n",
        "    tokenizer: AutoTokenizer,\n",
        "    max_token_len: int = 128\n",
        "\n",
        "  ):\n",
        "    self.tokenizer = tokenizer\n",
        "    self.data = data\n",
        "    self.max_token_len = max_token_len\n",
        "\n",
        "  def __len__(self):\n",
        "    return len(self.data)\n",
        "\n",
        "  def __getitem__(self, index: int):\n",
        "    data_row = self.data.iloc[index]\n",
        "    comment_text = data_row.comment_text\n",
        "    labels = data_row[LABEL_COLUMNS]\n",
        "\n",
        "    encoding = self.tokenizer.encode_plus(\n",
        "      comment_text,\n",
        "      add_special_tokens=True,\n",
        "      max_length=self.max_token_len,\n",
        "      return_token_type_ids=False,\n",
        "      padding=\"max_length\",\n",
        "      truncation=True,\n",
        "      return_attention_mask=True,\n",
        "      return_tensors='pt',\n",
        "\n",
        "    )\n",
        "\n",
        "    return dict(\n",
        "      # comment_text=comment_text, # don't put text here\n",
        "      input_ids=encoding[\"input_ids\"].flatten(),\n",
        "      attention_mask=encoding[\"attention_mask\"].flatten(),\n",
        "      labels = torch.IntTensor(labels)\n",
        "    #   labels=torch.FloatTensor(labels)\n",
        "    )"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "validations"
      ],
      "metadata": {
        "id": "r_rFYeyYUCVJ"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mMPUJ8inH2Jg",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c249cfe8-fcb7-41ea-f231-88acf13f20c6"
      },
      "source": [
        "# test\n",
        "train_dataset = ToxicCommentsDataset(\n",
        "  train_df,\n",
        "  tokenizer,\n",
        "  max_token_len=seq_len\n",
        ")\n",
        "\n",
        "sample_item = train_dataset[0]\n",
        "sample_item.keys()"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dict_keys(['input_ids', 'attention_mask', 'labels'])"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Vvh5OltLILMa",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2b70dde9-d1c1-4a9d-acfd-6056af638463"
      },
      "source": [
        "print(sample_item[\"input_ids\"], sample_item[\"labels\"])\n",
        "print(sample_item[\"input_ids\"].shape)"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([    0, 43086,  8446,    40,  6009,    23,    24,   456,   650, 25128,\n",
            "          856,  1073, 14118,  2489,   298, 43555, 43555, 43555,     2,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1,     1,     1,     1,     1,\n",
            "            1,     1,     1,     1,     1,     1]) tensor([1, 1, 1, 0, 1, 0], dtype=torch.int32)\n",
            "torch.Size([256])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# test label smoothing\n",
        "labels = sample_item[\"labels\"]\n",
        "print(labels)\n",
        "\n",
        "alpha = 0.1\n",
        "labels_smoothed = labels.to(dtype=torch.float32) * (1-alpha) + alpha * (1-labels.to(dtype=torch.float32))\n",
        "print(labels_smoothed)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pzZ0WNMWO8Lz",
        "outputId": "3bf51627-0d03-49c3-ddfb-5ab4a366438e"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([1, 1, 1, 0, 1, 0], dtype=torch.int32)\n",
            "tensor([0.9000, 0.9000, 0.9000, 0.1000, 0.9000, 0.1000])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# data sets\n",
        "train_dataset = ToxicCommentsDataset(\n",
        "    train_df,\n",
        "    tokenizer,\n",
        "    max_token_len=seq_len\n",
        ")\n",
        "\n",
        "val_dataset = ToxicCommentsDataset(\n",
        "    val_df,\n",
        "    tokenizer,\n",
        "    max_token_len=seq_len\n",
        ")\n",
        "\n"
      ],
      "metadata": {
        "id": "CJ7oQpcqDytP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Creating Dataloaders\n",
        "\n",
        "Using `koila` to determine max batch size that fits in gpu mem"
      ],
      "metadata": {
        "id": "QiEhWaF1UOXZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "train_dl = DataLoader(\n",
        "    train_dataset,\n",
        "    batch_size=bs,\n",
        "    shuffle=True,\n",
        "    num_workers=1\n",
        ")\n",
        "\n",
        "val_dl = DataLoader(\n",
        "    val_dataset,\n",
        "    batch_size=bs,\n",
        "    shuffle=False,\n",
        "    num_workers=1\n",
        ")\n",
        "\n",
        "print(len(train_dl), len(val_dl))"
      ],
      "metadata": {
        "id": "vUwZ1vyrUSeB",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "048ce7fb-1091-4464-d305-960a7f958438"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "42 42\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D1zQOd2AJl7l"
      },
      "source": [
        "# Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ad0Y-3mUXmRM"
      },
      "source": [
        "class BertModel(nn.Module):\n",
        "  def __init__(self, n_classes: int):\n",
        "    super().__init__()\n",
        "    self.model = model\n",
        "    self.classifier = nn.Linear(self.model.config.hidden_size, n_classes)\n",
        "\n",
        "  def forward(self, input_ids, attention_mask): \n",
        "    out = self.model(input_ids, attention_mask=attention_mask)\n",
        "    out = self.classifier(out.pooler_output)\n",
        "    return out"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7Eb5FtVEYWQQ"
      },
      "source": [
        "bert_model = BertModel(len(LABEL_COLUMNS))\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Validate model with dataloader"
      ],
      "metadata": {
        "id": "WPNBxWEGUcvS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sample_batch = next(iter(DataLoader(train_dataset, batch_size=bs, num_workers=1)))\n",
        "sample_batch[\"input_ids\"].shape, sample_batch[\"attention_mask\"].shape"
      ],
      "metadata": {
        "id": "4fC7NY5nUdYF",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5a4ce8ca-827c-4bb5-d474-0bf108a575b0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(torch.Size([12, 256]), torch.Size([12, 256]))"
            ]
          },
          "metadata": {},
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "bert_model(sample_batch[\"input_ids\"], sample_batch[\"attention_mask\"]).shape  # should be bs x 6"
      ],
      "metadata": {
        "id": "SXa015uyUhaL",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ee526c76-0bbb-437c-e505-e9b9cb381810"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([12, 6])"
            ]
          },
          "metadata": {},
          "execution_count": 22
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mwG-to8c4pUV"
      },
      "source": [
        "# Training enhancements"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## SAM optimizer"
      ],
      "metadata": {
        "id": "FZUmJgRYVy7a"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6pVZgHoG4qVE"
      },
      "source": [
        "class SAM(torch.optim.Optimizer):\n",
        "    def __init__(self, params, base_optimizer, rho=0.05, adaptive=False, **kwargs):\n",
        "        assert rho >= 0.0, f\"Invalid rho, should be non-negative: {rho}\"\n",
        "\n",
        "        defaults = dict(rho=rho, adaptive=adaptive, **kwargs)\n",
        "        super(SAM, self).__init__(params, defaults)\n",
        "\n",
        "        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)\n",
        "        self.param_groups = self.base_optimizer.param_groups\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def first_step(self, zero_grad=False):\n",
        "        grad_norm = self._grad_norm()\n",
        "        for group in self.param_groups:\n",
        "            scale = group[\"rho\"] / (grad_norm + 1e-12)\n",
        "\n",
        "            for p in group[\"params\"]:\n",
        "                if p.grad is None: continue\n",
        "                self.state[p][\"old_p\"] = p.data.clone()\n",
        "                e_w = (torch.pow(p, 2) if group[\"adaptive\"] else 1.0) * p.grad * scale.to(p)\n",
        "                p.add_(e_w)  # climb to the local maximum \"w + e(w)\"\n",
        "\n",
        "        if zero_grad: self.zero_grad()\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def second_step(self, zero_grad=False):\n",
        "        for group in self.param_groups:\n",
        "            for p in group[\"params\"]:\n",
        "                if p.grad is None: continue\n",
        "                p.data = self.state[p][\"old_p\"]  # get back to \"w\" from \"w + e(w)\"\n",
        "\n",
        "        self.base_optimizer.step()  # do the actual \"sharpness-aware\" update\n",
        "\n",
        "        if zero_grad: self.zero_grad()\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def step(self, closure=None):\n",
        "        assert closure is not None, \"Sharpness Aware Minimization requires closure, but it was not provided\"\n",
        "        closure = torch.enable_grad()(closure)  # the closure should do a full forward-backward pass\n",
        "\n",
        "        self.first_step(zero_grad=True)\n",
        "        closure()\n",
        "        self.second_step()\n",
        "\n",
        "    def _grad_norm(self):\n",
        "        shared_device = self.param_groups[0][\"params\"][0].device  # put everything on the same device, in case of model parallelism\n",
        "        norm = torch.norm(\n",
        "                    torch.stack([\n",
        "                        ((torch.abs(p) if group[\"adaptive\"] else 1.0) * p.grad).norm(p=2).to(shared_device)\n",
        "                        for group in self.param_groups for p in group[\"params\"]\n",
        "                        if p.grad is not None\n",
        "                    ]),\n",
        "                    p=2\n",
        "               )\n",
        "        return norm\n",
        "\n",
        "    def load_state_dict(self, state_dict):\n",
        "        super().load_state_dict(state_dict)\n",
        "        self.base_optimizer.param_groups = self.param_groups"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## model tracking (loss and metrics)"
      ],
      "metadata": {
        "id": "guIwYjJEV1aA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# keep running average of loss values\n",
        "class AverageMeter(object):\n",
        "    \"\"\"Computes and stores the average and current value\"\"\"\n",
        "    def __init__(self, name, fmt=':f'):\n",
        "        self.name = name\n",
        "        self.fmt = fmt\n",
        "        self.reset()\n",
        "\n",
        "    def reset(self):\n",
        "        self.val = 0\n",
        "        self.avg = 0\n",
        "        self.sum = 0\n",
        "        self.count = 0\n",
        "\n",
        "    def update(self, val, n=1):\n",
        "        self.val = val\n",
        "        self.sum += val * n\n",
        "        self.count += n\n",
        "        self.avg = self.sum / self.count\n",
        "\n",
        "    def __str__(self):\n",
        "        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'\n",
        "        return fmtstr.format(**self.__dict__)"
      ],
      "metadata": {
        "id": "orJ2pdoTUcSq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## LR scheduler"
      ],
      "metadata": {
        "id": "U--W43FFV5bq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# flat_cos scheduler\n",
        "def d(x): \n",
        "  \"\"\"\n",
        "  dummy function\n",
        "  \"\"\"\n",
        "  return 1\n",
        "    \n",
        "class ConcatLR(torch.optim.lr_scheduler._LRScheduler):\n",
        "    def __init__(self, optimizer, scheduler1, scheduler2, total_steps, pct_start=0.5, last_epoch=-1):\n",
        "        self.scheduler1 = scheduler1\n",
        "        self.scheduler2 = scheduler2\n",
        "        self.step_start = float(pct_start * total_steps) - 1\n",
        "        super(ConcatLR, self).__init__(optimizer, last_epoch)\n",
        "    \n",
        "    def step(self):\n",
        "        if self.last_epoch <= self.step_start:\n",
        "            self.scheduler1.step()\n",
        "        else:\n",
        "            self.scheduler2.step()\n",
        "        super().step()\n",
        "        \n",
        "    def get_lr(self):\n",
        "        if self.last_epoch <= self.step_start:\n",
        "            return self.scheduler1.get_lr()\n",
        "        else:\n",
        "            return self.scheduler2.get_lr()"
      ],
      "metadata": {
        "id": "9l6De166Hq8t"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Helper functions for configuring optimizers and lr scheduler"
      ],
      "metadata": {
        "id": "DjJj-2n9V9BD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# helper funcs\n",
        "\n",
        "def configure_optimizer(\n",
        "    model: nn.Module, \n",
        "    lr: float = 2e-5, \n",
        "    betas: Tuple[float, float] = (0.9, 0.999),\n",
        "    eps: float = 1e-8,\n",
        "    wd: float = 0.01,\n",
        "    use_sam: bool = False, \n",
        "    rho: float = 0.05, \n",
        "    asam: bool = False\n",
        "):\n",
        "  if not use_sam:\n",
        "    optimizer = AdamW(model.parameters(), lr=lr, betas=betas, eps=eps, weight_decay=wd)\n",
        "  else:\n",
        "    base_optimizer = AdamW\n",
        "    optimizer = SAM(\n",
        "        model.parameters(), \n",
        "        base_optimizer=base_optimizer, \n",
        "        lr=lr, \n",
        "        betas=betas,\n",
        "        weight_decay=wd,\n",
        "        eps=eps,\n",
        "        rho=rho,\n",
        "        adaptive=asam\n",
        "    )\n",
        "  return optimizer \n",
        "\n",
        "def configure_scheduler(\n",
        "    fit_func: str, \n",
        "    lr: float, \n",
        "    total_steps: int, \n",
        "    optimizer, \n",
        "    pct_start: float = 0.3, \n",
        "    use_sam: bool = False\n",
        "):\n",
        "  assert fit_func in [\"one_cycle\", \"flat_cos\"], f\"fit function {fit_func} not found\"\n",
        "  # if using sam, lr scheduler is on base optimizer\n",
        "  opt = optimizer.base_optimizer if use_sam else optimizer\n",
        "  if fit_func == \"one_cycle\":\n",
        "    scheduler = OneCycleLR(\n",
        "        optimizer=opt,\n",
        "        max_lr=lr,\n",
        "        pct_start=pct_start,\n",
        "        total_steps=total_steps\n",
        "      )\n",
        "  elif fit_func == \"flat_cos\":\n",
        "    dummy = LambdaLR(opt, d)\n",
        "    cosine = CosineAnnealingLR(opt, total_steps*(1-pct_start))\n",
        "    scheduler = ConcatLR(opt, dummy, cosine, total_steps, pct_start)\n",
        "\n",
        "  else:\n",
        "    raise ValueError(f\"fit_func {fit_func} not found\")\n",
        "\n",
        "  return scheduler"
      ],
      "metadata": {
        "id": "pbC9aoZpF38e"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Model saver"
      ],
      "metadata": {
        "id": "nAbyOKvjWAok"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class ModelSaver:\n",
        "  def __init__(self, save_path: str, mode: str = \"max\"):\n",
        "    \"\"\"\n",
        "    class used for saving models during training\n",
        "    \"\"\"\n",
        "    self.save_path = save_path\n",
        "    self.mode = mode\n",
        "    assert self.mode in [\"min\", \"max\"], f\"mode {mode} not found\"\n",
        "    # self.best_value = torch.tensor(float(\"inf\")) if mode == \"min\" else torch.tensor(float(\"-inf\"))\n",
        "    self.best_value = float(\"inf\") if mode == \"min\" else float(\"-inf\")\n",
        "\n",
        "  def save_model(self, epoch: int, model: nn.Module, current_value: float):\n",
        "    \"\"\"\n",
        "    compares current_value with self.best_value\n",
        "    if current_value is better then\n",
        "    1. save model weights to self.save_path\n",
        "    2. update self.best_value with current_value\n",
        "    \"\"\"\n",
        "    if (\n",
        "      (self.mode == \"min\" and current_value <= self.best_value)\n",
        "      or (self.mode == \"max\" and current_value >= self.best_value)\n",
        "    ):\n",
        "      torch.save(model.state_dict(), self.save_path)\n",
        "      print(f\"better model found at epoch {epoch} with value {current_value}\")\n",
        "      print(f\"model weights saved to '{self.save_path}', to load model weights, create new model and use new_model.load_state_dict(torch.load('{self.save_path}'))\")\n",
        "      self.best_value = current_value\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "_D_4Qe12omp3"
      },
      "execution_count": 57,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# test\n",
        "\n",
        "model_saver = ModelSaver(save_path=\"test.pth\", mode=\"max\")\n",
        "\n",
        "model_saver.save_model(0, model=bert_model, current_value=5)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Vq_B5tRGYLYZ",
        "outputId": "f07deb6f-e2ce-4f80-aace-c4aba2e5a28c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "better model found at epoch 0 with value 5\n",
            "model weights saved to 'test.pth', to load model weights, create new model and use new_model.load_state_dict(torch.load(test.pth))\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## LabelSmoothingBCE"
      ],
      "metadata": {
        "id": "K6vztCxIUJVG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class LabelSmoothingBCEWithLogitsLoss(nn.Module):\n",
        "  \"\"\" \n",
        "  BCE with label smoothing \n",
        "  \"\"\"\n",
        "  def __init__(\n",
        "          self, smoothing=0.1, target_threshold: Optional[float] = None, weight: Optional[torch.Tensor] = None,\n",
        "          reduction: str = 'mean', pos_weight: Optional[torch.Tensor] = None):\n",
        "      super(LabelSmoothingBCEWithLogitsLoss, self).__init__()\n",
        "      assert 0. <= smoothing < 1.0\n",
        "      self.smoothing = smoothing\n",
        "      self.target_threshold = target_threshold\n",
        "      self.reduction = reduction\n",
        "      self.register_buffer('weight', weight)\n",
        "      self.register_buffer('pos_weight', pos_weight)\n",
        "\n",
        "  def forward(self, x: torch.Tensor, target: torch.Tensor) -> torch.Tensor:\n",
        "      target = target * (1 - self.smoothing) + self.smoothing * (1 - target)\n",
        "      return F.binary_cross_entropy_with_logits(\n",
        "          x, target,\n",
        "          self.weight,\n",
        "          pos_weight=self.pos_weight,\n",
        "          reduction=self.reduction)\n"
      ],
      "metadata": {
        "id": "jwHsJF2sUL47"
      },
      "execution_count": 55,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "loss = nn.BCEWithLogitsLoss()\n",
        "smoothed_loss = LabelSmoothingBCEWithLogitsLoss(smoothing=0.1)\n",
        "loss(labels.to(dtype=torch.float32), labels.to(dtype=torch.float32)), smoothed_loss(labels.to(dtype=torch.float32), labels.to(dtype=torch.float32))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XyvvE3yMWC9q",
        "outputId": "3b3411d0-495d-4273-9eee-bb3557012bf5"
      },
      "execution_count": 56,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(tensor(0.4399), tensor(0.5066))"
            ]
          },
          "metadata": {},
          "execution_count": 56
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training Model"
      ],
      "metadata": {
        "id": "dsrERC6XEk8j"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "FP16 with SAM:\n",
        "\n",
        "discussion thread: https://github.com/davda54/sam/issues/7\n",
        "\n",
        "should be something like:\n",
        "```\n",
        "# first pass\n",
        "with torch.cuda.amp.autocast():\n",
        "    out = model(input_ids, attention_mask)\n",
        "    loss = criterion(out, labels.to(dtype=torch.float32))\n",
        "\n",
        "scaler.scale(loss).backward()\n",
        "scaler.unscale_(optimizer)\n",
        "optimizer.first_step(zero_grad=True)\n",
        "scaler.update()\n",
        "\n",
        "# 2nd pass\n",
        "with torch.cuda.amp.autocast():\n",
        "    out_2 = model(input_ids, attention_mask)\n",
        "    loss_2 = criterion(out_2, labels.to(dtype=torch.float32))\n",
        "\n",
        "scaler.scale(loss_2).backward()\n",
        "scaler.unscale_(optimizer)\n",
        "optimizer.second_step(zero_grad=True)\n",
        "scaler.update()\n",
        "```\n",
        "\n",
        "Gradient accumulation with SAM: https://github.com/davda54/sam/issues/3\n",
        "\n",
        "4 different types of training modes:\n",
        "1. FP32\n",
        "2. SAM\n",
        "3. FP16\n",
        "4. FP16 + SAM\n",
        "\n",
        "TODO:\n",
        "- [x] Saving best model (with/without EMA) based on measured metric\n",
        "- Early stopping? based on metric\n",
        "- [x] `ReduceLROnPlateau` configurable\n",
        "- [x] terminate training on NaN\n",
        "- [x] add gradient clipping (default value 1.0)\n",
        "- [x] label smoothing on bce loss (CE loss already has natic label smoothing)"
      ],
      "metadata": {
        "id": "6WDV26cuFs3s"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def train(\n",
        "    model: nn.Module,\n",
        "    train_dl: DataLoader,\n",
        "    val_dl: DataLoader,\n",
        "    n_epochs: int,\n",
        "    lr: float = 2e-5,\n",
        "    b1: float = 0.9,\n",
        "    b2: float = 0.999,\n",
        "    eps: float = 1e-7,\n",
        "    wd: float = 0.01,\n",
        "    max_norm: float=1.0,\n",
        "    use_sam: bool = False,\n",
        "    rho: float = 0.05,\n",
        "    asam: bool = False,\n",
        "    fit_func: str = \"one_cycle\",\n",
        "    pct_start: float = 0.3,\n",
        "    fp16: bool = False,\n",
        "    is_ddp: bool = False,\n",
        "    n_gpus: int = 1,\n",
        "    use_ema: bool = False,\n",
        "    ema_decay: float = 0.995,\n",
        "    grad_acc_batches: int = 1,\n",
        "    model_saver: Optional[ModelSaver] = None,\n",
        "    reduce_lr_on_plateau: bool = False,\n",
        "    reduce_lr_on_plateau_mode: str = \"min\",\n",
        "    reduce_lr_on_plateau_patience: int = 10,\n",
        "    reduce_lr_on_plateau_factor: float = 0.1,\n",
        "    terminate_on_nan: bool = False\n",
        "):\n",
        "  \"\"\"\n",
        "  args:\n",
        "      model: nn.Module the model to train\n",
        "      train_dl: DataLoader, train dataloader\n",
        "      val_dl: DataLoader, validation dataloader\n",
        "      n_epochs: int, number of epochs to train\n",
        "      lr: float = 2e-5, learning rate\n",
        "      b1: float = 0.9, beta1 for optimizers like adam, ranger etc.\n",
        "      b2: float = 0.999, beta2 for optimizers like adam, ranger etc.\n",
        "      eps: float = 1e-7, eps for optimizers (set to >1e-7 for fp16)\n",
        "      wd: float = 0.01, weight decay regularization\n",
        "      use_sam: bool = False, whether to use SAM with optimizer as base optimizer\n",
        "      rho: float = 0.05, neighborhood size for SAM (set to 10x larger for ASAM)\n",
        "      asam: bool = False, whether to use ASAM variant of SAM\n",
        "      fit_func: str = \"one_cycle\", what type of training, one_cycle or flat_cos\n",
        "      pct_start: float = 0.3, pct to start cosine annealing for fit func\n",
        "      fp16: bool = False, whether to use mixed precision training with AMP\n",
        "      is_ddp: bool = False, whether to use DDP (when n_gpus >1)\n",
        "      n_gpus: int = 1, number of gpus used for training\n",
        "      use_ema: bool = False, whether to use EMA to average model weights\n",
        "      ema_decay: float = 0.995, decay factor for EMA\n",
        "      grad_acc_batches: int = 1, number of batches to accumulate for gradient accumulation\n",
        "      model_saver: Optional[ModelSaver] = None, saving models\n",
        "      reduce_lr_on_plateau: bool = False, whether to use ReduceLROnPleateau scheduler\n",
        "      reduce_lr_on_plateau_patience: int = 10,\n",
        "      reduce_lr_on_plateau_factor: float = 0.1, factor to reduce lr by\n",
        "    \n",
        "  \"\"\"\n",
        "  # set device\n",
        "  device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "  # set model to device if haven't already done so\n",
        "  model.to(device)\n",
        "\n",
        "  # set loss\n",
        "  criterion = nn.BCEWithLogitsLoss()\n",
        "\n",
        "  # configure optimizer \n",
        "  optimizer = configure_optimizer(model, lr=lr, betas=(b1, b2), eps=eps, wd=wd, use_sam=use_sam, rho=rho, asam=asam)\n",
        "\n",
        "  # configure lr scheduler\n",
        "  total_steps = int(n_epochs * len(train_dl) / n_gpus / grad_acc_batches)\n",
        "  lr_schedule = configure_scheduler(fit_func, lr, total_steps, optimizer, pct_start, use_sam)\n",
        "\n",
        "  # configure ema\n",
        "  if use_ema:\n",
        "    ema = ExponentialMovingAverage(model.parameters(), decay=ema_decay)\n",
        "\n",
        "  # configure scaler for fp16\n",
        "  scaler = torch.cuda.amp.GradScaler() if fp16 else None\n",
        "\n",
        "  # used for SAM\n",
        "  input_list, attn_mask_list, labels_list = [], [], []\n",
        "\n",
        "  # reduce lr on plateau\n",
        "  reduce_lr_on_plateau_scheduler = ReduceLROnPlateau(\n",
        "      optimizer,\n",
        "      mode=reduce_lr_on_plateau_mode,\n",
        "      patience=reduce_lr_on_plateau_patience,\n",
        "      factor=reduce_lr_on_plateau_factor,\n",
        "      verbose=True\n",
        "  ) if reduce_lr_on_plateau else None\n",
        "  # for testing purposes\n",
        "  # lrs = []\n",
        "\n",
        "  # start training\n",
        "  for epoch in range(n_epochs):\n",
        "    #############################################\n",
        "    # training step\n",
        "    #############################################\n",
        "    model.train()\n",
        "    # initialize tracking variables\n",
        "    start = time.time()\n",
        "    losses = AverageMeter('Loss', ':.4e')\n",
        "    f1_metrics = AverageMeter(\"F1\", ':6.2f')\n",
        "\n",
        "    for batch_idx, batch in enumerate(train_dl):\n",
        "      input_ids = batch[\"input_ids\"].to(device)\n",
        "      attention_mask = batch[\"attention_mask\"].to(device)\n",
        "      labels = batch[\"labels\"].to(device)\n",
        "\n",
        "      # with fp16\n",
        "      if fp16 and not use_sam:\n",
        "        with torch.cuda.amp.autocast():\n",
        "          out = model(input_ids, attention_mask)\n",
        "          loss = criterion(out, labels.to(dtype=torch.float32))\n",
        "          # scale the loss by gradient accumulation batches\n",
        "          loss = loss / grad_acc_batches\n",
        "        # backwards step  \n",
        "        scaler.scale(loss).backward()\n",
        "        # optimizer\n",
        "        if (batch_idx + 1) % grad_acc_batches == 0:\n",
        "          scaler.unscale_(optimizer)\n",
        "          torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)\n",
        "          scaler.step(optimizer)\n",
        "          scaler.update()\n",
        "          optimizer.zero_grad()\n",
        "\n",
        "      # with fp16 + sam\n",
        "      elif fp16 and use_sam:\n",
        "        with torch.cuda.amp.autocast():\n",
        "          out = model(input_ids, attention_mask)\n",
        "\n",
        "          # save input and output for 2nd step\n",
        "          input_list.append(input_ids)\n",
        "          attn_mask_list.append(attention_mask)\n",
        "          labels_list.append(labels)\n",
        "\n",
        "          loss = criterion(out, labels.to(dtype=torch.float32))\n",
        "          # scale the loss by gradient accumulation batches\n",
        "          loss = loss / grad_acc_batches\n",
        "        # backwards step  \n",
        "        scaler.scale(loss).backward()\n",
        "        # optimizer step\n",
        "        if (batch_idx + 1) % grad_acc_batches == 0:\n",
        "          scaler.unscale_(optimizer)\n",
        "          torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)\n",
        "          optimizer.first_step(zero_grad=True)\n",
        "          scaler.update()\n",
        "\n",
        "          # 2nd forward pass with saved input_list and labels_list\n",
        "          # to get the accumulated gradients again\n",
        "          for (input_ids, attention_mask, labels) in list(zip(input_list, attn_mask_list, labels_list)):\n",
        "            with torch.cuda.amp.autocast():\n",
        "              out_2 = model(input_ids, attention_mask)\n",
        "              loss_2 = criterion(out_2, labels.to(dtype=torch.float32))\n",
        "              loss_2 = loss_2 / grad_acc_batches\n",
        "            # 2nd backwards step  \n",
        "            scaler.scale(loss_2).backward()\n",
        "\n",
        "          # 2nd optimizer step (outside the for loop)\n",
        "          scaler.unscale_(optimizer)\n",
        "          torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)\n",
        "          optimizer.second_step(zero_grad=True)\n",
        "          scaler.update()\n",
        "\n",
        "          # clear saved lists\n",
        "          input_list, attn_mask_list, labels_list = [], [], []\n",
        "      \n",
        "      # with fp32 + sam\n",
        "      elif not fp16 and use_sam:\n",
        "        # save input and output for 2nd step\n",
        "        input_list.append(input_ids)\n",
        "        attn_mask_list.append(attention_mask)\n",
        "        labels_list.append(labels)\n",
        "\n",
        "        out = model(input_ids, attention_mask)\n",
        "        loss = criterion(out, labels.to(dtype=torch.float32))\n",
        "        # scale the loss by gradient accumulation batches\n",
        "        loss = loss / grad_acc_batches\n",
        "        loss.backward()\n",
        "        \n",
        "        if (batch_idx + 1) % grad_acc_batches == 0:\n",
        "          torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)\n",
        "          optimizer.first_step(zero_grad=True)\n",
        "\n",
        "          # 2nd step\n",
        "          for (input_ids, attention_mask, labels) in list(zip(input_list, attn_mask_list, labels_list)):\n",
        "            out_2 = model(input_ids, attention_mask)\n",
        "            loss_2 = criterion(out_2, labels.to(dtype=torch.float32))\n",
        "            loss_2 = loss_2 / grad_acc_batches\n",
        "            loss_2.backward()\n",
        "          torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)\n",
        "          optimizer.second_step(zero_grad=True)\n",
        "          input_list, attn_mask_list, labels_list = [], [], []\n",
        "\n",
        "\n",
        "      # without fp16 or sam\n",
        "      else:\n",
        "        out = model(input_ids, attention_mask)\n",
        "        loss = criterion(out, labels.to(dtype=torch.float32))\n",
        "        # scale the loss by gradient accumulation batches\n",
        "        loss = loss / grad_acc_batches\n",
        "        # backwards step\n",
        "        loss.backward()\n",
        "        if (batch_idx + 1) % grad_acc_batches == 0:\n",
        "          torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)\n",
        "          optimizer.step()\n",
        "          optimizer.zero_grad()\n",
        "\n",
        "      # terminate training if nan encountered\n",
        "      if terminate_on_nan and torch.isnan(loss).item():\n",
        "        print(\"NaN encountered in training loss, terminating training.\")\n",
        "        break\n",
        "      \n",
        "      # log loss and metrics \n",
        "      losses.update(loss.item() * grad_acc_batches)  # scale loss back up by grad_acc_batches\n",
        "      f1_metrics.update(f1(torch.sigmoid(out), labels))\n",
        "\n",
        "      # update EMA and lr schedule\n",
        "      if (batch_idx + 1) % grad_acc_batches == 0:\n",
        "        # update ema\n",
        "        if use_ema:\n",
        "          ema.update()\n",
        "        lr_schedule.step()\n",
        "        # for tracking lr schedule\n",
        "        # lrs.append(optimizer.param_groups[0][\"lr\"])\n",
        "\n",
        "    #############################################\n",
        "    # validation step\n",
        "    #############################################\n",
        "    model.eval()\n",
        "    # initialize val variables\n",
        "    val_losses = AverageMeter('Loss', ':.4e')\n",
        "    val_f1_metrics = AverageMeter(\"F1\", ':6.2f')\n",
        "\n",
        "    with torch.no_grad():\n",
        "      for batch_idx, batch in enumerate(val_dl):\n",
        "        input_ids = batch[\"input_ids\"].to(device)\n",
        "        attention_mask = batch[\"attention_mask\"].to(device)\n",
        "        labels = batch[\"labels\"].to(device)\n",
        "\n",
        "        # eval with fp16\n",
        "        if fp16:\n",
        "          with torch.cuda.amp.autocast():\n",
        "            # use EMA\n",
        "            if use_ema:\n",
        "              with ema.average_parameters():\n",
        "                out = model(input_ids, attention_mask)\n",
        "            else:\n",
        "              out = model(input_ids, attention_mask)\n",
        "          \n",
        "        else:\n",
        "          # use EMA\n",
        "          if use_ema:\n",
        "            with ema.average_parameters():\n",
        "              out = model(input_ids, attention_mask)\n",
        "          else:\n",
        "            out = model(input_ids, attention_mask)\n",
        "\n",
        "        # calculate val loss and val_f1\n",
        "        loss = criterion(out, labels.to(dtype=torch.float32))\n",
        "        val_losses.update(loss.item())\n",
        "        val_f1_metrics.update(f1(torch.sigmoid(out), labels))\n",
        "\n",
        "    #############################################\n",
        "    # end of epoch\n",
        "    # logging training stats\n",
        "    # reduce_lr_on_plateau_scheduler.step() if is not None\n",
        "    # model saving step\n",
        "    #############################################\n",
        "\n",
        "    if model_saver:\n",
        "      if use_ema:\n",
        "        with ema.average_parameters():\n",
        "          model_saver.save_model(epoch+1, model, val_f1_metrics.avg)    \n",
        "      else:\n",
        "        model_saver.save_model(epoch+1, model, val_f1_metrics.avg)\n",
        "\n",
        "    if reduce_lr_on_plateau_scheduler:\n",
        "      reduce_lr_on_plateau_scheduler.step(val_f1_metrics.avg)\n",
        "    \n",
        "    end = time.time()\n",
        "    elapsed = end - start\n",
        "\n",
        "    # log relevant metrics at the end of epoch\n",
        "    print(f\"Epoch {epoch+1}: train loss: {losses.avg}, val loss: {val_losses.avg}, training f1: {f1_metrics.avg}, val f1: {val_f1_metrics.avg}, time: {elapsed}\")\n",
        "\n",
        "  # # saving model weights\n",
        "  # saved_model_pth = \"saved_model.pth\"\n",
        "\n",
        "  # # if using ema weights, copy those weights to model before saving\n",
        "  # if use_ema:\n",
        "  #   ema.copy_to(model.parameters())\n",
        "  # torch.save(model.state_dict(), saved_model_pth)\n",
        "  # print(f\"model weights saved to '{saved_model_pth}', to load model weights, create new model and use new_model.load_state_dict(torch.load('{saved_model_pth}'))\")\n",
        "  # # print(lrs)\n",
        "\n"
      ],
      "metadata": {
        "id": "gWk0p2uhEmLH"
      },
      "execution_count": 58,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "del bert_model\n",
        "bert_model = BertModel(len(LABEL_COLUMNS))"
      ],
      "metadata": {
        "id": "2WWEAN9uYrUm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "use_sam = False\n",
        "fp16 = False\n",
        "use_ema = True\n",
        "grad_acc_batches = 10\n",
        "model_saver = ModelSaver(save_path=\"test.pth\", mode=\"max\")\n",
        "\n",
        "train(\n",
        "    bert_model,\n",
        "    train_dl,\n",
        "    val_dl,\n",
        "    n_epochs=2,\n",
        "    fp16=fp16,\n",
        "    use_sam=use_sam,\n",
        "    use_ema=use_ema,\n",
        "    grad_acc_batches=grad_acc_batches,\n",
        "    model_saver=model_saver\n",
        ")"
      ],
      "metadata": {
        "id": "cGgbPcIyRxy5",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e98f3172-746c-4602-9b91-a6c52c1f5d4f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "before\n",
            "Parameter containing:\n",
            "tensor([[ 0.0104,  0.0750,  0.0364,  ..., -0.0239,  0.0762,  0.0593],\n",
            "        [-0.0002,  0.0319, -0.0069,  ...,  0.0129,  0.0892,  0.0641],\n",
            "        [ 0.0423,  0.0649, -0.0370,  ..., -0.0719, -0.0194,  0.1182],\n",
            "        [ 0.0335,  0.0879, -0.0925,  ..., -0.0433,  0.0145,  0.0385],\n",
            "        [-0.0205, -0.0736,  0.0050,  ...,  0.0224,  0.0882, -0.1193],\n",
            "        [ 0.0527,  0.0204, -0.0368,  ..., -0.0582,  0.0490,  0.0696]],\n",
            "       device='cuda:0', requires_grad=True)\n",
            "after\n",
            "Parameter containing:\n",
            "tensor([[ 0.0120,  0.0550,  0.0349,  ..., -0.0254,  0.0690,  0.0393],\n",
            "        [ 0.0018,  0.0281, -0.0090,  ...,  0.0109,  0.0740,  0.0603],\n",
            "        [ 0.0283,  0.0393, -0.0230,  ..., -0.0579, -0.0090,  0.0926],\n",
            "        [ 0.0270,  0.0793, -0.0860,  ..., -0.0368,  0.0075,  0.0299],\n",
            "        [-0.0265, -0.0490,  0.0109,  ...,  0.0284,  0.0875, -0.0947],\n",
            "        [ 0.0483,  0.0132, -0.0324,  ..., -0.0538,  0.0398,  0.0624]],\n",
            "       device='cuda:0', requires_grad=True)\n",
            "better model found at epoch 1 with value 0.0624995082616806\n",
            "model weights saved to 'test.pth', to load model weights, create new model and use new_model.load_state_dict(torch.load(test.pth))\n",
            "Epoch 1: train loss: 2.025878010317683, val loss: 4.2797179108574275, training f1: 0.2666379511356354, val f1: 0.0624995082616806, time: 38.91955280303955\n",
            "before\n",
            "Parameter containing:\n",
            "tensor([[ 0.0400,  0.0771,  0.0068,  ..., -0.0534,  0.1340,  0.0615],\n",
            "        [ 0.0492, -0.0174, -0.0564,  ..., -0.0365,  0.1481,  0.0148],\n",
            "        [ 0.0962,  0.0711, -0.0909,  ..., -0.1258,  0.0025,  0.1243],\n",
            "        [ 0.0853,  0.0527, -0.1442,  ..., -0.0951,  0.0685,  0.0033],\n",
            "        [-0.0082, -0.1190, -0.0074,  ...,  0.0100,  0.0460, -0.1646],\n",
            "        [ 0.1037, -0.0205, -0.0878,  ..., -0.1092,  0.1059,  0.0286]],\n",
            "       device='cuda:0', requires_grad=True)\n",
            "after\n",
            "Parameter containing:\n",
            "tensor([[ 0.0375,  0.0764,  0.0093,  ..., -0.0510,  0.1290,  0.0607],\n",
            "        [ 0.0454, -0.0136, -0.0525,  ..., -0.0327,  0.1428,  0.0186],\n",
            "        [ 0.0917,  0.0692, -0.0864,  ..., -0.1212,  0.0015,  0.1224],\n",
            "        [ 0.0808,  0.0551, -0.1398,  ..., -0.0906,  0.0639,  0.0057],\n",
            "        [-0.0091, -0.1147, -0.0065,  ...,  0.0109,  0.0497, -0.1603],\n",
            "        [ 0.0994, -0.0176, -0.0835,  ..., -0.1049,  0.1010,  0.0316]],\n",
            "       device='cuda:0', requires_grad=True)\n",
            "Epoch 2: train loss: 3.885560929775238, val loss: 0.5240822042382899, training f1: 0.18613344430923462, val f1: 0.0, time: 37.85096597671509\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# test model saver with EMA verified to work\n",
        "bert_model_2 = BertModel(len(LABEL_COLUMNS))\n",
        "bert_model_2.load_state_dict(torch.load(\"test.pth\"))\n",
        "\n",
        "print(bert_model.classifier.weight)\n",
        "\n",
        "print(bert_model_2.classifier.weight)\n",
        "\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uy6SKtH1Yn22",
        "outputId": "fc5ba30b-1cde-4e22-e4d8-9b50f668d58a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Parameter containing:\n",
            "tensor([[ 0.0400,  0.0771,  0.0068,  ..., -0.0534,  0.1340,  0.0615],\n",
            "        [ 0.0492, -0.0174, -0.0564,  ..., -0.0365,  0.1481,  0.0148],\n",
            "        [ 0.0962,  0.0711, -0.0909,  ..., -0.1258,  0.0025,  0.1243],\n",
            "        [ 0.0853,  0.0527, -0.1442,  ..., -0.0951,  0.0685,  0.0033],\n",
            "        [-0.0082, -0.1190, -0.0074,  ...,  0.0100,  0.0460, -0.1646],\n",
            "        [ 0.1037, -0.0205, -0.0878,  ..., -0.1092,  0.1059,  0.0286]],\n",
            "       device='cuda:0', requires_grad=True)\n",
            "Parameter containing:\n",
            "tensor([[ 0.0120,  0.0550,  0.0349,  ..., -0.0254,  0.0690,  0.0393],\n",
            "        [ 0.0018,  0.0281, -0.0090,  ...,  0.0109,  0.0740,  0.0603],\n",
            "        [ 0.0283,  0.0393, -0.0230,  ..., -0.0579, -0.0090,  0.0926],\n",
            "        [ 0.0270,  0.0793, -0.0860,  ..., -0.0368,  0.0075,  0.0299],\n",
            "        [-0.0265, -0.0490,  0.0109,  ...,  0.0284,  0.0875, -0.0947],\n",
            "        [ 0.0483,  0.0132, -0.0324,  ..., -0.0538,  0.0398,  0.0624]],\n",
            "       requires_grad=True)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Some results\n",
        "trained for 1 epoch only, gradient accumulation 12 batches (144 bs)\n",
        "\n",
        "\n",
        "baseline (no fp16, sam, ema):\n",
        "```\n",
        "Epoch 1: train loss: 0.6239880334053721, val loss: 0.5045745234404292, training f1: 0.4051700830459595, val f1: 0.3108885884284973, time: 55.422086238861084\n",
        "```\n",
        "\n",
        "fp16:\n",
        "```\n",
        "Epoch 1: train loss: 0.5987142457493714, val loss: 0.47237819823480787, training f1: 0.27484646439552307, val f1: 0.261325478553772, time: 110.72927689552307\n",
        "```\n",
        "sam:\n",
        "```\n",
        "Epoch 1: train loss: 0.46754554366426815, val loss: 0.3434885683513823, training f1: 0.4873872697353363, val f1: 0.3218401372432709, time: 93.9434859752655\n",
        "```\n",
        "\n",
        "ema:\n",
        "```\n",
        "Epoch 1: train loss: 0.5785652372453894, val loss: 0.47438419610261917, training f1: 0.5454299449920654, val f1: 0.41533538699150085, time: 56.51772093772888\n",
        "```\n",
        "\n",
        "\n",
        "sam + fp16\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "fp16 + ema\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "\n",
        "sam+ema\n",
        "```\n",
        "Epoch 1: train loss: 0.6115093609052045, val loss: 0.516309067606926, training f1: 0.3987903594970703, val f1: 0.2628635764122009, time: 94.63389563560486\n",
        "```\n",
        "\n",
        "fp16 + sam + ema\n",
        "```\n",
        "Epoch 1: train loss: 0.5853038478110517, val loss: 0.47891736775636673, training f1: 0.39333027601242065, val f1: 0.1278771311044693, time: 192.29889297485352\n",
        "```\n"
      ],
      "metadata": {
        "id": "p5XEVCDRHzWN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "device = \"cuda\"\n",
        "sample_batch = next(iter(DataLoader(val_dataset, batch_size=bs, num_workers=1)))\n",
        "input_ids = sample_batch[\"input_ids\"].to(device)\n",
        "attention_mask = sample_batch[\"attention_mask\"].to(device)\n",
        "labels = sample_batch[\"labels\"].to(device)\n",
        "\n",
        "out = bert_model(input_ids, attention_mask)\n",
        "\n",
        "print(f1(torch.sigmoid(out), labels))\n",
        "print(torch.sigmoid(out))\n",
        "print(labels)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "csyQbh_2BRcH",
        "outputId": "ef62f1f9-96c5-423b-d18e-1bd6d7489124"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor(0., device='cuda:0')\n",
            "tensor([[0.3604, 0.0749, 0.1837, 0.0692, 0.1929, 0.0806],\n",
            "        [0.3803, 0.0752, 0.1921, 0.0698, 0.2088, 0.0787],\n",
            "        [0.3515, 0.0763, 0.1836, 0.0712, 0.1901, 0.0817],\n",
            "        [0.3820, 0.0757, 0.1940, 0.0667, 0.1922, 0.0804],\n",
            "        [0.4409, 0.0791, 0.2198, 0.0593, 0.2180, 0.0834],\n",
            "        [0.3872, 0.0754, 0.1918, 0.0665, 0.1922, 0.0803],\n",
            "        [0.3724, 0.0750, 0.1772, 0.0717, 0.1873, 0.0847],\n",
            "        [0.3428, 0.0758, 0.1795, 0.0735, 0.1838, 0.0833],\n",
            "        [0.3711, 0.0760, 0.1883, 0.0677, 0.1894, 0.0823],\n",
            "        [0.3813, 0.0759, 0.1940, 0.0691, 0.1950, 0.0782],\n",
            "        [0.4030, 0.0771, 0.2016, 0.0631, 0.1977, 0.0796],\n",
            "        [0.3810, 0.0765, 0.1917, 0.0652, 0.1971, 0.0807]], device='cuda:0',\n",
            "       grad_fn=<SigmoidBackward0>)\n",
            "tensor([[0, 0, 0, 0, 0, 0],\n",
            "        [0, 0, 0, 0, 0, 0],\n",
            "        [0, 0, 0, 0, 0, 0],\n",
            "        [0, 0, 0, 0, 0, 0],\n",
            "        [0, 0, 0, 0, 0, 0],\n",
            "        [0, 0, 0, 0, 0, 0],\n",
            "        [0, 0, 0, 0, 0, 0],\n",
            "        [0, 0, 0, 0, 0, 0],\n",
            "        [0, 0, 0, 0, 0, 0],\n",
            "        [0, 0, 0, 0, 0, 0],\n",
            "        [0, 0, 1, 0, 1, 0],\n",
            "        [0, 0, 0, 0, 0, 0]], device='cuda:0', dtype=torch.int32)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(f1(F.sigmoid(out), labels))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8aWrc6z7DBrv",
        "outputId": "eae2339f-e564-405d-93eb-22f332535396"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor(0.4444, device='cuda:0')\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py:1806: UserWarning: nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\n",
            "  warnings.warn(\"nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\")\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z_Qa6mk9bB-V"
      },
      "source": [
        "# Saving and Exporting model"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# if using ema weights, copy those weights to model before saving\n",
        "bert1 = BertModel(len(LABEL_COLUMNS))\n",
        "print(\"bert1\")\n",
        "print(bert1.classifier.weight)\n",
        "ema = ExponentialMovingAverage(bert1.parameters(), decay=0.995)\n",
        "bert2 = BertModel(len(LABEL_COLUMNS))\n",
        "print(\"bert2\")\n",
        "print(bert2.classifier.weight)\n",
        "\n",
        "\n",
        "if use_ema:\n",
        "  ema.copy_to(bert2.parameters())\n",
        "\n",
        "print(\"bert2 after ema copy weight\")\n",
        "print(bert2.classifier.weight)\n",
        "\n",
        "\n",
        "# torch.save(bert_model.state_dict(), saved_model_pth)"
      ],
      "metadata": {
        "id": "R7TSQhy0SJVj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aSX1sSa1hsMI"
      },
      "source": [
        "new_bert_model = BertModel(6)\n",
        "\n",
        "new_bert_model.load_state_dict(torch.load(saved_model_pth))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ytqP_Olebu64"
      },
      "source": [
        "# check values are the same\n",
        "\n",
        "inf_preds = new_bert_model(sample_batch[\"input_ids\"], sample_batch[\"attention_mask\"])\n",
        "print(inf_preds.shape) # should be 8 x 6\n",
        "inf_preds"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_GEsu9KydC93"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}