colab_demo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "GPT-J-6B Inference Demo.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "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": {
        "06b17c299dbc470d84ccdc428d05e8eb": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "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_20ec90427edd47c3b3e91faffe1de76c",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_a99395f09b9d4ffa9dddf219fd75f614",
              "IPY_MODEL_5673aea1d57443b88ffae41a725531fd"
            ]
          }
        },
        "20ec90427edd47c3b3e91faffe1de76c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "a99395f09b9d4ffa9dddf219fd75f614": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_8ef4566445e44b08baee45d997708837",
            "_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_06b943526e604b0baaf4169348efafda"
          }
        },
        "5673aea1d57443b88ffae41a725531fd": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_2916fbd3264d439097e3b4db151ef090",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 0.99M/0.99M [00:10&lt;00:00, 104kB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_9a7660cf79b7486ebe2832c80297f1bf"
          }
        },
        "8ef4566445e44b08baee45d997708837": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "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"
          }
        },
        "06b943526e604b0baaf4169348efafda": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "2916fbd3264d439097e3b4db151ef090": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "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"
          }
        },
        "9a7660cf79b7486ebe2832c80297f1bf": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "48fabfd1f10549699d02eb9fc721a608": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "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_f0bc3d33955841fa914dddeb2759222c",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_b9440118187748a2af79bd8c1bb86391",
              "IPY_MODEL_18f3336cacba4ad9a3863dceef684d62"
            ]
          }
        },
        "f0bc3d33955841fa914dddeb2759222c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "b9440118187748a2af79bd8c1bb86391": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_5a043e80fbb243619e1ec11b14dfcf82",
            "_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_91ed7493ccb8424bab2372dadf2470f0"
          }
        },
        "18f3336cacba4ad9a3863dceef684d62": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_4c532da328fc4aec9a70dbc2a7b749f5",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 446k/446k [00:01&lt;00:00, 452kB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_e28c6a5968e944be8a832541d05cb336"
          }
        },
        "5a043e80fbb243619e1ec11b14dfcf82": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "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"
          }
        },
        "91ed7493ccb8424bab2372dadf2470f0": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "4c532da328fc4aec9a70dbc2a7b749f5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "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"
          }
        },
        "e28c6a5968e944be8a832541d05cb336": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "0cdfd04d211f4d32a82c4ef42f7e99f9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "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_781dc409b97146e68be675a1d18328b6",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_7538b4d51acf4c369558a1471f985ffe",
              "IPY_MODEL_1e0681e3be994e96b9d5853ce47040fa"
            ]
          }
        },
        "781dc409b97146e68be675a1d18328b6": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "7538b4d51acf4c369558a1471f985ffe": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_20a12a8118f74f65aa03403500594aaa",
            "_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_72f39605f935481f8f68dbf83a672306"
          }
        },
        "1e0681e3be994e96b9d5853ce47040fa": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_ec441bd40f124042bad3242e84bf9133",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 1.29M/1.29M [00:00&lt;00:00, 2.33MB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_e76ed536e65f43d9a1ff6f9e023657d8"
          }
        },
        "20a12a8118f74f65aa03403500594aaa": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "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"
          }
        },
        "72f39605f935481f8f68dbf83a672306": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "ec441bd40f124042bad3242e84bf9133": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "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"
          }
        },
        "e76ed536e65f43d9a1ff6f9e023657d8": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "22f8f50f14ad42eea06fc43473208741": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "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_41600b0a132b4f30bf1ac331055d8f06",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_9aee12e4d953453680599582ae215ba0",
              "IPY_MODEL_e012f372170844e18928ed98656bc310"
            ]
          }
        },
        "41600b0a132b4f30bf1ac331055d8f06": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "9aee12e4d953453680599582ae215ba0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_fc5441d719d045a4a1655e838b2437a9",
            "_dom_classes": [],
            "description": "Downloading: 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 665,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 665,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_600e8b54ce4543dba851586530ce43d9"
          }
        },
        "e012f372170844e18928ed98656bc310": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_fc2294cbfec7435c922aa2a67b0efa5d",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 665/665 [00:08&lt;00:00, 81.0B/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_8a5ba8d9802044f592e435251efd3495"
          }
        },
        "fc5441d719d045a4a1655e838b2437a9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "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"
          }
        },
        "600e8b54ce4543dba851586530ce43d9": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "fc2294cbfec7435c922aa2a67b0efa5d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "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"
          }
        },
        "8a5ba8d9802044f592e435251efd3495": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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/edenweb1/retroactive-query/blob/master/colab_demo.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": "fa17fe51-d8ea-452e-fde3-f2c4ac0736e3"
      },
      "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 tensorflow==2.5.0"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "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 37 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 (334 kB/s)\n",
            "Selecting previously unselected package zstd.\n",
            "(Reading database ... 155047 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-10-19 07:49:53--  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.244\n",
            "Connecting to the-eye.eu (the-eye.eu)|162.213.130.244|: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  60%[===========>        ]   5.28G  66.6MB/s    in 78s     \n",
            "\n",
            "2021-10-19 07:51:12 (69.3 MB/s) - Connection closed at byte 5666635560. Retrying.\n",
            "\n",
            "--2021-10-19 07:51:13--  (try: 2)  https://the-eye.eu/public/AI/GPT-J-6B/step_383500_slim.tar.zstd\n",
            "Connecting to the-eye.eu (the-eye.eu)|162.213.130.244|:443... connected.\n",
            "HTTP request sent, awaiting response... 206 Partial Content\n",
            "Length: 9414712325 (8.8G), 3748076765 (3.5G) remaining [application/octet-stream]\n",
            "Saving to: ‘step_383500_slim.tar.zstd’\n",
            "\n",
            "step_383500_slim.ta 100%[++++++++++++=======>]   8.77G  74.4MB/s    in 63s     \n",
            "\n",
            "2021-10-19 07:52:16 (56.7 MB/s) - ‘step_383500_slim.tar.zstd’ saved [9414712325/9414712325]\n",
            "\n",
            "\n",
            "real\t2m22.530s\n",
            "user\t0m8.752s\n",
            "sys\t0m45.408s\n",
            "\n",
            "real\t1m28.219s\n",
            "user\t0m34.611s\n",
            "sys\t0m30.258s\n",
            "Cloning into 'mesh-transformer-jax'...\n",
            "remote: Enumerating objects: 760, done.\u001b[K\n",
            "remote: Counting objects: 100% (400/400), done.\u001b[K\n",
            "remote: Compressing objects: 100% (121/121), done.\u001b[K\n",
            "remote: Total 760 (delta 342), reused 285 (delta 279), pack-reused 360\u001b[K\n",
            "Receiving objects: 100% (760/760), 242.15 KiB | 3.61 MiB/s, done.\n",
            "Resolving deltas: 100% (502/502), done.\n",
            "Collecting git+https://github.com/deepmind/dm-haiku (from -r mesh-transformer-jax/requirements.txt (line 8))\n",
            "  Cloning https://github.com/deepmind/dm-haiku to /tmp/pip-req-build-7qqzpay1\n",
            "  Running command git clone -q https://github.com/deepmind/dm-haiku /tmp/pip-req-build-7qqzpay1\n",
            "Collecting git+https://github.com/EleutherAI/lm-evaluation-harness/ (from -r mesh-transformer-jax/requirements.txt (line 9))\n",
            "  Cloning https://github.com/EleutherAI/lm-evaluation-harness/ to /tmp/pip-req-build-54yf2chn\n",
            "  Running command git clone -q https://github.com/EleutherAI/lm-evaluation-harness/ /tmp/pip-req-build-54yf2chn\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 tqdm~=4.45.0\n",
            "  Downloading tqdm-4.45.0-py2.py3-none-any.whl (60 kB)\n",
            "\u001b[K     |████████████████████████████████| 60 kB 3.4 MB/s \n",
            "\u001b[?25hCollecting wandb>=0.11.2\n",
            "  Downloading wandb-0.12.4-py2.py3-none-any.whl (1.7 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.7 MB 11.2 MB/s \n",
            "\u001b[?25hCollecting einops~=0.3.0\n",
            "  Downloading einops-0.3.2-py3-none-any.whl (25 kB)\n",
            "Collecting requests~=2.25.1\n",
            "  Downloading requests-2.25.1-py2.py3-none-any.whl (61 kB)\n",
            "\u001b[K     |████████████████████████████████| 61 kB 7.7 MB/s \n",
            "\u001b[?25hCollecting fabric~=2.6.0\n",
            "  Downloading fabric-2.6.0-py2.py3-none-any.whl (53 kB)\n",
            "\u001b[K     |████████████████████████████████| 53 kB 1.4 MB/s \n",
            "\u001b[?25hCollecting optax==0.0.6\n",
            "  Downloading optax-0.0.6-py3-none-any.whl (96 kB)\n",
            "\u001b[K     |████████████████████████████████| 96 kB 5.1 MB/s \n",
            "\u001b[?25hCollecting ray[default]==1.4.1\n",
            "  Downloading ray-1.4.1-cp37-cp37m-manylinux2014_x86_64.whl (51.6 MB)\n",
            "\u001b[K     |████████████████████████████████| 51.6 MB 92 kB/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 11)) (0.2.21)\n",
            "Requirement already satisfied: Flask~=1.1.2 in /usr/local/lib/python3.7/dist-packages (from -r mesh-transformer-jax/requirements.txt (line 12)) (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 13)) (1.3.0)\n",
            "Collecting tensorflow-cpu~=2.5.0\n",
            "  Downloading tensorflow_cpu-2.5.1-cp37-cp37m-manylinux2010_x86_64.whl (168.3 MB)\n",
            "\u001b[K     |████████████████████████████████| 168.3 MB 3.0 kB/s \n",
            "\u001b[?25hCollecting google-cloud-storage~=1.36.2\n",
            "  Downloading google_cloud_storage-1.36.2-py2.py3-none-any.whl (97 kB)\n",
            "\u001b[K     |████████████████████████████████| 97 kB 5.9 MB/s \n",
            "\u001b[?25hCollecting transformers\n",
            "  Downloading transformers-4.11.3-py3-none-any.whl (2.9 MB)\n",
            "\u001b[K     |████████████████████████████████| 2.9 MB 62.1 MB/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 17)) (5.2.1)\n",
            "Collecting func_timeout\n",
            "  Downloading func_timeout-4.3.5.tar.gz (44 kB)\n",
            "\u001b[K     |████████████████████████████████| 44 kB 2.2 MB/s \n",
            "\u001b[?25hCollecting ftfy\n",
            "  Downloading ftfy-6.0.3.tar.gz (64 kB)\n",
            "\u001b[K     |████████████████████████████████| 64 kB 2.5 MB/s \n",
            "\u001b[?25hCollecting fastapi\n",
            "  Downloading fastapi-0.70.0-py3-none-any.whl (51 kB)\n",
            "\u001b[K     |████████████████████████████████| 51 kB 478 kB/s \n",
            "\u001b[?25hCollecting uvicorn\n",
            "  Downloading uvicorn-0.15.0-py3-none-any.whl (54 kB)\n",
            "\u001b[K     |████████████████████████████████| 54 kB 2.6 MB/s \n",
            "\u001b[?25hCollecting lm_dataformat\n",
            "  Downloading lm_dataformat-0.0.20-py3-none-any.whl (5.8 kB)\n",
            "Collecting pathy\n",
            "  Downloading pathy-0.6.0-py3-none-any.whl (42 kB)\n",
            "\u001b[K     |████████████████████████████████| 42 kB 1.3 MB/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 8)) (0.12.0)\n",
            "Collecting jmp>=0.0.2\n",
            "  Downloading jmp-0.0.2-py3-none-any.whl (16 kB)\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 8)) (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 8)) (3.7.4.3)\n",
            "Collecting bleurt@ https://github.com/google-research/bleurt/archive/b610120347ef22b494b6d69b4316e303f5932516.zip#egg=bleurt\n",
            "  Downloading https://github.com/google-research/bleurt/archive/b610120347ef22b494b6d69b4316e303f5932516.zip (16.5 MB)\n",
            "\u001b[K     |████████████████████████████████| 16.5 MB 396 kB/s \n",
            "\u001b[?25hCollecting black==20.8b1\n",
            "  Downloading black-20.8b1.tar.gz (1.1 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.1 MB 59.5 MB/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 best_download-0.0.7-py3-none-any.whl (4.5 kB)\n",
            "Collecting datasets>=1.2.1\n",
            "  Downloading datasets-1.13.3-py3-none-any.whl (287 kB)\n",
            "\u001b[K     |████████████████████████████████| 287 kB 60.6 MB/s \n",
            "\u001b[?25hRequirement already satisfied: click>=7.1 in /usr/local/lib/python3.7/dist-packages (from lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (7.1.2)\n",
            "Collecting scikit-learn>=0.24.1\n",
            "  Downloading scikit_learn-1.0-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (23.1 MB)\n",
            "\u001b[K     |████████████████████████████████| 23.1 MB 1.1 MB/s \n",
            "\u001b[?25hRequirement already satisfied: torch>=1.7 in /usr/local/lib/python3.7/dist-packages (from lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (1.9.0+cu111)\n",
            "Collecting sqlitedict==1.6.0\n",
            "  Downloading sqlitedict-1.6.0.tar.gz (29 kB)\n",
            "Collecting pytablewriter==0.58.0\n",
            "  Downloading pytablewriter-0.58.0-py3-none-any.whl (96 kB)\n",
            "\u001b[K     |████████████████████████████████| 96 kB 5.6 MB/s \n",
            "\u001b[?25hCollecting sacrebleu==1.5.0\n",
            "  Downloading sacrebleu-1.5.0-py3-none-any.whl (65 kB)\n",
            "\u001b[K     |████████████████████████████████| 65 kB 3.8 MB/s \n",
            "\u001b[?25hCollecting pycountry==20.7.3\n",
            "  Downloading pycountry-20.7.3.tar.gz (10.1 MB)\n",
            "\u001b[K     |████████████████████████████████| 10.1 MB 55.1 MB/s \n",
            "\u001b[?25hCollecting numexpr==2.7.2\n",
            "  Downloading numexpr-2.7.2-cp37-cp37m-manylinux2010_x86_64.whl (471 kB)\n",
            "\u001b[K     |████████████████████████████████| 471 kB 62.0 MB/s \n",
            "\u001b[?25hCollecting lm_dataformat\n",
            "  Downloading lm_dataformat-0.0.19-py3-none-any.whl (5.4 kB)\n",
            "Collecting pytest==6.2.3\n",
            "  Downloading pytest-6.2.3-py3-none-any.whl (280 kB)\n",
            "\u001b[K     |████████████████████████████████| 280 kB 64.7 MB/s \n",
            "\u001b[?25hCollecting pybind11==2.6.2\n",
            "  Downloading pybind11-2.6.2-py2.py3-none-any.whl (191 kB)\n",
            "\u001b[K     |████████████████████████████████| 191 kB 68.0 MB/s \n",
            "\u001b[?25hCollecting tqdm-multiprocess==0.0.11\n",
            "  Downloading tqdm_multiprocess-0.0.11-py3-none-any.whl (9.8 kB)\n",
            "Collecting zstandard==0.15.2\n",
            "  Downloading zstandard-0.15.2-cp37-cp37m-manylinux2014_x86_64.whl (2.2 MB)\n",
            "\u001b[K     |████████████████████████████████| 2.2 MB 55.2 MB/s \n",
            "\u001b[?25hCollecting jsonlines==2.0.0\n",
            "  Downloading jsonlines-2.0.0-py3-none-any.whl (6.3 kB)\n",
            "Collecting mock==4.0.3\n",
            "  Downloading mock-4.0.3-py3-none-any.whl (28 kB)\n",
            "Collecting openai==0.6.4\n",
            "  Downloading openai-0.6.4.tar.gz (159 kB)\n",
            "\u001b[K     |████████████████████████████████| 159 kB 70.2 MB/s \n",
            "\u001b[?25hRequirement already satisfied: jieba==0.42.1 in /usr/local/lib/python3.7/dist-packages (from lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (0.42.1)\n",
            "Collecting nagisa==0.2.7\n",
            "  Downloading nagisa-0.2.7-cp37-cp37m-manylinux1_x86_64.whl (21.5 MB)\n",
            "\u001b[K     |████████████████████████████████| 21.5 MB 230 kB/s \n",
            "\u001b[?25hCollecting t5==0.7.1\n",
            "  Downloading t5-0.7.1-py3-none-any.whl (172 kB)\n",
            "\u001b[K     |████████████████████████████████| 172 kB 66.4 MB/s \n",
            "\u001b[?25hRequirement already satisfied: pandas in /usr/local/lib/python3.7/dist-packages (from bleurt@ https://github.com/google-research/bleurt/archive/b610120347ef22b494b6d69b4316e303f5932516.zip#egg=bleurt->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (1.1.5)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from bleurt@ https://github.com/google-research/bleurt/archive/b610120347ef22b494b6d69b4316e303f5932516.zip#egg=bleurt->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (1.4.1)\n",
            "Requirement already satisfied: tensorflow in /usr/local/lib/python3.7/dist-packages (from bleurt@ https://github.com/google-research/bleurt/archive/b610120347ef22b494b6d69b4316e303f5932516.zip#egg=bleurt->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (2.6.0)\n",
            "Collecting tf-slim>=1.1\n",
            "  Downloading tf_slim-1.1.0-py2.py3-none-any.whl (352 kB)\n",
            "\u001b[K     |████████████████████████████████| 352 kB 67.3 MB/s \n",
            "\u001b[?25hCollecting sentencepiece\n",
            "  Downloading sentencepiece-0.1.96-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.2 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.2 MB 61.2 MB/s \n",
            "\u001b[?25hCollecting chex>=0.0.4\n",
            "  Downloading chex-0.0.8-py3-none-any.whl (57 kB)\n",
            "\u001b[K     |████████████████████████████████| 57 kB 5.2 MB/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 7)) (0.1.71+cuda111)\n",
            "Collecting ujson\n",
            "  Downloading ujson-4.2.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (214 kB)\n",
            "\u001b[K     |████████████████████████████████| 214 kB 63.8 MB/s \n",
            "\u001b[?25hCollecting typed-ast>=1.4.0\n",
            "  Downloading typed_ast-1.4.3-cp37-cp37m-manylinux1_x86_64.whl (743 kB)\n",
            "\u001b[K     |████████████████████████████████| 743 kB 66.8 MB/s \n",
            "\u001b[?25hCollecting pathspec<1,>=0.6\n",
            "  Downloading pathspec-0.9.0-py2.py3-none-any.whl (31 kB)\n",
            "Requirement already satisfied: appdirs in /usr/local/lib/python3.7/dist-packages (from black==20.8b1->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (1.4.4)\n",
            "Requirement already satisfied: toml>=0.10.1 in /usr/local/lib/python3.7/dist-packages (from black==20.8b1->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (0.10.2)\n",
            "Collecting regex>=2020.1.8\n",
            "  Downloading regex-2021.10.8-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (748 kB)\n",
            "\u001b[K     |████████████████████████████████| 748 kB 65.7 MB/s \n",
            "\u001b[?25hCollecting mypy-extensions>=0.4.3\n",
            "  Downloading mypy_extensions-0.4.3-py2.py3-none-any.whl (4.5 kB)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from nagisa==0.2.7->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (1.15.0)\n",
            "Collecting DyNet\n",
            "  Downloading dyNET-2.1.2-cp37-cp37m-manylinux1_x86_64.whl (4.4 MB)\n",
            "\u001b[K     |████████████████████████████████| 4.4 MB 58.1 MB/s \n",
            "\u001b[?25hCollecting pathvalidate<3,>=2.3.0\n",
            "  Downloading pathvalidate-2.5.0-py3-none-any.whl (19 kB)\n",
            "Collecting msgfy<1,>=0.1.0\n",
            "  Downloading msgfy-0.2.0-py3-none-any.whl (4.3 kB)\n",
            "Requirement already satisfied: setuptools>=38.3.0 in /usr/local/lib/python3.7/dist-packages (from pytablewriter==0.58.0->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (57.4.0)\n",
            "Collecting typepy[datetime]<2,>=1.1.1\n",
            "  Downloading typepy-1.3.0-py3-none-any.whl (31 kB)\n",
            "Collecting tcolorpy<1,>=0.0.5\n",
            "  Downloading tcolorpy-0.1.1-py3-none-any.whl (8.0 kB)\n",
            "Collecting mbstrdecoder<2,>=1.0.0\n",
            "  Downloading mbstrdecoder-1.1.0-py3-none-any.whl (7.8 kB)\n",
            "Collecting tabledata<2,>=1.1.3\n",
            "  Downloading tabledata-1.3.0-py3-none-any.whl (11 kB)\n",
            "Collecting DataProperty<2,>=0.50.0\n",
            "  Downloading DataProperty-0.54.2-py3-none-any.whl (26 kB)\n",
            "Requirement already satisfied: attrs>=19.2.0 in /usr/local/lib/python3.7/dist-packages (from pytest==6.2.3->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (21.2.0)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from pytest==6.2.3->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (21.0)\n",
            "Requirement already satisfied: importlib-metadata>=0.12 in /usr/local/lib/python3.7/dist-packages (from pytest==6.2.3->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (4.8.1)\n",
            "Collecting pluggy<1.0.0a1,>=0.12\n",
            "  Downloading pluggy-0.13.1-py2.py3-none-any.whl (18 kB)\n",
            "Requirement already satisfied: py>=1.8.2 in /usr/local/lib/python3.7/dist-packages (from pytest==6.2.3->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (1.10.0)\n",
            "Requirement already satisfied: iniconfig in /usr/local/lib/python3.7/dist-packages (from pytest==6.2.3->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (1.1.1)\n",
            "Requirement already satisfied: protobuf>=3.15.3 in /usr/local/lib/python3.7/dist-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (3.17.3)\n",
            "Collecting aioredis\n",
            "  Downloading aioredis-2.0.0-py3-none-any.whl (69 kB)\n",
            "\u001b[K     |████████████████████████████████| 69 kB 7.2 MB/s \n",
            "\u001b[?25hCollecting gpustat\n",
            "  Downloading gpustat-0.6.0.tar.gz (78 kB)\n",
            "\u001b[K     |████████████████████████████████| 78 kB 6.5 MB/s \n",
            "\u001b[?25hRequirement already satisfied: filelock in /usr/local/lib/python3.7/dist-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (3.3.0)\n",
            "Collecting colorama\n",
            "  Downloading colorama-0.4.4-py2.py3-none-any.whl (16 kB)\n",
            "Collecting opencensus\n",
            "  Downloading opencensus-0.8.0-py2.py3-none-any.whl (128 kB)\n",
            "\u001b[K     |████████████████████████████████| 128 kB 57.2 MB/s \n",
            "\u001b[?25hRequirement already satisfied: prometheus-client>=0.7.1 in /usr/local/lib/python3.7/dist-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (0.11.0)\n",
            "Collecting aiohttp-cors\n",
            "  Downloading aiohttp_cors-0.7.0-py3-none-any.whl (27 kB)\n",
            "Requirement already satisfied: jsonschema in /usr/local/lib/python3.7/dist-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (2.6.0)\n",
            "Requirement already satisfied: msgpack<2.0.0,>=1.0.0 in /usr/local/lib/python3.7/dist-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (1.0.2)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.7/dist-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (3.13)\n",
            "Collecting redis>=3.5.0\n",
            "  Downloading redis-3.5.3-py2.py3-none-any.whl (72 kB)\n",
            "\u001b[K     |████████████████████████████████| 72 kB 499 kB/s \n",
            "\u001b[?25hCollecting aiohttp\n",
            "  Downloading aiohttp-3.7.4.post0-cp37-cp37m-manylinux2014_x86_64.whl (1.3 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.3 MB 59.4 MB/s \n",
            "\u001b[?25hRequirement already satisfied: grpcio>=1.28.1 in /usr/local/lib/python3.7/dist-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (1.41.0)\n",
            "Collecting py-spy>=0.2.0\n",
            "  Downloading py_spy-0.3.10-py2.py3-none-manylinux_2_5_x86_64.manylinux1_x86_64.whl (3.2 MB)\n",
            "\u001b[K     |████████████████████████████████| 3.2 MB 55.8 MB/s \n",
            "\u001b[?25hCollecting pydantic>=1.8\n",
            "  Downloading pydantic-1.8.2-cp37-cp37m-manylinux2014_x86_64.whl (10.1 MB)\n",
            "\u001b[K     |████████████████████████████████| 10.1 MB 61.4 MB/s \n",
            "\u001b[?25hCollecting colorful\n",
            "  Downloading colorful-0.5.4-py2.py3-none-any.whl (201 kB)\n",
            "\u001b[K     |████████████████████████████████| 201 kB 68.1 MB/s \n",
            "\u001b[?25hCollecting portalocker\n",
            "  Downloading portalocker-2.3.2-py2.py3-none-any.whl (15 kB)\n",
            "Requirement already satisfied: gin-config in /usr/local/lib/python3.7/dist-packages (from t5==0.7.1->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (0.4.0)\n",
            "Requirement already satisfied: nltk in /usr/local/lib/python3.7/dist-packages (from t5==0.7.1->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (3.2.5)\n",
            "Collecting mesh-tensorflow[transformer]>=0.1.13\n",
            "  Downloading mesh_tensorflow-0.1.19-py3-none-any.whl (366 kB)\n",
            "\u001b[K     |████████████████████████████████| 366 kB 64.3 MB/s \n",
            "\u001b[?25hCollecting tfds-nightly\n",
            "  Downloading tfds_nightly-4.4.0.dev202110180107-py3-none-any.whl (4.0 MB)\n",
            "\u001b[K     |████████████████████████████████| 4.0 MB 51.9 MB/s \n",
            "\u001b[?25hCollecting rouge-score\n",
            "  Downloading rouge_score-0.0.4-py2.py3-none-any.whl (22 kB)\n",
            "Requirement already satisfied: babel in /usr/local/lib/python3.7/dist-packages (from t5==0.7.1->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (2.9.1)\n",
            "Collecting tensorflow-text\n",
            "  Downloading tensorflow_text-2.6.0-cp37-cp37m-manylinux1_x86_64.whl (4.4 MB)\n",
            "\u001b[K     |████████████████████████████████| 4.4 MB 55.9 MB/s \n",
            "\u001b[?25hCollecting GitPython>=1.0.0\n",
            "  Downloading GitPython-3.1.24-py3-none-any.whl (180 kB)\n",
            "\u001b[K     |████████████████████████████████| 180 kB 63.3 MB/s \n",
            "\u001b[?25hCollecting yaspin>=1.0.0\n",
            "  Downloading yaspin-2.1.0-py3-none-any.whl (18 kB)\n",
            "Collecting configparser>=3.8.1\n",
            "  Downloading configparser-5.0.2-py3-none-any.whl (19 kB)\n",
            "Collecting docker-pycreds>=0.4.0\n",
            "  Downloading docker_pycreds-0.4.0-py2.py3-none-any.whl (9.0 kB)\n",
            "Requirement already satisfied: promise<3,>=2.0 in /usr/local/lib/python3.7/dist-packages (from wandb>=0.11.2->-r mesh-transformer-jax/requirements.txt (line 3)) (2.3)\n",
            "Requirement already satisfied: python-dateutil>=2.6.1 in /usr/local/lib/python3.7/dist-packages (from wandb>=0.11.2->-r mesh-transformer-jax/requirements.txt (line 3)) (2.8.2)\n",
            "Collecting pathtools\n",
            "  Downloading pathtools-0.1.2.tar.gz (11 kB)\n",
            "Collecting shortuuid>=0.5.0\n",
            "  Downloading shortuuid-1.0.1-py3-none-any.whl (7.5 kB)\n",
            "Collecting sentry-sdk>=1.0.0\n",
            "  Downloading sentry_sdk-1.4.3-py2.py3-none-any.whl (139 kB)\n",
            "\u001b[K     |████████████████████████████████| 139 kB 63.2 MB/s \n",
            "\u001b[?25hCollecting subprocess32>=3.5.3\n",
            "  Downloading subprocess32-3.5.4.tar.gz (97 kB)\n",
            "\u001b[K     |████████████████████████████████| 97 kB 6.4 MB/s \n",
            "\u001b[?25hRequirement already satisfied: psutil>=5.0.0 in /usr/local/lib/python3.7/dist-packages (from wandb>=0.11.2->-r mesh-transformer-jax/requirements.txt (line 3)) (5.4.8)\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 5)) (3.0.4)\n",
            "Requirement 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 5)) (2.10)\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 5)) (1.24.3)\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 5)) (2021.5.30)\n",
            "Collecting paramiko>=2.4\n",
            "  Downloading paramiko-2.8.0-py2.py3-none-any.whl (206 kB)\n",
            "\u001b[K     |████████████████████████████████| 206 kB 62.2 MB/s \n",
            "\u001b[?25hCollecting invoke<2.0,>=1.3\n",
            "  Downloading invoke-1.6.0-py3-none-any.whl (212 kB)\n",
            "\u001b[K     |████████████████████████████████| 212 kB 68.8 MB/s \n",
            "\u001b[?25hCollecting pathlib2\n",
            "  Downloading pathlib2-2.3.6-py2.py3-none-any.whl (17 kB)\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 11)) (3.3.0)\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 12)) (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 12)) (1.1.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 12)) (1.0.1)\n",
            "Collecting keras-nightly~=2.5.0.dev\n",
            "  Downloading keras_nightly-2.5.0.dev2021032900-py2.py3-none-any.whl (1.2 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.2 MB 57.8 MB/s \n",
            "\u001b[?25hCollecting tensorflow-estimator<2.6.0,>=2.5.0\n",
            "  Downloading tensorflow_estimator-2.5.0-py2.py3-none-any.whl (462 kB)\n",
            "\u001b[K     |████████████████████████████████| 462 kB 65.6 MB/s \n",
            "\u001b[?25hRequirement already satisfied: h5py~=3.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.5.0->-r mesh-transformer-jax/requirements.txt (line 14)) (3.1.0)\n",
            "Requirement already satisfied: astunparse~=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.5.0->-r mesh-transformer-jax/requirements.txt (line 14)) (1.6.3)\n",
            "Requirement already satisfied: termcolor~=1.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.5.0->-r mesh-transformer-jax/requirements.txt (line 14)) (1.1.0)\n",
            "Requirement already satisfied: tensorboard~=2.5 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.5.0->-r mesh-transformer-jax/requirements.txt (line 14)) (2.6.0)\n",
            "Collecting grpcio>=1.28.1\n",
            "  Downloading grpcio-1.34.1-cp37-cp37m-manylinux2014_x86_64.whl (4.0 MB)\n",
            "\u001b[K     |████████████████████████████████| 4.0 MB 56.0 MB/s \n",
            "\u001b[?25hRequirement already satisfied: wheel~=0.35 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.5.0->-r mesh-transformer-jax/requirements.txt (line 14)) (0.37.0)\n",
            "Requirement already satisfied: google-pasta~=0.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.5.0->-r mesh-transformer-jax/requirements.txt (line 14)) (0.2.0)\n",
            "Requirement already satisfied: gast==0.4.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.5.0->-r mesh-transformer-jax/requirements.txt (line 14)) (0.4.0)\n",
            "Requirement already satisfied: wrapt~=1.12.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.5.0->-r mesh-transformer-jax/requirements.txt (line 14)) (1.12.1)\n",
            "Requirement already satisfied: flatbuffers~=1.12.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.5.0->-r mesh-transformer-jax/requirements.txt (line 14)) (1.12)\n",
            "Requirement already satisfied: keras-preprocessing~=1.1.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow-cpu~=2.5.0->-r mesh-transformer-jax/requirements.txt (line 14)) (1.1.2)\n",
            "Collecting google-cloud-core<2.0dev,>=1.4.1\n",
            "  Downloading google_cloud_core-1.7.2-py2.py3-none-any.whl (28 kB)\n",
            "Requirement 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 15)) (1.35.0)\n",
            "Collecting google-resumable-media<2.0dev,>=1.2.0\n",
            "  Downloading google_resumable_media-1.3.3-py2.py3-none-any.whl (75 kB)\n",
            "\u001b[K     |████████████████████████████████| 75 kB 4.6 MB/s \n",
            "\u001b[?25hCollecting tokenizers<0.11,>=0.10.1\n",
            "  Downloading tokenizers-0.10.3-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (3.3 MB)\n",
            "\u001b[K     |████████████████████████████████| 3.3 MB 53.7 MB/s \n",
            "\u001b[?25hCollecting huggingface-hub>=0.0.17\n",
            "  Downloading huggingface_hub-0.0.19-py3-none-any.whl (56 kB)\n",
            "\u001b[K     |████████████████████████████████| 56 kB 4.5 MB/s \n",
            "\u001b[?25hCollecting sacremoses\n",
            "  Downloading sacremoses-0.0.46-py3-none-any.whl (895 kB)\n",
            "\u001b[K     |████████████████████████████████| 895 kB 66.0 MB/s \n",
            "\u001b[?25hCollecting pyyaml\n",
            "  Downloading PyYAML-6.0-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (596 kB)\n",
            "\u001b[K     |████████████████████████████████| 596 kB 67.5 MB/s \n",
            "\u001b[?25hCollecting rehash\n",
            "  Downloading rehash-1.0.0-py2.py3-none-any.whl (8.7 kB)\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 7)) (0.1.6)\n",
            "Requirement 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 7)) (0.11.1)\n",
            "Collecting xxhash\n",
            "  Downloading xxhash-2.0.2-cp37-cp37m-manylinux2010_x86_64.whl (243 kB)\n",
            "\u001b[K     |████████████████████████████████| 243 kB 65.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: dill in /usr/local/lib/python3.7/dist-packages (from datasets>=1.2.1->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (0.3.4)\n",
            "Collecting fsspec[http]>=2021.05.0\n",
            "  Downloading fsspec-2021.10.1-py3-none-any.whl (125 kB)\n",
            "\u001b[K     |████████████████████████████████| 125 kB 66.5 MB/s \n",
            "\u001b[?25hCollecting datasets>=1.2.1\n",
            "  Downloading datasets-1.13.2-py3-none-any.whl (287 kB)\n",
            "\u001b[K     |████████████████████████████████| 287 kB 62.8 MB/s \n",
            "\u001b[?25h  Downloading datasets-1.13.1-py3-none-any.whl (287 kB)\n",
            "\u001b[K     |████████████████████████████████| 287 kB 64.2 MB/s \n",
            "\u001b[?25h  Downloading datasets-1.13.0-py3-none-any.whl (285 kB)\n",
            "\u001b[K     |████████████████████████████████| 285 kB 65.8 MB/s \n",
            "\u001b[?25h  Downloading datasets-1.12.1-py3-none-any.whl (270 kB)\n",
            "\u001b[K     |████████████████████████████████| 270 kB 66.4 MB/s \n",
            "\u001b[?25h  Downloading datasets-1.12.0-py3-none-any.whl (269 kB)\n",
            "\u001b[K     |████████████████████████████████| 269 kB 66.9 MB/s \n",
            "\u001b[?25h  Downloading datasets-1.11.0-py3-none-any.whl (264 kB)\n",
            "\u001b[K     |████████████████████████████████| 264 kB 68.7 MB/s \n",
            "\u001b[?25hRequirement already satisfied: multiprocess in /usr/local/lib/python3.7/dist-packages (from datasets>=1.2.1->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (0.70.12.2)\n",
            "Requirement already satisfied: pyarrow!=4.0.0,>=1.0.0 in /usr/local/lib/python3.7/dist-packages (from datasets>=1.2.1->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (3.0.0)\n",
            "Collecting gitdb<5,>=4.0.1\n",
            "  Downloading gitdb-4.0.7-py3-none-any.whl (63 kB)\n",
            "\u001b[K     |████████████████████████████████| 63 kB 1.7 MB/s \n",
            "\u001b[?25hCollecting smmap<5,>=3.0.1\n",
            "  Downloading smmap-4.0.0-py2.py3-none-any.whl (24 kB)\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 15)) (0.2.8)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4 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 15)) (4.7.2)\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 15)) (4.2.4)\n",
            "Requirement already satisfied: google-api-core<2.0.0dev,>=1.21.0 in /usr/local/lib/python3.7/dist-packages (from google-cloud-core<2.0dev,>=1.4.1->google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 15)) (1.26.3)\n",
            "Requirement already satisfied: pytz in /usr/local/lib/python3.7/dist-packages (from google-api-core<2.0.0dev,>=1.21.0->google-cloud-core<2.0dev,>=1.4.1->google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 15)) (2018.9)\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.0dev,>=1.21.0->google-cloud-core<2.0dev,>=1.4.1->google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 15)) (1.53.0)\n",
            "Collecting google-crc32c<2.0dev,>=1.0\n",
            "  Downloading google_crc32c-1.3.0-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (38 kB)\n",
            "Requirement already satisfied: cached-property in /usr/local/lib/python3.7/dist-packages (from h5py~=3.1.0->tensorflow-cpu~=2.5.0->-r mesh-transformer-jax/requirements.txt (line 14)) (1.5.2)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata>=0.12->pytest==6.2.3->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (3.6.0)\n",
            "Requirement 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 12)) (2.0.1)\n",
            "Requirement already satisfied: future in /usr/local/lib/python3.7/dist-packages (from mesh-tensorflow[transformer]>=0.1.13->t5==0.7.1->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (0.16.0)\n",
            "Requirement already satisfied: tensorflow-datasets in /usr/local/lib/python3.7/dist-packages (from mesh-tensorflow[transformer]>=0.1.13->t5==0.7.1->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (4.0.1)\n",
            "Requirement already satisfied: pyparsing>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->pytest==6.2.3->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (2.4.7)\n",
            "Collecting cryptography>=2.5\n",
            "  Downloading cryptography-35.0.0-cp36-abi3-manylinux_2_24_x86_64.whl (3.5 MB)\n",
            "\u001b[K     |████████████████████████████████| 3.5 MB 54.4 MB/s \n",
            "\u001b[?25hCollecting pynacl>=1.0.1\n",
            "  Downloading PyNaCl-1.4.0-cp35-abi3-manylinux1_x86_64.whl (961 kB)\n",
            "\u001b[K     |████████████████████████████████| 961 kB 65.0 MB/s \n",
            "\u001b[?25hCollecting bcrypt>=3.1.3\n",
            "  Downloading bcrypt-3.2.0-cp36-abi3-manylinux2010_x86_64.whl (63 kB)\n",
            "\u001b[K     |████████████████████████████████| 63 kB 2.2 MB/s \n",
            "\u001b[?25hRequirement already satisfied: cffi>=1.1 in /usr/local/lib/python3.7/dist-packages (from bcrypt>=3.1.3->paramiko>=2.4->fabric~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 6)) (1.14.6)\n",
            "Requirement already satisfied: pycparser in /usr/local/lib/python3.7/dist-packages (from cffi>=1.1->bcrypt>=3.1.3->paramiko>=2.4->fabric~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 6)) (2.20)\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 15)) (0.4.8)\n",
            "Collecting threadpoolctl>=2.0.0\n",
            "  Downloading threadpoolctl-3.0.0-py3-none-any.whl (14 kB)\n",
            "Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.7/dist-packages (from scikit-learn>=0.24.1->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (1.0.1)\n",
            "Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow-cpu~=2.5.0->-r mesh-transformer-jax/requirements.txt (line 14)) (0.6.1)\n",
            "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow-cpu~=2.5.0->-r mesh-transformer-jax/requirements.txt (line 14)) (0.4.6)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow-cpu~=2.5.0->-r mesh-transformer-jax/requirements.txt (line 14)) (3.3.4)\n",
            "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow-cpu~=2.5.0->-r mesh-transformer-jax/requirements.txt (line 14)) (1.8.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.5->tensorflow-cpu~=2.5.0->-r mesh-transformer-jax/requirements.txt (line 14)) (1.3.0)\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.5->tensorflow-cpu~=2.5.0->-r mesh-transformer-jax/requirements.txt (line 14)) (3.1.1)\n",
            "Requirement already satisfied: wcwidth in /usr/local/lib/python3.7/dist-packages (from ftfy->-r mesh-transformer-jax/requirements.txt (line 19)) (0.2.5)\n",
            "Collecting starlette==0.16.0\n",
            "  Downloading starlette-0.16.0-py3-none-any.whl (61 kB)\n",
            "\u001b[K     |████████████████████████████████| 61 kB 228 kB/s \n",
            "\u001b[?25hCollecting anyio<4,>=3.0.0\n",
            "  Downloading anyio-3.3.4-py3-none-any.whl (78 kB)\n",
            "\u001b[K     |████████████████████████████████| 78 kB 5.9 MB/s \n",
            "\u001b[?25hCollecting sniffio>=1.1\n",
            "  Downloading sniffio-1.2.0-py3-none-any.whl (10 kB)\n",
            "Collecting h11>=0.8\n",
            "  Downloading h11-0.12.0-py3-none-any.whl (54 kB)\n",
            "\u001b[K     |████████████████████████████████| 54 kB 3.1 MB/s \n",
            "\u001b[?25hCollecting asgiref>=3.4.0\n",
            "  Downloading asgiref-3.4.1-py3-none-any.whl (25 kB)\n",
            "Collecting typer<1.0.0,>=0.3.0\n",
            "  Downloading typer-0.4.0-py3-none-any.whl (27 kB)\n",
            "Collecting yarl<2.0,>=1.0\n",
            "  Downloading yarl-1.7.0-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (271 kB)\n",
            "\u001b[K     |████████████████████████████████| 271 kB 63.4 MB/s \n",
            "\u001b[?25hCollecting multidict<7.0,>=4.5\n",
            "  Downloading multidict-5.2.0-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (160 kB)\n",
            "\u001b[K     |████████████████████████████████| 160 kB 64.4 MB/s \n",
            "\u001b[?25hCollecting async-timeout<4.0,>=3.0\n",
            "  Downloading async_timeout-3.0.1-py3-none-any.whl (8.2 kB)\n",
            "Requirement already satisfied: cython in /usr/local/lib/python3.7/dist-packages (from DyNet->nagisa==0.2.7->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (0.29.24)\n",
            "Requirement already satisfied: nvidia-ml-py3>=7.352.0 in /usr/local/lib/python3.7/dist-packages (from gpustat->ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (7.352.0)\n",
            "Collecting blessings>=1.6\n",
            "  Downloading blessings-1.7-py3-none-any.whl (18 kB)\n",
            "Collecting opencensus-context==0.1.2\n",
            "  Downloading opencensus_context-0.1.2-py2.py3-none-any.whl (4.4 kB)\n",
            "Collecting tensorflow\n",
            "  Downloading tensorflow-2.5.1-cp37-cp37m-manylinux2010_x86_64.whl (454.4 MB)\n",
            "\u001b[K     |████████████████████████████████| 454.4 MB 4.8 kB/s \n",
            "\u001b[?25hRequirement already satisfied: importlib-resources in /usr/local/lib/python3.7/dist-packages (from tensorflow-datasets->mesh-tensorflow[transformer]>=0.1.13->t5==0.7.1->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (5.2.2)\n",
            "Requirement already satisfied: tensorflow-metadata in /usr/local/lib/python3.7/dist-packages (from tensorflow-datasets->mesh-tensorflow[transformer]>=0.1.13->t5==0.7.1->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (1.2.0)\n",
            "Requirement already satisfied: tensorflow-hub>=0.8.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow-text->t5==0.7.1->lm-eval==0.0.1->-r mesh-transformer-jax/requirements.txt (line 9)) (0.12.0)\n",
            "INFO: pip is looking at multiple versions of tensorflow-text to determine which version is compatible with other requirements. This could take a while.\n",
            "Collecting tensorflow-text\n",
            "  Downloading tensorflow_text-2.5.0-cp37-cp37m-manylinux1_x86_64.whl (4.3 MB)\n",
            "\u001b[K     |████████████████████████████████| 4.3 MB 57.5 MB/s \n",
            "\u001b[?25hBuilding wheels for collected packages: dm-haiku, lm-eval, bleurt, black, openai, pycountry, sqlitedict, subprocess32, func-timeout, ftfy, gpustat, pathtools\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-py3-none-any.whl size=534377 sha256=2a5f251b119000d538c8f8133781f2a10926b6a5d52b6f6ad919c4fd135a4916\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-40yax3y5/wheels/06/28/69/ebaac5b2435641427299f29d88d005fb4e2627f4a108f0bdbc\n",
            "  Building wheel for lm-eval (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for lm-eval: filename=lm_eval-0.0.1-py3-none-any.whl size=108310 sha256=1af0d40cef5cc6ca6cc16526d51e5ee55880e7e10e30b941612b43901f0792da\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-40yax3y5/wheels/35/22/ec/71da337686fde904ec27f675d1bf73c4f253c48a6174d822b0\n",
            "  Building wheel for bleurt (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for bleurt: filename=BLEURT-0.0.2-py3-none-any.whl size=16454023 sha256=d1cf639ddc9cede5c89f98f629ea465260fb80ecf6e61f0ff71efb1646662697\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-40yax3y5/wheels/e4/25/b5/d3b124fc2717a794a8dcb176fac9a01bec46e9f3ff0dbf5d37\n",
            "  Building wheel for black (PEP 517) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for black: filename=black-20.8b1-py3-none-any.whl size=124196 sha256=e23c86ed7eee44b58661b020ee98fc660fc68cfa0f481abc56778cc3c2a3a2c7\n",
            "  Stored in directory: /root/.cache/pip/wheels/c5/85/79/f3af8daaf8037c0bf14beb3b7a1511a39b6e6902ca2aaf494e\n",
            "  Building wheel for openai (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for openai: filename=openai-0.6.4-py3-none-any.whl size=172215 sha256=351a77de7e77a0176dd697ca6b0ba3d3a94e57afcade72a4e0ce51b14b2b4928\n",
            "  Stored in directory: /root/.cache/pip/wheels/54/c4/02/aa519fe2aaf97a9bba197a8585182c8c07802760351afdb64a\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=e94d8e43e06b3f83de815fd22feac0de9fc818df6dc13360b75af147ea904a93\n",
            "  Stored in directory: /root/.cache/pip/wheels/57/e8/3f/120ccc1ff7541c108bc5d656e2a14c39da0d824653b62284c6\n",
            "  Building wheel for sqlitedict (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for sqlitedict: filename=sqlitedict-1.6.0-py3-none-any.whl size=14712 sha256=e4e9df4fc174c607940f98f4c68e0931919c90b08f69bf21c2bfa4e839800430\n",
            "  Stored in directory: /root/.cache/pip/wheels/58/dd/2c/0a57aadf6a7f26bec0af66d742c50af74d11967780f0bb7a7d\n",
            "  Building wheel for subprocess32 (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for subprocess32: filename=subprocess32-3.5.4-py3-none-any.whl size=6502 sha256=a5057a19db52d405557c3d48586de83e7bed7b10ed840eb168aad0442b65dfca\n",
            "  Stored in directory: /root/.cache/pip/wheels/50/ca/fa/8fca8d246e64f19488d07567547ddec8eb084e8c0d7a59226a\n",
            "  Building wheel for func-timeout (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for func-timeout: filename=func_timeout-4.3.5-py3-none-any.whl size=15095 sha256=cb312a59810e5b6629784a650ea14b917159bf5424c9eeaa99ce3ed6aedac5f5\n",
            "  Stored in directory: /root/.cache/pip/wheels/68/b5/a5/67c4364c354e141f5a1bd3ec568126f77877ab7554cf5af8cb\n",
            "  Building wheel for ftfy (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for ftfy: filename=ftfy-6.0.3-py3-none-any.whl size=41933 sha256=fb6b147eb1514da2d0a1619445f3afaeb7f50192424c8bad4642e31f6febe60b\n",
            "  Stored in directory: /root/.cache/pip/wheels/19/f5/38/273eb3b5e76dfd850619312f693716ac4518b498f5ffb6f56d\n",
            "  Building wheel for gpustat (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for gpustat: filename=gpustat-0.6.0-py3-none-any.whl size=12617 sha256=d015eb223d5b2d1d6212da2e772d46060ba19b0c8dad6703a4d293c152e16d0c\n",
            "  Stored in directory: /root/.cache/pip/wheels/e6/67/af/f1ad15974b8fd95f59a63dbf854483ebe5c7a46a93930798b8\n",
            "  Building wheel for pathtools (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pathtools: filename=pathtools-0.1.2-py3-none-any.whl size=8807 sha256=2d70770d479a0bc1e2dae7a76080ce65cfee013f8bfcaf66da850c43ca33c177\n",
            "  Stored in directory: /root/.cache/pip/wheels/3e/31/09/fa59cef12cdcfecc627b3d24273699f390e71828921b2cbba2\n",
            "Successfully built dm-haiku lm-eval bleurt black openai pycountry sqlitedict subprocess32 func-timeout ftfy gpustat pathtools\n",
            "Installing collected packages: requests, mbstrdecoder, typepy, multidict, grpcio, yarl, tqdm, tensorflow-estimator, regex, pyyaml, keras-nightly, async-timeout, tokenizers, threadpoolctl, tensorflow, sniffio, smmap, sacremoses, portalocker, opencensus-context, mesh-tensorflow, huggingface-hub, google-crc32c, DataProperty, blessings, aiohttp, zstandard, xxhash, ujson, typed-ast, transformers, tfds-nightly, tf-slim, tensorflow-text, tcolorpy, tabledata, sentencepiece, scikit-learn, sacrebleu, rouge-score, rehash, redis, pynacl, pydantic, py-spy, pluggy, pathvalidate, pathspec, opencensus, mypy-extensions, msgfy, jsonlines, gpustat, google-resumable-media, google-cloud-core, gitdb, fsspec, DyNet, cryptography, colorama, bcrypt, anyio, aioredis, aiohttp-cors, yaspin, typer, tqdm-multiprocess, t5, subprocess32, starlette, sqlitedict, shortuuid, sentry-sdk, ray, pytest, pytablewriter, pycountry, pybind11, pathtools, pathlib2, paramiko, openai, numexpr, nagisa, mock, lm-dataformat, jmp, invoke, h11, google-cloud-storage, GitPython, docker-pycreds, datasets, configparser, colorful, chex, bleurt, black, best-download, asgiref, wandb, uvicorn, tensorflow-cpu, pathy, optax, lm-eval, func-timeout, ftfy, fastapi, fabric, einops, dm-haiku\n",
            "  Attempting uninstall: requests\n",
            "    Found existing installation: requests 2.23.0\n",
            "    Uninstalling requests-2.23.0:\n",
            "      Successfully uninstalled requests-2.23.0\n",
            "  Attempting uninstall: grpcio\n",
            "    Found existing installation: grpcio 1.41.0\n",
            "    Uninstalling grpcio-1.41.0:\n",
            "      Successfully uninstalled grpcio-1.41.0\n",
            "  Attempting uninstall: tqdm\n",
            "    Found existing installation: tqdm 4.62.3\n",
            "    Uninstalling tqdm-4.62.3:\n",
            "      Successfully uninstalled tqdm-4.62.3\n",
            "  Attempting uninstall: tensorflow-estimator\n",
            "    Found existing installation: tensorflow-estimator 2.6.0\n",
            "    Uninstalling tensorflow-estimator-2.6.0:\n",
            "      Successfully uninstalled tensorflow-estimator-2.6.0\n",
            "  Attempting uninstall: regex\n",
            "    Found existing installation: regex 2019.12.20\n",
            "    Uninstalling regex-2019.12.20:\n",
            "      Successfully uninstalled regex-2019.12.20\n",
            "  Attempting uninstall: pyyaml\n",
            "    Found existing installation: PyYAML 3.13\n",
            "    Uninstalling PyYAML-3.13:\n",
            "      Successfully uninstalled PyYAML-3.13\n",
            "  Attempting uninstall: tensorflow\n",
            "    Found existing installation: tensorflow 2.6.0\n",
            "    Uninstalling tensorflow-2.6.0:\n",
            "      Successfully uninstalled tensorflow-2.6.0\n",
            "  Attempting uninstall: scikit-learn\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",
            "  Attempting uninstall: pluggy\n",
            "    Found existing installation: pluggy 0.7.1\n",
            "    Uninstalling pluggy-0.7.1:\n",
            "      Successfully uninstalled pluggy-0.7.1\n",
            "  Attempting uninstall: google-resumable-media\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",
            "  Attempting uninstall: google-cloud-core\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",
            "  Attempting uninstall: pytest\n",
            "    Found existing installation: pytest 3.6.4\n",
            "    Uninstalling pytest-3.6.4:\n",
            "      Successfully uninstalled pytest-3.6.4\n",
            "  Attempting uninstall: numexpr\n",
            "    Found existing installation: numexpr 2.7.3\n",
            "    Uninstalling numexpr-2.7.3:\n",
            "      Successfully uninstalled numexpr-2.7.3\n",
            "  Attempting uninstall: google-cloud-storage\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",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "panel 0.12.1 requires tqdm>=4.48.0, but you have tqdm 4.45.0 which is incompatible.\n",
            "google-colab 1.0.0 requires requests~=2.23.0, but you have requests 2.25.1 which is incompatible.\n",
            "google-cloud-bigquery 1.21.0 requires google-resumable-media!=0.4.0,<0.5.0dev,>=0.3.1, but you have google-resumable-media 1.3.3 which is incompatible.\n",
            "datascience 0.10.6 requires folium==0.2.1, but you have folium 0.8.3 which is incompatible.\u001b[0m\n",
            "Successfully installed DataProperty-0.54.2 DyNet-2.1.2 GitPython-3.1.24 aiohttp-3.7.4.post0 aiohttp-cors-0.7.0 aioredis-2.0.0 anyio-3.3.4 asgiref-3.4.1 async-timeout-3.0.1 bcrypt-3.2.0 best-download-0.0.7 black-20.8b1 blessings-1.7 bleurt-0.0.2 chex-0.0.8 colorama-0.4.4 colorful-0.5.4 configparser-5.0.2 cryptography-35.0.0 datasets-1.11.0 dm-haiku-0.0.5.dev0 docker-pycreds-0.4.0 einops-0.3.2 fabric-2.6.0 fastapi-0.70.0 fsspec-2021.10.1 ftfy-6.0.3 func-timeout-4.3.5 gitdb-4.0.7 google-cloud-core-1.7.2 google-cloud-storage-1.36.2 google-crc32c-1.3.0 google-resumable-media-1.3.3 gpustat-0.6.0 grpcio-1.34.1 h11-0.12.0 huggingface-hub-0.0.19 invoke-1.6.0 jmp-0.0.2 jsonlines-2.0.0 keras-nightly-2.5.0.dev2021032900 lm-dataformat-0.0.19 lm-eval-0.0.1 mbstrdecoder-1.1.0 mesh-tensorflow-0.1.19 mock-4.0.3 msgfy-0.2.0 multidict-5.2.0 mypy-extensions-0.4.3 nagisa-0.2.7 numexpr-2.7.2 openai-0.6.4 opencensus-0.8.0 opencensus-context-0.1.2 optax-0.0.6 paramiko-2.8.0 pathlib2-2.3.6 pathspec-0.9.0 pathtools-0.1.2 pathvalidate-2.5.0 pathy-0.6.0 pluggy-0.13.1 portalocker-2.3.2 py-spy-0.3.10 pybind11-2.6.2 pycountry-20.7.3 pydantic-1.8.2 pynacl-1.4.0 pytablewriter-0.58.0 pytest-6.2.3 pyyaml-6.0 ray-1.4.1 redis-3.5.3 regex-2021.10.8 rehash-1.0.0 requests-2.25.1 rouge-score-0.0.4 sacrebleu-1.5.0 sacremoses-0.0.46 scikit-learn-1.0 sentencepiece-0.1.96 sentry-sdk-1.4.3 shortuuid-1.0.1 smmap-4.0.0 sniffio-1.2.0 sqlitedict-1.6.0 starlette-0.16.0 subprocess32-3.5.4 t5-0.7.1 tabledata-1.3.0 tcolorpy-0.1.1 tensorflow-2.5.1 tensorflow-cpu-2.5.1 tensorflow-estimator-2.5.0 tensorflow-text-2.5.0 tf-slim-1.1.0 tfds-nightly-4.4.0.dev202110180107 threadpoolctl-3.0.0 tokenizers-0.10.3 tqdm-4.45.0 tqdm-multiprocess-0.0.11 transformers-4.11.3 typed-ast-1.4.3 typepy-1.3.0 typer-0.4.0 ujson-4.2.0 uvicorn-0.15.0 wandb-0.12.4 xxhash-2.0.2 yarl-1.7.0 yaspin-2.1.0 zstandard-0.15.2\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "google"
                ]
              }
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Processing ./mesh-transformer-jax\n",
            "\u001b[33m  DEPRECATION: A future pip version will change local packages to be built in-place without first copying to a temporary directory. We recommend you use --use-feature=in-tree-build to test your packages with this new behavior before it becomes the default.\n",
            "   pip 21.3 will remove support for this functionality. You can find discussion regarding this at https://github.com/pypa/pip/issues/7555.\u001b[0m\n",
            "Collecting jax==0.2.12\n",
            "  Downloading jax-0.2.12.tar.gz (590 kB)\n",
            "\u001b[K     |████████████████████████████████| 590 kB 5.3 MB/s \n",
            "\u001b[?25hCollecting tensorflow==2.5.0\n",
            "  Downloading tensorflow-2.5.0-cp37-cp37m-manylinux2010_x86_64.whl (454.3 MB)\n",
            "\u001b[K     |████████████████████████████████| 454.3 MB 11 kB/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: tensorflow-estimator<2.6.0,>=2.5.0rc0 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.5.0) (2.5.0)\n",
            "Requirement already satisfied: h5py~=3.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.5.0) (3.1.0)\n",
            "Requirement already satisfied: flatbuffers~=1.12.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.5.0) (1.12)\n",
            "Requirement already satisfied: wheel~=0.35 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.5.0) (0.37.0)\n",
            "Requirement already satisfied: gast==0.4.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.5.0) (0.4.0)\n",
            "Requirement already satisfied: termcolor~=1.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.5.0) (1.1.0)\n",
            "Requirement already satisfied: google-pasta~=0.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.5.0) (0.2.0)\n",
            "Requirement already satisfied: keras-nightly~=2.5.0.dev in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.5.0) (2.5.0.dev2021032900)\n",
            "Requirement already satisfied: wrapt~=1.12.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.5.0) (1.12.1)\n",
            "Requirement already satisfied: keras-preprocessing~=1.1.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.5.0) (1.1.2)\n",
            "Requirement already satisfied: six~=1.15.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.5.0) (1.15.0)\n",
            "Requirement already satisfied: astunparse~=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.5.0) (1.6.3)\n",
            "Requirement already satisfied: tensorboard~=2.5 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.5.0) (2.6.0)\n",
            "Requirement already satisfied: grpcio~=1.34.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.5.0) (1.34.1)\n",
            "Requirement already satisfied: protobuf>=3.9.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.5.0) (3.17.3)\n",
            "Requirement already satisfied: typing-extensions~=3.7.4 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.5.0) (3.7.4.3)\n",
            "Requirement already satisfied: cached-property in /usr/local/lib/python3.7/dist-packages (from h5py~=3.1.0->tensorflow==2.5.0) (1.5.2)\n",
            "Requirement already satisfied: google-auth<2,>=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow==2.5.0) (1.35.0)\n",
            "Requirement already satisfied: setuptools>=41.0.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow==2.5.0) (57.4.0)\n",
            "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow==2.5.0) (1.8.0)\n",
            "Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow==2.5.0) (0.6.1)\n",
            "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow==2.5.0) (0.4.6)\n",
            "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow==2.5.0) (2.25.1)\n",
            "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow==2.5.0) (1.0.1)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.5->tensorflow==2.5.0) (3.3.4)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard~=2.5->tensorflow==2.5.0) (0.2.8)\n",
            "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard~=2.5->tensorflow==2.5.0) (4.2.4)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard~=2.5->tensorflow==2.5.0) (4.7.2)\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.5->tensorflow==2.5.0) (1.3.0)\n",
            "Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.7/dist-packages (from markdown>=2.6.8->tensorboard~=2.5->tensorflow==2.5.0) (4.8.1)\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,>=1.6.3->tensorboard~=2.5->tensorflow==2.5.0) (0.4.8)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard~=2.5->tensorflow==2.5.0) (2021.5.30)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard~=2.5->tensorflow==2.5.0) (1.24.3)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard~=2.5->tensorflow==2.5.0) (3.0.4)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard~=2.5->tensorflow==2.5.0) (2.10)\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.5->tensorflow==2.5.0) (3.1.1)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->markdown>=2.6.8->tensorboard~=2.5->tensorflow==2.5.0) (3.6.0)\n",
            "Building wheels for collected packages: mesh-transformer, jax\n",
            "  Building wheel for mesh-transformer (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for mesh-transformer: filename=mesh_transformer-0.0.0-py3-none-any.whl size=26330 sha256=d3aa9d58a1da2380b8b20da07789df4b4e8b9db8f00744908048bd81b6476e9e\n",
            "  Stored in directory: /root/.cache/pip/wheels/56/bd/89/b1f6b2f3d6b938d0c5812ee97756a1afd32521bea293543863\n",
            "  Building wheel for jax (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for jax: filename=jax-0.2.12-py3-none-any.whl size=682487 sha256=0c52da00e15488b55a4023ed213e20554d3997c3b2b0add7552e24bf7278970e\n",
            "  Stored in directory: /root/.cache/pip/wheels/f6/4d/e5/73eec5070b77f25664c67bd793d4eb97f41bbf9be7afafd15e\n",
            "Successfully built mesh-transformer jax\n",
            "Installing collected packages: tensorflow, mesh-transformer, jax\n",
            "  Attempting uninstall: tensorflow\n",
            "    Found existing installation: tensorflow 2.5.1\n",
            "    Uninstalling tensorflow-2.5.1:\n",
            "      Successfully uninstalled tensorflow-2.5.1\n",
            "  Attempting uninstall: jax\n",
            "    Found existing installation: jax 0.2.21\n",
            "    Uninstalling jax-0.2.21:\n",
            "      Successfully uninstalled jax-0.2.21\n",
            "Successfully installed jax-0.2.12 mesh-transformer-0.0.0 tensorflow-2.5.0\n"
          ]
        }
      ]
    },
    {
      "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": 3,
      "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_lowmem\n",
        "from mesh_transformer.sampling import nucleaus_sample\n",
        "from mesh_transformer.transformer_shard import CausalTransformer"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QAgKq-X2kmba",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 213,
          "referenced_widgets": [
            "06b17c299dbc470d84ccdc428d05e8eb",
            "20ec90427edd47c3b3e91faffe1de76c",
            "a99395f09b9d4ffa9dddf219fd75f614",
            "5673aea1d57443b88ffae41a725531fd",
            "8ef4566445e44b08baee45d997708837",
            "06b943526e604b0baaf4169348efafda",
            "2916fbd3264d439097e3b4db151ef090",
            "9a7660cf79b7486ebe2832c80297f1bf",
            "48fabfd1f10549699d02eb9fc721a608",
            "f0bc3d33955841fa914dddeb2759222c",
            "b9440118187748a2af79bd8c1bb86391",
            "18f3336cacba4ad9a3863dceef684d62",
            "5a043e80fbb243619e1ec11b14dfcf82",
            "91ed7493ccb8424bab2372dadf2470f0",
            "4c532da328fc4aec9a70dbc2a7b749f5",
            "e28c6a5968e944be8a832541d05cb336",
            "0cdfd04d211f4d32a82c4ef42f7e99f9",
            "781dc409b97146e68be675a1d18328b6",
            "7538b4d51acf4c369558a1471f985ffe",
            "1e0681e3be994e96b9d5853ce47040fa",
            "20a12a8118f74f65aa03403500594aaa",
            "72f39605f935481f8f68dbf83a672306",
            "ec441bd40f124042bad3242e84bf9133",
            "e76ed536e65f43d9a1ff6f9e023657d8",
            "22f8f50f14ad42eea06fc43473208741",
            "41600b0a132b4f30bf1ac331055d8f06",
            "9aee12e4d953453680599582ae215ba0",
            "e012f372170844e18928ed98656bc310",
            "fc5441d719d045a4a1655e838b2437a9",
            "600e8b54ce4543dba851586530ce43d9",
            "fc2294cbfec7435c922aa2a67b0efa5d",
            "8a5ba8d9802044f592e435251efd3495"
          ]
        },
        "outputId": "52078521-3f25-4225-882b-89b6c22c4ab2"
      },
      "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": 5,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "06b17c299dbc470d84ccdc428d05e8eb",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=1042301.0, style=ProgressStyle(descript…"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "48fabfd1f10549699d02eb9fc721a608",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=456318.0, style=ProgressStyle(descripti…"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "0cdfd04d211f4d32a82c4ef42f7e99f9",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=1355256.0, style=ProgressStyle(descript…"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "22f8f50f14ad42eea06fc43473208741",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=665.0, style=ProgressStyle(description_…"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n"
          ]
        }
      ]
    },
    {
      "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": "ed150f15-e081-4eec-982f-dc2a8354c397"
      },
      "source": [
        "total_batch = per_replica_batch * jax.device_count() // cores_per_replica\n",
        "\n",
        "network = CausalTransformer(params)\n",
        "\n",
        "network.state = read_ckpt_lowmem(network.state, \"step_383500/\", devices.shape[1])\n",
        "\n",
        "network.state = network.move_xmap(network.state, np.zeros(cores_per_replica))"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "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"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "key shape (8, 2)\n",
            "in shape (1, 2048)\n",
            "dp 1\n",
            "mp 8\n",
            "Total parameters: 6053381344\n",
            "read from disk/gcs in 216.473s\n"
          ]
        }
      ]
    },
    {
      "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": "Wzma_IF4hqoG"
      },
      "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": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 340
        },
        "id": "ZVzs2TYlvYeX",
        "outputId": "7c19dbe9-e9a4-4209-9780-10959b1710e3"
      },
      "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": 8,
      "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": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "completion done in 56.72703957557678s\n",
            "\u001b[1mEleutherAI is\u001b[0m great for Melee, really great. The team is completely fine and everyone is really good, very competitive, and fun to watch. It's amazing to see Melee grow with the new technology and so many good players getting stronger and stronger. Once you get used to EVO, you really get into the groove of the tournament and it becomes much easier to get excited and hyped up for every upcoming tournament.\n",
            "\n",
            "Agree with this 100%, a lot of people who argue Evo is too busy or has too many tournaments don't understand the support EVO provides for Melee, not only to Smashers but for the game as a whole. EVO's growth is good for the scene and provides a great experience for the top players in the world to compete against the top talent at once and take a year to heal up and improve from.\n",
            "\n",
            "I think most of the time there's so much more excitement and hype about a couple of the large EVO events that we forget about the rest of the hundreds of smaller events around the world that provide some of the best competitive experience. There are many great smaller tournaments and even the I.S.A. hosted Evo's each year (3 last year).\n",
            "\n",
            "I would personally not mind if EVO was more busy, I just hope they focus their time more on providing quality experience for the players instead of multiple events each year.\n",
            "\n",
            "People can be really mean and competitive though, and EVO can be a tense, highly competitive event and you don't have to agree with the winning bracket or the decisions in order to be excited. I think EVO is a great time to see all the upsets happen and see how the top players are. Maybe it's also important to remember that in other games, major EVOs are more even than Smash, top level play is less prevalent in other games with the dominating NA scene. It's been a while since I've seen a big event, but personally I think EVO feels more unique and the results feel a lot more shocking than most of the other major events.\n",
            "\n",
            "I think people like myself have seen EVO a lot and just get really excited and hyped for the tournament. I haven't gone to many of these big tournaments since I stopped competing, but this year at EVO it was probably the most hype event I've seen in a while.\n",
            "\n",
            "To be honest, it's not necessarily bad to have more tournaments to attend. I just think they need to change up the flow a little bit and\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 323
        },
        "id": "nvlAK6RbCJYg",
        "outputId": "dc91490e-49c9-416c-d069-b5290c86f91b"
      },
      "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=512, context=context)[0])"
      ],
      "execution_count": 10,
      "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": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "completion done in 13.531424522399902s\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 The elusive creatures had been hiding from the researchers ever since they were discovered, and were presumed to be hostile towards mankind. When the scientists finally reached the valley of gold, they found a valley inhabited by strange creatures living in the surrounding mountains. While some of them looked like unicorns, others looked like enormous spiders. They seemed to be afraid of humans and always fled when the group reached the valley. The fact that these creatures talked perfect English did not seem to worry them. Most people would think that such creatures must be from a different planet. However, when they found out that there was a unicorn population in their very own country, they thought it must be one of the mythical creatures of legends. The scientists headed towards the valley where the unicorns lived and set up camp. The night seemed to be very quiet and most of the creatures seemed to be sleeping, however, the smallest and timid creatures seemed to be looking out for the intruders. When they started their experiments, they found that the unicorns had an interesting ability that none of the other creatures had – they could blend in with their surroundings. So what was this strange ability about? And what was this “golden valley”?\n",
            "\n",
            "A herd of unicorns has been discovered in the South American Andes Mountains. What is the meaning of “unicorn”?\n",
            "\n",
            "“A unicorn is a fabulous wild animal, about which there is a lot of folklore and myth about. A unicorn is a mythological creature from ancient Persia and elsewhere which is believed to be an animal which is made of a single horn on its head. People who see unicorns are often terrified. They believe that the unicorn is a creature made of pure gold and is a part of the nordic god Freyr. Unicorns are also believed to be living in England and were included in its folklore. Unicorns are also said to inhabit the Americas, particularly Mexico, but are very rare.”\n",
            "\n",
            "There are a lot of unicorns living in South America – how do you recognize one?\n",
            "\n",
            "“A unicorn looks like any other wild creature with an impressive horn on the top of its head. It has a long neck and four legs. Its coat is usually white, brown, or black, but it can be different colours too. The eyes are blue or yellow. It has four pointed, white hooves, a soft muzzle, and a long tail. The horn is made of yellow or white material. Unicorns are very rare animals, and\n"
          ]
        }
      ]
    }
  ]
}