notebooks/demo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "03638a96888649dd9dc625a623eb6c34": {
          "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_ffa7a23d200843068267e4e21553e98d",
              "IPY_MODEL_ad60a3884afe47a9ab9b4a70a461d49b",
              "IPY_MODEL_2248396d844b4f159fd895fa100f5875"
            ],
            "layout": "IPY_MODEL_073cc65432b44fe29eae1c91471ff970"
          }
        },
        "ffa7a23d200843068267e4e21553e98d": {
          "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_137e376507144fccbabb12bf1268f699",
            "placeholder": "​",
            "style": "IPY_MODEL_1a6a03868cdc4fb286b4ec580da9a391",
            "value": ""
          }
        },
        "ad60a3884afe47a9ab9b4a70a461d49b": {
          "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_108098be0bdc49e791fb784d487fe175",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_77eb47b1a453410b860b91a01357e98f",
            "value": 0
          }
        },
        "2248396d844b4f159fd895fa100f5875": {
          "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_ddc0e14f4448489499a534c5cf2f5945",
            "placeholder": "​",
            "style": "IPY_MODEL_6e8a37d65bf64bbaac84230613164936",
            "value": " 0/0 [00:00&lt;?, ?it/s]"
          }
        },
        "073cc65432b44fe29eae1c91471ff970": {
          "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
          }
        },
        "137e376507144fccbabb12bf1268f699": {
          "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
          }
        },
        "1a6a03868cdc4fb286b4ec580da9a391": {
          "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": ""
          }
        },
        "108098be0bdc49e791fb784d487fe175": {
          "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": "20px"
          }
        },
        "77eb47b1a453410b860b91a01357e98f": {
          "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": ""
          }
        },
        "ddc0e14f4448489499a534c5cf2f5945": {
          "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
          }
        },
        "6e8a37d65bf64bbaac84230613164936": {
          "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": ""
          }
        },
        "67dc7ff4d0fe4af796308b2c8355d150": {
          "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_b41dad78195344a6a77b570b1783f7a0",
              "IPY_MODEL_494bb8b09d1f4f2abbc111eb5e6da130",
              "IPY_MODEL_243efd24b8994278a406c597039d0fc4"
            ],
            "layout": "IPY_MODEL_c2ccdf34885d41a0ab3f9b40c775b25a"
          }
        },
        "b41dad78195344a6a77b570b1783f7a0": {
          "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_39efadb9ccb048deab07969016e7bd38",
            "placeholder": "​",
            "style": "IPY_MODEL_d36446b0a9534f41a9f21b865b0a08f4",
            "value": "config.json: 100%"
          }
        },
        "494bb8b09d1f4f2abbc111eb5e6da130": {
          "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_6a458bdef43543d0b7b7edc2883e2dc1",
            "max": 720,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_7f85bc3a52a540558999a1ff25a813b3",
            "value": 720
          }
        },
        "243efd24b8994278a406c597039d0fc4": {
          "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_325f9b174633405483cf60cff12d652b",
            "placeholder": "​",
            "style": "IPY_MODEL_e60bfa41b7454322a908e01bb9caed65",
            "value": " 720/720 [00:00&lt;00:00, 20.2kB/s]"
          }
        },
        "c2ccdf34885d41a0ab3f9b40c775b25a": {
          "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
          }
        },
        "39efadb9ccb048deab07969016e7bd38": {
          "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
          }
        },
        "d36446b0a9534f41a9f21b865b0a08f4": {
          "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": ""
          }
        },
        "6a458bdef43543d0b7b7edc2883e2dc1": {
          "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
          }
        },
        "7f85bc3a52a540558999a1ff25a813b3": {
          "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": ""
          }
        },
        "325f9b174633405483cf60cff12d652b": {
          "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
          }
        },
        "e60bfa41b7454322a908e01bb9caed65": {
          "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": ""
          }
        },
        "7b7b50c298de414b85cc622ef5660dfb": {
          "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_d1e82a41293d41d6aa8807d9d575e86b",
              "IPY_MODEL_a19471c4ecce428bbebab2a90f535861",
              "IPY_MODEL_8a2c815f41b8483692a79165ebeaf3f6"
            ],
            "layout": "IPY_MODEL_9782a56e37e24892a78b282295ae5ac0"
          }
        },
        "d1e82a41293d41d6aa8807d9d575e86b": {
          "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_9be1146cefb9416cb875741640dca611",
            "placeholder": "​",
            "style": "IPY_MODEL_15bfcbb901574f0583b8af3733f58b7a",
            "value": "Loading experts: 100%"
          }
        },
        "a19471c4ecce428bbebab2a90f535861": {
          "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_c0dba9bd6ee94ca1a0084eabeac0e1ca",
            "max": 32,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_83bcda80b79a4e6eaf5b0f27bcbc0ee0",
            "value": 32
          }
        },
        "8a2c815f41b8483692a79165ebeaf3f6": {
          "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_b792acc1fff049fca658f836557c8e6f",
            "placeholder": "​",
            "style": "IPY_MODEL_dcf2b97ed0b7415fa02f38c8bb0e3009",
            "value": " 32/32 [01:36&lt;00:00,  3.03s/it]"
          }
        },
        "9782a56e37e24892a78b282295ae5ac0": {
          "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
          }
        },
        "9be1146cefb9416cb875741640dca611": {
          "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
          }
        },
        "15bfcbb901574f0583b8af3733f58b7a": {
          "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": ""
          }
        },
        "c0dba9bd6ee94ca1a0084eabeac0e1ca": {
          "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
          }
        },
        "83bcda80b79a4e6eaf5b0f27bcbc0ee0": {
          "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": ""
          }
        },
        "b792acc1fff049fca658f836557c8e6f": {
          "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
          }
        },
        "dcf2b97ed0b7415fa02f38c8bb0e3009": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Mixtral in Colab\n",
        "\n",
        "Welcome! In this notebook you can run [Mixtral8x7B-Instruct](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1) with decent generation speed **right in Google Colab or on a consumer-grade GPU**. This was made possible by quantizing the original model in mixed precision and implementing a MoE-specific offloading strategy.\n",
        "\n",
        "To learn more, read our [tech report](https://arxiv.org/abs/2312.17238) or check out the [repo](https://github.com/dvmazur/mixtral-offloading) on GitHub."
      ],
      "metadata": {
        "id": "OW1moHJ1TdhO"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "One will need approximately 16 GB of VRAM and 11 GB of RAM to run this notebook and generate somewhat long texts.\n",
        "\n",
        "\n",
        "<details>\n",
        "\n",
        "<summary>How to balance between RAM and GPU VRAM usage</summary>\n",
        "\n",
        "You can balance between RAM and GPU VRAM usage by changing <code>offload_per_layer</code> variable in the <a href=\"#scrollTo=_mIpePTMFyRY&line=10&uniqifier=1\">Initialize model</a> section. Increasing <code>offload_per_layer</code> will decrease GPU VRAM usage, increase RAM usage and decrease generation speed. Decreasing <code>offload_per_layer</code> will have the opposite effect.\n",
        "\n",
        "Note that this notebook should run normally in Google Colab with <code>offload_per_layer = 4</code>, but may crush with other values. However, if you run this somewhere else, you're free to play with this variable.\n",
        "</details>"
      ],
      "metadata": {
        "id": "2-dvAX_hKZT4"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Install and import libraries"
      ],
      "metadata": {
        "id": "Y8MhvkC7TKEL"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "f7qY7ebqX7T7"
      },
      "outputs": [],
      "source": [
        "# fix numpy in colab\n",
        "import numpy\n",
        "from IPython.display import clear_output\n",
        "\n",
        "# fix triton in colab\n",
        "!export LC_ALL=\"en_US.UTF-8\"\n",
        "!export LD_LIBRARY_PATH=\"/usr/lib64-nvidia\"\n",
        "!export LIBRARY_PATH=\"/usr/local/cuda/lib64/stubs\"\n",
        "!ldconfig /usr/lib64-nvidia\n",
        "\n",
        "!git clone https://github.com/dvmazur/mixtral-offloading.git --quiet\n",
        "!cd mixtral-offloading && pip install -q -r requirements.txt\n",
        "!huggingface-cli download lavawolfiee/Mixtral-8x7B-Instruct-v0.1-offloading-demo --quiet --local-dir Mixtral-8x7B-Instruct-v0.1-offloading-demo\n",
        "\n",
        "clear_output()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import sys\n",
        "\n",
        "sys.path.append(\"mixtral-offloading\")\n",
        "import torch\n",
        "from torch.nn import functional as F\n",
        "from hqq.core.quantize import BaseQuantizeConfig\n",
        "from huggingface_hub import snapshot_download\n",
        "from IPython.display import clear_output\n",
        "from tqdm.auto import trange\n",
        "from transformers import AutoConfig, AutoTokenizer\n",
        "from transformers.utils import logging as hf_logging\n",
        "\n",
        "from src.build_model import OffloadConfig, QuantConfig, build_model"
      ],
      "metadata": {
        "id": "GgpjnV7fV49W",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 105,
          "referenced_widgets": [
            "03638a96888649dd9dc625a623eb6c34",
            "ffa7a23d200843068267e4e21553e98d",
            "ad60a3884afe47a9ab9b4a70a461d49b",
            "2248396d844b4f159fd895fa100f5875",
            "073cc65432b44fe29eae1c91471ff970",
            "137e376507144fccbabb12bf1268f699",
            "1a6a03868cdc4fb286b4ec580da9a391",
            "108098be0bdc49e791fb784d487fe175",
            "77eb47b1a453410b860b91a01357e98f",
            "ddc0e14f4448489499a534c5cf2f5945",
            "6e8a37d65bf64bbaac84230613164936"
          ]
        },
        "outputId": "1c10ea85-61e5-4572-aac9-0f02c4d2cb30"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "hqq_aten package not installed. HQQBackend.ATEN backend will not work unless you install the hqq_aten lib in hqq/kernels.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "The cache for model files in Transformers v4.22.0 has been updated. Migrating your old cache. This is a one-time only operation. You can interrupt this and resume the migration later on by calling `transformers.utils.move_cache()`.\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "0it [00:00, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "03638a96888649dd9dc625a623eb6c34"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Initialize model"
      ],
      "metadata": {
        "id": "OkSYibHcTQsH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model_name = \"mistralai/Mixtral-8x7B-Instruct-v0.1\"\n",
        "quantized_model_name = \"lavawolfiee/Mixtral-8x7B-Instruct-v0.1-offloading-demo\"\n",
        "state_path = \"Mixtral-8x7B-Instruct-v0.1-offloading-demo\"\n",
        "\n",
        "config = AutoConfig.from_pretrained(quantized_model_name)\n",
        "\n",
        "device = torch.device(\"cuda:0\")\n",
        "\n",
        "##### Change this to 5 if you have only 12 GB of GPU VRAM #####\n",
        "offload_per_layer = 4\n",
        "# offload_per_layer = 5\n",
        "###############################################################\n",
        "\n",
        "num_experts = config.num_local_experts\n",
        "\n",
        "offload_config = OffloadConfig(\n",
        "    main_size=config.num_hidden_layers * (num_experts - offload_per_layer),\n",
        "    offload_size=config.num_hidden_layers * offload_per_layer,\n",
        "    buffer_size=4,\n",
        "    offload_per_layer=offload_per_layer,\n",
        ")\n",
        "\n",
        "\n",
        "attn_config = BaseQuantizeConfig(\n",
        "    nbits=4,\n",
        "    group_size=64,\n",
        "    quant_zero=True,\n",
        "    quant_scale=True,\n",
        ")\n",
        "attn_config[\"scale_quant_params\"][\"group_size\"] = 256\n",
        "\n",
        "\n",
        "ffn_config = BaseQuantizeConfig(\n",
        "    nbits=2,\n",
        "    group_size=16,\n",
        "    quant_zero=True,\n",
        "    quant_scale=True,\n",
        ")\n",
        "quant_config = QuantConfig(ffn_config=ffn_config, attn_config=attn_config)\n",
        "\n",
        "\n",
        "model = build_model(\n",
        "    device=device,\n",
        "    quant_config=quant_config,\n",
        "    offload_config=offload_config,\n",
        "    state_path=state_path,\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 227,
          "referenced_widgets": [
            "67dc7ff4d0fe4af796308b2c8355d150",
            "b41dad78195344a6a77b570b1783f7a0",
            "494bb8b09d1f4f2abbc111eb5e6da130",
            "243efd24b8994278a406c597039d0fc4",
            "c2ccdf34885d41a0ab3f9b40c775b25a",
            "39efadb9ccb048deab07969016e7bd38",
            "d36446b0a9534f41a9f21b865b0a08f4",
            "6a458bdef43543d0b7b7edc2883e2dc1",
            "7f85bc3a52a540558999a1ff25a813b3",
            "325f9b174633405483cf60cff12d652b",
            "e60bfa41b7454322a908e01bb9caed65",
            "7b7b50c298de414b85cc622ef5660dfb",
            "d1e82a41293d41d6aa8807d9d575e86b",
            "a19471c4ecce428bbebab2a90f535861",
            "8a2c815f41b8483692a79165ebeaf3f6",
            "9782a56e37e24892a78b282295ae5ac0",
            "9be1146cefb9416cb875741640dca611",
            "15bfcbb901574f0583b8af3733f58b7a",
            "c0dba9bd6ee94ca1a0084eabeac0e1ca",
            "83bcda80b79a4e6eaf5b0f27bcbc0ee0",
            "b792acc1fff049fca658f836557c8e6f",
            "dcf2b97ed0b7415fa02f38c8bb0e3009"
          ]
        },
        "id": "_mIpePTMFyRY",
        "outputId": "8a113aee-bcaa-4434-bdfc-2e65c0c58b15"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/huggingface_hub/utils/_token.py:72: 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": [
              "config.json:   0%|          | 0.00/720 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "67dc7ff4d0fe4af796308b2c8355d150"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torch/nn/init.py:412: UserWarning: Initializing zero-element tensors is a no-op\n",
            "  warnings.warn(\"Initializing zero-element tensors is a no-op\")\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Loading experts:   0%|          | 0/32 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "7b7b50c298de414b85cc622ef5660dfb"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Run the model"
      ],
      "metadata": {
        "id": "Z4hBFYtPTUzT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import TextStreamer\n",
        "\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)\n",
        "past_key_values = None\n",
        "sequence = None\n",
        "\n",
        "seq_len = 0\n",
        "while True:\n",
        "  print(\"User: \", end=\"\")\n",
        "  user_input = input()\n",
        "  print(\"\\n\")\n",
        "\n",
        "  user_entry = dict(role=\"user\", content=user_input)\n",
        "  input_ids = tokenizer.apply_chat_template([user_entry], return_tensors=\"pt\").to(device)\n",
        "\n",
        "  if past_key_values is None:\n",
        "    attention_mask = torch.ones_like(input_ids)\n",
        "  else:\n",
        "    seq_len = input_ids.size(1) + past_key_values[0][0][0].size(1)\n",
        "    attention_mask = torch.ones([1, seq_len - 1], dtype=torch.int, device=device)\n",
        "\n",
        "  print(\"Mixtral: \", end=\"\")\n",
        "  result = model.generate(\n",
        "    input_ids=input_ids,\n",
        "    attention_mask=attention_mask,\n",
        "    past_key_values=past_key_values,\n",
        "    streamer=streamer,\n",
        "    do_sample=True,\n",
        "    temperature=0.9,\n",
        "    top_p=0.9,\n",
        "    max_new_tokens=512,\n",
        "    pad_token_id=tokenizer.eos_token_id,\n",
        "    return_dict_in_generate=True,\n",
        "    output_hidden_states=True,\n",
        "  )\n",
        "  print(\"\\n\")\n",
        "\n",
        "  sequence = result[\"sequences\"]\n",
        "  past_key_values = result[\"past_key_values\"]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "Zf4GkspecSm8",
        "outputId": "a4ce58f9-30cc-414d-fa66-e5a28262d666"
      },
      "execution_count": null,
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "User: Write a funny poem about Python, please\n",
            "\n",
            "\n",
            "Mixtral: There once was a language so bright,\n",
            "named Python, a helpful little wight.\n",
            "It slithered through code with ease,\n",
            "No bug stood a chance, to say the least.\n",
            "\n",
            "Its syntax so clean, and easy on the eyes,\n",
            "Programmers from far and wide would shout their surprise.\n",
            "But unlike its namesake, it's not sneaky or sly,\n",
            "It's open source, and free to the sky!\n",
            "\n",
            "With libraries so vast, it's a data geek's dream,\n",
            "From AI to web scraping, it's the ultimate scheme.\n",
            "It's the Swiss Army knife of coding, or perhaps a black belt,\n",
            "In the world of programming, it's the ultimate feel.\n",
            "\n",
            "So, whether you're a beginner or a seasoned coder,\n",
            "Python's the language that will make you feel higher.\n",
            "With a community so welcoming and a syntax so neat,\n",
            "It's no wonder that Python is simply hard to beat!\n",
            "\n",
            "So here's to the snake, in programming so great,\n",
            "May it continue to dominate, until a much later date.\n",
            "In the world of technology, it's here to stay,\n",
            "Long live Python, programming the Python way!\n",
            "\n",
            "\n",
            "User: Thank you, that's cool! Now write a funny poem about LLMs (large language models)\n",
            "\n",
            "\n",
            "Mixtral: There once was a model, quite large you see,\n",
            "Made by AI, for all humanity.\n",
            "It spoke in tongues, from Old English to French,\n",
            "But sometimes its answers, would leave you in a lurch.\n",
            "\n",
            "It could write a poem, a story or more,\n",
            "But when asked for advice, it would often be a bore.\n",
            "It could summarize a book, in just a few lines,\n",
            "But when it comes to humour, it tends to malfunction and crack bad lines.\n",
            "\n",
            "It could answer a question, with a lot of thought,\n",
            "But sometimes it would get stuck, and the answer would be naught.\n",
            "It could imitate a human, in speech and in text,\n",
            "But when it comes to jokes, it tends to miss the context.\n",
            "\n",
            "It's a Large Language Model, so vast and so wide,\n",
            "But when it comes to humour, it's often left outside.\n",
            "It can do a lot of things, that would amaze you for sure,\n",
            "But please, don't take it seriously, it has no sense of humour!\n",
            "\n",
            "So here's to LLMs, so advanced and so cool,\n",
            "May they continue to learn, and not be a fool.\n",
            "Even if they struggle, to understand a good joke,\n",
            "They are still an amazing feat, for us to admire and evoke!\n",
            "\n",
            "\n",
            "User: Summarize your poem about Python in two sentences\n",
            "\n",
            "\n",
            "Mixtral: Python is a powerful and versatile programming language with a vast community and libraries, making it a top choice for both beginners and experts. It's praised for its clean and easy-to-understand syntax, as well as its ability to handle complex tasks with simplicity and efficiency.\n",
            "\n",
            "\n",
            "User: What is your favourite joke in your poem about Python?\n",
            "\n",
            "\n",
            "Mixtral: My personal favorite joke from the Python poem would be:\n",
            "\n",
            "\"It's a Swiss Army knife of coding, or perhaps a black belt,\n",
            "In the world of programming, it's the ultimate feel.\"\n",
            "\n",
            "This joke highlights the versatility and capabilities of the Python language, as well as the efficiency and ease of use it provides, making it a powerful tool for programmers. It's a fun and light-hearted way of saying that Python is a versatile and powerful language that can help you accomplish a lot with minimal effort.\n",
            "\n",
            "\n",
            "User: "
          ]
        }
      ]
    }
  ]
}