TraitGym.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "354b6ae4307840b5a73c4cb10f976d7d": {
          "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_079b1e09f86c40c1936b276c93e14271",
              "IPY_MODEL_6059d28fa74845b9a7ec65878187f978",
              "IPY_MODEL_93ac776b9acf4b1c81566c16d2408e8c"
            ],
            "layout": "IPY_MODEL_0b42e95b4d98494893bce8b38ec46ada"
          }
        },
        "079b1e09f86c40c1936b276c93e14271": {
          "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_9f9af5e0b00f4d419927290074e398fe",
            "placeholder": "​",
            "style": "IPY_MODEL_4a63ea2b9cae4e979426205b1312f39e",
            "value": "README.md: 100%"
          }
        },
        "6059d28fa74845b9a7ec65878187f978": {
          "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_e47ef5863d5847cf8667106d9fc95057",
            "max": 563,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_d55c9de46b6749deb738cc1b59904ff8",
            "value": 563
          }
        },
        "93ac776b9acf4b1c81566c16d2408e8c": {
          "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_6e19e94413f946daa10f3f8f2ce8d895",
            "placeholder": "​",
            "style": "IPY_MODEL_7afe78cb30a94ce2ac4a2e6587e13764",
            "value": " 563/563 [00:00&lt;00:00, 18.2kB/s]"
          }
        },
        "0b42e95b4d98494893bce8b38ec46ada": {
          "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
          }
        },
        "9f9af5e0b00f4d419927290074e398fe": {
          "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
          }
        },
        "4a63ea2b9cae4e979426205b1312f39e": {
          "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": ""
          }
        },
        "e47ef5863d5847cf8667106d9fc95057": {
          "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
          }
        },
        "d55c9de46b6749deb738cc1b59904ff8": {
          "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": ""
          }
        },
        "6e19e94413f946daa10f3f8f2ce8d895": {
          "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
          }
        },
        "7afe78cb30a94ce2ac4a2e6587e13764": {
          "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": ""
          }
        },
        "5690f364ff504f6bb6d89aa074794e10": {
          "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_cc50a4e9a2404208bdc8f3fd5f9ba2eb",
              "IPY_MODEL_146ad95d109f4203828a42688c582711",
              "IPY_MODEL_53e2168f0f494113a28d251fa1c61051"
            ],
            "layout": "IPY_MODEL_75281a09c57e4eaeb230a9737a5dba7a"
          }
        },
        "cc50a4e9a2404208bdc8f3fd5f9ba2eb": {
          "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_688266106cb54dffbe066bb3d8b33b99",
            "placeholder": "​",
            "style": "IPY_MODEL_9ebbe4ddd37f4f2ebdd7d7ea3b5c3339",
            "value": "test.parquet: 100%"
          }
        },
        "146ad95d109f4203828a42688c582711": {
          "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_07a56cf6486b4f338f967e249a89a470",
            "max": 50205,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_aa033b8bd0bc461490fa597ecb1f2b29",
            "value": 50205
          }
        },
        "53e2168f0f494113a28d251fa1c61051": {
          "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_d0f4f73e87104e8686a1f410c7aad9a4",
            "placeholder": "​",
            "style": "IPY_MODEL_0e5176ab3e054e3bba1eacb239931176",
            "value": " 50.2k/50.2k [00:00&lt;00:00, 1.26MB/s]"
          }
        },
        "75281a09c57e4eaeb230a9737a5dba7a": {
          "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
          }
        },
        "688266106cb54dffbe066bb3d8b33b99": {
          "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
          }
        },
        "9ebbe4ddd37f4f2ebdd7d7ea3b5c3339": {
          "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": ""
          }
        },
        "07a56cf6486b4f338f967e249a89a470": {
          "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
          }
        },
        "aa033b8bd0bc461490fa597ecb1f2b29": {
          "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": ""
          }
        },
        "d0f4f73e87104e8686a1f410c7aad9a4": {
          "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
          }
        },
        "0e5176ab3e054e3bba1eacb239931176": {
          "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": ""
          }
        },
        "ef77fe79ff3a4bd6b9815d12c1b1ef42": {
          "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_7b5c67519b744a769d43f16e0296c8f8",
              "IPY_MODEL_a68e1e3900794fbb8969e54714ec4a83",
              "IPY_MODEL_90a329c6bb4e43bbb6efb108a5fa9805"
            ],
            "layout": "IPY_MODEL_61d106a61b474176a0d565dea76d19f4"
          }
        },
        "7b5c67519b744a769d43f16e0296c8f8": {
          "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_4073c99a09424482be157783b5142567",
            "placeholder": "​",
            "style": "IPY_MODEL_3f044b6801c444d992abffbf42b7912d",
            "value": "Generating test split: 100%"
          }
        },
        "a68e1e3900794fbb8969e54714ec4a83": {
          "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_381192f45b1a4bdea86af2b810bb9c45",
            "max": 3380,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_7357aaf41476452aab915e52ddf6bf7b",
            "value": 3380
          }
        },
        "90a329c6bb4e43bbb6efb108a5fa9805": {
          "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_fde154d03164486fb4160cf1dcc516a8",
            "placeholder": "​",
            "style": "IPY_MODEL_443f5b80dd534c88aaf2a17c3e1a893b",
            "value": " 3380/3380 [00:00&lt;00:00, 32442.93 examples/s]"
          }
        },
        "61d106a61b474176a0d565dea76d19f4": {
          "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
          }
        },
        "4073c99a09424482be157783b5142567": {
          "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
          }
        },
        "3f044b6801c444d992abffbf42b7912d": {
          "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": ""
          }
        },
        "381192f45b1a4bdea86af2b810bb9c45": {
          "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
          }
        },
        "7357aaf41476452aab915e52ddf6bf7b": {
          "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": ""
          }
        },
        "fde154d03164486fb4160cf1dcc516a8": {
          "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
          }
        },
        "443f5b80dd534c88aaf2a17c3e1a893b": {
          "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": ""
          }
        },
        "0ade8e14a15944b695d390db199b9dd9": {
          "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_c53c155008af4c4e9ae6c21598e3b0f1",
              "IPY_MODEL_ff9dd793206542ffa9f31a48b7e0314d",
              "IPY_MODEL_7146c456cbaa41498f9122fd0d29b1d2"
            ],
            "layout": "IPY_MODEL_dcfa4d369495418295baf7fbe9568dce"
          }
        },
        "c53c155008af4c4e9ae6c21598e3b0f1": {
          "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_151e2df6bea14a43ad2705f678c99195",
            "placeholder": "​",
            "style": "IPY_MODEL_6dbf91bea24840fea03b861144a3946e",
            "value": "tokenizer_config.json: 100%"
          }
        },
        "ff9dd793206542ffa9f31a48b7e0314d": {
          "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_d1431f375a2e46838a5970cc80f61200",
            "max": 784,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_dcfa667e270c40d79dd93309da9f77fb",
            "value": 784
          }
        },
        "7146c456cbaa41498f9122fd0d29b1d2": {
          "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_32418a73eb3b41f3be335fb89d2b877c",
            "placeholder": "​",
            "style": "IPY_MODEL_fa0e9d33904e4900984fafcdcdf6e6e7",
            "value": " 784/784 [00:00&lt;00:00, 26.8kB/s]"
          }
        },
        "dcfa4d369495418295baf7fbe9568dce": {
          "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
          }
        },
        "151e2df6bea14a43ad2705f678c99195": {
          "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
          }
        },
        "6dbf91bea24840fea03b861144a3946e": {
          "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": ""
          }
        },
        "d1431f375a2e46838a5970cc80f61200": {
          "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
          }
        },
        "dcfa667e270c40d79dd93309da9f77fb": {
          "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": ""
          }
        },
        "32418a73eb3b41f3be335fb89d2b877c": {
          "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
          }
        },
        "fa0e9d33904e4900984fafcdcdf6e6e7": {
          "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": ""
          }
        },
        "3737b2439dbd4e568d366a9318c260a4": {
          "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_c02a2073a45d4b93bc30bc5ed6be8805",
              "IPY_MODEL_c5874ac66fd94a0b9cf43762c741118f",
              "IPY_MODEL_51b43b7ce3d04b548361ddcb011ac227"
            ],
            "layout": "IPY_MODEL_134c2b6d1ac94d1ebd41517d5f16d9b6"
          }
        },
        "c02a2073a45d4b93bc30bc5ed6be8805": {
          "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_7faa6f4bb40c4f49b21b8cdd999105be",
            "placeholder": "​",
            "style": "IPY_MODEL_3a3bf3efb5334a41a7d1412a746c2e1c",
            "value": "tokenizer.json: 100%"
          }
        },
        "c5874ac66fd94a0b9cf43762c741118f": {
          "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_525edb5bb9194355b730557663f00a43",
            "max": 1145,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_8d88a2fedae6498bb3a5f57457333bd9",
            "value": 1145
          }
        },
        "51b43b7ce3d04b548361ddcb011ac227": {
          "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_c08096b9eb584ca4870d13f3e5824a9b",
            "placeholder": "​",
            "style": "IPY_MODEL_74fc154af2304c2fb0ba48d14ecc85c7",
            "value": " 1.15k/1.15k [00:00&lt;00:00, 37.6kB/s]"
          }
        },
        "134c2b6d1ac94d1ebd41517d5f16d9b6": {
          "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
          }
        },
        "7faa6f4bb40c4f49b21b8cdd999105be": {
          "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
          }
        },
        "3a3bf3efb5334a41a7d1412a746c2e1c": {
          "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": ""
          }
        },
        "525edb5bb9194355b730557663f00a43": {
          "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
          }
        },
        "8d88a2fedae6498bb3a5f57457333bd9": {
          "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": ""
          }
        },
        "c08096b9eb584ca4870d13f3e5824a9b": {
          "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
          }
        },
        "74fc154af2304c2fb0ba48d14ecc85c7": {
          "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": ""
          }
        },
        "565dd0e46c434aa7a73d6bb908bd0680": {
          "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_2b3fb166766c4f79a0efc8626c93b246",
              "IPY_MODEL_baf99987c9b54fe881574d800ee09f14",
              "IPY_MODEL_c8f25334ab334b5fac6797ab7d309f38"
            ],
            "layout": "IPY_MODEL_a86710668b3540d0aa8bb4a36dce8ceb"
          }
        },
        "2b3fb166766c4f79a0efc8626c93b246": {
          "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_18028875cdaa4988bf19419280e0279a",
            "placeholder": "​",
            "style": "IPY_MODEL_e22fdcb3b13c44638b6ac147cd5a4f8c",
            "value": "special_tokens_map.json: 100%"
          }
        },
        "baf99987c9b54fe881574d800ee09f14": {
          "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_d2d8169da2df4bf78acd960b4e3676b0",
            "max": 419,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_0c3e11bcd3ac4f628676209f92949f14",
            "value": 419
          }
        },
        "c8f25334ab334b5fac6797ab7d309f38": {
          "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_cc29c662fd6f4be0b49889c75655414d",
            "placeholder": "​",
            "style": "IPY_MODEL_acb44fc0488843469bb1f83a0c856331",
            "value": " 419/419 [00:00&lt;00:00, 8.70kB/s]"
          }
        },
        "a86710668b3540d0aa8bb4a36dce8ceb": {
          "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
          }
        },
        "18028875cdaa4988bf19419280e0279a": {
          "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
          }
        },
        "e22fdcb3b13c44638b6ac147cd5a4f8c": {
          "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": ""
          }
        },
        "d2d8169da2df4bf78acd960b4e3676b0": {
          "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
          }
        },
        "0c3e11bcd3ac4f628676209f92949f14": {
          "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": ""
          }
        },
        "cc29c662fd6f4be0b49889c75655414d": {
          "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
          }
        },
        "acb44fc0488843469bb1f83a0c856331": {
          "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": ""
          }
        },
        "1656307b55e241078531ac0983ee96fb": {
          "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_b2956c54417446cf9346ad51e3808075",
              "IPY_MODEL_df36a32b3b42485db9a27cec3cd21e31",
              "IPY_MODEL_77988a600cb64c7a9773f84c1725a21b"
            ],
            "layout": "IPY_MODEL_93c39b9dab7f4fa5ac7bc5f2c03e9175"
          }
        },
        "b2956c54417446cf9346ad51e3808075": {
          "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_10b81b4b60a34853a9663e78bb2e1abe",
            "placeholder": "​",
            "style": "IPY_MODEL_7b4126a076cb4dcbba25b04d19334b0b",
            "value": "config.json: 100%"
          }
        },
        "df36a32b3b42485db9a27cec3cd21e31": {
          "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_eae24a291baa4c04a1bc8d1966a3ccae",
            "max": 1165,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_c649b21f7c1642ebb880d3a62b8281ed",
            "value": 1165
          }
        },
        "77988a600cb64c7a9773f84c1725a21b": {
          "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_46f44a143d7f4536a93db59e033ce665",
            "placeholder": "​",
            "style": "IPY_MODEL_ed4930de91c440c0a1a7239a4ed06e5d",
            "value": " 1.17k/1.17k [00:00&lt;00:00, 21.6kB/s]"
          }
        },
        "93c39b9dab7f4fa5ac7bc5f2c03e9175": {
          "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
          }
        },
        "10b81b4b60a34853a9663e78bb2e1abe": {
          "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
          }
        },
        "7b4126a076cb4dcbba25b04d19334b0b": {
          "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": ""
          }
        },
        "eae24a291baa4c04a1bc8d1966a3ccae": {
          "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
          }
        },
        "c649b21f7c1642ebb880d3a62b8281ed": {
          "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": ""
          }
        },
        "46f44a143d7f4536a93db59e033ce665": {
          "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
          }
        },
        "ed4930de91c440c0a1a7239a4ed06e5d": {
          "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": ""
          }
        },
        "a31b47baeb624307b0bb270d81949504": {
          "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_7459011748494d4a9939220746108716",
              "IPY_MODEL_c923d77c086d4abb93fabc6f8d20b5b7",
              "IPY_MODEL_c8922ee8c0224221b2673f7618701da4"
            ],
            "layout": "IPY_MODEL_abd5f43ea691418898820e68d36fdd8f"
          }
        },
        "7459011748494d4a9939220746108716": {
          "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_2a38d3b1f0df48fb96b2e9256a0c564b",
            "placeholder": "​",
            "style": "IPY_MODEL_3a408803721449a1aa63a65cd51f9a1d",
            "value": "model.safetensors: 100%"
          }
        },
        "c923d77c086d4abb93fabc6f8d20b5b7": {
          "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_9a39944759c64b57a27f5c4e5ef935d7",
            "max": 608802168,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_816ad4562cbc4ff1934da1f5cd3b69c7",
            "value": 608802168
          }
        },
        "c8922ee8c0224221b2673f7618701da4": {
          "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_197fe37ca48b4dfeb674e8a0688102d0",
            "placeholder": "​",
            "style": "IPY_MODEL_0faeff78cc7446bab7c32806805cce5a",
            "value": " 609M/609M [00:14&lt;00:00, 42.5MB/s]"
          }
        },
        "abd5f43ea691418898820e68d36fdd8f": {
          "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
          }
        },
        "2a38d3b1f0df48fb96b2e9256a0c564b": {
          "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
          }
        },
        "3a408803721449a1aa63a65cd51f9a1d": {
          "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": ""
          }
        },
        "9a39944759c64b57a27f5c4e5ef935d7": {
          "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
          }
        },
        "816ad4562cbc4ff1934da1f5cd3b69c7": {
          "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": ""
          }
        },
        "197fe37ca48b4dfeb674e8a0688102d0": {
          "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
          }
        },
        "0faeff78cc7446bab7c32806805cce5a": {
          "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": [
        "# TraitGym [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/songlab-cal/TraitGym/blob/main/TraitGym.ipynb)\n",
        "In this example we will load the Mendelian traits dataset and run variant effect prediction based on euclidean distance of GPN-Animal-Promoter embeddings of the reference and alternate sequences."
      ],
      "metadata": {
        "id": "muWargPyv5Zi"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setup"
      ],
      "metadata": {
        "id": "5V9APVrtwLlr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q pyfaidx s3fs git+https://github.com/songlab-cal/gpn.git\n",
        "!pip install -q -U transformers datasets"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "II5CCc_mxEXj",
        "outputId": "847f69d3-de6d-48b2-81b1-5d2ba02751d5"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m183.9/183.9 kB\u001b[0m \u001b[31m8.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m77.7/77.7 kB\u001b[0m \u001b[31m6.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m151.5/151.5 kB\u001b[0m \u001b[31m12.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.3/3.3 MB\u001b[0m \u001b[31m67.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m519.3/519.3 kB\u001b[0m \u001b[31m27.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m185.8/185.8 kB\u001b[0m \u001b[31m11.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m181.4/181.4 kB\u001b[0m \u001b[31m13.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m13.3/13.3 MB\u001b[0m \u001b[31m75.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m115.3/115.3 kB\u001b[0m \u001b[31m9.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m8.9/8.9 MB\u001b[0m \u001b[31m73.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m135.4/135.4 kB\u001b[0m \u001b[31m10.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m194.8/194.8 kB\u001b[0m \u001b[31m13.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m53.7/53.7 kB\u001b[0m \u001b[31m4.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Building wheel for gpn (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for pandarallel (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "gcsfs 2024.10.0 requires fsspec==2024.10.0, but you have fsspec 2024.12.0 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m44.4/44.4 kB\u001b[0m \u001b[31m2.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m9.7/9.7 MB\u001b[0m \u001b[31m34.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m480.6/480.6 kB\u001b[0m \u001b[31m16.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m179.3/179.3 kB\u001b[0m \u001b[31m14.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "s3fs 2024.12.0 requires fsspec==2024.12.0.*, but you have fsspec 2024.9.0 which is incompatible.\n",
            "gcsfs 2024.10.0 requires fsspec==2024.10.0, but you have fsspec 2024.9.0 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0m"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from Bio.Seq import Seq\n",
        "from datasets import load_dataset\n",
        "import fsspec\n",
        "import gpn.model  # to register AutoModel\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import polars as pl\n",
        "from pyfaidx import Fasta\n",
        "import seaborn as sns\n",
        "from sklearn.metrics import average_precision_score\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "from transformers import AutoTokenizer, AutoModel, TrainingArguments, Trainer\n",
        "import tempfile"
      ],
      "metadata": {
        "id": "yFgUHIlJxs2P"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset_path = \"songlab/TraitGym\"\n",
        "dataset_config = \"mendelian_traits\"\n",
        "model_path = \"songlab/gpn-animal-promoter\"\n",
        "# in the paper we average with the predictions with the reverse complement\n",
        "# however for a quick evaluation you can omit this\n",
        "average_rc = False"
      ],
      "metadata": {
        "id": "G9JKGeFXOB5n"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load dataset"
      ],
      "metadata": {
        "id": "39wqWU2nw0th"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 287,
          "referenced_widgets": [
            "354b6ae4307840b5a73c4cb10f976d7d",
            "079b1e09f86c40c1936b276c93e14271",
            "6059d28fa74845b9a7ec65878187f978",
            "93ac776b9acf4b1c81566c16d2408e8c",
            "0b42e95b4d98494893bce8b38ec46ada",
            "9f9af5e0b00f4d419927290074e398fe",
            "4a63ea2b9cae4e979426205b1312f39e",
            "e47ef5863d5847cf8667106d9fc95057",
            "d55c9de46b6749deb738cc1b59904ff8",
            "6e19e94413f946daa10f3f8f2ce8d895",
            "7afe78cb30a94ce2ac4a2e6587e13764",
            "5690f364ff504f6bb6d89aa074794e10",
            "cc50a4e9a2404208bdc8f3fd5f9ba2eb",
            "146ad95d109f4203828a42688c582711",
            "53e2168f0f494113a28d251fa1c61051",
            "75281a09c57e4eaeb230a9737a5dba7a",
            "688266106cb54dffbe066bb3d8b33b99",
            "9ebbe4ddd37f4f2ebdd7d7ea3b5c3339",
            "07a56cf6486b4f338f967e249a89a470",
            "aa033b8bd0bc461490fa597ecb1f2b29",
            "d0f4f73e87104e8686a1f410c7aad9a4",
            "0e5176ab3e054e3bba1eacb239931176",
            "ef77fe79ff3a4bd6b9815d12c1b1ef42",
            "7b5c67519b744a769d43f16e0296c8f8",
            "a68e1e3900794fbb8969e54714ec4a83",
            "90a329c6bb4e43bbb6efb108a5fa9805",
            "61d106a61b474176a0d565dea76d19f4",
            "4073c99a09424482be157783b5142567",
            "3f044b6801c444d992abffbf42b7912d",
            "381192f45b1a4bdea86af2b810bb9c45",
            "7357aaf41476452aab915e52ddf6bf7b",
            "fde154d03164486fb4160cf1dcc516a8",
            "443f5b80dd534c88aaf2a17c3e1a893b"
          ]
        },
        "id": "6p3Vk6zRvPhQ",
        "outputId": "a39b7fb8-70ea-4b09-f97d-24132efcff78"
      },
      "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": [
              "README.md:   0%|          | 0.00/563 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "354b6ae4307840b5a73c4cb10f976d7d"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "test.parquet:   0%|          | 0.00/50.2k [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "5690f364ff504f6bb6d89aa074794e10"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Generating test split:   0%|          | 0/3380 [00:00<?, ? examples/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "ef77fe79ff3a4bd6b9815d12c1b1ef42"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Dataset({\n",
              "    features: ['chrom', 'pos', 'ref', 'alt', 'OMIM', 'consequence', 'label', 'tss_dist', 'match_group'],\n",
              "    num_rows: 3380\n",
              "})"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ],
      "source": [
        "dataset = load_dataset(dataset_path, dataset_config, split=\"test\")\n",
        "dataset"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dataset[0]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QW1mpMscyOUD",
        "outputId": "965889ee-c94e-4e04-e298-01fdc83c4ee5"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'chrom': '1',\n",
              " 'pos': 1425822,\n",
              " 'ref': 'C',\n",
              " 'alt': 'G',\n",
              " 'OMIM': None,\n",
              " 'consequence': 'PLS',\n",
              " 'label': False,\n",
              " 'tss_dist': 48,\n",
              " 'match_group': 'PLS_4'}"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# if you just want a dataframe you can load it directly:\n",
        "# pd.read_parquet(\"hf://datasets/songlab/TraitGym/mendelian_traits_matched_9/test.parquet\")\n",
        "V = dataset.to_pandas()\n",
        "V"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 424
        },
        "id": "1DBiScBWyVis",
        "outputId": "56e39155-d354-4d34-e456-4ead360311e7"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "     chrom        pos ref alt  OMIM            consequence  label  tss_dist  \\\n",
              "0        1    1425822   C   G  None                    PLS  False        48   \n",
              "1        1    1615869   C   T  None                    PLS  False        35   \n",
              "2        1    1659060   G   A  None                    PLS  False        47   \n",
              "3        1    1659114   A   G  None                    PLS  False       101   \n",
              "4        1    2050958   T   C  None    5_prime_UTR_variant  False       149   \n",
              "...    ...        ...  ..  ..   ...                    ...    ...       ...   \n",
              "3375     X  155613005   C   T  None                    PLS  False        52   \n",
              "3376     X  155719093   C   A  None    5_prime_UTR_variant  False         4   \n",
              "3377     X  155881342   A   C  None                    PLS  False         2   \n",
              "3378     X  155881414   C   T  None    5_prime_UTR_variant  False        35   \n",
              "3379     X  156020793   A   G  None  upstream_gene_variant  False       167   \n",
              "\n",
              "                  match_group  \n",
              "0                       PLS_4  \n",
              "1                       PLS_0  \n",
              "2                       PLS_4  \n",
              "3                       PLS_5  \n",
              "4       5_prime_UTR_variant_7  \n",
              "...                       ...  \n",
              "3375                   PLS_52  \n",
              "3376  5_prime_UTR_variant_101  \n",
              "3377                   PLS_57  \n",
              "3378  5_prime_UTR_variant_110  \n",
              "3379  upstream_gene_variant_9  \n",
              "\n",
              "[3380 rows x 9 columns]"
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-a0a98cc9-1dcc-4728-bc4f-c007019eb67d\" class=\"colab-df-container\">\n",
              "    <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>chrom</th>\n",
              "      <th>pos</th>\n",
              "      <th>ref</th>\n",
              "      <th>alt</th>\n",
              "      <th>OMIM</th>\n",
              "      <th>consequence</th>\n",
              "      <th>label</th>\n",
              "      <th>tss_dist</th>\n",
              "      <th>match_group</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>1</td>\n",
              "      <td>1425822</td>\n",
              "      <td>C</td>\n",
              "      <td>G</td>\n",
              "      <td>None</td>\n",
              "      <td>PLS</td>\n",
              "      <td>False</td>\n",
              "      <td>48</td>\n",
              "      <td>PLS_4</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>1</td>\n",
              "      <td>1615869</td>\n",
              "      <td>C</td>\n",
              "      <td>T</td>\n",
              "      <td>None</td>\n",
              "      <td>PLS</td>\n",
              "      <td>False</td>\n",
              "      <td>35</td>\n",
              "      <td>PLS_0</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>1</td>\n",
              "      <td>1659060</td>\n",
              "      <td>G</td>\n",
              "      <td>A</td>\n",
              "      <td>None</td>\n",
              "      <td>PLS</td>\n",
              "      <td>False</td>\n",
              "      <td>47</td>\n",
              "      <td>PLS_4</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>1</td>\n",
              "      <td>1659114</td>\n",
              "      <td>A</td>\n",
              "      <td>G</td>\n",
              "      <td>None</td>\n",
              "      <td>PLS</td>\n",
              "      <td>False</td>\n",
              "      <td>101</td>\n",
              "      <td>PLS_5</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>1</td>\n",
              "      <td>2050958</td>\n",
              "      <td>T</td>\n",
              "      <td>C</td>\n",
              "      <td>None</td>\n",
              "      <td>5_prime_UTR_variant</td>\n",
              "      <td>False</td>\n",
              "      <td>149</td>\n",
              "      <td>5_prime_UTR_variant_7</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>...</th>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3375</th>\n",
              "      <td>X</td>\n",
              "      <td>155613005</td>\n",
              "      <td>C</td>\n",
              "      <td>T</td>\n",
              "      <td>None</td>\n",
              "      <td>PLS</td>\n",
              "      <td>False</td>\n",
              "      <td>52</td>\n",
              "      <td>PLS_52</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3376</th>\n",
              "      <td>X</td>\n",
              "      <td>155719093</td>\n",
              "      <td>C</td>\n",
              "      <td>A</td>\n",
              "      <td>None</td>\n",
              "      <td>5_prime_UTR_variant</td>\n",
              "      <td>False</td>\n",
              "      <td>4</td>\n",
              "      <td>5_prime_UTR_variant_101</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3377</th>\n",
              "      <td>X</td>\n",
              "      <td>155881342</td>\n",
              "      <td>A</td>\n",
              "      <td>C</td>\n",
              "      <td>None</td>\n",
              "      <td>PLS</td>\n",
              "      <td>False</td>\n",
              "      <td>2</td>\n",
              "      <td>PLS_57</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3378</th>\n",
              "      <td>X</td>\n",
              "      <td>155881414</td>\n",
              "      <td>C</td>\n",
              "      <td>T</td>\n",
              "      <td>None</td>\n",
              "      <td>5_prime_UTR_variant</td>\n",
              "      <td>False</td>\n",
              "      <td>35</td>\n",
              "      <td>5_prime_UTR_variant_110</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3379</th>\n",
              "      <td>X</td>\n",
              "      <td>156020793</td>\n",
              "      <td>A</td>\n",
              "      <td>G</td>\n",
              "      <td>None</td>\n",
              "      <td>upstream_gene_variant</td>\n",
              "      <td>False</td>\n",
              "      <td>167</td>\n",
              "      <td>upstream_gene_variant_9</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "<p>3380 rows × 9 columns</p>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-a0a98cc9-1dcc-4728-bc4f-c007019eb67d')\"\n",
              "            title=\"Convert this dataframe to an interactive table.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "    <script>\n",
              "      const buttonEl =\n",
              "        document.querySelector('#df-a0a98cc9-1dcc-4728-bc4f-c007019eb67d button.colab-df-convert');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      async function convertToInteractive(key) {\n",
              "        const element = document.querySelector('#df-a0a98cc9-1dcc-4728-bc4f-c007019eb67d');\n",
              "        const dataTable =\n",
              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                    [key], {});\n",
              "        if (!dataTable) return;\n",
              "\n",
              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "          + ' to learn more about interactive tables.';\n",
              "        element.innerHTML = '';\n",
              "        dataTable['output_type'] = 'display_data';\n",
              "        await google.colab.output.renderOutput(dataTable, element);\n",
              "        const docLink = document.createElement('div');\n",
              "        docLink.innerHTML = docLinkHtml;\n",
              "        element.appendChild(docLink);\n",
              "      }\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "\n",
              "<div id=\"df-66457b7e-eb07-42d0-b572-724f6511440f\">\n",
              "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-66457b7e-eb07-42d0-b572-724f6511440f')\"\n",
              "            title=\"Suggest charts\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "     width=\"24px\">\n",
              "    <g>\n",
              "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
              "    </g>\n",
              "</svg>\n",
              "  </button>\n",
              "\n",
              "<style>\n",
              "  .colab-df-quickchart {\n",
              "      --bg-color: #E8F0FE;\n",
              "      --fill-color: #1967D2;\n",
              "      --hover-bg-color: #E2EBFA;\n",
              "      --hover-fill-color: #174EA6;\n",
              "      --disabled-fill-color: #AAA;\n",
              "      --disabled-bg-color: #DDD;\n",
              "  }\n",
              "\n",
              "  [theme=dark] .colab-df-quickchart {\n",
              "      --bg-color: #3B4455;\n",
              "      --fill-color: #D2E3FC;\n",
              "      --hover-bg-color: #434B5C;\n",
              "      --hover-fill-color: #FFFFFF;\n",
              "      --disabled-bg-color: #3B4455;\n",
              "      --disabled-fill-color: #666;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart {\n",
              "    background-color: var(--bg-color);\n",
              "    border: none;\n",
              "    border-radius: 50%;\n",
              "    cursor: pointer;\n",
              "    display: none;\n",
              "    fill: var(--fill-color);\n",
              "    height: 32px;\n",
              "    padding: 0;\n",
              "    width: 32px;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart:hover {\n",
              "    background-color: var(--hover-bg-color);\n",
              "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "    fill: var(--button-hover-fill-color);\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart-complete:disabled,\n",
              "  .colab-df-quickchart-complete:disabled:hover {\n",
              "    background-color: var(--disabled-bg-color);\n",
              "    fill: var(--disabled-fill-color);\n",
              "    box-shadow: none;\n",
              "  }\n",
              "\n",
              "  .colab-df-spinner {\n",
              "    border: 2px solid var(--fill-color);\n",
              "    border-color: transparent;\n",
              "    border-bottom-color: var(--fill-color);\n",
              "    animation:\n",
              "      spin 1s steps(1) infinite;\n",
              "  }\n",
              "\n",
              "  @keyframes spin {\n",
              "    0% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "      border-left-color: var(--fill-color);\n",
              "    }\n",
              "    20% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    30% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    40% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    60% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    80% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "    90% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "  }\n",
              "</style>\n",
              "\n",
              "  <script>\n",
              "    async function quickchart(key) {\n",
              "      const quickchartButtonEl =\n",
              "        document.querySelector('#' + key + ' button');\n",
              "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
              "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
              "      try {\n",
              "        const charts = await google.colab.kernel.invokeFunction(\n",
              "            'suggestCharts', [key], {});\n",
              "      } catch (error) {\n",
              "        console.error('Error during call to suggestCharts:', error);\n",
              "      }\n",
              "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
              "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
              "    }\n",
              "    (() => {\n",
              "      let quickchartButtonEl =\n",
              "        document.querySelector('#df-66457b7e-eb07-42d0-b572-724f6511440f button');\n",
              "      quickchartButtonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "    })();\n",
              "  </script>\n",
              "</div>\n",
              "\n",
              "  <div id=\"id_5b0fcc4e-3a0d-4ee0-9cc4-351759ccc559\">\n",
              "    <style>\n",
              "      .colab-df-generate {\n",
              "        background-color: #E8F0FE;\n",
              "        border: none;\n",
              "        border-radius: 50%;\n",
              "        cursor: pointer;\n",
              "        display: none;\n",
              "        fill: #1967D2;\n",
              "        height: 32px;\n",
              "        padding: 0 0 0 0;\n",
              "        width: 32px;\n",
              "      }\n",
              "\n",
              "      .colab-df-generate:hover {\n",
              "        background-color: #E2EBFA;\n",
              "        box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "        fill: #174EA6;\n",
              "      }\n",
              "\n",
              "      [theme=dark] .colab-df-generate {\n",
              "        background-color: #3B4455;\n",
              "        fill: #D2E3FC;\n",
              "      }\n",
              "\n",
              "      [theme=dark] .colab-df-generate:hover {\n",
              "        background-color: #434B5C;\n",
              "        box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "        filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "        fill: #FFFFFF;\n",
              "      }\n",
              "    </style>\n",
              "    <button class=\"colab-df-generate\" onclick=\"generateWithVariable('V')\"\n",
              "            title=\"Generate code using this dataframe.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "       width=\"24px\">\n",
              "    <path d=\"M7,19H8.4L18.45,9,17,7.55,7,17.6ZM5,21V16.75L18.45,3.32a2,2,0,0,1,2.83,0l1.4,1.43a1.91,1.91,0,0,1,.58,1.4,1.91,1.91,0,0,1-.58,1.4L9.25,21ZM18.45,9,17,7.55Zm-12,3A5.31,5.31,0,0,0,4.9,8.1,5.31,5.31,0,0,0,1,6.5,5.31,5.31,0,0,0,4.9,4.9,5.31,5.31,0,0,0,6.5,1,5.31,5.31,0,0,0,8.1,4.9,5.31,5.31,0,0,0,12,6.5,5.46,5.46,0,0,0,6.5,12Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "    <script>\n",
              "      (() => {\n",
              "      const buttonEl =\n",
              "        document.querySelector('#id_5b0fcc4e-3a0d-4ee0-9cc4-351759ccc559 button.colab-df-generate');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      buttonEl.onclick = () => {\n",
              "        google.colab.notebook.generateWithVariable('V');\n",
              "      }\n",
              "      })();\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "    </div>\n",
              "  </div>\n"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "dataframe",
              "variable_name": "V",
              "summary": "{\n  \"name\": \"V\",\n  \"rows\": 3380,\n  \"fields\": [\n    {\n      \"column\": \"chrom\",\n      \"properties\": {\n        \"dtype\": \"category\",\n        \"num_unique_values\": 19,\n        \"samples\": [\n          \"1\",\n          \"7\",\n          \"13\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"pos\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 56516547,\n        \"min\": 90661,\n        \"max\": 248591989,\n        \"num_unique_values\": 3354,\n        \"samples\": [\n          233464121,\n          119367787,\n          63474189\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"ref\",\n      \"properties\": {\n        \"dtype\": \"category\",\n        \"num_unique_values\": 4,\n        \"samples\": [\n          \"G\",\n          \"T\",\n          \"C\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"alt\",\n      \"properties\": {\n        \"dtype\": \"category\",\n        \"num_unique_values\": 4,\n        \"samples\": [\n          \"T\",\n          \"C\",\n          \"G\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"OMIM\",\n      \"properties\": {\n        \"dtype\": \"category\",\n        \"num_unique_values\": 113,\n        \"samples\": [\n          \"MIM 613978\",\n          \"MIM 602771\",\n          \"MIM 188740\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"consequence\",\n      \"properties\": {\n        \"dtype\": \"category\",\n        \"num_unique_values\": 12,\n        \"samples\": [\n          \"upstream_gene_variant\",\n          \"pELS_flank\",\n          \"PLS\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"label\",\n      \"properties\": {\n        \"dtype\": \"boolean\",\n        \"num_unique_values\": 2,\n        \"samples\": [\n          true,\n          false\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"tss_dist\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 31507,\n        \"min\": 0,\n        \"max\": 371628,\n        \"num_unique_values\": 1387,\n        \"samples\": [\n          101764,\n          25171\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"match_group\",\n      \"properties\": {\n        \"dtype\": \"category\",\n        \"num_unique_values\": 338,\n        \"samples\": [\n          \"non_coding_transcript_exon_variant_66\",\n          \"3_prime_UTR_variant_28\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}"
            }
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load model"
      ],
      "metadata": {
        "id": "PyvRpy3-x9UH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Wrapper to compute a VEP score\n",
        "class ModelVEP(torch.nn.Module):\n",
        "    def __init__(self, model_path):\n",
        "        super().__init__()\n",
        "        self.model = AutoModel.from_pretrained(\n",
        "            model_path,\n",
        "            trust_remote_code=True,\n",
        "        )\n",
        "\n",
        "    def forward(\n",
        "        self,\n",
        "        input_ids_ref=None,\n",
        "        input_ids_alt=None,\n",
        "    ):\n",
        "        embed_ref = self.model(input_ids=input_ids_ref).last_hidden_state.reshape(len(input_ids_ref), -1)\n",
        "        embed_alt = self.model(input_ids=input_ids_alt).last_hidden_state.reshape(len(input_ids_ref), -1)\n",
        "        return F.pairwise_distance(embed_ref, embed_alt)\n",
        "\n",
        "\n",
        "class ModelVEPAverageRC(torch.nn.Module):\n",
        "    def __init__(self, model_path):\n",
        "        super().__init__()\n",
        "        self.model = AutoModel.from_pretrained(\n",
        "            model_path,\n",
        "            trust_remote_code=True,\n",
        "        )\n",
        "\n",
        "    def get_scores(self, input_ids_ref, input_ids_alt):\n",
        "        embed_ref = self.model(input_ids=input_ids_ref).last_hidden_state.reshape(len(input_ids_ref), -1)\n",
        "        embed_alt = self.model(input_ids=input_ids_alt).last_hidden_state.reshape(len(input_ids_ref), -1)\n",
        "        return F.pairwise_distance(embed_ref, embed_alt)\n",
        "\n",
        "    def forward(\n",
        "        self,\n",
        "        input_ids_ref_fwd=None,\n",
        "        input_ids_alt_fwd=None,\n",
        "        input_ids_ref_rev=None,\n",
        "        input_ids_alt_rev=None,\n",
        "    ):\n",
        "        fwd = self.get_scores(input_ids_ref_fwd, input_ids_alt_fwd)\n",
        "        rev = self.get_scores(input_ids_ref_rev, input_ids_alt_rev)\n",
        "        return (fwd + rev) / 2"
      ],
      "metadata": {
        "id": "JID0TXVj5lv-"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(model_path)\n",
        "window_size = 512\n",
        "if average_rc:\n",
        "    model = ModelVEPAverageRC(model_path)\n",
        "else:\n",
        "    model = ModelVEP(model_path)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 214,
          "referenced_widgets": [
            "0ade8e14a15944b695d390db199b9dd9",
            "c53c155008af4c4e9ae6c21598e3b0f1",
            "ff9dd793206542ffa9f31a48b7e0314d",
            "7146c456cbaa41498f9122fd0d29b1d2",
            "dcfa4d369495418295baf7fbe9568dce",
            "151e2df6bea14a43ad2705f678c99195",
            "6dbf91bea24840fea03b861144a3946e",
            "d1431f375a2e46838a5970cc80f61200",
            "dcfa667e270c40d79dd93309da9f77fb",
            "32418a73eb3b41f3be335fb89d2b877c",
            "fa0e9d33904e4900984fafcdcdf6e6e7",
            "3737b2439dbd4e568d366a9318c260a4",
            "c02a2073a45d4b93bc30bc5ed6be8805",
            "c5874ac66fd94a0b9cf43762c741118f",
            "51b43b7ce3d04b548361ddcb011ac227",
            "134c2b6d1ac94d1ebd41517d5f16d9b6",
            "7faa6f4bb40c4f49b21b8cdd999105be",
            "3a3bf3efb5334a41a7d1412a746c2e1c",
            "525edb5bb9194355b730557663f00a43",
            "8d88a2fedae6498bb3a5f57457333bd9",
            "c08096b9eb584ca4870d13f3e5824a9b",
            "74fc154af2304c2fb0ba48d14ecc85c7",
            "565dd0e46c434aa7a73d6bb908bd0680",
            "2b3fb166766c4f79a0efc8626c93b246",
            "baf99987c9b54fe881574d800ee09f14",
            "c8f25334ab334b5fac6797ab7d309f38",
            "a86710668b3540d0aa8bb4a36dce8ceb",
            "18028875cdaa4988bf19419280e0279a",
            "e22fdcb3b13c44638b6ac147cd5a4f8c",
            "d2d8169da2df4bf78acd960b4e3676b0",
            "0c3e11bcd3ac4f628676209f92949f14",
            "cc29c662fd6f4be0b49889c75655414d",
            "acb44fc0488843469bb1f83a0c856331",
            "1656307b55e241078531ac0983ee96fb",
            "b2956c54417446cf9346ad51e3808075",
            "df36a32b3b42485db9a27cec3cd21e31",
            "77988a600cb64c7a9773f84c1725a21b",
            "93c39b9dab7f4fa5ac7bc5f2c03e9175",
            "10b81b4b60a34853a9663e78bb2e1abe",
            "7b4126a076cb4dcbba25b04d19334b0b",
            "eae24a291baa4c04a1bc8d1966a3ccae",
            "c649b21f7c1642ebb880d3a62b8281ed",
            "46f44a143d7f4536a93db59e033ce665",
            "ed4930de91c440c0a1a7239a4ed06e5d",
            "a31b47baeb624307b0bb270d81949504",
            "7459011748494d4a9939220746108716",
            "c923d77c086d4abb93fabc6f8d20b5b7",
            "c8922ee8c0224221b2673f7618701da4",
            "abd5f43ea691418898820e68d36fdd8f",
            "2a38d3b1f0df48fb96b2e9256a0c564b",
            "3a408803721449a1aa63a65cd51f9a1d",
            "9a39944759c64b57a27f5c4e5ef935d7",
            "816ad4562cbc4ff1934da1f5cd3b69c7",
            "197fe37ca48b4dfeb674e8a0688102d0",
            "0faeff78cc7446bab7c32806805cce5a"
          ]
        },
        "id": "_-hrLbqwxqxH",
        "outputId": "ea4249ed-1ba8-4b30-9fcd-3f0efc17a93b"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "tokenizer_config.json:   0%|          | 0.00/784 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "0ade8e14a15944b695d390db199b9dd9"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "tokenizer.json:   0%|          | 0.00/1.15k [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "3737b2439dbd4e568d366a9318c260a4"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "special_tokens_map.json:   0%|          | 0.00/419 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "565dd0e46c434aa7a73d6bb908bd0680"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "config.json:   0%|          | 0.00/1.17k [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "1656307b55e241078531ac0983ee96fb"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "model.safetensors:   0%|          | 0.00/609M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "a31b47baeb624307b0bb270d81949504"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "dilation_schedule=[1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Tokenize dataset"
      ],
      "metadata": {
        "id": "BmoMd7eyx_eS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Genome:\n",
        "    def __init__(self, path):\n",
        "        self.data = Fasta(fsspec.open(path, anon=True))\n",
        "\n",
        "    def __call__(self, chrom, start, end, strand=\"+\"):\n",
        "        res = self.data[chrom][start:end]\n",
        "        if strand == \"-\":\n",
        "            res = res.reverse.complement\n",
        "        return str(res)"
      ],
      "metadata": {
        "id": "qiLlhmcT1BxK"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "genome = Genome(\"s3://broad-references/hg38/v0/Homo_sapiens_assembly38.fasta\")"
      ],
      "metadata": {
        "id": "50tERpTd1jGW"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def tokenize(seqs):\n",
        "    return tokenizer(\n",
        "        seqs,\n",
        "        padding=False,\n",
        "        truncation=False,\n",
        "        return_token_type_ids=False,\n",
        "        return_attention_mask=False,\n",
        "        return_special_tokens_mask=False,\n",
        "        return_tensors=\"pt\",\n",
        "    )[\"input_ids\"]\n",
        "\n",
        "def get_tokenized_seq(vs):\n",
        "    # we convert from 1-based coordinate (standard in VCF) to\n",
        "    # 0-based, to use with Genome\n",
        "    chrom = np.array(vs[\"chrom\"])\n",
        "    n = len(chrom)\n",
        "    pos = np.array(vs[\"pos\"]) - 1\n",
        "    start = pos - window_size // 2\n",
        "    end = pos + window_size // 2\n",
        "    seq_fwd = [genome(\"chr\" + chrom[i], start[i], end[i]) for i in range(n)]\n",
        "    seq_fwd = np.array([list(seq.upper()) for seq in seq_fwd], dtype=\"object\")\n",
        "    assert seq_fwd.shape[1] == window_size\n",
        "    ref_fwd = np.array(vs[\"ref\"])\n",
        "    alt_fwd = np.array(vs[\"alt\"])\n",
        "    pos_fwd = window_size // 2\n",
        "\n",
        "    def prepare_output(seq, pos, ref, alt):\n",
        "        assert (seq[:, pos] == ref).all(), f\"{seq[:, pos]}, {ref}\"\n",
        "        seq_ref = seq\n",
        "        seq_alt = seq.copy()\n",
        "        seq_alt[:, pos] = alt\n",
        "        return (\n",
        "            tokenize([\"\".join(x) for x in seq_ref]),\n",
        "            tokenize([\"\".join(x) for x in seq_alt]),\n",
        "        )\n",
        "\n",
        "    res = {}\n",
        "    res[\"input_ids_ref\"], res[\"input_ids_alt\"] = prepare_output(seq_fwd, pos_fwd, ref_fwd, alt_fwd)\n",
        "    return res\n",
        "\n",
        "def get_tokenized_seq_average_rc(vs):\n",
        "    # we convert from 1-based coordinate (standard in VCF) to\n",
        "    # 0-based, to use with Genome\n",
        "    chrom = np.array(vs[\"chrom\"])\n",
        "    n = len(chrom)\n",
        "    pos = np.array(vs[\"pos\"]) - 1\n",
        "    start = pos - window_size // 2\n",
        "    end = pos + window_size // 2\n",
        "    seq_fwd = [genome(\"chr\" + chrom[i], start[i], end[i]) for i in range(n)]\n",
        "    seq_rev = [str(Seq(x).reverse_complement()) for x in seq_fwd]\n",
        "    seq_fwd = np.array([list(seq.upper()) for seq in seq_fwd], dtype=\"object\")\n",
        "    seq_rev = np.array([list(seq.upper()) for seq in seq_rev], dtype=\"object\")\n",
        "    assert seq_fwd.shape[1] == window_size\n",
        "    assert seq_rev.shape[1] == window_size\n",
        "    ref_fwd = np.array(vs[\"ref\"])\n",
        "    alt_fwd = np.array(vs[\"alt\"])\n",
        "    ref_rev = np.array([str(Seq(x).reverse_complement()) for x in ref_fwd])\n",
        "    alt_rev = np.array([str(Seq(x).reverse_complement()) for x in alt_fwd])\n",
        "    pos_fwd = window_size // 2\n",
        "    pos_rev = pos_fwd - 1 if window_size % 2 == 0 else pos_fwd\n",
        "\n",
        "    def prepare_output(seq, pos, ref, alt):\n",
        "        assert (seq[:, pos] == ref).all(), f\"{seq[:, pos]}, {ref}\"\n",
        "        seq_ref = seq\n",
        "        seq_alt = seq.copy()\n",
        "        seq_alt[:, pos] = alt\n",
        "        return (\n",
        "            tokenize([\"\".join(x) for x in seq_ref]),\n",
        "            tokenize([\"\".join(x) for x in seq_alt]),\n",
        "        )\n",
        "\n",
        "    res = {}\n",
        "    res[\"input_ids_ref_fwd\"], res[\"input_ids_alt_fwd\"] = prepare_output(seq_fwd, pos_fwd, ref_fwd, alt_fwd)\n",
        "    res[\"input_ids_ref_rev\"], res[\"input_ids_alt_rev\"] = prepare_output(seq_rev, pos_rev, ref_rev, alt_rev)\n",
        "    return res"
      ],
      "metadata": {
        "id": "2Nsgbz8A2Fo3"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "if average_rc:\n",
        "    dataset.set_transform(get_tokenized_seq_average_rc)\n",
        "else:\n",
        "    dataset.set_transform(get_tokenized_seq)"
      ],
      "metadata": {
        "id": "zj4W7vHQ35Q3"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset[0]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GLDcWQj94hu0",
        "outputId": "dff9d8ca-6262-496e-9d62-c9ef0d7ae151"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'input_ids_ref': tensor([4, 4, 6, 5, 5, 4, 3, 5, 4, 4, 6, 5, 4, 6, 5, 3, 5, 6, 5, 4, 4, 6, 5, 4,\n",
              "         4, 5, 6, 5, 6, 3, 4, 5, 5, 5, 4, 4, 4, 4, 3, 4, 3, 4, 4, 4, 6, 5, 4, 5,\n",
              "         5, 4, 6, 4, 3, 4, 5, 4, 5, 4, 3, 6, 4, 4, 4, 3, 5, 5, 6, 4, 4, 4, 4, 6,\n",
              "         4, 6, 4, 3, 6, 4, 4, 3, 4, 6, 4, 6, 5, 3, 5, 3, 4, 5, 4, 3, 5, 5, 4, 5,\n",
              "         5, 3, 5, 5, 3, 4, 3, 4, 4, 5, 3, 4, 3, 3, 5, 5, 6, 5, 6, 6, 4, 3, 5, 3,\n",
              "         6, 5, 3, 5, 5, 5, 3, 5, 4, 3, 5, 4, 5, 5, 5, 5, 4, 6, 5, 5, 5, 5, 4, 6,\n",
              "         5, 5, 5, 3, 4, 3, 5, 5, 5, 5, 5, 6, 5, 5, 5, 5, 5, 5, 4, 6, 5, 5, 4, 4,\n",
              "         6, 4, 4, 3, 5, 5, 6, 5, 5, 4, 3, 5, 4, 6, 3, 4, 6, 5, 3, 6, 5, 5, 5, 5,\n",
              "         4, 3, 5, 4, 6, 3, 4, 4, 5, 5, 3, 5, 6, 4, 3, 3, 5, 4, 4, 4, 6, 6, 6, 5,\n",
              "         3, 3, 3, 6, 5, 3, 5, 5, 6, 4, 6, 6, 6, 5, 6, 5, 6, 5, 5, 5, 5, 5, 3, 5,\n",
              "         4, 6, 5, 4, 3, 5, 4, 4, 4, 4, 5, 4, 4, 4, 5, 5, 4, 3, 4, 3, 5, 5, 5, 3,\n",
              "         5, 5, 5, 3, 5, 5, 4, 5, 5, 4, 4, 4, 4, 5, 5, 6, 5, 5, 5, 3, 3, 5, 5, 5,\n",
              "         4, 4, 4, 6, 5, 4, 3, 3, 5, 5, 4, 4, 6, 5, 4, 4, 4, 3, 5, 3, 5, 4, 3, 4,\n",
              "         4, 3, 5, 5, 3, 5, 4, 6, 5, 5, 4, 6, 5, 5, 5, 5, 4, 6, 5, 5, 3, 5, 3, 4,\n",
              "         6, 5, 5, 4, 5, 4, 6, 5, 4, 6, 5, 5, 3, 5, 4, 6, 5, 4, 3, 5, 4, 6, 5, 4,\n",
              "         4, 6, 5, 5, 4, 3, 5, 3, 5, 5, 5, 5, 3, 3, 3, 4, 6, 5, 3, 5, 5, 6, 5, 5,\n",
              "         5, 5, 5, 6, 5, 6, 5, 5, 4, 6, 4, 6, 5, 5, 5, 3, 4, 3, 5, 3, 4, 3, 6, 4,\n",
              "         3, 5, 5, 5, 4, 3, 5, 5, 3, 4, 4, 6, 4, 4, 4, 6, 3, 5, 4, 4, 3, 5, 5, 5,\n",
              "         6, 4, 4, 3, 4, 5, 5, 6, 4, 6, 5, 5, 4, 6, 5, 5, 5, 5, 5, 3, 6, 5, 6, 5,\n",
              "         5, 3, 4, 6, 5, 5, 4, 6, 5, 4, 3, 3, 4, 5, 6, 5, 5, 5, 5, 4, 5, 5, 5, 5,\n",
              "         6, 6, 3, 3, 3, 5, 5, 6, 4, 6, 6, 4, 4, 3, 5, 5, 5, 4, 3, 5, 4, 4, 4, 6,\n",
              "         5, 5, 3, 5, 3, 5, 4, 3]),\n",
              " 'input_ids_alt': tensor([4, 4, 6, 5, 5, 4, 3, 5, 4, 4, 6, 5, 4, 6, 5, 3, 5, 6, 5, 4, 4, 6, 5, 4,\n",
              "         4, 5, 6, 5, 6, 3, 4, 5, 5, 5, 4, 4, 4, 4, 3, 4, 3, 4, 4, 4, 6, 5, 4, 5,\n",
              "         5, 4, 6, 4, 3, 4, 5, 4, 5, 4, 3, 6, 4, 4, 4, 3, 5, 5, 6, 4, 4, 4, 4, 6,\n",
              "         4, 6, 4, 3, 6, 4, 4, 3, 4, 6, 4, 6, 5, 3, 5, 3, 4, 5, 4, 3, 5, 5, 4, 5,\n",
              "         5, 3, 5, 5, 3, 4, 3, 4, 4, 5, 3, 4, 3, 3, 5, 5, 6, 5, 6, 6, 4, 3, 5, 3,\n",
              "         6, 5, 3, 5, 5, 5, 3, 5, 4, 3, 5, 4, 5, 5, 5, 5, 4, 6, 5, 5, 5, 5, 4, 6,\n",
              "         5, 5, 5, 3, 4, 3, 5, 5, 5, 5, 5, 6, 5, 5, 5, 5, 5, 5, 4, 6, 5, 5, 4, 4,\n",
              "         6, 4, 4, 3, 5, 5, 6, 5, 5, 4, 3, 5, 4, 6, 3, 4, 6, 5, 3, 6, 5, 5, 5, 5,\n",
              "         4, 3, 5, 4, 6, 3, 4, 4, 5, 5, 3, 5, 6, 4, 3, 3, 5, 4, 4, 4, 6, 6, 6, 5,\n",
              "         3, 3, 3, 6, 5, 3, 5, 5, 6, 4, 6, 6, 6, 5, 6, 5, 6, 5, 5, 5, 5, 5, 3, 5,\n",
              "         4, 6, 5, 4, 3, 5, 4, 4, 4, 4, 5, 4, 4, 4, 5, 5, 5, 3, 4, 3, 5, 5, 5, 3,\n",
              "         5, 5, 5, 3, 5, 5, 4, 5, 5, 4, 4, 4, 4, 5, 5, 6, 5, 5, 5, 3, 3, 5, 5, 5,\n",
              "         4, 4, 4, 6, 5, 4, 3, 3, 5, 5, 4, 4, 6, 5, 4, 4, 4, 3, 5, 3, 5, 4, 3, 4,\n",
              "         4, 3, 5, 5, 3, 5, 4, 6, 5, 5, 4, 6, 5, 5, 5, 5, 4, 6, 5, 5, 3, 5, 3, 4,\n",
              "         6, 5, 5, 4, 5, 4, 6, 5, 4, 6, 5, 5, 3, 5, 4, 6, 5, 4, 3, 5, 4, 6, 5, 4,\n",
              "         4, 6, 5, 5, 4, 3, 5, 3, 5, 5, 5, 5, 3, 3, 3, 4, 6, 5, 3, 5, 5, 6, 5, 5,\n",
              "         5, 5, 5, 6, 5, 6, 5, 5, 4, 6, 4, 6, 5, 5, 5, 3, 4, 3, 5, 3, 4, 3, 6, 4,\n",
              "         3, 5, 5, 5, 4, 3, 5, 5, 3, 4, 4, 6, 4, 4, 4, 6, 3, 5, 4, 4, 3, 5, 5, 5,\n",
              "         6, 4, 4, 3, 4, 5, 5, 6, 4, 6, 5, 5, 4, 6, 5, 5, 5, 5, 5, 3, 6, 5, 6, 5,\n",
              "         5, 3, 4, 6, 5, 5, 4, 6, 5, 4, 3, 3, 4, 5, 6, 5, 5, 5, 5, 4, 5, 5, 5, 5,\n",
              "         6, 6, 3, 3, 3, 5, 5, 6, 4, 6, 6, 4, 4, 3, 5, 5, 5, 4, 3, 5, 4, 4, 4, 6,\n",
              "         5, 5, 3, 5, 3, 5, 4, 3])}"
            ]
          },
          "metadata": {},
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Run inference"
      ],
      "metadata": {
        "id": "-lz_YC-G539T"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "training_args = TrainingArguments(\n",
        "    output_dir=tempfile.TemporaryDirectory().name,\n",
        "    per_device_eval_batch_size=128,\n",
        "    # pyfaidx does not allow multiple workers\n",
        "    # for longer jobs you can use the gpn.data.Genome class,\n",
        "    # which loads a local fasta file into memory and can use multiple\n",
        "    # workers, see e.g.\n",
        "    # https://github.com/songlab-cal/gpn/blob/main/gpn/ss/run_vep_embed_dist.py\n",
        "    dataloader_num_workers=0,\n",
        "    remove_unused_columns=False,\n",
        "    torch_compile=False,\n",
        "    fp16=True,\n",
        "    report_to=\"none\",\n",
        ")\n",
        "trainer = Trainer(model=model, args=training_args)\n",
        "preds = trainer.predict(test_dataset=dataset).predictions\n",
        "preds.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "id": "z7cS7RDC5DWJ",
        "outputId": "db12461a-41ac-4539-8bc6-5f781cf10d03"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": []
          },
          "metadata": {}
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(3380,)"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Compute metrics"
      ],
      "metadata": {
        "id": "ijChgnmI55fn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "V[\"score\"] = preds"
      ],
      "metadata": {
        "id": "84DkGrmPN5nN"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# let's compare to the precomputed scores\n",
        "# this is using RC averaging\n",
        "V[\"precomputed_score\"] = pd.read_parquet(\"hf://datasets/songlab/TraitGym/mendelian_traits_matched_9/features/GPN_final_EuclideanDistance.parquet\").score.values\n",
        "V[[\"score\", \"precomputed_score\"]].corr()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 112
        },
        "id": "HheYRetBSher",
        "outputId": "5a9dccc1-bb6f-4bc1-8c71-d8dedd776022"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "                      score  precomputed_score\n",
              "score              1.000000           0.956949\n",
              "precomputed_score  0.956949           1.000000"
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-64577e70-3d8f-4692-a29e-77ca0b01c5d3\" class=\"colab-df-container\">\n",
              "    <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>score</th>\n",
              "      <th>precomputed_score</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>score</th>\n",
              "      <td>1.000000</td>\n",
              "      <td>0.956949</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>precomputed_score</th>\n",
              "      <td>0.956949</td>\n",
              "      <td>1.000000</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-64577e70-3d8f-4692-a29e-77ca0b01c5d3')\"\n",
              "            title=\"Convert this dataframe to an interactive table.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "    <script>\n",
              "      const buttonEl =\n",
              "        document.querySelector('#df-64577e70-3d8f-4692-a29e-77ca0b01c5d3 button.colab-df-convert');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      async function convertToInteractive(key) {\n",
              "        const element = document.querySelector('#df-64577e70-3d8f-4692-a29e-77ca0b01c5d3');\n",
              "        const dataTable =\n",
              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                    [key], {});\n",
              "        if (!dataTable) return;\n",
              "\n",
              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "          + ' to learn more about interactive tables.';\n",
              "        element.innerHTML = '';\n",
              "        dataTable['output_type'] = 'display_data';\n",
              "        await google.colab.output.renderOutput(dataTable, element);\n",
              "        const docLink = document.createElement('div');\n",
              "        docLink.innerHTML = docLinkHtml;\n",
              "        element.appendChild(docLink);\n",
              "      }\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "\n",
              "<div id=\"df-aefb2118-307a-4fe8-9070-7058f05b7a16\">\n",
              "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-aefb2118-307a-4fe8-9070-7058f05b7a16')\"\n",
              "            title=\"Suggest charts\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "     width=\"24px\">\n",
              "    <g>\n",
              "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
              "    </g>\n",
              "</svg>\n",
              "  </button>\n",
              "\n",
              "<style>\n",
              "  .colab-df-quickchart {\n",
              "      --bg-color: #E8F0FE;\n",
              "      --fill-color: #1967D2;\n",
              "      --hover-bg-color: #E2EBFA;\n",
              "      --hover-fill-color: #174EA6;\n",
              "      --disabled-fill-color: #AAA;\n",
              "      --disabled-bg-color: #DDD;\n",
              "  }\n",
              "\n",
              "  [theme=dark] .colab-df-quickchart {\n",
              "      --bg-color: #3B4455;\n",
              "      --fill-color: #D2E3FC;\n",
              "      --hover-bg-color: #434B5C;\n",
              "      --hover-fill-color: #FFFFFF;\n",
              "      --disabled-bg-color: #3B4455;\n",
              "      --disabled-fill-color: #666;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart {\n",
              "    background-color: var(--bg-color);\n",
              "    border: none;\n",
              "    border-radius: 50%;\n",
              "    cursor: pointer;\n",
              "    display: none;\n",
              "    fill: var(--fill-color);\n",
              "    height: 32px;\n",
              "    padding: 0;\n",
              "    width: 32px;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart:hover {\n",
              "    background-color: var(--hover-bg-color);\n",
              "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "    fill: var(--button-hover-fill-color);\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart-complete:disabled,\n",
              "  .colab-df-quickchart-complete:disabled:hover {\n",
              "    background-color: var(--disabled-bg-color);\n",
              "    fill: var(--disabled-fill-color);\n",
              "    box-shadow: none;\n",
              "  }\n",
              "\n",
              "  .colab-df-spinner {\n",
              "    border: 2px solid var(--fill-color);\n",
              "    border-color: transparent;\n",
              "    border-bottom-color: var(--fill-color);\n",
              "    animation:\n",
              "      spin 1s steps(1) infinite;\n",
              "  }\n",
              "\n",
              "  @keyframes spin {\n",
              "    0% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "      border-left-color: var(--fill-color);\n",
              "    }\n",
              "    20% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    30% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    40% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    60% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    80% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "    90% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "  }\n",
              "</style>\n",
              "\n",
              "  <script>\n",
              "    async function quickchart(key) {\n",
              "      const quickchartButtonEl =\n",
              "        document.querySelector('#' + key + ' button');\n",
              "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
              "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
              "      try {\n",
              "        const charts = await google.colab.kernel.invokeFunction(\n",
              "            'suggestCharts', [key], {});\n",
              "      } catch (error) {\n",
              "        console.error('Error during call to suggestCharts:', error);\n",
              "      }\n",
              "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
              "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
              "    }\n",
              "    (() => {\n",
              "      let quickchartButtonEl =\n",
              "        document.querySelector('#df-aefb2118-307a-4fe8-9070-7058f05b7a16 button');\n",
              "      quickchartButtonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "    })();\n",
              "  </script>\n",
              "</div>\n",
              "\n",
              "    </div>\n",
              "  </div>\n"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "dataframe",
              "summary": "{\n  \"name\": \"V[[\\\"score\\\", \\\"precomputed_score\\\"]]\",\n  \"rows\": 2,\n  \"fields\": [\n    {\n      \"column\": \"score\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 0.030441660786716115,\n        \"min\": 0.9569489904542648,\n        \"max\": 1.0,\n        \"num_unique_values\": 2,\n        \"samples\": [\n          0.9569489904542648,\n          1.0\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"precomputed_score\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 0.030441660786716115,\n        \"min\": 0.9569489904542648,\n        \"max\": 1.0,\n        \"num_unique_values\": 2,\n        \"samples\": [\n          1.0,\n          0.9569489904542648\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}"
            }
          },
          "metadata": {},
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "plt.figure(figsize=(2, 2))\n",
        "sns.histplot(\n",
        "    data=V, x=\"score\", bins=30, hue=\"label\", stat=\"density\",\n",
        "    common_norm=False, common_bins=True,\n",
        ")\n",
        "sns.despine();"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 234
        },
        "id": "4y9p5gEtQpcg",
        "outputId": "74aca40a-2123-4a50-c0bd-a062e8cbfb15"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 200x200 with 1 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAOIAAADZCAYAAAA9rey5AAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQAAJLNJREFUeJzt3XlUFFfaB+Bfg3QDyiKyqiAiKpIoKIjhqIMGRszibuIYRwGjmURJVI4mEifgkknjxqAOESdGicYoMVGzmDgSFBIUF1AWEYmK2EZZg+yr9P3+4KOkpVFoursKeJ9z+hz71q2qt01eq+rWXUSMMQZCCK90+A6AEEKJSIggUCISIgCUiIQIACUiIQJAiUiIAFAiEiIAlIiECECPS0TGGMrLy0H9GIiQ9LhErKiogImJCSoqKvgOhRBOj0tEQoSIEpEQARBEIkZGRsLe3h76+voYN24cLl261K79jhw5ApFIhJkzZ2o2QEI0jPdEjImJQVBQEEJDQ3HlyhW4uLjA19cXhYWFT90vNzcXq1evxsSJE7UUKSGaw3sihoeHY+nSpQgICICzszOioqJgaGiIffv2tblPY2MjFixYgA0bNsDBwUGL0RKiGb34PHl9fT1SUlIQHBzMleno6MDHxwdJSUlt7rdx40ZYWlrizTffxG+//fbUc9TV1aGuro77Xl5e3um4a2trkZycrFDm7u4OfX39Th9bGxhjePToERobG/kOpcvT09ODrq5up4/DayIWFxejsbERVlZWCuVWVla4ceOG0n0SExPx+eefIzU1tV3nkEql2LBhQ2dDVZCcnIyImFjYOAwHAOTlZGMlgAkTJqj1PJpQX1+PvLw8VFdX8x1KtyASiTBw4ED06dOnU8fhNRE7qqKiAgsXLsRnn30Gc3Pzdu0THByMoKAg7nt5eTlsbW07HYuNw3A4PO/e6eNok1wux507d6Crq4v+/ftDLBZDJBLxHVaXxRhDUVER/vjjDwwdOrRTV0ZeE9Hc3By6urooKChQKC8oKIC1tXWr+rdv30Zubi6mTZvGlcnlcgBAr169kJ2djSFDhijsI5FIIJFINBB911NfXw+5XA5bW1sYGhryHU63YGFhgdzcXDQ0NHQqEXltrBGLxXBzc0NcXBxXJpfLERcXB09Pz1b1nZyckJGRgdTUVO4zffp0TJ48GampqWq50vUEOjq8t9F1G+q6o+D91jQoKAh+fn5wd3eHh4cHIiIiUFVVhYCAAADAokWLMGDAAEilUujr6+P5559X2N/U1BQAWpUT0pXwnojz5s1DUVERQkJCkJ+fD1dXV5w6dYprwJHJZLz+C66shTQ9PR2NjUY8RUS6I94TEQACAwMRGBiodFt8fPxT942OjlZ/QC082UIKABmJl2E51AVDNXrmrm/SpElwdXVFRETEM+vGx8dj8uTJePjwIXeXowp7e3usXLkSK1euVPkYfBBEIgrdky2kD3KyeYyGdEf01E6IAFAiEq04ePAg3N3dYWRkBGtra7zxxhtK+xOfO3cOo0aNgr6+Pl544QVcu3ZNYXtiYiImTpwIAwMD2Nra4r333kNVVZW2fobGUCISrWhoaMCmTZuQlpaGEydOIDc3F/7+/q3qrVmzBtu3b8fly5dhYWGBadOmoaGhAUDTe+SpU6dizpw5SE9PR0xMDBITE9tsX+hK6BmRaMXixYu5Pzs4OGDnzp0YO3YsKisrFbqHhYaG4q9//SsA4IsvvsDAgQNx/PhxvP7665BKpViwYAHXEDN06FDs3LkTXl5e2L17d5fp66sMXRGJVqSkpGDatGmws7ODkZERvLy8ADS9nmqpZUcOMzMzDB8+HFlZWQCAtLQ0REdHo0+fPtzH19eX67rXldEVkWhcVVUVfH194evri0OHDsHCwgIymQy+vr6or69v93EqKyvxj3/8A++9916rbXZ2duoMWesoEYnG3bhxA3/++SfCwsK4bohPdpJoduHCBS6pHj58iN9//x0jRowAAIwZMwbXr1+Ho6OjdgLXIro1JRpnZ2cHsViMXbt2IScnB99//z02bdqktO7GjRsRFxeHa9euwd/fH+bm5txUKB988AHOnz+PwMBApKam4ubNm/juu++6RWMNJSLROAsLC0RHR+Po0aNwdnZGWFgYtm3bprRuWFgYVqxYATc3N+Tn5+OHH36AWCwGAIwaNQoJCQn4/fffMXHiRIwePRohISHo37+/Nn+ORtCtKdGYlt0T58+fj/nz5ytsbznJ86RJk7jvr776apvHHDt2LE6fPt3m9tzcXNWC5RldEQkRAEpEQgSAEpEQAaBEJEQAKBEJEQBKREIEgBKREAGgRCREAOiFPgHQNAqiuLhYa+czNzfXSkft6OhorFy5EqWlpRo/V2dQIhLIZDI4jRiBGi1Ow29gaIgbWVntTkZ/f3988cUXrcpv3rzZLTqBUyISFBcXo6a6Ggs+2AoruyHP3qGTCmS3cWjzGhQXF3foqjh16lTs379foczCwkLd4fGCEpFwrOyGYODQ5/gOo00SiaTVUgzh4eHYv38/cnJyYGZmhmnTpmHLli1tLgqTlpaGlStXIjk5GSKRCEOHDsWePXvg7t40S19iYiKCg4ORnJwMc3NzzJo1C1KpFL1799bob6PGGtKl6ejoYOfOncjMzMQXX3yBM2fO4P3332+z/oIFCzBw4EBcvnwZKSkpWLt2LfT09ADwOycOXRFJl/Hjjz8qXOleeuklHD16lPtub2+Pjz/+GG+//TY+/fRTpceQyWRYs2YNnJycADTNe9OMzzlxKBFJlzF58mTs3r2b+967d2/88ssvkEqluHHjBsrLy/Ho0SPU1taiurpa6YpXQUFBWLJkCQ4ePAgfHx+89tpr3ApiaWlpSE9Px6FDh7j6jDFuTpzmmQI0QRC3ppGRkbC3t4e+vj7GjRuHS5cutVn32LFjcHd3h6mpKXr37g1XV1ccPHhQi9ESvvTu3RuOjo7cp66uDq+++ipGjRqFb7/9FikpKYiMjASANufCWb9+PTIzM/HKK6/gzJkzcHZ2xvHjxwE8nhOn5WpjaWlpuHnzZqvl/tSN9ytiTEwMgoKCEBUVhXHjxiEiIgK+vr7Izs6GpaVlq/pmZmZYt24dnJycIBaL8eOPPyIgIACWlpbw9fXl4RcQvqSkpEAul2P79u3cQkVff/31M/cbNmwYhg0bhlWrVmH+/PnYv38/Zs2axeucOLxfEcPDw7F06VIEBATA2dkZUVFRMDQ0xL59+5TWnzRpEmbNmoURI0ZgyJAhWLFiBUaNGoXExEQtR0745ujoiIaGBm4unIMHDyIqKqrN+jU1NQgMDER8fDzu3r2Lc+fO4fLly9wtJ59z4vB6Rayvr0dKSgqCg4O5Mh0dHfj4+CApKemZ+zPGcObMGWRnZ2Pz5s1K69TV1aGuro77Xl5e3vnAu6kC2e0udR4XFxeEh4dj8+bNCA4Oxl/+8hdIpVIsWrRIaX1dXV38+eefWLRoEQoKCmBubo7Zs2djw4YNAB7PibNu3TpMnDgRjDEMGTIE8+bNU0u8T8NrIhYXF6OxsZFbC7GZlZUVbty40eZ+ZWVlGDBgAOrq6qCrq4tPP/2Umx36SVKplPuLJsqZm5vDwNAQhzav0do5DQwNYW5u3u76bS2/t2rVKqxatUqhbOHChdyf/f39uan9xWIxDh8+/NTzPGtOHE3h/RlRFUZGRkhNTUVlZSXi4uIQFBQEBwcHTJo0qVXd4OBgBAUFcd/Ly8tpie8n2NnZ4UZWVrfsa9pV8JqI5ubm0NXVRUFBgUJ5QUFBqx4ULeno6HAP1K6ursjKyoJUKlWaiBKJBBKJRK1xd0d2dnaUGDzitbFGLBbDzc0NcXFxXJlcLkdcXJzCGgjPIpfLFZ4DCelqeL81DQoKgp+fH9zd3eHh4YGIiAhUVVUhICAAALBo0SIMGDAAUqkUQNMzn7u7O4YMGYK6ujr89NNPOHjwoMKLXkK6Gt4Tcd68eSgqKkJISAjy8/Ph6uqKU6dOcQ04MpmMe0cENC1osmzZMvzxxx8wMDCAk5MTvvzyS620bBGiKbwnIgAEBga2+a6m5WzRAPDxxx/j448/1kJUhGgP7y/0CSGUiIQIgkq3pjk5OXBwcFB3LIJUV1eHgoJC9DJ+vLLtw4cPYWZgxGNUpLtRKREdHR3h5eWFN998E3Pnzu3Sa5c/S3Z2Nh7mZmBAHzlXVpF3G7qGpvwFpQHddfKorkKlRLxy5Qr279+PoKAgBAYGYt68eXjzzTfh4eGh7vgEoa+RAeysTLnvRobdq4OATCbDiBFOqK6u0do5DQ0NkJV1o13JKBKJnro9NDQU69evV1Nk/FApEV1dXbFjxw5s374d33//PaKjozFhwgQMGzYMixcvxsKFC7vNpD49QXFxMaqra/Dlh69jhJ3m/7tlyYrw90++bvfkUXl5edyfY2JiEBISguzsbK6s5ah9xhgaGxvRq5cgXgi0W6ei7dWrF2bPno1XXnkFn376KYKDg7F69Wp8+OGHeP3117F582bY2NioK1aiYSPsLDBm2AC+w2ilZXdHExMTiEQiriw+Ph6TJ0/GTz/9hH/+85/IyMjA6dOnER0djdLSUpw4cYLbd+XKlUhNTeVeicnlcmzevBn//e9/kZ+fj2HDhuGjjz7C3LlztfnzAHSy1TQ5ORnLli2DjY0NwsPDsXr1aty+fRuxsbF48OABZsyYoa44BUUuB8oryiGTySCTyVBQUEhd7Hi2du1ahIWFISsrC6NGjWrXPlKpFAcOHEBUVBQyMzOxatUq/P3vf0dCQoKGo21NpSti8xR22dnZePnll3HgwAG8/PLLXA+YwYMHIzo6Gvb29uqMVTCqauvQKH8A5KUDAB7mZiI7Wwxvb2+eI+u5Nm7c2OZQOGXq6urwySef4JdffuH6NTs4OCAxMRF79uyBl5eXpkJVSqVE3L17NxYvXgx/f/82bz0tLS3x+eefdyo4IetjIOYacO7fNeA3GMLNS9pet27dQnV1davkra+vx+jRo9UZWruolIixsbGws7NT6AMKND0o37t3D3Z2dhCLxfDz81NLkIQ8y5MTAOvo6IAxplDW0NDA/bmyshIAcPLkSQwYoPhczMewOZUScciQIcjLy2s1uVNJSQkGDx6MxsZGtQRHiKosLCxw7do1hbLU1FRuMmFnZ2dIJBLIZDKt34Yqo1IiPvkvTbPKyspu/XKfdB0vvvgitm7digMHDsDT0xNffvklrl27xt12GhkZYfXq1Vi1ahXkcjkmTJiAsrIynDt3DsbGxlq/m+tQIjZPOSESiRASEqIwgWtjYyMuXrwIV1dXtQZItCdLVtRtzuPr64uPPvoI77//Pmpra7F48WIsWrQIGRkZXJ1NmzbBwsICUqkUOTk5MDU1xZgxY/Dhhx9qPL4ndSgRr169CqDpipiRkQGxWMxtE4vFcHFxwerVq9UbIdE4c3NzGBoa4O+fPHtOUHUxNDTo0ORRzVpOBgU0Ta/Z1h3ahg0bnjpxmEgkwooVK7BixYoOx6FuHUrEs2fPAgACAgKwY8cOGBsbayQool12dnbIyrpBfU15pNIz4pNr1JGujyaP4le7E3H27NmIjo6GsbExZs+e/dS6x44d63RghPQk7U7E5j5+zX8mhKhPuxOx5e0o3ZoSol4qdfquqalBdXU19/3u3buIiIjgZapy0nFttTKSjlPX36VKiThjxgwcOHAAAFBaWgoPDw9s374dM2bMoPlFBay5V0nLf0RJ5zSvw6irq9up46g8Qv/f//43AOCbb76BtbU1rl69im+//RYhISF45513OhUU0QxdXV2YmpqisLAQAGBoaPjM0e+kbXK5HEVFRTA0NOz0QGSV9q6uroaRUdPkSadPn8bs2bOho6ODF154AXfv3u1UQESzmgfUNicj6RwdHR3Y2dl1+h80lSePOnHiBGbNmoX//e9/3LJYhYWF9JJf4EQiEWxsbGBpaakwGoGoRiwWtxqFpAqVEjEkJARvvPEGVq1aBW9vb25g5enTp1UayxUZGYmtW7ciPz8fLi4u2LVrV5sTUX322Wc4cOAA17Pezc0Nn3zySbeduEpTdHV1O/1cQ9RHpVSeO3cuZDIZkpOTcerUKa7c29ube3Zsr5iYGAQFBSE0NBRXrlyBi4sLfH1927x1io+Px/z583H27FkkJSXB1tYWU6ZMwf3791X5KYQIgsrXVGtra4wePVrhsuzh4QEnJ6cOHSc8PBxLly5FQEAAnJ2dERUVBUNDQ+zbt09p/UOHDmHZsmVwdXWFk5MT9u7dyy3lRkhXpdKtaVVVFcLCwhAXF4fCwkLI5XKF7Tk5Oe06Tn19PVJSUhAcHMyV6ejowMfHB0lJSe06RnV1NRoaGmBmZqZ0e11dncLETuXl5e06LiHapFIiLlmyBAkJCVi4cCFsbGxUbjEqLi5GY2MjtwRbMysrK9y4caNdx/jggw/Qv39/+Pj4KN0ulUqfOhSGECFQKRF//vlnnDx5EuPHj1d3PB0SFhaGI0eOID4+vs2ZAYKDg7kBzUDTFdHW1lZbIRLSLiolYt++fdu8FewIc3Nz6OrqoqCgQKG8oKBAYVJZZbZt24awsDD88ssvT53HUiKR8DIZECEdoVJjzaZNmxASEtLprlJisRhubm4KDS3NDS/Nr0SU2bJlCzZt2oRTp051eBo9QoRIpSvi9u3bcfv2bVhZWcHe3p7rw9jsypUr7T5WUFAQ/Pz84O7uDg8PD0RERKCqqgoBAQEAgEWLFmHAgAGQSqUAgM2bNyMkJARfffUV7O3tkZ+fD6Bp/YOWayAQ0pWolIgzZ85UWwDz5s1DUVERQkJCkJ+fD1dXV5w6dYprwJHJZAqvSHbv3o36+vpW6xN0hxWBSM+lUiKGhoaqNYjAwEAEBgYq3da8YEiz3NxctZ6bECFQ+YV+aWkp9u7di+DgYJSUlABouiWlHi6EdJxKV8T09HT4+PjAxMQEubm5WLp0KczMzHDs2DHIZDJurGJPIZcz3Lp1C4mJiQrl7u7uNOEyaReVEjEoKAj+/v7YsmULNxwKAF5++WW88cYbaguuqyirqkXF/d8A+yquLD0nH8AaTJgwgb/ASJehUiJevnwZe/bsaVU+YMAArhWzpxlkaYIJI+35DoN0USo9I0okEqV9Nn///XdaspsQFaiUiNOnT8fGjRu5gaUikQgymQwffPAB5syZo9YACekJVErE7du3o7KyEhYWFqipqYGXlxccHR1hZGSEf/3rX+qOkZBuT6VnRBMTE8TGxuLcuXNIS0tDZWUlxowZ0+YICELI03U4EeVyOaKjo3Hs2DHk5uZCJBJh8ODBsLa2BmOMZgUjRAUdujVljGH69OlYsmQJ7t+/j5EjR+K5557D3bt34e/vj1mzZmkqTkK6tQ5dEaOjo/Hrr78iLi4OkydPVth25swZzJw5EwcOHMCiRYvUGiQh3V2HroiHDx/Ghx9+2CoJgaalkteuXYtDhw6pLThCeooOJWJ6ejqmTp3a5vaXXnoJaWlpnQ6KkJ6mQ4lYUlLSan6ZlqysrPDw4cNOB0VIT9OhRGxsbHzqHP+6urp49OhRp4MipKfpUGMNYwz+/v5tzgHTctpCQkj7dSgR/fz8nlmHWkwJ6bgOJSKtFEyIZnR+GRtCSKdRIhIiAJSIhAgAJSIhAkCJSIgAUCISIgCUiIQIACUiIQLAeyJGRkbC3t4e+vr6GDduHC5dutRm3czMTMyZMwf29vYQiUSIiIjQXqBP0dj4CHcLy5CYkct90m/nU5c/0m68JmJMTAyCgoIQGhqKK1euwMXFBb6+vigsLFRav7q6Gg4ODggLC3vm+onaVFFSjEs1Noj5w4z7fJPTC9nZ2XyHRroIXhMxPDwcS5cuRUBAAJydnREVFQVDQ0Ps27dPaf2xY8di69at+Nvf/ia4xUdN+lnAYehQ7mNubcN3SKQL4S0R6+vrkZKSojDzm46ODnx8fJCUlKS289TV1aG8vFzhQ4jQ8JaIxcXFaGxsbDXQ2MrKSq3T9kulUpiYmHAfW1tbtR2bEHXhvbFG04KDg1FWVsZ97t27x3dIhLSi0gTD6mBubg5dXV0UFBQolBcUFKi1IUYikQjueZKQJ/F2RRSLxXBzc0NcXBxXJpfLERcXB09PT77CIoQXvF0RgaZ1Fv38/ODu7g4PDw9ERESgqqoKAQEBAJpG+w8YMABSqRRAUwPP9evXuT/fv38fqamp6NOnDxwdHXn7HYR0Fq+JOG/ePBQVFSEkJAT5+flwdXXFqVOnuAYcmUwGHZ3HF+0HDx5g9OjR3Pdt27Zh27Zt8PLyQnx8vLbDJ0RteE1EAAgMDERgYKDSbU8ml729PRhjWoiKEO3q9q2mhHQFlIiECAAlIiECQIlIiADw3ljTXcnlDLdu3UJiYqJCubu7O/T19XmKiggVJaKGlFXVouL+b4B9FVeWnpMPYA0mTJjAX2BEkCgRNWiQpQkmjLTnOwzSBdAzIiECQFdELapreITs9PRW5fTcSCgRtShbVoySh4eBR49XVabnRgJQImrdCNt+9NxIWqFnREIEgBKREAGgW9MWamtrkZycrFB269YtyOU04oNoFiViC8nJyYiIiYWNw3Cu7Nw1GfqZmXT4WC0nHW5268Gf6OvQV6GespZUakXteSgRn2DjMBwOz7tz3zNTLgCPqp6yh3IVJcW49MgG+n+YcWUXSkxhZlCqUO/JllRqRe2ZKBE1qHnS4WY3ZIUAylrVo5ZUQo01hAgAJSIhAkCJSIgA0DNiN6bsdQxArbJCRInYTShLuvT0dDRm/YzRQ/s/LqNWWUGiRNSi9rxbVPZesXnB0yeXDmh5ZUtOTkb6N1sxyuHxcgW3fsvE+GHm1CLbBVAialF73i0qG6Fx8rdMmBmIMMndmStTdmUb5WCtkHTpt9W3qhbRLEpELWvPu8Un3yum386HVR8RXdm6MUrEFurq6lBQUIhexjKurKK8HCYGmmtcfvJ2VVk3ONL9CeL1RWRkJOzt7aGvr49x48bh0qVLT61/9OhRODk5QV9fHyNHjsRPP/2kljiys7PxMDcDyEvnPlV/PkB9Q71ajq9MRUkxLtXYIOYPM8T8YYbfSkxxO69UY+cjwsT7FTEmJgZBQUGIiorCuHHjEBERAV9fX2RnZ8PS0rJV/fPnz2P+/PmQSqV49dVX8dVXX2HmzJm4cuUKnn/++U7H09fIAHZWptz33vp6nT7ms7S8Xc288wB3C2Ud7iyenp4O58bGZ56rK3cyV9Yy3J7YVd1Pm3hPxPDwcCxdupRbii0qKgonT57Evn37sHbt2lb1d+zYgalTp2LNmjUAgE2bNiE2Nhb/+c9/EBUVpdXYNUFZg875IiNUVNyF7YDHryG+O5eFiopLQLYDACAh/Q76egwCMOSpx3+yMSjl9/tIT38Zo0aN4uq0p5VW2f/cyvZrb1l7EuPJluH2xv7ka5z2vsLR5ntYXhOxvr4eKSkpCA4O5sp0dHTg4+ODpKQkpfskJSUhKChIoczX1xcnTpxQWr+uro77jwMAZWVNDSPl5eWt6tbU1ODWzZtoqH98K1qQnwdxWQV+TbzQZll76nToWAZ98EfBn1ydoqIC/FjXiIwTd7iyWzcLoCfRR1ZjDQDg3n2gLO4aHpQ9viqey7wLIwmQ+2d9qzKD6/e476W/bkZyfwuuznVZIQx7iWDfoiz/YQUGu/vAwaEp8XNycnAn+RdY9zV66n7tKXvy2G3JyckB7hegprahQ7FflxVieH8j1Dc0/d3k5Jcg8/BhXLx48Znne/I3Kot16dKlTz2OkZERRCLRU+uA8ej+/fsMADt//rxC+Zo1a5iHh4fSffT09NhXX32lUBYZGcksLS2V1g8NDWUA6EMf3j5lZWXPzAXeb001LTg4WOEKKpfLUVJSgn79+j37XymelJeXw9bWFvfu3YOxsTHf4XRYV45fE7EbGRk9sw6viWhubg5dXV0UFBQolBcUFMDa2lrpPtbW1h2qL5FIWj3rmJqaqh60FhkbG3e5/5Fb6srxazt2Xl9fiMViuLm5IS4ujiuTy+WIi4uDp6en0n08PT0V6gNAbGxsm/UJ6RLa+TinMUeOHGESiYRFR0ez69evs7feeouZmpqy/Px8xhhjCxcuZGvXruXqnzt3jvXq1Ytt27aNZWVlsdDQUKanp8cyMjL4+glqV1ZW1u5nCyHqyvHzFTvvicgYY7t27WJ2dnZMLBYzDw8PduHCBW6bl5cX8/PzU6j/9ddfs2HDhjGxWMyee+45dvLkSS1HrFm1tbUsNDSU1dbW8h2KSrpy/HzFLmKMMZ4vyoT0eILo4kZIT0eJSIgAUCISIgCUiIQIACUij3799VdMmzYN/fv3h0gkatVfljGGkJAQ2NjYwMDAAD4+Prh58yY/wT5BKpVi7NixMDIygqWlJWbOnIns7GyFOrW1tVi+fDn69euHPn36YM6cOa06Y/Bl9+7dGDVqFPfi3tPTEz///DO3XduxUyLyqKqqCi4uLoiMjFS6fcuWLdi5cyeioqJw8eJF9O7dG76+vqitrdVypK0lJCRg+fLluHDhAmJjY9HQ0IApU6agqurx8gSrVq3CDz/8gKNHjyIhIQEPHjzA7NmzeYz6sYEDByIsLAwpKSlITk7Giy++iBkzZiAzMxMAD7Fr9WUJaRMAdvz4ce67XC5n1tbWbOvWrVxZaWkpk0gk7PDhwzxE+HSFhYUMAEtISGCMNcWqp6fHjh49ytXJyspiAFhSUhJfYT5V37592d69e3mJna6IAnXnzh3k5+fDx8eHKzMxMcG4cePaHCLGp+bhZWZmTeMoU1JS0NDQoBC/k5MT7OzsBBd/Y2Mjjhw5gqqqKnh6evISe7cffdFV5ec3zcBmZWWlUG5lZcVtEwq5XI6VK1di/Pjx3CwJ+fn5EIvFrTrYCyn+jIwMeHp6ora2Fn369MHx48fh7OyM1NRUrcdOiUg6bfny5bh27RoSExP5DqVDhg8fjtTUVJSVleGbb76Bn58fEhISeImFbk0FqnlYV0eGfPEhMDAQP/74I86ePYuBAwdy5dbW1qivr0dpaalCfSHFLxaL4ejoCDc3N0ilUri4uGDHjh28xE6JKFCDBw+GtbW1wpCv8vJyXLx4URBDvhhjCAwMxPHjx3HmzBkMHjxYYbubmxv09PQU4s/OzoZMJhNE/MrI5XLU1dXxE7tGmoBIu1RUVLCrV6+yq1evMgAsPDycXb16ld29e5cxxlhYWBgzNTVl3333HUtPT2czZsxggwcPZjU1NTxHztg777zDTExMWHx8PMvLy+M+1dXVXJ23336b2dnZsTNnzrDk5GTm6enJPD09eYz6sbVr17KEhAR2584dlp6eztauXctEIhE7ffo0Y0z7sVMi8ujs2bNK5zhpHvYll8vZRx99xKysrJhEImHe3t4sOzub36D/n7K4AbD9+/dzdWpqatiyZctY3759maGhIZs1axbLy8vjL+gWFi9ezAYNGsTEYjGzsLBg3t7eXBIypv3YaRgUIQJAz4iECAAlIiECQIlIiABQIhIiAJSIhAgAJSIhAkCJSIgAUCISIgCUiIQIACUiUav6es0tc96dUSL2EN988w1GjhwJAwMD9OvXDz4+Ptz8Mvv27cNzzz0HiUQCGxsbBAYGcvvJZDLMmDEDffr0gbGxMV5//XWFoVnr16+Hq6sr9u7di8GDB3Mr6ZaWlmLJkiWwsLCAsbExXnzxRaSlpWn3R3chlIg9QF5eHubPn4/FixcjKysL8fHxmD17Nhhj2L17N5YvX4633noLGRkZ+P777+Ho6AigaVjQjBkzUFJSgoSEBMTGxiInJwfz5s1TOP6tW7fw7bff4tixY0hNTQUAvPbaaygsLMTPP/+MlJQUjBkzBt7e3igpKdH2z+8aNNadnAhGSkoKA8Byc3Nbbevfvz9bt26d0v1Onz7NdHV1mUwm48oyMzMZAHbp0iXGGONW4yosLOTq/Pbbb8zY2LjVQi5Dhgxhe/bsUcdP6nZoqowewMXFBd7e3hg5ciR8fX0xZcoUzJ07Fw0NDXjw4AG8vb2V7peVlQVbW1vY2tpyZc7OzjA1NUVWVhbGjh0LABg0aBAsLB6vWZ+WlobKykr069dP4Xg1NTW4ffu2Bn5h10eJ2APo6uoiNjYW58+fx+nTp7Fr1y6sW7eu1YKvqurdu7fC98rKStjY2CA+Pr5V3a6yWrO2USL2ECKRCOPHj8f48eMREhKCQYMGITY2Fvb29oiLi8PkyZNb7TNixAjcu3cP9+7d466K169fR2lpKZydnds815gxY5Cfn49evXrB3t5eUz+pW6FE7AEuXryIuLg4TJkyBZaWlrh48SKKioowYsQIrF+/Hm+//TYsLS3x0ksvoaKiAufOncO7774LHx8fjBw5EgsWLEBERAQePXqEZcuWwcvLC+7u7m2ez8fHB56enpg5cya2bNmCYcOG4cGDBzh58iRmzZr11H17LL4fUonmXb9+nfn6+jILCwsmkUjYsGHD2K5du7jtUVFRbPjw4UxPT4/Z2Niwd999l9t29+5dNn36dNa7d29mZGTEXnvtNW5ZdcaaGmtcXFxanbO8vJy9++67rH///kxPT4/Z2tqyBQsWKDT8kMdoqgxCBIDeIxIiAJSIhAgAJSIhAkCJSIgAUCISIgCUiIQIACUiIQJAiUiIAFAiEiIAlIiECAAlIiECQIlIiAD8H0tfstftWNSuAAAAAElFTkSuQmCC\n"
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# global AUPRC\n",
        "average_precision_score(V.label, V.score)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "REMKfC6MTXiO",
        "outputId": "1cbf068a-6b7c-4d6f-d865-ac30f347703f"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.3050859084098214"
            ]
          },
          "metadata": {},
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# AUPRC by chrom\n",
        "res_by_chrom = []\n",
        "for chrom in V.chrom.unique():\n",
        "    V_chrom = V[V.chrom == chrom]\n",
        "    res_by_chrom.append([chrom, len(V_chrom), average_precision_score(V_chrom.label, V_chrom.score)])\n",
        "res_by_chrom = pd.DataFrame(res_by_chrom, columns=[\"chrom\", \"n\", \"AUPRC\"])\n",
        "res_by_chrom"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 645
        },
        "id": "KC4NnIFPUbVe",
        "outputId": "27bf63e3-3b5d-4ee8-e046-94ae74314293"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "   chrom    n     AUPRC\n",
              "0      1  210  0.285684\n",
              "1      2  230  0.317321\n",
              "2      3  310  0.151283\n",
              "3      5   20  1.000000\n",
              "4      6   30  0.347619\n",
              "5      7  210  0.192152\n",
              "6      8   70  0.120540\n",
              "7      9  240  0.086625\n",
              "8     10  190  0.183487\n",
              "9     11  480  0.285151\n",
              "10    12   30  0.500000\n",
              "11    13  210  0.466352\n",
              "12    14   40  0.125368\n",
              "13    16   80  0.467061\n",
              "14    17   60  0.558547\n",
              "15    19  400  0.324191\n",
              "16    20   50  0.441667\n",
              "17    22   20  0.500000\n",
              "18     X  500  0.629526"
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-c6d01f56-dccd-4a4c-98be-e4428e260aff\" class=\"colab-df-container\">\n",
              "    <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>chrom</th>\n",
              "      <th>n</th>\n",
              "      <th>AUPRC</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>1</td>\n",
              "      <td>210</td>\n",
              "      <td>0.285684</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>2</td>\n",
              "      <td>230</td>\n",
              "      <td>0.317321</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>3</td>\n",
              "      <td>310</td>\n",
              "      <td>0.151283</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>5</td>\n",
              "      <td>20</td>\n",
              "      <td>1.000000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>6</td>\n",
              "      <td>30</td>\n",
              "      <td>0.347619</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>5</th>\n",
              "      <td>7</td>\n",
              "      <td>210</td>\n",
              "      <td>0.192152</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>6</th>\n",
              "      <td>8</td>\n",
              "      <td>70</td>\n",
              "      <td>0.120540</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>7</th>\n",
              "      <td>9</td>\n",
              "      <td>240</td>\n",
              "      <td>0.086625</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>8</th>\n",
              "      <td>10</td>\n",
              "      <td>190</td>\n",
              "      <td>0.183487</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>9</th>\n",
              "      <td>11</td>\n",
              "      <td>480</td>\n",
              "      <td>0.285151</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>10</th>\n",
              "      <td>12</td>\n",
              "      <td>30</td>\n",
              "      <td>0.500000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>11</th>\n",
              "      <td>13</td>\n",
              "      <td>210</td>\n",
              "      <td>0.466352</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>12</th>\n",
              "      <td>14</td>\n",
              "      <td>40</td>\n",
              "      <td>0.125368</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>13</th>\n",
              "      <td>16</td>\n",
              "      <td>80</td>\n",
              "      <td>0.467061</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>14</th>\n",
              "      <td>17</td>\n",
              "      <td>60</td>\n",
              "      <td>0.558547</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>15</th>\n",
              "      <td>19</td>\n",
              "      <td>400</td>\n",
              "      <td>0.324191</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>16</th>\n",
              "      <td>20</td>\n",
              "      <td>50</td>\n",
              "      <td>0.441667</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>17</th>\n",
              "      <td>22</td>\n",
              "      <td>20</td>\n",
              "      <td>0.500000</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>18</th>\n",
              "      <td>X</td>\n",
              "      <td>500</td>\n",
              "      <td>0.629526</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-c6d01f56-dccd-4a4c-98be-e4428e260aff')\"\n",
              "            title=\"Convert this dataframe to an interactive table.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "    <script>\n",
              "      const buttonEl =\n",
              "        document.querySelector('#df-c6d01f56-dccd-4a4c-98be-e4428e260aff button.colab-df-convert');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      async function convertToInteractive(key) {\n",
              "        const element = document.querySelector('#df-c6d01f56-dccd-4a4c-98be-e4428e260aff');\n",
              "        const dataTable =\n",
              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                    [key], {});\n",
              "        if (!dataTable) return;\n",
              "\n",
              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "          + ' to learn more about interactive tables.';\n",
              "        element.innerHTML = '';\n",
              "        dataTable['output_type'] = 'display_data';\n",
              "        await google.colab.output.renderOutput(dataTable, element);\n",
              "        const docLink = document.createElement('div');\n",
              "        docLink.innerHTML = docLinkHtml;\n",
              "        element.appendChild(docLink);\n",
              "      }\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "\n",
              "<div id=\"df-64a08cc9-3838-46c7-83bc-134a7a153504\">\n",
              "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-64a08cc9-3838-46c7-83bc-134a7a153504')\"\n",
              "            title=\"Suggest charts\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "     width=\"24px\">\n",
              "    <g>\n",
              "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
              "    </g>\n",
              "</svg>\n",
              "  </button>\n",
              "\n",
              "<style>\n",
              "  .colab-df-quickchart {\n",
              "      --bg-color: #E8F0FE;\n",
              "      --fill-color: #1967D2;\n",
              "      --hover-bg-color: #E2EBFA;\n",
              "      --hover-fill-color: #174EA6;\n",
              "      --disabled-fill-color: #AAA;\n",
              "      --disabled-bg-color: #DDD;\n",
              "  }\n",
              "\n",
              "  [theme=dark] .colab-df-quickchart {\n",
              "      --bg-color: #3B4455;\n",
              "      --fill-color: #D2E3FC;\n",
              "      --hover-bg-color: #434B5C;\n",
              "      --hover-fill-color: #FFFFFF;\n",
              "      --disabled-bg-color: #3B4455;\n",
              "      --disabled-fill-color: #666;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart {\n",
              "    background-color: var(--bg-color);\n",
              "    border: none;\n",
              "    border-radius: 50%;\n",
              "    cursor: pointer;\n",
              "    display: none;\n",
              "    fill: var(--fill-color);\n",
              "    height: 32px;\n",
              "    padding: 0;\n",
              "    width: 32px;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart:hover {\n",
              "    background-color: var(--hover-bg-color);\n",
              "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "    fill: var(--button-hover-fill-color);\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart-complete:disabled,\n",
              "  .colab-df-quickchart-complete:disabled:hover {\n",
              "    background-color: var(--disabled-bg-color);\n",
              "    fill: var(--disabled-fill-color);\n",
              "    box-shadow: none;\n",
              "  }\n",
              "\n",
              "  .colab-df-spinner {\n",
              "    border: 2px solid var(--fill-color);\n",
              "    border-color: transparent;\n",
              "    border-bottom-color: var(--fill-color);\n",
              "    animation:\n",
              "      spin 1s steps(1) infinite;\n",
              "  }\n",
              "\n",
              "  @keyframes spin {\n",
              "    0% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "      border-left-color: var(--fill-color);\n",
              "    }\n",
              "    20% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    30% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    40% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    60% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    80% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "    90% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "  }\n",
              "</style>\n",
              "\n",
              "  <script>\n",
              "    async function quickchart(key) {\n",
              "      const quickchartButtonEl =\n",
              "        document.querySelector('#' + key + ' button');\n",
              "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
              "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
              "      try {\n",
              "        const charts = await google.colab.kernel.invokeFunction(\n",
              "            'suggestCharts', [key], {});\n",
              "      } catch (error) {\n",
              "        console.error('Error during call to suggestCharts:', error);\n",
              "      }\n",
              "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
              "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
              "    }\n",
              "    (() => {\n",
              "      let quickchartButtonEl =\n",
              "        document.querySelector('#df-64a08cc9-3838-46c7-83bc-134a7a153504 button');\n",
              "      quickchartButtonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "    })();\n",
              "  </script>\n",
              "</div>\n",
              "\n",
              "  <div id=\"id_cea20d60-6b19-40da-9258-c54ba82ea648\">\n",
              "    <style>\n",
              "      .colab-df-generate {\n",
              "        background-color: #E8F0FE;\n",
              "        border: none;\n",
              "        border-radius: 50%;\n",
              "        cursor: pointer;\n",
              "        display: none;\n",
              "        fill: #1967D2;\n",
              "        height: 32px;\n",
              "        padding: 0 0 0 0;\n",
              "        width: 32px;\n",
              "      }\n",
              "\n",
              "      .colab-df-generate:hover {\n",
              "        background-color: #E2EBFA;\n",
              "        box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "        fill: #174EA6;\n",
              "      }\n",
              "\n",
              "      [theme=dark] .colab-df-generate {\n",
              "        background-color: #3B4455;\n",
              "        fill: #D2E3FC;\n",
              "      }\n",
              "\n",
              "      [theme=dark] .colab-df-generate:hover {\n",
              "        background-color: #434B5C;\n",
              "        box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "        filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "        fill: #FFFFFF;\n",
              "      }\n",
              "    </style>\n",
              "    <button class=\"colab-df-generate\" onclick=\"generateWithVariable('res_by_chrom')\"\n",
              "            title=\"Generate code using this dataframe.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "       width=\"24px\">\n",
              "    <path d=\"M7,19H8.4L18.45,9,17,7.55,7,17.6ZM5,21V16.75L18.45,3.32a2,2,0,0,1,2.83,0l1.4,1.43a1.91,1.91,0,0,1,.58,1.4,1.91,1.91,0,0,1-.58,1.4L9.25,21ZM18.45,9,17,7.55Zm-12,3A5.31,5.31,0,0,0,4.9,8.1,5.31,5.31,0,0,0,1,6.5,5.31,5.31,0,0,0,4.9,4.9,5.31,5.31,0,0,0,6.5,1,5.31,5.31,0,0,0,8.1,4.9,5.31,5.31,0,0,0,12,6.5,5.46,5.46,0,0,0,6.5,12Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "    <script>\n",
              "      (() => {\n",
              "      const buttonEl =\n",
              "        document.querySelector('#id_cea20d60-6b19-40da-9258-c54ba82ea648 button.colab-df-generate');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      buttonEl.onclick = () => {\n",
              "        google.colab.notebook.generateWithVariable('res_by_chrom');\n",
              "      }\n",
              "      })();\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "    </div>\n",
              "  </div>\n"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "dataframe",
              "variable_name": "res_by_chrom",
              "summary": "{\n  \"name\": \"res_by_chrom\",\n  \"rows\": 19,\n  \"fields\": [\n    {\n      \"column\": \"chrom\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 19,\n        \"samples\": [\n          \"1\",\n          \"7\",\n          \"13\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"n\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 155,\n        \"min\": 20,\n        \"max\": 500,\n        \"num_unique_values\": 15,\n        \"samples\": [\n          40,\n          60,\n          210\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"AUPRC\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 0.22281536138652086,\n        \"min\": 0.08662464205313769,\n        \"max\": 1.0,\n        \"num_unique_values\": 18,\n        \"samples\": [\n          0.28568421491122165,\n          0.3173206756653251,\n          0.18348692401355995\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}"
            }
          },
          "metadata": {},
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Weighted average\n",
        "\n",
        "def stat(df):\n",
        "    weight = df[\"n\"] / df[\"n\"].sum()\n",
        "    return (df[\"AUPRC\"] * weight).sum()\n",
        "\n",
        "stat(res_by_chrom)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8B9kgGTlVDHX",
        "outputId": "80ece55e-4180-49c8-f635-24cb765dede5"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.3304637185530955"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def bootstrap_se(df, stat, n_bootstraps=1000):\n",
        "    df = pl.DataFrame(df)\n",
        "    return (\n",
        "        pl.Series([\n",
        "            stat(df.sample(len(df), with_replacement=True, seed=i))\n",
        "            for i in range(n_bootstraps)]\n",
        "        )\n",
        "        .std()\n",
        "    )\n",
        "\n",
        "bootstrap_se(res_by_chrom, stat)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "chbtzAYSVkag",
        "outputId": "5477c618-e026-4258-f3c0-a18d5b3406eb"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.05398842833086511"
            ]
          },
          "metadata": {},
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# compare with official results (which use RC averaging)\n",
        "pd.read_csv(\"hf://datasets/songlab/TraitGym/mendelian_traits_matched_9/AUPRC_by_chrom_weighted_average/all/GPN_final_EuclideanDistance.plus.score.csv\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 81
        },
        "id": "HgvdWsZkV7sR",
        "outputId": "d7538db1-0234-4387-f1fe-7551a7c53db7"
      },
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "                                    model metric     score        se\n",
              "0  GPN_final_EuclideanDistance.plus.score  AUPRC  0.346963  0.056745"
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-389a580f-0541-4778-9141-6f8d4383fa2b\" class=\"colab-df-container\">\n",
              "    <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>model</th>\n",
              "      <th>metric</th>\n",
              "      <th>score</th>\n",
              "      <th>se</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>GPN_final_EuclideanDistance.plus.score</td>\n",
              "      <td>AUPRC</td>\n",
              "      <td>0.346963</td>\n",
              "      <td>0.056745</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-389a580f-0541-4778-9141-6f8d4383fa2b')\"\n",
              "            title=\"Convert this dataframe to an interactive table.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "    <script>\n",
              "      const buttonEl =\n",
              "        document.querySelector('#df-389a580f-0541-4778-9141-6f8d4383fa2b button.colab-df-convert');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      async function convertToInteractive(key) {\n",
              "        const element = document.querySelector('#df-389a580f-0541-4778-9141-6f8d4383fa2b');\n",
              "        const dataTable =\n",
              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                    [key], {});\n",
              "        if (!dataTable) return;\n",
              "\n",
              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "          + ' to learn more about interactive tables.';\n",
              "        element.innerHTML = '';\n",
              "        dataTable['output_type'] = 'display_data';\n",
              "        await google.colab.output.renderOutput(dataTable, element);\n",
              "        const docLink = document.createElement('div');\n",
              "        docLink.innerHTML = docLinkHtml;\n",
              "        element.appendChild(docLink);\n",
              "      }\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "\n",
              "    </div>\n",
              "  </div>\n"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "dataframe",
              "summary": "{\n  \"name\": \"pd\",\n  \"rows\": 1,\n  \"fields\": [\n    {\n      \"column\": \"model\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 1,\n        \"samples\": [\n          \"GPN_final_EuclideanDistance.plus.score\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"metric\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 1,\n        \"samples\": [\n          \"AUPRC\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"score\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": null,\n        \"min\": 0.3469633009736218,\n        \"max\": 0.3469633009736218,\n        \"num_unique_values\": 1,\n        \"samples\": [\n          0.3469633009736218\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"se\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": null,\n        \"min\": 0.056745382840026,\n        \"max\": 0.056745382840026,\n        \"num_unique_values\": 1,\n        \"samples\": [\n          0.056745382840026\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}"
            }
          },
          "metadata": {},
          "execution_count": 22
        }
      ]
    }
  ]
}