test3_jax_multiple_machine_tpu.ipynb

{
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.10.13",
      "mimetype": "text/x-python",
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "pygments_lexer": "ipython3",
      "nbconvert_exporter": "python",
      "file_extension": ".py"
    },
    "kaggle": {
      "accelerator": "nvidiaTeslaT4",
      "dataSources": [],
      "dockerImageVersionId": 30716,
      "isInternetEnabled": true,
      "language": "python",
      "sourceType": "notebook",
      "isGpuEnabled": true
    },
    "colab": {
      "provenance": [],
      "gpuType": "V28"
    },
    "accelerator": "TPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "e263afc2c9fe4a1583a5a5c2b845a50f": {
          "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_f3aff971ae2d43cf872a0d6bbbf8a392",
              "IPY_MODEL_8b756212a3dd4e169012bf2d62770ccd",
              "IPY_MODEL_0cdbbf0e98e7411bb562f22e2436888a"
            ],
            "layout": "IPY_MODEL_808a2bddd1814f7aa90d751880375c52"
          }
        },
        "f3aff971ae2d43cf872a0d6bbbf8a392": {
          "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_4621f2282c484f2b8ae16b70cc73ddcb",
            "placeholder": "​",
            "style": "IPY_MODEL_f2f53f150d77481c81df52b26c55819f",
            "value": "Downloading readme: 100%"
          }
        },
        "8b756212a3dd4e169012bf2d62770ccd": {
          "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_8a4c5272a7e74bc5ae31b4706ce444ad",
            "max": 28064,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_04ac49562024426c914a7f5a214a84f2",
            "value": 28064
          }
        },
        "0cdbbf0e98e7411bb562f22e2436888a": {
          "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_a34604ca1828446b8ea7ce216829ab99",
            "placeholder": "​",
            "style": "IPY_MODEL_07e57c3c4601459086d23b71d0cff4b0",
            "value": " 28.1k/28.1k [00:00&lt;00:00, 2.45MB/s]"
          }
        },
        "808a2bddd1814f7aa90d751880375c52": {
          "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
          }
        },
        "4621f2282c484f2b8ae16b70cc73ddcb": {
          "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
          }
        },
        "f2f53f150d77481c81df52b26c55819f": {
          "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": ""
          }
        },
        "8a4c5272a7e74bc5ae31b4706ce444ad": {
          "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
          }
        },
        "04ac49562024426c914a7f5a214a84f2": {
          "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": ""
          }
        },
        "a34604ca1828446b8ea7ce216829ab99": {
          "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
          }
        },
        "07e57c3c4601459086d23b71d0cff4b0": {
          "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": ""
          }
        },
        "1ba66d0b45574a778eda36d1f41d07c5": {
          "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_f48e45e5402940df80ff9f46645d97bd",
              "IPY_MODEL_704c96aa0cd9440f9bb1b5d5017e8309",
              "IPY_MODEL_7134f384e6384a5f9d95426efb511738"
            ],
            "layout": "IPY_MODEL_b7b2053c16a3442ea95a34ba6a6d9232"
          }
        },
        "f48e45e5402940df80ff9f46645d97bd": {
          "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_05724b4292044397bb320ebdea2a7de7",
            "placeholder": "​",
            "style": "IPY_MODEL_4552e08385344020923c52fe653c5dbd",
            "value": "Downloading data: 100%"
          }
        },
        "704c96aa0cd9440f9bb1b5d5017e8309": {
          "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_fb54ccc429ed4ad79f3b0aa8896abe06",
            "max": 23065198,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_e7f161c57eac4cd0a1b59d7dbe2764c4",
            "value": 23065198
          }
        },
        "7134f384e6384a5f9d95426efb511738": {
          "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_7c8b32d8b924415cae6e422209d32f3c",
            "placeholder": "​",
            "style": "IPY_MODEL_f9398ec1118745949bb012f97239ce44",
            "value": " 23.1M/23.1M [00:00&lt;00:00, 67.1MB/s]"
          }
        },
        "b7b2053c16a3442ea95a34ba6a6d9232": {
          "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
          }
        },
        "05724b4292044397bb320ebdea2a7de7": {
          "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
          }
        },
        "4552e08385344020923c52fe653c5dbd": {
          "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": ""
          }
        },
        "fb54ccc429ed4ad79f3b0aa8896abe06": {
          "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
          }
        },
        "e7f161c57eac4cd0a1b59d7dbe2764c4": {
          "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": ""
          }
        },
        "7c8b32d8b924415cae6e422209d32f3c": {
          "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
          }
        },
        "f9398ec1118745949bb012f97239ce44": {
          "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": ""
          }
        },
        "09b9ab71b57649818b331ac69b9eaad2": {
          "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_72a86b2670bb4859b6cd1403d3241b5e",
              "IPY_MODEL_a37f09b7daf24e4c82f5dc7d2494cfc3",
              "IPY_MODEL_cf74ec8b71a441828654076db4cc6ffd"
            ],
            "layout": "IPY_MODEL_eefde40d5e174570b38d22672c374fc6"
          }
        },
        "72a86b2670bb4859b6cd1403d3241b5e": {
          "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_9ffc7ae56b0347f4bcb44dfa103f277e",
            "placeholder": "​",
            "style": "IPY_MODEL_653050d956124b56ae4caf67d635269a",
            "value": "Generating train split: 100%"
          }
        },
        "a37f09b7daf24e4c82f5dc7d2494cfc3": {
          "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_b3ce06aa2e8542f0b3313ce890832807",
            "max": 137151,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_84a959004d7344a99a004d54531bcf1a",
            "value": 137151
          }
        },
        "cf74ec8b71a441828654076db4cc6ffd": {
          "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_cba26d225fd34dd792718309f4e1c1d1",
            "placeholder": "​",
            "style": "IPY_MODEL_9a915ab5ab854c27bc5fbcbded6e6927",
            "value": " 137151/137151 [00:00&lt;00:00, 584908.62 examples/s]"
          }
        },
        "eefde40d5e174570b38d22672c374fc6": {
          "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
          }
        },
        "9ffc7ae56b0347f4bcb44dfa103f277e": {
          "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
          }
        },
        "653050d956124b56ae4caf67d635269a": {
          "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": ""
          }
        },
        "b3ce06aa2e8542f0b3313ce890832807": {
          "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
          }
        },
        "84a959004d7344a99a004d54531bcf1a": {
          "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": ""
          }
        },
        "cba26d225fd34dd792718309f4e1c1d1": {
          "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
          }
        },
        "9a915ab5ab854c27bc5fbcbded6e6927": {
          "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": ""
          }
        },
        "0f894258311c42599845b808ee59922a": {
          "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_7136464b8afb43c2a259da17778aee50",
              "IPY_MODEL_5496d0a320804de5990bbcc3d493e8d1",
              "IPY_MODEL_c6b19c931c1f4f29b433628b157a7bdc"
            ],
            "layout": "IPY_MODEL_114e7e6c4864433eaa73d9732410cc79"
          }
        },
        "7136464b8afb43c2a259da17778aee50": {
          "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_b0e347980dff4e0481b3f6d2ef48edc6",
            "placeholder": "​",
            "style": "IPY_MODEL_027edd6a9b8b4213841f38d2a965f931",
            "value": "tokenizer_config.json: 100%"
          }
        },
        "5496d0a320804de5990bbcc3d493e8d1": {
          "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_c449eea4ca344e3f8a8a56023ef0c7f0",
            "max": 42,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_40a43dad986c4fb1859417b4eb29ab50",
            "value": 42
          }
        },
        "c6b19c931c1f4f29b433628b157a7bdc": {
          "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_697dfce9641e4120b7b0ef7963ea293a",
            "placeholder": "​",
            "style": "IPY_MODEL_532a301e893a44bdb3a893f553174738",
            "value": " 42.0/42.0 [00:00&lt;00:00, 4.07kB/s]"
          }
        },
        "114e7e6c4864433eaa73d9732410cc79": {
          "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
          }
        },
        "b0e347980dff4e0481b3f6d2ef48edc6": {
          "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
          }
        },
        "027edd6a9b8b4213841f38d2a965f931": {
          "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": ""
          }
        },
        "c449eea4ca344e3f8a8a56023ef0c7f0": {
          "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
          }
        },
        "40a43dad986c4fb1859417b4eb29ab50": {
          "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": ""
          }
        },
        "697dfce9641e4120b7b0ef7963ea293a": {
          "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
          }
        },
        "532a301e893a44bdb3a893f553174738": {
          "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": ""
          }
        },
        "4f808ac47f8442b0ae85c50ed5c3d413": {
          "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_6e6bc82a81774489977a9fac31de042f",
              "IPY_MODEL_61633f6fd20c42fb86720c62a98c41d4",
              "IPY_MODEL_4ed93dfd117e4aadb56280f6e805acba"
            ],
            "layout": "IPY_MODEL_7f32f52c085847b982635a251846232c"
          }
        },
        "6e6bc82a81774489977a9fac31de042f": {
          "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_3d1727222e794c95bb6f2488ef8b5c37",
            "placeholder": "​",
            "style": "IPY_MODEL_044d409b3cf740c8b15e03184a6f2c2f",
            "value": "source.spm: 100%"
          }
        },
        "61633f6fd20c42fb86720c62a98c41d4": {
          "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_d420de1bee8e41fbab51cb3441f5b853",
            "max": 792289,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_b2b111f54336469bb4614c039e046ca6",
            "value": 792289
          }
        },
        "4ed93dfd117e4aadb56280f6e805acba": {
          "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_e7dd99cbbde74113ad88ee74d555f61a",
            "placeholder": "​",
            "style": "IPY_MODEL_adda0e06c2f64ea79795981b50d2fe0e",
            "value": " 792k/792k [00:00&lt;00:00, 18.0MB/s]"
          }
        },
        "7f32f52c085847b982635a251846232c": {
          "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
          }
        },
        "3d1727222e794c95bb6f2488ef8b5c37": {
          "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
          }
        },
        "044d409b3cf740c8b15e03184a6f2c2f": {
          "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": ""
          }
        },
        "d420de1bee8e41fbab51cb3441f5b853": {
          "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
          }
        },
        "b2b111f54336469bb4614c039e046ca6": {
          "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": ""
          }
        },
        "e7dd99cbbde74113ad88ee74d555f61a": {
          "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
          }
        },
        "adda0e06c2f64ea79795981b50d2fe0e": {
          "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": ""
          }
        },
        "8689ec2286e5414d989fa3f223484ae7": {
          "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_144b417137674551a895482d147a8249",
              "IPY_MODEL_2de396a5e6ca4a10af298a2ee61b2d37",
              "IPY_MODEL_fffa7984d49c4c38accf7311236c8152"
            ],
            "layout": "IPY_MODEL_2da4eed1a1204bed8e6b4c87e3bcdc3e"
          }
        },
        "144b417137674551a895482d147a8249": {
          "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_1a112d908efa4ae082274ee21fffaee0",
            "placeholder": "​",
            "style": "IPY_MODEL_13ea6f7789864999bed305937a18a359",
            "value": "target.spm: 100%"
          }
        },
        "2de396a5e6ca4a10af298a2ee61b2d37": {
          "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_eee3c0d438044662980440d112f2d3ba",
            "max": 849914,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_49dad883ecab4bad983167a32d55da58",
            "value": 849914
          }
        },
        "fffa7984d49c4c38accf7311236c8152": {
          "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_2f22a464dcf1477e8f1f8d39ff1ebc4d",
            "placeholder": "​",
            "style": "IPY_MODEL_aea3e71ef7e44d56941a1682aa5e95b9",
            "value": " 850k/850k [00:00&lt;00:00, 12.9MB/s]"
          }
        },
        "2da4eed1a1204bed8e6b4c87e3bcdc3e": {
          "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
          }
        },
        "1a112d908efa4ae082274ee21fffaee0": {
          "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
          }
        },
        "13ea6f7789864999bed305937a18a359": {
          "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": ""
          }
        },
        "eee3c0d438044662980440d112f2d3ba": {
          "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
          }
        },
        "49dad883ecab4bad983167a32d55da58": {
          "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": ""
          }
        },
        "2f22a464dcf1477e8f1f8d39ff1ebc4d": {
          "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
          }
        },
        "aea3e71ef7e44d56941a1682aa5e95b9": {
          "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": ""
          }
        },
        "2c0557789f39449b83937fa970c735b7": {
          "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_8f449f75e1d94e5e843c0c77bcd92446",
              "IPY_MODEL_0bda9d055aa84c04aa1ae75122d302e1",
              "IPY_MODEL_b4d214644cd3453481a813863ca6a8ae"
            ],
            "layout": "IPY_MODEL_db25847587334e64ad383464705b9e0c"
          }
        },
        "8f449f75e1d94e5e843c0c77bcd92446": {
          "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_3bf041ab1dc3475da1ef23a1e23d50ab",
            "placeholder": "​",
            "style": "IPY_MODEL_84b88cb3f2b24e4f9d01d5ecbb0db169",
            "value": "vocab.json: 100%"
          }
        },
        "0bda9d055aa84c04aa1ae75122d302e1": {
          "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_d377b0f3c85042acaa3e8d894007293c",
            "max": 1565798,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_9f03d4c5db7e470180972e669da12574",
            "value": 1565798
          }
        },
        "b4d214644cd3453481a813863ca6a8ae": {
          "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_ee9420c3efe74f419eda70edd28a9ec4",
            "placeholder": "​",
            "style": "IPY_MODEL_143d3eb4d22b4d81a3108b5743ff334b",
            "value": " 1.57M/1.57M [00:00&lt;00:00, 17.4MB/s]"
          }
        },
        "db25847587334e64ad383464705b9e0c": {
          "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
          }
        },
        "3bf041ab1dc3475da1ef23a1e23d50ab": {
          "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
          }
        },
        "84b88cb3f2b24e4f9d01d5ecbb0db169": {
          "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": ""
          }
        },
        "d377b0f3c85042acaa3e8d894007293c": {
          "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
          }
        },
        "9f03d4c5db7e470180972e669da12574": {
          "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": ""
          }
        },
        "ee9420c3efe74f419eda70edd28a9ec4": {
          "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
          }
        },
        "143d3eb4d22b4d81a3108b5743ff334b": {
          "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": ""
          }
        },
        "88162af0d5f3467983e9724e195c0bff": {
          "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_c6583e04b06c4bccbe5a4678fa484c7c",
              "IPY_MODEL_a339ffda3dbf46cfbc78a08f7ecf158b",
              "IPY_MODEL_8e19f83d25ce4e81a728cdf054790065"
            ],
            "layout": "IPY_MODEL_425cd5b013394e399dc2298e9682d1d8"
          }
        },
        "c6583e04b06c4bccbe5a4678fa484c7c": {
          "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_365ac66278b3400a842538087c82a28c",
            "placeholder": "​",
            "style": "IPY_MODEL_dd13ef35ec3d451e8ad65eb8ac9e1711",
            "value": "config.json: 100%"
          }
        },
        "a339ffda3dbf46cfbc78a08f7ecf158b": {
          "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_20870934ea084da49fd0a18ddd457f60",
            "max": 1381,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_d9c5d6f2331047be9a836918d7f2a82a",
            "value": 1381
          }
        },
        "8e19f83d25ce4e81a728cdf054790065": {
          "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_b81cc242a2bd4cd5b1808bb20258b9b9",
            "placeholder": "​",
            "style": "IPY_MODEL_bf95a897d59f459d883fcdd93bb36fa7",
            "value": " 1.38k/1.38k [00:00&lt;00:00, 148kB/s]"
          }
        },
        "425cd5b013394e399dc2298e9682d1d8": {
          "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
          }
        },
        "365ac66278b3400a842538087c82a28c": {
          "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
          }
        },
        "dd13ef35ec3d451e8ad65eb8ac9e1711": {
          "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": ""
          }
        },
        "20870934ea084da49fd0a18ddd457f60": {
          "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
          }
        },
        "d9c5d6f2331047be9a836918d7f2a82a": {
          "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": ""
          }
        },
        "b81cc242a2bd4cd5b1808bb20258b9b9": {
          "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
          }
        },
        "bf95a897d59f459d883fcdd93bb36fa7": {
          "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": ""
          }
        },
        "682773a903be4be393c205216798fb76": {
          "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_3429a1a9c7714781a6b29f7657a42cd1",
              "IPY_MODEL_bb6341d0e3b7442791900959c73ee2e8",
              "IPY_MODEL_e4b47402c556494b98d37f595002e92a"
            ],
            "layout": "IPY_MODEL_44608d1b30d94a70aea4d2ad0cdf11f9"
          }
        },
        "3429a1a9c7714781a6b29f7657a42cd1": {
          "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_ec8cb16c18d543ff9f3f904ce0817610",
            "placeholder": "​",
            "style": "IPY_MODEL_7c5a94f7bb0f4b0eb1caf122df3dae1f",
            "value": "pytorch_model.bin: 100%"
          }
        },
        "bb6341d0e3b7442791900959c73ee2e8": {
          "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_0f3ad835d88a4393a481db2fd190060f",
            "max": 307000101,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_4b038c5622e24e20ae41bb7334e0a771",
            "value": 307000101
          }
        },
        "e4b47402c556494b98d37f595002e92a": {
          "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_a8f6b50aa87c4165b35cb63b015f718e",
            "placeholder": "​",
            "style": "IPY_MODEL_11db1eb4bca245d7aae36ea5e9bcceb5",
            "value": " 307M/307M [00:01&lt;00:00, 202MB/s]"
          }
        },
        "44608d1b30d94a70aea4d2ad0cdf11f9": {
          "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
          }
        },
        "ec8cb16c18d543ff9f3f904ce0817610": {
          "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
          }
        },
        "7c5a94f7bb0f4b0eb1caf122df3dae1f": {
          "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": ""
          }
        },
        "0f3ad835d88a4393a481db2fd190060f": {
          "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
          }
        },
        "4b038c5622e24e20ae41bb7334e0a771": {
          "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": ""
          }
        },
        "a8f6b50aa87c4165b35cb63b015f718e": {
          "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
          }
        },
        "11db1eb4bca245d7aae36ea5e9bcceb5": {
          "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": ""
          }
        },
        "dc6f89f1d1e848a2adf3dc07a549bc95": {
          "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_bc43a390fa2c4c8faf679b23fdc2a564",
              "IPY_MODEL_453e46b83ab94d4ea8e47d63c27c0fca",
              "IPY_MODEL_10e15538afbe4ce3a811280c43ac5755"
            ],
            "layout": "IPY_MODEL_78b5cae51e8c478db152d3f61dc38081"
          }
        },
        "bc43a390fa2c4c8faf679b23fdc2a564": {
          "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_2e66e571995346869b5d9a81072f45d5",
            "placeholder": "​",
            "style": "IPY_MODEL_00e26680684943339bcd02e84e2773cc",
            "value": "generation_config.json: 100%"
          }
        },
        "453e46b83ab94d4ea8e47d63c27c0fca": {
          "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_df7174fd535e433091faf53a54b6426b",
            "max": 293,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_26441d4b8330450094506f2a5bc89743",
            "value": 293
          }
        },
        "10e15538afbe4ce3a811280c43ac5755": {
          "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_751dfe82fd5348188e9870a6c430ea4d",
            "placeholder": "​",
            "style": "IPY_MODEL_b12f0e2f4674457f9b26cfae802c55be",
            "value": " 293/293 [00:00&lt;00:00, 29.3kB/s]"
          }
        },
        "78b5cae51e8c478db152d3f61dc38081": {
          "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
          }
        },
        "2e66e571995346869b5d9a81072f45d5": {
          "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
          }
        },
        "00e26680684943339bcd02e84e2773cc": {
          "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": ""
          }
        },
        "df7174fd535e433091faf53a54b6426b": {
          "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
          }
        },
        "26441d4b8330450094506f2a5bc89743": {
          "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": ""
          }
        },
        "751dfe82fd5348188e9870a6c430ea4d": {
          "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
          }
        },
        "b12f0e2f4674457f9b26cfae802c55be": {
          "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": ""
          }
        },
        "be5ac7d6a07d4128863a7ab279c8024a": {
          "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_3a587589d98c475e8504f22080b42694",
              "IPY_MODEL_773646b21ccd4cfcb4549900d507fd94",
              "IPY_MODEL_034d59af89444a53bfb94c853c37d6a8"
            ],
            "layout": "IPY_MODEL_df14e319dd704b0f860f1a9182922e64"
          }
        },
        "3a587589d98c475e8504f22080b42694": {
          "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_4414556419b84a448d5c99281a033e0a",
            "placeholder": "​",
            "style": "IPY_MODEL_3e9d1b51fdb14544a52a09d2539ad3c7",
            "value": "Map: 100%"
          }
        },
        "773646b21ccd4cfcb4549900d507fd94": {
          "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_f697b7d1861e4683a1943f9a92ae3f0e",
            "max": 109720,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_91d54976ce7c42bcbed276f0707b78e1",
            "value": 109720
          }
        },
        "034d59af89444a53bfb94c853c37d6a8": {
          "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_dc08fcfee29340c6afe258c499682ca0",
            "placeholder": "​",
            "style": "IPY_MODEL_aac4d7efafae4a879cf67d7101163ba5",
            "value": " 109720/109720 [00:32&lt;00:00, 3512.29 examples/s]"
          }
        },
        "df14e319dd704b0f860f1a9182922e64": {
          "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
          }
        },
        "4414556419b84a448d5c99281a033e0a": {
          "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
          }
        },
        "3e9d1b51fdb14544a52a09d2539ad3c7": {
          "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": ""
          }
        },
        "f697b7d1861e4683a1943f9a92ae3f0e": {
          "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
          }
        },
        "91d54976ce7c42bcbed276f0707b78e1": {
          "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": ""
          }
        },
        "dc08fcfee29340c6afe258c499682ca0": {
          "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
          }
        },
        "aac4d7efafae4a879cf67d7101163ba5": {
          "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": ""
          }
        },
        "690df136207644daaf42380f181c7ced": {
          "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_92f8dc6b5a114972b1c5ae821845832d",
              "IPY_MODEL_53d9d3750940472cb77e60d206de54f3",
              "IPY_MODEL_9de6043ecc6d469f88e798f578d020bb"
            ],
            "layout": "IPY_MODEL_d9effe34749c4948a2098698e1dc1f88"
          }
        },
        "92f8dc6b5a114972b1c5ae821845832d": {
          "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_481e16c45c2d477db4689891ede0b1d4",
            "placeholder": "​",
            "style": "IPY_MODEL_a495a7c59e7d4771bd6ecdb1235cab03",
            "value": "Map: 100%"
          }
        },
        "53d9d3750940472cb77e60d206de54f3": {
          "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_33e8b0cf32b44e3f93d6cdf403b3966b",
            "max": 13716,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_6441dfea809648d0a6707d03ddd3909a",
            "value": 13716
          }
        },
        "9de6043ecc6d469f88e798f578d020bb": {
          "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_2c1e043eb32a45c0a814bf8268a70642",
            "placeholder": "​",
            "style": "IPY_MODEL_e8a97c273fc943e497390e3052892ae1",
            "value": " 13716/13716 [00:03&lt;00:00, 3500.16 examples/s]"
          }
        },
        "d9effe34749c4948a2098698e1dc1f88": {
          "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
          }
        },
        "481e16c45c2d477db4689891ede0b1d4": {
          "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
          }
        },
        "a495a7c59e7d4771bd6ecdb1235cab03": {
          "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": ""
          }
        },
        "33e8b0cf32b44e3f93d6cdf403b3966b": {
          "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
          }
        },
        "6441dfea809648d0a6707d03ddd3909a": {
          "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": ""
          }
        },
        "2c1e043eb32a45c0a814bf8268a70642": {
          "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
          }
        },
        "e8a97c273fc943e497390e3052892ae1": {
          "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": ""
          }
        },
        "5da90955287e42cc83ed1bbba5d12e11": {
          "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_d43bfd54a13644869754525c046802e4",
              "IPY_MODEL_cdbc5fa699f5457eb0843a565e7881f7",
              "IPY_MODEL_6848df35d3eb4814bacff411f8a473ac"
            ],
            "layout": "IPY_MODEL_43a0a2392e7d48bfb8321c3c8e8ffed6"
          }
        },
        "d43bfd54a13644869754525c046802e4": {
          "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_ad6cadadebf946a6824106e1796615ea",
            "placeholder": "​",
            "style": "IPY_MODEL_a50d44c51be84b4897ae94d0313c1ee0",
            "value": "Map: 100%"
          }
        },
        "cdbc5fa699f5457eb0843a565e7881f7": {
          "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_5d08d65ea1fc41eab5a98d3807f6aa91",
            "max": 13715,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_2f99c6a66a6146c0923fe8ee107f9281",
            "value": 13715
          }
        },
        "6848df35d3eb4814bacff411f8a473ac": {
          "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_a8fb09476729415794e89c6fe910514d",
            "placeholder": "​",
            "style": "IPY_MODEL_6879f59127654029b1d9d93d876d891c",
            "value": " 13715/13715 [00:04&lt;00:00, 3413.18 examples/s]"
          }
        },
        "43a0a2392e7d48bfb8321c3c8e8ffed6": {
          "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
          }
        },
        "ad6cadadebf946a6824106e1796615ea": {
          "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
          }
        },
        "a50d44c51be84b4897ae94d0313c1ee0": {
          "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": ""
          }
        },
        "5d08d65ea1fc41eab5a98d3807f6aa91": {
          "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
          }
        },
        "2f99c6a66a6146c0923fe8ee107f9281": {
          "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": ""
          }
        },
        "a8fb09476729415794e89c6fe910514d": {
          "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
          }
        },
        "6879f59127654029b1d9d93d876d891c": {
          "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": ""
          }
        }
      }
    }
  },
  "nbformat_minor": 0,
  "nbformat": 4,
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!pip install --upgrade pip"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:23:10.048531Z",
          "iopub.execute_input": "2024-06-04T00:23:10.049339Z",
          "iopub.status.idle": "2024-06-04T00:23:24.183486Z",
          "shell.execute_reply.started": "2024-06-04T00:23:10.049302Z",
          "shell.execute_reply": "2024-06-04T00:23:24.182312Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "67SrDdUevSPy",
        "outputId": "ebaa875d-4c81-4d09-cd1d-8370096bbdae"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: pip in /usr/local/lib/python3.10/dist-packages (23.1.2)\n",
            "Collecting pip\n",
            "  Downloading pip-24.0-py3-none-any.whl (2.1 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.1/2.1 MB\u001b[0m \u001b[31m8.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: pip\n",
            "  Attempting uninstall: pip\n",
            "    Found existing installation: pip 23.1.2\n",
            "    Uninstalling pip-23.1.2:\n",
            "      Successfully uninstalled pip-23.1.2\n",
            "Successfully installed pip-24.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html -q"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "i9xC3btLxIDZ",
        "outputId": "96fa7ee1-1f86-4f84-d31b-6a8db5dc34c9"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
            "\u001b[0m"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "print('JAX is running on', jax.lib.xla_bridge.get_backend().platform)\n",
        "print(jax.devices())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tR8zzMVWxuy-",
        "outputId": "ef4bdc77-05af-4435-f650-692801f1e589"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "JAX is running on tpu\n",
            "[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1), TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1), TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1), TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "!pip install datasets\n",
        "!pip install git+https://github.com/huggingface/transformers.git\n",
        "\n",
        "!pip install flax\n",
        "!pip install git+https://github.com/deepmind/optax.git"
      ],
      "metadata": {
        "id": "rj0U3jaSvqZT"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install --upgrade jax jaxlib optax jax-dataloader -q"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:25:22.875137Z",
          "iopub.execute_input": "2024-06-04T00:25:22.875615Z",
          "iopub.status.idle": "2024-06-04T00:25:36.934338Z",
          "shell.execute_reply.started": "2024-06-04T00:25:22.875569Z",
          "shell.execute_reply": "2024-06-04T00:25:36.933150Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Xo3pXIATvSP1",
        "outputId": "d610f14e-1807-4397-cc33-eca23f6ded8b"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.9/1.9 MB\u001b[0m \u001b[31m9.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m77.5/77.5 MB\u001b[0m \u001b[31m13.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m917.8/917.8 kB\u001b[0m \u001b[31m39.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
            "\u001b[0m"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import flax.linen as nn\n",
        "import flax\n",
        "from jax import jit, random\n",
        "import optax\n",
        "from transformers import MarianTokenizer, FlaxMarianMTModel, AutoConfig\n",
        "from datasets import load_dataset, DatasetDict\n",
        "from flax.training import train_state\n",
        "from flax.training.common_utils import get_metrics, onehot, shard, shard_prng_key\n",
        "from functools import partial\n",
        "from tqdm import tqdm\n",
        "import numpy as np\n",
        "from time import time\n",
        "from typing import Callable\n",
        "import jax_dataloader as jdl"
      ],
      "metadata": {
        "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
        "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
        "execution": {
          "iopub.status.busy": "2024-06-04T00:25:36.937758Z",
          "iopub.execute_input": "2024-06-04T00:25:36.938202Z",
          "iopub.status.idle": "2024-06-04T00:25:44.799741Z",
          "shell.execute_reply.started": "2024-06-04T00:25:36.938156Z",
          "shell.execute_reply": "2024-06-04T00:25:44.798703Z"
        },
        "trusted": true,
        "id": "uKplYYC9vSP2"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_config():\n",
        "    config = {\n",
        "        'model_name': 'Helsinki-NLP/opus-mt-en-hu', # 't5-small',\n",
        "        'max_length': 64,\n",
        "        'batch_size': 64,\n",
        "        'lr': 10 ** -5,\n",
        "        'epochs': 10,\n",
        "        'seed': 42,\n",
        "        'metric_name': 'sacrebleu',\n",
        "        'per_device_batch_size': 64\n",
        "\n",
        "    }\n",
        "    return config"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:25:45.142258Z",
          "iopub.execute_input": "2024-06-04T00:25:45.142669Z",
          "iopub.status.idle": "2024-06-04T00:25:45.153079Z",
          "shell.execute_reply.started": "2024-06-04T00:25:45.142632Z",
          "shell.execute_reply": "2024-06-04T00:25:45.152242Z"
        },
        "trusted": true,
        "id": "0j_9ZsSpvSP3"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "config = get_config()"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:25:45.154369Z",
          "iopub.execute_input": "2024-06-04T00:25:45.154693Z",
          "iopub.status.idle": "2024-06-04T00:25:45.166822Z",
          "shell.execute_reply.started": "2024-06-04T00:25:45.154667Z",
          "shell.execute_reply": "2024-06-04T00:25:45.165805Z"
        },
        "trusted": true,
        "id": "HfTrb26NvSP3"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Load dataset and metric\n",
        "dataset = load_dataset('Helsinki-NLP/opus_books', 'en-hu')\n",
        "val_test_set = dataset['train'].train_test_split(test_size=0.2, seed=42)\n",
        "test_set = val_test_set['test'].train_test_split(test_size=0.5, seed=42)\n",
        "\n",
        "dataset = DatasetDict({\n",
        "    'train': val_test_set['train'],\n",
        "    'val': test_set['test'],\n",
        "    'test': test_set['train']\n",
        "})"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:25:45.168090Z",
          "iopub.execute_input": "2024-06-04T00:25:45.168507Z",
          "iopub.status.idle": "2024-06-04T00:25:46.814026Z",
          "shell.execute_reply.started": "2024-06-04T00:25:45.168456Z",
          "shell.execute_reply": "2024-06-04T00:25:46.813243Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 242,
          "referenced_widgets": [
            "e263afc2c9fe4a1583a5a5c2b845a50f",
            "f3aff971ae2d43cf872a0d6bbbf8a392",
            "8b756212a3dd4e169012bf2d62770ccd",
            "0cdbbf0e98e7411bb562f22e2436888a",
            "808a2bddd1814f7aa90d751880375c52",
            "4621f2282c484f2b8ae16b70cc73ddcb",
            "f2f53f150d77481c81df52b26c55819f",
            "8a4c5272a7e74bc5ae31b4706ce444ad",
            "04ac49562024426c914a7f5a214a84f2",
            "a34604ca1828446b8ea7ce216829ab99",
            "07e57c3c4601459086d23b71d0cff4b0",
            "1ba66d0b45574a778eda36d1f41d07c5",
            "f48e45e5402940df80ff9f46645d97bd",
            "704c96aa0cd9440f9bb1b5d5017e8309",
            "7134f384e6384a5f9d95426efb511738",
            "b7b2053c16a3442ea95a34ba6a6d9232",
            "05724b4292044397bb320ebdea2a7de7",
            "4552e08385344020923c52fe653c5dbd",
            "fb54ccc429ed4ad79f3b0aa8896abe06",
            "e7f161c57eac4cd0a1b59d7dbe2764c4",
            "7c8b32d8b924415cae6e422209d32f3c",
            "f9398ec1118745949bb012f97239ce44",
            "09b9ab71b57649818b331ac69b9eaad2",
            "72a86b2670bb4859b6cd1403d3241b5e",
            "a37f09b7daf24e4c82f5dc7d2494cfc3",
            "cf74ec8b71a441828654076db4cc6ffd",
            "eefde40d5e174570b38d22672c374fc6",
            "9ffc7ae56b0347f4bcb44dfa103f277e",
            "653050d956124b56ae4caf67d635269a",
            "b3ce06aa2e8542f0b3313ce890832807",
            "84a959004d7344a99a004d54531bcf1a",
            "cba26d225fd34dd792718309f4e1c1d1",
            "9a915ab5ab854c27bc5fbcbded6e6927"
          ]
        },
        "id": "E3bEpJV-vSP3",
        "outputId": "598f8f0c-1ed1-41e9-cd4a-016bccc3ecac"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/huggingface_hub/utils/_token.py:89: 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": [
              "Downloading readme:   0%|          | 0.00/28.1k [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "e263afc2c9fe4a1583a5a5c2b845a50f"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading data:   0%|          | 0.00/23.1M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "1ba66d0b45574a778eda36d1f41d07c5"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Generating train split:   0%|          | 0/137151 [00:00<?, ? examples/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "09b9ab71b57649818b331ac69b9eaad2"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer = MarianTokenizer.from_pretrained(config['model_name'])\n",
        "# Load the model\n",
        "model_config = AutoConfig.from_pretrained(config['model_name'])\n",
        "model = FlaxMarianMTModel.from_pretrained(config['model_name'], config=model_config, from_pt=True)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:25:46.817071Z",
          "iopub.execute_input": "2024-06-04T00:25:46.817379Z",
          "iopub.status.idle": "2024-06-04T00:25:54.948554Z",
          "shell.execute_reply.started": "2024-06-04T00:25:46.817354Z",
          "shell.execute_reply": "2024-06-04T00:25:54.947720Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 351,
          "referenced_widgets": [
            "0f894258311c42599845b808ee59922a",
            "7136464b8afb43c2a259da17778aee50",
            "5496d0a320804de5990bbcc3d493e8d1",
            "c6b19c931c1f4f29b433628b157a7bdc",
            "114e7e6c4864433eaa73d9732410cc79",
            "b0e347980dff4e0481b3f6d2ef48edc6",
            "027edd6a9b8b4213841f38d2a965f931",
            "c449eea4ca344e3f8a8a56023ef0c7f0",
            "40a43dad986c4fb1859417b4eb29ab50",
            "697dfce9641e4120b7b0ef7963ea293a",
            "532a301e893a44bdb3a893f553174738",
            "4f808ac47f8442b0ae85c50ed5c3d413",
            "6e6bc82a81774489977a9fac31de042f",
            "61633f6fd20c42fb86720c62a98c41d4",
            "4ed93dfd117e4aadb56280f6e805acba",
            "7f32f52c085847b982635a251846232c",
            "3d1727222e794c95bb6f2488ef8b5c37",
            "044d409b3cf740c8b15e03184a6f2c2f",
            "d420de1bee8e41fbab51cb3441f5b853",
            "b2b111f54336469bb4614c039e046ca6",
            "e7dd99cbbde74113ad88ee74d555f61a",
            "adda0e06c2f64ea79795981b50d2fe0e",
            "8689ec2286e5414d989fa3f223484ae7",
            "144b417137674551a895482d147a8249",
            "2de396a5e6ca4a10af298a2ee61b2d37",
            "fffa7984d49c4c38accf7311236c8152",
            "2da4eed1a1204bed8e6b4c87e3bcdc3e",
            "1a112d908efa4ae082274ee21fffaee0",
            "13ea6f7789864999bed305937a18a359",
            "eee3c0d438044662980440d112f2d3ba",
            "49dad883ecab4bad983167a32d55da58",
            "2f22a464dcf1477e8f1f8d39ff1ebc4d",
            "aea3e71ef7e44d56941a1682aa5e95b9",
            "2c0557789f39449b83937fa970c735b7",
            "8f449f75e1d94e5e843c0c77bcd92446",
            "0bda9d055aa84c04aa1ae75122d302e1",
            "b4d214644cd3453481a813863ca6a8ae",
            "db25847587334e64ad383464705b9e0c",
            "3bf041ab1dc3475da1ef23a1e23d50ab",
            "84b88cb3f2b24e4f9d01d5ecbb0db169",
            "d377b0f3c85042acaa3e8d894007293c",
            "9f03d4c5db7e470180972e669da12574",
            "ee9420c3efe74f419eda70edd28a9ec4",
            "143d3eb4d22b4d81a3108b5743ff334b",
            "88162af0d5f3467983e9724e195c0bff",
            "c6583e04b06c4bccbe5a4678fa484c7c",
            "a339ffda3dbf46cfbc78a08f7ecf158b",
            "8e19f83d25ce4e81a728cdf054790065",
            "425cd5b013394e399dc2298e9682d1d8",
            "365ac66278b3400a842538087c82a28c",
            "dd13ef35ec3d451e8ad65eb8ac9e1711",
            "20870934ea084da49fd0a18ddd457f60",
            "d9c5d6f2331047be9a836918d7f2a82a",
            "b81cc242a2bd4cd5b1808bb20258b9b9",
            "bf95a897d59f459d883fcdd93bb36fa7",
            "682773a903be4be393c205216798fb76",
            "3429a1a9c7714781a6b29f7657a42cd1",
            "bb6341d0e3b7442791900959c73ee2e8",
            "e4b47402c556494b98d37f595002e92a",
            "44608d1b30d94a70aea4d2ad0cdf11f9",
            "ec8cb16c18d543ff9f3f904ce0817610",
            "7c5a94f7bb0f4b0eb1caf122df3dae1f",
            "0f3ad835d88a4393a481db2fd190060f",
            "4b038c5622e24e20ae41bb7334e0a771",
            "a8f6b50aa87c4165b35cb63b015f718e",
            "11db1eb4bca245d7aae36ea5e9bcceb5",
            "dc6f89f1d1e848a2adf3dc07a549bc95",
            "bc43a390fa2c4c8faf679b23fdc2a564",
            "453e46b83ab94d4ea8e47d63c27c0fca",
            "10e15538afbe4ce3a811280c43ac5755",
            "78b5cae51e8c478db152d3f61dc38081",
            "2e66e571995346869b5d9a81072f45d5",
            "00e26680684943339bcd02e84e2773cc",
            "df7174fd535e433091faf53a54b6426b",
            "26441d4b8330450094506f2a5bc89743",
            "751dfe82fd5348188e9870a6c430ea4d",
            "b12f0e2f4674457f9b26cfae802c55be"
          ]
        },
        "id": "PEfRF6zPvSP4",
        "outputId": "6053e13a-c761-479d-e477-96ab079205a9"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "tokenizer_config.json:   0%|          | 0.00/42.0 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "0f894258311c42599845b808ee59922a"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "source.spm:   0%|          | 0.00/792k [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "4f808ac47f8442b0ae85c50ed5c3d413"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "target.spm:   0%|          | 0.00/850k [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "8689ec2286e5414d989fa3f223484ae7"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "vocab.json:   0%|          | 0.00/1.57M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "2c0557789f39449b83937fa970c735b7"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "config.json:   0%|          | 0.00/1.38k [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "88162af0d5f3467983e9724e195c0bff"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/transformers/models/marian/tokenization_marian.py:175: UserWarning: Recommended: pip install sacremoses.\n",
            "  warnings.warn(\"Recommended: pip install sacremoses.\")\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "pytorch_model.bin:   0%|          | 0.00/307M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "682773a903be4be393c205216798fb76"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Some weights of the model checkpoint at Helsinki-NLP/opus-mt-en-hu were not used when initializing FlaxMarianMTModel: {('model', 'encoder', 'embed_positions', 'kernel'), ('model', 'decoder', 'embed_positions', 'kernel'), ('model', 'encoder', 'embed_tokens', 'kernel'), ('model', 'decoder', 'embed_tokens', 'kernel')}\n",
            "- This IS expected if you are initializing FlaxMarianMTModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
            "- This IS NOT expected if you are initializing FlaxMarianMTModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "generation_config.json:   0%|          | 0.00/293 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "dc6f89f1d1e848a2adf3dc07a549bc95"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def preprocess_fn(dataset: DatasetDict):\n",
        "    inputs = [ex['en'] for ex in dataset['translation']]\n",
        "    targets = [ex['hu'] for ex in dataset['translation']]\n",
        "    model_inputs = tokenizer(inputs, max_length=config['max_length'], truncation=True, padding='longest')\n",
        "\n",
        "    # tokenizer targets\n",
        "    #with tokenizer.as_target_tokenizer():\n",
        "    labels =  tokenizer(targets, max_length=config['max_length'], truncation=True, padding='longest')\n",
        "\n",
        "    model_inputs['labels'] = labels.input_ids\n",
        "\n",
        "    return model_inputs"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:25:54.949744Z",
          "iopub.execute_input": "2024-06-04T00:25:54.950049Z",
          "iopub.status.idle": "2024-06-04T00:25:54.956554Z",
          "shell.execute_reply.started": "2024-06-04T00:25:54.950023Z",
          "shell.execute_reply": "2024-06-04T00:25:54.955588Z"
        },
        "trusted": true,
        "id": "hObjtz9mvSP4"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokenized_datasets = dataset.map(preprocess_fn, batched=True, remove_columns=dataset[\"train\"].column_names)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:25:54.957743Z",
          "iopub.execute_input": "2024-06-04T00:25:54.957998Z",
          "iopub.status.idle": "2024-06-04T00:25:59.240189Z",
          "shell.execute_reply.started": "2024-06-04T00:25:54.957976Z",
          "shell.execute_reply": "2024-06-04T00:25:59.239126Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 113,
          "referenced_widgets": [
            "be5ac7d6a07d4128863a7ab279c8024a",
            "3a587589d98c475e8504f22080b42694",
            "773646b21ccd4cfcb4549900d507fd94",
            "034d59af89444a53bfb94c853c37d6a8",
            "df14e319dd704b0f860f1a9182922e64",
            "4414556419b84a448d5c99281a033e0a",
            "3e9d1b51fdb14544a52a09d2539ad3c7",
            "f697b7d1861e4683a1943f9a92ae3f0e",
            "91d54976ce7c42bcbed276f0707b78e1",
            "dc08fcfee29340c6afe258c499682ca0",
            "aac4d7efafae4a879cf67d7101163ba5",
            "690df136207644daaf42380f181c7ced",
            "92f8dc6b5a114972b1c5ae821845832d",
            "53d9d3750940472cb77e60d206de54f3",
            "9de6043ecc6d469f88e798f578d020bb",
            "d9effe34749c4948a2098698e1dc1f88",
            "481e16c45c2d477db4689891ede0b1d4",
            "a495a7c59e7d4771bd6ecdb1235cab03",
            "33e8b0cf32b44e3f93d6cdf403b3966b",
            "6441dfea809648d0a6707d03ddd3909a",
            "2c1e043eb32a45c0a814bf8268a70642",
            "e8a97c273fc943e497390e3052892ae1",
            "5da90955287e42cc83ed1bbba5d12e11",
            "d43bfd54a13644869754525c046802e4",
            "cdbc5fa699f5457eb0843a565e7881f7",
            "6848df35d3eb4814bacff411f8a473ac",
            "43a0a2392e7d48bfb8321c3c8e8ffed6",
            "ad6cadadebf946a6824106e1796615ea",
            "a50d44c51be84b4897ae94d0313c1ee0",
            "5d08d65ea1fc41eab5a98d3807f6aa91",
            "2f99c6a66a6146c0923fe8ee107f9281",
            "a8fb09476729415794e89c6fe910514d",
            "6879f59127654029b1d9d93d876d891c"
          ]
        },
        "id": "Fe0eCj0FvSP4",
        "outputId": "43f71084-888a-4cd0-cf80-e0ace3437d37"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Map:   0%|          | 0/109720 [00:00<?, ? examples/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "be5ac7d6a07d4128863a7ab279c8024a"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Map:   0%|          | 0/13716 [00:00<?, ? examples/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "690df136207644daaf42380f181c7ced"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Map:   0%|          | 0/13715 [00:00<?, ? examples/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "5da90955287e42cc83ed1bbba5d12e11"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "tokenized_datasets"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:25:59.241353Z",
          "iopub.execute_input": "2024-06-04T00:25:59.241656Z",
          "iopub.status.idle": "2024-06-04T00:25:59.251123Z",
          "shell.execute_reply.started": "2024-06-04T00:25:59.241623Z",
          "shell.execute_reply": "2024-06-04T00:25:59.250148Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lEdQG6TDvSP4",
        "outputId": "1a28fe7e-78d7-40ec-c434-16d8619dcbed"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DatasetDict({\n",
              "    train: Dataset({\n",
              "        features: ['input_ids', 'attention_mask', 'labels'],\n",
              "        num_rows: 109720\n",
              "    })\n",
              "    val: Dataset({\n",
              "        features: ['input_ids', 'attention_mask', 'labels'],\n",
              "        num_rows: 13716\n",
              "    })\n",
              "    test: Dataset({\n",
              "        features: ['input_ids', 'attention_mask', 'labels'],\n",
              "        num_rows: 13715\n",
              "    })\n",
              "})"
            ]
          },
          "metadata": {},
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Dataloader"
      ],
      "metadata": {
        "id": "gbsyKlCdvSP5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "total_batch_size = config['per_device_batch_size'] * jax.local_device_count()\n",
        "print(\"The overall batch size (both for training and eval) is\", total_batch_size)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:25:59.252233Z",
          "iopub.execute_input": "2024-06-04T00:25:59.252532Z",
          "iopub.status.idle": "2024-06-04T00:25:59.260982Z",
          "shell.execute_reply.started": "2024-06-04T00:25:59.252508Z",
          "shell.execute_reply": "2024-06-04T00:25:59.260028Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZIn7kzEMvSP6",
        "outputId": "72ed01a0-6dbb-4362-c5b2-55bd3be2812f"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The overall batch size (both for training and eval) is 512\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "train_loader = jdl.DataLoader(tokenized_datasets['train'], 'jax', batch_size=total_batch_size, shuffle=True)\n",
        "val_loader = jdl.DataLoader(tokenized_datasets['val'], 'jax', batch_size=total_batch_size, shuffle=False, drop_last=True)\n",
        "test_loader = jdl.DataLoader(tokenized_datasets['test'], 'jax', batch_size=total_batch_size, shuffle=False)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:25:59.262154Z",
          "iopub.execute_input": "2024-06-04T00:25:59.262477Z",
          "iopub.status.idle": "2024-06-04T00:25:59.280703Z",
          "shell.execute_reply.started": "2024-06-04T00:25:59.262452Z",
          "shell.execute_reply": "2024-06-04T00:25:59.279638Z"
        },
        "trusted": true,
        "id": "4kTsoP6LvSP6"
      },
      "execution_count": 42,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "len_train_loader = len(train_loader)\n",
        "len_val_loader = len(val_loader)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:25:59.282077Z",
          "iopub.execute_input": "2024-06-04T00:25:59.282504Z",
          "iopub.status.idle": "2024-06-04T00:25:59.289022Z",
          "shell.execute_reply.started": "2024-06-04T00:25:59.282468Z",
          "shell.execute_reply": "2024-06-04T00:25:59.288148Z"
        },
        "trusted": true,
        "id": "8e5Vl_LFvSP6"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Tokenizer and model"
      ],
      "metadata": {
        "id": "DPV3j444vSP7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "num_train_steps = len(dataset['train']) // total_batch_size * config['epochs']\n",
        "\n",
        "#learning_rate_function = optax.linear_schedule(init_value=config['lr'], end_value=0, transition_steps=num_train_steps)\n",
        "learning_rate_function = optax.cosine_onecycle_schedule(transition_steps=num_train_steps, peak_value=config['lr'], pct_start=0.1, )"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:25:59.290422Z",
          "iopub.execute_input": "2024-06-04T00:25:59.290813Z",
          "iopub.status.idle": "2024-06-04T00:25:59.304263Z",
          "shell.execute_reply.started": "2024-06-04T00:25:59.290779Z",
          "shell.execute_reply": "2024-06-04T00:25:59.303343Z"
        },
        "trusted": true,
        "id": "GZ0ex1pavSP7"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class TrainState(train_state.TrainState):\n",
        "    loss_function: Callable = flax.struct.field(pytree_node=False)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:25:59.305542Z",
          "iopub.execute_input": "2024-06-04T00:25:59.305918Z",
          "iopub.status.idle": "2024-06-04T00:25:59.319108Z",
          "shell.execute_reply.started": "2024-06-04T00:25:59.305883Z",
          "shell.execute_reply": "2024-06-04T00:25:59.318083Z"
        },
        "trusted": true,
        "id": "VVCsd5LGvSP7"
      },
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "@jit\n",
        "def loss_function(logits, labels):\n",
        "    # Compute the cross-entropy loss\n",
        "    loss = optax.softmax_cross_entropy(logits, onehot(labels, num_classes=logits.shape[-1]))\n",
        "    return jnp.mean(loss)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:28:01.634637Z",
          "iopub.execute_input": "2024-06-04T00:28:01.635083Z",
          "iopub.status.idle": "2024-06-04T00:28:01.641162Z",
          "shell.execute_reply.started": "2024-06-04T00:28:01.635052Z",
          "shell.execute_reply": "2024-06-04T00:28:01.640103Z"
        },
        "trusted": true,
        "id": "cI5OpYtMvSP7"
      },
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "state = TrainState.create(\n",
        "    apply_fn=model.__call__,\n",
        "    params=model.params,\n",
        "    tx=optax.adamw(learning_rate=learning_rate_function),\n",
        "    loss_function=loss_function,\n",
        ")"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:28:06.588057Z",
          "iopub.execute_input": "2024-06-04T00:28:06.588724Z",
          "iopub.status.idle": "2024-06-04T00:28:06.968856Z",
          "shell.execute_reply.started": "2024-06-04T00:28:06.588688Z",
          "shell.execute_reply": "2024-06-04T00:28:06.967786Z"
        },
        "trusted": true,
        "id": "tYmZ8WtIvSP7"
      },
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def train_step(state, batch, dropout_rng):\n",
        "    dropout_rng, new_dropout_rng = random.split(dropout_rng)\n",
        "    def compute_loss(params):\n",
        "        labels = batch.pop(\"labels\")\n",
        "        logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]\n",
        "        loss = state.loss_function(logits, labels)\n",
        "        return loss\n",
        "\n",
        "    loss, grads = jax.value_and_grad(compute_loss)(state.params)\n",
        "    grads = jax.lax.pmean(grads, axis_name='batch')\n",
        "    loss = jax.lax.pmean(loss, axis_name='batch')\n",
        "    new_state = state.apply_gradients(grads=grads)\n",
        "    metrics = jax.lax.pmean({\"loss\": loss}, axis_name=\"batch\")\n",
        "    return new_state, metrics, new_dropout_rng"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:28:12.315328Z",
          "iopub.execute_input": "2024-06-04T00:28:12.316108Z",
          "iopub.status.idle": "2024-06-04T00:28:12.323380Z",
          "shell.execute_reply.started": "2024-06-04T00:28:12.316068Z",
          "shell.execute_reply": "2024-06-04T00:28:12.322340Z"
        },
        "trusted": true,
        "id": "4c7K4univSP7"
      },
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "parallel_train_step = jax.pmap(train_step, axis_name=\"batch\", donate_argnums=(0,))"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:28:13.704849Z",
          "iopub.execute_input": "2024-06-04T00:28:13.705653Z",
          "iopub.status.idle": "2024-06-04T00:28:13.710149Z",
          "shell.execute_reply.started": "2024-06-04T00:28:13.705612Z",
          "shell.execute_reply": "2024-06-04T00:28:13.709124Z"
        },
        "trusted": true,
        "id": "W4ZveuiEvSP7"
      },
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def eval_step(state, batch):\n",
        "    labels = batch.pop(\"labels\")\n",
        "    logits = state.apply_fn(**batch, params=state.params, train=False)[0]\n",
        "    loss = state.loss_function(logits, labels)\n",
        "    metrics = jax.lax.pmean({\"loss\": loss}, axis_name=\"batch\")\n",
        "    return metrics"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:28:16.498338Z",
          "iopub.execute_input": "2024-06-04T00:28:16.499639Z",
          "iopub.status.idle": "2024-06-04T00:28:16.506036Z",
          "shell.execute_reply.started": "2024-06-04T00:28:16.499587Z",
          "shell.execute_reply": "2024-06-04T00:28:16.504922Z"
        },
        "trusted": true,
        "id": "JQAIx_sUvSP8"
      },
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "parallel_eval_step = jax.pmap(eval_step, axis_name=\"batch\")"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:28:17.945339Z",
          "iopub.execute_input": "2024-06-04T00:28:17.946119Z",
          "iopub.status.idle": "2024-06-04T00:28:17.950539Z",
          "shell.execute_reply.started": "2024-06-04T00:28:17.946084Z",
          "shell.execute_reply": "2024-06-04T00:28:17.949574Z"
        },
        "trusted": true,
        "id": "GtNAq2bzvSP8"
      },
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "state = flax.jax_utils.replicate(state)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:28:19.278045Z",
          "iopub.execute_input": "2024-06-04T00:28:19.278884Z",
          "iopub.status.idle": "2024-06-04T00:28:19.718911Z",
          "shell.execute_reply.started": "2024-06-04T00:28:19.278844Z",
          "shell.execute_reply": "2024-06-04T00:28:19.717602Z"
        },
        "trusted": true,
        "id": "BAI5dXDzvSP8"
      },
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training"
      ],
      "metadata": {
        "id": "Gs4oHRrdvSP8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def shard_data(xs):\n",
        "    num_devices = jax.device_count()\n",
        "    return {k: np.reshape(v, (num_devices, -1) + v.shape[1:]) for k, v in xs.items()}"
      ],
      "metadata": {
        "id": "xVJg429p6V-x"
      },
      "execution_count": 43,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "rng = jax.random.PRNGKey(0)\n",
        "dropout_rngs = jax.random.split(rng, jax.local_device_count())"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:28:27.359036Z",
          "iopub.execute_input": "2024-06-04T00:28:27.359959Z",
          "iopub.status.idle": "2024-06-04T00:28:27.365807Z",
          "shell.execute_reply.started": "2024-06-04T00:28:27.359923Z",
          "shell.execute_reply": "2024-06-04T00:28:27.364816Z"
        },
        "trusted": true,
        "id": "KIRwE5s4vSP9"
      },
      "execution_count": 28,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "total_batch_size = config['per_device_batch_size'] * jax.local_device_count()\n",
        "print(\"The overall batch size (both for training and eval) is\", total_batch_size)\n",
        "\n",
        "start_time = time()\n",
        "print('======== Start training ======== ')\n",
        "for epoch in range(config['epochs']):\n",
        "\n",
        "    with tqdm(total=len_train_loader, desc=\"Training...\", leave=False) as progress_bar_train:\n",
        "        for batch in train_loader:\n",
        "            batch = shard_data(batch)\n",
        "            state, train_metric, dropout_rngs = parallel_train_step(state, batch, dropout_rngs)\n",
        "            progress_bar_train.update(1)\n",
        "        progress_bar_train.write(\n",
        "              f\"Train... ({epoch}/{config['epochs']} | Loss: {round(train_metric['loss'].mean(), 3)})\"\n",
        "        )\n",
        "\n",
        "    eval_metrics = []\n",
        "    with tqdm(total=len_val_loader, desc=\"Evaluating...\", leave=False) as progress_bar_eval:\n",
        "        for batch in val_loader:\n",
        "            batch = shard_data(batch)\n",
        "            eval_metric = parallel_eval_step(state, batch)\n",
        "            eval_metrics.append(eval_metric)\n",
        "            progress_bar_eval.update(1)\n",
        "        eval_metrics = get_metrics(eval_metrics)\n",
        "        eval_metrics = jax.tree_map(jnp.mean, eval_metrics)\n",
        "        progress_bar_eval.write(\n",
        "            f\"Eval... ({epoch}/{config['epochs']} | Loss: {eval_metrics['loss']})\"\n",
        "        )\n",
        "\n",
        "\n",
        "print('======== End training ========')\n",
        "total_training_time = time() - start_time\n",
        "print(f\"Total training time: {total_training_time}\")"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-06-04T00:28:27.556260Z",
          "iopub.execute_input": "2024-06-04T00:28:27.556653Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VLrDiVKnvSP9",
        "outputId": "df81584a-1db2-4063-9344-9e02afad040b"
      },
      "execution_count": 44,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The overall batch size (both for training and eval) is 512\n",
            "======== Start training ======== \n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train... (0/10 | Loss: 3.809000253677368)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Evaluating...:  81%|████████▏ | 22/27 [00:00<00:00, 68.81it/s]<ipython-input-44-0599d877754a>:33: DeprecationWarning: jax.tree_map is deprecated: use jax.tree.map (jax v0.4.25 or newer) or jax.tree_util.tree_map (any JAX version).\n",
            "  eval_metrics = jax.tree_map(jnp.mean, eval_metrics)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Eval... (0/10 | Loss: 3.8092541694641113)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train... (1/10 | Loss: 3.9230000972747803)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Evaluating...:  93%|█████████▎| 25/27 [00:00<00:00, 61.81it/s]<ipython-input-44-0599d877754a>:33: DeprecationWarning: jax.tree_map is deprecated: use jax.tree.map (jax v0.4.25 or newer) or jax.tree_util.tree_map (any JAX version).\n",
            "  eval_metrics = jax.tree_map(jnp.mean, eval_metrics)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Eval... (1/10 | Loss: 3.7923882007598877)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train... (2/10 | Loss: 3.69100022315979)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Eval... (2/10 | Loss: 3.781726121902466)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train... (3/10 | Loss: 3.757000207901001)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Eval... (3/10 | Loss: 3.777029275894165)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train... (4/10 | Loss: 3.9900002479553223)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Eval... (4/10 | Loss: 3.774854898452759)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train... (5/10 | Loss: 3.994000196456909)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Eval... (5/10 | Loss: 3.774083137512207)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train... (6/10 | Loss: 3.7810001373291016)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Eval... (6/10 | Loss: 3.774078845977783)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train... (7/10 | Loss: 3.7370002269744873)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Eval... (7/10 | Loss: 3.7740955352783203)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train... (8/10 | Loss: 3.9600002765655518)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Eval... (8/10 | Loss: 3.7741010189056396)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": []
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Train... (9/10 | Loss: 3.886000156402588)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "                                                              "
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Eval... (9/10 | Loss: 3.774094343185425)\n",
            "======== End training ========\n",
            "Total training time: 371.52358317375183\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r"
          ]
        }
      ]
    }
  ]
}