Pytorch_Lightning_BERT_Huggingface_w_SAM.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Pytorch Lightning BERT Huggingface w SAM.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": {
        "a8a32ce0327042569dc04173a441b408": {
          "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_9d9004c8434744b7addca2a7ae0e7af8",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_429fafd9c7bb4ef59baf13127ca15882",
              "IPY_MODEL_19d4885702d6409b928508637925c8e8",
              "IPY_MODEL_f65e7965b2b949e8acc6d887b7dd19e1"
            ]
          }
        },
        "9d9004c8434744b7addca2a7ae0e7af8": {
          "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
          }
        },
        "429fafd9c7bb4ef59baf13127ca15882": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_6233f14fdb7c4f39875c8638b3622b08",
            "_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_9da8011ee9ff4ffab8ac7c55b5188e1f"
          }
        },
        "19d4885702d6409b928508637925c8e8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_47b251b9688046348683a0e6c7e7609b",
            "_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_853556db1eb54f2cb04ef9038b109b97"
          }
        },
        "f65e7965b2b949e8acc6d887b7dd19e1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_3a24c29a9ba84fb1a9b219b8c88eed65",
            "_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, 7.37kB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_3af3716d879242d3b6f1aa47a0ac8c14"
          }
        },
        "6233f14fdb7c4f39875c8638b3622b08": {
          "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"
          }
        },
        "9da8011ee9ff4ffab8ac7c55b5188e1f": {
          "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
          }
        },
        "47b251b9688046348683a0e6c7e7609b": {
          "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"
          }
        },
        "853556db1eb54f2cb04ef9038b109b97": {
          "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
          }
        },
        "3a24c29a9ba84fb1a9b219b8c88eed65": {
          "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"
          }
        },
        "3af3716d879242d3b6f1aa47a0ac8c14": {
          "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
          }
        },
        "e076db28e270432abac53fb2455a7302": {
          "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_9d27a36da3d44ef384f666c8d6037955",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_9757ea20f6d84163b7a7958258f152d1",
              "IPY_MODEL_07b1b21bc5eb46e6b698e06c69cd2cb5",
              "IPY_MODEL_1d5c59e03b014bde8ecdc341454b9c1e"
            ]
          }
        },
        "9d27a36da3d44ef384f666c8d6037955": {
          "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
          }
        },
        "9757ea20f6d84163b7a7958258f152d1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_e7cda65dbcb349aebf5725abe2070cc3",
            "_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_b5d5fbda9acf4384b1738bc163de9c4c"
          }
        },
        "07b1b21bc5eb46e6b698e06c69cd2cb5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_39143ee069214e0796238d67a27cea23",
            "_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_3d26e328034b45d99cb9533111a25d09"
          }
        },
        "1d5c59e03b014bde8ecdc341454b9c1e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_384eefcb754f44779616ae9ee521d021",
            "_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, 855kB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_8a1ca69360f0432ba11b14a74dc1de7a"
          }
        },
        "e7cda65dbcb349aebf5725abe2070cc3": {
          "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"
          }
        },
        "b5d5fbda9acf4384b1738bc163de9c4c": {
          "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
          }
        },
        "39143ee069214e0796238d67a27cea23": {
          "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"
          }
        },
        "3d26e328034b45d99cb9533111a25d09": {
          "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
          }
        },
        "384eefcb754f44779616ae9ee521d021": {
          "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"
          }
        },
        "8a1ca69360f0432ba11b14a74dc1de7a": {
          "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
          }
        },
        "3137906a7b314838903216175f285e23": {
          "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_6aeb2dfd08bb4d42b908d64d77748b6a",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_703cedc024334cfe873e389b79c2ebe6",
              "IPY_MODEL_5063bc7095194f16828b543bff0bbd80",
              "IPY_MODEL_fd97baa873af4e18bde0545e13b3359b"
            ]
          }
        },
        "6aeb2dfd08bb4d42b908d64d77748b6a": {
          "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
          }
        },
        "703cedc024334cfe873e389b79c2ebe6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_5cd837ed8ecb4da9bfa5eb2d188cb22f",
            "_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_04aebe5f69324b64a79c9771f8983b9a"
          }
        },
        "5063bc7095194f16828b543bff0bbd80": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_700270f1885c4fb3a34924b883a996f7",
            "_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_9cdc707479064582a0961d1a648db1c1"
          }
        },
        "fd97baa873af4e18bde0545e13b3359b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_f345edf149494215a4d777309c576cca",
            "_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, 796kB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_275dcab49b8f4af7b81e18d0b979bd6f"
          }
        },
        "5cd837ed8ecb4da9bfa5eb2d188cb22f": {
          "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"
          }
        },
        "04aebe5f69324b64a79c9771f8983b9a": {
          "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
          }
        },
        "700270f1885c4fb3a34924b883a996f7": {
          "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"
          }
        },
        "9cdc707479064582a0961d1a648db1c1": {
          "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
          }
        },
        "f345edf149494215a4d777309c576cca": {
          "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"
          }
        },
        "275dcab49b8f4af7b81e18d0b979bd6f": {
          "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
          }
        },
        "ec4591cde3634a038361c336e214b0f8": {
          "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_8ba157478f274a769c4d23891cd0b32e",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_f5b72f65de104c9e89db3227404a6110",
              "IPY_MODEL_67e3b810c66948b8ad06d3ca969cc481",
              "IPY_MODEL_0838f83f77f945fca8e15ab2da99a616"
            ]
          }
        },
        "8ba157478f274a769c4d23891cd0b32e": {
          "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
          }
        },
        "f5b72f65de104c9e89db3227404a6110": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_d71be4f0fa6145baa6ed12a7cba8a9c1",
            "_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_18637366266e472d81897fc148405a32"
          }
        },
        "67e3b810c66948b8ad06d3ca969cc481": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_8d3e924ef6d743648fd5f891b647bd94",
            "_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_48bc56d52b0d43b194cb43bba423da6e"
          }
        },
        "0838f83f77f945fca8e15ab2da99a616": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_a690b4d97ec74f5d9d2e9baca9719da4",
            "_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.12MB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_1134e0d0633e486899c30173060ad980"
          }
        },
        "d71be4f0fa6145baa6ed12a7cba8a9c1": {
          "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"
          }
        },
        "18637366266e472d81897fc148405a32": {
          "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
          }
        },
        "8d3e924ef6d743648fd5f891b647bd94": {
          "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"
          }
        },
        "48bc56d52b0d43b194cb43bba423da6e": {
          "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
          }
        },
        "a690b4d97ec74f5d9d2e9baca9719da4": {
          "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"
          }
        },
        "1134e0d0633e486899c30173060ad980": {
          "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
          }
        },
        "8aa7b5874ba8411b92da1441104117e8": {
          "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_af86948d2f99452c939b68072894f69e",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_7f33aaaf0e4743959dd3f55df6904e58",
              "IPY_MODEL_2e1b8f796d3b439894c5966a2c24efd5",
              "IPY_MODEL_5900d2869dcf43e0ab934cf7622d8c2b"
            ]
          }
        },
        "af86948d2f99452c939b68072894f69e": {
          "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
          }
        },
        "7f33aaaf0e4743959dd3f55df6904e58": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_f8d73a69ccba4da789fe61f1f65e40b0",
            "_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_582aa898370343779eea8b004a07b697"
          }
        },
        "2e1b8f796d3b439894c5966a2c24efd5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_925f7ad5f06849df84680009dc1321ad",
            "_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_848f85580cb14902b83c296d86eefee7"
          }
        },
        "5900d2869dcf43e0ab934cf7622d8c2b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_5e0e5450b8db4fcabae4540843ba248b",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 316M/316M [00:11&lt;00:00, 28.0MB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_70a592b047ec4abebcc6d4455c6f8e7c"
          }
        },
        "f8d73a69ccba4da789fe61f1f65e40b0": {
          "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"
          }
        },
        "582aa898370343779eea8b004a07b697": {
          "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
          }
        },
        "925f7ad5f06849df84680009dc1321ad": {
          "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"
          }
        },
        "848f85580cb14902b83c296d86eefee7": {
          "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
          }
        },
        "5e0e5450b8db4fcabae4540843ba248b": {
          "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"
          }
        },
        "70a592b047ec4abebcc6d4455c6f8e7c": {
          "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
          }
        },
        "7aea1bf3de1e43ef890f7c1aad966b04": {
          "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_c32531cc302b4a59884a5c85519a8957",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_24f6d3e223d44068b0701daafaa07f62",
              "IPY_MODEL_f58951b5683d4554a2df7f6bd74b3730",
              "IPY_MODEL_39efe0ead32c4d90b5d48dcf1a490b1b"
            ]
          }
        },
        "c32531cc302b4a59884a5c85519a8957": {
          "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": "row wrap",
            "width": "100%",
            "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": "inline-flex",
            "left": null
          }
        },
        "24f6d3e223d44068b0701daafaa07f62": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_74800a8f29c2457e8553540b8aedde44",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": "Validation sanity check:   0%",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_39acd4da787a4bf5baa8d468e21f7715"
          }
        },
        "f58951b5683d4554a2df7f6bd74b3730": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_af33f16f86504041b1b2f340b5a4e182",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "danger",
            "max": 2,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 0,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_1527d28c976f44468a0b791cf81be271"
          }
        },
        "39efe0ead32c4d90b5d48dcf1a490b1b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_0bc850aa8d81406491b930c54ee81a09",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 0/2 [00:00&lt;?, ?it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_c9538fe593ec481c878fb5a8c58f7850"
          }
        },
        "74800a8f29c2457e8553540b8aedde44": {
          "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"
          }
        },
        "39acd4da787a4bf5baa8d468e21f7715": {
          "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
          }
        },
        "af33f16f86504041b1b2f340b5a4e182": {
          "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"
          }
        },
        "1527d28c976f44468a0b791cf81be271": {
          "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": "2",
            "_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
          }
        },
        "0bc850aa8d81406491b930c54ee81a09": {
          "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"
          }
        },
        "c9538fe593ec481c878fb5a8c58f7850": {
          "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
          }
        },
        "c758b12eef464d3089f32fce401a65c9": {
          "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_8cd4f69584754f7fa029f4971d65fa91",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_d668bbda43964950a0eec05c4bfa5da1",
              "IPY_MODEL_1dcd44a8e2574c7d9ff1b891259f4c9c",
              "IPY_MODEL_91dbcff7475f46c3ba6e6a585414a2bf"
            ]
          }
        },
        "8cd4f69584754f7fa029f4971d65fa91": {
          "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": "row wrap",
            "width": "100%",
            "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": "inline-flex",
            "left": null
          }
        },
        "d668bbda43964950a0eec05c4bfa5da1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_d9b6105e50a447e5b9c896a53a32a6a4",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": "Epoch 0:   0%",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_b926716c578646a3a77045a3194427ec"
          }
        },
        "1dcd44a8e2574c7d9ff1b891259f4c9c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_ec262b68cdb44ba1b2e9d6bda4d0a83b",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "",
            "max": 1668,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 0,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_a9796fae209d4cebb13011450172e701"
          }
        },
        "91dbcff7475f46c3ba6e6a585414a2bf": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_400d011bb4d242dc95d6792bb17c3dfd",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 0/1668 [00:00&lt;?, ?it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_7257193458d94e3d9c977a636b042139"
          }
        },
        "d9b6105e50a447e5b9c896a53a32a6a4": {
          "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"
          }
        },
        "b926716c578646a3a77045a3194427ec": {
          "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
          }
        },
        "ec262b68cdb44ba1b2e9d6bda4d0a83b": {
          "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"
          }
        },
        "a9796fae209d4cebb13011450172e701": {
          "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": "2",
            "_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
          }
        },
        "400d011bb4d242dc95d6792bb17c3dfd": {
          "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"
          }
        },
        "7257193458d94e3d9c977a636b042139": {
          "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_Lightning_BERT_Huggingface_w_SAM.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",
        "https://curiousily.com/posts/multi-label-text-classification-with-bert-and-pytorch-lightning/\n",
        "\n",
        "https://github.com/mgrankin/over9000/blob/master/train.py\n",
        "\n",
        "\n",
        "https://medium.com/pytorch/getting-started-with-ray-lightning-easy-multi-node-pytorch-lightning-training-e639031aff8b\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1o9pMmzzoq7e",
        "outputId": "bfde5e04-71ea-4bba-e13d-33b5d0b095be"
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Sun Nov 28 18:16:08 2021       \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   70C    P8    35W / 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 pytorch_lightning\n",
        "!pip install torchmetrics\n",
        "!pip install torch-ema"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zDygGNnC_VGv"
      },
      "source": [
        "import pandas as pd\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torch.optim.lr_scheduler import OneCycleLR # , ReduceLROnPlateau\n",
        "from torch.optim import AdamW\n",
        "\n",
        "from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModel\n",
        "  \n",
        "import pytorch_lightning as pl\n",
        "from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping, LearningRateMonitor, StochasticWeightAveraging\n",
        "from pytorch_lightning.loggers import TensorBoardLogger\n",
        "import torchmetrics\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"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 252,
          "referenced_widgets": [
            "a8a32ce0327042569dc04173a441b408",
            "9d9004c8434744b7addca2a7ae0e7af8",
            "429fafd9c7bb4ef59baf13127ca15882",
            "19d4885702d6409b928508637925c8e8",
            "f65e7965b2b949e8acc6d887b7dd19e1",
            "6233f14fdb7c4f39875c8638b3622b08",
            "9da8011ee9ff4ffab8ac7c55b5188e1f",
            "47b251b9688046348683a0e6c7e7609b",
            "853556db1eb54f2cb04ef9038b109b97",
            "3a24c29a9ba84fb1a9b219b8c88eed65",
            "3af3716d879242d3b6f1aa47a0ac8c14",
            "e076db28e270432abac53fb2455a7302",
            "9d27a36da3d44ef384f666c8d6037955",
            "9757ea20f6d84163b7a7958258f152d1",
            "07b1b21bc5eb46e6b698e06c69cd2cb5",
            "1d5c59e03b014bde8ecdc341454b9c1e",
            "e7cda65dbcb349aebf5725abe2070cc3",
            "b5d5fbda9acf4384b1738bc163de9c4c",
            "39143ee069214e0796238d67a27cea23",
            "3d26e328034b45d99cb9533111a25d09",
            "384eefcb754f44779616ae9ee521d021",
            "8a1ca69360f0432ba11b14a74dc1de7a",
            "3137906a7b314838903216175f285e23",
            "6aeb2dfd08bb4d42b908d64d77748b6a",
            "703cedc024334cfe873e389b79c2ebe6",
            "5063bc7095194f16828b543bff0bbd80",
            "fd97baa873af4e18bde0545e13b3359b",
            "5cd837ed8ecb4da9bfa5eb2d188cb22f",
            "04aebe5f69324b64a79c9771f8983b9a",
            "700270f1885c4fb3a34924b883a996f7",
            "9cdc707479064582a0961d1a648db1c1",
            "f345edf149494215a4d777309c576cca",
            "275dcab49b8f4af7b81e18d0b979bd6f",
            "ec4591cde3634a038361c336e214b0f8",
            "8ba157478f274a769c4d23891cd0b32e",
            "f5b72f65de104c9e89db3227404a6110",
            "67e3b810c66948b8ad06d3ca969cc481",
            "0838f83f77f945fca8e15ab2da99a616",
            "d71be4f0fa6145baa6ed12a7cba8a9c1",
            "18637366266e472d81897fc148405a32",
            "8d3e924ef6d743648fd5f891b647bd94",
            "48bc56d52b0d43b194cb43bba423da6e",
            "a690b4d97ec74f5d9d2e9baca9719da4",
            "1134e0d0633e486899c30173060ad980",
            "8aa7b5874ba8411b92da1441104117e8",
            "af86948d2f99452c939b68072894f69e",
            "7f33aaaf0e4743959dd3f55df6904e58",
            "2e1b8f796d3b439894c5966a2c24efd5",
            "5900d2869dcf43e0ab934cf7622d8c2b",
            "f8d73a69ccba4da789fe61f1f65e40b0",
            "582aa898370343779eea8b004a07b697",
            "925f7ad5f06849df84680009dc1321ad",
            "848f85580cb14902b83c296d86eefee7",
            "5e0e5450b8db4fcabae4540843ba248b",
            "70a592b047ec4abebcc6d4455c6f8e7c"
          ]
        },
        "id": "SWso8EH2IDcU",
        "outputId": "e4ce72cc-60b7-41e5-b206-3e66dbca4327"
      },
      "source": [
        "# pretrained model and tokenizer\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"distilroberta-base\")\n",
        "model = AutoModel.from_pretrained(\"distilroberta-base\")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "a8a32ce0327042569dc04173a441b408",
              "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": "e076db28e270432abac53fb2455a7302",
              "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": "3137906a7b314838903216175f285e23",
              "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": "ec4591cde3634a038361c336e214b0f8",
              "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": "8aa7b5874ba8411b92da1441104117e8",
              "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.bias', 'lm_head.layer_norm.weight', 'lm_head.layer_norm.bias', 'lm_head.dense.bias', 'lm_head.decoder.weight', 'lm_head.dense.weight']\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": [
        "# Data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7WLrlT3NGZGT",
        "outputId": "452ec1a5-2a59-4e85-983e-63a7ad70d2fc"
      },
      "source": [
        "!gdown --id 1VuQ-U7TtggShMeuRSA_hzC8qGDl2LRkr"
      ],
      "execution_count": null,
      "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:01<00:00, 67.7MB/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 206
        },
        "id": "xuywc9Y5Gbf1",
        "outputId": "79f71b08-fb15-4e95-d387-f0174372028e"
      },
      "source": [
        "df = pd.read_csv(\"toxic_comments.csv\")\n",
        "\n",
        "df.head()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/html": [
              "<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>"
            ],
            "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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3ro0J6WPGv69",
        "outputId": "3dc60ff5-931d-41c2-9d70-988e9c51be3a"
      },
      "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": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((28744, 8), (23936, 8))"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EF3tWEZ74Vww",
        "outputId": "48a48e7f-9a95-40a5-a134-dd0eac42432b"
      },
      "source": [
        "# take only a subsample of each train_df and val_df for faster iterations\n",
        "train_df = train_df.sample(10000)\n",
        "val_df = val_df.sample(10000)\n",
        "\n",
        "train_df.shape, val_df.shape"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((10000, 8), (10000, 8))"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dn-Q0BjUoB9Y"
      },
      "source": [
        "## Creating Dataset and Lightning Data Module"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QV-deDgoo1Rx"
      },
      "source": [
        "# set batch size and max seq_len\n",
        "bs = 12\n",
        "seq_len = 256"
      ],
      "execution_count": null,
      "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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mMPUJ8inH2Jg",
        "outputId": "f85cd218-a549-4431-a74a-d95ff9354db0"
      },
      "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": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dict_keys(['input_ids', 'attention_mask', 'labels'])"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Vvh5OltLILMa",
        "outputId": "a047c183-424a-4f52-f85b-3f7408deffe0"
      },
      "source": [
        "print(sample_item[\"input_ids\"], sample_item[\"labels\"])\n",
        "print(sample_item[\"input_ids\"].shape)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([    0,   113, 50140,  3401,   109,    45, 24604,  2072,  6052,     6,\n",
            "           25,    47,   222,    19,    42, 17668,     7,   587,   501,     4,\n",
            "          318,    47,   535,     7,   109,    98,     6,    47,    40,    28,\n",
            "         4953,    31,  5390,     4,  1437, 43614,  1437, 46927, 23171,   298,\n",
            "        44721, 47264, 43662, 20024,   139, 43662, 15264,   139, 45704,  5782,\n",
            "        46927, 23171, 44721, 41335, 18400, 47264,   116,    22,     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]) tensor([0, 0, 0, 0, 0, 0], dtype=torch.int32)\n",
            "torch.Size([256])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "k0gIaph4IkSu"
      },
      "source": [
        "class ToxicCommentsDataModule(pl.LightningDataModule):\n",
        "  def __init__(self, train_df, test_df, tokenizer, batch_size=8, max_token_len=512):\n",
        "    super().__init__()\n",
        "    self.train_df, self.test_df = train_df, test_df\n",
        "    self.tokenizer = tokenizer\n",
        "    self.batch_size = batch_size\n",
        "    self.max_token_len = max_token_len\n",
        "\n",
        "  def setup(self, stage=None):\n",
        "    self.train_dataset = ToxicCommentsDataset(\n",
        "        self.train_df,\n",
        "        self.tokenizer,\n",
        "        self.max_token_len\n",
        "    )\n",
        "\n",
        "    self.test_dataset = ToxicCommentsDataset(\n",
        "        self.test_df,\n",
        "        self.tokenizer,\n",
        "        self.max_token_len\n",
        "    )\n",
        "\n",
        "  def train_dataloader(self):\n",
        "    return DataLoader(\n",
        "        self.train_dataset,\n",
        "        batch_size = self.batch_size,\n",
        "        shuffle = True,\n",
        "        num_workers=1\n",
        "    )\n",
        "\n",
        "  def val_dataloader(self):\n",
        "    return DataLoader(\n",
        "        self.test_dataset,\n",
        "        batch_size = self.batch_size,\n",
        "        shuffle = False,\n",
        "        num_workers=1\n",
        "    )\n",
        "\n",
        "  def test_dataloader(self):\n",
        "    return DataLoader(\n",
        "        self.test_dataset,\n",
        "        batch_size = self.batch_size,\n",
        "        shuffle = False,\n",
        "        num_workers=1\n",
        "    )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fwxqBy-FIjf5"
      },
      "source": [
        "data_module = ToxicCommentsDataModule(\n",
        "    train_df,\n",
        "    val_df,\n",
        "    tokenizer,\n",
        "    batch_size = bs,\n",
        "    max_token_len = seq_len\n",
        ")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8zCGQwkFqZD8"
      },
      "source": [
        "data_module.setup() # call this before getting len of dataloader"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iM2A5D6tqteX",
        "outputId": "eb780db6-7ebe-4d96-fe16-70d333e553e8"
      },
      "source": [
        "len(data_module.train_dataloader())"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "834"
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mwG-to8c4pUV"
      },
      "source": [
        "# SAM"
      ]
    },
    {
      "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",
      "metadata": {
        "id": "D1zQOd2AJl7l"
      },
      "source": [
        "# Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "a99EheTsnJyK",
        "outputId": "7a694e45-b874-4fad-c402-e7ec7bbb17a7"
      },
      "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"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(torch.Size([12, 256]), torch.Size([12, 256]))"
            ]
          },
          "metadata": {},
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u8LOwh55Okdz"
      },
      "source": [
        "# test lr schedule with https://www.kaggle.com/isbhargav/guide-to-pytorch-learning-rate-scheduling\n",
        "# test wtih flat cos lr"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "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))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Uc6_SAgKYmmC"
      },
      "source": [
        "# bert_model(sample_batch[\"input_ids\"], sample_batch[\"attention_mask\"]).shape  # should be bs x 6"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XgaReoIsJh0M"
      },
      "source": [
        "class ToxicCommentClassifier(pl.LightningModule):\n",
        "  def __init__(self, pytorch_model: nn.Module, total_steps: int, lr: float = 2e-5, fit_func: str = \"one_cycle\", is_ddp: bool = False):\n",
        "    super().__init__()\n",
        "    self.lr = lr\n",
        "    self.total_steps = total_steps # for lr schedule\n",
        "    self.fit_func = fit_func\n",
        "    self.is_ddp = is_ddp \n",
        "\n",
        "    self.pytorch_model = pytorch_model\n",
        "    self.criterion = nn.BCEWithLogitsLoss() # bce for multi-label\n",
        "\n",
        "    # metrics \n",
        "    self.accuracy = torchmetrics.Accuracy()\n",
        "\n",
        "    # manually define opt step\n",
        "    self.automatic_optimization = False\n",
        "\n",
        "  def forward(self, input_ids, attention_mask):\n",
        "    \"\"\"\n",
        "    forward step for lightning module\n",
        "    \"\"\"\n",
        "    out = self.pytorch_model(input_ids, attention_mask=attention_mask)\n",
        "    return out\n",
        "\n",
        "  def training_step(self, batch, batch_idx):\n",
        "    \"\"\"\n",
        "    training step\n",
        "    define manual optimization with opt.step()\n",
        "    and learning rate scheduler step\n",
        "    \"\"\"\n",
        "    # batch comes from dataset\n",
        "    input_ids = batch[\"input_ids\"]\n",
        "    attention_mask = batch[\"attention_mask\"]\n",
        "    labels = batch[\"labels\"]\n",
        "\n",
        "    opt, opt_adamw = self.optimizers()\n",
        "    print(opt.param_groups[0][\"lr\"], opt_adamw.param_groups[0][\"lr\"]) # check that lrs are the same (is true)\n",
        "    out = self(input_ids, attention_mask)\n",
        "\n",
        "\n",
        "    loss = self.criterion(out, labels.to(dtype=torch.float32)) # cast to float because labels is in Int\n",
        "    acc = self.accuracy(out, labels)\n",
        "\n",
        "    # using closure\n",
        "    # def closure():\n",
        "    #   out = self(input_ids, attention_mask)\n",
        "    #   loss = self.criterion(out, labels.to(dtype=torch.float32))\n",
        "    #   self.manual_backward(loss)\n",
        "    #   return loss\n",
        "    \n",
        "    # self.manual_backward(loss)\n",
        "    # opt.step(closure)\n",
        "    # opt.zero_grad()\n",
        "\n",
        "\n",
        "    if self.is_ddp:\n",
        "      # with self.pytorch_model.no_sync():\n",
        "      # maybe should be (test with both on ddp)\n",
        "      with self.trainer.model.no_sync():\n",
        "        self.manual_backward(loss)\n",
        "    else:\n",
        "      self.manual_backward(loss)\n",
        "    opt.first_step(zero_grad=True)\n",
        "    # 2nd pass\n",
        "    out2 = self(input_ids, attention_mask)\n",
        "    loss2 = self.criterion(out2, labels.to(dtype=torch.float32))\n",
        "    self.manual_backward(loss2)\n",
        "    opt.second_step(zero_grad=True)\n",
        "\n",
        "    # lr schedule\n",
        "    lr_sch = self.lr_schedulers()\n",
        "    lr_sch.step()\n",
        "    # log loss and metrics each k steps and each epoch\n",
        "    log_values = {\"train_loss\": loss, \"train_acc\": acc}\n",
        "    self.log_dict(log_values, prog_bar=True, on_step=True, on_epoch=True)\n",
        "    return {\"loss\": loss, \"accuracy\": acc, \"predictions\": out, \"labels\": labels}\n",
        "\n",
        "  def validation_step(self, batch, batch_idx):\n",
        "    \"\"\"\n",
        "    validation step\n",
        "    get validation loss and accuracy metrics\n",
        "    \"\"\"\n",
        "    # batch comes from dataset\n",
        "    input_ids = batch[\"input_ids\"]\n",
        "    attention_mask = batch[\"attention_mask\"]\n",
        "    labels = batch[\"labels\"]\n",
        "\n",
        "    out = self(input_ids, attention_mask)\n",
        "\n",
        "    # compute loss + metrics\n",
        "    loss = self.criterion(out, labels.to(dtype=torch.float32)) # cast to float because labels is in Int\n",
        "    acc = self.accuracy(out, labels)\n",
        "    # log loss and metrics each k steps and each epoch\n",
        "    log_values = {\"val_loss\": loss, \"val_acc\": acc}\n",
        "    self.log_dict(log_values, sync_dist=True, prog_bar=True, on_step=True, on_epoch=True)\n",
        "    return {\"val_loss\": loss, \"val_accuracy\": acc}\n",
        "\n",
        "  def training_epoch_end(self, outputs):\n",
        "    \"\"\"\n",
        "    after every training epoch, log train metrics\n",
        "\n",
        "    note: this might not be needed if logging is done for each step and epoch\n",
        "    \"\"\"\n",
        "    # log validation loss and metrics\n",
        "    avg_loss = torch.stack([x[\"loss\"] for x in outputs]).mean()\n",
        "    avg_acc = torch.stack([x[\"accuracy\"] for x in outputs]).mean()\n",
        "\n",
        "    self.log(\"avg_train_loss\", avg_loss, sync_dist=True, prog_bar=True)\n",
        "    self.log(\"train_accuracy\", avg_acc, sync_dist=True, prog_bar=True)\n",
        "\n",
        "  def validation_epoch_end(self, outputs):\n",
        "    \"\"\"\n",
        "    after every training epoch, log val metrics\n",
        "\n",
        "    note: this might not be needed if logging is done for each step and epoch\n",
        "    \"\"\"\n",
        "    # log validation loss and metrics\n",
        "    avg_loss = torch.stack([x[\"val_loss\"] for x in outputs]).mean()\n",
        "    avg_acc = torch.stack([x[\"val_accuracy\"] for x in outputs]).mean()\n",
        "    self.log(\"avg_val_loss\", avg_loss, sync_dist=True, prog_bar=True)\n",
        "    self.log(\"val_accuracy\", avg_acc, sync_dist=True, prog_bar=True)\n",
        "\n",
        "  def configure_optimizers(self):\n",
        "    # TODO: \n",
        "    # 1. automatically find steps for onecycleLR\n",
        "    # 2. add flat cos\n",
        "    # 3. make lr scheduler configurable(one_cycle vs flat_cos)\n",
        "\n",
        "    # trying SAM\n",
        "    base_optimizer = AdamW\n",
        "    optimizer = SAM(\n",
        "        self.pytorch_model.parameters(), \n",
        "        base_optimizer=base_optimizer, \n",
        "        lr=self.lr, \n",
        "        betas=(0.9, 0.99)\n",
        "    )\n",
        "    # lr schedule with sam (seems to work fine with 2 step?)\n",
        "    scheduler = OneCycleLR(\n",
        "      optimizer=optimizer,\n",
        "      max_lr=self.lr,\n",
        "      pct_start=0.3,\n",
        "      total_steps=self.total_steps\n",
        "    )\n",
        "\n",
        "    # lr schedule with base optimizer\n",
        "    # scheduler = OneCycleLR(\n",
        "    #   optimizer=optimizer.base_optimizer,\n",
        "    #   max_lr=self.lr,\n",
        "    #   pct_start=0.3,\n",
        "    #   total_steps=self.total_steps\n",
        "    # )\n",
        "    return [optimizer, optimizer.base_optimizer], [scheduler]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZHhFIJkGUTU4",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c42d5f33-de00-4022-82c7-fd2cf4c66d05"
      },
      "source": [
        "# training epoch related hyperparams\n",
        "steps_per_epoch = len(data_module.train_dataloader())\n",
        "n_epochs = 5\n",
        "\n",
        "total_steps = steps_per_epoch * n_epochs\n",
        "print(f\"total_steps: {total_steps}\")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "total_steps: 4170\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YPhpGLWyV0eL"
      },
      "source": [
        "Note: if using multiple gpus, total steps calculation is something like\n",
        "```\n",
        "steps_per_epoch = len(data_module.train_dataloader())\n",
        "\n",
        "# for gpu or tpu cores\n",
        "num_devices = 4\n",
        "if tpu_cores:\n",
        "    num_devices = max(num_devices, tpu_cores)\n",
        "\n",
        "accumulate_grad_batches = 1 # for no gradient accumulation\n",
        "\n",
        "effective_accum = accumulate_grad_batches * num_devices\n",
        "total_steps (steps_per_epoch // effective_accum) * n_epochs\n",
        "\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O5o32pPDKTTi"
      },
      "source": [
        "toxic_comment_model = ToxicCommentClassifier(bert_model, total_steps = total_steps)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IS16h4W9Okgw"
      },
      "source": [
        "sam, adamw = toxic_comment_model.configure_optimizers()[0][0], toxic_comment_model.configure_optimizers()[0][1]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cSMfgHDgOvwx",
        "outputId": "064d8041-0826-4521-d068-bcf8d7ecf900"
      },
      "source": [
        "sam.param_groups[0][\"lr\"], adamw.param_groups[0][\"lr\"]"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(7.999999999999991e-07, 7.999999999999991e-07)"
            ]
          },
          "metadata": {},
          "execution_count": 49
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZR27IbG2Y8o2"
      },
      "source": [
        "# toxic_comment_model(sample_batch[\"input_ids\"], sample_batch[\"attention_mask\"])[1].shape  # should be bs x 6"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "li0_cLa82LXx"
      },
      "source": [
        "# Trainer\n",
        "\n",
        "defining the trainer and running `trainer.fit` "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XrBKqeKnKWy-"
      },
      "source": [
        "# callbacks\n",
        "lr_monitor_cb = LearningRateMonitor(logging_interval='step')\n",
        "\n",
        "logger = TensorBoardLogger(\"lightning_logs\", name=\"toxic-comments\")\n",
        "\n",
        "callbacks = [lr_monitor_cb]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cZdITHvZhBSF",
        "outputId": "88ebfd2f-0e81-454c-979a-7d7565e92fd4"
      },
      "source": [
        "trainer = pl.Trainer(\n",
        "    logger=logger,\n",
        "    callbacks = callbacks,\n",
        "    precision=16,\n",
        "    max_epochs = n_epochs,\n",
        "    accelerator=\"gpu\",\n",
        "    devices=1,\n",
        ")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Using 16bit native Automatic Mixed Precision (AMP)\n",
            "GPU available: True, used: True\n",
            "TPU available: False, using: 0 TPU cores\n",
            "IPU available: False, using: 0 IPUs\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FSIidqFNT5n5"
      },
      "source": [
        "# optional: currently not working?\n",
        "# lr_finder = trainer.tuner.lr_find(model=toxic_comment_model, datamodule=data_module)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 579,
          "referenced_widgets": [
            "7aea1bf3de1e43ef890f7c1aad966b04",
            "c32531cc302b4a59884a5c85519a8957",
            "24f6d3e223d44068b0701daafaa07f62",
            "f58951b5683d4554a2df7f6bd74b3730",
            "39efe0ead32c4d90b5d48dcf1a490b1b",
            "74800a8f29c2457e8553540b8aedde44",
            "39acd4da787a4bf5baa8d468e21f7715",
            "af33f16f86504041b1b2f340b5a4e182",
            "1527d28c976f44468a0b791cf81be271",
            "0bc850aa8d81406491b930c54ee81a09",
            "c9538fe593ec481c878fb5a8c58f7850",
            "c758b12eef464d3089f32fce401a65c9",
            "8cd4f69584754f7fa029f4971d65fa91",
            "d668bbda43964950a0eec05c4bfa5da1",
            "1dcd44a8e2574c7d9ff1b891259f4c9c",
            "91dbcff7475f46c3ba6e6a585414a2bf",
            "d9b6105e50a447e5b9c896a53a32a6a4",
            "b926716c578646a3a77045a3194427ec",
            "ec262b68cdb44ba1b2e9d6bda4d0a83b",
            "a9796fae209d4cebb13011450172e701",
            "400d011bb4d242dc95d6792bb17c3dfd",
            "7257193458d94e3d9c977a636b042139"
          ]
        },
        "id": "Jfl3uth_hSRg",
        "outputId": "de361270-87bd-40a7-cc1c-6bdbcd6e1913"
      },
      "source": [
        "trainer.fit(toxic_comment_model, data_module)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/core/datamodule.py:470: LightningDeprecationWarning: DataModule.setup has already been called, so it will not be called again. In v1.6 this behavior will change to always call DataModule.setup.\n",
            "  f\"DataModule.{name} has already been called, so it will not be called again. \"\n",
            "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
            "\n",
            "  | Name          | Type              | Params\n",
            "----------------------------------------------------\n",
            "0 | pytorch_model | BertModel         | 82.1 M\n",
            "1 | criterion     | BCEWithLogitsLoss | 0     \n",
            "2 | accuracy      | Accuracy          | 0     \n",
            "----------------------------------------------------\n",
            "82.1 M    Trainable params\n",
            "0         Non-trainable params\n",
            "82.1 M    Total params\n",
            "164.246   Total estimated model params size (MB)\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "7aea1bf3de1e43ef890f7c1aad966b04",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "Validation sanity check: 0it [00:00, ?it/s]"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "c758b12eef464d3089f32fce401a65c9",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "Training: 0it [00:00, ?it/s]"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "7.999999999999991e-07 7.999999999999991e-07\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/torch/optim/lr_scheduler.py:134: UserWarning: Detected call of `lr_scheduler.step()` before `optimizer.step()`. In PyTorch 1.1.0 and later, you should call them in the opposite order: `optimizer.step()` before `lr_scheduler.step()`.  Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate\n",
            "  \"https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate\", UserWarning)\n",
            "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/optimization/closure.py:36: LightningDeprecationWarning: One of the returned values {'accuracy', 'loss', 'predictions', 'labels'} has a `grad_fn`. We will detach it automatically but this behaviour will change in v1.6. Please detach it manually: `return {'loss': ..., 'something': something.detach()}`\n",
            "  f\"One of the returned values {set(extra.keys())} has a `grad_fn`. We will detach it automatically\"\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "8.000303194087598e-07 8.000303194087598e-07\n",
            "8.001212774435278e-07 8.001212774435278e-07\n",
            "8.002728735297606e-07 8.002728735297606e-07\n",
            "8.004851067099044e-07 8.004851067099044e-07\n",
            "8.007579756433636e-07 8.007579756433636e-07\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py:685: UserWarning: Detected KeyboardInterrupt, attempting graceful shutdown...\n",
            "  rank_zero_warn(\"Detected KeyboardInterrupt, attempting graceful shutdown...\")\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AtRPI8ocIRai",
        "outputId": "4ccfe560-dafa-4bb4-fe1f-9010c5190f61"
      },
      "source": [
        "trainer.logged_metrics"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'train_accuracy': 0.9270631074905396,\n",
              " 'train_loss': 0.17185573279857635,\n",
              " 'val_accuracy': 0.9817500114440918,\n",
              " 'val_loss': 0.05154530331492424}"
            ]
          },
          "metadata": {},
          "execution_count": 128
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KTYeeBp4LDuy"
      },
      "source": [
        "%reload_ext tensorboard\n",
        "%tensorboard --logdir lightning_logs/\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O5q2zeY3zLzG"
      },
      "source": [
        "# Saving only the pytorch model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GeNvSmBCwWpN"
      },
      "source": [
        "# saving model weights\n",
        "saved_model_pth = \"saved_model.pth\"\n",
        "\n",
        "torch.save(toxic_comment_model.pytorch_model.state_dict(), saved_model_pth)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z_Qa6mk9bB-V"
      },
      "source": [
        "# Exporting model as pure pytorch "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aSX1sSa1hsMI",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "07c25701-2ed7-4f5f-f484-0a2f66775467"
      },
      "source": [
        "new_bert_model = BertModel(6)\n",
        "\n",
        "new_bert_model.load_state_dict(torch.load(saved_model_pth))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<All keys matched successfully>"
            ]
          },
          "metadata": {},
          "execution_count": 126
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ytqP_Olebu64",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d58519ce-5b10-418d-e41b-2310e479cf42"
      },
      "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": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([8, 6])\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[ 3.3002, -2.6973,  2.7221, -5.2559,  0.7052, -3.9604],\n",
              "        [ 4.4156, -0.5025,  3.8863, -3.7945,  2.5557, -2.4770],\n",
              "        [ 3.4849, -2.9588,  0.6141, -4.2794,  1.5000, -1.5067],\n",
              "        [ 1.9291, -5.8827, -1.8313, -5.9951, -1.6118, -4.0301],\n",
              "        [ 0.1542, -5.0230, -2.0128, -4.9881, -2.5206, -4.2967],\n",
              "        [ 2.1084, -3.1527, -2.2322,  0.3907, -2.0325, -3.0269],\n",
              "        [ 0.9454, -6.3798, -3.6319, -5.6751, -1.6457, -4.8039],\n",
              "        [ 1.1208, -5.3446, -0.0648, -5.9639, -2.0950, -5.3173]],\n",
              "       grad_fn=<AddmmBackward0>)"
            ]
          },
          "metadata": {},
          "execution_count": 127
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Q0vlXp_dc94H"
      },
      "source": [
        "toxic_comment_model(sample_batch[\"input_ids\"], sample_batch[\"attention_mask\"])[1]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_GEsu9KydC93"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}