TriBlank_Public.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "TriBlank-Public.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "71d6745389204b64bc67363c87221894": {
          "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_4494c479c65c4de8a15765ddfbd33bed",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_b4fe3a34532f41b3afdd7dfa90399503",
              "IPY_MODEL_f98efbd5cd534269b18749f0d76787cd"
            ]
          }
        },
        "4494c479c65c4de8a15765ddfbd33bed": {
          "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
          }
        },
        "b4fe3a34532f41b3afdd7dfa90399503": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_bf2a9750442b418eb5cdc2c037826f5a",
            "_dom_classes": [],
            "description": "100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "",
            "max": 3053,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 3053,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_46e19b037182480b8f190f5ba3156ae4"
          }
        },
        "f98efbd5cd534269b18749f0d76787cd": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_c86b52a74ff544c29828ddbad6971b2d",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 3053/3053 [00:02&lt;00:00, 1135.28it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_ffc05d3685fb4f7eb63239003a68574c"
          }
        },
        "bf2a9750442b418eb5cdc2c037826f5a": {
          "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"
          }
        },
        "46e19b037182480b8f190f5ba3156ae4": {
          "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
          }
        },
        "c86b52a74ff544c29828ddbad6971b2d": {
          "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"
          }
        },
        "ffc05d3685fb4f7eb63239003a68574c": {
          "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
          }
        },
        "bfc937a826f44915953bd80a8d872a61": {
          "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_f49c90fd95344af5958658724c0f1801",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_0523b7e53b8a4fc090a40de851a4a234",
              "IPY_MODEL_d3f492143a094c139080e03ef3ed6bc8"
            ]
          }
        },
        "f49c90fd95344af5958658724c0f1801": {
          "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
          }
        },
        "0523b7e53b8a4fc090a40de851a4a234": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_448ccbdce0d140b5870647818fc14100",
            "_dom_classes": [],
            "description": "100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "",
            "max": 3053,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 3053,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_8fa03e4ecb244735ba3d7ede167e7d09"
          }
        },
        "d3f492143a094c139080e03ef3ed6bc8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_cc58ab62616049e099734d16e431529a",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 3053/3053 [00:02&lt;00:00, 1090.36it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_1bfa767eb73f46dc820eb2599ab71342"
          }
        },
        "448ccbdce0d140b5870647818fc14100": {
          "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"
          }
        },
        "8fa03e4ecb244735ba3d7ede167e7d09": {
          "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
          }
        },
        "cc58ab62616049e099734d16e431529a": {
          "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"
          }
        },
        "1bfa767eb73f46dc820eb2599ab71342": {
          "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
          }
        },
        "ad7548f36e53491792341002cb67db3d": {
          "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_f06cb614e85b44bc9c32708a938c7e81",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_3f6245f27c374e109c0f23df7514a67e",
              "IPY_MODEL_0946c6febbfa41aab03a6747d102dc5b"
            ]
          }
        },
        "f06cb614e85b44bc9c32708a938c7e81": {
          "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
          }
        },
        "3f6245f27c374e109c0f23df7514a67e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_4b6c536cc482463199256e4ec1776b69",
            "_dom_classes": [],
            "description": "100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "",
            "max": 1000,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 1000,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_bb44a5e32e6b45b3a1c4f045ef1856c1"
          }
        },
        "0946c6febbfa41aab03a6747d102dc5b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_46508c2a18bc403d94c69c5cf8c2addd",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 1000/1000 [00:01&lt;00:00, 996.75it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_7ee0e319c7b14619b4661892ff29b27c"
          }
        },
        "4b6c536cc482463199256e4ec1776b69": {
          "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"
          }
        },
        "bb44a5e32e6b45b3a1c4f045ef1856c1": {
          "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
          }
        },
        "46508c2a18bc403d94c69c5cf8c2addd": {
          "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"
          }
        },
        "7ee0e319c7b14619b4661892ff29b27c": {
          "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
          }
        },
        "1d60366ee78140b2b14ef7979f1ce50f": {
          "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_ee9dae7c7fc14059a4ce9a0243a629f8",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_ba4e7df66c2544cd93eb33b808e50608",
              "IPY_MODEL_7565cedda838479db980cb9ff10162e2"
            ]
          }
        },
        "ee9dae7c7fc14059a4ce9a0243a629f8": {
          "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
          }
        },
        "ba4e7df66c2544cd93eb33b808e50608": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_bab7cb5804ae4f8f92dab5d2904c17b9",
            "_dom_classes": [],
            "description": "100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "",
            "max": 1000,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 1000,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_a4d1a5553600459b952d3f8b762aee1c"
          }
        },
        "7565cedda838479db980cb9ff10162e2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_e3f739495c38431eacd40db10f5ec50a",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 1000/1000 [00:00&lt;00:00, 1083.55it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_900edad658894be0a28c458b33535af4"
          }
        },
        "bab7cb5804ae4f8f92dab5d2904c17b9": {
          "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"
          }
        },
        "a4d1a5553600459b952d3f8b762aee1c": {
          "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
          }
        },
        "e3f739495c38431eacd40db10f5ec50a": {
          "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"
          }
        },
        "900edad658894be0a28c458b33535af4": {
          "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
          }
        },
        "91522f7fccd8465d981bcdd131d2b632": {
          "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_7c06a37ac987432c9a7b465dac5c8ade",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_988ccc3c53a84fb695c06c204b11e360",
              "IPY_MODEL_6ba4a78c39ff44faa84284f847703b35"
            ]
          }
        },
        "7c06a37ac987432c9a7b465dac5c8ade": {
          "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
          }
        },
        "988ccc3c53a84fb695c06c204b11e360": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_7343f62e67564244ab00622810726c84",
            "_dom_classes": [],
            "description": "100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "",
            "max": 1000,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 1000,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_b9050348749142ecb032ccced1512394"
          }
        },
        "6ba4a78c39ff44faa84284f847703b35": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_c13e13b40adf4b77bd608c47db058cc1",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 1000/1000 [00:00&lt;00:00, 1016.07it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_c452d3c4613d48b8899a6f045e0d5c4b"
          }
        },
        "7343f62e67564244ab00622810726c84": {
          "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"
          }
        },
        "b9050348749142ecb032ccced1512394": {
          "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
          }
        },
        "c13e13b40adf4b77bd608c47db058cc1": {
          "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"
          }
        },
        "c452d3c4613d48b8899a6f045e0d5c4b": {
          "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": "zchaASifUMGH"
      },
      "source": [
        "# Initialization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q3FeZq0nA2lw"
      },
      "source": [
        "%%capture\n",
        "!pip install datasets transformers\n",
        "!pip install torch torchvision torchaudio"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E8_JRDgfBZTJ"
      },
      "source": [
        "import datasets\n",
        "\n",
        "from transformers import BertTokenizer, BertModel, BertConfig\n",
        "\n",
        "import torch\n",
        "from torch import nn\n",
        "from torch import optim\n",
        "import torch.nn.functional as F\n",
        "\n",
        "import random\n",
        "\n",
        "import pprint as pprintmodule\n",
        "pp = pprintmodule.PrettyPrinter(indent=2)\n",
        "pprint = pp.pprint\n",
        "\n",
        "from tqdm.notebook import tqdm\n",
        "\n",
        "import gc\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wXJArIOEw5G1"
      },
      "source": [
        "DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ciE0y3m4UIHI"
      },
      "source": [
        "# Google Drive / Persistence"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uiiJEchBWBfL",
        "outputId": "8f022209-3153-44da-a6ca-1b639ddb8d37"
      },
      "source": [
        "### Google Colab Model Persistence\n",
        "\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "\n",
        "MODEL_PERSIST_PREFIX = 'drive/MyDrive/CS6740/Final Project/saved_models/'"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Mounted at /content/drive\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ia2vzBFwWiqD"
      },
      "source": [
        "from datetime import datetime\n",
        "import os\n",
        "\n",
        "def get_path(model_name):\n",
        "    if 'MODEL_PERSIST_PREFIX' in globals():\n",
        "        path = os.path.join(MODEL_PERSIST_PREFIX, model_name)\n",
        "    else:\n",
        "        path = model_name\n",
        "    return path\n",
        "\n",
        "def get_default_name(model):\n",
        "    return type(model).__name__ + '-' + datetime.now().strftime('%m-%d@%H-%M')\n",
        "\n",
        "def save_model(model, name=None):\n",
        "    name = name or (get_default_name(model) + '.pt')\n",
        "    path = get_path(name)\n",
        "\n",
        "    print('Saving', name)\n",
        "    torch.save(model.state_dict(), path)\n",
        "    return model\n",
        "\n",
        "def load_model(model, name):\n",
        "    path = get_path(name)\n",
        "    model.load_state_dict(torch.load(path, map_location=torch.device(DEVICE)))\n",
        "    return model\n",
        "\n",
        "def get_default_checkpoint_name(model, epoch):\n",
        "    return 'Checkpoint-{}-e{}.pt'.format(get_default_name(model), epoch)\n",
        "\n",
        "def save_checkpoint(model, optim, epoch):\n",
        "    name = get_default_checkpoint_name(model, epoch)\n",
        "    path = get_path(name)\n",
        "\n",
        "    print('Saving', name)\n",
        "    torch.save(\n",
        "        {\n",
        "            'epoch': epoch,\n",
        "            'model': model.state_dict(),\n",
        "            'optim': optim.state_dict()\n",
        "        }, \n",
        "        path\n",
        "    )\n",
        "    return model\n",
        "\n",
        "def load_checkpoint(model, optim, name):\n",
        "    path = get_path(name)\n",
        "\n",
        "    checkpoint = torch.load(path, map_location=torch.device(DEVICE))\n",
        "\n",
        "    model.load_state_dict(checkpoint['model'])\n",
        "    optim.load_state_dict(checkpoint['optim'])\n",
        "    return checkpoint['epoch']\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rhhd9TDQUQsS"
      },
      "source": [
        "# The Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ju3pdBcyBf0Q"
      },
      "source": [
        "%%capture\n",
        "DOCRED = datasets.load_dataset('docred')\n",
        "DOCRED_TRAIN = DOCRED['train_annotated']\n",
        "DOCRED_VALID = DOCRED['validation']\n",
        "DOCRED_TEST = DOCRED['test']"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aFuZMjT0FLRb",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17,
          "referenced_widgets": [
            "71d6745389204b64bc67363c87221894",
            "4494c479c65c4de8a15765ddfbd33bed",
            "b4fe3a34532f41b3afdd7dfa90399503",
            "f98efbd5cd534269b18749f0d76787cd",
            "bf2a9750442b418eb5cdc2c037826f5a",
            "46e19b037182480b8f190f5ba3156ae4",
            "c86b52a74ff544c29828ddbad6971b2d",
            "ffc05d3685fb4f7eb63239003a68574c",
            "bfc937a826f44915953bd80a8d872a61",
            "f49c90fd95344af5958658724c0f1801",
            "0523b7e53b8a4fc090a40de851a4a234",
            "d3f492143a094c139080e03ef3ed6bc8",
            "448ccbdce0d140b5870647818fc14100",
            "8fa03e4ecb244735ba3d7ede167e7d09",
            "cc58ab62616049e099734d16e431529a",
            "1bfa767eb73f46dc820eb2599ab71342",
            "ad7548f36e53491792341002cb67db3d",
            "f06cb614e85b44bc9c32708a938c7e81",
            "3f6245f27c374e109c0f23df7514a67e",
            "0946c6febbfa41aab03a6747d102dc5b",
            "4b6c536cc482463199256e4ec1776b69",
            "bb44a5e32e6b45b3a1c4f045ef1856c1",
            "46508c2a18bc403d94c69c5cf8c2addd",
            "7ee0e319c7b14619b4661892ff29b27c",
            "1d60366ee78140b2b14ef7979f1ce50f",
            "ee9dae7c7fc14059a4ce9a0243a629f8",
            "ba4e7df66c2544cd93eb33b808e50608",
            "7565cedda838479db980cb9ff10162e2",
            "bab7cb5804ae4f8f92dab5d2904c17b9",
            "a4d1a5553600459b952d3f8b762aee1c",
            "e3f739495c38431eacd40db10f5ec50a",
            "900edad658894be0a28c458b33535af4",
            "91522f7fccd8465d981bcdd131d2b632",
            "7c06a37ac987432c9a7b465dac5c8ade",
            "988ccc3c53a84fb695c06c204b11e360",
            "6ba4a78c39ff44faa84284f847703b35",
            "7343f62e67564244ab00622810726c84",
            "b9050348749142ecb032ccced1512394",
            "c13e13b40adf4b77bd608c47db058cc1",
            "c452d3c4613d48b8899a6f045e0d5c4b"
          ]
        },
        "outputId": "fc85b37f-c783-4272-d009-57e3ecec0340"
      },
      "source": [
        "def get_all_relation_ids(dataset):\n",
        "    ids = set()\n",
        "    for example in tqdm(dataset, leave=False):\n",
        "        ids |= set(example['labels']['relation_id'])\n",
        "    ids = list(ids)\n",
        "    ids.sort()\n",
        "    return ids\n",
        "\n",
        "def relation_ids_to_text(dataset):\n",
        "    rel2txt = {}\n",
        "    for example in tqdm(dataset, leave=False):\n",
        "        labels = example['labels']\n",
        "        for rel, txt in zip(labels['relation_id'], labels['relation_text']):\n",
        "            rel2txt[rel] = txt\n",
        "    return rel2txt\n",
        "\n",
        "ALL_RELATION_IDS = get_all_relation_ids(DOCRED_TRAIN)\n",
        "RELATION_ID_TO_TEXT = relation_ids_to_text(DOCRED_TRAIN)\n",
        "\n",
        "assert set(get_all_relation_ids(DOCRED['validation'])).issubset(set(ALL_RELATION_IDS))\n",
        "assert set(get_all_relation_ids(DOCRED['test'])).issubset(set(ALL_RELATION_IDS))\n",
        "assert set(get_all_relation_ids(DOCRED['train_distant'])).issubset(set(ALL_RELATION_IDS))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "71d6745389204b64bc67363c87221894",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=3053.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\r"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "bfc937a826f44915953bd80a8d872a61",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=3053.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\r"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "ad7548f36e53491792341002cb67db3d",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=1000.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\r"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "1d60366ee78140b2b14ef7979f1ce50f",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=1000.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\r"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "91522f7fccd8465d981bcdd131d2b632",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=1000.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\r"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QCvmga9cm4Fr"
      },
      "source": [
        "# Tokenizer and Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CPHBNQaY1Gly"
      },
      "source": [
        "ENT_BGN = ['[E0]', '[E1]', '[E2]']\n",
        "ENT_END = ['[/E0]', '[/E1]', '[/E2]']\n",
        "ENT_BLN = '[BLANK]'\n",
        "\n",
        "def create_tokenizer_bert_model():\n",
        "    tokenizer = BertTokenizer.from_pretrained('bert-large-uncased')\n",
        "    bert_model = BertModel.from_pretrained('bert-large-uncased').to(DEVICE)\n",
        "\n",
        "    # add the entity markers and blank\n",
        "\n",
        "    tokenizer.add_special_tokens({\n",
        "        'additional_special_tokens': ENT_BGN + ENT_END + [ENT_BLN]\n",
        "    })\n",
        "    bert_model.resize_token_embeddings(len(tokenizer))\n",
        "\n",
        "    return tokenizer, bert_model\n",
        "\n",
        "def ent_token_ids(tokenizer):\n",
        "    return {\n",
        "        'bgn': tokenizer.additional_special_tokens_ids[:3],\n",
        "        'end': tokenizer.additional_special_tokens_ids[3:6],\n",
        "        'bln': tokenizer.additional_special_tokens_ids[6]\n",
        "    }"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "V6lDHLtEsqIX",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 146
        },
        "outputId": "241f1b2e-43d3-49cd-b2f7-8162ecd917b2"
      },
      "source": [
        "def format_example(\n",
        "    example, \n",
        "    eid_and_bln\n",
        "):\n",
        "    # flatten example['sents']\n",
        "    sents = []\n",
        "    sentence_offsets = [0]\n",
        "    for sentence in example['sents']:\n",
        "        sents += sentence\n",
        "        sentence_offsets.append(sentence_offsets[-1] + len(sentence))\n",
        "\n",
        "    entities = example['vertexSet']\n",
        "\n",
        "    fmt = lambda entity, index: (\n",
        "        sentence_offsets[entity['sent_id']] + entity['pos'][0], \n",
        "        sentence_offsets[entity['sent_id']] + entity['pos'][1], \n",
        "        index\n",
        "    )\n",
        "\n",
        "    # an associated list from (start, end) positions to entity index { 0, 1, 2 }\n",
        "    pos_to_ek = []\n",
        "    for i, (eid, _) in enumerate(eid_and_bln):\n",
        "        for entity in entities[eid]:\n",
        "            pos_to_ek.append(fmt(entity, i))\n",
        "    pos_to_ek.sort(key=lambda x: x[0])\n",
        "\n",
        "    # remove overlap\n",
        "    indexes = set()\n",
        "    for i, (bgn, end, _) in reversed(list(enumerate(pos_to_ek))):\n",
        "        if len(set(range(bgn, end)) & indexes) != 0:\n",
        "            del pos_to_ek[i]\n",
        "        else:\n",
        "            indexes |= set(range(bgn, end))\n",
        "    \n",
        "    # assert no overlap\n",
        "    # indexes = set()\n",
        "    # for bgn, end, _ in pos_to_ek:\n",
        "    #     assert len(set(range(bgn, end)) & indexes) == 0, (bgn, end, indexes)\n",
        "    #     indexes |= set(range(bgn, end))\n",
        "\n",
        "    # insert entity markers\n",
        "    for bgn, end, ei in reversed(pos_to_ek):\n",
        "        sents.insert(end, ENT_END[ei])\n",
        "        sents.insert(bgn, ENT_BGN[ei])\n",
        "\n",
        "        if eid_and_bln[ei][1]:\n",
        "            # blank the entity\n",
        "            sents[bgn + 1] = ENT_BLN\n",
        "            for i in reversed(range(bgn + 2, end + 1)):\n",
        "                del sents[i]\n",
        "\n",
        "    return ' '.join(sents)\n",
        "\n",
        "format_example(DOCRED_TRAIN[25], [(0, False), (1, False), (2, False)])\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'[E0] Chelsea [/E0] was an early [E1] 1970s [/E1] band from [E2] New York City [/E2] , best known for being the band of drummer Peter Criss before he joined Kiss . They released one album , the self - titled album Chelsea in 1971 and then collapsed during the recording of their unreleased second album . In August 1971 , the band became Lips ( a trio consisting of Criss and his [E0] Chelsea [/E0] bandmates Michael Benvenga and Stan Penridge ) . By the spring of 1973 , Lips was just the duo of Criss and Penridge and eventually disbanded completely . Their sound has been compared to the Moody Blues and Procol Harum . In 1973 , Pete Shepley & Mike Brand recorded an unreleased album which included post - Chelsea Michael Benvenga , a pre - Kiss Peter Criss , and on two songs Gene Simmons as session musicians . It was titled Captain Sanity .'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r1lkt4WrE4ei"
      },
      "source": [
        "class BertWithEntityStartPooling(nn.Module):\n",
        "\n",
        "    def __init__(self, bert, ent_token_ids):\n",
        "        super(BertWithEntityStartPooling, self).__init__()\n",
        "\n",
        "        self._bert = bert\n",
        "        self._ent_bgn_ids = ent_token_ids['bgn']\n",
        "        self._h = self._bert.config.hidden_size\n",
        "        self.hidden_size = 2 * self._h\n",
        "\n",
        "    def forward(self, input):\n",
        "        batched_hidden_states = self._bert(**input).last_hidden_state\n",
        "        batched_result = []\n",
        "\n",
        "        for input_ids, masks, hidden_states in zip(\n",
        "            input['input_ids'], \n",
        "            input['attention_mask'],\n",
        "            batched_hidden_states\n",
        "        ):\n",
        "            ei_to_h = [[], [], []]\n",
        "\n",
        "            for input_id, mask, h in zip(input_ids, masks, hidden_states):\n",
        "                if not mask:\n",
        "                    continue\n",
        "\n",
        "                try:\n",
        "                    ei = self._ent_bgn_ids.index(input_id)\n",
        "                    ei_to_h[ei].append(h)\n",
        "                except:\n",
        "                    pass\n",
        "\n",
        "            for ei, hs in enumerate(ei_to_h):\n",
        "                if hs: \n",
        "                    ei_to_h[ei] = torch.cat([\n",
        "                        h.unsqueeze(0) \n",
        "                        for h in hs\n",
        "                    ])\n",
        "                else:\n",
        "                    ei_to_h[ei] = None\n",
        "\n",
        "            for ei, hs in enumerate(ei_to_h):\n",
        "                if hs != None:\n",
        "                    hs = torch.transpose(hs, 0, 1).unsqueeze(0)\n",
        "                    ei_to_h[ei] = F.max_pool1d(hs, hs.shape[-1]).squeeze()\n",
        "\n",
        "                    assert ei_to_h[ei].shape == (self._h,)\n",
        "                else:\n",
        "                    ei_to_h[ei] = torch.zeros(self._h).to(DEVICE)\n",
        "\n",
        "            output = []\n",
        "            for ei, ej in [(0, 1), (0, 2), (1, 2)]:\n",
        "                output.append(torch.cat((ei_to_h[ei], ei_to_h[ej])).unsqueeze(0))\n",
        "\n",
        "            output = torch.cat(output).unsqueeze(0)\n",
        "            batched_result.append(output)\n",
        "\n",
        "        return torch.cat(batched_result)\n",
        "\n",
        "\n",
        "class FullyConnectedLayer(nn.Module):\n",
        "\n",
        "    \"\"\" \n",
        "    A fully connected layer with an optional activation function.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, input, hidden, output, activation_fn=None):\n",
        "        super(FullyConnectedLayer, self).__init__()\n",
        "\n",
        "        self._linear1 = nn.Linear(input, hidden)\n",
        "        self._activation_fn = activation_fn or nn.Identity()\n",
        "        self._linear2 = nn.Linear(hidden, output)\n",
        "\n",
        "    def forward(self, input):\n",
        "        return self._linear2(self._activation_fn(self._linear1(input)))\n",
        "\n",
        "\n",
        "class PerClassScore(nn.Module):\n",
        "\n",
        "    def __init__(self, pcr_size):\n",
        "        super(PerClassScore, self).__init__()\n",
        "        self._pcr = nn.Parameter(torch.randn(pcr_size, len(ALL_RELATION_IDS)))\n",
        "    \n",
        "    def forward(self, input):\n",
        "        result = torch.matmul(input, self._pcr)\n",
        "        result = F.log_softmax(result, dim=2)\n",
        "        return result\n",
        "\n",
        "\n",
        "class TriBlank(nn.Module):\n",
        "\n",
        "    def __init__(self, bert, ent_token_ids):\n",
        "        super(TriBlank, self).__init__()\n",
        "\n",
        "        bwesp = BertWithEntityStartPooling(bert, ent_token_ids)\n",
        "        h = bwesp.hidden_size\n",
        "        fcl = FullyConnectedLayer(h, h // 2, h // 2)\n",
        "        pcr = PerClassScore(h // 2)\n",
        "\n",
        "        self._seq = nn.Sequential(\n",
        "            bwesp,\n",
        "            fcl, \n",
        "            pcr\n",
        "        )\n",
        "\n",
        "    def forward(self, input):\n",
        "        return self._seq(input)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WU3xLl94mEN6"
      },
      "source": [
        "def run_model(model, tokenizer, batch, eid_and_bln=[(0, False), (1, False), (2, True)]):\n",
        "    model.eval()\n",
        "    \n",
        "    input = tokenizer(\n",
        "        [\n",
        "            format_example(example, eid_and_bln)\n",
        "            for example in batch\n",
        "        ],\n",
        "        padding=True,\n",
        "        return_tensors='pt'\n",
        "    ).to(DEVICE)\n",
        "    \n",
        "    output = model(input)\n",
        "    print(output)\n",
        "    print()\n",
        "    for i in range(len(batch)):\n",
        "        example = batch[i]\n",
        "        for ei in range(len(eid_and_bln) - 1):\n",
        "            for ej in range(ei + 1, len(eid_and_bln)):\n",
        "                assert ei + ej - 1 in (0, 1, 2)\n",
        "                pred = torch.argmax(output[i][ei + ej - 1]).item()\n",
        "                print(example['vertexSet'][eid_and_bln[ei][0]][0]['name'])\n",
        "                print(RELATION_ID_TO_TEXT[ALL_RELATION_IDS[pred]])\n",
        "                print(example['vertexSet'][eid_and_bln[ej][0]][0]['name'])\n",
        "                print()\n",
        "\n",
        "def namespace():\n",
        "    tokenizer, bert_model = create_tokenizer_bert_model()\n",
        "    model = TriBlank(bert_model, ent_token_ids(tokenizer)).to(DEVICE)\n",
        "    run_model(model, tokenizer, [DOCRED_TRAIN[5], DOCRED_TRAIN[6]])\n",
        "\n",
        "    torch.cuda.empty_cache()\n",
        "    gc.collect()\n",
        "\n",
        "# namespace()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KcKwzlymOeQO"
      },
      "source": [
        "# Re2\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DbzTD2rjePlN"
      },
      "source": [
        "def map_to_re2_dataset(dataset):\n",
        "    re2_dataset = []\n",
        "\n",
        "    for i, example in enumerate(tqdm(dataset)):\n",
        "        labels = example['labels']\n",
        "\n",
        "        for head, rel, tail in zip(labels['head'], labels['relation_id'], labels['tail']):\n",
        "            re2_dataset.append((i, head, ALL_RELATION_IDS.index(rel), tail))\n",
        "\n",
        "    return re2_dataset\n",
        "\n",
        "def form_re2_batch(docred_dataset, re2_dataset, tokenizer, start_index, batch_size, blank_alpha=None):\n",
        "    blank_alpha = blank_alpha or 0.7\n",
        "    end_index = min(start_index + batch_size, len(re2_dataset))\n",
        "    \n",
        "    examples = [\n",
        "        format_example(\n",
        "            docred_dataset[i], \n",
        "            [\n",
        "                (e0, random.random() < blank_alpha), \n",
        "                (e1, random.random() < blank_alpha)\n",
        "            ]\n",
        "        )\n",
        "        for (i, e0, _, e1) in re2_dataset[start_index:end_index]\n",
        "    ]\n",
        "\n",
        "    examples = tokenizer(examples, padding=True, return_tensors='pt', truncation=True).to(DEVICE)\n",
        "\n",
        "    gold = torch.tensor([x[2] for x in re2_dataset[start_index:end_index]]).to(DEVICE)\n",
        "\n",
        "    return (examples, gold)\n",
        "\n",
        "def iter_re2_batches(docred_dataset, re2_dataset, tokenizer, batch_size, blank_alpha=None):\n",
        "    for start_index in tqdm(range(0, len(re2_dataset), batch_size)):\n",
        "        yield form_re2_batch(\n",
        "            docred_dataset, \n",
        "            re2_dataset, \n",
        "            tokenizer, \n",
        "            start_index, \n",
        "            batch_size, \n",
        "            blank_alpha=blank_alpha\n",
        "        )\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6IRZhZq-ad3_"
      },
      "source": [
        "def train_epoch_re2(\n",
        "    model, \n",
        "    tokenizer, \n",
        "    optim, \n",
        "    docred_dataset, \n",
        "    re2_dataset, \n",
        "    batch_size, \n",
        "    max_grad_norm=None,\n",
        "    blank_alpha=None\n",
        "):\n",
        "    model.train()\n",
        "    \n",
        "    for examples, gold in iter_re2_batches(\n",
        "            docred_dataset, \n",
        "            re2_dataset, \n",
        "            tokenizer, \n",
        "            batch_size, \n",
        "            blank_alpha=blank_alpha\n",
        "        ):\n",
        "\n",
        "        optim.zero_grad()\n",
        "        output = torch.transpose(model(examples), 0, 1)\n",
        "        loss = F.nll_loss(output[0], gold)\n",
        "        loss.backward()\n",
        "        if max_grad_norm:\n",
        "            nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)\n",
        "        optim.step()\n",
        "\n",
        "        del output\n",
        "        del loss\n",
        "        torch.cuda.empty_cache()\n",
        "        gc.collect()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "du7zzvUjIoYb"
      },
      "source": [
        "def eval_contingency_table_re2(model, tokenizer, dataset, re2_dataset, blank=False):\n",
        "    model.eval()\n",
        "\n",
        "    contingency_table = [[0] * len(ALL_RELATION_IDS) for _ in range(len(ALL_RELATION_IDS))]\n",
        "\n",
        "    with torch.no_grad():\n",
        "        for i, e0, rel, e1 in tqdm(re2_dataset):\n",
        "            input = tokenizer(\n",
        "                [format_example(dataset[i], [(e0, blank), (e1, blank)])],\n",
        "                padding=True,\n",
        "                return_tensors='pt',\n",
        "                truncation=True\n",
        "            ).to(DEVICE)\n",
        "    \n",
        "            output = model(input).squeeze(0)\n",
        "            pred = torch.argmax(output[0]).item()\n",
        "            contingency_table[rel][pred] += 1\n",
        "\n",
        "    return contingency_table"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j7yib8sNaX8z"
      },
      "source": [
        "# Re3"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "akMxdYDKadsr"
      },
      "source": [
        "def all_index_triples(iter):\n",
        "    for i in range(len(iter)):\n",
        "        for j in range(i + 1, len(iter)):\n",
        "            for k in range(j + 1, len(iter)):\n",
        "                yield (i, j, k)\n",
        "\n",
        "def is_related_set(example):\n",
        "    related_indexes = set()\n",
        "    label = example['labels']\n",
        "    for head, tail in zip(label['head'], label['tail']):\n",
        "        related_indexes.add((head, tail))\n",
        "        related_indexes.add((tail, head))\n",
        "    return related_indexes\n",
        "\n",
        "def extract_triples(example):\n",
        "    triples = set()\n",
        "    is_related = is_related_set(example)\n",
        "    for i, j, k in all_index_triples(example['vertexSet']):\n",
        "        if (i, j) in is_related and (j, k) in is_related:\n",
        "            triples.add((i, j, k))\n",
        "    return triples\n",
        "\n",
        "def extract_triway(example):\n",
        "    triples = set()\n",
        "    is_related = is_related_set(example)\n",
        "    for i, j, k in all_index_triples(example['vertexSet']):\n",
        "        if (i, j) in is_related and (j, k) in is_related and (i, k) in is_related:\n",
        "            triples.add((i, j, k))\n",
        "    return triples\n",
        "\n",
        "def print_relation(example, index_pair):\n",
        "    vertex_set = example['vertexSet']\n",
        "    labels = example['labels']\n",
        "    i, j = index_pair\n",
        "    for head, relation, tail in zip(labels['head'], labels['relation_text'], labels['tail']):\n",
        "        head_rep = vertex_set[head][0]['name']\n",
        "        tail_rep = vertex_set[tail][0]['name']\n",
        "\n",
        "        if (i, j) == (head, tail):\n",
        "            print('{} {} {}'.format(head_rep, relation, tail_rep))\n",
        "            \n",
        "        elif (j, i) == (head, tail):\n",
        "            print('{} {} {}'.format(tail_rep, relation, head_rep))\n",
        "\n",
        "def print_triple_relation(example, index_triple):\n",
        "    i, j, k = index_triple\n",
        "    print_relation(example, (i, j))\n",
        "    print_relation(example, (j, k))\n",
        "    print_relation(example, (i, k))\n",
        "            \n",
        "def print_extract_triway(example):\n",
        "    for index_triple in extract_triway(example):\n",
        "        print(index_triple)\n",
        "        print_triple_relation(example, index_triple)\n",
        "        print()\n",
        "\n",
        "# extract_triway(DOCRED_TRAIN[0])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oTUVFRvNVQEN"
      },
      "source": [
        "def extract_labeled_edges(example):\n",
        "    labels = example['labels']\n",
        "    labeled_edges = {}\n",
        "    \n",
        "    for head, relation, tail in zip(labels['head'], labels['relation_id'], labels['tail']):\n",
        "        key = (head, tail)\n",
        "\n",
        "        if key not in labeled_edges:\n",
        "            labeled_edges[key] = []\n",
        "\n",
        "        labeled_edges[key].append(relation)\n",
        "\n",
        "    return labeled_edges\n",
        "\n",
        "def map_to_re3_dataset(dataset):\n",
        "    re3_dataset = []\n",
        "\n",
        "    for dataset_index, example in enumerate(tqdm(dataset)):\n",
        "        labeled_edges = extract_labeled_edges(example)\n",
        "\n",
        "        for i in range(len(example['vertexSet'])):\n",
        "            for j in range(len(example['vertexSet'])):\n",
        "                if i == j:\n",
        "                    continue\n",
        "                \n",
        "                for k in range(len(example['vertexSet'])):\n",
        "                    if i == k or j == k:\n",
        "                        continue\n",
        "\n",
        "                    if (i, j) not in labeled_edges or (j, k) not in labeled_edges:\n",
        "                        continue\n",
        "                        \n",
        "                    for rel_ij in labeled_edges[(i, j)]:\n",
        "                        for rel_jk in labeled_edges[(j, k)]:\n",
        "                            re3_dataset.append((\n",
        "                                dataset_index, \n",
        "                                i, \n",
        "                                ALL_RELATION_IDS.index(rel_ij),\n",
        "                                j,\n",
        "                                ALL_RELATION_IDS.index(rel_jk),\n",
        "                                k\n",
        "                            ))\n",
        "\n",
        "    return re3_dataset\n",
        "\n",
        "def form_re3_batch(docred_dataset, re3_dataset, tokenizer, start_index, batch_size, blank_alpha=None):\n",
        "    blank_alpha = blank_alpha or 0.7\n",
        "    end_index = min(start_index + batch_size, len(re3_dataset))\n",
        "    \n",
        "    examples = [\n",
        "        format_example(\n",
        "            docred_dataset[i], \n",
        "            [\n",
        "                (e0, random.random() < blank_alpha), \n",
        "                (e1, random.random() < blank_alpha),\n",
        "                (e2, random.random() < blank_alpha)\n",
        "            ]\n",
        "        )\n",
        "        for (i, e0, _, e1, _, e2) in re3_dataset[start_index:end_index]\n",
        "    ]\n",
        "\n",
        "    examples = tokenizer(examples, padding=True, return_tensors='pt', truncation=True).to(DEVICE)\n",
        "\n",
        "    gold_01 = torch.tensor([x[2] for x in re3_dataset[start_index:end_index]]).to(DEVICE)\n",
        "    gold_12 = torch.tensor([x[4] for x in re3_dataset[start_index:end_index]]).to(DEVICE)\n",
        "\n",
        "    return (examples, gold_01, gold_12)\n",
        "\n",
        "def iter_re3_batches(docred_dataset, re3_dataset, tokenizer, batch_size, blank_alpha=None):\n",
        "    for start_index in tqdm(range(0, len(re3_dataset), batch_size)):\n",
        "        yield form_re3_batch(\n",
        "            docred_dataset, \n",
        "            re3_dataset, \n",
        "            tokenizer, \n",
        "            start_index, \n",
        "            batch_size, \n",
        "            blank_alpha=blank_alpha\n",
        "        )\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c2f1kV9RBKf2"
      },
      "source": [
        "def train_epoch_re3(\n",
        "    model, \n",
        "    tokenizer, \n",
        "    optim, \n",
        "    docred_dataset, \n",
        "    re3_dataset, \n",
        "    batch_size, \n",
        "    max_grad_norm=None,\n",
        "    blank_alpha=None\n",
        "):\n",
        "    model.train()\n",
        "    \n",
        "    for examples, gold_01, gold_12 in iter_re3_batches(\n",
        "            docred_dataset, \n",
        "            re3_dataset, \n",
        "            tokenizer, \n",
        "            batch_size, \n",
        "            blank_alpha=blank_alpha\n",
        "        ):\n",
        "\n",
        "        optim.zero_grad()\n",
        "        output = torch.transpose(model(examples), 0, 1)\n",
        "        loss = F.nll_loss(output[0], gold_01) + F.nll_loss(output[2], gold_12)\n",
        "        loss.backward()\n",
        "        if max_grad_norm:\n",
        "            nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)\n",
        "        optim.step()\n",
        "\n",
        "        del output\n",
        "        del loss\n",
        "        torch.cuda.empty_cache()\n",
        "        gc.collect()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_hdI_EUtOiCr"
      },
      "source": [
        "def eval_contingency_table_re3(model, tokenizer, dataset, re3_dataset, blank=False):\n",
        "    model.eval()\n",
        "\n",
        "    contingency_table = [[0] * len(ALL_RELATION_IDS) for _ in range(len(ALL_RELATION_IDS))]\n",
        "\n",
        "    with torch.no_grad():\n",
        "        for i, e0, rel_01, e1, rel_12, e2 in tqdm(re3_dataset):\n",
        "            input = tokenizer(\n",
        "                [format_example(dataset[i], [(e0, blank), (e1, blank), (e2, blank)])],\n",
        "                padding=True,\n",
        "                return_tensors='pt',\n",
        "                truncation=True\n",
        "            ).to(DEVICE)\n",
        "    \n",
        "            output = model(input).squeeze(0)\n",
        "    \n",
        "            pred = torch.argmax(output[0]).item()\n",
        "            contingency_table[rel_01][pred] += 1\n",
        "    \n",
        "            pred = torch.argmax(output[2]).item()\n",
        "            contingency_table[rel_12][pred] += 1\n",
        "\n",
        "    return contingency_table\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NWFyjp5U6ylY"
      },
      "source": [
        "# Tri"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "edFy9UrL60wp"
      },
      "source": [
        "def map_to_tri_dataset(dataset):\n",
        "    tri_dataset = []\n",
        "\n",
        "    for dataset_index, example in enumerate(tqdm(dataset)):\n",
        "        labeled_edges = extract_labeled_edges(example)\n",
        "\n",
        "        for i in range(len(example['vertexSet'])):\n",
        "            for j in range(len(example['vertexSet'])):\n",
        "                if i == j:\n",
        "                    continue\n",
        "                \n",
        "                for k in range(len(example['vertexSet'])):\n",
        "                    if i == k or j == k:\n",
        "                        continue\n",
        "\n",
        "                    if (i, j) not in labeled_edges or (i, k) not in labeled_edges or (j, k) not in labeled_edges:\n",
        "                        continue\n",
        "                        \n",
        "                    for rel_ij in labeled_edges[(i, j)]:\n",
        "                        for rel_ik in labeled_edges[(i, k)]:\n",
        "                            for rel_jk in labeled_edges[(j, k)]:\n",
        "                                tri_dataset.append((\n",
        "                                    dataset_index, \n",
        "                                    i, \n",
        "                                    j, \n",
        "                                    k,\n",
        "                                    ALL_RELATION_IDS.index(rel_ij),\n",
        "                                    ALL_RELATION_IDS.index(rel_ik),\n",
        "                                    ALL_RELATION_IDS.index(rel_jk),\n",
        "                                ))\n",
        "\n",
        "    return tri_dataset\n",
        "\n",
        "def form_tri_batch(docred_dataset, tri_dataset, tokenizer, start_index, batch_size, blank_alpha=None):\n",
        "    blank_alpha = blank_alpha or 0.7\n",
        "    end_index = min(start_index + batch_size, len(tri_dataset))\n",
        "    \n",
        "    examples = [\n",
        "        format_example(\n",
        "            docred_dataset[i], \n",
        "            [\n",
        "                (e0, random.random() < blank_alpha), \n",
        "                (e1, random.random() < blank_alpha),\n",
        "                (e2, random.random() < blank_alpha)\n",
        "            ]\n",
        "        )\n",
        "        for (i, e0, e1, e2, _, _, _) in tri_dataset[start_index:end_index]\n",
        "    ]\n",
        "\n",
        "    examples = tokenizer(examples, padding=True, return_tensors='pt', truncation=True).to(DEVICE)\n",
        "\n",
        "    gold_01 = torch.tensor([x[4] for x in tri_dataset[start_index:end_index]]).to(DEVICE)\n",
        "    gold_02 = torch.tensor([x[5] for x in tri_dataset[start_index:end_index]]).to(DEVICE)\n",
        "    gold_12 = torch.tensor([x[6] for x in tri_dataset[start_index:end_index]]).to(DEVICE)\n",
        "\n",
        "    return (examples, gold_01, gold_02, gold_12)\n",
        "\n",
        "def iter_tri_batches(docred_dataset, tri_dataset, tokenizer, batch_size, blank_alpha=None):\n",
        "    for start_index in tqdm(range(0, len(tri_dataset), batch_size)):\n",
        "        yield form_tri_batch(\n",
        "            docred_dataset, \n",
        "            tri_dataset, \n",
        "            tokenizer, \n",
        "            start_index, \n",
        "            batch_size, \n",
        "            blank_alpha=blank_alpha\n",
        "        )\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4u1DwE6060wr"
      },
      "source": [
        "def train_epoch_tri(\n",
        "    model, \n",
        "    tokenizer, \n",
        "    optim, \n",
        "    docred_dataset, \n",
        "    tri_dataset, \n",
        "    batch_size, \n",
        "    max_grad_norm=None,\n",
        "    blank_alpha=None\n",
        "):\n",
        "    model.train()\n",
        "    \n",
        "    for examples, gold_01, gold_02, gold_12 in iter_tri_batches(\n",
        "            docred_dataset, \n",
        "            tri_dataset, \n",
        "            tokenizer, \n",
        "            batch_size, \n",
        "            blank_alpha=blank_alpha\n",
        "        ):\n",
        "\n",
        "        optim.zero_grad()\n",
        "        output = torch.transpose(model(examples), 0, 1)\n",
        "        loss = F.nll_loss(output[0], gold_01) + F.nll_loss(output[1], gold_02) + F.nll_loss(output[2], gold_12)\n",
        "        loss.backward()\n",
        "        if max_grad_norm:\n",
        "            nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)\n",
        "        optim.step()\n",
        "\n",
        "        del output\n",
        "        del loss\n",
        "        torch.cuda.empty_cache()\n",
        "        gc.collect()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Yo3e-NG-60wr"
      },
      "source": [
        "def eval_contingency_table_tri(model, tokenizer, dataset, tri_dataset, blank=False):\n",
        "    model.eval()\n",
        "\n",
        "    contingency_table = [[0] * len(ALL_RELATION_IDS) for _ in range(len(ALL_RELATION_IDS))]\n",
        "\n",
        "    with torch.no_grad():\n",
        "        for i, e0, e1, e2, rel_01, rel_02, rel_12 in tqdm(tri_dataset):\n",
        "            input = tokenizer(\n",
        "                [format_example(dataset[i], [(e0, blank), (e1, blank), (e2, blank)])],\n",
        "                padding=True,\n",
        "                return_tensors='pt',\n",
        "                truncation=True\n",
        "            ).to(DEVICE)\n",
        "    \n",
        "            output = model(input).squeeze(0)\n",
        "    \n",
        "            pred = torch.argmax(output[0]).item()\n",
        "            contingency_table[rel_01][pred] += 1\n",
        "    \n",
        "            pred = torch.argmax(output[1]).item()\n",
        "            contingency_table[rel_02][pred] += 1\n",
        "    \n",
        "            pred = torch.argmax(output[2]).item()\n",
        "            contingency_table[rel_12][pred] += 1\n",
        "\n",
        "    return contingency_table\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Gqs0Qb3tcFo3"
      },
      "source": [
        "def eval_contingency_table_all(model, tokenizer, dataset, blank=False):\n",
        "    def compute_acc(tbl):\n",
        "        return sum(tbl[i][i] for i in range(len(tbl))) / sum(tbl[i][j] for i in range(len(tbl)) for j in range(len(tbl)))\n",
        "        \n",
        "    re2_dataset = map_to_re2_dataset(dataset)\n",
        "    re2_tbl = eval_contingency_table_re2(model, tokenizer, dataset, re2_dataset, blank=blank)\n",
        "    re2_acc = compute_acc(re2_tbl)\n",
        "\n",
        "    re3_dataset = map_to_re3_dataset(dataset)\n",
        "    re3_tbl = eval_contingency_table_re3(model, tokenizer, dataset, re3_dataset, blank=blank)\n",
        "    re3_acc = compute_acc(re3_tbl)\n",
        "\n",
        "    tri_dataset = map_to_tri_dataset(dataset)\n",
        "    tri_tbl = eval_contingency_table_tri(model, tokenizer, dataset, tri_dataset, blank=blank)\n",
        "    tri_acc = compute_acc(tri_tbl)\n",
        "\n",
        "    print(re2_acc, re3_acc, tri_acc)\n",
        "\n",
        "    return {\n",
        "        're2': re2_tbl,\n",
        "        're3': re3_tbl,\n",
        "        'tri': tri_tbl\n",
        "    }\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}