GPT-J-Inference.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "GPT-J-6B Inference Demo.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "TPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "a0b6bd601ad84c18af2fa1b8554abf94": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_a4f34a2cfe4a412bbdc65a70ac0b89a8",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_916b27dbeaa3461cb695f51daf190402",
              "IPY_MODEL_5e420d160a484a658b03a6a8fffc20cc"
            ]
          }
        },
        "a4f34a2cfe4a412bbdc65a70ac0b89a8": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "916b27dbeaa3461cb695f51daf190402": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_951c8411294642a09d6648823effe71f",
            "_dom_classes": [],
            "description": "Downloading: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 1042301,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 1042301,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_97b662ebe7a748c8af6479c1f0c4bc9f"
          }
        },
        "5e420d160a484a658b03a6a8fffc20cc": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_a7cfdce4a1a64cc6bcd2842dce96303f",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 1.04M/1.04M [00:04&lt;00:00, 215kB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_cdee93968c5f42cba4b7a866f8900c97"
          }
        },
        "951c8411294642a09d6648823effe71f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "97b662ebe7a748c8af6479c1f0c4bc9f": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "a7cfdce4a1a64cc6bcd2842dce96303f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "cdee93968c5f42cba4b7a866f8900c97": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "b4939af44ff14f7185c4ee219da80efa": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_a48492a16fd1452aaa519301e0a230d0",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_613e272e33914065bb89d38ddc6e6678",
              "IPY_MODEL_aa655ed683694c158af8f9fbffea478f"
            ]
          }
        },
        "a48492a16fd1452aaa519301e0a230d0": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "613e272e33914065bb89d38ddc6e6678": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_fda271a910e84563bc7542ae7a1ea966",
            "_dom_classes": [],
            "description": "Downloading: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 456318,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 456318,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_32b6ed93fac94d6bbf0556d952b303c3"
          }
        },
        "aa655ed683694c158af8f9fbffea478f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_266629ad54db4b959adf0ea15f81e0d7",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 456k/456k [00:01&lt;00:00, 255kB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_0656f2f4abc94a17b0c0ca532312ae3e"
          }
        },
        "fda271a910e84563bc7542ae7a1ea966": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "32b6ed93fac94d6bbf0556d952b303c3": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "266629ad54db4b959adf0ea15f81e0d7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "0656f2f4abc94a17b0c0ca532312ae3e": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "736045507c53436390c8d1b88472f770": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_6f6a6457a8604be8859d9de0f082def4",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_805c07778d3542bb8dfbe3daf3423d50",
              "IPY_MODEL_5c8e194b8fb44395b4b03a560e7f4c0e"
            ]
          }
        },
        "6f6a6457a8604be8859d9de0f082def4": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "805c07778d3542bb8dfbe3daf3423d50": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_2c6e73aff92649f181faa96e746779a4",
            "_dom_classes": [],
            "description": "Downloading: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 1355256,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 1355256,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_c447b1b8e99b4e2b97a69e035ccf7f49"
          }
        },
        "5c8e194b8fb44395b4b03a560e7f4c0e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_8b7b7e1da4c74bbcbb09f949d8bf25e5",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 1.36M/1.36M [00:01&lt;00:00, 1.26MB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_068220f70575498d85a03bde26fb4651"
          }
        },
        "2c6e73aff92649f181faa96e746779a4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "c447b1b8e99b4e2b97a69e035ccf7f49": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "8b7b7e1da4c74bbcbb09f949d8bf25e5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "068220f70575498d85a03bde26fb4651": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/srulikbd/Deep-Learning/blob/master/GPT-J-Inference.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pHIJVqHsh4An"
      },
      "source": [
        "# GPT-J-6B Inference Demo\n",
        "\n",
        "<a href=\"http://colab.research.google.com/github/kingoflolz/mesh-transformer-jax/blob/master/colab_demo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
        "\n",
        "This notebook demonstrates how to run the [GPT-J-6B model](https://github.com/kingoflolz/mesh-transformer-jax/#GPT-J-6B). See the link for more details about the model, including evaluation metrics and credits."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8CMw_dSQKfhT"
      },
      "source": [
        "## Install Dependencies\n",
        "\n",
        "First we download the model and install some dependencies. This step takes at least 5 minutes (possibly longer depending on server load).\n",
        "\n",
        "!!! **Make sure you are using a TPU runtime!** !!!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "n7xAFw-LOYfe",
        "outputId": "1915f110-ab3a-498a-9e4b-d22effa0fbaa"
      },
      "source": [
        "!apt install zstd\n",
        "\n",
        "# the \"slim\" version contain only bf16 weights and no optimizer parameters, which minimizes bandwidth and memory\n",
        "!time wget -c https://the-eye.eu/public/AI/GPT-J-6B/step_383500_slim.tar.zstd\n",
        "\n",
        "!time tar -I zstd -xf step_383500_slim.tar.zstd\n",
        "\n",
        "!git clone https://github.com/kingoflolz/mesh-transformer-jax.git\n",
        "!pip install -r mesh-transformer-jax/requirements.txt\n",
        "\n",
        "# jax 0.2.12 is required due to a regression with xmap in 0.2.13\n",
        "!pip install mesh-transformer-jax/ jax==0.2.12"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Reading package lists... Done\n",
            "Building dependency tree       \n",
            "Reading state information... Done\n",
            "The following NEW packages will be installed:\n",
            "  zstd\n",
            "0 upgraded, 1 newly installed, 0 to remove and 39 not upgraded.\n",
            "Need to get 278 kB of archives.\n",
            "After this operation, 1,141 kB of additional disk space will be used.\n",
            "Get:1 http://archive.ubuntu.com/ubuntu bionic-updates/universe amd64 zstd amd64 1.3.3+dfsg-2ubuntu1.2 [278 kB]\n",
            "Fetched 278 kB in 1s (374 kB/s)\n",
            "Selecting previously unselected package zstd.\n",
            "(Reading database ... 160772 files and directories currently installed.)\n",
            "Preparing to unpack .../zstd_1.3.3+dfsg-2ubuntu1.2_amd64.deb ...\n",
            "Unpacking zstd (1.3.3+dfsg-2ubuntu1.2) ...\n",
            "Setting up zstd (1.3.3+dfsg-2ubuntu1.2) ...\n",
            "Processing triggers for man-db (2.8.3-2ubuntu0.1) ...\n",
            "--2021-06-29 12:34:13--  https://the-eye.eu/public/AI/GPT-J-6B/step_383500_slim.tar.zstd\n",
            "Resolving the-eye.eu (the-eye.eu)... 162.213.130.242\n",
            "Connecting to the-eye.eu (the-eye.eu)|162.213.130.242|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 9414712325 (8.8G) [application/octet-stream]\n",
            "Saving to: ‘step_383500_slim.tar.zstd’\n",
            "\n",
            "step_383500_slim.ta 100%[===================>]   8.77G   117MB/s    in 78s     \n",
            "\n",
            "2021-06-29 12:35:32 (115 MB/s) - ‘step_383500_slim.tar.zstd’ saved [9414712325/9414712325]\n",
            "\n",
            "\n",
            "real\t1m18.238s\n",
            "user\t0m6.011s\n",
            "sys\t0m25.333s\n",
            "\n",
            "real\t1m30.133s\n",
            "user\t0m34.852s\n",
            "sys\t0m29.786s\n",
            "Cloning into 'mesh-transformer-jax'...\n",
            "remote: Enumerating objects: 549, done.\u001b[K\n",
            "remote: Counting objects: 100% (148/148), done.\u001b[K\n",
            "remote: Compressing objects: 100% (78/78), done.\u001b[K\n",
            "remote: Total 549 (delta 101), reused 94 (delta 70), pack-reused 401\u001b[K\n",
            "Receiving objects: 100% (549/549), 186.82 KiB | 2.71 MiB/s, done.\n",
            "Resolving deltas: 100% (349/349), done.\n",
            "Collecting git+https://github.com/deepmind/dm-haiku (from -r mesh-transformer-jax/requirements.txt (line 10))\n",
            "  Cloning https://github.com/deepmind/dm-haiku to /tmp/pip-req-build-0e3whn02\n",
            "  Running command git clone -q https://github.com/deepmind/dm-haiku /tmp/pip-req-build-0e3whn02\n",
            "Collecting git+https://github.com/EleutherAI/lm-evaluation-harness/ (from -r mesh-transformer-jax/requirements.txt (line 11))\n",
            "  Cloning https://github.com/EleutherAI/lm-evaluation-harness/ to /tmp/pip-req-build-wmtgqql1\n",
            "  Running command git clone -q https://github.com/EleutherAI/lm-evaluation-harness/ /tmp/pip-req-build-wmtgqql1\n",
            "Requirement already satisfied: numpy~=1.19.5 in /usr/local/lib/python3.7/dist-packages (from -r mesh-transformer-jax/requirements.txt (line 1)) (1.19.5)\n",
            "Collecting transformers~=4.4.2\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/ed/d5/f4157a376b8a79489a76ce6cfe147f4f3be1e029b7144fa7b8432e8acb26/transformers-4.4.2-py3-none-any.whl (2.0MB)\n",
            "\u001b[K     |████████████████████████████████| 2.0MB 6.9MB/s \n",
            "\u001b[?25hCollecting tqdm~=4.45.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/4a/1c/6359be64e8301b84160f6f6f7936bbfaaa5e9a4eab6cbc681db07600b949/tqdm-4.45.0-py2.py3-none-any.whl (60kB)\n",
            "\u001b[K     |████████████████████████████████| 61kB 6.9MB/s \n",
            "\u001b[?25hCollecting setuptools~=51.3.3\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/b2/81/509db0082c0d2ca2af307c6652ea422865de1f83c14b1e1f3549e415cfac/setuptools-51.3.3-py3-none-any.whl (786kB)\n",
            "\u001b[K     |████████████████████████████████| 788kB 23.3MB/s \n",
            "\u001b[?25hCollecting wandb~=0.10.22\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/e0/b4/9d92953d8cddc8450c859be12e3dbdd4c7754fb8def94c28b3b351c6ee4e/wandb-0.10.32-py2.py3-none-any.whl (1.8MB)\n",
            "\u001b[K     |████████████████████████████████| 1.8MB 29.4MB/s \n",
            "\u001b[?25hCollecting einops~=0.3.0\n",
            "  Downloading https://files.pythonhosted.org/packages/5d/a0/9935e030634bf60ecd572c775f64ace82ceddf2f504a5fd3902438f07090/einops-0.3.0-py2.py3-none-any.whl\n",
            "Collecting requests~=2.25.1\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/29/c1/24814557f1d22c56d50280771a17307e6bf87b70727d975fd6b2ce6b014a/requests-2.25.1-py2.py3-none-any.whl (61kB)\n",
            "\u001b[K     |████████████████████████████████| 61kB 7.6MB/s \n",
            "\u001b[?25hCollecting fabric~=2.6.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/c1/9d/59df62b620985871a4ba7d8b509b84340bbd1573257e55a427ae2df2d56e/fabric-2.6.0-py2.py3-none-any.whl (53kB)\n",
            "\u001b[K     |████████████████████████████████| 61kB 7.7MB/s \n",
            "\u001b[?25hCollecting optax==0.0.6\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/ec/7a/6259edd319ee7fa94dd23c54f15eff667f599d179e889af90fe0c204612c/optax-0.0.6-py3-none-any.whl (96kB)\n",
            "\u001b[K     |████████████████████████████████| 102kB 11.1MB/s \n",
            "\u001b[?25hCollecting ray~=1.2.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/11/14/15d0f0aec20a4674a996429160565a071688f27f49f789327ebed8188ffb/ray-1.2.0-cp37-cp37m-manylinux2014_x86_64.whl (47.5MB)\n",
            "\u001b[K     |████████████████████████████████| 47.5MB 123kB/s \n",
            "\u001b[?25hRequirement already satisfied: jax~=0.2.12 in /usr/local/lib/python3.7/dist-packages (from -r mesh-transformer-jax/requirements.txt (line 13)) (0.2.13)\n",
            "Requirement already satisfied: Flask~=1.1.2 in /usr/local/lib/python3.7/dist-packages (from -r mesh-transformer-jax/requirements.txt (line 14)) (1.1.4)\n",
            "Requirement already satisfied: cloudpickle~=1.3.0 in /usr/local/lib/python3.7/dist-packages (from -r mesh-transformer-jax/requirements.txt (line 15)) (1.3.0)\n",
            "Collecting tensorflow-cpu~=2.4.1\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/69/99/c4ffb7dec86049dcbe871262622caef477ff248963d0398d07c8d726a6de/tensorflow_cpu-2.4.2-cp37-cp37m-manylinux2010_x86_64.whl (144.3MB)\n",
            "\u001b[K     |████████████████████████████████| 144.3MB 50kB/s \n",
            "\u001b[?25hCollecting google-cloud-storage~=1.36.2\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/f2/0e/da07ffa511daa559bcc209f9344d71a90ba4d7b391fb795e6282f86d2935/google_cloud_storage-1.36.2-py2.py3-none-any.whl (97kB)\n",
            "\u001b[K     |████████████████████████████████| 102kB 11.2MB/s \n",
            "\u001b[?25hRequirement already satisfied: smart_open[gcs] in /usr/local/lib/python3.7/dist-packages (from -r mesh-transformer-jax/requirements.txt (line 18)) (5.1.0)\n",
            "Collecting func_timeout\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/b3/0d/bf0567477f7281d9a3926c582bfef21bff7498fc0ffd3e9de21811896a0b/func_timeout-4.3.5.tar.gz (44kB)\n",
            "\u001b[K     |████████████████████████████████| 51kB 6.2MB/s \n",
            "\u001b[?25hCollecting ftfy\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/af/da/d215a091986e5f01b80f5145cff6f22e2dc57c6b048aab2e882a07018473/ftfy-6.0.3.tar.gz (64kB)\n",
            "\u001b[K     |████████████████████████████████| 71kB 7.2MB/s \n",
            "\u001b[?25hRequirement already satisfied: absl-py>=0.7.1 in /usr/local/lib/python3.7/dist-packages (from dm-haiku==0.0.5.dev0->-r mesh-transformer-jax/requirements.txt (line 10)) (0.12.0)\n",
            "Collecting jmp>=0.0.2\n",
            "  Downloading https://files.pythonhosted.org/packages/ff/5c/1482f4a4a502e080af2ca54d7f80a60b5d4735f464c151666d583b78c226/jmp-0.0.2-py3-none-any.whl\n",
            "Requirement already satisfied: tabulate>=0.8.9 in /usr/local/lib/python3.7/dist-packages (from dm-haiku==0.0.5.dev0->-r mesh-transformer-jax/requirements.txt (line 10)) (0.8.9)\n",
            "Requirement already satisfied: typing_extensions in /usr/local/lib/python3.7/dist-packages (from dm-haiku==0.0.5.dev0->-r mesh-transformer-jax/requirements.txt (line 10)) (3.7.4.3)\n",
            "Collecting black==20.8b1\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/dc/7b/5a6bbe89de849f28d7c109f5ea87b65afa5124ad615f3419e71beb29dc96/black-20.8b1.tar.gz (1.1MB)\n",
            "\u001b[K     |████████████████████████████████| 1.1MB 50.8MB/s \n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "    Preparing wheel metadata ... \u001b[?25l\u001b[?25hdone\n",
            "Collecting best_download>=0.0.6\n",
            "  Downloading https://files.pythonhosted.org/packages/cd/8f/d8f2145cb0fd42ca55a6c6420d512156b6ff5d9857b3dd9cd0d679377cc5/best_download-0.0.7-py3-none-any.whl\n",
            "Collecting datasets>=1.2.1\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/08/a2/d4e1024c891506e1cee8f9d719d20831bac31cb5b7416983c4d2f65a6287/datasets-1.8.0-py3-none-any.whl (237kB)\n",
            "\u001b[K     |████████████████████████████████| 245kB 52.0MB/s \n",
            "\u001b[?25hRequirement already satisfied: click>=7.1 in /usr/local/lib/python3.7/dist-packages (from lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (7.1.2)\n",
            "Collecting scikit-learn>=0.24.1\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/a8/eb/a48f25c967526b66d5f1fa7a984594f0bf0a5afafa94a8c4dbc317744620/scikit_learn-0.24.2-cp37-cp37m-manylinux2010_x86_64.whl (22.3MB)\n",
            "\u001b[K     |████████████████████████████████| 22.3MB 1.2MB/s \n",
            "\u001b[?25hRequirement already satisfied: torch>=1.7 in /usr/local/lib/python3.7/dist-packages (from lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (1.9.0+cu102)\n",
            "Collecting sqlitedict==1.6.0\n",
            "  Downloading https://files.pythonhosted.org/packages/0f/1c/c757b93147a219cf1e25cef7e1ad9b595b7f802159493c45ce116521caff/sqlitedict-1.6.0.tar.gz\n",
            "Collecting pytablewriter==0.58.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/fd/e2/62b208cdb8771dee1849bd2b4ed129284e1efff7669985697e4c124c1000/pytablewriter-0.58.0-py3-none-any.whl (96kB)\n",
            "\u001b[K     |████████████████████████████████| 102kB 11.6MB/s \n",
            "\u001b[?25hCollecting sacrebleu==1.5.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/3b/7f/4fd83db8570288c3899d8e57666c2841403c15659f3d792a3cb8dc1c6689/sacrebleu-1.5.0-py3-none-any.whl (65kB)\n",
            "\u001b[K     |████████████████████████████████| 71kB 8.9MB/s \n",
            "\u001b[?25hCollecting pycountry==20.7.3\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/76/73/6f1a412f14f68c273feea29a6ea9b9f1e268177d32e0e69ad6790d306312/pycountry-20.7.3.tar.gz (10.1MB)\n",
            "\u001b[K     |████████████████████████████████| 10.1MB 52.1MB/s \n",
            "\u001b[?25hCollecting numexpr==2.7.2\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/9c/f4/fa8755c1aa44b431267aa019922f6cc9ec099cef0c6fc0ead0f9a2aa59e5/numexpr-2.7.2-cp37-cp37m-manylinux2010_x86_64.whl (471kB)\n",
            "\u001b[K     |████████████████████████████████| 481kB 51.1MB/s \n",
            "\u001b[?25hCollecting lm_dataformat==0.0.19\n",
            "  Downloading https://files.pythonhosted.org/packages/83/b5/8d10bf5a8082921792bb09c9d591dfd622cf4a16fbb7e283cc921c5ffc50/lm_dataformat-0.0.19-py3-none-any.whl\n",
            "Collecting pytest==6.2.3\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/76/4d/9c00146923da9f1cabd1878209d71b1380d537ec331a1a613e8f4b9d7985/pytest-6.2.3-py3-none-any.whl (280kB)\n",
            "\u001b[K     |████████████████████████████████| 286kB 54.7MB/s \n",
            "\u001b[?25hCollecting pybind11==2.6.2\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/8d/43/7339dbabbc2793718d59703aace4166f53c29ee1c202f6ff5bf8a26c4d91/pybind11-2.6.2-py2.py3-none-any.whl (191kB)\n",
            "\u001b[K     |████████████████████████████████| 194kB 58.2MB/s \n",
            "\u001b[?25hCollecting tqdm-multiprocess==0.0.11\n",
            "  Downloading https://files.pythonhosted.org/packages/25/7e/0d889fc6c84e3df6b69aaafe893fc77f69b3d968ac9ce574d1c62c688050/tqdm_multiprocess-0.0.11-py3-none-any.whl\n",
            "Collecting zstandard==0.15.2\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/5b/56/dc2a85d06e973f2ad96584b0e5b876d063135d449cb040aeaadd22a910f9/zstandard-0.15.2-cp37-cp37m-manylinux2014_x86_64.whl (2.2MB)\n",
            "\u001b[K     |████████████████████████████████| 2.2MB 52.1MB/s \n",
            "\u001b[?25hCollecting jsonlines==2.0.0\n",
            "  Downloading https://files.pythonhosted.org/packages/d4/58/06f430ff7607a2929f80f07bfd820acbc508a4e977542fefcc522cde9dff/jsonlines-2.0.0-py3-none-any.whl\n",
            "Collecting mock==4.0.3\n",
            "  Downloading https://files.pythonhosted.org/packages/5c/03/b7e605db4a57c0f6fba744b11ef3ddf4ddebcada35022927a2b5fc623fdf/mock-4.0.3-py3-none-any.whl\n",
            "Collecting openai==0.6.4\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/00/6b/ec780cfdd44ce0d7e56aae12477fe163ad152078cfdb744dfccab4ddae4b/openai-0.6.4.tar.gz (159kB)\n",
            "\u001b[K     |████████████████████████████████| 163kB 57.7MB/s \n",
            "\u001b[?25hCollecting tokenizers<0.11,>=0.10.1\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/d4/e2/df3543e8ffdab68f5acc73f613de9c2b155ac47f162e725dcac87c521c11/tokenizers-0.10.3-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (3.3MB)\n",
            "\u001b[K     |████████████████████████████████| 3.3MB 52.9MB/s \n",
            "\u001b[?25hRequirement already satisfied: importlib-metadata; python_version < \"3.8\" in /usr/local/lib/python3.7/dist-packages (from transformers~=4.4.2->-r mesh-transformer-jax/requirements.txt (line 2)) (4.5.0)\n",
            "Collecting sacremoses\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/75/ee/67241dc87f266093c533a2d4d3d69438e57d7a90abb216fa076e7d475d4a/sacremoses-0.0.45-py3-none-any.whl (895kB)\n",
            "\u001b[K     |████████████████████████████████| 901kB 45.2MB/s \n",
            "\u001b[?25hRequirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from transformers~=4.4.2->-r mesh-transformer-jax/requirements.txt (line 2)) (20.9)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.7/dist-packages (from transformers~=4.4.2->-r mesh-transformer-jax/requirements.txt (line 2)) (2019.12.20)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.7/dist-packages (from transformers~=4.4.2->-r mesh-transformer-jax/requirements.txt (line 2)) (3.0.12)\n",
            "Requirement already satisfied: protobuf>=3.12.0 in /usr/local/lib/python3.7/dist-packages (from wandb~=0.10.22->-r mesh-transformer-jax/requirements.txt (line 5)) (3.12.4)\n",
            "Requirement already satisfied: six>=1.13.0 in /usr/local/lib/python3.7/dist-packages (from wandb~=0.10.22->-r mesh-transformer-jax/requirements.txt (line 5)) (1.15.0)\n",
            "Collecting pathtools\n",
            "  Downloading https://files.pythonhosted.org/packages/e7/7f/470d6fcdf23f9f3518f6b0b76be9df16dcc8630ad409947f8be2eb0ed13a/pathtools-0.1.2.tar.gz\n",
            "Requirement already satisfied: promise<3,>=2.0 in /usr/local/lib/python3.7/dist-packages (from wandb~=0.10.22->-r mesh-transformer-jax/requirements.txt (line 5)) (2.3)\n",
            "Collecting shortuuid>=0.5.0\n",
            "  Downloading https://files.pythonhosted.org/packages/25/a6/2ecc1daa6a304e7f1b216f0896b26156b78e7c38e1211e9b798b4716c53d/shortuuid-1.0.1-py3-none-any.whl\n",
            "Collecting GitPython>=1.0.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/bc/91/b38c4fabb6e5092ab23492ded4f318ab7299b19263272b703478038c0fbc/GitPython-3.1.18-py3-none-any.whl (170kB)\n",
            "\u001b[K     |████████████████████████████████| 174kB 53.7MB/s \n",
            "\u001b[?25hCollecting sentry-sdk>=0.4.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/1c/4a/a54b254f67d8f4052338d54ebe90126f200693440a93ef76d254d581e3ec/sentry_sdk-1.1.0-py2.py3-none-any.whl (131kB)\n",
            "\u001b[K     |████████████████████████████████| 133kB 53.3MB/s \n",
            "\u001b[?25hRequirement already satisfied: PyYAML in /usr/local/lib/python3.7/dist-packages (from wandb~=0.10.22->-r mesh-transformer-jax/requirements.txt (line 5)) (3.13)\n",
            "Requirement already satisfied: python-dateutil>=2.6.1 in /usr/local/lib/python3.7/dist-packages (from wandb~=0.10.22->-r mesh-transformer-jax/requirements.txt (line 5)) (2.8.1)\n",
            "Collecting docker-pycreds>=0.4.0\n",
            "  Downloading https://files.pythonhosted.org/packages/f5/e8/f6bd1eee09314e7e6dee49cbe2c5e22314ccdb38db16c9fc72d2fa80d054/docker_pycreds-0.4.0-py2.py3-none-any.whl\n",
            "Collecting configparser>=3.8.1\n",
            "  Downloading https://files.pythonhosted.org/packages/fd/01/ff260a18caaf4457eb028c96eeb405c4a230ca06c8ec9c1379f813caa52e/configparser-5.0.2-py3-none-any.whl\n",
            "Requirement already satisfied: psutil>=5.0.0 in /usr/local/lib/python3.7/dist-packages (from wandb~=0.10.22->-r mesh-transformer-jax/requirements.txt (line 5)) (5.4.8)\n",
            "Collecting subprocess32>=3.5.3\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/32/c8/564be4d12629b912ea431f1a50eb8b3b9d00f1a0b1ceff17f266be190007/subprocess32-3.5.4.tar.gz (97kB)\n",
            "\u001b[K     |████████████████████████████████| 102kB 11.5MB/s \n",
            "\u001b[?25hRequirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests~=2.25.1->-r mesh-transformer-jax/requirements.txt (line 7)) (2.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests~=2.25.1->-r mesh-transformer-jax/requirements.txt (line 7)) (2021.5.30)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests~=2.25.1->-r mesh-transformer-jax/requirements.txt (line 7)) (1.24.3)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests~=2.25.1->-r mesh-transformer-jax/requirements.txt (line 7)) (3.0.4)\n",
            "Collecting pathlib2\n",
            "  Downloading https://files.pythonhosted.org/packages/e9/45/9c82d3666af4ef9f221cbb954e1d77ddbb513faf552aea6df5f37f1a4859/pathlib2-2.3.5-py2.py3-none-any.whl\n",
            "Collecting paramiko>=2.4\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/95/19/124e9287b43e6ff3ebb9cdea3e5e8e88475a873c05ccdf8b7e20d2c4201e/paramiko-2.7.2-py2.py3-none-any.whl (206kB)\n",
            "\u001b[K     |████████████████████████████████| 215kB 54.5MB/s \n",
            "\u001b[?25hCollecting invoke<2.0,>=1.3\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/87/8f/c153d7db091f342da6bc97f7bedd1b2ce2867c4a8b0aab40fbba85a05e33/invoke-1.5.0-py3-none-any.whl (211kB)\n",
            "\u001b[K     |████████████████████████████████| 215kB 55.4MB/s \n",
            "\u001b[?25hCollecting chex>=0.0.4\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/f5/b9/445eb59ec23249acffc5322c79b07e20b12dbff45b9c1da6cdae9e947685/chex-0.0.7-py3-none-any.whl (52kB)\n",
            "\u001b[K     |████████████████████████████████| 61kB 7.9MB/s \n",
            "\u001b[?25hRequirement already satisfied: jaxlib>=0.1.37 in /usr/local/lib/python3.7/dist-packages (from optax==0.0.6->-r mesh-transformer-jax/requirements.txt (line 9)) (0.1.66+cuda110)\n",
            "Collecting aiohttp-cors\n",
            "  Downloading https://files.pythonhosted.org/packages/13/e7/e436a0c0eb5127d8b491a9b83ecd2391c6ff7dcd5548dfaec2080a2340fd/aiohttp_cors-0.7.0-py3-none-any.whl\n",
            "Collecting opencensus\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/18/59/12044123133d000f705383ad98579aeb0dd82d66b33a254a21b54bf0d6bb/opencensus-0.7.13-py2.py3-none-any.whl (127kB)\n",
            "\u001b[K     |████████████████████████████████| 133kB 58.5MB/s \n",
            "\u001b[?25hCollecting redis>=3.5.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/a7/7c/24fb0511df653cf1a5d938d8f5d19802a88cef255706fdda242ff97e91b7/redis-3.5.3-py2.py3-none-any.whl (72kB)\n",
            "\u001b[K     |████████████████████████████████| 81kB 10.0MB/s \n",
            "\u001b[?25hRequirement already satisfied: jsonschema in /usr/local/lib/python3.7/dist-packages (from ray~=1.2.0->-r mesh-transformer-jax/requirements.txt (line 12)) (2.6.0)\n",
            "Collecting py-spy>=0.2.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/9d/4d/1a9cbe9a0b543e6733cb38afe26451522a9ef8e4897b59e74cc76838f245/py_spy-0.3.7-py2.py3-none-manylinux1_x86_64.whl (3.1MB)\n",
            "\u001b[K     |████████████████████████████████| 3.1MB 48.8MB/s \n",
            "\u001b[?25hRequirement already satisfied: grpcio>=1.28.1 in /usr/local/lib/python3.7/dist-packages (from ray~=1.2.0->-r mesh-transformer-jax/requirements.txt (line 12)) (1.34.1)\n",
            "Collecting gpustat\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/b4/69/d8c849715171aeabd61af7da080fdc60948b5a396d2422f1f4672e43d008/gpustat-0.6.0.tar.gz (78kB)\n",
            "\u001b[K     |████████████████████████████████| 81kB 9.5MB/s \n",
            "\u001b[?25hCollecting colorama\n",
            "  Downloading https://files.pythonhosted.org/packages/44/98/5b86278fbbf250d239ae0ecb724f8572af1c91f4a11edf4d36a206189440/colorama-0.4.4-py2.py3-none-any.whl\n",
            "Requirement already satisfied: prometheus-client>=0.7.1 in /usr/local/lib/python3.7/dist-packages (from ray~=1.2.0->-r mesh-transformer-jax/requirements.txt (line 12)) (0.11.0)\n",
            "Collecting aioredis\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/b0/64/1b1612d0a104f21f80eb4c6e1b6075f2e6aba8e228f46f229cfd3fdac859/aioredis-1.3.1-py3-none-any.whl (65kB)\n",
            "\u001b[K     |████████████████████████████████| 71kB 8.7MB/s \n",
            "\u001b[?25hCollecting aiohttp\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/88/c0/5890b4c8b04a79b7360e8fe4490feb0bb3ab179743f199f0e6220cebd568/aiohttp-3.7.4.post0-cp37-cp37m-manylinux2014_x86_64.whl (1.3MB)\n",
            "\u001b[K     |████████████████████████████████| 1.3MB 51.0MB/s \n",
            "\u001b[?25hCollecting colorful\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/b0/8e/e386e248266952d24d73ed734c2f5513f34d9557032618c8910e605dfaf6/colorful-0.5.4-py2.py3-none-any.whl (201kB)\n",
            "\u001b[K     |████████████████████████████████| 204kB 56.7MB/s \n",
            "\u001b[?25hRequirement already satisfied: msgpack<2.0.0,>=1.0.0 in /usr/local/lib/python3.7/dist-packages (from ray~=1.2.0->-r mesh-transformer-jax/requirements.txt (line 12)) (1.0.2)\n",
            "Requirement already satisfied: opt-einsum in /usr/local/lib/python3.7/dist-packages (from jax~=0.2.12->-r mesh-transformer-jax/requirements.txt (line 13)) (3.3.0)\n",
            "Requirement already satisfied: Werkzeug<2.0,>=0.15 in /usr/local/lib/python3.7/dist-packages (from Flask~=1.1.2->-r mesh-transformer-jax/requirements.txt (line 14)) (1.0.1)\n",
            "Requirement already satisfied: Jinja2<3.0,>=2.10.1 in /usr/local/lib/python3.7/dist-packages (from Flask~=1.1.2->-r mesh-transformer-jax/requirements.txt (line 14)) (2.11.3)\n",
            "Requirement already satisfied: itsdangerous<2.0,>=0.24 in /usr/local/lib/python3.7/dist-packages (from Flask~=1.1.2->-r mesh-transformer-jax/requirements.txt (line 14)) (1.1.0)\n",
            "Collecting gast==0.3.3\n",
            "  Downloading https://files.pythonhosted.org/packages/d6/84/759f5dd23fec8ba71952d97bcc7e2c9d7d63bdc582421f3cd4be845f0c98/gast-0.3.3-py2.py3-none-any.whl\n",
            "Collecting tensorflow-estimator<2.5.0,>=2.4.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/74/7e/622d9849abf3afb81e482ffc170758742e392ee129ce1540611199a59237/tensorflow_estimator-2.4.0-py2.py3-none-any.whl (462kB)\n",
            "\u001b[K     |████████████████████████████████| 471kB 54.6MB/s \n",
            "\u001b[?25hRequirement already satisfied: termcolor~=1.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.4.1->-r mesh-transformer-jax/requirements.txt (line 16)) (1.1.0)\n",
            "Requirement already satisfied: wheel~=0.35 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.4.1->-r mesh-transformer-jax/requirements.txt (line 16)) (0.36.2)\n",
            "Requirement already satisfied: google-pasta~=0.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.4.1->-r mesh-transformer-jax/requirements.txt (line 16)) (0.2.0)\n",
            "Requirement already satisfied: keras-preprocessing~=1.1.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.4.1->-r mesh-transformer-jax/requirements.txt (line 16)) (1.1.2)\n",
            "Collecting h5py~=2.10.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/3f/c0/abde58b837e066bca19a3f7332d9d0493521d7dd6b48248451a9e3fe2214/h5py-2.10.0-cp37-cp37m-manylinux1_x86_64.whl (2.9MB)\n",
            "\u001b[K     |████████████████████████████████| 2.9MB 53.9MB/s \n",
            "\u001b[?25hRequirement already satisfied: wrapt~=1.12.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.4.1->-r mesh-transformer-jax/requirements.txt (line 16)) (1.12.1)\n",
            "Requirement already satisfied: tensorboard~=2.4 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.4.1->-r mesh-transformer-jax/requirements.txt (line 16)) (2.5.0)\n",
            "Requirement already satisfied: flatbuffers~=1.12.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.4.1->-r mesh-transformer-jax/requirements.txt (line 16)) (1.12)\n",
            "Requirement already satisfied: astunparse~=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.4.1->-r mesh-transformer-jax/requirements.txt (line 16)) (1.6.3)\n",
            "Collecting google-cloud-core<2.0dev,>=1.4.1\n",
            "  Downloading https://files.pythonhosted.org/packages/88/a7/b74266a6fd888d91a6f517c574c17425731181fe44e2df1e414d4b77fbe8/google_cloud_core-1.7.1-py2.py3-none-any.whl\n",
            "Collecting google-resumable-media<2.0dev,>=1.2.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/11/96/4360dc70bef5559b3faf3deeda97aae7d10ff7660d41fd233eb792e7d09f/google_resumable_media-1.3.1-py2.py3-none-any.whl (75kB)\n",
            "\u001b[K     |████████████████████████████████| 81kB 9.9MB/s \n",
            "\u001b[?25hRequirement already satisfied: google-auth<2.0dev,>=1.11.0 in /usr/local/lib/python3.7/dist-packages (from google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 17)) (1.31.0)\n",
            "Requirement already satisfied: wcwidth in /usr/local/lib/python3.7/dist-packages (from ftfy->-r mesh-transformer-jax/requirements.txt (line 20)) (0.2.5)\n",
            "Requirement already satisfied: toml>=0.10.1 in /usr/local/lib/python3.7/dist-packages (from black==20.8b1->lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (0.10.2)\n",
            "Collecting mypy-extensions>=0.4.3\n",
            "  Downloading https://files.pythonhosted.org/packages/5c/eb/975c7c080f3223a5cdaff09612f3a5221e4ba534f7039db34c35d95fa6a5/mypy_extensions-0.4.3-py2.py3-none-any.whl\n",
            "Collecting typed-ast>=1.4.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/65/b3/573d2f1fecbbe8f82a8d08172e938c247f99abe1be3bef3da2efaa3810bf/typed_ast-1.4.3-cp37-cp37m-manylinux1_x86_64.whl (743kB)\n",
            "\u001b[K     |████████████████████████████████| 747kB 54.2MB/s \n",
            "\u001b[?25hCollecting pathspec<1,>=0.6\n",
            "  Downloading https://files.pythonhosted.org/packages/29/29/a465741a3d97ea3c17d21eaad4c64205428bde56742360876c4391f930d4/pathspec-0.8.1-py2.py3-none-any.whl\n",
            "Requirement already satisfied: appdirs in /usr/local/lib/python3.7/dist-packages (from black==20.8b1->lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (1.4.4)\n",
            "Collecting rehash\n",
            "  Downloading https://files.pythonhosted.org/packages/c9/e4/30db193232d9e9c8e123764d84f0807535677548833ca251556ad6134c24/rehash-1.0.0-py2.py3-none-any.whl\n",
            "Requirement already satisfied: multiprocess in /usr/local/lib/python3.7/dist-packages (from datasets>=1.2.1->lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (0.70.12.2)\n",
            "Collecting fsspec\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/0e/3a/666e63625a19883ae8e1674099e631f9737bd5478c4790e5ad49c5ac5261/fsspec-2021.6.1-py3-none-any.whl (115kB)\n",
            "\u001b[K     |████████████████████████████████| 122kB 57.9MB/s \n",
            "\u001b[?25hCollecting huggingface-hub<0.1.0\n",
            "  Downloading https://files.pythonhosted.org/packages/35/03/071adc023c0a7e540cf4652fa9cad13ab32e6ae469bf0cc0262045244812/huggingface_hub-0.0.13-py3-none-any.whl\n",
            "Collecting xxhash\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/7d/4f/0a862cad26aa2ed7a7cd87178cbbfa824fc1383e472d63596a0d018374e7/xxhash-2.0.2-cp37-cp37m-manylinux2010_x86_64.whl (243kB)\n",
            "\u001b[K     |████████████████████████████████| 245kB 55.5MB/s \n",
            "\u001b[?25hRequirement already satisfied: pyarrow<4.0.0,>=1.0.0 in /usr/local/lib/python3.7/dist-packages (from datasets>=1.2.1->lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (3.0.0)\n",
            "Requirement already satisfied: dill in /usr/local/lib/python3.7/dist-packages (from datasets>=1.2.1->lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (0.3.4)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.7/dist-packages (from datasets>=1.2.1->lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (1.1.5)\n",
            "Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.7/dist-packages (from scikit-learn>=0.24.1->lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (1.0.1)\n",
            "Requirement already satisfied: scipy>=0.19.1 in /usr/local/lib/python3.7/dist-packages (from scikit-learn>=0.24.1->lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (1.4.1)\n",
            "Collecting threadpoolctl>=2.0.0\n",
            "  Downloading https://files.pythonhosted.org/packages/f7/12/ec3f2e203afa394a149911729357aa48affc59c20e2c1c8297a60f33f133/threadpoolctl-2.1.0-py3-none-any.whl\n",
            "Collecting mbstrdecoder<2,>=1.0.0\n",
            "  Downloading https://files.pythonhosted.org/packages/e8/f6/0e6bb50c3c6380a4982c87d80e70b2f6e366523a57a0c58594aea472206d/mbstrdecoder-1.0.1-py3-none-any.whl\n",
            "Collecting pathvalidate<3,>=2.3.0\n",
            "  Downloading https://files.pythonhosted.org/packages/87/55/7d63b78986f1f8764180b84ee3e8a47c583ec059d32c98d8fba7fc0dc1ae/pathvalidate-2.4.1-py3-none-any.whl\n",
            "Collecting typepy[datetime]<2,>=1.1.1\n",
            "  Downloading https://files.pythonhosted.org/packages/60/3a/1239e59924250d9c2dd1d5b84748da82d15aaa241b3ceeffa08aa5eba589/typepy-1.1.5-py3-none-any.whl\n",
            "Collecting tabledata<2,>=1.1.3\n",
            "  Downloading https://files.pythonhosted.org/packages/85/93/4c695da7e6589e1e4b513c02d5b562dcc5afacb8a2f6cac8eb2ac2e88833/tabledata-1.1.4-py3-none-any.whl\n",
            "Collecting msgfy<1,>=0.1.0\n",
            "  Downloading https://files.pythonhosted.org/packages/48/52/c4441871514276e7c4cb51c122e663b5ef19dc20030f6ab7723071118464/msgfy-0.1.0-py3-none-any.whl\n",
            "Collecting DataProperty<2,>=0.50.0\n",
            "  Downloading https://files.pythonhosted.org/packages/b9/5f/c773c362fcba227d6a4021225cb0213b51849c7dc9c93004d34d9004078b/DataProperty-0.50.1-py3-none-any.whl\n",
            "Collecting tcolorpy<1,>=0.0.5\n",
            "  Downloading https://files.pythonhosted.org/packages/96/73/2a73a7d53df3708636dae4e817814d07e455efd53897476f3863925cf0af/tcolorpy-0.1.1-py3-none-any.whl\n",
            "Collecting portalocker\n",
            "  Downloading https://files.pythonhosted.org/packages/68/33/cb524f4de298509927b90aa5ee34767b9a2b93e663cf354b2a3efa2b4acd/portalocker-2.3.0-py2.py3-none-any.whl\n",
            "Collecting ujson\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/17/4e/50e8e4cf5f00b537095711c2c86ac4d7191aed2b4fffd5a19f06898f6929/ujson-4.0.2-cp37-cp37m-manylinux1_x86_64.whl (179kB)\n",
            "\u001b[K     |████████████████████████████████| 184kB 58.5MB/s \n",
            "\u001b[?25hRequirement already satisfied: attrs>=19.2.0 in /usr/local/lib/python3.7/dist-packages (from pytest==6.2.3->lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (21.2.0)\n",
            "Requirement already satisfied: iniconfig in /usr/local/lib/python3.7/dist-packages (from pytest==6.2.3->lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (1.1.1)\n",
            "Requirement already satisfied: py>=1.8.2 in /usr/local/lib/python3.7/dist-packages (from pytest==6.2.3->lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (1.10.0)\n",
            "Collecting pluggy<1.0.0a1,>=0.12\n",
            "  Downloading https://files.pythonhosted.org/packages/a0/28/85c7aa31b80d150b772fbe4a229487bc6644da9ccb7e427dd8cc60cb8a62/pluggy-0.13.1-py2.py3-none-any.whl\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata; python_version < \"3.8\"->transformers~=4.4.2->-r mesh-transformer-jax/requirements.txt (line 2)) (3.4.1)\n",
            "Requirement already satisfied: pyparsing>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->transformers~=4.4.2->-r mesh-transformer-jax/requirements.txt (line 2)) (2.4.7)\n",
            "Collecting gitdb<5,>=4.0.1\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/ea/e8/f414d1a4f0bbc668ed441f74f44c116d9816833a48bf81d22b697090dba8/gitdb-4.0.7-py3-none-any.whl (63kB)\n",
            "\u001b[K     |████████████████████████████████| 71kB 8.8MB/s \n",
            "\u001b[?25hCollecting pynacl>=1.0.1\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/9d/57/2f5e6226a674b2bcb6db531e8b383079b678df5b10cdaa610d6cf20d77ba/PyNaCl-1.4.0-cp35-abi3-manylinux1_x86_64.whl (961kB)\n",
            "\u001b[K     |████████████████████████████████| 962kB 51.1MB/s \n",
            "\u001b[?25hCollecting bcrypt>=3.1.3\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/26/70/6d218afbe4c73538053c1016dd631e8f25fffc10cd01f5c272d7acf3c03d/bcrypt-3.2.0-cp36-abi3-manylinux2010_x86_64.whl (63kB)\n",
            "\u001b[K     |████████████████████████████████| 71kB 9.2MB/s \n",
            "\u001b[?25hCollecting cryptography>=2.5\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/b2/26/7af637e6a7e87258b963f1731c5982fb31cd507f0d90d91836e446955d02/cryptography-3.4.7-cp36-abi3-manylinux2014_x86_64.whl (3.2MB)\n",
            "\u001b[K     |████████████████████████████████| 3.2MB 53.9MB/s \n",
            "\u001b[?25hRequirement already satisfied: toolz>=0.9.0 in /usr/local/lib/python3.7/dist-packages (from chex>=0.0.4->optax==0.0.6->-r mesh-transformer-jax/requirements.txt (line 9)) (0.11.1)\n",
            "Requirement already satisfied: dm-tree>=0.1.5 in /usr/local/lib/python3.7/dist-packages (from chex>=0.0.4->optax==0.0.6->-r mesh-transformer-jax/requirements.txt (line 9)) (0.1.6)\n",
            "Requirement already satisfied: google-api-core<2.0.0,>=1.0.0 in /usr/local/lib/python3.7/dist-packages (from opencensus->ray~=1.2.0->-r mesh-transformer-jax/requirements.txt (line 12)) (1.26.3)\n",
            "Collecting opencensus-context==0.1.2\n",
            "  Downloading https://files.pythonhosted.org/packages/f1/33/990f1bd9e7ee770fc8d3c154fc24743a96f16a0e49e14e1b7540cc2fdd93/opencensus_context-0.1.2-py2.py3-none-any.whl\n",
            "Requirement already satisfied: nvidia-ml-py3>=7.352.0 in /usr/local/lib/python3.7/dist-packages (from gpustat->ray~=1.2.0->-r mesh-transformer-jax/requirements.txt (line 12)) (7.352.0)\n",
            "Collecting blessings>=1.6\n",
            "  Downloading https://files.pythonhosted.org/packages/03/74/489f85a78247609c6b4f13733cbf3ba0d864b11aa565617b645d6fdf2a4a/blessings-1.7-py3-none-any.whl\n",
            "Collecting hiredis\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/ed/33/290cea35b09c80b4634773ad5572a8030a87b5d39736719f698f521d2a13/hiredis-2.0.0-cp37-cp37m-manylinux2010_x86_64.whl (85kB)\n",
            "\u001b[K     |████████████████████████████████| 92kB 11.3MB/s \n",
            "\u001b[?25hCollecting async-timeout\n",
            "  Downloading https://files.pythonhosted.org/packages/e1/1e/5a4441be21b0726c4464f3f23c8b19628372f606755a9d2e46c187e65ec4/async_timeout-3.0.1-py3-none-any.whl\n",
            "Collecting multidict<7.0,>=4.5\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/7c/a6/4123b8165acbe773d1a8dc8e3f0d1edea16d29f7de018eda769abb56bd30/multidict-5.1.0-cp37-cp37m-manylinux2014_x86_64.whl (142kB)\n",
            "\u001b[K     |████████████████████████████████| 143kB 57.3MB/s \n",
            "\u001b[?25hCollecting yarl<2.0,>=1.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/f1/62/046834c5fc998c88ab2ef722f5d42122230a632212c8afa76418324f53ff/yarl-1.6.3-cp37-cp37m-manylinux2014_x86_64.whl (294kB)\n",
            "\u001b[K     |████████████████████████████████| 296kB 59.2MB/s \n",
            "\u001b[?25hRequirement already satisfied: MarkupSafe>=0.23 in /usr/local/lib/python3.7/dist-packages (from Jinja2<3.0,>=2.10.1->Flask~=1.1.2->-r mesh-transformer-jax/requirements.txt (line 14)) (2.0.1)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.4->tensorflow-cpu~=2.4.1->-r mesh-transformer-jax/requirements.txt (line 16)) (3.3.4)\n",
            "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.4->tensorflow-cpu~=2.4.1->-r mesh-transformer-jax/requirements.txt (line 16)) (0.4.4)\n",
            "Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.4->tensorflow-cpu~=2.4.1->-r mesh-transformer-jax/requirements.txt (line 16)) (0.6.1)\n",
            "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.4->tensorflow-cpu~=2.4.1->-r mesh-transformer-jax/requirements.txt (line 16)) (1.8.0)\n",
            "Collecting google-crc32c<2.0dev,>=1.0; python_version >= \"3.5\"\n",
            "  Downloading https://files.pythonhosted.org/packages/fc/ae/b6efa1019e18c6c791f0f5cd93b2ff40f8f06696dbf04db39ec0f5591b1e/google_crc32c-1.1.2-cp37-cp37m-manylinux2014_x86_64.whl\n",
            "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from google-auth<2.0dev,>=1.11.0->google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 17)) (4.2.2)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from google-auth<2.0dev,>=1.11.0->google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 17)) (0.2.8)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4; python_version >= \"3.6\" in /usr/local/lib/python3.7/dist-packages (from google-auth<2.0dev,>=1.11.0->google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 17)) (4.7.2)\n",
            "Requirement already satisfied: pytz>=2017.2 in /usr/local/lib/python3.7/dist-packages (from pandas->datasets>=1.2.1->lm-eval-harness==0.0.1->-r mesh-transformer-jax/requirements.txt (line 11)) (2018.9)\n",
            "Collecting smmap<5,>=3.0.1\n",
            "  Downloading https://files.pythonhosted.org/packages/68/ee/d540eb5e5996eb81c26ceffac6ee49041d473bc5125f2aa995cf51ec1cf1/smmap-4.0.0-py2.py3-none-any.whl\n",
            "Requirement already satisfied: cffi>=1.4.1 in /usr/local/lib/python3.7/dist-packages (from pynacl>=1.0.1->paramiko>=2.4->fabric~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 8)) (1.14.5)\n",
            "Requirement already satisfied: googleapis-common-protos<2.0dev,>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from google-api-core<2.0.0,>=1.0.0->opencensus->ray~=1.2.0->-r mesh-transformer-jax/requirements.txt (line 12)) (1.53.0)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.7/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard~=2.4->tensorflow-cpu~=2.4.1->-r mesh-transformer-jax/requirements.txt (line 16)) (1.3.0)\n",
            "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.7/dist-packages (from pyasn1-modules>=0.2.1->google-auth<2.0dev,>=1.11.0->google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 17)) (0.4.8)\n",
            "Requirement already satisfied: pycparser in /usr/local/lib/python3.7/dist-packages (from cffi>=1.4.1->pynacl>=1.0.1->paramiko>=2.4->fabric~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 8)) (2.20)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.7/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard~=2.4->tensorflow-cpu~=2.4.1->-r mesh-transformer-jax/requirements.txt (line 16)) (3.1.1)\n",
            "Building wheels for collected packages: black\n",
            "  Building wheel for black (PEP 517) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for black: filename=black-20.8b1-cp37-none-any.whl size=124195 sha256=0ee9e82ba444917fd7b83b79ef01d7a2d61ae9644cf8cfcf7d1cafb9e1168ea3\n",
            "  Stored in directory: /root/.cache/pip/wheels/6e/10/b5/edf7359c2edd0305cce7e3f96e07daf7ce55dceac9d3ce3373\n",
            "Successfully built black\n",
            "Building wheels for collected packages: func-timeout, ftfy, dm-haiku, lm-eval-harness, sqlitedict, pycountry, openai, pathtools, subprocess32, gpustat\n",
            "  Building wheel for func-timeout (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for func-timeout: filename=func_timeout-4.3.5-cp37-none-any.whl size=15097 sha256=8d924cb214efe476e9472ba7d3597809fe657d298bb9f0521fa62a1c1de65436\n",
            "  Stored in directory: /root/.cache/pip/wheels/46/7c/4f/24f1d2d5bbff92219debe7ea19af84f76ddeb90dd4ec544f26\n",
            "  Building wheel for ftfy (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for ftfy: filename=ftfy-6.0.3-cp37-none-any.whl size=41935 sha256=2babe41c535360b4ccfdb94551c4337791b4ed1b93675a6d251e1010f43aa954\n",
            "  Stored in directory: /root/.cache/pip/wheels/99/2c/e6/109c8a28fef7a443f67ba58df21fe1d0067ac3322e75e6b0b7\n",
            "  Building wheel for dm-haiku (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for dm-haiku: filename=dm_haiku-0.0.5.dev0-cp37-none-any.whl size=530361 sha256=acabf390fd440111e819e6107e67b8ba3cdedf56d653242696b2793f3b9da8e1\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-d_dltioq/wheels/97/0f/e9/17f34e377f8d4060fa88a7e82bee5d8afbf7972384768a5499\n",
            "  Building wheel for lm-eval-harness (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for lm-eval-harness: filename=lm_eval_harness-0.0.1-cp37-none-any.whl size=100737 sha256=148982c790babfd1be2d5d166d6d00fb1625ef2b0b9590f9bee3397195a983c1\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-d_dltioq/wheels/a8/db/b4/32ca8efd6b64f9187fefbabd636d7c94cd2150a657262ee22a\n",
            "  Building wheel for sqlitedict (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for sqlitedict: filename=sqlitedict-1.6.0-cp37-none-any.whl size=14713 sha256=09c5498db026d0da132f7ec495ca36bac8ac8f9566c0347d3fb25f7ba7c1202b\n",
            "  Stored in directory: /root/.cache/pip/wheels/bd/57/d3/907c3ee02d35e66f674ad0106e61f06eeeb98f6ee66a6cc3fe\n",
            "  Building wheel for pycountry (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pycountry: filename=pycountry-20.7.3-py2.py3-none-any.whl size=10746883 sha256=641699e41a985b077d70892dca41c0001aefc813f3a8563ea9cf85ddc041ef56\n",
            "  Stored in directory: /root/.cache/pip/wheels/33/4e/a6/be297e6b83567e537bed9df4a93f8590ec01c1acfbcd405348\n",
            "  Building wheel for openai (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for openai: filename=openai-0.6.4-cp37-none-any.whl size=172215 sha256=f466c3783b49c41c36c1c30fb047fb97838733117e92250bc1015393c62bad77\n",
            "  Stored in directory: /root/.cache/pip/wheels/f9/a3/f5/00a714fabfbe3389a25242c2f91b626dd1d8fbd63b35a96730\n",
            "  Building wheel for pathtools (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pathtools: filename=pathtools-0.1.2-cp37-none-any.whl size=8807 sha256=670362f3ea3169f7a5a80c182d241831fd709307e33d420a7e60b4d08ec0c020\n",
            "  Stored in directory: /root/.cache/pip/wheels/0b/04/79/c3b0c3a0266a3cb4376da31e5bfe8bba0c489246968a68e843\n",
            "  Building wheel for subprocess32 (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for subprocess32: filename=subprocess32-3.5.4-cp37-none-any.whl size=6502 sha256=3418156727719c0c80e139c7f0e576e54a9f7906ab9934067091be976e407c91\n",
            "  Stored in directory: /root/.cache/pip/wheels/68/39/1a/5e402bdfdf004af1786c8b853fd92f8c4a04f22aad179654d1\n",
            "  Building wheel for gpustat (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for gpustat: filename=gpustat-0.6.0-cp37-none-any.whl size=12621 sha256=9d324abef86dcbdaeb320311155517b55c187136373e5262d1a13fc0a33539b1\n",
            "  Stored in directory: /root/.cache/pip/wheels/48/b4/d5/fb5b7f1d040f2ff20687e3bad6867d63155dbde5a7c10f4293\n",
            "Successfully built func-timeout ftfy dm-haiku lm-eval-harness sqlitedict pycountry openai pathtools subprocess32 gpustat\n",
            "\u001b[31mERROR: tensorflow 2.5.0 has requirement gast==0.4.0, but you'll have gast 0.3.3 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: tensorflow 2.5.0 has requirement h5py~=3.1.0, but you'll have h5py 2.10.0 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: tensorflow 2.5.0 has requirement tensorflow-estimator<2.6.0,>=2.5.0rc0, but you'll have tensorflow-estimator 2.4.0 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: google-colab 1.0.0 has requirement requests~=2.23.0, but you'll have requests 2.25.1 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: google-cloud-bigquery 1.21.0 has requirement google-resumable-media!=0.4.0,<0.5.0dev,>=0.3.1, but you'll have google-resumable-media 1.3.1 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: datascience 0.10.6 has requirement folium==0.2.1, but you'll have folium 0.8.3 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: tensorflow-cpu 2.4.2 has requirement grpcio~=1.32.0, but you'll have grpcio 1.34.1 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: black 20.8b1 has requirement regex>=2020.1.8, but you'll have regex 2019.12.20 which is incompatible.\u001b[0m\n",
            "Installing collected packages: tqdm, tokenizers, sacremoses, requests, transformers, setuptools, pathtools, shortuuid, smmap, gitdb, GitPython, sentry-sdk, docker-pycreds, configparser, subprocess32, wandb, einops, pathlib2, pynacl, bcrypt, cryptography, paramiko, invoke, fabric, chex, optax, multidict, yarl, async-timeout, aiohttp, aiohttp-cors, opencensus-context, opencensus, redis, py-spy, blessings, gpustat, colorama, hiredis, aioredis, colorful, ray, gast, tensorflow-estimator, h5py, tensorflow-cpu, google-cloud-core, google-crc32c, google-resumable-media, google-cloud-storage, func-timeout, ftfy, jmp, dm-haiku, mypy-extensions, typed-ast, pathspec, black, rehash, best-download, fsspec, huggingface-hub, xxhash, datasets, threadpoolctl, scikit-learn, sqlitedict, mbstrdecoder, pathvalidate, typepy, DataProperty, tabledata, msgfy, tcolorpy, pytablewriter, portalocker, sacrebleu, pycountry, numexpr, jsonlines, ujson, zstandard, lm-dataformat, pluggy, pytest, pybind11, tqdm-multiprocess, mock, openai, lm-eval-harness\n",
            "  Found existing installation: tqdm 4.41.1\n",
            "    Uninstalling tqdm-4.41.1:\n",
            "      Successfully uninstalled tqdm-4.41.1\n",
            "  Found existing installation: requests 2.23.0\n",
            "    Uninstalling requests-2.23.0:\n",
            "      Successfully uninstalled requests-2.23.0\n",
            "  Found existing installation: setuptools 57.0.0\n",
            "    Uninstalling setuptools-57.0.0:\n",
            "      Successfully uninstalled setuptools-57.0.0\n",
            "  Found existing installation: gast 0.4.0\n",
            "    Uninstalling gast-0.4.0:\n",
            "      Successfully uninstalled gast-0.4.0\n",
            "  Found existing installation: tensorflow-estimator 2.5.0\n",
            "    Uninstalling tensorflow-estimator-2.5.0:\n",
            "      Successfully uninstalled tensorflow-estimator-2.5.0\n",
            "  Found existing installation: h5py 3.1.0\n",
            "    Uninstalling h5py-3.1.0:\n",
            "      Successfully uninstalled h5py-3.1.0\n",
            "  Found existing installation: google-cloud-core 1.0.3\n",
            "    Uninstalling google-cloud-core-1.0.3:\n",
            "      Successfully uninstalled google-cloud-core-1.0.3\n",
            "  Found existing installation: google-resumable-media 0.4.1\n",
            "    Uninstalling google-resumable-media-0.4.1:\n",
            "      Successfully uninstalled google-resumable-media-0.4.1\n",
            "  Found existing installation: google-cloud-storage 1.18.1\n",
            "    Uninstalling google-cloud-storage-1.18.1:\n",
            "      Successfully uninstalled google-cloud-storage-1.18.1\n",
            "  Found existing installation: scikit-learn 0.22.2.post1\n",
            "    Uninstalling scikit-learn-0.22.2.post1:\n",
            "      Successfully uninstalled scikit-learn-0.22.2.post1\n",
            "  Found existing installation: numexpr 2.7.3\n",
            "    Uninstalling numexpr-2.7.3:\n",
            "      Successfully uninstalled numexpr-2.7.3\n",
            "  Found existing installation: pluggy 0.7.1\n",
            "    Uninstalling pluggy-0.7.1:\n",
            "      Successfully uninstalled pluggy-0.7.1\n",
            "  Found existing installation: pytest 3.6.4\n",
            "    Uninstalling pytest-3.6.4:\n",
            "      Successfully uninstalled pytest-3.6.4\n",
            "Successfully installed DataProperty-0.50.1 GitPython-3.1.18 aiohttp-3.7.4.post0 aiohttp-cors-0.7.0 aioredis-1.3.1 async-timeout-3.0.1 bcrypt-3.2.0 best-download-0.0.7 black-20.8b1 blessings-1.7 chex-0.0.7 colorama-0.4.4 colorful-0.5.4 configparser-5.0.2 cryptography-3.4.7 datasets-1.8.0 dm-haiku-0.0.5.dev0 docker-pycreds-0.4.0 einops-0.3.0 fabric-2.6.0 fsspec-2021.6.1 ftfy-6.0.3 func-timeout-4.3.5 gast-0.3.3 gitdb-4.0.7 google-cloud-core-1.7.1 google-cloud-storage-1.36.2 google-crc32c-1.1.2 google-resumable-media-1.3.1 gpustat-0.6.0 h5py-2.10.0 hiredis-2.0.0 huggingface-hub-0.0.13 invoke-1.5.0 jmp-0.0.2 jsonlines-2.0.0 lm-dataformat-0.0.19 lm-eval-harness-0.0.1 mbstrdecoder-1.0.1 mock-4.0.3 msgfy-0.1.0 multidict-5.1.0 mypy-extensions-0.4.3 numexpr-2.7.2 openai-0.6.4 opencensus-0.7.13 opencensus-context-0.1.2 optax-0.0.6 paramiko-2.7.2 pathlib2-2.3.5 pathspec-0.8.1 pathtools-0.1.2 pathvalidate-2.4.1 pluggy-0.13.1 portalocker-2.3.0 py-spy-0.3.7 pybind11-2.6.2 pycountry-20.7.3 pynacl-1.4.0 pytablewriter-0.58.0 pytest-6.2.3 ray-1.2.0 redis-3.5.3 rehash-1.0.0 requests-2.25.1 sacrebleu-1.5.0 sacremoses-0.0.45 scikit-learn-0.24.2 sentry-sdk-1.1.0 setuptools-51.3.3 shortuuid-1.0.1 smmap-4.0.0 sqlitedict-1.6.0 subprocess32-3.5.4 tabledata-1.1.4 tcolorpy-0.1.1 tensorflow-cpu-2.4.2 tensorflow-estimator-2.4.0 threadpoolctl-2.1.0 tokenizers-0.10.3 tqdm-4.45.0 tqdm-multiprocess-0.0.11 transformers-4.4.2 typed-ast-1.4.3 typepy-1.1.5 ujson-4.0.2 wandb-0.10.32 xxhash-2.0.2 yarl-1.6.3 zstandard-0.15.2\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "google",
                  "pkg_resources"
                ]
              }
            }
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "Processing ./mesh-transformer-jax\n",
            "Collecting jax==0.2.12\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/9a/67/d1a9c94104c559b49bbcb72e9efc33859e982d741ea4902d2a00e66e09d9/jax-0.2.12.tar.gz (590kB)\n",
            "\u001b[K     |████████████████████████████████| 593kB 5.1MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy>=1.12 in /usr/local/lib/python3.7/dist-packages (from jax==0.2.12) (1.19.5)\n",
            "Requirement already satisfied: absl-py in /usr/local/lib/python3.7/dist-packages (from jax==0.2.12) (0.12.0)\n",
            "Requirement already satisfied: opt_einsum in /usr/local/lib/python3.7/dist-packages (from jax==0.2.12) (3.3.0)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from absl-py->jax==0.2.12) (1.15.0)\n",
            "Building wheels for collected packages: jax, mesh-transformer\n",
            "  Building wheel for jax (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for jax: filename=jax-0.2.12-cp37-none-any.whl size=682484 sha256=20190cac202464f0e703eb4c161974f3568753399aa008e8e51cbc63144d8db3\n",
            "  Stored in directory: /root/.cache/pip/wheels/cf/00/88/75c2043dff473f58e892c7e6adfd2c44ccefb6111fcc021e5b\n",
            "  Building wheel for mesh-transformer (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for mesh-transformer: filename=mesh_transformer-0.0.0-cp37-none-any.whl size=21447 sha256=33e1f678630e467a1429c9dd0f7863778b6e912ebe804a80ae48bfa34e79f2eb\n",
            "  Stored in directory: /root/.cache/pip/wheels/de/a9/d2/2be3e25299342b60fca7965d4e416264ff8b6d8a7e8def76da\n",
            "Successfully built jax mesh-transformer\n",
            "Installing collected packages: jax, mesh-transformer\n",
            "  Found existing installation: jax 0.2.13\n",
            "    Uninstalling jax-0.2.13:\n",
            "      Successfully uninstalled jax-0.2.13\n",
            "Successfully installed jax-0.2.12 mesh-transformer-0.0.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aO1UXepF-0Uq"
      },
      "source": [
        "## Setup Model\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ex0qJgaueZtJ"
      },
      "source": [
        "import os\n",
        "import requests \n",
        "from jax.config import config\n",
        "\n",
        "colab_tpu_addr = os.environ['COLAB_TPU_ADDR'].split(':')[0]\n",
        "url = f'http://{colab_tpu_addr}:8475/requestversion/tpu_driver0.1_dev20210607'\n",
        "requests.post(url)\n",
        "\n",
        "# The following is required to use TPU Driver as JAX's backend.\n",
        "config.FLAGS.jax_xla_backend = \"tpu_driver\"\n",
        "config.FLAGS.jax_backend_target = \"grpc://\" + os.environ['COLAB_TPU_ADDR']"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NIgUVdFLe4A8"
      },
      "source": [
        "Sometimes the next step errors for some reason, just run it again ¯\\\\\\_(ツ)\\_/¯"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-A5IGYSaeze3"
      },
      "source": [
        "import time\n",
        "\n",
        "import jax\n",
        "from jax.experimental import maps\n",
        "import numpy as np\n",
        "import optax\n",
        "import transformers\n",
        "\n",
        "from mesh_transformer.checkpoint import read_ckpt\n",
        "from mesh_transformer.sampling import nucleaus_sample\n",
        "from mesh_transformer.transformer_shard import CausalTransformer"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QAgKq-X2kmba",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 163,
          "referenced_widgets": [
            "a0b6bd601ad84c18af2fa1b8554abf94",
            "a4f34a2cfe4a412bbdc65a70ac0b89a8",
            "916b27dbeaa3461cb695f51daf190402",
            "5e420d160a484a658b03a6a8fffc20cc",
            "951c8411294642a09d6648823effe71f",
            "97b662ebe7a748c8af6479c1f0c4bc9f",
            "a7cfdce4a1a64cc6bcd2842dce96303f",
            "cdee93968c5f42cba4b7a866f8900c97",
            "b4939af44ff14f7185c4ee219da80efa",
            "a48492a16fd1452aaa519301e0a230d0",
            "613e272e33914065bb89d38ddc6e6678",
            "aa655ed683694c158af8f9fbffea478f",
            "fda271a910e84563bc7542ae7a1ea966",
            "32b6ed93fac94d6bbf0556d952b303c3",
            "266629ad54db4b959adf0ea15f81e0d7",
            "0656f2f4abc94a17b0c0ca532312ae3e",
            "736045507c53436390c8d1b88472f770",
            "6f6a6457a8604be8859d9de0f082def4",
            "805c07778d3542bb8dfbe3daf3423d50",
            "5c8e194b8fb44395b4b03a560e7f4c0e",
            "2c6e73aff92649f181faa96e746779a4",
            "c447b1b8e99b4e2b97a69e035ccf7f49",
            "8b7b7e1da4c74bbcbb09f949d8bf25e5",
            "068220f70575498d85a03bde26fb4651"
          ]
        },
        "outputId": "70ee756a-d406-4920-cfca-c9b715db66f1"
      },
      "source": [
        "params = {\n",
        "  \"layers\": 28,\n",
        "  \"d_model\": 4096,\n",
        "  \"n_heads\": 16,\n",
        "  \"n_vocab\": 50400,\n",
        "  \"norm\": \"layernorm\",\n",
        "  \"pe\": \"rotary\",\n",
        "  \"pe_rotary_dims\": 64,\n",
        "\n",
        "  \"seq\": 2048,\n",
        "  \"cores_per_replica\": 8,\n",
        "  \"per_replica_batch\": 1,\n",
        "}\n",
        "\n",
        "per_replica_batch = params[\"per_replica_batch\"]\n",
        "cores_per_replica = params[\"cores_per_replica\"]\n",
        "seq = params[\"seq\"]\n",
        "\n",
        "\n",
        "params[\"sampler\"] = nucleaus_sample\n",
        "\n",
        "# here we \"remove\" the optimizer parameters from the model (as we don't need them for inference)\n",
        "params[\"optimizer\"] = optax.scale(0)\n",
        "\n",
        "mesh_shape = (jax.device_count() // cores_per_replica, cores_per_replica)\n",
        "devices = np.array(jax.devices()).reshape(mesh_shape)\n",
        "\n",
        "maps.thread_resources.env = maps.ResourceEnv(maps.Mesh(devices, ('dp', 'mp')))\n",
        "\n",
        "tokenizer = transformers.GPT2TokenizerFast.from_pretrained('gpt2')"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "a0b6bd601ad84c18af2fa1b8554abf94",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=1042301.0, style=ProgressStyle(descript…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "b4939af44ff14f7185c4ee219da80efa",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=456318.0, style=ProgressStyle(descripti…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "736045507c53436390c8d1b88472f770",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=1355256.0, style=ProgressStyle(descript…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yFgRkUgfiNdA"
      },
      "source": [
        "Here we create the network and load the parameters from the downloaded files. Expect this to take around 5 minutes."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lwNETD2Uk8nu",
        "outputId": "f22b3366-bea7-4030-d1a9-7127d8506e8b"
      },
      "source": [
        "total_batch = per_replica_batch * jax.device_count() // cores_per_replica\n",
        "\n",
        "network = CausalTransformer(params)\n",
        "\n",
        "network.state = read_ckpt(network.state, \"step_383500/\", devices.shape[1])\n",
        "\n",
        "network.state = network.move_xmap(network.state, np.zeros(cores_per_replica))"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/jax/experimental/maps.py:412: UserWarning: xmap is an experimental feature and probably has bugs!\n",
            "  warn(\"xmap is an experimental feature and probably has bugs!\")\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "key shape (8, 2)\n",
            "in shape (1, 2048)\n",
            "dp 1\n",
            "mp 8\n",
            "read from disk/gcs in 20.3309s\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A-eT7Sw6if4J"
      },
      "source": [
        "## Run Model\n",
        "\n",
        "Finally, we are ready to infer with the model! The first sample takes around a minute due to compilation, but after that it should only take about 10 seconds per sample.\n",
        "\n",
        "Feel free to mess with the different sampling parameters (top_p and temp), as well as the length of the generations (gen_len, causes a recompile when changed).\n",
        "\n",
        "You can also change other things like per_replica_batch in the previous cells to change how many generations are done in parallel. A larger batch has higher latency but higher throughput when measured in tokens generated/s. This is useful for doing things like best-of-n cherry picking.\n",
        "\n",
        "*Tip for best results: Make sure your prompt does not have any trailing spaces, which tend to confuse the model due to the BPE tokenization used during training.*"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8TiclqEvwoCL"
      },
      "source": [
        "# allow text wrapping in generated output: https://stackoverflow.com/a/61401455\n",
        "from IPython.display import HTML, display\n",
        "\n",
        "def set_css():\n",
        "  display(HTML('''\n",
        "  <style>\n",
        "    pre {\n",
        "        white-space: pre-wrap;\n",
        "    }\n",
        "  </style>\n",
        "  '''))\n",
        "get_ipython().events.register('pre_run_cell', set_css)"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 468
        },
        "id": "ZVzs2TYlvYeX",
        "outputId": "3161289f-fc5d-43f3-cb94-d18f8483a8b5"
      },
      "source": [
        "def infer(context, top_p=0.9, temp=1.0, gen_len=512):\n",
        "    tokens = tokenizer.encode(context)\n",
        "\n",
        "    provided_ctx = len(tokens)\n",
        "    pad_amount = seq - provided_ctx\n",
        "\n",
        "    padded_tokens = np.pad(tokens, ((pad_amount, 0),)).astype(np.uint32)\n",
        "    batched_tokens = np.array([padded_tokens] * total_batch)\n",
        "    length = np.ones(total_batch, dtype=np.uint32) * len(tokens)\n",
        "\n",
        "    start = time.time()\n",
        "    output = network.generate(batched_tokens, length, gen_len, {\"top_p\": np.ones(total_batch) * top_p, \"temp\": np.ones(total_batch) * temp})\n",
        "\n",
        "    samples = []\n",
        "    decoded_tokens = output[1][0]\n",
        "\n",
        "    for o in decoded_tokens[:, :, 0]:\n",
        "      samples.append(f\"\\033[1m{context}\\033[0m{tokenizer.decode(o)}\")\n",
        "\n",
        "    print(f\"completion done in {time.time() - start:06}s\")\n",
        "    return samples\n",
        "\n",
        "print(infer(\"EleutherAI is\")[0])"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "completion done in 60.493677616119385s\n",
            "\u001b[1mEleutherAI is\u001b[0m a special AI for all the mythical 'elves', and is divided in many dungeons. Each dungeon has a different type of monsters. They can only be defeated when you have a specified level. To progress the dungeon, you must find the labyrinth, activate the secret door, open the door of the labyrinth and defeat the monsters on the way.\n",
            "\n",
            "I have no idea why the story is about elves, but it has a variety of levels with varying gameplay mechanics. There are a lot of random items, but there is a great amount of randomness. There are also many references, such as name references to good mithril items. In the mean time I tried out the mechanic of random items that you need to have.\n",
            "\n",
            "The gameplay is a mix of twitch reflexes, luck, and determination. Each stage has different types of abilities and each of them affects the entire game. You do have a chosen class, but it isn't easy to tell which one is better as the dungeon world is big and there are different types of monsters. You have four items with a passive ability and two active abilities. The passive one is generally weak, so you need to have a lot of items.\n",
            "\n",
            "There is a lot of randomness, and you can die if you encounter many monsters, but it's not impossible to come back to life with a time vortex. There are no hard time restrictions or limited time as such. I felt like I was playing an RPG, but it wasn't too long before I felt like I was playing an adventure game, because you need to find and activate the secret door. I also liked the monster variety. Although each dungeon can be harder than the last, it will always be interesting to see which type of monsters you'll encounter.\n",
            "\n",
            "Since this game is about the experience of an AI player, I need to give some tips. You need to beat all the dungeons to start, since you need a power level, or it won't let you progress. You also need to spend the money you earned to buy equipment to raise the power level. There are plenty of items, and you can earn more by defeating monsters, winning treasure chests, and finding various items. Each item you find will have a chance to affect your experience level, so it's not guaranteed that you'll be able to buy a new item when you meet a certain power level.\n",
            "\n",
            "This game is better for newcomers to the Minecraft world. You can only fight monsters, and not humans. You do need to have a\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "nvlAK6RbCJYg",
        "outputId": "a6492702-28df-4c5d-c2a0-46d6ab89a262"
      },
      "source": [
        "#@title  { form-width: \"300px\" }\n",
        "top_p = 0.9 #@param {type:\"slider\", min:0, max:1, step:0.1}\n",
        "temp = 1 #@param {type:\"slider\", min:0, max:1, step:0.1}\n",
        "\n",
        "context = \"\"\"In a shocking finding, scientist discovered a herd of unicorns living in a remote, previously unexplored valley, in the Andes Mountains. Even more surprising to the researchers was the fact that the unicorns spoke perfect English.\"\"\"\n",
        "\n",
        "print(infer(top_p=top_p, temp=temp, gen_len=1024, context=context)[0])"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "completion done in 67.64067959785461s\n",
            "\u001b[1mIn a shocking finding, scientist discovered a herd of unicorns living in a remote, previously unexplored valley, in the Andes Mountains. Even more surprising to the researchers was the fact that the unicorns spoke perfect English.\u001b[0m\n",
            "\n",
            "\"I'm now seriously into ghosts. I'm also into monsters, black magic, and the reality that we're all getting fatter every day. It's a fact of life.\"\n",
            "\n",
            "— Author Hunter S. Thompson\n",
            "\n",
            "There's just one problem: his ideas aren't the same as the world's.\n",
            "\n",
            "What the hell am I talking about? To prove to the world that the two opinions are equally absurd, each time Stephen Hawking or Albert Einstein said something outrageous, the people who most agreed with the statements were scientists. They even voted that the two of them are perfect examples of Nobel prize-winning individuals.\n",
            "\n",
            "In the table below you'll see a breakdown of the opinions held by various groups (a relative majority indicates that most people agreed with the statement). Each row represents one group and each column represents a statement by Hawking or Einstein.\n",
            "\n",
            "The number of opinions a statement has indicates the percentage of the people who believed that it was true.\n",
            "\n",
            "Summary of the findings\n",
            "\n",
            "As you can see, groups that agreed with Hawking or Einstein the most were overwhelmingly scientists. While the results are shocking, they are not entirely surprising.\n",
            "\n",
            "Scientific findings are largely derived from the ideas of scientists, which can be compared to an iceberg. While the vast majority of the scientific understanding is derived from things we call ice, icebergs are predominantly water. This means that almost nothing of real value comes from the part of the iceberg that we don't see. Instead of the entire theory of relativity being dependent on Isaac Newton's equations for the motion of bodies, it's mostly dependent on Albert Einstein's theory of relativity.\n",
            "\n",
            "Similarly, a lot of the people who agree with Hawking or Einstein probably have their opinions based on a lot of learning, but a lot of that learning comes from the theories that they themselves have put forward. For this reason, it's not surprising that most people agree with these two people. After all, most people do agree that they are two of the greatest scientists to have ever lived.\n",
            "\n",
            "It's also interesting to note that scientists hold the most conservative opinions of any group. This is no doubt due to their knowledge of the scientific facts. If the findings of scientists are shown to be absurd, it will probably be deemed as an error rather than a lack of fact.\n",
            "\n",
            "The table above also shows that we should be careful about using Hawking's and Einstein's ideas to predict the future. While they have had very large and profound influences on the world, the majority of people in the world still follow Newton's ideas when it comes to most of the things that they learn about.\n",
            "\n",
            "One can also see that the \"hawkish\" and \"elegant\" ideas were pretty unpopular. While those were widely considered truths before, a large number of people in the world now seem to agree with the theories that the popular public think are foolish.\n",
            "\n",
            "This may be in part due to the fact that they weren't able to fit them into their world view, but most people agreed with the ideas when Hawking and Einstein were first coming up with them. The same people who thought them ridiculous then also thought them ridiculous years later, even though they have continued to grow into scientific truths.\n",
            "\n",
            "Meanwhile, the vast majority of those who agree with the \"elegant\" ideas now have gained greater respect for them as time has gone on.\n",
            "\n",
            "Conclusion\n",
            "\n",
            "Before the table, I said that it is shocking that the most popular individuals in the world are two Nobel prize-winning scientists. However, it's not the fact that they are scientists that are shocking, but that they seem to hold common ideas among all of us. It's also shocking that, while we accept their ideas, they are widely rejected by the public.\n",
            "\n",
            "If it seems like a lot of people want to blame Hawking and Einstein for their unpopular ideas, it is important to remember that many of their ideas were considered unlikely at the time. In the early 1900s, it was claimed that the Sun revolved around the Earth and that the Moon was a flat disc that revolves around the Earth. Both of these theories, while odd at the time, were eventually shown to be very good theories.\n",
            "\n",
            "When we think of these theories in the light of the known facts, we see that these were two of the boldest ideas ever put forward by scientists. Even more important, as time went on, they were proven to be right. As a result, even though they had a large number of people who agreed with them when they were first put forward, it doesn't mean that they are still right. While scientists are very important in a certain sense, one should realize that we have a better sense of what's true or not because of people like Hawking and Einstein.\n",
            "\n",
            "With that in mind, it's not hard to see that, while the ideas of these two individuals are ridiculous, the people who take up these ideas also do so for reasons that have more to do with their own mental make-up. Some of\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}