notebooks/subcellular_localization_linear_probe.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "L4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "4d33ce73985b4d7f99699eb8bc496137": {
          "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_b77649b94b5a4eb3b859e9dc7ea3a8c9",
              "IPY_MODEL_7dca7d901d55482aa93c12853c47a8cc",
              "IPY_MODEL_4d9d6620338343aca3a941da33d2d230"
            ],
            "layout": "IPY_MODEL_06a076f0632d4cd690da4579a073f986"
          }
        },
        "b77649b94b5a4eb3b859e9dc7ea3a8c9": {
          "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_1cbe6087915447e4a5872fa307f1d572",
            "placeholder": "​",
            "style": "IPY_MODEL_112d1030ef8e4d2fbb5ab64014b279da",
            "value": "tokenizer_config.json: 100%"
          }
        },
        "7dca7d901d55482aa93c12853c47a8cc": {
          "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_527b171c978643ada5c33cc51a748739",
            "max": 95,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_096306bac75a404c97d2c90ec22a211e",
            "value": 95
          }
        },
        "4d9d6620338343aca3a941da33d2d230": {
          "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_6ac7e73b375d4f9497db77e4af815462",
            "placeholder": "​",
            "style": "IPY_MODEL_376f923f62dd48f683abb86ab3f76a17",
            "value": " 95.0/95.0 [00:00&lt;00:00, 7.16kB/s]"
          }
        },
        "06a076f0632d4cd690da4579a073f986": {
          "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
          }
        },
        "1cbe6087915447e4a5872fa307f1d572": {
          "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
          }
        },
        "112d1030ef8e4d2fbb5ab64014b279da": {
          "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": ""
          }
        },
        "527b171c978643ada5c33cc51a748739": {
          "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
          }
        },
        "096306bac75a404c97d2c90ec22a211e": {
          "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": ""
          }
        },
        "6ac7e73b375d4f9497db77e4af815462": {
          "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
          }
        },
        "376f923f62dd48f683abb86ab3f76a17": {
          "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": ""
          }
        },
        "44929914f97d46e28677eb1495265bbe": {
          "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_0d5b443197d3445fbc41d32f85b4ad92",
              "IPY_MODEL_1077bc34487e4d2a9f7c4f3c78cf6231",
              "IPY_MODEL_a5514d85aa974d5fb099d1c7d05893d7"
            ],
            "layout": "IPY_MODEL_c2fe8c18f4984739a71c9bf8d88489da"
          }
        },
        "0d5b443197d3445fbc41d32f85b4ad92": {
          "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_b32b708749344858b3c90a4e61f31c1a",
            "placeholder": "​",
            "style": "IPY_MODEL_6327e62190f14ac193fc031a96cd1202",
            "value": "vocab.txt: 100%"
          }
        },
        "1077bc34487e4d2a9f7c4f3c78cf6231": {
          "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_d1ee79aacd4f4628954b4e3a00db89f9",
            "max": 93,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_cbe8d646dd7446b6ace524ff5009fa4a",
            "value": 93
          }
        },
        "a5514d85aa974d5fb099d1c7d05893d7": {
          "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_8dc2017da87e4145adb7e4883eed1947",
            "placeholder": "​",
            "style": "IPY_MODEL_cf1b3303b3114977ba6a29528e5967e4",
            "value": " 93.0/93.0 [00:00&lt;00:00, 8.15kB/s]"
          }
        },
        "c2fe8c18f4984739a71c9bf8d88489da": {
          "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
          }
        },
        "b32b708749344858b3c90a4e61f31c1a": {
          "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
          }
        },
        "6327e62190f14ac193fc031a96cd1202": {
          "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": ""
          }
        },
        "d1ee79aacd4f4628954b4e3a00db89f9": {
          "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
          }
        },
        "cbe8d646dd7446b6ace524ff5009fa4a": {
          "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": ""
          }
        },
        "8dc2017da87e4145adb7e4883eed1947": {
          "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
          }
        },
        "cf1b3303b3114977ba6a29528e5967e4": {
          "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": ""
          }
        },
        "15eb4bb0e076459995e0c7147004bf33": {
          "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_be016dc56862487893d10ce2f63819ad",
              "IPY_MODEL_4c2722d7c2a0442990639ea2649be48a",
              "IPY_MODEL_7c1e0aea0103422bbc7a6a8416a439f0"
            ],
            "layout": "IPY_MODEL_d9a1df056ebe42b399e1ab21e4d2f812"
          }
        },
        "be016dc56862487893d10ce2f63819ad": {
          "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_22f5957820fe4bd7ac26d56a024ccb07",
            "placeholder": "​",
            "style": "IPY_MODEL_5ff378f1d8fc4f07b586610f5bfc9f9b",
            "value": "special_tokens_map.json: 100%"
          }
        },
        "4c2722d7c2a0442990639ea2649be48a": {
          "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_d04c3282c81b40bea370d9c04078fa9e",
            "max": 125,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_719985ce8bb24dafb11fb3ed808f381e",
            "value": 125
          }
        },
        "7c1e0aea0103422bbc7a6a8416a439f0": {
          "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_1c7da8e5658247cdb5110db2feef0bfb",
            "placeholder": "​",
            "style": "IPY_MODEL_65765ad302f14f9ba6511fc0385df9c0",
            "value": " 125/125 [00:00&lt;00:00, 11.4kB/s]"
          }
        },
        "d9a1df056ebe42b399e1ab21e4d2f812": {
          "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
          }
        },
        "22f5957820fe4bd7ac26d56a024ccb07": {
          "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
          }
        },
        "5ff378f1d8fc4f07b586610f5bfc9f9b": {
          "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": ""
          }
        },
        "d04c3282c81b40bea370d9c04078fa9e": {
          "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
          }
        },
        "719985ce8bb24dafb11fb3ed808f381e": {
          "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": ""
          }
        },
        "1c7da8e5658247cdb5110db2feef0bfb": {
          "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
          }
        },
        "65765ad302f14f9ba6511fc0385df9c0": {
          "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": ""
          }
        },
        "a61bee1afcec48fe907048b3308d5513": {
          "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_5b8c0a6ee7824d3083528e77872e9e27",
              "IPY_MODEL_9b0bc31a91c6489b93ba249b10fa324d",
              "IPY_MODEL_99aa79ea8614459c84c101089d31d71b"
            ],
            "layout": "IPY_MODEL_441d60f6afbf4ab1be22fc6cc2cc57d5"
          }
        },
        "5b8c0a6ee7824d3083528e77872e9e27": {
          "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_d941b05c14324544ae309cd452228a62",
            "placeholder": "​",
            "style": "IPY_MODEL_27f45ab4d1a94977b1f402bf3d4661a7",
            "value": "config.json: 100%"
          }
        },
        "9b0bc31a91c6489b93ba249b10fa324d": {
          "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_e254ec98096243c6a5fbd98ce67f5a6f",
            "max": 724,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_7ee2d1463f5541279ac6ee5bd4d83afc",
            "value": 724
          }
        },
        "99aa79ea8614459c84c101089d31d71b": {
          "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_744c02939fad4be9ad414b5acb943da8",
            "placeholder": "​",
            "style": "IPY_MODEL_3776d7114c7e43008069d16a8d06d569",
            "value": " 724/724 [00:00&lt;00:00, 67.7kB/s]"
          }
        },
        "441d60f6afbf4ab1be22fc6cc2cc57d5": {
          "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
          }
        },
        "d941b05c14324544ae309cd452228a62": {
          "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
          }
        },
        "27f45ab4d1a94977b1f402bf3d4661a7": {
          "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": ""
          }
        },
        "e254ec98096243c6a5fbd98ce67f5a6f": {
          "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
          }
        },
        "7ee2d1463f5541279ac6ee5bd4d83afc": {
          "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": ""
          }
        },
        "744c02939fad4be9ad414b5acb943da8": {
          "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
          }
        },
        "3776d7114c7e43008069d16a8d06d569": {
          "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": ""
          }
        },
        "622dd8b5f7444233959ab5c2845e841f": {
          "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_c32e1e99a5eb44a9abfe086b7bb6200e",
              "IPY_MODEL_e7c288801de644c2bdeeaa60340f12a2",
              "IPY_MODEL_2df384ea3fef41fd906ef4e3d33dc309"
            ],
            "layout": "IPY_MODEL_ee57d1bf6bd949db9aaf762d759014b5"
          }
        },
        "c32e1e99a5eb44a9abfe086b7bb6200e": {
          "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_6960bf58ecdd4306a9b5d6aa22de71e6",
            "placeholder": "​",
            "style": "IPY_MODEL_b5201c0311cc4012a4ea072223baa2b2",
            "value": "model.safetensors: 100%"
          }
        },
        "e7c288801de644c2bdeeaa60340f12a2": {
          "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_68b91f656b6246ce96e0b4846856a5fa",
            "max": 2609506392,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_bf6a7595b91f409db35e136773699ecd",
            "value": 2609506392
          }
        },
        "2df384ea3fef41fd906ef4e3d33dc309": {
          "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_8e023b2ec83446cfb8180be4983ada45",
            "placeholder": "​",
            "style": "IPY_MODEL_e0fe36770acf41609e8a0f0d24735a50",
            "value": " 2.61G/2.61G [00:11&lt;00:00, 214MB/s]"
          }
        },
        "ee57d1bf6bd949db9aaf762d759014b5": {
          "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
          }
        },
        "6960bf58ecdd4306a9b5d6aa22de71e6": {
          "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
          }
        },
        "b5201c0311cc4012a4ea072223baa2b2": {
          "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": ""
          }
        },
        "68b91f656b6246ce96e0b4846856a5fa": {
          "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
          }
        },
        "bf6a7595b91f409db35e136773699ecd": {
          "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": ""
          }
        },
        "8e023b2ec83446cfb8180be4983ada45": {
          "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
          }
        },
        "e0fe36770acf41609e8a0f0d24735a50": {
          "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": ""
          }
        },
        "f3a236ca8b6b43fb86e026a0cc28e6ee": {
          "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_8ad615fc3bb6463eb7f94246e3a3e626",
              "IPY_MODEL_9e17402df2cb485da8eda46ea7161af4",
              "IPY_MODEL_67fb910605474bc2b3cd2042abe5e6dd"
            ],
            "layout": "IPY_MODEL_3ba98d08fa404235a13496e26c0cd17b"
          }
        },
        "8ad615fc3bb6463eb7f94246e3a3e626": {
          "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_158f4fd5d8ca46c2bbb7698557b43ff7",
            "placeholder": "​",
            "style": "IPY_MODEL_2c133a764418444eb424755de94ba659",
            "value": "esm2_plm1280_l28_sae4096.safetensors: 100%"
          }
        },
        "9e17402df2cb485da8eda46ea7161af4": {
          "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_6c0dd248a3664fea9fd2ad239214300d",
            "max": 41997696,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_17da7c0c322042a2808b61b39ab76037",
            "value": 41997696
          }
        },
        "67fb910605474bc2b3cd2042abe5e6dd": {
          "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_58ffbb372d154922aa5ae75ec64143f9",
            "placeholder": "​",
            "style": "IPY_MODEL_500c037a889c4b32ac60605c29cfba0d",
            "value": " 42.0M/42.0M [00:00&lt;00:00, 42.6MB/s]"
          }
        },
        "3ba98d08fa404235a13496e26c0cd17b": {
          "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
          }
        },
        "158f4fd5d8ca46c2bbb7698557b43ff7": {
          "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
          }
        },
        "2c133a764418444eb424755de94ba659": {
          "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": ""
          }
        },
        "6c0dd248a3664fea9fd2ad239214300d": {
          "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
          }
        },
        "17da7c0c322042a2808b61b39ab76037": {
          "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": ""
          }
        },
        "58ffbb372d154922aa5ae75ec64143f9": {
          "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
          }
        },
        "500c037a889c4b32ac60605c29cfba0d": {
          "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": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Setup\n",
        "\n",
        "Install InterProt, load ESM and SAE"
      ],
      "metadata": {
        "id": "7Eaz3UmJF7S4"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "EFOR7aiGFbtW"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install git+https://github.com/etowahadams/interprot.git"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import csv\n",
        "import os\n",
        "import torch\n",
        "import warnings\n",
        "\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import polars as pl\n",
        "\n",
        "from huggingface_hub import hf_hub_download\n",
        "from safetensors.torch import load_file\n",
        "from sklearn.linear_model import LogisticRegression\n",
        "from sklearn.model_selection import GridSearchCV, PredefinedSplit\n",
        "from sklearn.metrics import accuracy_score\n",
        "from tqdm import tqdm\n",
        "from transformers import AutoTokenizer, EsmModel\n",
        "\n",
        "from interprot.sae_model import SparseAutoencoder\n",
        "from interprot.utils import get_layer_activations\n",
        "\n",
        "ESM_DIM = 1280\n",
        "SAE_DIM = 4096\n",
        "LAYER = 28\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "# Load ESM model\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"facebook/esm2_t33_650M_UR50D\")\n",
        "esm_model = EsmModel.from_pretrained(\"facebook/esm2_t33_650M_UR50D\")\n",
        "esm_model.to(device)\n",
        "esm_model.eval()\n",
        "\n",
        "# Load SAE model\n",
        "checkpoint_path = hf_hub_download(\n",
        "    repo_id=\"liambai/InterProt-ESM2-SAEs\",\n",
        "    filename=\"esm2_plm1280_l28_sae4096.safetensors\"\n",
        ")\n",
        "sae_model = SparseAutoencoder(ESM_DIM, SAE_DIM)\n",
        "sae_model.load_state_dict(load_file(checkpoint_path))\n",
        "sae_model.to(device)\n",
        "sae_model.eval()"
      ],
      "metadata": {
        "id": "GP2EorV7GEMG",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 390,
          "referenced_widgets": [
            "4d33ce73985b4d7f99699eb8bc496137",
            "b77649b94b5a4eb3b859e9dc7ea3a8c9",
            "7dca7d901d55482aa93c12853c47a8cc",
            "4d9d6620338343aca3a941da33d2d230",
            "06a076f0632d4cd690da4579a073f986",
            "1cbe6087915447e4a5872fa307f1d572",
            "112d1030ef8e4d2fbb5ab64014b279da",
            "527b171c978643ada5c33cc51a748739",
            "096306bac75a404c97d2c90ec22a211e",
            "6ac7e73b375d4f9497db77e4af815462",
            "376f923f62dd48f683abb86ab3f76a17",
            "44929914f97d46e28677eb1495265bbe",
            "0d5b443197d3445fbc41d32f85b4ad92",
            "1077bc34487e4d2a9f7c4f3c78cf6231",
            "a5514d85aa974d5fb099d1c7d05893d7",
            "c2fe8c18f4984739a71c9bf8d88489da",
            "b32b708749344858b3c90a4e61f31c1a",
            "6327e62190f14ac193fc031a96cd1202",
            "d1ee79aacd4f4628954b4e3a00db89f9",
            "cbe8d646dd7446b6ace524ff5009fa4a",
            "8dc2017da87e4145adb7e4883eed1947",
            "cf1b3303b3114977ba6a29528e5967e4",
            "15eb4bb0e076459995e0c7147004bf33",
            "be016dc56862487893d10ce2f63819ad",
            "4c2722d7c2a0442990639ea2649be48a",
            "7c1e0aea0103422bbc7a6a8416a439f0",
            "d9a1df056ebe42b399e1ab21e4d2f812",
            "22f5957820fe4bd7ac26d56a024ccb07",
            "5ff378f1d8fc4f07b586610f5bfc9f9b",
            "d04c3282c81b40bea370d9c04078fa9e",
            "719985ce8bb24dafb11fb3ed808f381e",
            "1c7da8e5658247cdb5110db2feef0bfb",
            "65765ad302f14f9ba6511fc0385df9c0",
            "a61bee1afcec48fe907048b3308d5513",
            "5b8c0a6ee7824d3083528e77872e9e27",
            "9b0bc31a91c6489b93ba249b10fa324d",
            "99aa79ea8614459c84c101089d31d71b",
            "441d60f6afbf4ab1be22fc6cc2cc57d5",
            "d941b05c14324544ae309cd452228a62",
            "27f45ab4d1a94977b1f402bf3d4661a7",
            "e254ec98096243c6a5fbd98ce67f5a6f",
            "7ee2d1463f5541279ac6ee5bd4d83afc",
            "744c02939fad4be9ad414b5acb943da8",
            "3776d7114c7e43008069d16a8d06d569",
            "622dd8b5f7444233959ab5c2845e841f",
            "c32e1e99a5eb44a9abfe086b7bb6200e",
            "e7c288801de644c2bdeeaa60340f12a2",
            "2df384ea3fef41fd906ef4e3d33dc309",
            "ee57d1bf6bd949db9aaf762d759014b5",
            "6960bf58ecdd4306a9b5d6aa22de71e6",
            "b5201c0311cc4012a4ea072223baa2b2",
            "68b91f656b6246ce96e0b4846856a5fa",
            "bf6a7595b91f409db35e136773699ecd",
            "8e023b2ec83446cfb8180be4983ada45",
            "e0fe36770acf41609e8a0f0d24735a50",
            "f3a236ca8b6b43fb86e026a0cc28e6ee",
            "8ad615fc3bb6463eb7f94246e3a3e626",
            "9e17402df2cb485da8eda46ea7161af4",
            "67fb910605474bc2b3cd2042abe5e6dd",
            "3ba98d08fa404235a13496e26c0cd17b",
            "158f4fd5d8ca46c2bbb7698557b43ff7",
            "2c133a764418444eb424755de94ba659",
            "6c0dd248a3664fea9fd2ad239214300d",
            "17da7c0c322042a2808b61b39ab76037",
            "58ffbb372d154922aa5ae75ec64143f9",
            "500c037a889c4b32ac60605c29cfba0d"
          ]
        },
        "outputId": "da5932dc-9d94-41d4-abd5-efa1e36e972c"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_auth.py:94: 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": [
              "tokenizer_config.json:   0%|          | 0.00/95.0 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "4d33ce73985b4d7f99699eb8bc496137"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "vocab.txt:   0%|          | 0.00/93.0 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "44929914f97d46e28677eb1495265bbe"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "special_tokens_map.json:   0%|          | 0.00/125 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "15eb4bb0e076459995e0c7147004bf33"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "config.json:   0%|          | 0.00/724 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "a61bee1afcec48fe907048b3308d5513"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "model.safetensors:   0%|          | 0.00/2.61G [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "622dd8b5f7444233959ab5c2845e841f"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Some weights of EsmModel were not initialized from the model checkpoint at facebook/esm2_t33_650M_UR50D and are newly initialized: ['esm.pooler.dense.bias', 'esm.pooler.dense.weight']\n",
            "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "esm2_plm1280_l28_sae4096.safetensors:   0%|          | 0.00/42.0M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "f3a236ca8b6b43fb86e026a0cc28e6ee"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "SparseAutoencoder()"
            ]
          },
          "metadata": {},
          "execution_count": 2
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Load the subcellular localization dataset"
      ],
      "metadata": {
        "id": "ScDQGD5wWyNG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Download subcellular localization dataset from https://zenodo.org/records/10631963\n",
        "!gdown https://drive.google.com/uc?id=1BG91Eu80t546q-9wIDxCaSYpzGh4lPHX\n",
        "\n",
        "data_path = \"balanced.csv\"\n",
        "df = pl.read_csv(data_path)\n",
        "df = df.filter(pl.col(\"sequence\").str.len_chars() <= 1000)\n",
        "\n",
        "train_df = df.filter(pl.col(\"set\") == \"train\")\n",
        "test_df = df.filter(pl.col(\"set\") != \"train\")\n",
        "\n",
        "df.head()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 326
        },
        "id": "shgBcNlAWwFt",
        "outputId": "a4ebaaf0-c97a-43bb-e79e-01b8bca1df6b"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=1BG91Eu80t546q-9wIDxCaSYpzGh4lPHX\n",
            "To: /content/balanced.csv\n",
            "100% 6.25M/6.25M [00:00<00:00, 20.3MB/s]\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "shape: (5, 4)\n",
              "┌─────────────────────────────────┬───────────────┬───────┬────────────┐\n",
              "│ sequence                        ┆ target        ┆ set   ┆ validation │\n",
              "│ ---                             ┆ ---           ┆ ---   ┆ ---        │\n",
              "│ str                             ┆ str           ┆ str   ┆ str        │\n",
              "╞═════════════════════════════════╪═══════════════╪═══════╪════════════╡\n",
              "│ MEVLEEPAPGPGGADAAERRGLRRLLLSGF… ┆ Cell membrane ┆ train ┆ null       │\n",
              "│ MMKTLSSGNCTLNVPAKNSYRMVVLGASRV… ┆ Cell membrane ┆ train ┆ null       │\n",
              "│ MAKRTFSNLETFLIFLLVMMSAITVALLSL… ┆ Cell membrane ┆ train ┆ null       │\n",
              "│ MGNCQAGHNLHLCLAHHPPLVCATLILLLL… ┆ Cell membrane ┆ train ┆ null       │\n",
              "│ MDPSKQGTLNRVENSVYRTAFKLRSVQTLC… ┆ Cell membrane ┆ train ┆ null       │\n",
              "└─────────────────────────────────┴───────────────┴───────┴────────────┘"
            ],
            "text/html": [
              "<div><style>\n",
              ".dataframe > thead > tr,\n",
              ".dataframe > tbody > tr {\n",
              "  text-align: right;\n",
              "  white-space: pre-wrap;\n",
              "}\n",
              "</style>\n",
              "<small>shape: (5, 4)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>sequence</th><th>target</th><th>set</th><th>validation</th></tr><tr><td>str</td><td>str</td><td>str</td><td>str</td></tr></thead><tbody><tr><td>&quot;MEVLEEPAPGPGGADAAERRGLRRLLLSGF…</td><td>&quot;Cell membrane&quot;</td><td>&quot;train&quot;</td><td>null</td></tr><tr><td>&quot;MMKTLSSGNCTLNVPAKNSYRMVVLGASRV…</td><td>&quot;Cell membrane&quot;</td><td>&quot;train&quot;</td><td>null</td></tr><tr><td>&quot;MAKRTFSNLETFLIFLLVMMSAITVALLSL…</td><td>&quot;Cell membrane&quot;</td><td>&quot;train&quot;</td><td>null</td></tr><tr><td>&quot;MGNCQAGHNLHLCLAHHPPLVCATLILLLL…</td><td>&quot;Cell membrane&quot;</td><td>&quot;train&quot;</td><td>null</td></tr><tr><td>&quot;MDPSKQGTLNRVENSVYRTAFKLRSVQTLC…</td><td>&quot;Cell membrane&quot;</td><td>&quot;train&quot;</td><td>null</td></tr></tbody></table></div>"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define some helpers for linear classified training, evaluation, and inspection."
      ],
      "metadata": {
        "id": "THGgSK42Wksr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def train_classifier(seq_acts):\n",
        "    \"\"\"Trains a linear classifier with a predefined validation split.\n",
        "\n",
        "    Uses grid search over regularization strengths to find the best classifier.\n",
        "    \"\"\"\n",
        "    X_train = []\n",
        "    y_train = []\n",
        "    test_fold = []\n",
        "\n",
        "    for row in train_df.iter_rows(named=True):\n",
        "        X_train.append(seq_acts[row['sequence']].cpu().detach().numpy())\n",
        "        y_train.append(row['target'])\n",
        "        if row['validation']:\n",
        "            test_fold.append(0)\n",
        "        else:\n",
        "            test_fold.append(-1)\n",
        "\n",
        "    X_train = np.array(X_train)\n",
        "    y_train = np.array(y_train)\n",
        "\n",
        "    classifier = LogisticRegression(multi_class='ovr')\n",
        "    ps = PredefinedSplit(test_fold=test_fold)\n",
        "    param_grid = {'C': [0.01, 0.1, 1, 10, 100]}\n",
        "    grid_search = GridSearchCV(classifier, param_grid, cv=ps, scoring='accuracy')\n",
        "    grid_search.fit(X_train, y_train)\n",
        "\n",
        "    best_classifier = grid_search.best_estimator_\n",
        "    return best_classifier, grid_search.best_params_['C'], grid_search.best_score_\n",
        "\n",
        "\n",
        "def evaluate_classifier(classifier, seq_acts):\n",
        "    \"\"\"Computes the accuracy of the classifier over a test set\"\"\"\n",
        "    X_test = []\n",
        "    y_test = []\n",
        "    for row in test_df.iter_rows(named=True):\n",
        "        X_test.append(seq_acts[row['sequence']].cpu().detach().numpy())\n",
        "        y_test.append(row['target'])\n",
        "    X_test = np.array(X_test)\n",
        "    y_test = np.array(y_test)\n",
        "\n",
        "    y_pred = classifier.predict(X_test)\n",
        "\n",
        "    accuracy = accuracy_score(y_test, y_pred)\n",
        "    return accuracy\n",
        "\n",
        "\n",
        "def write_sorted_weights(classifier, output_dir):\n",
        "    \"\"\"Writes model weights to CSV files in sorted order.\n",
        "\n",
        "    For multi-class models, writes one CSV per class.\n",
        "    Each CSV has columns \"Index\" and \"Weight\" sorted descending.\n",
        "\n",
        "    Args:\n",
        "        output_dir: Directory to write CSV files to\n",
        "    \"\"\"\n",
        "    os.makedirs(output_dir, exist_ok=True)\n",
        "\n",
        "    for i, class_label in enumerate(classifier.classes_):\n",
        "        class_label = class_label.replace(\"/\", \"_\").lower()\n",
        "        output_file = os.path.join(output_dir, f\"class_{class_label}_weights.csv\")\n",
        "\n",
        "        print(f\"Class: {class_label}\")\n",
        "        class_weights = classifier.coef_[i]\n",
        "        sorted_weights = sorted(enumerate(class_weights), key=lambda x: x[1], reverse=True)\n",
        "        with open(output_file, 'w', newline='') as csvfile:\n",
        "            writer = csv.writer(csvfile)\n",
        "            writer.writerow(['Index', 'Weight'])\n",
        "            for index, weight in sorted_weights:\n",
        "                writer.writerow([index, weight])\n",
        "\n",
        "        print(pd.read_csv(output_file).head())"
      ],
      "metadata": {
        "id": "_U4huDDaWjkT"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# ESM & SAE Inference\n",
        "\n",
        "Compute and store mean-pooled embeddings for both ESM and SAE"
      ],
      "metadata": {
        "id": "KkGwi_p-IL-g"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "seq_esm_acts = {}\n",
        "seq_sae_acts = {}\n",
        "\n",
        "for seq_idx, row in tqdm(enumerate(df.iter_rows(named=True)), total=len(df)):\n",
        "    seq = row['sequence']\n",
        "    esm_layer_acts = get_layer_activations(\n",
        "        tokenizer=tokenizer, plm=esm_model, seqs=[seq], layer=LAYER\n",
        "    )[0][1:-1] # Trim off BoS and EoS tokens\n",
        "\n",
        "    seq_esm_acts[seq] = torch.mean(esm_layer_acts, axis=0)\n",
        "\n",
        "    sae_acts = sae_model.get_acts(esm_layer_acts)\n",
        "    seq_sae_acts[seq] = torch.mean(sae_acts, axis=0)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WEKYZ1bpGaM6",
        "outputId": "9c9f6d46-1706-4632-e093-ec037d363cc9"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 10461/10461 [17:57<00:00,  9.71it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Linear classifiers\n",
        "\n",
        "- Train classifier on ESM embedding (for accuracy comparison)\n",
        "- Train classifier on SAE embedding\n",
        "- Inspect and write the weights of the SAE classifier"
      ],
      "metadata": {
        "id": "7_NvR1ZklRNe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with warnings.catch_warnings():\n",
        "    warnings.simplefilter(\"ignore\")\n",
        "\n",
        "    print(\"Training ESM classifier\")\n",
        "    esm_classifier, _, _ = train_classifier(seq_esm_acts)\n",
        "    esm_accuracy = evaluate_classifier(esm_classifier, seq_esm_acts)\n",
        "    print(f\"ESM accuracy: {esm_accuracy}\")\n",
        "\n",
        "    print(\"Training SAE classifier\")\n",
        "    sae_classifier, _, _ = train_classifier(seq_sae_acts)\n",
        "    sae_accuracy = evaluate_classifier(sae_classifier, seq_sae_acts)\n",
        "    print(f\"SAE accuracy: {sae_accuracy}\")\n",
        "    write_sorted_weights(sae_classifier, 'weights')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HkDwM75sXmGA",
        "outputId": "bf2b2a2b-3a00-437d-8933-1cde91db587b"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training ESM classifier\n",
            "ESM accuracy: 0.6194690265486725\n",
            "Training SAE classifier\n",
            "SAE accuracy: 0.640117994100295\n",
            "Class: cell membrane\n",
            "   Index    Weight\n",
            "0   3815  2.481609\n",
            "1    550  1.947456\n",
            "2   2818  1.482108\n",
            "3   2000  1.283794\n",
            "4   3154  1.073169\n",
            "Class: cytoplasm\n",
            "   Index    Weight\n",
            "0   3274  1.052095\n",
            "1   1849  0.893690\n",
            "2   2966  0.867804\n",
            "3   2253  0.771584\n",
            "4   2554  0.768400\n",
            "Class: endoplasmic reticulum\n",
            "   Index    Weight\n",
            "0   1053  1.239194\n",
            "1   1298  1.229100\n",
            "2    848  1.148793\n",
            "3   1055  0.984984\n",
            "4    565  0.968559\n",
            "Class: extracellular\n",
            "   Index    Weight\n",
            "0   1541  1.293407\n",
            "1   1470  1.264098\n",
            "2   1555  1.151713\n",
            "3   2111  1.135436\n",
            "4   2472  1.035609\n",
            "Class: golgi apparatus\n",
            "   Index    Weight\n",
            "0   4086  1.601550\n",
            "1   1635  1.204602\n",
            "2    880  1.179289\n",
            "3   3729  1.169152\n",
            "4     75  1.096873\n",
            "Class: lysosome_vacuole\n",
            "   Index    Weight\n",
            "0   2354  1.203674\n",
            "1   3592  1.142584\n",
            "2   3021  1.105611\n",
            "3   1683  1.094772\n",
            "4   2867  0.992007\n",
            "Class: mitochondrion\n",
            "   Index    Weight\n",
            "0   3277  2.362949\n",
            "1   1948  1.114703\n",
            "2   1690  1.081239\n",
            "3   2745  1.058991\n",
            "4    326  1.025347\n",
            "Class: nucleus\n",
            "   Index    Weight\n",
            "0   2375  2.823197\n",
            "1   2267  1.975491\n",
            "2   3396  1.618289\n",
            "3   3764  1.310228\n",
            "4    383  1.121367\n",
            "Class: peroxisome\n",
            "   Index    Weight\n",
            "0    655  1.142600\n",
            "1    755  0.931766\n",
            "2   1212  0.918675\n",
            "3   1801  0.917598\n",
            "4   1053  0.822834\n",
            "Class: plastid\n",
            "   Index    Weight\n",
            "0    880  2.123226\n",
            "1   2601  1.570855\n",
            "2   3389  1.446348\n",
            "3   2024  1.272865\n",
            "4    179  1.213332\n"
          ]
        }
      ]
    }
  ]
}