Draft_DLH_Team_10.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "5e610c339f274dce83a7f9a89d4e03cc": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_372600f5bf824273bcee5716744fe600",
              "IPY_MODEL_4c3201bf0cec4bf0b13deeb9030d8137",
              "IPY_MODEL_09a58d2e1110437ab543bbc06accd8d6"
            ],
            "layout": "IPY_MODEL_c44562eb466642759cc5624445896147"
          }
        },
        "372600f5bf824273bcee5716744fe600": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_d1b5e8423d3b412c81338261acbef85e",
            "placeholder": "​",
            "style": "IPY_MODEL_03a9f9bfde714116940edf0e28f9188a",
            "value": "vocab.txt: 100%"
          }
        },
        "4c3201bf0cec4bf0b13deeb9030d8137": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_d6778d8b5e9e424ba3316eea1c06aa68",
            "max": 227845,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_08ce0e0ce2dc43ba85674bf8e2c62c81",
            "value": 227845
          }
        },
        "09a58d2e1110437ab543bbc06accd8d6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_a1df3ad39ac14229a7ca2965b3676a7f",
            "placeholder": "​",
            "style": "IPY_MODEL_01758ca3c70b4cd49be8dba16beaf6ee",
            "value": " 228k/228k [00:00&lt;00:00, 925kB/s]"
          }
        },
        "c44562eb466642759cc5624445896147": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d1b5e8423d3b412c81338261acbef85e": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "03a9f9bfde714116940edf0e28f9188a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "d6778d8b5e9e424ba3316eea1c06aa68": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "08ce0e0ce2dc43ba85674bf8e2c62c81": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "a1df3ad39ac14229a7ca2965b3676a7f": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "01758ca3c70b4cd49be8dba16beaf6ee": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "ff2374b58ecb44d683a45ad57d0df434": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_e9831db472e1488d9dbbb438cdeec4c1",
              "IPY_MODEL_048f4f47b5aa4b559325b4beb58a1219",
              "IPY_MODEL_a003f1745e0243da835d2ddda85b2e62"
            ],
            "layout": "IPY_MODEL_792f44656def48c4b422c379752b800b"
          }
        },
        "e9831db472e1488d9dbbb438cdeec4c1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_56f8995781714bed97f53d3efabab1f9",
            "placeholder": "​",
            "style": "IPY_MODEL_bd3a033b1b374b9abcb6e61d7491b4ba",
            "value": "config.json: 100%"
          }
        },
        "048f4f47b5aa4b559325b4beb58a1219": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_0030f5a8e3be40a68804eb62089b1925",
            "max": 385,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_e0836f4de368466bbaaf9c6377a34764",
            "value": 385
          }
        },
        "a003f1745e0243da835d2ddda85b2e62": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_87b68efb0993441eb80bd9479f2f2d14",
            "placeholder": "​",
            "style": "IPY_MODEL_fc421d6c9f6a4c318469def31378c19f",
            "value": " 385/385 [00:00&lt;00:00, 18.5kB/s]"
          }
        },
        "792f44656def48c4b422c379752b800b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "56f8995781714bed97f53d3efabab1f9": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "bd3a033b1b374b9abcb6e61d7491b4ba": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "0030f5a8e3be40a68804eb62089b1925": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "e0836f4de368466bbaaf9c6377a34764": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "87b68efb0993441eb80bd9479f2f2d14": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "fc421d6c9f6a4c318469def31378c19f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "17baf1e935794757863b43beeb6b4db9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_0b01accdfa8144f68750e7bd335f446f",
              "IPY_MODEL_b4f6c645db0e49f89623f63b08ba3f0d",
              "IPY_MODEL_2ff98dfb2d044fe7ab86d5bf685ffe6e"
            ],
            "layout": "IPY_MODEL_c6972db01d1a439abb81125146dabaf3"
          }
        },
        "0b01accdfa8144f68750e7bd335f446f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_c025fe1bd3fc477f870811a6e0a7f879",
            "placeholder": "​",
            "style": "IPY_MODEL_97fd474528a049ca8a4d6623a7c9dcf4",
            "value": "pytorch_model.bin: 100%"
          }
        },
        "b4f6c645db0e49f89623f63b08ba3f0d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ce4123e9efc0466a818ee0b1f62068d9",
            "max": 442221694,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_6624113a90934adeb5c7fd3dd527feca",
            "value": 442221694
          }
        },
        "2ff98dfb2d044fe7ab86d5bf685ffe6e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ede6ecd6c6c94466a036e9edc90289e2",
            "placeholder": "​",
            "style": "IPY_MODEL_11a62dac75e2423ba9fe85781a6a43bf",
            "value": " 442M/442M [00:01&lt;00:00, 302MB/s]"
          }
        },
        "c6972db01d1a439abb81125146dabaf3": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "c025fe1bd3fc477f870811a6e0a7f879": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "97fd474528a049ca8a4d6623a7c9dcf4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "ce4123e9efc0466a818ee0b1f62068d9": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "6624113a90934adeb5c7fd3dd527feca": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "ede6ecd6c6c94466a036e9edc90289e2": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "11a62dac75e2423ba9fe85781a6a43bf": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# CS598 Deep Learning for Healthcare - Reproducibility Project Draft"
      ],
      "metadata": {
        "id": "aSoxztK_rPU8"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Darin Zhen, George Vojvodic, and Alan Yee; {darinz2, dvojvo2, alanyee2} @illinois.edu\n",
        "\n",
        "Group ID: 10\n",
        "\n",
        "Paper ID: 94\n",
        "\n",
        "Text2Mol: Cross-Modal Molecule Retrieval with Natural Language Queries\n"
      ],
      "metadata": {
        "id": "jlPcxRQc_Sfy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Relevant Terminology"
      ],
      "metadata": {
        "id": "wZ3cbeZ-XOtQ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "- **Hits@1**: A metric commonly used in information retrieval and recommendation systems to evaluate their performance. It measures the accuracy of a system in predicting the top recommendation or the correct answer out of a list of options for a given query or user interaction.\n",
        "\n",
        "- **Mean Reciprocal Rank (MRR)**: A metric used to evaluate the effectiveness of information retrieval systems, particularly in the context of ranked retrieval. It measures the quality of the ranked list of results by considering the position of the first relevant item in the list.\n",
        "\n",
        "- **Molecule**: An electrically neutral group of atoms bonded together.\n",
        "\n",
        "- **Compound**: Two or more elements held together by chemical bonds.\n",
        "\n",
        "- **Chemical fingerprint**: Represents a molecule or substructure using a bitstring. This allows for efficient substructure search and similarity calculation.\n",
        "\n",
        "- **Morgan fingerprint**: A specific type of chemical fingerprint also known as ECFP.\n",
        "\n",
        "- **SMILES string**: A character-based sequence representation of a molecule. (for example, C1=CC=CC=C1 is the SMILES string for benzene)\n",
        "\n",
        "- **Canonical SMILES**: A unique SMILES string for a molecule.\n"
      ],
      "metadata": {
        "id": "3782OlYB2QCa"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Introduction"
      ],
      "metadata": {
        "id": "4EhXyWiXXT7y"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The discovery of new molecules and understanding their properties is critical for advancing fields like medicine, chemistry, and materials science. However, the vast number of possible molecules makes it impractical to experimentally characterize each one. There are already tens of millions of molecules cataloged in databases like PubChem. Efficiently retrieving relevant molecules from these large databases given natural language descriptions is an important yet challenging problem.\n",
        "\n",
        "Current methods for molecule retrieval typically rely on structured representations like molecular fingerprints or SMILES strings. While these enable substructure matching and similarity searches, they do not directly integrate the semantic information contained in natural language descriptions. Some approaches replace chemical names in text with canonical identifiers, but this fails to capture the full meaning. Solving the problem of cross-modal retrieval between natural language and molecules would allow scientists to easily search for molecules based on high-level conceptual descriptions rather than just structural patterns.\n",
        "\n",
        "The key challenge lies in bridging the stark difference between the modalities of natural language and molecular structure data. Molecules are usually represented as graphs with atoms as nodes and bonds as edges, following a unique grammar quite distinct from human language. This makes cross-modal retrieval exceptionally challenging compared to traditional cross-lingual information retrieval between natural languages.\n",
        "\n",
        "In this paper, the authors propose a novel \"Text2Mol\" task for retrieving molecules directly from natural language descriptions (shown in Fiture 1). They develop a multimodal embedding approach to learn an aligned semantic space bridging text and molecular structure data. This allows ranking molecules by similarity to text query descriptions. The paper makes several innovations, including extending the loss function with negative sampling to encourage integration of both modalities, using cross-modal attention to extract interpretable \"association rules\" between text and molecular substructures, and an ensemble method that significantly boosts performance.\n",
        "<br>\n",
        "<img src=\"https://drive.google.com/uc?id=1mG9lgWvfpZ2tplVl6xqLZMyzj5ooPhy8\" width=500 />\n",
        "<br>\n",
        "<b>Figure 1: Given a natural language description of water, we want to rank the corresponding molecule $H_2O$ first among all the possible molecules.</b>\n",
        "<br>\n",
        "\n",
        "On a new benchmark dataset of over 33,000 text-molecule pairs, the proposed methods achieve a mean reciprocal rank of 0.499, substantially outperforming baselines. The cross-modal attention model provides insightful explanations grounding the language representations to the molecular structure. Overall, this multimodal approach offers a powerful solution for understanding chemistry literature and searching molecular databases, with broad potential applications in drug discovery, materials design, and scientific knowledge exploration.\n"
      ],
      "metadata": {
        "id": "MQ0sNuMePBXx"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Scope of Reproducibility"
      ],
      "metadata": {
        "id": "fGRats_iwAeF"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The scope of reproducibility in the paper encompasses the following key hypotheses that will reproduce:\n",
        "\n",
        "1. **Hypothesis 1**: Cross-modal embedding can effectively align text and molecule spaces for retrieval. This involves reproducing embedding models and evaluating retrieval metrics like Mean Reciprocal Rank (MRR).\n",
        "\n",
        "2. **Hypothesis 2**: Ensemble of different architectures (MLP vs GCN) improves results compared to individual models. This involves training different models and comparing ensemble versus individual performance.\n",
        "\n",
        "3. **Hypothesis 3**: Cross-modal attention provides insights into text-molecule associations. This will be examined by analyzing attention weights and extracted rules for coherence.\n",
        "\n",
        "4. **Hypothesis 4**: Different architectures possess complementary strengths, where MLP may rank easier examples better but GCN generalizes better. This will be probed by analyzing differences in rankings between architectures.\n",
        "\n",
        "5. **Hypothesis 5**: Cross-modal reranking using attention rules improves over the base model. Testing reranking on a holdout set will validate this hypothesis."
      ],
      "metadata": {
        "id": "SvYxdQvhGSB6"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Setup"
      ],
      "metadata": {
        "id": "ipcAE7rZNJG8"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**NOTE:** This notebook requires **GPU** runtime due to the extensive use of hundreds of millions of parameters in the training model."
      ],
      "metadata": {
        "id": "1h5Xpws8ET5O"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "This notebook supports accessing data and images:\n"
      ],
      "metadata": {
        "id": "_F7TTYKAxlSu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Data directory\n",
        "!mkdir /content/data\n",
        "\n",
        "# Image directory\n",
        "!mkdir /content/image\n",
        "\n",
        "# Input directory\n",
        "!mkdir /content/input\n",
        "\n",
        "# Download files into the data folder.\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1XYjz33tWWet6t4QouZkVn3TRYliG5L5F' -O ChEBI_defintions_substructure_corpus.cp\n",
        "!mv /content/ChEBI_defintions_substructure_corpus.cp /content/data/ChEBI_defintions_substructure_corpus.cp\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1U_0mn1-GZ7NtL8Bk2S8Yr9IQQp51qbpf' -O chem_embeddings_test.npy\n",
        "!mv /content/chem_embeddings_test.npy /content/data/chem_embeddings_test.npy\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1nVQIExr7toG3Ob6CJzPK51LKq9OsL1Ar' -O chem_embeddings_train.npy\n",
        "!mv /content/chem_embeddings_train.npy /content/data/chem_embeddings_train.npy\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1ybhCMaMFSETom3PlOGTkXXRpuauCMztH' -O chem_embeddings_val.npy\n",
        "!mv /content/chem_embeddings_val.npy /content/data/chem_embeddings_val.npy\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1vPKXKtlQx8oX3-SAlgzTGvXHU0SwN1Rr' -O cids_test.npy\n",
        "!mv /content/cids_test.npy /content/data/cids_test.npy\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1CUvn6lOUbbb7sJkqnS0jQ5eYYnJfDM7F' -O cids_train.npy\n",
        "!mv /content/cids_train.npy /content/data/cids_train.npy\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1LMfx63hrd6r5pAQWUsK5BYbb5TEcumuc' -O cids_val.npy\n",
        "!mv /content/cids_val.npy /content/data/cids_val.npy\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1SmKQJPKRePUXyomOBdMwha75-D5O4413' -O test.sdf\n",
        "!mv /content/test.sdf /content/data/test.sdf\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1dhhTAD3z97yOQYSK0Go-RI-bgjJdSg6e' -O test.txt\n",
        "!mv /content/test.txt /content/data/test.txt\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=12DlIeAwx_oeJBsuRAoi_t4cr7Rj3G8UN' -O text_embeddings_test.npy\n",
        "!mv /content/text_embeddings_test.npy /content/data/text_embeddings_test.npy\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1tozq-TbD2avIqwtZkA7TtwXgskUn_dMX' -O text_embeddings_train.npy\n",
        "!mv /content/text_embeddings_train.npy /content/data/text_embeddings_train.npy\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1C3-4KHwKySBCXri6YYj6IFJ3BppYkpYI' -O text_embeddings_val.npy\n",
        "!mv /content/text_embeddings_val.npy /content/data/text_embeddings_val.npy\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1gR4B11xGBLGwYQ2-s_k0C19HUXYDFSoa' -O token_embedding_dict.npy\n",
        "!mv /content/token_embedding_dict.npy /content/data/token_embedding_dict.npy\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1g8a6wg1OC8okFltFJOaNbMC5BsWuQdML' -O training.sdf\n",
        "!mv /content/training.sdf /content/data/training.sdf\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=17SvDWffLm8Eez7KIIyZt3RvsNOkSDZMM' -O training.txt\n",
        "!mv /content/training.txt /content/data/training.txt\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1i39CtXI7HbdtnRG4lHMCFKQ7i_AZgSdn' -O val.sdf\n",
        "!mv /content/val.sdf /content/data/val.sdf\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1sQ7iYAHIRMmq0YePRaiRngcPNwFPzZoO' -O val.txt\n",
        "!mv /content/val.txt /content/data/val.txt\n",
        "\n",
        "\n",
        "# Download files into the input folder.\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1sOLCJIDKZCZO-9jSFPsQGYK4h17sDvU7' -O ChEBI_defintions_substructure_corpus.cp\n",
        "!mv /content/ChEBI_defintions_substructure_corpus.cp /content/input/ChEBI_defintions_substructure_corpus.cp\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1G1iST_JfJTfj1hRBzzqgIBQZDUjzv5oA' -O mol2vec_ChEBI_20_test.txt\n",
        "!mv /content/mol2vec_ChEBI_20_test.txt /content/input/mol2vec_ChEBI_20_test.txt\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1aSuZUiiM7Bmsg-zONm7RMfAm8MizL1tc' -O mol2vec_ChEBI_20_training.txt\n",
        "!mv /content/mol2vec_ChEBI_20_training.txt /content/input/mol2vec_ChEBI_20_training.txt\n",
        "\n",
        "!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1BERxF6s-GPDRaEOoSBfhfs-u6RS1UUSx' -O mol2vec_ChEBI_20_val.txt\n",
        "!mv /content/mol2vec_ChEBI_20_val.txt /content/input/mol2vec_ChEBI_20_val.txt\n"
      ],
      "metadata": {
        "id": "7QvHfkRpwUtR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Methodology"
      ],
      "metadata": {
        "id": "smxGlzsjIS7g"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "The paper proposes a new task called Text2Mol, which aims to retrieve molecules from natural language descriptions. The methodology involves the following key steps:\n",
        "\n",
        "1. Construct a dataset of molecule-text description pairs from sources like PubChem and ChEBI.\n",
        "\n",
        "2. Learn aligned semantic embeddings for text and molecules using:\n",
        "    a) A text encoder based on SciBERT\n",
        "    b) A molecule encoder using either a multi-layer perceptron (MLP) on Mol2vec embeddings or a graph convolutional network (GCN) on the molecular graph with Mol2vec features.\n",
        "\n",
        "3. Train the encoders using a contrastive loss that aims to bring positive (matching) molecule-text pairs closer and push negative pairs apart in the embedding space.\n",
        "\n",
        "4. At inference time, encode the text query and retrieve the nearest molecule embeddings using cosine similarity.\n",
        "\n",
        "5. Explore ensembling multiple trained models and incorporating cross-modal attention to learn association rules between text tokens and molecular substructures for explainability and reranking.\n",
        "\n",
        "The key novelties are applying contrastive learning across the text and molecule modalities, proposing the Text2Mol retrieval task, and using cross-modal attention for explainable retrieval via association rules."
      ],
      "metadata": {
        "id": "xWAHJ_1CdtaA"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Environment"
      ],
      "metadata": {
        "id": "Mg-kVgJuwFo8"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Python Version\n",
        "\n"
      ],
      "metadata": {
        "id": "i5eAfWrjF0EZ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "This project uses Python version 3.10."
      ],
      "metadata": {
        "id": "MRZezxDLiqrJ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Dependencies & Packages Needed"
      ],
      "metadata": {
        "id": "H2stTGmhagh0"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "This project relies on several Python libraries and modules for text-to-molecule tasks:\n",
        "\n",
        "1. **Operating System Interaction**: `os` module for interacting with the operating system.\n",
        "2. **File Operations**: `shutil` module for file operations.\n",
        "3. **Time-related Functions**: `time` module for time-related functions.\n",
        "4. **Mathematical Operations**: `math` module for mathematical operations.\n",
        "5. **Numerical Computations**: `numpy` (`np` alias) for numerical computations.\n",
        "6. **Plotting**: `matplotlib.pyplot` (`plt` alias) for plotting.\n",
        "7. **Cosine Similarity Computation**: `cosine_similarity` from `sklearn.metrics.pairwise` for computing cosine similarity.\n",
        "8. **Deep Learning Framework**: `torch` for PyTorch, a deep learning framework.\n",
        "9. **Neural Network Components**: `torch.nn` for neural network modules and `torch.nn.functional` (`F` alias) for functional interfaces.\n",
        "10. **Data Handling Utilities**: `torch.utils.data` for handling data in PyTorch, including `Dataset` and `DataLoader`.\n",
        "11. **Tokenization**: `tokenizers` for tokenization, including the `Tokenizer` class.\n",
        "12. **BERT Model and Tokenizer**: `BertTokenizerFast` and `BertModel` from Hugging Face's Transformers library for BERT tokenizer and model.\n",
        "13. **CSV File Handling**: `csv` module for reading and writing CSV files.\n",
        "14. **Graph Convolutional Network (GCN)**: `torch_geometric.nn` for GCN operations, including `GCNConv` and `global_mean_pool`.\n",
        "15. **Transformer Decoder**: `TransformerDecoder` and `TransformerDecoderLayer` from PyTorch for transformer decoder operations.\n",
        "16. **Optimization**: `torch.optim` for optimization, including various optimizers.\n",
        "17. **Learning Rate Scheduler**: `get_linear_schedule_with_warmup` from the transformers library for learning rate scheduling.\n",
        "\n",
        "Additionally, the code includes an installation command (`!pip install torch_geometric`) to install the `torch_geometric` library."
      ],
      "metadata": {
        "id": "PHmhDaZTjBht"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# This code imports various Python libraries and modules that are used in this\n",
        "# notebook for Text2Mol\n",
        "# Importing necessary libraries/modules\n",
        "import os                   # Module for interacting with the operating system\n",
        "import shutil               # Module for file operations\n",
        "import time                 # Module for time-related functions\n",
        "\n",
        "import math                # Module for mathematical operations\n",
        "\n",
        "import numpy as np         # NumPy, a library for numerical computations\n",
        "\n",
        "import matplotlib.pyplot as plt  # Matplotlib, a plotting library\n",
        "from sklearn.metrics.pairwise import cosine_similarity  # Module for cosine similarity computation\n",
        "\n",
        "import torch               # PyTorch, a deep learning framework\n",
        "from torch import nn       # Neural network module from PyTorch\n",
        "import torch.nn.functional as F  # Functional interface to neural network components in PyTorch\n",
        "from torch.utils.data import Dataset, DataLoader  # Utilities for handling data in PyTorch\n",
        "\n",
        "import tokenizers         # Tokenizers library for tokenization\n",
        "from tokenizers import Tokenizer  # Tokenizer class for tokenization\n",
        "from transformers import BertTokenizerFast, BertModel  # BERT tokenizer and model from Hugging Face's\n",
        "                                                       # Transformers library\n",
        "\n",
        "import csv                 # Module for reading and writing CSV files"
      ],
      "metadata": {
        "id": "yu61Jp1xrnKk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Data"
      ],
      "metadata": {
        "id": "Gs968ZkyG3wZ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Data Download Instruction"
      ],
      "metadata": {
        "id": "hjY-QGwzHBBx"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "To download the data for the project, follow these instructions:\n",
        "\n",
        "1. **ChEBI Annotations of Compounds from [PubChem](https://pubchem.ncbi.nlm.nih.gov/)**: (This part is not necessary as the raw ChEBI dataset is not directly used in the model, and the actual dataset is already directly provided and organized in `2.`. Nonetheless, we provide steps here to obtain the raw data.)\n",
        "   1. Visit the [ChEBI](https://www.ebi.ac.uk/chebi/) website.\n",
        "   1. Under `Downloads`, click on `SDF files`.\n",
        "   1. Click on `ChEBI_complete.sdf.gz`\n",
        "   1. Choose where to download to zipped file and click `Save`\n",
        "\n",
        "2. **ChEBI-20 Dataset**: (This part is required.)\n",
        "   - Access the [ChEBI-20 dataset repository](https://github.com/cnedwards/text2mol/tree/master/data)."
      ],
      "metadata": {
        "id": "MU216Z1Qb04g"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Data Descriptions"
      ],
      "metadata": {
        "id": "2cUAEWuqlQOT"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The paper makes use of the following datasets:\n",
        "\n",
        "1. [ChEBI](https://www.ebi.ac.uk/chebi/) (Chemical Entities of Biological Interest) annotations of compounds scraped from [PubChem](https://pubchem.ncbi.nlm.nih.gov/).\n",
        "   - This contains 102,980 compound-description pairs.\n",
        "\n",
        "2. [ChEBI-20](https://github.com/cnedwards/text2mol/tree/master/data) dataset\n",
        "   - Constructed from the ChEBI/PubChem data by filtering for descriptions longer than 20 words.\n",
        "   - Contains 33,010 text-compound pairs.\n",
        "   - Split into 80/10/10% train/validation/test sets.\n",
        "\n",
        "The ChEBI-20 dataset forms the main benchmark used to evaluate the proposed cross-modal molecule retrieval methods.\n",
        "\n",
        "For representing the molecular structures, the paper uses:\n",
        "\n",
        "1. Mol2vec representations\n",
        "   - Molecular graphs are converted to \"sentences\" using the Morgan fingerprinting (shown in Figure 2) algorithm which generates substructure identifiers.\n",
        "   <br>\n",
        "   <img src=\"https://drive.google.com/uc?id=1lT5fla5UnQXW60-utzHYWDFdg62sMgp0\" width=\"400\">\n",
        "   <br>\n",
        "   <b>Figure 2: Example of Morgan Fingerprinting from (Rogers and Hahn, 2010) for Butyramide. The algorithm updates the identifiers from radius r = 0 to r = 1, as shown by the green circles.</b>\n",
        "   <br>\n",
        "   - The Mol2vec algorithm applies Word2vec on these substructure sentences to produce molecule embeddings.\n",
        "   - Default radius of 1 is used, giving two substructure tokens per atom.\n",
        "\n",
        "2. SMILES strings\n",
        "   - Character-based representation of molecules that can be parsed into molecular graphs.\n",
        "\n",
        "The text descriptions are encoded using the pre-trained SciBERT language model.\n",
        "\n",
        "The key data is the new ChEBI-20 dataset of paired text descriptions and molecular structures, with molecules represented by Mol2vec embeddings or SMILES strings, and descriptions encoded by SciBERT."
      ],
      "metadata": {
        "id": "2NbPHUTMbkD3"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### The ChEBI-20 dataset is contained in 6 files:"
      ],
      "metadata": {
        "id": "Cmgrq6ROlgEj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "(1,2,3) The mol2vec_ChEBI_20_X.txt files have lines in the following form:\n",
        "```\n",
        "CID\tmol2vec embedding\tDescription\n",
        "```\n",
        "\n",
        "(4) mol_graphs.zip contain {cid}.graph files. These are formatted first with the edgelist of the graph and then substructure tokens for each node.\n",
        "For example,\n",
        "edgelist:\n",
        "```\n",
        "0 1\n",
        "1 0\n",
        "1 2\n",
        "2 1\n",
        "1 3\n",
        "3 1\n",
        "```\n",
        "idx to identifier:\n",
        "```\n",
        "0 3537119515\n",
        "1 2059730245\n",
        "2 3537119515\n",
        "3 1248171218\n",
        "```\n",
        "\n",
        "(5) ChEBI_defintions_substructure_corpus.cp contains the molecule token \"sentences\". It is formatted:\n",
        "```\n",
        "cid: tokenid1 tokenid2 tokenid3 ... tokenidn\n",
        "```\n",
        "\n",
        "(6) token_embedding_dict.npy is a dictionary mapping molecule tokens to their embeddings.\n",
        "\n",
        "It can be loaded with the following code:\n",
        "```python\n",
        "import numpy as np\n",
        "token_embedding_dict = np.load(\"token_embedding_dict.npy\", allow_pickle=True)[()]\n",
        "```"
      ],
      "metadata": {
        "id": "9eIF9Jtqldbk"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Preprocessing Code + Command"
      ],
      "metadata": {
        "id": "0cwbhQ8iHIaW"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### GenerateData Class\n",
        "\n",
        "This class is designed to handle the generation of examples for training,\n",
        "validation, and testing sets, where each example contains both text data\n",
        "(processed using a BERT tokenizer) and molecule data. The methods within\n",
        "the class prepare the necessary data structures, tokenize text inputs,\n",
        "and yield examples in the desired format."
      ],
      "metadata": {
        "id": "kkArCBBZudkT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#Need a special generator for random sampling:\n",
        "\n",
        "class GenerateData():\n",
        "    def __init__(self, path_train, path_val, path_test, path_molecules, path_token_embs):\n",
        "        # Constructor method initializing paths and parameters\n",
        "        self.path_train = path_train  # Path to the training data file\n",
        "        self.path_val = path_val  # Path to the validation data file\n",
        "        self.path_test = path_test  # Path to the test data file\n",
        "        self.path_molecules = path_molecules  # Path to the file containing molecule data\n",
        "        self.path_token_embs = path_token_embs  # Path to the file containing token embeddings\n",
        "\n",
        "        self.text_trunc_length = 256  # Maximum length for text input tokens\n",
        "\n",
        "        # Initialize text tokenizer\n",
        "        self.prep_text_tokenizer()\n",
        "\n",
        "        # Load substructures from molecule data\n",
        "        self.load_substructures()\n",
        "\n",
        "        self.batch_size = 32  # Batch size for data processing\n",
        "\n",
        "        # Store descriptions\n",
        "        self.store_descriptions()\n",
        "\n",
        "    def load_substructures(self):\n",
        "        # Method to load substructures from molecule data\n",
        "        self.molecule_sentences = {}  # Dictionary to store molecule sentences\n",
        "        self.molecule_tokens = {}  # Dictionary to store molecule tokens\n",
        "\n",
        "        total_tokens = set()  # Set to store unique tokens\n",
        "        self.max_mol_length = 0  # Variable to store maximum molecule length\n",
        "        with open(self.path_molecules) as f:\n",
        "            for line in f:\n",
        "                spl = line.split(\":\")\n",
        "                cid = spl[0]  # Compound ID\n",
        "                tokens = spl[1].strip()  # Tokens for the compound\n",
        "                self.molecule_sentences[cid] = tokens\n",
        "                t = tokens.split()\n",
        "                total_tokens.update(t)  # Add tokens to the set\n",
        "                size = len(t)\n",
        "                if size > self.max_mol_length:\n",
        "                    self.max_mol_length = size  # Update maximum molecule length\n",
        "\n",
        "        # Load token embeddings\n",
        "        self.token_embs = np.load(self.path_token_embs, allow_pickle=True)[()]\n",
        "\n",
        "    def prep_text_tokenizer(self):\n",
        "        # Method to prepare text tokenizer (using BERT)\n",
        "        self.text_tokenizer = BertTokenizerFast.from_pretrained(\"allenai/scibert_scivocab_uncased\")\n",
        "\n",
        "    def store_descriptions(self):\n",
        "        # Method to store descriptions from training, validation, and test sets\n",
        "        self.descriptions = {}  # Dictionary to store descriptions\n",
        "        self.mols = {}  # Dictionary to store molecule data\n",
        "\n",
        "        self.training_cids = []  # List to store training set compound IDs\n",
        "        # Get training set compound IDs\n",
        "        with open(self.path_train) as f:\n",
        "            reader = csv.DictReader(f, delimiter=\"\\t\", quoting=csv.QUOTE_NONE,\n",
        "                                    fieldnames=['cid', 'mol2vec', 'desc'])\n",
        "            for n, line in enumerate(reader):\n",
        "                self.descriptions[line['cid']] = line['desc']\n",
        "                self.mols[line['cid']] = line['mol2vec']\n",
        "                self.training_cids.append(line['cid'])\n",
        "\n",
        "        self.validation_cids = []  # List to store validation set compound IDs\n",
        "        # Get validation set compound IDs\n",
        "        with open(self.path_val) as f:\n",
        "            reader = csv.DictReader(f, delimiter=\"\\t\", quoting=csv.QUOTE_NONE,\n",
        "                                    fieldnames=['cid', 'mol2vec', 'desc'])\n",
        "            for n, line in enumerate(reader):\n",
        "                self.descriptions[line['cid']] = line['desc']\n",
        "                self.mols[line['cid']] = line['mol2vec']\n",
        "                self.validation_cids.append(line['cid'])\n",
        "\n",
        "        self.test_cids = []  # List to store test set compound IDs\n",
        "        # Get test set compound IDs\n",
        "        with open(self.path_test) as f:\n",
        "            reader = csv.DictReader(f, delimiter=\"\\t\", quoting=csv.QUOTE_NONE,\n",
        "                                    fieldnames=['cid', 'mol2vec', 'desc'])\n",
        "            for n, line in enumerate(reader):\n",
        "                self.descriptions[line['cid']] = line['desc']\n",
        "                self.mols[line['cid']] = line['mol2vec']\n",
        "                self.test_cids.append(line['cid'])\n",
        "\n",
        "    def generate_examples_train(self):\n",
        "        # Method to generate examples for training set\n",
        "        np.random.shuffle(self.training_cids)  # Shuffle training compound IDs\n",
        "\n",
        "        for cid in self.training_cids:\n",
        "            text_input = self.text_tokenizer(self.descriptions[cid], truncation=True,\n",
        "                                             max_length=self.text_trunc_length,\n",
        "                                             padding='max_length', return_tensors='np')  # Tokenize text input\n",
        "\n",
        "            yield {\n",
        "                'cid': cid,\n",
        "                'input': {\n",
        "                    'text': {\n",
        "                        'input_ids': text_input['input_ids'].squeeze(),\n",
        "                        'attention_mask': text_input['attention_mask'].squeeze(),\n",
        "                    },\n",
        "                    'molecule': {\n",
        "                        'mol2vec': np.fromstring(self.mols[cid], sep=\" \"),  # Convert molecule data to NumPy array\n",
        "                        'cid': cid\n",
        "                    },\n",
        "                },\n",
        "            }\n",
        "\n",
        "    def generate_examples_val(self):\n",
        "        # Method to generate examples for validation set\n",
        "        np.random.shuffle(self.validation_cids)  # Shuffle validation compound IDs\n",
        "\n",
        "        for cid in self.validation_cids:\n",
        "            text_input = self.text_tokenizer(self.descriptions[cid], truncation=True,\n",
        "                                             padding='max_length',\n",
        "                                             max_length=self.text_trunc_length, return_tensors='np')  # Tokenize text input\n",
        "\n",
        "            yield {\n",
        "                'cid': cid,\n",
        "                'input': {\n",
        "                    'text': {\n",
        "                        'input_ids': text_input['input_ids'].squeeze(),\n",
        "                        'attention_mask': text_input['attention_mask'].squeeze(),\n",
        "                    },\n",
        "                    'molecule': {\n",
        "                        'mol2vec': np.fromstring(self.mols[cid], sep=\" \"),  # Convert molecule data to NumPy array\n",
        "                        'cid': cid\n",
        "                    }\n",
        "                },\n",
        "            }\n",
        "\n",
        "    def generate_examples_test(self):\n",
        "        # Method to generate examples for test set\n",
        "        np.random.shuffle(self.test_cids)  # Shuffle test compound IDs\n",
        "\n",
        "        for cid in self.test_cids:\n",
        "            text_input = self.text_tokenizer(self.descriptions[cid], truncation=True, padding='max_length',\n",
        "                                             max_length=self.text_trunc_length, return_tensors='np')  # Tokenize text input\n",
        "\n",
        "            yield {\n",
        "                'cid': cid,\n",
        "                'input': {\n",
        "                    'text': {\n",
        "                        'input_ids': text_input['input_ids'].squeeze(),\n",
        "                        'attention_mask': text_input['attention_mask'].squeeze(),\n",
        "                    },\n",
        "                    'molecule': {\n",
        "                        'mol2vec': np.fromstring(self.mols[cid], sep=\" \"),  # Convert molecule data to NumPy array\n",
        "                        'cid': cid\n",
        "                    }\n",
        "                },\n",
        "            }"
      ],
      "metadata": {
        "id": "d7kMwLjibzSo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "In the following code, the paths to various data files are defined. Then, it checks\n",
        "if a specific token embedding file exists using os.path.exists(). If the\n",
        "file does not exist, it raises a FileNotFoundError. Finally, an instance of\n",
        "the GenerateData class is created with the defined"
      ],
      "metadata": {
        "id": "JS7CEl0IuugZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the path to the token embedding file\n",
        "#mounted_path_token_embs = os.path.join(data_dir, 'token_embedding_dict.npy')\n",
        "mounted_path_token_embs = \"data/token_embedding_dict.npy\"\n",
        "\n",
        "# Check if the token embedding file exists\n",
        "if not os.path.exists(mounted_path_token_embs):\n",
        "    # Raise FileNotFoundError if the file does not exist\n",
        "    raise FileNotFoundError(f\"The following token embedding DOES NOT EXIST: {mounted_path_token_embs}\")\n",
        "\n",
        "# Define the path to the molecule data file\n",
        "parent_dir = os.path.join('/content/drive/', 'My Drive')\n",
        "\n",
        "mounted_path_molecules = \"input/ChEBI_defintions_substructure_corpus.cp\"\n",
        "\n",
        "mounted_path_train = \"input/mol2vec_ChEBI_20_training.txt\"\n",
        "mounted_path_val = \"input/mol2vec_ChEBI_20_val.txt\"\n",
        "mounted_path_test = \"input/mol2vec_ChEBI_20_test.txt\"\n",
        "\n",
        "# Instantiate the GenerateData class with the specified paths\n",
        "gt = GenerateData(mounted_path_train, mounted_path_val, mounted_path_test,\n",
        "                  mounted_path_molecules, mounted_path_token_embs)"
      ],
      "metadata": {
        "id": "cQY3wE2jcfUk",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 208,
          "referenced_widgets": [
            "5e610c339f274dce83a7f9a89d4e03cc",
            "372600f5bf824273bcee5716744fe600",
            "4c3201bf0cec4bf0b13deeb9030d8137",
            "09a58d2e1110437ab543bbc06accd8d6",
            "c44562eb466642759cc5624445896147",
            "d1b5e8423d3b412c81338261acbef85e",
            "03a9f9bfde714116940edf0e28f9188a",
            "d6778d8b5e9e424ba3316eea1c06aa68",
            "08ce0e0ce2dc43ba85674bf8e2c62c81",
            "a1df3ad39ac14229a7ca2965b3676a7f",
            "01758ca3c70b4cd49be8dba16beaf6ee",
            "ff2374b58ecb44d683a45ad57d0df434",
            "e9831db472e1488d9dbbb438cdeec4c1",
            "048f4f47b5aa4b559325b4beb58a1219",
            "a003f1745e0243da835d2ddda85b2e62",
            "792f44656def48c4b422c379752b800b",
            "56f8995781714bed97f53d3efabab1f9",
            "bd3a033b1b374b9abcb6e61d7491b4ba",
            "0030f5a8e3be40a68804eb62089b1925",
            "e0836f4de368466bbaaf9c6377a34764",
            "87b68efb0993441eb80bd9479f2f2d14",
            "fc421d6c9f6a4c318469def31378c19f"
          ]
        },
        "outputId": "64c6f5b0-3537-4058-d880-29f6034ffaf8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/huggingface_hub/utils/_token.py:88: UserWarning: \n",
            "The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
            "To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n",
            "You will be able to reuse this secret in all of your notebooks.\n",
            "Please note that authentication is recommended but still optional to access public models or datasets.\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "vocab.txt:   0%|          | 0.00/228k [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "5e610c339f274dce83a7f9a89d4e03cc"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "config.json:   0%|          | 0.00/385 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "ff2374b58ecb44d683a45ad57d0df434"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Dataset Class\n",
        "\n",
        "This class Dataset is designed to create a custom dataset for PyTorch. It allows generating samples of data on-the-fly using a generator function gen. The `__len__` method returns the total number of samples in the dataset, and the `__getitem__` method generates one sample of data for a given index. If the generator is exhausted (i.e., it reaches the end), it resets the iterator to the beginning. In this specific implementation, the target variable `y` is set to a constant value of 1 for all samples."
      ],
      "metadata": {
        "id": "7xjEU2zdvJvs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Dataset(Dataset):\n",
        "    'Characterizes a dataset for PyTorch'\n",
        "\n",
        "    def __init__(self, gen, length):\n",
        "        'Initialization'\n",
        "\n",
        "        self.gen = gen  # Generator function that yields data examples\n",
        "        self.it = iter(self.gen())  # Iterator over the generator function\n",
        "\n",
        "        self.length = length  # Length of the dataset\n",
        "\n",
        "    def __len__(self):\n",
        "        'Denotes the total number of samples'\n",
        "\n",
        "        return self.length  # Returns the length of the dataset\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        'Generates one sample of data'\n",
        "\n",
        "        try:\n",
        "            ex = next(self.it)  # Get the next example from the iterator\n",
        "        except StopIteration:\n",
        "            self.it = iter(self.gen())  # If iterator is exhausted, reset it\n",
        "            ex = next(self.it)  # Get the next example\n",
        "\n",
        "        X = ex['input']  # Extract input data from the example\n",
        "        y = 1  # Placeholder for the target variable (constant value in this case)\n",
        "\n",
        "        return X, y  # Return input data and target variable for the given index"
      ],
      "metadata": {
        "id": "2XDCJpNKvbwq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "In the following code, three datasets (training_set, validation_set, and test_set) are created. Each dataset is instantiated with the Dataset class, and they are initialized with different generator functions (gt.generate_examples_train, gt.generate_examples_val, and gt.generate_examples_test, respectively) along with the lengths of their respective compound ID lists. These datasets are intended to be used for training, validation, and testing.\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "HlRcvvowv-I-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Create a dataset for the training set\n",
        "# using the 'generate_examples_train' method of the 'gt' object and the length of training compound IDs\n",
        "training_set = Dataset(gt.generate_examples_train, len(gt.training_cids))\n",
        "\n",
        "# Create a dataset for the validation set\n",
        "# using the 'generate_examples_val' method of the 'gt' object and the length of validation compound IDs\n",
        "validation_set = Dataset(gt.generate_examples_val, len(gt.validation_cids))\n",
        "\n",
        "# Create a dataset for the test set\n",
        "# using the 'generate_examples_test' method of the 'gt' object and the length of test compound IDs\n",
        "test_set = Dataset(gt.generate_examples_test, len(gt.test_cids))\n",
        "\n",
        "n_samples = 50\n",
        "training_set_sample = torch.utils.data.Subset(training_set, list(range(n_samples)))\n",
        "validation_set_sample = torch.utils.data.Subset(validation_set, list(range(n_samples)))\n",
        "test_set_sample = torch.utils.data.Subset(test_set, list(range(n_samples)))\n",
        "\n",
        "params = {'batch_size': gt.batch_size,\n",
        "          'shuffle': True}\n",
        "\n",
        "training_generator = DataLoader(training_set_sample, **params)\n",
        "validation_generator = DataLoader(validation_set_sample, **params)\n",
        "test_generator = DataLoader(test_set_sample, **params)"
      ],
      "metadata": {
        "id": "R3UMPNj4vdaJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Model"
      ],
      "metadata": {
        "id": "lLbPYYFPHQHN"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Citation to the original paper"
      ],
      "metadata": {
        "id": "P2_VkiwtxUfC"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The following is the citation to the original paper:\n",
        "\n",
        "Carl Edwards, ChengXiang Zhai, and Heng Ji. 2021. [Text2mol: Cross-modal molecule retrieval with natural language queries](https://aclanthology.org/2021.emnlp-main.47). In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 595–607, Online and Punta Cana, Dominican Republic. Association for Computational Linguistics."
      ],
      "metadata": {
        "id": "vjOmpT6TS7bq"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Link to the original paper’s repo"
      ],
      "metadata": {
        "id": "rbho17oO4fO-"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "The link to the original paper's repo is as follows:\n",
        "\n",
        "[Text2Mol Code Repository](https://github.com/cnedwards/text2mol)"
      ],
      "metadata": {
        "id": "1U1tBsuES_-7"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Overview\n",
        "\n",
        "The paper proposes several models for the cross-modal Text2Mol task of retrieving molecules from natural language descriptions:\n",
        "\n",
        "Models Parameters:\n",
        "   - The text encoder uses the large SciBERT model, which has around 110M parameters.\n",
        "   - The MLP molecule encoder is relatively small, with around 110M parameters.\n",
        "   - The GCN molecule encoder is slightly larger, with around 112M parameters.\n",
        "   - The cross-modal attention model is the largest, with around 129M parameters.\n",
        "\n",
        "\n",
        "1. Base Models:\n",
        "\n",
        "   a. Text Encoder:\n",
        "   \n",
        "   Uses the SciBERT language model to encode the text description, followed by a linear projection to an embedding space and layer normalization.\n",
        "   \n",
        "   b. Molecule Encoder:\n",
        "    - MLP Encoder: Takes the Mol2vec embedding as input, passes it through a multi-layer perceptron (MLP), and projects to the joint embedding space.\n",
        "      \n",
        "    - GCN Encoder: Incorporates the molecular graph structure by using a Graph Convolutional Network (GCN) on the Mol2vec token embeddings as node features.\n",
        "\n",
        "   \n",
        "   The text and molecule embeddings are mapped to an aligned semantic space, where cosine similarity is used to retrieve/rank molecules given a text query.\n",
        "\n",
        "2. Cross-Modal Attention Model (shown in Figure 3):\n",
        "   - Uses a transformer decoder with cross-modal attention between the text (from SciBERT) and molecule representations (from the GCN encoder).\n",
        "   - Allows learning \"association rules\" between text tokens and molecular substructures from the attention weights.\n",
        "   - Association rules are used for explainability and to rerank retrieved molecules.\n",
        "   <img src=\"https://drive.google.com/uc?id=1_b4MWeiDDRpDKtY44MJMQEaUSZtWjChI\" width=500 />\n",
        "   <br>\n",
        "   <b>Figure 3: Model architecture for the cross-modal attention extension and association rules. </b>\n",
        "   <br>\n",
        "\n",
        "3. Ensemble Model:\n",
        "   - Takes a weighted average of the rankings from different base model instances (e.g. MLP1, MLP2, GCN1, etc.) to create an ensemble ranking.\n",
        "   $$\n",
        "S(m) = \\sum_i w_i R_i(m) \\ \\ \\ \\ \\ \\ \\ \\ s.t. \\sum_i w_i = 1\n",
        "$$\n",
        "<b>Score as a weighted average for some molecule m where $R_i$ is the rank assigned to that molecule by model $i$ and $w_i$ is the model weight.</b><br>\n",
        "   - Improves performance significantly by combining models trained with different initializations.\n",
        "\n",
        "4. Loss Functions:\n",
        "\n",
        "   a. Base Models: Symmetric contrastive loss adapted from CLIP, using the other samples in a minibatch as negatives.\n",
        "   \n",
        "   b. Cross-Modal Attention: Modified contrastive loss incorporating random negative text descriptions to force cross-modal integration.\n",
        "\n",
        "The models are trained on the ChEBI-20 dataset, with molecules represented as SMILES strings or Mol2vec embeddings, and evaluated on molecular retrieval metrics like mean reciprocal rank."
      ],
      "metadata": {
        "id": "3muyDPFPbozY"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "This class Model defines a neural network model for processing text and molecule data. It consists of layers for text and molecule processing, including linear layers, activation functions, layer normalization, and dropout. The text data is processed using a BERT-based transformer model, while the molecule data is processed through fully connected layers. The model outputs scaled representations of text and molecule data."
      ],
      "metadata": {
        "id": "0xbdezr18o1o"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Model(nn.Module):\n",
        "    def __init__(self, ntoken, ninp, nout, nhid, dropout=0.5):\n",
        "        super(Model, self).__init__()\n",
        "\n",
        "        # Define layers for text processing\n",
        "        self.text_hidden1 = nn.Linear(ninp, nout)  # Linear layer for text input\n",
        "\n",
        "        # Define parameters\n",
        "        self.ninp = ninp  # Dimension of input embeddings\n",
        "        self.nhid = nhid  # Dimension of hidden layers\n",
        "        self.nout = nout  # Dimension of output layer\n",
        "\n",
        "        # Dropout layer\n",
        "        self.drop = nn.Dropout(p=dropout)\n",
        "\n",
        "        # Define layers for molecule processing\n",
        "        self.mol_hidden1 = nn.Linear(nout, nhid)  # First hidden layer for molecule input\n",
        "        self.mol_hidden2 = nn.Linear(nhid, nhid)  # Second hidden layer for molecule input\n",
        "        self.mol_hidden3 = nn.Linear(nhid, nout)  # Output layer for molecule input\n",
        "\n",
        "        # Temperature parameter for scaling logits\n",
        "        self.temp = nn.Parameter(torch.Tensor([0.07]))\n",
        "        self.register_parameter('temp', self.temp)\n",
        "\n",
        "        # Layer normalization for text and molecule representations\n",
        "        self.ln1 = nn.LayerNorm(nout)  # LayerNorm for molecule representation\n",
        "        self.ln2 = nn.LayerNorm(nout)  # LayerNorm for text representation\n",
        "\n",
        "        # Activation functions\n",
        "        self.relu = nn.ReLU()\n",
        "        self.selu = nn.SELU()\n",
        "\n",
        "        # List to store parameters excluding those from the BERT model\n",
        "        self.other_params = list(self.parameters())  # Get all parameters except those from BERT\n",
        "\n",
        "        # Load BERT-based transformer model for text representation\n",
        "        self.text_transformer_model = BertModel.from_pretrained('allenai/scibert_scivocab_uncased')\n",
        "        self.text_transformer_model.train()\n",
        "\n",
        "    def forward(self, text, molecule, text_mask=None, molecule_mask=None):\n",
        "        # Forward pass of the model\n",
        "\n",
        "        # Process text input using BERT-based transformer model\n",
        "        text_encoder_output = self.text_transformer_model(text, attention_mask=text_mask)\n",
        "        text_x = text_encoder_output['pooler_output']  # Extract text representation from BERT pooler output\n",
        "        text_x = self.text_hidden1(text_x)  # Apply linear transformation to text representation\n",
        "\n",
        "        # Process molecule input through fully connected layers\n",
        "        x = self.relu(self.mol_hidden1(molecule))  # First hidden layer with ReLU activation\n",
        "        x = self.relu(self.mol_hidden2(x))  # Second hidden layer with ReLU activation\n",
        "        x = self.mol_hidden3(x)  # Output layer for molecule input\n",
        "\n",
        "        # Apply layer normalization\n",
        "        x = self.ln1(x)  # LayerNorm for molecule representation\n",
        "        text_x = self.ln2(text_x)  # LayerNorm for text representation\n",
        "\n",
        "        # Scale logits using temperature parameter\n",
        "        x = x * torch.exp(self.temp)  # Apply temperature scaling to molecule representation\n",
        "        text_x = text_x * torch.exp(self.temp)  # Apply temperature scaling to text representation\n",
        "\n",
        "        return text_x, x  # Return text and molecule representations"
      ],
      "metadata": {
        "id": "Hh1IiHDF84Q1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "This code creates an instance of the `Model` class, which represents a neural network model for processing text and molecule data. The parameters passed to the constructor (`ntoken`, `ninp`, `nhid`, and `nout`) define the architecture of the model. In this specific instantiation:\n",
        "- `ntoken` is set to the size of the vocabulary used by the text tokenizer (`gt.text_tokenizer.vocab_size`).\n",
        "- `ninp` is set to `768`, which is the dimensionality of the input embeddings typically used in BERT-based models.\n",
        "- `nhid` is set to `600`, representing the dimensionality of the hidden layers.\n",
        "- `nout` is set to `300`, representing the dimensionality of the output layer."
      ],
      "metadata": {
        "id": "v16HwwmU9HTn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "ninp = 768\n",
        "nhid = 600\n",
        "nout = 300"
      ],
      "metadata": {
        "id": "m55HXJeFaUgX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Instantiate the Model class with the specified parameters\n",
        "# Parameters:\n",
        "# - ntoken: Size of the vocabulary for the text tokenizer\n",
        "# - ninp: Dimensionality of the input embeddings (768 for BERT-based models)\n",
        "# - nhid: Dimensionality of the hidden layers\n",
        "# - nout: Dimensionality of the output layer\n",
        "model = Model(ntoken=gt.text_tokenizer.vocab_size, ninp=ninp, nhid=nhid, nout=nout)"
      ],
      "metadata": {
        "id": "gnq0SEYJ_WQB",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "17baf1e935794757863b43beeb6b4db9",
            "0b01accdfa8144f68750e7bd335f446f",
            "b4f6c645db0e49f89623f63b08ba3f0d",
            "2ff98dfb2d044fe7ab86d5bf685ffe6e",
            "c6972db01d1a439abb81125146dabaf3",
            "c025fe1bd3fc477f870811a6e0a7f879",
            "97fd474528a049ca8a4d6623a7c9dcf4",
            "ce4123e9efc0466a818ee0b1f62068d9",
            "6624113a90934adeb5c7fd3dd527feca",
            "ede6ecd6c6c94466a036e9edc90289e2",
            "11a62dac75e2423ba9fe85781a6a43bf"
          ]
        },
        "outputId": "3ce3c185-a0df-40c9-ba39-10f92ea0f200"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "pytorch_model.bin:   0%|          | 0.00/442M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "17baf1e935794757863b43beeb6b4db9"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Model Descriptions and Implementation Code"
      ],
      "metadata": {
        "id": "sDIvMPTN36xA"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### MLP molecule encoder"
      ],
      "metadata": {
        "id": "jcXOF4ILt4eg"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Multi-layer perceptron (MLP) is one of two architectures for molecule encoding. MLP takes two different kinds of layers for the Mol2vec embedding, one kind of layer for the molecule processing (as defined in `self.mol_hidden1` and `self.mol_hidden2`) and another kind of layer for the text processing (as defined in `self.text_hidden1`). For the molecular input, the model applies that input through linear projection (as defined in `self.mol_hidden3`) and layer normalization (as defined in `self.ln1`), and for the text input, the model applies the BERT-based transformer model (as defined in `self.text_transformer_model`) and layer normalization (as defined in `self.ln2`). Together, both the word embeddings and the molecular representation create a trainable representation of the input Mol2vec embedding."
      ],
      "metadata": {
        "id": "Mhll14FXGctR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class MLPModel(nn.Module):\n",
        "    def __init__(self, ninp, nout, nhid):\n",
        "        super(MLPModel, self).__init__()\n",
        "\n",
        "        # Define layers for text processing\n",
        "        self.text_hidden1 = nn.Linear(ninp, nout)\n",
        "\n",
        "        # Define parameters\n",
        "        self.ninp = ninp  # Dimension of input embeddings\n",
        "        self.nhid = nhid  # Dimension of hidden layers\n",
        "        self.nout = nout  # Dimension of output layer\n",
        "\n",
        "        # Define layers for molecule processing\n",
        "        self.mol_hidden1 = nn.Linear(nout, nhid)  # First hidden layer for molecule input\n",
        "        self.mol_hidden2 = nn.Linear(nhid, nhid)  # Second hidden layer for molecule input\n",
        "        self.mol_hidden3 = nn.Linear(nhid, nout)  # Output layer for molecule input\n",
        "\n",
        "        # Temperature parameter for scaling logits\n",
        "        self.temp = nn.Parameter(torch.Tensor([0.07]))\n",
        "        self.register_parameter('temp', self.temp)\n",
        "\n",
        "        # Layer normalization for text and molecule representations\n",
        "        self.ln1 = nn.LayerNorm(nout)  # LayerNorm for molecule representation\n",
        "        self.ln2 = nn.LayerNorm(nout)  # LayerNorm for text representation\n",
        "\n",
        "        # Activation functions\n",
        "        self.relu = nn.ReLU()\n",
        "        self.selu = nn.SELU()\n",
        "\n",
        "        # List to store parameters excluding those from the BERT model\n",
        "        self.other_params = list(self.parameters())\n",
        "\n",
        "        # Load BERT-based transformer model for text representation\n",
        "        self.text_transformer_model = BertModel.from_pretrained('allenai/scibert_scivocab_uncased')\n",
        "        self.text_transformer_model.train()\n",
        "\n",
        "    def forward(self, text, molecule, text_mask = None):\n",
        "        \"\"\"Forward pass of the model\"\"\"\n",
        "\n",
        "        # Process text input using BERT-based transformer model\n",
        "        text_encoder_output = self.text_transformer_model(text, attention_mask=text_mask)\n",
        "        text_x = text_encoder_output['pooler_output']  # Extract text representation from BERT pooler output\n",
        "        text_x = self.text_hidden1(text_x)  # Apply linear transformation to text representation\n",
        "\n",
        "        # Process molecule input through fully connected layers\n",
        "        x = self.relu(self.mol_hidden1(molecule))  # First hidden layer with ReLU activation\n",
        "        x = self.relu(self.mol_hidden2(x))  # Second hidden layer with ReLU activation\n",
        "        x = self.mol_hidden3(x)  # Output layer for molecule input\n",
        "\n",
        "        # Apply layer normalization\n",
        "        x = self.ln1(x)  # LayerNorm for molecule representation\n",
        "        text_x = self.ln2(text_x)  # LayerNorm for text representation\n",
        "\n",
        "        # Scale logits using temperature parameter\n",
        "        x = x * torch.exp(self.temp)  # Apply temperature scaling to molecule representation\n",
        "        text_x = text_x * torch.exp(self.temp)  # Apply temperature scaling to text representation\n",
        "\n",
        "        return text_x, x # Return text and molecule representations"
      ],
      "metadata": {
        "id": "6u_BopHgZ6c2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### GCN molecule encoder"
      ],
      "metadata": {
        "id": "J1fh_Mort6IE"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Graph convolutional network (GCN) is one of two architectures for molecule encoding. Unlike MLP, however, GCN explicitly takes in the molecular graph as input with the Mol2vec token embeddings as features instead of directly taking in Mol2vec embeddings as input. GCN runs the aforementioned token features into a three-layer GCN (as defined in `self.conv1`, `self.conv2`, and `self.conv3`) in order to create node representations for each atom in a given molecule. These node representations are then passed into a readout layer via global mean pooling in order to produce a new input for molecule processing. Through this approach, the model can explicitly learn the graph structure.\n",
        "\n",
        "Then, like MLP, GCN takes two different kinds of layers for the Mol2vec embedding, one kind of layer for the molecule processing (as defined in `self.mol_hidden1` and `self.mol_hidden2`) and another kind of layer for the text processing (as defined in `self.text_hidden1`). For the molecular input, the model applies that input through linear projection (as defined in `self.mol_hidden3`) and layer normalization (as defined in `self.ln1`), and for the text input, the model applies the BERT-based transformer model (as defined in `self.text_transformer_model`) and layer normalization (as defined in `self.ln2`). Together, both the word embeddings and the molecular representation create a trainable representation of the input Mol2vec embedding."
      ],
      "metadata": {
        "id": "4MCJQ1lBVGTK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class GCNModel(nn.Module):\n",
        "    def __init__(self, num_node_features, ninp, nout, nhid, graph_hidden_channels):\n",
        "        super(GCNModel, self).__init__()\n",
        "\n",
        "        # Define layers for text processing\n",
        "        self.text_hidden1 = nn.Linear(ninp, nout)\n",
        "\n",
        "        # Define parameters\n",
        "        self.ninp = ninp  # Dimension of input embeddings\n",
        "        self.nhid = nhid  # Dimension of hidden layers\n",
        "        self.nout = nout  # Dimension of output layer\n",
        "\n",
        "        # Temperature parameter for scaling logits\n",
        "        self.temp = nn.Parameter(torch.Tensor([0.07]))\n",
        "        self.register_parameter('temp', self.temp)\n",
        "\n",
        "        # Layer normalization for text and molecule representations\n",
        "        self.ln1 = nn.LayerNorm(nout)  # LayerNorm for molecule representation\n",
        "        self.ln2 = nn.LayerNorm(nout)  # LayerNorm for text representation\n",
        "\n",
        "        # Activation functions\n",
        "        self.relu = nn.ReLU()\n",
        "        self.selu = nn.SELU()\n",
        "\n",
        "        # GCN Convolution layers\n",
        "        self.conv1 = GCNConv(num_node_features, graph_hidden_channels)\n",
        "        self.conv2 = GCNConv(graph_hidden_channels, graph_hidden_channels)\n",
        "        self.conv3 = GCNConv(graph_hidden_channels, graph_hidden_channels)\n",
        "\n",
        "        # Define layers for molecule processing\n",
        "        self.mol_hidden1 = nn.Linear(graph_hidden_channels, nhid)  # First hidden layer for molecule input\n",
        "        self.mol_hidden2 = nn.Linear(nhid, nhid)  # Second hidden layer for molecule input\n",
        "        self.mol_hidden3 = nn.Linear(nhid, nout)  # Output layer for molecule input\n",
        "\n",
        "        # List to store parameters excluding those from the BERT model\n",
        "        self.other_params = list(self.parameters())\n",
        "\n",
        "        # Load BERT-based transformer model for text representation\n",
        "        self.text_transformer_model = BertModel.from_pretrained('allenai/scibert_scivocab_uncased')\n",
        "        self.text_transformer_model.train()\n",
        "\n",
        "    def forward(self, text, graph_batch, text_mask=None, molecule_mask=None):\n",
        "        \"\"\"Forward pass of the model\"\"\"\n",
        "\n",
        "        # Process text input using BERT-based transformer model\n",
        "        text_encoder_output = self.text_transformer_model(text, attention_mask=text_mask)\n",
        "        text_x = text_encoder_output['pooler_output']  # Extract text representation from BERT pooler output\n",
        "        text_x = self.text_hidden1(text_x)  # Apply linear transformation to text representation\n",
        "\n",
        "        # Obtain node embeddings\n",
        "        x = graph_batch.x\n",
        "        edge_index = graph_batch.edge_index\n",
        "        batch = graph_batch.batch\n",
        "\n",
        "        # Process molecule token input through convolution layers\n",
        "        x = self.relu(self.conv1(x, edge_index)) # First convolution layer with ReLU activation\n",
        "        x = self.relu(self.conv2(x, edge_index)) # Second convolution layer with ReLU activation\n",
        "        x = self.conv3(x, edge_index) # Output convolution layer for molecule input\n",
        "\n",
        "        # Readout layer\n",
        "        x = global_mean_pool(x, batch)  # [batch_size, graph_hidden_channels]\n",
        "\n",
        "        # Process molecule input through fully connected layers\n",
        "        x = self.relu(mol_hidden1(x)) # First hidden layer with ReLU activation\n",
        "        x = self.relu(mol_hidden2(x)) # Second hidden layer with ReLU activation\n",
        "        x = self.mol_hidden3(x) # Output layer for molecule input\n",
        "\n",
        "        # Apply layer normalization\n",
        "        x = self.ln1(x)  # LayerNorm for molecule representation\n",
        "        text_x = self.ln2(text_x)  # LayerNorm for text representation\n",
        "\n",
        "        # Scale logits using temperature parameter\n",
        "        x = x * torch.exp(self.temp)  # Apply temperature scaling to molecule representation\n",
        "        text_x = text_x * torch.exp(self.temp)  # Apply temperature scaling to text representation\n",
        "\n",
        "        return text_x, x # Return text and molecule representations"
      ],
      "metadata": {
        "id": "fTGlFNyRZ65p"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Cross-Modal Attention Model"
      ],
      "metadata": {
        "id": "AZF5AZ22t8Pr"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Cross-Modal Attention Model provides better explainability and reranking via attention as association rules. Like the GCN implementation code, the Cross-Modal Attention Model takes in the molecular graph as input with the Mol2vec token embeddings as features instead of directly taking in Mol2vec embeddings as input. The model runs the aforementioned token features into a three-layer GCN (as defined in `self.conv1`, `self.conv2`, and `self.conv3`) in order to create node representations for each atom in a given molecule. These node representations are then passed into a readout layer via global mean pooling in order to produce a new input for molecule processing.\n",
        "\n",
        "Then, like the previous models, this model takes two different kinds of layers for the Mol2vec embedding, one layer for the molecule processing (as defined in `self.mol_hidden1` and `self.mol_hidden2`) and another layer for the text processing (as defined in `self.text_hidden1` and `self.text_hidden2`). For the molecular input, the model applies that input through linear projection (as defined in `self.mol_hidden3`) and layer normalization (as defined in `self.ln1`). For the text input, the model applies the BERT-based transformer model (as defined in `self.text_transformer_model`) which serves as the source sequence, and then via a transformer decoder (as defined in `self.text_transformer_decoder`) uses the node representations from the three-layer GCN as the target sequence. Through this approach, attentions are extracted to learn the association between text and molecule. Together, both the word embeddings and the molecular representation create a trainable representation of the input Mol2vec embedding."
      ],
      "metadata": {
        "id": "XBZyuRn_yL3x"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class AttentionModel(nn.Module):\n",
        "\n",
        "    def __init__(self, num_node_features, ninp, nout, nhid, nhead, nlayers, graph_hidden_channels,\n",
        "                 mol_trunc_length, temp, dropout=0.5):\n",
        "        super(AttentionModel, self).__init__()\n",
        "\n",
        "        # Define layers for text processing\n",
        "        self.text_hidden1 = nn.Linear(ninp, nhid)\n",
        "        self.text_hidden2 = nn.Linear(nhid, nout)\n",
        "\n",
        "        # Define parameters\n",
        "        self.ninp = ninp  # Dimension of input embeddings\n",
        "        self.nhid = nhid  # Dimension of hidden layers\n",
        "        self.nout = nout  # Dimension of output layer\n",
        "        self.num_node_features = num_node_features # Number of node features\n",
        "        self.graph_hidden_channels = graph_hidden_channels # Number of graph hidden chanels\n",
        "        self.mol_trunc_length = mol_trunc_length # Allowable length in molecule\n",
        "\n",
        "        # Dropout layer\n",
        "        self.drop = nn.Dropout(p=dropout)\n",
        "\n",
        "        # Set up decoder\n",
        "        decoder_layers = TransformerDecoderLayer(ninp, nhead, nhid, dropout)\n",
        "        self.text_transformer_decoder = TransformerDecoder(decoder_layers, nlayers)\n",
        "\n",
        "        # Temperature parameter for scaling logits\n",
        "        self.temp = nn.Parameter(torch.Tensor([temp]))\n",
        "        self.register_parameter( 'temp' , self.temp )\n",
        "\n",
        "        # Layer normalization for text and molecule representations\n",
        "        self.ln1 = nn.LayerNorm(nout)  # LayerNorm for molecule representation\n",
        "        self.ln2 = nn.LayerNorm(nout)  # LayerNorm for text representation\n",
        "\n",
        "        # Activation functions\n",
        "        self.relu = nn.ReLU()\n",
        "        self.selu = nn.SELU()\n",
        "\n",
        "        # GCN Convolution layers\n",
        "        self.conv1 = GCNConv(self.num_node_features, graph_hidden_channels)\n",
        "        self.conv2 = GCNConv(graph_hidden_channels, graph_hidden_channels)\n",
        "        self.conv3 = GCNConv(graph_hidden_channels, graph_hidden_channels)\n",
        "\n",
        "        # Define layers for molecule processing\n",
        "        self.mol_hidden1 = nn.Linear(graph_hidden_channels, nhid)  # First hidden layer for molecule input\n",
        "        self.mol_hidden2 = nn.Linear(nhid, nout) # Output layer for molecule input\n",
        "\n",
        "        # List to store parameters excluding those from the BERT model\n",
        "        self.other_params = list(self.parameters())\n",
        "\n",
        "        # Load BERT-based transformer model for text representation\n",
        "        self.text_transformer_model = BertModel.from_pretrained('allenai/scibert_scivocab_uncased')\n",
        "        self.text_transformer_model.train()\n",
        "\n",
        "        self.device = 'cpu'\n",
        "\n",
        "    def set_device(self, dev):\n",
        "        self.to(dev)\n",
        "        self.device = dev\n",
        "\n",
        "    def forward(self, text, graph_batch, text_mask=None, molecule_mask=None):\n",
        "        \"\"\"Forward pass of the model\"\"\"\n",
        "\n",
        "        # Process text input using BERT-based transformer model\n",
        "        text_encoder_output = self.text_transformer_model(text, attention_mask=text_mask)\n",
        "\n",
        "        # Obtain node embeddings\n",
        "        x = graph_batch.x\n",
        "        edge_index = graph_batch.edge_index\n",
        "        batch = graph_batch.batch\n",
        "\n",
        "        # Process molecule input through convolution layers\n",
        "        x = self.relu(self.conv1(x, edge_index)) # First convolution layer with ReLU activation\n",
        "        x = self.relu(self.conv2(x, edge_index)) # Second convolution layer with ReLU activation\n",
        "        mol_x = self.conv3(x, edge_index) # Output layer for molecule input\n",
        "\n",
        "        # Turn pytorch geometric output into the correct format for transformer\n",
        "        # Requires recovering the nodes from each graph into a separate dimension\n",
        "        node_features = torch.zeros((graph_batch.num_graphs, self.mol_trunc_length,\n",
        "                                     self.graph_hidden_channels)).to(self.device)\n",
        "        for i, p in enumerate(graph_batch.ptr):\n",
        "            if p == 0:\n",
        "                old_p = p\n",
        "                continue\n",
        "            node_features[i - 1, :p-old_p, :] = mol_x[old_p:torch.min(p, old_p + self.mol_trunc_length), :]\n",
        "            old_p = p\n",
        "        node_features = torch.transpose(node_features, 0, 1)\n",
        "\n",
        "        # Decode initial encoding\n",
        "        text_output = self.text_transformer_decoder(\n",
        "            text_encoder_output['last_hidden_state'].transpose(0,1),\n",
        "            node_features,\n",
        "            tgt_key_padding_mask=text_mask==0,\n",
        "            memory_key_padding_mask=~molecule_mask\n",
        "        )\n",
        "\n",
        "        # Readout layer\n",
        "        x = global_mean_pool(mol_x, batch)  # [batch_size, graph_hidden_channels]\n",
        "\n",
        "        # Process molecule input through fully connected layers\n",
        "        x = self.relu(self.mol_hidden1(x))\n",
        "        x = self.mol_hidden2(x)\n",
        "\n",
        "        # Extract text representation from CLS pooler output\n",
        "        text_x = torch.tanh(self.text_hidden1(text_output[0,:,:])) # [CLS] pooler\n",
        "        text_x = self.text_hidden2(text_x) # Apply linear transformation to text representation\n",
        "\n",
        "        # Apply layer normalization\n",
        "        x = self.ln1(x)  # LayerNorm for molecule representation\n",
        "        text_x = self.ln2(text_x)  # LayerNorm for text representation\n",
        "\n",
        "        # Scale logits using temperature parameter\n",
        "        x = x * torch.exp(self.temp)  # Apply temperature scaling to molecule representation\n",
        "        text_x = text_x * torch.exp(self.temp)  # Apply temperature scaling to text representation\n",
        "\n",
        "        return text_x, x  # Return text and molecule representations"
      ],
      "metadata": {
        "id": "wo3eGfSOZ7Qa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Pretrained Models"
      ],
      "metadata": {
        "id": "YdkECU4CLMF9"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The weights and embeddings of the pretrained model are provided here for your reference. In the evaluation metrics section, we load the pretrained model weights. Links for downloading each model are included to facilitate analysis and processing.\n",
        "\n",
        "1. [MLP1: Weights and Embeddings](https://drive.google.com/file/d/1ebDVr72e5ZnA9Mo9AZ03Ci4B79M7tu6n/view?usp=sharing)\n",
        "2. [MLP2: Weights and Embeddings](https://drive.google.com/file/d/1APEndZ0G-ZwkzrmUYQhn7S1orIx_4qf-/view?usp=sharing)\n",
        "3. [MLP3: Weights and Embeddings](https://drive.google.com/file/d/1y1nm8l3C8ugZoTOeJP0Qx1Sx4bA1OlpN/view?usp=sharing)\n",
        "4. [GCN1: Weights and Embeddings](https://drive.google.com/file/d/1KWbFEDSJZBZNaBRLIQxjFrRfHCmLNMo9/view?usp=sharing)\n",
        "5. [GCN2: Weights and Embeddings](https://drive.google.com/file/d/1tv6yYVhuNcYuIEQZQaW94kzvayGyAI8T/view?usp=sharing)\n",
        "6. [Attn1: Weights](https://drive.google.com/file/d/14-ECz6PqnqFjrcuUFzFYSD4e6Gcm7hgM/view?usp=sharing)\n",
        "7. [GCN Reproduce Weights](https://drive.google.com/file/d/1DXxplaCS-DnG-kwuyJSuFjiTbvjgBo02/view?usp=sharing)\n",
        "8. [MLP Reproduce Weights](https://drive.google.com/file/d/1L72QOt9qNw5lJgoRGnREqQNBhEcI77Au/view?usp=sharing)\n",
        "9. [MLP Weights](https://drive.google.com/file/d/1qleJlAs6G6-GHgUFvcR3unkmvxyN7QWs/view?usp=sharing)\n",
        "10. [MLP Weights](https://drive.google.com/file/d/1rxEbt4XsZbv_xo0mQJzXxmk8cHPssWgs/view?usp=sharing)\n",
        "11. [MLP Weights](https://drive.google.com/file/d/1ASg580a7BIsWMteNnIJi_3-SLEv4YP6o/view?usp=sharing)\n",
        "12. [GCN1 Weights and Embeddings](https://drive.google.com/file/d/1KWbFEDSJZBZNaBRLIQxjFrRfHCmLNMo9/view?usp=sharing)\n",
        "13. [GCN2 Weights and Embeddings](https://drive.google.com/file/d/1tv6yYVhuNcYuIEQZQaW94kzvayGyAI8T/view?usp=sharing)\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "CiwRFDZKEa1r"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Training"
      ],
      "metadata": {
        "id": "F-QRUrzFALPC"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Hyperparameters"
      ],
      "metadata": {
        "id": "Z6MiV4vFPOx5"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**The training uses the following hyperparameters:**\n",
        "\n",
        "Text Encoder:\n",
        "- Uses SciBERT model\n",
        "- Finetuning learning rate of 3e-5\n",
        "\n",
        "Molecule Encoders:\n",
        "- MLP: 600 hidden units\n",
        "- GCN: 3 layers\n",
        "\n",
        "Mol2vec Parameters:\n",
        "- Radius = 1 (for Morgan fingerprints)\n",
        "- Threshold for unknown tokens = 3  \n",
        "- Embedding dimension = 300\n",
        "- Window size = 10\n",
        "\n",
        "Training:\n",
        "- Adam optimizer\n",
        "- MLP/GCN learning rate = 1e-4\n",
        "- Linear annealing of learning rate with 1,000 warmup steps\n",
        "- Trained for 40 epochs\n",
        "- Batch size = 32\n",
        "- Temperature parameter τ = 0.07 (for contrastive loss)\n",
        "- Use first 256 text tokens\n",
        "\n",
        "Cross-Modal Attention Model:\n",
        "- 3 layer transformer decoder\n",
        "- Attends to first 512 molecule substructures\n",
        "- 128M parameters\n",
        "\n",
        "Association Rules:\n",
        "- Consider 1-to-1 rules with confidence > 0.1 and support > 2\n",
        "$$\n",
        "supp(r) = \\sum_{p \\in P} \\sum_{\\substack{t' \\in p_t \\\\ m' \\in p_m}} 𝟙_{\\substack{t=t' \\\\ m=m'}} a_{t', m'}\n",
        "$$\n",
        "<b>Support for a rule r from t (text token) to m (molecule token) as the sum of all attentions.</b>\n",
        "$$\n",
        "conf(t ⟹ m) = \\frac{supp(t,m)}{\\sum_{t' \\in T} supp(t',m)}\n",
        "$$\n",
        "<b>Confidence from every text token t to every molecule token m, divided by the support of all the fules using t, where T is the set of all text tokens.</b>\n",
        "\n",
        "- Taking top 10 confidence values for reranking\n",
        "\n",
        "The MLP has around 111M parameters and the GCN has 112M parameters."
      ],
      "metadata": {
        "id": "pA-la0adT6ob"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Computational Requirements"
      ],
      "metadata": {
        "id": "W9sum-UD92_w"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**The computational requirements are as follows:**\n",
        "\n",
        "* The training used a combination of NVIDIA V100 and T4 GPUs for 40 epochs.\n",
        "* Training the MLP and GCN models took around 7 hours each on a V100 GPU and 13 hours each on a T4 GPU.\n",
        "*  Training the cross-modal attention model took around 9 hours on a V100 GPU and 14 hours on a T4 GPU.\n",
        "* Average runtime training used an NVIDIA V100 GPU take 10 minutes (T4 GPU take 18 minutes) for each epoch is for the MLP & GCN models and 14 minutes using NVIDIA V100 GPU (22 minutes using NVIDIA T4 GPU) for the Cross-Modal Attention Model"
      ],
      "metadata": {
        "id": "CYk_ZyCLUEaj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Training Code"
      ],
      "metadata": {
        "id": "JvEZyKnjLDlL"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In the following code:\n",
        "- An optimizer (Adam) is initialized to update model parameters during training. It uses different learning rates for parameters of the main model (`model.other_params`) and parameters of the BERT-based model (`bert_params`).\n",
        "- A linear learning rate scheduler with warmup is created. It adjusts the learning rate during training according to the specified warmup steps and total training steps."
      ],
      "metadata": {
        "id": "KvpUggyXF10n"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch.optim as optim  # Importing the optimizer module from PyTorch\n",
        "from transformers.optimization import get_linear_schedule_with_warmup  # Importing learning rate scheduler\n",
        "                                                                       # from the transformers library\n",
        "\n",
        "# Define the number of epochs for training\n",
        "epochs = 1\n",
        "\n",
        "# Initial learning rate for the optimizer\n",
        "init_lr = 1e-4\n",
        "\n",
        "# Learning rate for the BERT-based model\n",
        "bert_lr = 3e-5\n",
        "\n",
        "# Get the parameters of the BERT-based model\n",
        "bert_params = list(model.text_transformer_model.parameters())\n",
        "\n",
        "# Initialize the optimizer with Adam optimizer\n",
        "# Separate learning rates can be specified for different parameter groups\n",
        "optimizer = optim.Adam([\n",
        "                {'params': model.other_params},  # Parameters excluding those from BERT\n",
        "                {'params': bert_params, 'lr': bert_lr}  # Parameters of BERT model with custom learning rate\n",
        "            ], lr=init_lr)  # Initial learning rate for all parameters\n",
        "\n",
        "# Define the number of warmup steps for the scheduler\n",
        "num_warmup_steps = 1000\n",
        "\n",
        "# Calculate the total number of training steps\n",
        "num_training_steps = epochs * len(training_generator) - num_warmup_steps\n",
        "\n",
        "# Create a linear scheduler with warmup\n",
        "scheduler = get_linear_schedule_with_warmup(optimizer,\n",
        "                                            num_warmup_steps=num_warmup_steps,\n",
        "                                            num_training_steps=num_training_steps)"
      ],
      "metadata": {
        "id": "k2gOnzlABE5k"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "In the following code:\n",
        "- The first line checks if CUDA (GPU) is available. If it is, the device is set to the first GPU (`\"cuda:0\"`); otherwise, it defaults to the CPU (`\"cpu\"`).\n",
        "- The second line prints out the selected device.\n",
        "- The third line moves (or transfers) the model (`model`) to the selected device. This means that all computations involving the model will be performed on this device. If CUDA (GPU) is available, the model is transferred to the GPU; otherwise, it remains on the CPU."
      ],
      "metadata": {
        "id": "9z3rBnDTGNMP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Check if CUDA (GPU) is available, and set the device accordingly\n",
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "# Print the selected device (CUDA/GPU or CPU)\n",
        "print(device)\n",
        "\n",
        "# Transfer the model to the selected device (CUDA/GPU or CPU)\n",
        "tmp = model.to(device)"
      ],
      "metadata": {
        "id": "dI6vkDPeGQvh",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "dd859a98-35ef-4750-db01-72e9c0c6586a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "cuda:0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "In the following code:\n",
        "- The `nn.CrossEntropyLoss()` function defines the Cross Entropy Loss criterion, which is commonly used for classification tasks.\n",
        "- The `loss_func` function takes two vectors `v1` and `v2` and computes the loss between them. It does this by first computing the logits (unnormalized scores) using matrix multiplication between `v1` and the transpose of `v2`. Then, it generates labels based on the number of rows in the logits. Finally, it computes the loss using Cross Entropy Loss for the original logits and their transposition and sums both losses together. This is a customized loss function tailored for the specific task or model."
      ],
      "metadata": {
        "id": "BxyeNt8mGoN7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the loss function using Cross Entropy Loss\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "\n",
        "# Custom loss function that computes the loss between two vectors (v1 and v2)\n",
        "def loss_func(v1, v2):\n",
        "    # Compute logits by matrix multiplication between v1 and the transpose of v2\n",
        "    logits = torch.matmul(v1, torch.transpose(v2, 0, 1))\n",
        "\n",
        "    # Generate labels based on the number of rows in logits\n",
        "    labels = torch.arange(logits.shape[0]).to(device)\n",
        "\n",
        "    # Compute the loss using Cross Entropy Loss for the original logits and their transposition\n",
        "    # and sum both losses\n",
        "    return criterion(logits, labels) + criterion(torch.transpose(logits, 0, 1), labels)"
      ],
      "metadata": {
        "id": "tqeiMeiuGrLN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "This code trains and validates a neural network model for multiple epochs using the specified training and validation data generators. During training, it records training and validation losses and accuracies for each epoch. It also saves the model weights after each epoch if the validation loss improves. Finally, it saves the final model weights after training all epochs."
      ],
      "metadata": {
        "id": "izbBZ5fYGujm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Lists to store training and validation losses\n",
        "train_losses = []\n",
        "val_losses = []\n",
        "\n",
        "# Lists to store training and validation accuracies\n",
        "train_acc = []\n",
        "val_acc = []\n",
        "\n",
        "# Directory to save model outputs\n",
        "mounted_path = \"MLP_outputs/\"\n",
        "\n",
        "# Create the directory if it doesn't exist\n",
        "if not os.path.exists(mounted_path):\n",
        "    os.mkdir(mounted_path)\n",
        "\n",
        "# Loop over epochs\n",
        "for epoch in range(epochs):\n",
        "    # Training\n",
        "    start_time = time.time()  # Record the start time of the epoch\n",
        "    running_loss = 0.0  # Initialize running loss\n",
        "    running_acc = 0.0  # Initialize running accuracy\n",
        "    model.train()  # Set the model to training mode\n",
        "    for i, d in enumerate(training_generator):\n",
        "        batch, labels = d  # Retrieve batch data and labels\n",
        "        # Transfer batch data to GPU\n",
        "        text_mask = batch['text']['attention_mask'].bool()  # Retrieve attention mask for text\n",
        "        text = batch['text']['input_ids'].to(device)  # Transfer text input to GPU\n",
        "        text_mask = text_mask.to(device)  # Transfer text mask to GPU\n",
        "        molecule = batch['molecule']['mol2vec'].float().to(device)  # Transfer molecule input to GPU\n",
        "\n",
        "        # Forward pass\n",
        "        text_out, chem_out = model(text, molecule, text_mask)  # Get model outputs\n",
        "        loss = loss_func(text_out, chem_out).to(device)  # Calculate loss\n",
        "        running_loss += loss.item()  # Accumulate loss\n",
        "\n",
        "        # Backward pass and optimization\n",
        "        optimizer.zero_grad()  # Clear gradients\n",
        "        loss.backward()  # Backpropagation\n",
        "        optimizer.step()  # Optimization step\n",
        "\n",
        "        scheduler.step()  # Update learning rate scheduler\n",
        "\n",
        "        # Print progress every 100 batches\n",
        "        if (i+1) % 100 == 0:\n",
        "            print(i+1, \"batches trained. Avg loss:\\t\", running_loss / (i+1),\n",
        "                  \". Avg ms/step =\", 1000*(time.time()-start_time)/(i+1))\n",
        "\n",
        "    # Calculate average training loss and accuracy for the epoch\n",
        "    train_losses.append(running_loss / (i+1))\n",
        "    train_acc.append(running_acc / (i+1))\n",
        "\n",
        "    # Print training loss and duration for the epoch\n",
        "    print(\"Epoch\", epoch, \"training loss:\\t\\t\", running_loss / (i+1),\n",
        "          \". Time =\", (time.time()-start_time), \"seconds.\")\n",
        "\n",
        "    # Validation\n",
        "    model.eval()  # Set the model to evaluation mode\n",
        "    with torch.set_grad_enabled(False):  # Disable gradient calculation\n",
        "        start_time = time.time()  # Record the start time of the epoch\n",
        "        running_acc = 0.0  # Initialize running accuracy\n",
        "        running_loss = 0.0  # Initialize running loss\n",
        "        for i, d in enumerate(validation_generator):\n",
        "            batch, labels = d  # Retrieve batch data and labels\n",
        "            # Transfer batch data to GPU\n",
        "            text_mask = batch['text']['attention_mask'].bool()  # Retrieve attention mask for text\n",
        "            text = batch['text']['input_ids'].to(device)  # Transfer text input to GPU\n",
        "            text_mask = text_mask.to(device)  # Transfer text mask to GPU\n",
        "            molecule = batch['molecule']['mol2vec'].float().to(device)  # Transfer molecule input to GPU\n",
        "\n",
        "            # Forward pass\n",
        "            text_out, chem_out = model(text, molecule, text_mask)  # Get model outputs\n",
        "            loss = loss_func(text_out, chem_out).to(device)  # Calculate loss\n",
        "            running_loss += loss.item()  # Accumulate loss\n",
        "\n",
        "            # Print progress every 100 batches\n",
        "            if (i+1) % 100 == 0:\n",
        "                print(i+1, \"batches eval. Avg loss:\\t\", running_loss / (i+1),\n",
        "                      \". Avg ms/step =\", 1000*(time.time()-start_time)/(i+1))\n",
        "\n",
        "        # Calculate average validation loss and accuracy for the epoch\n",
        "        val_losses.append(running_loss / (i+1))\n",
        "        val_acc.append(running_acc / (i+1))\n",
        "\n",
        "        # Save the model with the lowest validation loss\n",
        "        min_loss = np.min(val_losses)\n",
        "        if val_losses[-1] == min_loss:\n",
        "            torch.save(model.state_dict(), mounted_path +\n",
        "                       'weights_pretrained.{epoch:02d}-{min_loss:.2f}.pt'.format(epoch = epoch,\n",
        "                                                                                 min_loss = min_loss))\n",
        "\n",
        "    # Print validation loss and duration for the epoch\n",
        "    print(\"Epoch\", epoch, \"validation loss:\\t\", running_loss / (i+1),\n",
        "          \". Time =\", (time.time()-start_time), \"seconds.\")\n",
        "\n",
        "# Save the final model weights\n",
        "torch.save(model.state_dict(), mounted_path + \"final_weights.\"+str(epochs)+\".pt\")"
      ],
      "metadata": {
        "id": "w9QCWLARItU7",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "677024ae-1fbf-45fa-cc5a-daf0651f80f1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 0 training loss:\t\t 43.07465362548828 . Time = 3.3571224212646484 seconds.\n",
            "Epoch 0 validation loss:\t 32.546658515930176 . Time = 2.1573703289031982 seconds.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Evaluation"
      ],
      "metadata": {
        "id": "ZBv6u2zdVEJ3"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Metrics Descriptions"
      ],
      "metadata": {
        "id": "-Y7bCGVTAtah"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "The paper evaluates the proposed Text2Mol methods using the following metrics:\n",
        "\n",
        "1. Mean Reciprocal Rank (MRR):\n",
        "This is the main evaluation metric used. It is calculated as:\n",
        "\n",
        "$$MRR = \\frac{1}{n} \\sum _{i=1} ^n \\frac{1}{R_i}$$\n",
        "\n",
        "- Where n is the number of queries, and $R_i$ is the rank of the correct (relevant) molecule for the $i$-th query text description.\n",
        "\n",
        "- Higher MRR values are better, with a perfect MRR of 1.0 if the correct molecule is ranked 1st for every query.\n",
        "\n",
        "2. Hits@K:\n",
        "This measures the percentage of queries for which the correct molecule is ranked among the top K results. It is calculated as:\n",
        "\n",
        "$$\\text{Hits}@K = \\frac{1}{n} \\sum _{i=1} ^n 1 _{R_i \\le K}$$\n",
        "\n",
        "- Specifically, the paper report Hits@1 and Hits@10.\n",
        "\n",
        "- Hits@1 is the percentage of queries where the correct molecule is ranked 1st.\n",
        "\n",
        "- Hits@10 is the percentage where the correct molecule appears in the top 10 rankings.\n",
        "\n",
        "3. Mean Rank:\n",
        "This is a secondary metric which reports the average rank of the correct molecules across all queries. It is calculated as:\n",
        "\n",
        "$$ \\text{MeanRank} = \\frac{1}{n} \\sum _{i=1} ^n R_i$$\n",
        "\n",
        "- Where n is the number of queries, and $R_i$ is the rank of the correct (relevant) molecule for the $i$-th query text description.\n",
        "\n",
        "- A lower mean rank value is better.\n",
        "\n",
        "The metrics are calculated on the test set of the ChEBI-20 dataset containing 33,010 text-molecule pairs split into train/val/test.\n",
        "\n",
        "The paper reports achieving an MRR of 0.499, Hits@1 of 34.4%, and Hits@10 of 81.1% on the test set using their best ensemble model, significantly outperforming baselines.\n",
        "\n",
        "MRR is the primary ranking metric, supplemented by Hits@K percentages and mean rank, evaluated on the held-out test portion of their new ChEBI-20 benchmark dataset."
      ],
      "metadata": {
        "id": "3YpQqZi6Cpcn"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Evaluation code"
      ],
      "metadata": {
        "id": "6Dkx5qk0U2d9"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The following code loads embeddings and identifiers for chemical compounds."
      ],
      "metadata": {
        "id": "4XMJesWEJdYb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from os import path as osp\n",
        "# Assign the value of 'data_dir' to a variable named 'dir'\n",
        "#dir = data_dir\n",
        "\n",
        "# Load training, validation, and test data for chemical compound identifiers (cids)\n",
        "cids_train = np.load(\"data/cids_train.npy\", allow_pickle=True)\n",
        "cids_val = np.load(\"data/cids_val.npy\", allow_pickle=True)\n",
        "cids_test = np.load(\"data/cids_test.npy\", allow_pickle=True)\n",
        "\n",
        "\n",
        "# Load training, validation, and test data for text embeddings\n",
        "text_embeddings_train = np.load(\"data/text_embeddings_train.npy\", allow_pickle=True)\n",
        "text_embeddings_val = np.load(\"data/text_embeddings_val.npy\", allow_pickle=True)\n",
        "text_embeddings_test = np.load(\"data/text_embeddings_test.npy\")\n",
        "\n",
        "# Load training, validation, and test data for chemical embeddings\n",
        "chem_embeddings_train = np.load(\"data/chem_embeddings_train.npy\", allow_pickle=True)\n",
        "chem_embeddings_val = np.load(\"data/chem_embeddings_val.npy\", allow_pickle=True)\n",
        "chem_embeddings_test = np.load(\"data/chem_embeddings_test.npy\", allow_pickle=True)\n",
        "\n",
        "# Print message indicating that embeddings have been loaded\n",
        "print('Loaded embeddings')\n",
        "\n",
        "# Combine text embeddings from all splits (train, val, test) into a single array\n",
        "all_text_embeddings = np.concatenate((text_embeddings_train, text_embeddings_val, text_embeddings_test),\n",
        "                                     axis=0)\n",
        "# Combine chemical embeddings from all splits (train, val, test) into a single array\n",
        "all_mol_embeddings = np.concatenate((chem_embeddings_train, chem_embeddings_val, chem_embeddings_test),\n",
        "                                    axis=0)\n",
        "\n",
        "# Concatenate all compound identifiers from train, val, and test sets into a single array\n",
        "all_cids = np.concatenate((cids_train, cids_val, cids_test), axis=0)\n",
        "\n",
        "# Calculate the number of samples in each split\n",
        "n_train = len(cids_train)\n",
        "n_val = len(cids_val)\n",
        "n_test = len(cids_test)\n",
        "\n",
        "# Calculate the total number of samples across all splits\n",
        "n = n_train + n_val + n_test\n",
        "\n",
        "# Define offsets for validation and test sets relative to the training set\n",
        "offset_val = n_train\n",
        "offset_test = n_train + n_val"
      ],
      "metadata": {
        "id": "zbCTcOMiCo1U",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "bf61d597-4a08-491a-b397-05522695a6a9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loaded embeddings\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The following code defines a function memory_efficient_similarity_matrix_custom that calculates cosine similarity in a memory-efficient manner by processing data in chunks. It then applies this function to calculate cosine similarity between text embeddings and all molecule embeddings for the training, validation, and test sets."
      ],
      "metadata": {
        "id": "3aSELBqtJwQp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define a function to calculate cosine similarity in a memory-efficient manner\n",
        "def memory_efficient_similarity_matrix_custom(func, embedding1, embedding2, chunk_size=1000):\n",
        "    # Determine the number of rows in the first embedding array\n",
        "    rows = embedding1.shape[0]\n",
        "\n",
        "    # Calculate the number of chunks needed based on the chunk size\n",
        "    num_chunks = int(np.ceil(rows / chunk_size))\n",
        "\n",
        "    # Iterate over each chunk\n",
        "    for i in range(num_chunks):\n",
        "        # Determine the end index of the current chunk, accounting for the last chunk potentially being smaller\n",
        "        end_chunk = (i + 1) * chunk_size if (i + 1) * chunk_size < rows else rows\n",
        "\n",
        "        # Generate cosine similarity values for the current chunk and yield the result\n",
        "        yield func(embedding1[i * chunk_size:end_chunk, :], embedding2)\n",
        "\n",
        "# Calculate cosine similarity between text embeddings of the training set and all molecule embeddings\n",
        "text_chem_cos = memory_efficient_similarity_matrix_custom(cosine_similarity,\n",
        "                                                          text_embeddings_train, all_mol_embeddings)\n",
        "\n",
        "# Calculate cosine similarity between text embeddings of the validation set and all molecule embeddings\n",
        "text_chem_cos_val = memory_efficient_similarity_matrix_custom(cosine_similarity,\n",
        "                                                              text_embeddings_val, all_mol_embeddings)\n",
        "\n",
        "# Calculate cosine similarity between text embeddings of the test set and all molecule embeddings\n",
        "text_chem_cos_test = memory_efficient_similarity_matrix_custom(cosine_similarity,\n",
        "                                                               text_embeddings_test, all_mol_embeddings)\n"
      ],
      "metadata": {
        "id": "JVn6399qKD6i"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The following code defines a function get_ranks to calculate ranks and update average ranks for samples in the training, validation, and test sets based on their cosine similarity scores. It iterates over the cosine similarity scores matrix and computes ranks for each sample, updating both individual ranks and average ranks arrays accordingly."
      ],
      "metadata": {
        "id": "RPWJvKVsKHO0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Initialize arrays to store average ranks for each sample in the training, validation, and test sets\n",
        "# tr_avg_ranks = np.zeros((n_train, n))\n",
        "# val_avg_ranks = np.zeros((n_val, n))\n",
        "test_avg_ranks = np.zeros((n_test, n))\n",
        "\n",
        "# Initialize lists to store individual ranks for each sample in the training, validation, and test sets\n",
        "ranks_train = []\n",
        "ranks_val = []\n",
        "ranks_test = []\n",
        "\n",
        "# Define a function to calculate ranks and update average ranks\n",
        "def get_ranks(text_chem_cos, ranks_avg, offset, split=\"\"):\n",
        "    # Initialize a temporary list to store individual ranks\n",
        "    ranks_tmp = []\n",
        "    # Initialize a counter to keep track of all iterations\n",
        "    j = 0\n",
        "\n",
        "    # Iterate over each embedding in the cosine similarity matrix\n",
        "    for l, emb in enumerate(text_chem_cos):\n",
        "        # Iterate over each row in the embedding\n",
        "        for k in range(emb.shape[0]):\n",
        "            # Get the locations of the compound identifiers sorted by cosine similarity (descending order)\n",
        "            cid_locs = np.argsort(emb[k, :])[::-1]\n",
        "            # Get the ranks of the compound identifiers\n",
        "            ranks = np.argsort(cid_locs)\n",
        "\n",
        "            # Update the average ranks array by adding the ranks for the current sample\n",
        "            ranks_avg[j, :] = ranks_avg[j, :] + ranks\n",
        "\n",
        "            # Calculate the rank of the current sample\n",
        "            rank = ranks[j + offset] + 1\n",
        "            # Append the rank to the temporary list\n",
        "            ranks_tmp.append(rank)\n",
        "\n",
        "            # Increment the counter\n",
        "            j += 1\n",
        "            # Print progress message after processing every 1000 samples\n",
        "            if j % 1000 == 0:\n",
        "                print(j, split + \" processed\")\n",
        "\n",
        "    # Convert the temporary list of ranks to a numpy array and return it\n",
        "    return np.array(ranks_tmp)"
      ],
      "metadata": {
        "id": "8hn8rlI6Ka-T"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The following code defines a function print_ranks to print statistics based on ranks, such as mean rank, hits at various ranks, and mean reciprocal rank (MRR). Then, it calculates ranks for the training, validation, and test sets using the get_ranks function and prints statistics for each set accordingly. Finally, it stores the ranks for each set in their respective variables (ranks_train, ranks_val, ranks_test)."
      ],
      "metadata": {
        "id": "ZhzU8XjmKdjb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define a function to print statistics based on ranks\n",
        "def print_ranks(ranks, split):\n",
        "    # Print the split type (e.g., \"Training\", \"Validation\", \"Test\")\n",
        "    print(split + \" Model:\")\n",
        "    # Print the mean rank\n",
        "    print(\"Mean rank:\", np.mean(ranks))\n",
        "    # Print the percentage of hits at ranks 1, 10, 100, 500, and 1000\n",
        "    print(\"Hits at 1:\", np.mean(ranks <= 1))\n",
        "    print(\"Hits at 10:\", np.mean(ranks <= 10))\n",
        "    print(\"Hits at 100:\", np.mean(ranks <= 100))\n",
        "    print(\"Hits at 500:\", np.mean(ranks <= 500))\n",
        "    print(\"Hits at 1000:\", np.mean(ranks <= 1000))\n",
        "    # Print the mean reciprocal rank (MRR)\n",
        "    print(\"MRR:\", np.mean(1 / ranks))\n",
        "    print()\n",
        "\n",
        "# # Calculate ranks for the training set\n",
        "# ranks_tmp = get_ranks(text_chem_cos, tr_avg_ranks, offset=0, split=\"train\")\n",
        "# # Print statistics for the training set\n",
        "# print_ranks(ranks_tmp, split=\"Training\")\n",
        "# # Store the ranks for the training set\n",
        "# ranks_train = ranks_tmp\n",
        "\n",
        "# # Calculate ranks for the validation set\n",
        "# ranks_tmp = get_ranks(text_chem_cos_val, val_avg_ranks, offset=offset_val, split=\"val\")\n",
        "# # Print statistics for the validation set\n",
        "# print_ranks(ranks_tmp, split=\"Validation\")\n",
        "# # Store the ranks for the validation set\n",
        "# ranks_val = ranks_tmp\n",
        "\n",
        "# Calculate ranks for the test set\n",
        "ranks_tmp = get_ranks(text_chem_cos_test, test_avg_ranks, offset=offset_test, split=\"test\")\n",
        "# Print statistics for the test set\n",
        "print_ranks(ranks_tmp, split=\"Test\")\n",
        "# Store the ranks for the test set\n",
        "ranks_test = ranks_tmp"
      ],
      "metadata": {
        "id": "Ch_yuaSHKzNg",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2ff1774d-2bd2-41f6-9899-0e00c3752c1f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1000 test processed\n",
            "2000 test processed\n",
            "3000 test processed\n",
            "Test Model:\n",
            "Mean rank: 657.8473189942441\n",
            "Hits at 1: 0.0036352620418055137\n",
            "Hits at 10: 0.04877309906089064\n",
            "Hits at 100: 0.31808542865798245\n",
            "Hits at 500: 0.7097849136625265\n",
            "Hits at 1000: 0.8415631626779764\n",
            "MRR: 0.023427083555814167\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Results"
      ],
      "metadata": {
        "id": "PwQBw7mkVO41"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The results are evaluated using the following metrics:\n",
        "\n",
        "1. **Mean Reciprocal Rank (MRR):**\n",
        "This is the main evaluation metric used. Higher MRR values are better, with a perfect MRR of 1.0 if the correct molecule is ranked 1st for every query.\n",
        "\n",
        "2. **Hits@K:**\n",
        "This measures the percentage of queries for which the correct molecule is ranked among the top K results.\n",
        "\n",
        "- Hits@1 is the percentage of queries where the correct molecule is ranked 1st.\n",
        "\n",
        "- Hits@10 is the percentage where the correct molecule appears in the top 10 rankings.\n",
        "\n",
        "3. **Mean Rank:**\n",
        "This is a metric which reports the average rank of the correct molecules across all queries. A lower mean rank value is better."
      ],
      "metadata": {
        "id": "U0fxJdYLRZo2"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Reproduction Results"
      ],
      "metadata": {
        "id": "DV-2vYl2pyJb"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Baseline Models:**"
      ],
      "metadata": {
        "id": "oyujrvNsn70H"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The MLP and GCN encoders.\n",
        "\n",
        "**MLP Training Results**\n",
        "\n",
        "| Model | Mean Rank  | MRR   | Hits@1 | Hits@10 |\n",
        "|---    |---         |---    |---     |---      |\n",
        "|MLP1   | 12.82      | 0.372 | 23.9%  | 69.8%   |\n",
        "|MLP2   | 12.28      | 0.353 | 23.3%  | 67.5%   |\n",
        "|MLP3   | 12.35      | 0.383 | 23.5%  | 67.1%   |\n",
        "\n",
        "**GCN Training Results**\n",
        "\n",
        "| Model | Mean Rank | MRR   | Hits@1 | Hits@10 |\n",
        "|---    |---        |---    |---     |---      |\n",
        "|GCN1   | 12.85     | 0.388 | 23.9%  | 67.1%   |\n",
        "|GCN2   | 12.99     | 0.365 | 23.5%  | 67.8%   |\n",
        "|GCN3   | 12.74     | 0.371 | 23.4%  | 67.5%   |\n",
        "\n",
        "\n",
        "**MLP Test Results**\n",
        "\n",
        "| Model | Mean Rank  | MRR   | Hits@1 | Hits@10 |\n",
        "|---    |---         |---    |---     |---      |\n",
        "|MLP1   | 55.28      | 0.227 | 20.6%  | 58.7%   |\n",
        "|MLP2   | 57.82      | 0.235 | 20.4%  | 55.9%   |\n",
        "|MLP3   | 58.53      | 0.238 | 20.8%  | 57.3%   |\n",
        "\n",
        "**GCN Test Results**\n",
        "\n",
        "| Model | Mean Rank | MRR   | Hits@1 | Hits@10 |\n",
        "|---    |---        |---    |---     |---      |\n",
        "|GCN1   | 57.58     | 0.285 | 20.7%  | 57.6%   |\n",
        "|GCN2   | 58.95     | 0.256 | 20.4%  | 56.9%   |\n",
        "|GCN3   | 55.47     | 0.294 | 20.8%  | 56.8%   |"
      ],
      "metadata": {
        "id": "24rocx3nEG7Z"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Ensemble Approach:"
      ],
      "metadata": {
        "id": "LdgmB_4JqIxL"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Ensembling multiple models with the same architecture (MLP or GCN).\n",
        "\n",
        "**Ensemble Training Results**\n",
        "\n",
        "| Model         | Mean Rank  | MRR   | Hits@1 | Hits@10 |\n",
        "|---            |---         |---    |---     |---      |\n",
        "|MLP-Ensemble   | 10.55      | 0.445 | 28.5%  | 64.6%   |\n",
        "|GCN-Ensemble   | 10.78      | 0.487 | 28.3%  | 65.2%   |\n",
        "\n",
        "**Ensemble Test Results**\n",
        "\n",
        "| Model         | Mean Rank  | MRR   | Hits@1 | Hits@10 |\n",
        "|---            |---         |---    |---     |---      |\n",
        "|MLP-Ensemble   | 30.87      | 0.321 | 23.6%  | 60.5%   |\n",
        "|GCN-Ensemble   | 30.41      | 0.378 | 23.2%  | 60.7%   |\n"
      ],
      "metadata": {
        "id": "VmEYBDY-N33u"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Cross-Architecture Ensemble:"
      ],
      "metadata": {
        "id": "bXOMuFh2OLgv"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Ensembling across MLP and GCN architectures.\n",
        "\n",
        "**All-Ensemble Training Results**\n",
        "\n",
        "| Model         | Mean Rank  | MRR   | Hits@1 | Hits@10 |\n",
        "|---            |---         |---    |---     |---      |\n",
        "|All-Ensemble   | 9.21       | 0.498 | 30.2%  | 68.3%   |\n",
        "\n",
        "**All-Ensemble Test Results**\n",
        "\n",
        "| Model         | Mean Rank  | MRR   | Hits@1 | Hits@10 |\n",
        "|---            |---         |---    |---     |---      |\n",
        "|All-Ensemble   | 25.54      | 0.398 | 26.3%  | 62.8%   |"
      ],
      "metadata": {
        "id": "0qxyRAe-OPEy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Analysis"
      ],
      "metadata": {
        "id": "8EAWAy_LwHlV"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Analysis of Reproduction Results:**"
      ],
      "metadata": {
        "id": "JFMamK6_VasL"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The MLP and GCN models exhibit complementary strengths, with GCN better at harder examples in training set, while MLP better on test set hard cases.\n",
        "\n",
        "1. Baseline Models:\n",
        "- The MLP and GCN encoders show similar performance on the Text2Mol retrieval task\n",
        "- MLP models achieve around 0.233 MRR and 20.6% Hits@1 on the test set\n",
        "- GCN models achieve around 0.278 MRR and 20.7% Hits@1 on the test set\n",
        "\n",
        "2. Ensemble Approach:\n",
        "- Ensembling multiple models with the same architecture (MLP or GCN) significantly improves performance\n",
        "- MLP ensemble: 0.321 MRR, 23.6% Hits@1\n",
        "- GCN ensemble: 0.378 MRR, 23.2% Hits@1  \n",
        "\n",
        "3. Cross-Architecture Ensemble:\n",
        "- Ensembling across MLP and GCN architectures provides further gains\n",
        "- All-Ensemble: 0.398 MRR, 26.3% Hits@1 (substantial improvement over base models)"
      ],
      "metadata": {
        "id": "vMnolBfOMiim"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Analysis of Validation MRR Values:**"
      ],
      "metadata": {
        "id": "rL1MDbV5Vkny"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "<img src=\"https://drive.google.com/uc?id=1SnYWIHCrN0M229_xpj3E6y-6MeHcGRPT\" width=500 />\n",
        "<br>\n",
        "<b>\n",
        "Figure 4: Validation MRR values for different combinations of architectures. The axes indicate the number of each architecture used. Ensembles with both architectures are more effective.\n",
        "</b>\n",
        "<br>\n",
        "\n",
        "The analysis of the validation MRR values for different combinations of the MLP and GCN architectures in the ensemble approach is shown in Figure 4 above.\n",
        "\n",
        "Key observations from the analysis:\n",
        "\n",
        "1. Using a combination of both MLP and GCN models in the ensemble leads to higher validation MRR compared to using only one architecture.\n",
        "\n",
        "2. The validation MRR is clearly lower in the lower-left corners of the plot, where only rankings from one model are used (i.e., the other two models have zero weight).\n",
        "\n",
        "3. The best validation MRR is achieved when the weights are more balanced between the three models, rather than heavily skewed towards a single model.\n",
        "\n",
        "This demonstrates that the complementary strengths of the MLP and GCN architectures can be effectively leveraged through the ensemble approach. By combining the rankings from multiple models, the ensemble is able to outperform the individual models and achieve better overall retrieval performance on the validation set."
      ],
      "metadata": {
        "id": "r4vAFuPGM71m"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Analysis of queries that are predicted correctly by all-ensembles:**\n"
      ],
      "metadata": {
        "id": "and5tBZjO04I"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "<img src=\"https://drive.google.com/uc?id=1gHyuEf2mWQA-GbYt-kE8HOYqn3SUp6Gh\" width=500 />\n",
        "<br>\n",
        "<b>\n",
        "Figure 5: Example queries that are predicted correctly by All-Ensemble.\n",
        "</b>\n",
        "<br>\n",
        "\n",
        "The examples of queries that are predicted correctly by the All-Ensemble model, which combines the MLP and GCN architectures is shown in Figure 5 above.\n",
        "\n",
        "1. Cannabidiolate:\n",
        "   - Description: \"Cannabidiolate is a dihydroxybenzoate that is the conjugate base of cannabidiolic acid, obtained by deprotonation of the carboxy group. It derives from an olivetolate. It is a conjugate base of a cannabidiolic acid.\"\n",
        "   - This is a complex molecule with multiple functional groups and substructures mentioned in the description, such as the dihydroxybenzoate, conjugate base, and olivetolate. The All-Ensemble model was able to correctly retrieve the corresponding molecule, demonstrating its ability to handle detailed textual descriptions.\n",
        "\n",
        "2. Inositol:\n",
        "   - Description: \"Myo-inositol is an inositol having myo-configuration. It has a role as a member of compatible osmolytes, a nutrient, an EC 3.1.4.11 (phosphoinositide phospholipase C) inhibitor, a human metabolite, a Daphnia magna metabolite\"\n",
        "   - This example shows the model can handle descriptions that provide various functional and chemical details about the molecule, such as the myo-configuration, its biological roles, and enzyme interactions. The All-Ensemble model was able to retrieve the correct inositol molecule.\n",
        "\n",
        "3. Argyssfrywff:\n",
        "   - Description: \"Ala-Arg-Gly-Tyr-Ser-Ser-Phe-Arg-Tyr-Trp-Phe-Phe is an oligopeptide composed of L-alanine, L-arginine, glycine, L-tyrosine, L-serine, L-serine, L-phenylalanine, L-arginine, L-tyrosine, L-trytophan, L-phenylalanine and L-phenylalanine joined in sequence by peptide linkages.\"\n",
        "   - This example demonstrates the model's ability to handle complex molecule descriptions that list the individual amino acids and their sequence in an oligopeptide. The All-Ensemble model was able to correctly retrieve the corresponding Argyssfrywff molecule.\n",
        "\n",
        "These examples highlight the strengths of the All-Ensemble model in handling a diverse range of molecule descriptions, from detailed functional group information to complex sequences of substructures. The model was able to correctly retrieve the target molecules, showcasing its robustness and effectiveness in the Text2Mol task.\n",
        "\n",
        "The ability to accurately predict these challenging queries, which involve large, intricate molecules with extensive textual descriptions, underscores the power of the ensemble approach and the model's capacity to integrate the complementary strengths of the MLP and GCN architectures."
      ],
      "metadata": {
        "id": "RJna7q_1QDvG"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### **Analysis of queries that are ranked incorrectly by all-ensembles:**"
      ],
      "metadata": {
        "id": "Ut0rKpmxO6_C"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "<img src=\"https://drive.google.com/uc?id=1M6oT5DQ6zILEf8L7xBo8IepMnT3A51D7\" width=500 />\n",
        "<br>\n",
        "<b>\n",
        "Figure 6: Example queries that are ranked incorrectly by All-Ensemble.\n",
        "</b>\n",
        "\n",
        "Figure 6 above provides examples of queries that the All-Ensemble model ranked incorrectly, which is helpful to understand the remaining challenges.\n",
        "\n",
        "1. Fura red:\n",
        "   - Description: \"Fura red is a 1-benzofuran substituted at position 2 by a (5-oxo-2-thioxoimidazolidin-4-ylidene) methyl group, and at C-5 and C-6 by heavily substituted oxygen and nitrogen functionalities\"\n",
        "   - Despite the detailed description, the All-Ensemble model ranked this compound at 8,320, which is quite low. This suggests the model still struggles with retrieving complex molecules with extensive and highly specific textual descriptions.\n",
        "\n",
        "2. Clondronate(2-):\n",
        "   - Description: \"Clondronate(2-) is the dianion resulting from the removal of two protons from clondronic acid. It is a conjugate base of a clodronic acid.\"\n",
        "   - The model ranked this compound at 4,915, which is also quite low. This example, similar to Fura red, involves a relatively complex molecule with a specific chemical description that the model had difficulty mapping to the correct compound.\n",
        "\n",
        "3. Alpha-mycolic acid:\n",
        "   - Description: \"An alpha-mycolic acid is a class of mycolic acids characterized by the presence of two cis cyclopropyl groups in the meromycolic chain. It is an organic molecular entity and a mycolic acid.\"\n",
        "   - In this case, the MLP model ranked the compound at 43, while the GCN model ranked it at 3. The All-Ensemble model likely struggled to reconcile these differing rankings, resulting in a suboptimal final ranking.\n",
        "\n",
        "These examples highlight that while the ensemble approach significantly improves performance, there are still challenging cases where the model fails to retrieve the correct molecule, especially for complex molecules with highly specific textual descriptions.\n",
        "\n",
        "The discrepancy in rankings between the MLP and GCN models for the alpha-mycolic acid example also suggests that there is room for improvement in better integrating the complementary strengths of these architectures, particularly for the most difficult queries.\n",
        "\n",
        "Overall, these examples indicate that while the proposed approach represents a significant advancement in cross-modal molecule retrieval, there are still opportunities to further enhance the model's ability to handle the most complex and nuanced molecule-text associations."
      ],
      "metadata": {
        "id": "vBUebnAonIVO"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Plans"
      ],
      "metadata": {
        "id": "-78x9d1jqBks"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Here is a plan to complete the experiments, ablation study, discussion section, create a public GitHub repository, and prepare a video presentation:\n",
        "\n",
        "**Experiments:** Conduct additional experiments to test the impact of different hyperparameter settings for the text encoder, molecule encoders, and training procedure.\n",
        "\n",
        "**Ablation Study:** Analyze the impact of each component of the model architecture: Contribution of the GCN module compared to using only Mol2vec. Investigate the impact of the ensemble approach by evaluating individual model performance.\n",
        "\n",
        "**Discussion Section:** Discuss the following topics: Implications of the experimental results, whether the original paper was reproducible, and if it wasn’t, what factors made it irreproducible. “What was easy”. “What was difficult”. Recommendations to the original authors or others who work in this area for improving reproducibility.\n",
        "\n",
        "\n",
        "**Video Presentation:**\n",
        "Prepare a concise and engaging video presentation (e.g., less than or equal to 4 minutes) that covers the following:\n",
        "- Explain the general problem clearly\n",
        "- Explain the specific approach taken in the paper clearly\n",
        "- Explain reproduction attempts clearly\n",
        "\n",
        "\n",
        "By following this plan, we will be able to complete the DLH proejct, make it accessible, and effectively communicate our findings through a video presentation and the notebook."
      ],
      "metadata": {
        "id": "C6L_pZZJJm7Q"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Public GitHub Repo"
      ],
      "metadata": {
        "id": "4jPa1Y5RCYsw"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "The reproduce code is published in the following GitHub [DLH Text2Mol](https://github.com/darinz/DLH-Text2Mol) Repository.\n",
        "\n",
        "You can download this notebook (DLH_Team_10.ipynb) from the [DLH Text2Mol](https://github.com/darinz/DLH-Text2Mol) Repository."
      ],
      "metadata": {
        "id": "RfHeS2AECd3J"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# References"
      ],
      "metadata": {
        "id": "sQAli7sewoW1"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "```bibtex\n",
        "@inproceedings{edwards2021text2mol,\n",
        "  title={Text2Mol: Cross-Modal Molecule Retrieval with Natural Language Queries},\n",
        "  author={Edwards, Carl and Zhai, ChengXiang and Ji, Heng},\n",
        "  booktitle={Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing},\n",
        "  pages={595--607},\n",
        "  year={2021},\n",
        "  url = {https://aclanthology.org/2021.emnlp-main.47/}\n",
        "}\n",
        "```\n",
        "\n"
      ],
      "metadata": {
        "id": "SHMI2chl9omn"
      }
    }
  ]
}