CurtisASmith · July 5, 2021 13:44 · CurtisASmith · Jul 5, 2021
diff --git a/gpt-j-6b-inference.ipynb b/gpt-j-6b-inference.ipynb
 {
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "GPT-J-6B Inference.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "TPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "7a21e484a4d949a7ab2bc6b73893a79d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_e126a4a4ea614c41afa3061b07e9f6d7",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_f3cc3ebcdd134c4386f49d6053feed2b",
              "IPY_MODEL_1b6cb9571254456aa5725c9e479dbd13"
            ]
          }
        },
        "e126a4a4ea614c41afa3061b07e9f6d7": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "f3cc3ebcdd134c4386f49d6053feed2b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_b1c08c9f18f745afa7d15e02e3e22474",
            "_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_588fc7ccb6ac4539a241486ee43c2b47"
          }
        },
        "1b6cb9571254456aa5725c9e479dbd13": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_29fbc82354164f029331d9b9b64edd82",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 1.04M/1.04M [00:11&lt;00:00, 90.6kB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_d6576d87930c4f378bd6ba901df0ab42"
          }
        },
        "b1c08c9f18f745afa7d15e02e3e22474": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "588fc7ccb6ac4539a241486ee43c2b47": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "29fbc82354164f029331d9b9b64edd82": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "d6576d87930c4f378bd6ba901df0ab42": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "4b0bdb3502ed496ba6ba5dd71b34734d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_1481cbe5efd740b0a9b5cfdb070a4170",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_2ec7b9c8b993411da7f9c257270c5767",
              "IPY_MODEL_cf99535de7024be29bb825f479854a1e"
            ]
          }
        },
        "1481cbe5efd740b0a9b5cfdb070a4170": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "2ec7b9c8b993411da7f9c257270c5767": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_531dd0be2fa843019f88991187ec1429",
            "_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_89659dcf535e43068384a3a814021d81"
          }
        },
        "cf99535de7024be29bb825f479854a1e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_e5097b53318e4c13b45b478c8ae4a8ef",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 456k/456k [00:00&lt;00:00, 1.23MB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_c35fc2317e684278988f5794c9dd2bd7"
          }
        },
        "531dd0be2fa843019f88991187ec1429": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "89659dcf535e43068384a3a814021d81": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "e5097b53318e4c13b45b478c8ae4a8ef": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "c35fc2317e684278988f5794c9dd2bd7": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "e7482c6643ef47f3830c5adfaf74a55b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_9e5313c3d5f9447fbe4e5fffbece22ee",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_527b4e5e36fd464d8c1ac0d6b8c72ab1",
              "IPY_MODEL_bc9cbfaff7dc4510a4c09cd544171ba6"
            ]
          }
        },
        "9e5313c3d5f9447fbe4e5fffbece22ee": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "527b4e5e36fd464d8c1ac0d6b8c72ab1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_cc4e2dbeb7ee4b048594ff044a55dea1",
            "_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_aa813e731d9842dca235cd358bc67402"
          }
        },
        "bc9cbfaff7dc4510a4c09cd544171ba6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_92e05af2b8e54df0922abce5bf1948ba",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 1.36M/1.36M [00:10&lt;00:00, 126kB/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_6c84dd3065714825ae20c08a391b42e4"
          }
        },
        "cc4e2dbeb7ee4b048594ff044a55dea1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "aa813e731d9842dca235cd358bc67402": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "92e05af2b8e54df0922abce5bf1948ba": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "6c84dd3065714825ae20c08a391b42e4": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/CurtisASmith/71601b3ae6b10548ef7d9c98926e3881/gpt-j-6b-inference.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pHIJVqHsh4An"
      },
      "source": [
        "# GPT-J-6B Inference Demo\n",
        "\n",
        "<a href=\"http://colab.research.google.com/github/kingoflolz/mesh-transformer-jax/blob/master/colab_demo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
        "\n",
        "This notebook demonstrates how to run the [GPT-J-6B model](https://github.com/kingoflolz/mesh-transformer-jax/#GPT-J-6B). See the link for more details about the model, including evaluation metrics and credits."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8CMw_dSQKfhT"
      },
      "source": [
        "## Install Dependencies\n",
        "\n",
        "First we download the model and install some dependencies. This step takes at least 5 minutes (possibly longer depending on server load).\n",
        "\n",
        "!!! **Make sure you are using a TPU runtime!** !!!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Mn1FCCWepvWT"
      },
      "source": [
        "!printenv"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n7xAFw-LOYfe"
      },
      "source": [
        "!apt install zstd\n",
        "\n",
        "# the \"slim\" version contain only bf16 weights and no optimizer parameters, which minimizes bandwidth and memory\n",
        "!time wget -c https://the-eye.eu/public/AI/GPT-J-6B/step_383500_slim.tar.zstd\n",
        "\n",
        "!time tar -I zstd -xf step_383500_slim.tar.zstd\n",
        "\n",
        "!git clone https://github.com/kingoflolz/mesh-transformer-jax.git\n",
        "!pip install -r mesh-transformer-jax/requirements.txt\n",
        "\n",
        "# jax 0.2.12 is required due to a regression with xmap in 0.2.13\n",
        "!pip install mesh-transformer-jax/ jax==0.2.12"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "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\n",
        "from mesh_transformer.sampling import nucleaus_sample\n",
        "from mesh_transformer.transformer_shard import CausalTransformer"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QAgKq-X2kmba",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 164,
          "referenced_widgets": [
            "7a21e484a4d949a7ab2bc6b73893a79d",
            "e126a4a4ea614c41afa3061b07e9f6d7",
            "f3cc3ebcdd134c4386f49d6053feed2b",
            "1b6cb9571254456aa5725c9e479dbd13",
            "b1c08c9f18f745afa7d15e02e3e22474",
            "588fc7ccb6ac4539a241486ee43c2b47",
            "29fbc82354164f029331d9b9b64edd82",
            "d6576d87930c4f378bd6ba901df0ab42",
            "4b0bdb3502ed496ba6ba5dd71b34734d",
            "1481cbe5efd740b0a9b5cfdb070a4170",
            "2ec7b9c8b993411da7f9c257270c5767",
            "cf99535de7024be29bb825f479854a1e",
            "531dd0be2fa843019f88991187ec1429",
            "89659dcf535e43068384a3a814021d81",
            "e5097b53318e4c13b45b478c8ae4a8ef",
            "c35fc2317e684278988f5794c9dd2bd7",
            "e7482c6643ef47f3830c5adfaf74a55b",
            "9e5313c3d5f9447fbe4e5fffbece22ee",
            "527b4e5e36fd464d8c1ac0d6b8c72ab1",
            "bc9cbfaff7dc4510a4c09cd544171ba6",
            "cc4e2dbeb7ee4b048594ff044a55dea1",
            "aa813e731d9842dca235cd358bc67402",
            "92e05af2b8e54df0922abce5bf1948ba",
            "6c84dd3065714825ae20c08a391b42e4"
          ]
        },
        "outputId": "078cbb95-70a6-4a07-edb2-4037647fc63e"
      },
      "source": [
        "params = {\n",
        "  \"layers\": 28,\n",
        "  \"d_model\": 4096,\n",
        "  \"n_heads\": 16,\n",
        "  \"n_vocab\": 50400,\n",
        "  \"norm\": \"layernorm\",\n",
        "  \"pe\": \"rotary\",\n",
        "  \"pe_rotary_dims\": 64,\n",
        "\n",
        "  \"seq\": 2048,\n",
        "  \"cores_per_replica\": 8,\n",
        "  \"per_replica_batch\": 1,\n",
        "}\n",
        "\n",
        "per_replica_batch = params[\"per_replica_batch\"]\n",
        "cores_per_replica = params[\"cores_per_replica\"]\n",
        "seq = params[\"seq\"]\n",
        "\n",
        "\n",
        "params[\"sampler\"] = nucleaus_sample\n",
        "\n",
        "# here we \"remove\" the optimizer parameters from the model (as we don't need them for inference)\n",
        "params[\"optimizer\"] = optax.scale(0)\n",
        "\n",
        "mesh_shape = (jax.device_count() // cores_per_replica, cores_per_replica)\n",
        "devices = np.array(jax.devices()).reshape(mesh_shape)\n",
        "\n",
        "maps.thread_resources.env = maps.ResourceEnv(maps.Mesh(devices, ('dp', 'mp')))\n",
        "\n",
        "tokenizer = transformers.GPT2TokenizerFast.from_pretrained('gpt2')"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "7a21e484a4d949a7ab2bc6b73893a79d",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=1042301.0, style=ProgressStyle(descript…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "4b0bdb3502ed496ba6ba5dd71b34734d",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=456318.0, style=ProgressStyle(descripti…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "e7482c6643ef47f3830c5adfaf74a55b",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=1355256.0, style=ProgressStyle(descript…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yFgRkUgfiNdA"
      },
      "source": [
        "Here we create the network and load the parameters from the downloaded files. Expect this to take around 5 minutes."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lwNETD2Uk8nu",
        "outputId": "8fbe14f0-610f-4883-8011-d77e5f76a252"
      },
      "source": [
        "total_batch = per_replica_batch * jax.device_count() // cores_per_replica\n",
        "\n",
        "network = CausalTransformer(params)\n",
        "\n",
        "network.state = read_ckpt(network.state, \"step_383500/\", devices.shape[1])\n",
        "\n",
        "network.state = network.move_xmap(network.state, np.zeros(cores_per_replica))"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/jax/experimental/maps.py:412: UserWarning: xmap is an experimental feature and probably has bugs!\n",
            "  warn(\"xmap is an experimental feature and probably has bugs!\")\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "key shape (8, 2)\n",
            "in shape (1, 2048)\n",
            "dp 1\n",
            "mp 8\n",
            "read from disk/gcs in 19.5013s\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A-eT7Sw6if4J"
      },
      "source": [
        "## Run Model\n",
        "\n",
        "Finally, we are ready to infer with the model! The first sample takes around a minute due to compilation, but after that it should only take about 10 seconds per sample.\n",
        "\n",
        "Feel free to mess with the different sampling parameters (top_p and temp), as well as the length of the generations (gen_len, causes a recompile when changed).\n",
        "\n",
        "You can also change other things like per_replica_batch in the previous cells to change how many generations are done in parallel. A larger batch has higher latency but higher throughput when measured in tokens generated/s. This is useful for doing things like best-of-n cherry picking.\n",
        "\n",
        "*Tip for best results: Make sure your prompt does not have any trailing spaces, which tend to confuse the model due to the BPE tokenization used during training.*"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2egssW4KsqgO"
      },
      "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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZVzs2TYlvYeX"
      },
      "source": [
        "def infer(context, top_p=0.8, temp=1.0, gen_len=2048,ptoken=\">>prompt>>\",gtoken=\">>generation>>\"):\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\"\\n{ptoken}{context}{gtoken}{tokenizer.decode(o)}\") \n",
        "    isecs = time.time() - start\n",
        "    isecs = round(isecs,2)\n",
        "    print(f\"completion done in {isecs}s\")\n",
        "    return samples\n",
        "\n",
        "### you can uncomment to compile and run the inference once below\n",
        "#print(infer(\"Now that you've mined your own rare earth minerals, you can start building your homemade semiconductor fabricator. First, you'll need to\")[0])"
      ],
      "execution_count": 81,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "cellView": "form",
        "id": "TR0xwc6V5hO3",
        "outputId": "5e3796c6-79a5-4bc7-a7d8-beec00be369c"
      },
      "source": [
        "#@title Infer multiple to file { form-width: \"300px\" }\n",
        "from time import strftime, localtime\n",
        "\n",
        "fullt = time.time()\n",
        "top_p = 0.9 #@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "temp = 0.8 #@param {type:\"slider\", min:0, max:2, step:0.05}\n",
        "gen_len = 2048 #@param {type:\"slider\", min:1, max:2048, step:1}\n",
        "samplecount = 1 #@param {type:\"slider\", min:1, max:50, step:1}\n",
        "\n",
        "#@markdown This form is nice for erasing the string but if your prompt is complex\n",
        "#@markdown and has line breaks you'll need to add escapes like `\\n` to the code.\n",
        "#@markdown Trying to add that here in the form will add `\\\\n` to the code; \"\\n\" to the parsed prompt.\n",
        "context = \"A Brief History of Natural Language Processing with Machine Learning:\\n\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown These prompt and generation tokens/markers are intended as an easy way to search generated files.\n",
        "#@markdown You can delete or change them as you please.\n",
        "ptoken=\">>prompt>>\" #@param {type:\"string\"}\n",
        "gtoken=\">>generation>>\" #@param {type:\"string\"}\n",
        "\n",
        "### look up strftime if you want to customize the timestamp. I prefer to delete %S for seconds when doing bulk generations longer than a minute\n",
        "timestamp = strftime(\"%Y-%d%b-%H%M-%S\",localtime()) \n",
        "\n",
        "#@markdown This section customizes the filename.<br>\n",
        "#@markdown By default and to ensure filenames are always unique, the name format is\n",
        "#@markdown `s[numberofsamples]-[year]-[day][month]-[hour][minute]-[seconds]-[suffix].txt`.\n",
        "#@markdown Toggle the `simplesave` option to simply save the file as you've named it.\n",
        "#@markdown New generations will simply append to the last (I regret not making this option immediately)\n",
        "simplesave = True #@param {type:\"boolean\"}\n",
        "suffix = \"ABriefHistoryNLP\" #@param {type:\"string\"}\n",
        "if simplesave==True: new = f\"{suffix}.txt\"\n",
        "if simplesave==False: new = f\"s{samplecount}-{timestamp}-{suffix}.txt\" \n",
        "\n",
        "# Now loop through the samples and write them to a new file\n",
        "with open(new,'a',encoding='utf8') as newfile:\n",
        "  s = 0\n",
        "  while s < samplecount:\n",
        "    s+=1\n",
        "    result=infer(top_p=top_p, temp=temp, gen_len=gen_len, context=context)[0]\n",
        "    #result=\"test!\" #to test changes without running a full inference cycle\n",
        "    newfile.write(result)\n",
        "    print(f\"Wrote sample {s}!\")\n",
        "\n",
        "plu = ''\n",
        "if samplecount>1: plu = 's'\n",
        "secs = time.time() - fullt\n",
        "secs = round(secs,2)\n",
        "\n",
        "print(f\"Wrote {samplecount} sample{plu} to {new} in {secs} seconds.\")\n",
        "\n",
        "### if you only requested one generation it will print below\n",
        "if samplecount==1: \n",
        "  with open(new,'r',encoding='utf8') as rn:\n",
        "    p = rn.read()\n",
        "    p = p.replace(ptoken,\"\\033[1m\")\n",
        "    p = p.replace(gtoken,\"\\033[0m\")\n",
        "    print(f\"\\nPrompt (in bold) and generation below.\\n{p}\")\n",
        "    \n"
      ],
      "execution_count": 85,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "completion done in 53.36s\n",
            "Wrote sample 1!\n",
            "Wrote 1 sample to ABriefHistoryNLP.txt in 53.37 seconds.\n",
            "\n",
            "Prompt (in bold) and generation below.\n",
            "\n",
            "\u001b[1mA Brief History of Natural Language Processing with Machine Learning:\n",
            "\u001b[0m\n",
            "The Beginning of Natural Language Processing with the Penn Treebank Corpus\n",
            "\n",
            "The Penn Treebank Corpus is an English treebank. A treebank is a text corpus where the words in the text are arranged as a tree with words being the nodes and the words connected to each other being the edges. The labels on the edges tell us the type of the words connected to each other. For example, the word “the” is connected to the word “a”, which indicates that they are the same type of word. The word “a” is connected to the word “lot” which indicates that they are both nouns. The treebank is arranged in such a way that we can get a better idea about the context of each word.\n",
            "\n",
            "One of the most important tasks of a treebank is to analyze the sentence. The goal of the treebank is to parse the sentences into syntactic trees, which are either noun phrase, verb phrase or prepositional phrases. The syntactic trees allow us to understand the relationship between the words and the sentences, and it allows us to analyze the meaning of the sentences.\n",
            "\n",
            "The Penn Treebank Corpus is a very well-known treebank and is used extensively in all the NLP projects. You can find more information about the Penn Treebank Corpus at this URL, and the purpose of the treebank is described on this URL. The treebank is the biggest treebank, with a size of about 900,000 words.\n",
            "\n",
            "There are some interesting facts about the treebank that you may be interested in.\n",
            "\n",
            "The Penn Treebank Corpus was created in 1990, when the first few versions of the treebank were released. It was the largest corpus for English language at that time. In the current day, the treebank is still being actively developed and has a size of about 1.2 billion words.\n",
            "\n",
            "The Penn Treebank Corpus can be used to train a language model, which is a kind of machine learning model that is trained on the treebank. There are two types of language models:\n",
            "\n",
            "Hidden Markov Models (HMMs)\n",
            "\n",
            "Markov Models (MMs)\n",
            "\n",
            "In a hidden Markov model, the probability of a word is calculated based on the probability of the previous word and the probability of the current word. For example, if the current word is “a” and the previous word is “the”, then the probability of the word “a” is 0.1, because “the” is the most probable word and “a” is the next most probable word. In Markov models, the probability of a word is calculated based on the probability of the current word and the probability of the previous words. For example, if the current word is “the” and the previous word is “a”, then the probability of the word “a” is 0.1, because “the” is the most probable word and “a” is the next most probable word.\n",
            "\n",
            "If you are interested in building a language model, you can find more information about the Penn Treebank Corpus in this URL and about HMMs and MMs in this URL.\n",
            "\n",
            "Why Natural Language Processing?\n",
            "\n",
            "Natural Language Processing (NLP) is the process of analyzing natural language. The goal of NLP is to automate the process of analyzing natural language. In the process of analyzing natural language, you can perform tasks like text summarization, sentiment analysis, summarization, translation, etc. Many people have started doing NLP projects, because it is a very popular field, and many people are working on NLP related projects.\n",
            "\n",
            "I am not an expert in NLP, but I want to make this blog post about a very interesting NLP project that I want to share with you.\n",
            "\n",
            "Natural Language Processing with a Machine Learning Model:\n",
            "\n",
            "Let’s start with a very simple example, where I ask you to type a sentence, and I then tell you whether the sentence is positive, negative or neutral. For example, I type the sentence “apple is better than banana” and then ask you to give me a feedback whether the sentence is positive, negative or neutral.\n",
            "\n",
            "Okay, let’s dive right in.\n",
            "\n",
            "Step 1. Training the Model:\n",
            "\n",
            "The first step to building a machine learning model is to train the model. There are two ways to train the model.\n",
            "\n",
            "You can use the Penn Treebank Corpus, which is a very big treebank, and it is a well-known corpus. You can find more information about the Penn Treebank Corpus in this URL.\n",
            "\n",
            "You can use the Corpus available on the Internet. There are lots of corpora available on the Internet, which you can use to train the model.\n",
            "\n",
            "For the purpose of this tutorial, I am using the small treebank available on the Internet, and I am using the corpus available on the Internet, and I am training the model on the training data.\n",
            "\n",
            "I will share the code for training the model, so you can train your own model and see how the model works.\n",
            "\n",
            "In this tutorial, I am going to give you an example of training the model. For the purpose of this tutorial, I am using the small treebank available on the Internet. You can get the source code for the treebank in this URL.\n",
            "\n",
            "Before we start training the model, let’s create a new directory for the code, and let’s create a file called train.py and save it inside the directory.\n",
            "\n",
            "This is the code for the train.py file, which will be responsible for training the model.\n",
            "\n",
            "I am using the treebank created by the Penn Treebank Corpus, which has a size of about 1.2 billion words. In this code, I am using the last 50,000 words from the treebank. These are the 50,000 most frequent words, which are the words that occur more frequently in the treebank.\n",
            "\n",
            "As you can see, this is a simple code for training the model. I am using the Python libraries and the Scikit-Learn library. I am using the KNeighborsClassifier model.\n",
            "\n",
            "Okay, let’s start running the code.\n",
            "\n",
            "If you want to know about the different types of libraries, you can check out the Scikit-Learn library website. I am using scikit-learn 0.20.1, which is the latest version of scikit-learn. If you are interested in installing scikit-learn, you can check out this URL, and if you are interested in installing scikit-learn 0.20.1, you can check out this URL.\n",
            "\n",
            "The first line of the code above is importing the libraries that we are going to use.\n",
            "\n",
            "Then, we are creating a list with the train and test data. The train data are the 50,000 most frequent words, and the test data are the remaining words. The number of rows in the train data are equal to the number of words in the train data, and the number of rows in the test data are equal to the number of words in the test data.\n",
            "\n",
            "Then, we are initializing the model with the train data.\n",
            "\n",
            "We are creating the model, and the KNeighborsClassifier model is used for this example. The KNeighborsClassifier model is one of the most popular models in Scikit-Learn. You can find more information about the KNeighborsClassifier model in the documentation of the KNeighborsClassifier model.\n",
            "\n",
            "Then, we are creating the model and the dataframe from the data in the train data.\n",
            "\n",
            "This code uses the TfidfVectorizer method to create a tf-idf matrix from the train data. The word frequency is calculated from the tf-idf matrix.\n",
            "\n",
            "Then, we are removing the stop words from the data. The stop words are the words that occur less than 3 times in the dataset.\n",
            "\n",
            "The next two lines of the code above are removing the words that have more than 2 letters, because they are too long.\n",
            "\n",
            "Then, we are calculating the length of the words, and we are only keeping the words that have a length of 2 to 15.\n",
            "\n",
            "After that, we are converting the words into lower case, so that we can use the lower case words as features.\n",
            "\n",
            "We are now converting the words into a vector of the length of 2000, because we will be using the first 2000 words as features.\n",
            "\n",
            "And then, we are creating a new dataframe, which is the training data, from the dataframe from the train data.\n",
            "\n",
            "The next two lines of the code above are creating the label for the train data. The labels for the train data are the labels for the test data. This is the label column of the training data, and this is the label column of the test data. The labels for the test data are the labels for the train data, and the labels for the train data are the labels for the test data.\n",
            "\n",
            "We are initializing a list with the train data.\n",
            "\n",
            "Then, we are extracting the features from the train data and the labels from the train data.\n",
            "\n",
            "The next few lines of the code above are calculating the mean and the variance of the features, and then we are calculating the variances from the test data and the training data.\n",
            "\n",
            "The next few lines of the code above are removing the words that are not in the list and creating the new dataframe with the words that are not in the list.\n",
            "\n",
            "And then, we are converting the features into a matrix, and we are also converting the labels into a matrix. The labels are 1-dimensional and are the same type as the features, which\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
 }