Colab_training_Triples.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Colab training: Triples.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "PQcKXNcYb-kO",
        "siH0EUgEcLej",
        "EXUHopxMcn06",
        "u_OB6Da_eF0e",
        "rznh0cJedlvB",
        "tmbJXpnddo-E",
        "XB5Fvq41eXCX",
        "exVOKdshgaCR"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "aa60b9b9f6414c2faa84755f12274c19": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "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_bc9ad37a618e484bb3c18413314c2594",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_b480cf74784749bbb695b00ccfb79d43",
              "IPY_MODEL_a6665a5398004fa2bc1437ac69b138b2"
            ]
          }
        },
        "bc9ad37a618e484bb3c18413314c2594": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "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
          }
        },
        "b480cf74784749bbb695b00ccfb79d43": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_d6f346e5f82c46aa81e4bede571581f6",
            "_dom_classes": [],
            "description": "Downloading: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 570,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 570,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_606df5aca1a3410ebc694cfcbc6a2df6"
          }
        },
        "a6665a5398004fa2bc1437ac69b138b2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_76a55d15ebe94e8488e43fa0169bf1d9",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 570/570 [00:15&lt;00:00, 37.8B/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_c0164eadd2254d84bd6b20a5f98daeec"
          }
        },
        "d6f346e5f82c46aa81e4bede571581f6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_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"
          }
        },
        "606df5aca1a3410ebc694cfcbc6a2df6": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "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
          }
        },
        "76a55d15ebe94e8488e43fa0169bf1d9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "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"
          }
        },
        "c0164eadd2254d84bd6b20a5f98daeec": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "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
          }
        },
        "35d6ab8b7cb44cbb8ae6e4a0e621a15b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "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_3dbb39a901cc4c3098251a09c62a0430",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_50f61fae634e476db0d0ba98def033b9",
              "IPY_MODEL_47b060ccd3084d078d2026f9103aa880"
            ]
          }
        },
        "3dbb39a901cc4c3098251a09c62a0430": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "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
          }
        },
        "50f61fae634e476db0d0ba98def033b9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_aa08522dfc654b30b6ef401d4c69bfae",
            "_dom_classes": [],
            "description": "Downloading: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 435779157,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 435779157,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_65d01253bfdb48cd9a6488a02edb34db"
          }
        },
        "47b060ccd3084d078d2026f9103aa880": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_c6258879d4244f17865487cb681ccbcc",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 436M/436M [00:14&lt;00:00, 29.8MB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_cb9130e18201498eaae12aecb12feccb"
          }
        },
        "aa08522dfc654b30b6ef401d4c69bfae": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_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"
          }
        },
        "65d01253bfdb48cd9a6488a02edb34db": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "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
          }
        },
        "c6258879d4244f17865487cb681ccbcc": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "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"
          }
        },
        "cb9130e18201498eaae12aecb12feccb": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "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
          }
        },
        "1c598e906ba343c59a6b4e907c4d0bf3": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "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_0dd4924490b447409073a6299501a8a0",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_687724e9265f4ab0887ad07b1a82cdfa",
              "IPY_MODEL_6d344bab148b482ab01ccc9be5cf9636"
            ]
          }
        },
        "0dd4924490b447409073a6299501a8a0": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "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
          }
        },
        "687724e9265f4ab0887ad07b1a82cdfa": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_fe2ce756730246db9b9cd3530c350cea",
            "_dom_classes": [],
            "description": "Epoch 1/15:   0%",
            "_model_name": "FloatProgressModel",
            "bar_style": "",
            "max": 15,
            "_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_692ebc7123b449f7852ea2eca3a4fb97"
          }
        },
        "6d344bab148b482ab01ccc9be5cf9636": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_0fba8ac13caa4218962ebf4e5550216e",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 0/15 [00:00&lt;?, ?sections/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_a3e62f3381a04a0ead5230cb3109a96c"
          }
        },
        "fe2ce756730246db9b9cd3530c350cea": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_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"
          }
        },
        "692ebc7123b449f7852ea2eca3a4fb97": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "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
          }
        },
        "0fba8ac13caa4218962ebf4e5550216e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "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"
          }
        },
        "a3e62f3381a04a0ead5230cb3109a96c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "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
          }
        },
        "dfda4c1ea0a64b76b6a1994fcddc0f93": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "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_3f10cc1852964e99935be6139d7024da",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_ad8567e9f8ff4d30ad73382bb5117ea0",
              "IPY_MODEL_c0ece211fc7d48db871b7a830e2f3230"
            ]
          }
        },
        "3f10cc1852964e99935be6139d7024da": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "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
          }
        },
        "ad8567e9f8ff4d30ad73382bb5117ea0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_b62dd92a7bb743b18f8541f072936db7",
            "_dom_classes": [],
            "description": "training iterations:  22%",
            "_model_name": "FloatProgressModel",
            "bar_style": "",
            "max": 290,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 64,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_5ee56bd3938a44378aaa320c5220fd9b"
          }
        },
        "c0ece211fc7d48db871b7a830e2f3230": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_9e1787b7201f4ef9815c5f0fb1dfd23f",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 64/290 [00:17&lt;00:59,  3.79batch/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_0caf142e0aad4a2ba4581d0a7eff5593"
          }
        },
        "b62dd92a7bb743b18f8541f072936db7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_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"
          }
        },
        "5ee56bd3938a44378aaa320c5220fd9b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "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
          }
        },
        "9e1787b7201f4ef9815c5f0fb1dfd23f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "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"
          }
        },
        "0caf142e0aad4a2ba4581d0a7eff5593": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "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": "PQcKXNcYb-kO"
      },
      "source": [
        "# One-timers"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yDktLi3sb_44",
        "outputId": "db223dc2-c3f9-410d-931f-910e2e9daa80"
      },
      "source": [
        "!nvidia-smi\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive', force_remount=True)\n",
        "!mkdir dataset\n",
        "\n",
        "!cp /content/drive/MyDrive/Research/triples/data/*.csv dataset\n",
        "!cp -r /content/drive/MyDrive/Research/triples/HuggingFace dataset\n",
        "!ls dataset\n",
        "!pip3 install -q transformers tensorboard_logger seqeval sentencepiece tokenizers sentence_transformers"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Tue Jul 27 13:00:20 2021       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 470.42.01    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 P100-PCIE...  Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   35C    P0    26W / 250W |      0MiB / 16280MiB |      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",
            "Mounted at /content/drive\n",
            "20NewsGroups.csv  most_frequent.csv  squad_frequents.csv  triples.csv\n",
            "frequents.csv\t  results.csv\t     squad_train.csv\n",
            "HuggingFace\t  squad1.csv\t     squad_triples.csv\n",
            "\u001b[K     |████████████████████████████████| 2.6 MB 16.6 MB/s \n",
            "\u001b[K     |████████████████████████████████| 43 kB 2.8 MB/s \n",
            "\u001b[K     |████████████████████████████████| 1.2 MB 66.6 MB/s \n",
            "\u001b[K     |████████████████████████████████| 3.3 MB 47.2 MB/s \n",
            "\u001b[K     |████████████████████████████████| 85 kB 6.0 MB/s \n",
            "\u001b[K     |████████████████████████████████| 895 kB 55.7 MB/s \n",
            "\u001b[K     |████████████████████████████████| 636 kB 58.5 MB/s \n",
            "\u001b[?25h  Building wheel for seqeval (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for sentence-transformers (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "siH0EUgEcLej"
      },
      "source": [
        "# Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WDFduBq7cGAd"
      },
      "source": [
        "import os\n",
        "os.environ['CUDA_LAUNCH_BLOCKING'] = \"1\" \n",
        "import time, torch, random, glob, re, gc, datetime, tokenizers, pdb\n",
        "\n",
        "import numpy as np\n",
        "import transformers\n",
        "import pandas as pd\n",
        "import torch.nn as nn\n",
        "import seaborn as sns\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "from tokenizers import *\n",
        "from transformers import *\n",
        "from functools import partial\n",
        "from pathlib import Path\n",
        "from tqdm.notebook import tqdm\n",
        "from torch.nn import functional as F\n",
        "from itertools import cycle, chain\n",
        "from torch.utils.data import Dataset, DataLoader, IterableDataset, TensorDataset\n",
        "# from sklearn.model_selection import GroupKFold\n",
        "from sklearn.metrics import f1_score, precision_score, recall_score, accuracy_score\n",
        "from sklearn.model_selection import train_test_split, RepeatedKFold, KFold\n",
        "from ast import literal_eval as eval\n",
        "\n",
        "# from transformers import RobertaForSequenceClassification, RobertaConfig, RobertaTokenizer, RobertaForTokenClassification"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tk9_g3N1cU4A"
      },
      "source": [
        "import sys\n",
        "DRIVE_DIR=\"/content/drive/My Drive/Research/triples/\"\n",
        "sys.path.insert(0, DRIVE_DIR)\n",
        "from utils import seed_everything, count_params"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EXUHopxMcn06"
      },
      "source": [
        "# Globals and Config"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "68syWH6Vcm_C"
      },
      "source": [
        "class Config:\n",
        "\n",
        "    random_state=2021\n",
        "    k=5 #folds\n",
        "\n",
        "    device=\"cuda\"\n",
        "    selected_folds=list(range(k))\n",
        "    # selected_folds=[3,4]\n",
        "    seed = 2021\n",
        "    model=\"bert-base-cased\"\n",
        "    # checkpoints=[f\"/content/drive/MyDrive/Research/triples/2021-07-24/NER_bert-base-cased_ fold - {i+1}_.pt\" for i in range(k)]\n",
        "    checkpoint=[] # if this list has any checkpoint, the model uses that checkpoint to be the starting point and then finetunes over it.\n",
        "    pretrained=True\n",
        "    lowercase = False\n",
        "\n",
        "    num_labels=4\n",
        "    batch_size = 32\n",
        "    batch_size_val = int(batch_size * 1.5)\n",
        "    weight_decay =0.001\n",
        "    \n",
        "    epochs = 15\n",
        "    lr =5e-5\n",
        "    warmup_prop = 0.1\n",
        "    freeze_main=False #this parameters controls whether we want to freeze the main bert and train only the classifier (True) or train the whole model\n",
        "    max_len=128\n",
        "    save_every_epoch=list(range(5, epochs, 3)) #list specifies which checkpoints to save\n",
        "    model_names=[\"roberta\", \"bert\", \"albert\", \"transformer\", \"distilbert\"]\n",
        "\n",
        "\n",
        "CP_DIR=Path(\"/content/drive/MyDrive/Research/triples\")\n",
        "NUM_WORKERS = 2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P2Y2kcgzcp3B"
      },
      "source": [
        "TRANSFORMERS_DIR=Path(\"dataset/HuggingFace/\")\n",
        "TRANSFORMERS={\n",
        "    \"roberta-base\":{\n",
        "        \"model_config\":(RobertaModel, RobertaConfig),\n",
        "        \"tokenizer\":RobertaTokenizer,\n",
        "    },\n",
        "    \"bert-base-cased\":{\n",
        "        \"model_config\":(BertModel, BertConfig),\n",
        "        \"tokenizer\":BertWordPieceTokenizer,\n",
        "    },\n",
        "    \"bert-base-uncased\":{\n",
        "        \"model_config\":(BertModel, BertConfig),\n",
        "        \"tokenizer\":BertWordPieceTokenizer,\n",
        "    },\n",
        "    \"albert-base-v2\":{\n",
        "        \"model_config\":(AlbertModel,AlbertConfig),\n",
        "        \"tokenizer\":AlbertTokenizer,\n",
        "    },\n",
        "    \"gpt2\":{\n",
        "        \"model_config\":(GPT2Model, GPT2Config),\n",
        "        \"tokenizer\":GPT2Tokenizer,\n",
        "    },\n",
        "    \"distilbert-base-cased\":{\n",
        "        \"model_config\":(DistilBertModel, DistilBertConfig),\n",
        "        \"tokenizer\":DistilBertTokenizer,\n",
        "    }\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u_OB6Da_eF0e"
      },
      "source": [
        "# Function and Helpers"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Tty9esMHeH5_"
      },
      "source": [
        "def get_checkpoint_dir():\n",
        "  today=str(datetime.date.today())\n",
        "  checkpoint_dir=CP_DIR/today\n",
        "\n",
        "  if not os.path.exists(checkpoint_dir):\n",
        "    os.makedirs(checkpoint_dir)\n",
        "  return checkpoint_dir\n",
        "\n",
        "def checkpoint_name():\n",
        "  return Config.task+\"_\"+Config.model\n",
        "\n",
        "def save_model_weights(model, filename, verbose=1, cp_folder=\"\"):\n",
        "    if verbose:\n",
        "        print(f\"\\n -> Saving weights to {os.path.join(cp_folder, filename)}\\n\")\n",
        "    torch.save(model.state_dict(), os.path.join(cp_folder, filename))\n",
        "\n",
        "\n",
        "def save_log(list_, logdir):\n",
        "    if os.path.exists(logdir):\n",
        "        mode=\"a\"\n",
        "    else:\n",
        "        mode=\"w\"\n",
        "    with open(logdir, mode) as f:\n",
        "        f.writelines(\"\\n\".join(list_))\n",
        "        f.writelines(\"\\n\")\n",
        "    \n",
        "\n",
        "# def get_scores(truths, preds):\n",
        "#     f1=f1_score(truths, preds)\n",
        "#     recall=recall_score(truths, preds)\n",
        "#     precision=precision_score(truths, preds)\n",
        "#     accuracy=accuracy_score(truths, preds)\n",
        "    \n",
        "#     return f1, precision, recall, accuracy\n",
        "\n",
        "def load(model, with_checkpoint=None):\n",
        "    model=Transformer(model)\n",
        "    if with_checkpoint:\n",
        "        checkpoint=torch.load(with_checkpoint, map_location=\"cpu\")\n",
        "        model.load_state_dict(checkpoint)\n",
        "        print(\"Checkpoint loaded!\", end=\"\\r\")\n",
        "    return model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1SUuB3aOl2vi"
      },
      "source": [
        "def modify_label(text, label_):\n",
        "\n",
        "  \"the modify relations function\"\n",
        "    label=label_.split(\" \")\n",
        "    loc1=text.find(label[0])\n",
        "    loc2=text.find(label[-1])+len(label[-1])\n",
        "\n",
        "    if loc1>loc2:\n",
        "      return label_\n",
        "    return text[loc1:loc2]\n",
        "\n",
        "\n",
        "def locate_label_string(text, label_, fill_value=1):\n",
        "\n",
        "    \"\"\"\n",
        "    Finds the label in the text\n",
        "    \"\"\"\n",
        "    if not label_ in text:\n",
        "        label=modify_label(text, label_)\n",
        "    else:\n",
        "        label=label_\n",
        "    len_label = len(label) - 1\n",
        "    candidates_idx = [i for i, e in enumerate(text) if e == label[1]]\n",
        "\n",
        "\n",
        "    for idx in candidates_idx:\n",
        "        if \" \" + text[idx: idx + len_label] == label:\n",
        "            idx_start = idx\n",
        "            idx_end = idx + len_label\n",
        "            break\n",
        "\n",
        "    assert (\n",
        "        text[idx_start:idx_end] == label[1:]\n",
        "    ), f'\"{text[idx_start: idx_end]}\" instead of \"{label}\" in \"{text}\"'\n",
        "\n",
        "    char_targets = np.zeros(len(text))\n",
        "    char_targets[idx_start:idx_end] = fill_value\n",
        "\n",
        "    return idx_start, idx_end, char_targets\n",
        "\n",
        "def locate_label_tokens(offsets, char_targets):\n",
        "    \"\"\"\n",
        "    Finds the tokens corresponding to the found labels\n",
        "    \"\"\"\n",
        "    \n",
        "\n",
        "    target_idx = []\n",
        "    for idx, (offset1, offset2) in enumerate(offsets):\n",
        "        if sum(char_targets[offset1:offset2]) > 0:\n",
        "            target_idx.append(idx)\n",
        "\n",
        "    if not len(target_idx):\n",
        "        for idx, (offset1, offset2) in enumerate(offsets):\n",
        "            if sum(char_targets[offset1:offset2]) > 0:\n",
        "                target_idx.append(idx)\n",
        "\n",
        "    return target_idx[0], target_idx[-1]\n",
        "\n",
        "def preprocess(text, entities, tokenizer, tokens, max_len=128):\n",
        "\n",
        "  \"\"\"Preprocessing required for the input to transformer and to reconstruct output from it.\"\"\"\n",
        "    text = \" \" + \" \".join(str(text).split())\n",
        "    entities=[\" \"+i for i in entities]\n",
        "    label_locations=[]\n",
        "    for idx, l in enumerate(entities):\n",
        "        label_locations.append(locate_label_string(text, l, fill_value=1))    \n",
        "\n",
        "    tokenized = tokenizer.encode(text)\n",
        "    input_ids_text = tokenized.ids[1:-1]\n",
        "    offsets = tokenized.offsets[1:-1]\n",
        "    label_location_in_tokens=[]\n",
        "    start_=0\n",
        "    end_=0\n",
        "    for label_location in label_locations:\n",
        "        target_start, target_end=locate_label_tokens(offsets, label_location[2])\n",
        "\n",
        "        start_=min(start_, target_start)\n",
        "        end_=max(end_, target_end)\n",
        "\n",
        "        label_location_in_tokens.append((target_start, target_end))\n",
        "\n",
        "    input_ids = (\n",
        "            [tokens[\"cls\"]]\n",
        "            + input_ids_text[:max_len - 2]\n",
        "            + [tokens[\"sep\"]])\n",
        "    text_offsets = [(0, 0)] + offsets[:max_len - 2] + [(0, 0)]\n",
        "    label_location_in_tokens=[(i[0]+1, i[1]+1) for i in label_location_in_tokens]\n",
        "    padding_length = max_len - len(input_ids)\n",
        "    if padding_length > 0:\n",
        "        input_ids = input_ids + ([tokens[\"pad\"]] * padding_length)\n",
        "        text_offsets = text_offsets + ([(0, 0)] * padding_length)\n",
        "        \n",
        "    return {\n",
        "        \"input_ids\":input_ids,\n",
        "        \"label_location_in_tokens\":label_location_in_tokens,\n",
        "        \"text\":text,\n",
        "        \"label\":entities,\n",
        "        \"offsets\":text_offsets,\n",
        "    }"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rznh0cJedlvB"
      },
      "source": [
        "# Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e_f32rfDc2YJ"
      },
      "source": [
        "class SentenceDataset(Dataset):\n",
        "    def __init__(self, df, tokenizer, tokens, max_len=128):\n",
        "        self.tokenizer = tokenizer\n",
        "        self.tokens=tokens\n",
        "        self.max_len = max_len\n",
        "\n",
        "        self.texts = df[\"sentence\"]\n",
        "        self.labels = df[\"triple\"].map(eval)\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.texts)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        data=preprocess(self.texts.iloc[idx], self.labels.iloc[idx], tokenizer=self.tokenizer, tokens=self.tokens, max_len=self.max_len)\n",
        "        label=[0]*len(data['input_ids'])\n",
        "        for idx, location in enumerate(data[\"label_location_in_tokens\"]):\n",
        "            if location[0]==location[1]:\n",
        "                label[location[0]]=idx+1\n",
        "            else:\n",
        "                label[location[0]:location[1]]=[idx+1]*(location[1]-location[0])\n",
        "                            \n",
        "        return {\n",
        "            \"input_ids\":torch.tensor(data[\"input_ids\"], dtype=torch.long),\n",
        "            \"label\":torch.nn.functional.one_hot(torch.tensor(label), num_classes=4).float(),\n",
        "            \"text\":data['text'],\n",
        "            \"offsets\":data['offsets']\n",
        "        }"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tmbJXpnddo-E"
      },
      "source": [
        "# Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "k2x6XWQ-dpvD"
      },
      "source": [
        "class Transformer(nn.Module):\n",
        "    def __init__(self, model, maxlen=128):\n",
        "        super().__init__()\n",
        "        self.name = model\n",
        "        model_type, config_type=TRANSFORMERS[model]['model_config']\n",
        "        if Config.pretrained:\n",
        "            self.transformer=model_type.from_pretrained(model, output_hidden_states=True, num_labels=Config.num_labels)\n",
        "        else:\n",
        "            config_file=TRANSFORMERS[model]['config']\n",
        "            config=config_type.from_json_file(config_file)\n",
        "            config.num_labels=Config.num_labels\n",
        "            config.output_hidden_states=True\n",
        "            self.transformer=model_type(config)\n",
        "            \n",
        "        self.nb_features = self.transformer.pooler.dense.out_features\n",
        "        if \"roberta\" in self.name:\n",
        "            self.pad_idx=1\n",
        "        else:\n",
        "            self.pad_idx=0\n",
        "        self.logits = nn.Sequential(\n",
        "            nn.Linear(self.nb_features, self.nb_features),\n",
        "            nn.Tanh(),\n",
        "            nn.Linear(self.nb_features, Config.num_labels),\n",
        "        )\n",
        "\n",
        "    def forward(self, input_ids, attention_mask=None):        \n",
        "        hidden_states = self.transformer(\n",
        "            input_ids,\n",
        "            attention_mask=(input_ids != self.pad_idx).long(),\n",
        "        )[-1]\n",
        "\n",
        "        features = hidden_states[-1]\n",
        "        logits = torch.sigmoid(self.logits(features))\n",
        "        \n",
        "        return logits"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XB5Fvq41eXCX"
      },
      "source": [
        "# Fitting"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KjETCrIrdsMM"
      },
      "source": [
        "def fit(model,train_dataset,val_dataset, fold, epochs,batch_size, weight_decay=0,warmup_prop=0.0,lr=5e-4):\n",
        "\n",
        "  \"\"\"Batchwise training and validation iterations.\"\"\"\n",
        "    train_loader = DataLoader(train_dataset, batch_size=batch_size, num_workers=NUM_WORKERS)\n",
        "    val_loader = DataLoader(val_dataset, batch_size=batch_size, num_workers=NUM_WORKERS)\n",
        "\n",
        "    opt_params = []\n",
        "    no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n",
        "    for n, p in model.named_parameters():\n",
        "        wd = 0 if any(nd in n for nd in no_decay) else weight_decay\n",
        "        opt_params.append(\n",
        "            {\"params\": [p], \"weight_decay\": wd, \"lr\": lr}\n",
        "        )\n",
        "\n",
        "    optimizer = AdamW(opt_params, lr=lr, betas=(0.5, 0.999))\n",
        "\n",
        "    n_steps=epochs*len(train_loader)\n",
        "    num_warmup_steps = int(warmup_prop * n_steps)\n",
        "    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps, n_steps)\n",
        "    # scheduler=ReduceLROnPlateau(optimizer, mode='min', factor=0.03, patience=2, threshold=0.001, threshold_mode='rel', cooldown=0, min_lr=1e-6, eps=1e-08, verbose=False)\n",
        "\n",
        "    total_steps = 0\n",
        "    epoch=0\n",
        "\n",
        "    loss_function=nn.BCELoss()\n",
        "\n",
        "    save_log([\"\\n\",str(datetime.datetime.now()).split(\".\")[0],\"\\n\", checkpoint_name()+f\"_fold_{fold+1}\"], logdir=get_checkpoint_dir()/f\"log_{CHECKPOINT_KEYWORD}.txt\")\n",
        "    with tqdm(total=epochs, desc=\"Epoch {}/{}\".format(epoch + 1, epochs), unit=\"sections\", position=0,leave=True) as pbar:\n",
        "        for epoch in range(epochs):\n",
        "            model.train()\n",
        "            start_time = time.time()\n",
        "            optimizer.zero_grad()\n",
        "            avg_loss = 0\n",
        "\n",
        "            with tqdm(total=len(train_loader), desc=\"training iterations\", unit=\"batch\", position=1, leave=True) as pbar2:\n",
        "                for step, data in enumerate(train_loader):\n",
        "                  total_steps+=1\n",
        "                  input_ids=data['input_ids']\n",
        "                  labels=data['label']\n",
        "\n",
        "                  logits=model(input_ids=input_ids.to(Config.device))                  \n",
        "                  loss=loss_function(logits, labels.to(Config.device))\n",
        "  \n",
        "                  avg_loss += loss.item() / len(train_loader)\n",
        "                  nn.utils.clip_grad_norm_(model.parameters(), 10.0)\n",
        "                  loss.backward()\n",
        "                  optimizer.step()\n",
        "                  scheduler.step()\n",
        "                  model.zero_grad()\n",
        "                  pbar2.update()\n",
        "\n",
        "\n",
        "            model.eval()\n",
        "            avg_val_loss = 0.\n",
        "            preds, truths = [], []\n",
        "            with torch.no_grad():\n",
        "                with tqdm(total=len(val_loader), desc=\"validation iterations\", unit=\"batch\", position=2, leave=True) as pbar3:\n",
        "                    for idx_val, data in enumerate(val_loader):\n",
        "                      input_ids=data['input_ids']\n",
        "                      labels=data['label']\n",
        "                      logits=model(input_ids=input_ids.to(Config.device))\n",
        "                      loss=loss_function(logits, labels.to(Config.device))\n",
        "                      avg_val_loss += loss.item() / len(val_loader)\n",
        "                      pbar3.update()\n",
        "                      \n",
        "            dt = time.time() - start_time\n",
        "            lr = scheduler.get_lr()[0]\n",
        "            # lr = optimizer.param_groups[0]['lr']\n",
        "            # scheduler.step(avg_val_loss)\n",
        "            \n",
        "            if epoch+1 in Config.save_every_epoch:\n",
        "                save_model_weights(model, f'{checkpoint_name()}_fold-{fold+1}_epoch-{epoch+1}_{CHECKPOINT_KEYWORD}.pt', cp_folder=get_checkpoint_dir())\n",
        "\n",
        "            log_lr=f\"Epoch {epoch + 1}/{epochs} \\t lr={lr:.1e} \\t t={dt:.0f}s \\t \\n\"\n",
        "            print(log_lr)\n",
        "            log_score=f\"loss={avg_loss:.3f}\\t val_loss={avg_val_loss:.3f}  \\n\"\n",
        "            print(log_score)\n",
        "            save_log([log_lr, log_score], logdir=get_checkpoint_dir()/f\"log_{CHECKPOINT_KEYWORD}.txt\")\n",
        "            pbar.update()\n",
        "\n",
        "\n",
        "    del loss, data, avg_val_loss, avg_loss, train_loader, val_loader\n",
        "    if Config.device != \"cpu\":\n",
        "        torch.cuda.empty_cache()\n",
        "    gc.collect()\n",
        "\n",
        "    return preds\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "exVOKdshgaCR"
      },
      "source": [
        "# K-fold"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "52ZpeGZgeWKy"
      },
      "source": [
        "def k_fold(df, save=True, config=None):\n",
        "\n",
        "    \"\"\"K-fold training\"\"\"\n",
        "  \n",
        "    kf=KFold(n_splits=config.k)\n",
        "    X=list(range(len(df)))\n",
        "    folds=list(kf.split(X=X))\n",
        "\n",
        "    tokenizer = BertWordPieceTokenizer(\n",
        "                \"dataset/HuggingFace/Bert/bert_base_uncased_vocab.txt\",\n",
        "                lowercase=False\n",
        "            )\n",
        "    tokens = {\n",
        "            'cls': tokenizer.token_to_id('[CLS]'),\n",
        "            'sep': tokenizer.token_to_id('[SEP]'),\n",
        "            'pad': tokenizer.token_to_id('[PAD]'),\n",
        "        }\n",
        "\n",
        "    seed_everything(config.seed)      \n",
        "    for fold, (train_idx, val_idx) in enumerate(folds):\n",
        "      if fold in config.selected_folds:\n",
        "        score = 0\n",
        "        model=Transformer(Config.model).to(config.device)\n",
        "\n",
        "        # if not model.transformer.config.pad_token_id:\n",
        "        #     model.transformer.config.pad_token_id=tokenizer.eos_token_id\n",
        "\n",
        "        if Config.checkpoints:\n",
        "            if len(Config.checkpoints)==1:\n",
        "                print(\"Loading common checkpoint\")\n",
        "                checkpoint=torch.load(Config.checkpoints[0], map_location=\"cpu\")\n",
        "\n",
        "            elif len(Config.checkpoints)>1:\n",
        "                print(\"Loading fold checkpoint\")\n",
        "                checkpoint=torch.load(Config.checkpoints[fold], map_location=\"cpu\")\n",
        "\n",
        "            model.load_state_dict(checkpoint)\n",
        "            print(\"loaded checkpoint\")\n",
        "\n",
        "        else:\n",
        "            print(\"No checkpoint provided..!\")\n",
        "            pass\n",
        "\n",
        "        \n",
        "        if Config.freeze_main:\n",
        "            frozen=0\n",
        "            for name, param in model.named_parameters():\n",
        "                if any(model_name in name for model_name in Config.model_names):\n",
        "                    param.requires_grad=False\n",
        "                    frozen+=1\n",
        "            print(f\"{frozen} layers frozen!\")\n",
        "\n",
        "                    \n",
        "        model.zero_grad()\n",
        "        print(f\"\\n-------------   Fold {fold + 1} / {len(folds)}  -------------\\n\")\n",
        "        train_dataset=SentenceDataset(df=df.iloc[train_idx], tokenizer=tokenizer, tokens=tokens, max_len=Config.max_len)\n",
        "        val_dataset=SentenceDataset(df=df.iloc[val_idx], tokenizer=tokenizer, tokens=tokens, max_len=Config.max_len)\n",
        "        n_parameters = count_params(model)\n",
        "\n",
        "        print(f\">>{len(train_dataset)} training texts<<\")\n",
        "        print(f\">>{len(val_dataset)} validation texts<<\")\n",
        "        print(f\">>{n_parameters} trainable parameters<<\\n\")  \n",
        "\n",
        "        preds = fit(model, train_dataset, val_dataset,epochs=config.epochs, batch_size=config.batch_size, weight_decay=config.weight_decay, lr=config.lr, warmup_prop=config.warmup_prop, fold=fold)\n",
        "\n",
        "        \n",
        "\n",
        "        if save:\n",
        "            save_model_weights(model, f'{checkpoint_name()}_ fold - {fold+1}_{CHECKPOINT_KEYWORD}.pt', cp_folder=get_checkpoint_dir())\n",
        "\n",
        "        del model, train_dataset, val_dataset\n",
        "        if Config.device!=\"cpu\":\n",
        "            torch.cuda.empty_cache()\n",
        "        gc.collect()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nVa5bWAZjmCq"
      },
      "source": [
        "# train"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QXwSy9X2hrYt",
        "outputId": "c269a85b-0a0d-400b-9292-ed3db1d1b97d"
      },
      "source": [
        "df=pd.read_csv(\"dataset/squad_train.csv\")\n",
        "len(df)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "11598"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "referenced_widgets": [
            "aa60b9b9f6414c2faa84755f12274c19",
            "bc9ad37a618e484bb3c18413314c2594",
            "b480cf74784749bbb695b00ccfb79d43",
            "a6665a5398004fa2bc1437ac69b138b2",
            "d6f346e5f82c46aa81e4bede571581f6",
            "606df5aca1a3410ebc694cfcbc6a2df6",
            "76a55d15ebe94e8488e43fa0169bf1d9",
            "c0164eadd2254d84bd6b20a5f98daeec",
            "35d6ab8b7cb44cbb8ae6e4a0e621a15b",
            "3dbb39a901cc4c3098251a09c62a0430",
            "50f61fae634e476db0d0ba98def033b9",
            "47b060ccd3084d078d2026f9103aa880",
            "aa08522dfc654b30b6ef401d4c69bfae",
            "65d01253bfdb48cd9a6488a02edb34db",
            "c6258879d4244f17865487cb681ccbcc",
            "cb9130e18201498eaae12aecb12feccb",
            "1c598e906ba343c59a6b4e907c4d0bf3",
            "0dd4924490b447409073a6299501a8a0",
            "687724e9265f4ab0887ad07b1a82cdfa",
            "6d344bab148b482ab01ccc9be5cf9636",
            "fe2ce756730246db9b9cd3530c350cea",
            "692ebc7123b449f7852ea2eca3a4fb97",
            "0fba8ac13caa4218962ebf4e5550216e",
            "a3e62f3381a04a0ead5230cb3109a96c",
            "dfda4c1ea0a64b76b6a1994fcddc0f93",
            "3f10cc1852964e99935be6139d7024da",
            "ad8567e9f8ff4d30ad73382bb5117ea0",
            "c0ece211fc7d48db871b7a830e2f3230",
            "b62dd92a7bb743b18f8541f072936db7",
            "5ee56bd3938a44378aaa320c5220fd9b",
            "9e1787b7201f4ef9815c5f0fb1dfd23f",
            "0caf142e0aad4a2ba4581d0a7eff5593"
          ]
        },
        "id": "4v4UrMHWhxaR",
        "outputId": "841e6272-3f84-4eaa-fb04-002006d65f7b"
      },
      "source": [
        "import pdb\n",
        "CHECKPOINT_KEYWORD=\"cls_only\"\n",
        "folds=k_fold(df,save=True, config=Config)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "aa60b9b9f6414c2faa84755f12274c19",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=570.0, style=ProgressStyle(description_…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "35d6ab8b7cb44cbb8ae6e4a0e621a15b",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=435779157.0, style=ProgressStyle(descri…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "Some weights of the model checkpoint at bert-base-cased were not used when initializing BertModel: ['cls.predictions.bias', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.decoder.weight', 'cls.seq_relationship.bias', 'cls.seq_relationship.weight']\n",
            "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
            "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Loading fold checkpoint\n",
            "loaded checkpoint\n",
            "199 layers frozen!\n",
            "\n",
            "-------------   Fold 1 / 5  -------------\n",
            "\n",
            ">>9278 training texts<<\n",
            ">>2320 validation texts<<\n",
            ">>593668 trainable parameters<<\n",
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "1c598e906ba343c59a6b4e907c4d0bf3",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Epoch 1/15', max=15.0, style=ProgressStyle(description_wi…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "dfda4c1ea0a64b76b6a1994fcddc0f93",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='training iterations', max=290.0, style=ProgressStyle(desc…"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IgtfHc3Tguun"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-A4j6Z3NerDn"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}