Skip to content

Instantly share code, notes, and snippets.

@buttercutter

buttercutter/GPT_J_6B_Inference_Demo.ipynb

Created April 7, 2023 17:03
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save buttercutter/ee4e6e2eddb658949d87f9fdd437b01b to your computer and use it in GitHub Desktop.
Save buttercutter/ee4e6e2eddb658949d87f9fdd437b01b to your computer and use it in GitHub Desktop.
Download ZIP
Display the source blob
Display the rendered blob
Raw
GPT_J_6B_Inference_Demo.ipynb
{
"cells": [
{
"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",
"execution_count": 13,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 1000
},
"id": "n7xAFw-LOYfe",
"outputId": "7bd7fd83-a40c-41a2-8a14-9a11fc2f150c"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Reading package lists... Done\n",
"Building dependency tree \n",
"Reading state information... Done\n",
"zstd is already the newest version (1.4.4+dfsg-3ubuntu0.1).\n",
"The following packages were automatically installed and are no longer required:\n",
" libatasmart4 libblockdev-fs2 libblockdev-loop2 libblockdev-part-err2\n",
" libblockdev-part2 libblockdev-swap2 libblockdev-utils2 libblockdev2\n",
" libmm-glib0 libnspr4 libnss3 libnuma1 libparted-fs-resize0 libudisks2-0\n",
" python-pip-whl python3-wheel usb-modeswitch usb-modeswitch-data\n",
"Use 'sudo apt autoremove' to remove them.\n",
"0 upgraded, 0 newly installed, 0 to remove and 51 not upgraded.\n",
"--2023-04-07 16:42:45-- 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.6\n",
"Connecting to the-eye.eu (the-eye.eu)|162.213.130.6|:443... connected.\n",
"HTTP request sent, awaiting response... 416 Requested Range Not Satisfiable\n",
"\n",
" The file is already fully retrieved; nothing to do.\n",
"\n",
"\n",
"real\t0m0.433s\n",
"user\t0m0.008s\n",
"sys\t0m0.000s\n",
"\n",
"real\t1m43.422s\n",
"user\t0m15.993s\n",
"sys\t0m25.143s\n",
"fatal: destination path 'mesh-transformer-jax' already exists and is not an empty directory.\n",
"Defaulting to user installation because normal site-packages is not writeable\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-xvu_1u9j\n",
" Running command git clone --filter=blob:none --quiet https://github.com/EleutherAI/lm-evaluation-harness/ /tmp/pip-req-build-xvu_1u9j\n",
" Resolved https://github.com/EleutherAI/lm-evaluation-harness/ to commit 50eb80ba5f23bde87cb80bbe1f1633bb98038e58\n",
" Preparing metadata (setup.py) ... \u001b[?25ldone\n",
"\u001b[?25hCollecting numpy~=1.19.5\n",
" Using cached numpy-1.19.5-cp38-cp38-manylinux2010_x86_64.whl (14.9 MB)\n",
"Requirement already satisfied: tqdm>=4.45.0 in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 2)) (4.64.1)\n",
"Requirement already satisfied: wandb>=0.11.2 in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 3)) (0.14.1)\n",
"Requirement already satisfied: einops~=0.3.0 in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 4)) (0.3.2)\n",
"Requirement already satisfied: requests~=2.25.1 in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 5)) (2.25.1)\n",
"Requirement already satisfied: fabric~=2.6.0 in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 6)) (2.6.0)\n",
"Collecting optax==0.0.9\n",
" Using cached optax-0.0.9-py3-none-any.whl (118 kB)\n",
"Requirement already satisfied: dm-haiku==0.0.5 in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 8)) (0.0.5)\n",
"Requirement already satisfied: ray[default]==1.4.1 in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 10)) (1.4.1)\n",
"Collecting jax~=0.2.12\n",
" Using cached jax-0.2.28-py3-none-any.whl\n",
"Requirement already satisfied: Flask~=1.1.2 in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 12)) (1.1.4)\n",
"Requirement already satisfied: cloudpickle~=1.3.0 in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 13)) (1.3.0)\n",
"Requirement already satisfied: tensorflow-cpu~=2.6.0 in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 14)) (2.6.5)\n",
"Requirement already satisfied: google-cloud-storage~=1.36.2 in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 15)) (1.36.2)\n",
"Requirement already satisfied: transformers in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 16)) (4.26.1)\n",
"Requirement already satisfied: smart_open[gcs] in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 17)) (6.2.0)\n",
"Requirement already satisfied: func_timeout in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 18)) (4.3.5)\n",
"Requirement already satisfied: ftfy in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 19)) (6.1.1)\n",
"Requirement already satisfied: fastapi in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 20)) (0.95.0)\n",
"Requirement already satisfied: uvicorn in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 21)) (0.21.1)\n",
"Requirement already satisfied: lm_dataformat in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 22)) (0.0.20)\n",
"Requirement already satisfied: pathy in /home/moe/.local/lib/python3.8/site-packages (from -r mesh-transformer-jax/requirements.txt (line 23)) (0.10.1)\n",
"Requirement already satisfied: absl-py>=0.7.1 in /home/moe/.local/lib/python3.8/site-packages (from optax==0.0.9->-r mesh-transformer-jax/requirements.txt (line 7)) (0.15.0)\n",
"Requirement already satisfied: chex>=0.0.4 in /home/moe/.local/lib/python3.8/site-packages (from optax==0.0.9->-r mesh-transformer-jax/requirements.txt (line 7)) (0.1.7)\n",
"Requirement already satisfied: jaxlib>=0.1.37 in /home/moe/.local/lib/python3.8/site-packages (from optax==0.0.9->-r mesh-transformer-jax/requirements.txt (line 7)) (0.1.76)\n",
"Requirement already satisfied: tabulate>=0.8.9 in /home/moe/.local/lib/python3.8/site-packages (from dm-haiku==0.0.5->-r mesh-transformer-jax/requirements.txt (line 8)) (0.9.0)\n",
"Requirement already satisfied: jmp>=0.0.2 in /home/moe/.local/lib/python3.8/site-packages (from dm-haiku==0.0.5->-r mesh-transformer-jax/requirements.txt (line 8)) (0.0.4)\n",
"Requirement already satisfied: aioredis in /home/moe/.local/lib/python3.8/site-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (2.0.1)\n",
"Requirement already satisfied: msgpack<2.0.0,>=1.0.0 in /home/moe/.local/lib/python3.8/site-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (1.0.4)\n",
"Requirement already satisfied: grpcio>=1.28.1 in /home/moe/.local/lib/python3.8/site-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (1.53.0)\n",
"Requirement already satisfied: colorama in /usr/lib/python3/dist-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (0.4.3)\n",
"Requirement already satisfied: protobuf>=3.15.3 in /home/moe/.local/lib/python3.8/site-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (3.19.6)\n",
"Requirement already satisfied: aiohttp-cors in /home/moe/.local/lib/python3.8/site-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (0.7.0)\n",
"Requirement already satisfied: gpustat in /home/moe/.local/lib/python3.8/site-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (1.1)\n",
"Requirement already satisfied: aiohttp in /home/moe/.local/lib/python3.8/site-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (3.8.3)\n",
"Requirement already satisfied: click>=7.0 in /home/moe/.local/lib/python3.8/site-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (7.1.2)\n",
"Requirement already satisfied: opencensus in /home/moe/.local/lib/python3.8/site-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (0.11.2)\n",
"Requirement already satisfied: filelock in /usr/local/lib/python3.8/dist-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (3.7.1)\n",
"Requirement already satisfied: pydantic>=1.8 in /home/moe/.local/lib/python3.8/site-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (1.9.2)\n",
"Requirement already satisfied: redis>=3.5.0 in /home/moe/.local/lib/python3.8/site-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (4.5.4)\n",
"Requirement already satisfied: jsonschema in /usr/lib/python3/dist-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (3.2.0)\n",
"Requirement already satisfied: py-spy>=0.2.0 in /home/moe/.local/lib/python3.8/site-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (0.3.14)\n",
"Requirement already satisfied: pyyaml in /usr/local/lib/python3.8/dist-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (5.4.1)\n",
"Requirement already satisfied: prometheus-client>=0.7.1 in /home/moe/.local/lib/python3.8/site-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (0.16.0)\n",
"Requirement already satisfied: colorful in /home/moe/.local/lib/python3.8/site-packages (from ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (0.5.5)\n",
"Requirement already satisfied: psutil>=5.0.0 in /home/moe/.local/lib/python3.8/site-packages (from wandb>=0.11.2->-r mesh-transformer-jax/requirements.txt (line 3)) (5.9.4)\n",
"Requirement already satisfied: setproctitle in /home/moe/.local/lib/python3.8/site-packages (from wandb>=0.11.2->-r mesh-transformer-jax/requirements.txt (line 3)) (1.3.2)\n",
"Requirement already satisfied: typing-extensions in /home/moe/.local/lib/python3.8/site-packages (from wandb>=0.11.2->-r mesh-transformer-jax/requirements.txt (line 3)) (4.5.0)\n",
"Requirement already satisfied: sentry-sdk>=1.0.0 in /home/moe/.local/lib/python3.8/site-packages (from wandb>=0.11.2->-r mesh-transformer-jax/requirements.txt (line 3)) (1.19.1)\n",
"Requirement already satisfied: GitPython!=3.1.29,>=1.0.0 in /home/moe/.local/lib/python3.8/site-packages (from wandb>=0.11.2->-r mesh-transformer-jax/requirements.txt (line 3)) (3.1.31)\n",
"Requirement already satisfied: docker-pycreds>=0.4.0 in /home/moe/.local/lib/python3.8/site-packages (from wandb>=0.11.2->-r mesh-transformer-jax/requirements.txt (line 3)) (0.4.0)\n",
"Requirement already satisfied: setuptools in /usr/local/lib/python3.8/dist-packages (from wandb>=0.11.2->-r mesh-transformer-jax/requirements.txt (line 3)) (62.3.2)\n",
"Requirement already satisfied: appdirs>=1.4.3 in /home/moe/.local/lib/python3.8/site-packages (from wandb>=0.11.2->-r mesh-transformer-jax/requirements.txt (line 3)) (1.4.4)\n",
"Requirement already satisfied: pathtools in /home/moe/.local/lib/python3.8/site-packages (from wandb>=0.11.2->-r mesh-transformer-jax/requirements.txt (line 3)) (0.1.2)\n",
"Requirement already satisfied: idna<3,>=2.5 in /usr/lib/python3/dist-packages (from requests~=2.25.1->-r mesh-transformer-jax/requirements.txt (line 5)) (2.8)\n",
"Requirement already satisfied: certifi>=2017.4.17 in /usr/lib/python3/dist-packages (from requests~=2.25.1->-r mesh-transformer-jax/requirements.txt (line 5)) (2019.11.28)\n",
"Requirement already satisfied: chardet<5,>=3.0.2 in /usr/lib/python3/dist-packages (from requests~=2.25.1->-r mesh-transformer-jax/requirements.txt (line 5)) (3.0.4)\n",
"Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/moe/.local/lib/python3.8/site-packages (from requests~=2.25.1->-r mesh-transformer-jax/requirements.txt (line 5)) (1.26.15)\n",
"Requirement already satisfied: paramiko>=2.4 in /home/moe/.local/lib/python3.8/site-packages (from fabric~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 6)) (3.1.0)\n",
"Requirement already satisfied: invoke<2.0,>=1.3 in /home/moe/.local/lib/python3.8/site-packages (from fabric~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 6)) (1.7.3)\n",
"Requirement already satisfied: pathlib2 in /home/moe/.local/lib/python3.8/site-packages (from fabric~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 6)) (2.3.7.post1)\n",
"Requirement already satisfied: datasets>=2.0.0 in /home/moe/.local/lib/python3.8/site-packages (from lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (2.8.0)\n",
"Requirement already satisfied: jsonlines in /home/moe/.local/lib/python3.8/site-packages (from lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (3.1.0)\n",
"Requirement already satisfied: numexpr in /home/moe/.local/lib/python3.8/site-packages (from lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (2.8.4)\n",
"Requirement already satisfied: openai>=0.6.4 in /home/moe/.local/lib/python3.8/site-packages (from lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (0.27.4)\n",
"Requirement already satisfied: omegaconf>=2.2 in /home/moe/.local/lib/python3.8/site-packages (from lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (2.3.0)\n",
"Requirement already satisfied: pybind11>=2.6.2 in /home/moe/.local/lib/python3.8/site-packages (from lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (2.10.4)\n",
"Requirement already satisfied: pycountry in /home/moe/.local/lib/python3.8/site-packages (from lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (22.3.5)\n",
"Requirement already satisfied: pytablewriter in /home/moe/.local/lib/python3.8/site-packages (from lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (0.64.2)\n",
"Requirement already satisfied: rouge-score>=0.0.4 in /home/moe/.local/lib/python3.8/site-packages (from lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (0.1.2)\n",
"Requirement already satisfied: sacrebleu==1.5.0 in /home/moe/.local/lib/python3.8/site-packages (from lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (1.5.0)\n",
"Requirement already satisfied: scikit-learn>=0.24.1 in /home/moe/.local/lib/python3.8/site-packages (from lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (1.2.2)\n",
"Requirement already satisfied: sqlitedict in /home/moe/.local/lib/python3.8/site-packages (from lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (2.1.0)\n",
"Requirement already satisfied: torch>=1.7 in /home/moe/.local/lib/python3.8/site-packages (from lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (1.13.1)\n",
"Requirement already satisfied: tqdm-multiprocess in /home/moe/.local/lib/python3.8/site-packages (from lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (0.0.11)\n",
"Requirement already satisfied: zstandard in /home/moe/.local/lib/python3.8/site-packages (from lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (0.20.0)\n",
"Requirement already satisfied: portalocker in /home/moe/.local/lib/python3.8/site-packages (from sacrebleu==1.5.0->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (2.7.0)\n",
"Requirement already satisfied: scipy>=1.2.1 in /home/moe/.local/lib/python3.8/site-packages (from jax~=0.2.12->-r mesh-transformer-jax/requirements.txt (line 11)) (1.10.0)\n",
"Requirement already satisfied: opt-einsum in /home/moe/.local/lib/python3.8/site-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/lib/python3/dist-packages (from Flask~=1.1.2->-r mesh-transformer-jax/requirements.txt (line 12)) (2.10.1)\n",
"Requirement already satisfied: itsdangerous<2.0,>=0.24 in /home/moe/.local/lib/python3.8/site-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 /home/moe/.local/lib/python3.8/site-packages (from Flask~=1.1.2->-r mesh-transformer-jax/requirements.txt (line 12)) (1.0.1)\n",
"Requirement already satisfied: wheel~=0.35 in /home/moe/.local/lib/python3.8/site-packages (from tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (0.40.0)\n",
"Requirement already satisfied: wrapt~=1.12.1 in /home/moe/.local/lib/python3.8/site-packages (from tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (1.12.1)\n",
"Requirement already satisfied: termcolor~=1.1.0 in /home/moe/.local/lib/python3.8/site-packages (from tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (1.1.0)\n",
"Requirement already satisfied: clang~=5.0 in /home/moe/.local/lib/python3.8/site-packages (from tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (5.0)\n",
"Requirement already satisfied: six~=1.15.0 in /home/moe/.local/lib/python3.8/site-packages (from tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (1.15.0)\n",
"Requirement already satisfied: h5py~=3.1.0 in /home/moe/.local/lib/python3.8/site-packages (from tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (3.1.0)\n",
"Requirement already satisfied: tensorflow-estimator<2.7,>=2.6.0 in /home/moe/.local/lib/python3.8/site-packages (from tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (2.6.0)\n",
"Requirement already satisfied: astunparse~=1.6.3 in /home/moe/.local/lib/python3.8/site-packages (from tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (1.6.3)\n",
"Requirement already satisfied: gast==0.4.0 in /home/moe/.local/lib/python3.8/site-packages (from tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (0.4.0)\n",
"Collecting typing-extensions\n",
" Using cached typing_extensions-3.10.0.2-py3-none-any.whl (26 kB)\n",
"Requirement already satisfied: tensorboard<2.7,>=2.6.0 in /home/moe/.local/lib/python3.8/site-packages (from tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (2.6.0)\n",
"Requirement already satisfied: google-pasta~=0.2 in /home/moe/.local/lib/python3.8/site-packages (from tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (0.2.0)\n",
"Requirement already satisfied: keras-preprocessing~=1.1.2 in /home/moe/.local/lib/python3.8/site-packages (from tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (1.1.2)\n",
"Requirement already satisfied: flatbuffers~=1.12.0 in /home/moe/.local/lib/python3.8/site-packages (from tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (1.12)\n",
"Requirement already satisfied: keras<2.7,>=2.6.0 in /home/moe/.local/lib/python3.8/site-packages (from tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (2.6.0)\n",
"Requirement already satisfied: google-cloud-core<2.0dev,>=1.4.1 in /home/moe/.local/lib/python3.8/site-packages (from google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 15)) (1.7.3)\n",
"Requirement already satisfied: google-auth<2.0dev,>=1.11.0 in /home/moe/.local/lib/python3.8/site-packages (from google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 15)) (1.35.0)\n",
"Requirement already satisfied: google-resumable-media<2.0dev,>=1.2.0 in /home/moe/.local/lib/python3.8/site-packages (from google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 15)) (1.3.3)\n",
"Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /home/moe/.local/lib/python3.8/site-packages (from transformers->-r mesh-transformer-jax/requirements.txt (line 16)) (0.13.2)\n",
"Requirement already satisfied: packaging>=20.0 in /home/moe/.local/lib/python3.8/site-packages (from transformers->-r mesh-transformer-jax/requirements.txt (line 16)) (23.0)\n",
"Requirement already satisfied: huggingface-hub<1.0,>=0.11.0 in /home/moe/.local/lib/python3.8/site-packages (from transformers->-r mesh-transformer-jax/requirements.txt (line 16)) (0.11.1)\n",
"Requirement already satisfied: regex!=2019.12.17 in /home/moe/.local/lib/python3.8/site-packages (from transformers->-r mesh-transformer-jax/requirements.txt (line 16)) (2022.10.31)\n",
"Requirement already satisfied: wcwidth>=0.2.5 in /home/moe/.local/lib/python3.8/site-packages (from ftfy->-r mesh-transformer-jax/requirements.txt (line 19)) (0.2.5)\n",
"Requirement already satisfied: starlette<0.27.0,>=0.26.1 in /home/moe/.local/lib/python3.8/site-packages (from fastapi->-r mesh-transformer-jax/requirements.txt (line 20)) (0.26.1)\n",
"Requirement already satisfied: h11>=0.8 in /home/moe/.local/lib/python3.8/site-packages (from uvicorn->-r mesh-transformer-jax/requirements.txt (line 21)) (0.14.0)\n",
"Requirement already satisfied: ujson in /home/moe/.local/lib/python3.8/site-packages (from lm_dataformat->-r mesh-transformer-jax/requirements.txt (line 22)) (5.7.0)\n",
"Requirement already satisfied: typer<1.0.0,>=0.3.0 in /home/moe/.local/lib/python3.8/site-packages (from pathy->-r mesh-transformer-jax/requirements.txt (line 23)) (0.7.0)\n",
"Collecting chex>=0.0.4\n",
" Using cached chex-0.1.6-py3-none-any.whl (87 kB)\n",
" Using cached chex-0.1.5-py3-none-any.whl (85 kB)\n",
"Requirement already satisfied: toolz>=0.9.0 in /home/moe/.local/lib/python3.8/site-packages (from chex>=0.0.4->optax==0.0.9->-r mesh-transformer-jax/requirements.txt (line 7)) (0.12.0)\n",
"Requirement already satisfied: dm-tree>=0.1.5 in /home/moe/.local/lib/python3.8/site-packages (from chex>=0.0.4->optax==0.0.9->-r mesh-transformer-jax/requirements.txt (line 7)) (0.1.8)\n",
"Requirement already satisfied: xxhash in /home/moe/.local/lib/python3.8/site-packages (from datasets>=2.0.0->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (3.2.0)\n",
"Requirement already satisfied: pyarrow>=6.0.0 in /home/moe/.local/lib/python3.8/site-packages (from datasets>=2.0.0->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (10.0.1)\n",
"Requirement already satisfied: multiprocess in /home/moe/.local/lib/python3.8/site-packages (from datasets>=2.0.0->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (0.70.14)\n",
"Requirement already satisfied: pandas in /home/moe/.local/lib/python3.8/site-packages (from datasets>=2.0.0->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (1.4.4)\n",
"Requirement already satisfied: responses<0.19 in /home/moe/.local/lib/python3.8/site-packages (from datasets>=2.0.0->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (0.18.0)\n",
"Requirement already satisfied: fsspec[http]>=2021.11.1 in /home/moe/.local/lib/python3.8/site-packages (from datasets>=2.0.0->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (2022.11.0)\n",
"Requirement already satisfied: dill<0.3.7 in /home/moe/.local/lib/python3.8/site-packages (from datasets>=2.0.0->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (0.3.6)\n",
"Requirement already satisfied: gitdb<5,>=4.0.1 in /home/moe/.local/lib/python3.8/site-packages (from GitPython!=3.1.29,>=1.0.0->wandb>=0.11.2->-r mesh-transformer-jax/requirements.txt (line 3)) (4.0.10)\n",
"Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/lib/python3/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.1)\n",
"Requirement already satisfied: cachetools<5.0,>=2.0.0 in /home/moe/.local/lib/python3.8/site-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: rsa<5,>=3.1.4 in /home/moe/.local/lib/python3.8/site-packages (from google-auth<2.0dev,>=1.11.0->google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 15)) (4.9)\n",
"Requirement already satisfied: google-api-core<3.0.0dev,>=1.21.0 in /home/moe/.local/lib/python3.8/site-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.34.0)\n",
"Requirement already satisfied: google-crc32c<2.0dev,>=1.0 in /home/moe/.local/lib/python3.8/site-packages (from google-resumable-media<2.0dev,>=1.2.0->google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 15)) (1.5.0)\n",
"Requirement already satisfied: antlr4-python3-runtime==4.9.* in /home/moe/.local/lib/python3.8/site-packages (from omegaconf>=2.2->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (4.9.3)\n",
"Requirement already satisfied: bcrypt>=3.2 in /home/moe/.local/lib/python3.8/site-packages (from paramiko>=2.4->fabric~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 6)) (4.0.1)\n",
"Requirement already satisfied: cryptography>=3.3 in /home/moe/.local/lib/python3.8/site-packages (from paramiko>=2.4->fabric~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 6)) (40.0.1)\n",
"Requirement already satisfied: pynacl>=1.5 in /home/moe/.local/lib/python3.8/site-packages (from paramiko>=2.4->fabric~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 6)) (1.5.0)\n",
"Requirement already satisfied: async-timeout>=4.0.2 in /home/moe/.local/lib/python3.8/site-packages (from redis>=3.5.0->ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (4.0.2)\n",
"Requirement already satisfied: nltk in /home/moe/.local/lib/python3.8/site-packages (from rouge-score>=0.0.4->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (3.8.1)\n",
"Requirement already satisfied: joblib>=1.1.1 in /home/moe/.local/lib/python3.8/site-packages (from scikit-learn>=0.24.1->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (1.2.0)\n",
"Requirement already satisfied: threadpoolctl>=2.0.0 in /home/moe/.local/lib/python3.8/site-packages (from scikit-learn>=0.24.1->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (3.1.0)\n",
"Requirement already satisfied: anyio<5,>=3.4.0 in /home/moe/.local/lib/python3.8/site-packages (from starlette<0.27.0,>=0.26.1->fastapi->-r mesh-transformer-jax/requirements.txt (line 20)) (3.6.2)\n",
"Requirement already satisfied: markdown>=2.6.8 in /home/moe/.local/lib/python3.8/site-packages (from tensorboard<2.7,>=2.6.0->tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (3.4.1)\n",
"Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /home/moe/.local/lib/python3.8/site-packages (from tensorboard<2.7,>=2.6.0->tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (0.4.6)\n",
"Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /home/moe/.local/lib/python3.8/site-packages (from tensorboard<2.7,>=2.6.0->tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (1.8.1)\n",
"Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /home/moe/.local/lib/python3.8/site-packages (from tensorboard<2.7,>=2.6.0->tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (0.6.1)\n",
"Requirement already satisfied: nvidia-cudnn-cu11==8.5.0.96 in /home/moe/.local/lib/python3.8/site-packages (from torch>=1.7->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (8.5.0.96)\n",
"Requirement already satisfied: nvidia-cuda-runtime-cu11==11.7.99 in /home/moe/.local/lib/python3.8/site-packages (from torch>=1.7->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (11.7.99)\n",
"Requirement already satisfied: nvidia-cublas-cu11==11.10.3.66 in /home/moe/.local/lib/python3.8/site-packages (from torch>=1.7->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (11.10.3.66)\n",
"Requirement already satisfied: nvidia-cuda-nvrtc-cu11==11.7.99 in /home/moe/.local/lib/python3.8/site-packages (from torch>=1.7->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (11.7.99)\n",
"Requirement already satisfied: yarl<2.0,>=1.0 in /home/moe/.local/lib/python3.8/site-packages (from aiohttp->ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (1.8.2)\n",
"Requirement already satisfied: attrs>=17.3.0 in /usr/lib/python3/dist-packages (from aiohttp->ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (19.3.0)\n",
"Requirement already satisfied: multidict<7.0,>=4.5 in /home/moe/.local/lib/python3.8/site-packages (from aiohttp->ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (6.0.4)\n",
"Requirement already satisfied: charset-normalizer<3.0,>=2.0 in /usr/local/lib/python3.8/dist-packages (from aiohttp->ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (2.0.12)\n",
"Requirement already satisfied: frozenlist>=1.1.1 in /home/moe/.local/lib/python3.8/site-packages (from aiohttp->ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (1.3.3)\n",
"Requirement already satisfied: aiosignal>=1.1.2 in /home/moe/.local/lib/python3.8/site-packages (from aiohttp->ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (1.3.1)\n",
"Requirement already satisfied: blessed>=1.17.1 in /home/moe/.local/lib/python3.8/site-packages (from gpustat->ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (1.20.0)\n",
"Requirement already satisfied: nvidia-ml-py>=11.450.129 in /home/moe/.local/lib/python3.8/site-packages (from gpustat->ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (11.525.112)\n",
"Requirement already satisfied: opencensus-context>=0.1.3 in /home/moe/.local/lib/python3.8/site-packages (from opencensus->ray[default]==1.4.1->-r mesh-transformer-jax/requirements.txt (line 10)) (0.1.3)\n",
"Requirement already satisfied: mbstrdecoder<2,>=1.0.0 in /home/moe/.local/lib/python3.8/site-packages (from pytablewriter->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (1.1.2)\n",
"Requirement already satisfied: DataProperty<2,>=0.55.0 in /home/moe/.local/lib/python3.8/site-packages (from pytablewriter->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (0.55.0)\n",
"Requirement already satisfied: typepy[datetime]<2,>=1.2.0 in /home/moe/.local/lib/python3.8/site-packages (from pytablewriter->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (1.3.0)\n",
"Requirement already satisfied: tabledata<2,>=1.3.0 in /home/moe/.local/lib/python3.8/site-packages (from pytablewriter->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (1.3.1)\n",
"Requirement already satisfied: pathvalidate<3,>=2.3.0 in /home/moe/.local/lib/python3.8/site-packages (from pytablewriter->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (2.5.2)\n",
"Requirement already satisfied: tcolorpy<1,>=0.0.5 in /home/moe/.local/lib/python3.8/site-packages (from pytablewriter->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (0.1.2)\n",
"Requirement already satisfied: sniffio>=1.1 in /home/moe/.local/lib/python3.8/site-packages (from anyio<5,>=3.4.0->starlette<0.27.0,>=0.26.1->fastapi->-r mesh-transformer-jax/requirements.txt (line 20)) (1.3.0)\n",
"Requirement already satisfied: cffi>=1.12 in /home/moe/.local/lib/python3.8/site-packages (from cryptography>=3.3->paramiko>=2.4->fabric~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 6)) (1.15.1)\n",
"Requirement already satisfied: smmap<6,>=3.0.1 in /home/moe/.local/lib/python3.8/site-packages (from gitdb<5,>=4.0.1->GitPython!=3.1.29,>=1.0.0->wandb>=0.11.2->-r mesh-transformer-jax/requirements.txt (line 3)) (5.0.0)\n",
"Requirement already satisfied: googleapis-common-protos<2.0dev,>=1.56.2 in /home/moe/.local/lib/python3.8/site-packages (from google-api-core<3.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.58.0)\n",
"Requirement already satisfied: requests-oauthlib>=0.7.0 in /home/moe/.local/lib/python3.8/site-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.7,>=2.6.0->tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (1.3.1)\n",
"Requirement already satisfied: importlib-metadata>=4.4 in /home/moe/.local/lib/python3.8/site-packages (from markdown>=2.6.8->tensorboard<2.7,>=2.6.0->tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (6.0.0)\n",
"Requirement already satisfied: pyasn1>=0.1.3 in /usr/lib/python3/dist-packages (from rsa<5,>=3.1.4->google-auth<2.0dev,>=1.11.0->google-cloud-storage~=1.36.2->-r mesh-transformer-jax/requirements.txt (line 15)) (0.4.2)\n",
"Requirement already satisfied: python-dateutil<3.0.0,>=2.8.0 in /home/moe/.local/lib/python3.8/site-packages (from typepy[datetime]<2,>=1.2.0->pytablewriter->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (2.8.2)\n",
"Requirement already satisfied: pytz>=2018.9 in /home/moe/.local/lib/python3.8/site-packages (from typepy[datetime]<2,>=1.2.0->pytablewriter->lm-eval==0.3.0->-r mesh-transformer-jax/requirements.txt (line 9)) (2022.7)\n",
"Requirement already satisfied: pycparser in /home/moe/.local/lib/python3.8/site-packages (from cffi>=1.12->cryptography>=3.3->paramiko>=2.4->fabric~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 6)) (2.21)\n",
"Requirement already satisfied: zipp>=0.5 in /home/moe/.local/lib/python3.8/site-packages (from importlib-metadata>=4.4->markdown>=2.6.8->tensorboard<2.7,>=2.6.0->tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (3.11.0)\n",
"Requirement already satisfied: oauthlib>=3.0.0 in /usr/lib/python3/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.7,>=2.6.0->tensorflow-cpu~=2.6.0->-r mesh-transformer-jax/requirements.txt (line 14)) (3.1.0)\n",
"Installing collected packages: typing-extensions, numpy, jax, chex, optax\n",
" Attempting uninstall: typing-extensions\n",
" Found existing installation: typing_extensions 4.5.0\n",
" Uninstalling typing_extensions-4.5.0:\n",
" Successfully uninstalled typing_extensions-4.5.0\n",
" Attempting uninstall: numpy\n",
" Found existing installation: numpy 1.24.2\n",
" Uninstalling numpy-1.24.2:\n",
" Successfully uninstalled numpy-1.24.2\n",
" Attempting uninstall: jax\n",
" Found existing installation: jax 0.4.8\n",
" Uninstalling jax-0.4.8:\n",
" Successfully uninstalled jax-0.4.8\n",
" Attempting uninstall: chex\n",
" Found existing installation: chex 0.1.7\n",
" Uninstalling chex-0.1.7:\n",
" Successfully uninstalled chex-0.1.7\n",
" Attempting uninstall: optax\n",
" Found existing installation: optax 0.1.5.dev0\n",
" Uninstalling optax-0.1.5.dev0:\n",
" Successfully uninstalled optax-0.1.5.dev0\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",
"tensorflow 2.5.0 requires grpcio~=1.34.0, but you have grpcio 1.53.0 which is incompatible.\n",
"tensorflow 2.5.0 requires tensorflow-estimator<2.6.0,>=2.5.0rc0, but you have tensorflow-estimator 2.6.0 which is incompatible.\n",
"tensorflow 2.5.0 requires typing-extensions~=3.7.4, but you have typing-extensions 3.10.0.2 which is incompatible.\n",
"rich 13.0.1 requires typing-extensions<5.0,>=4.0.0; python_version < \"3.9\", but you have typing-extensions 3.10.0.2 which is incompatible.\n",
"ml-dtypes 0.0.4 requires numpy>1.20, but you have numpy 1.19.5 which is incompatible.\n",
"flax 0.6.3 requires jax>=0.3.16, but you have jax 0.2.28 which is incompatible.\n",
"flax 0.6.3 requires typing_extensions>=4.1.1, but you have typing-extensions 3.10.0.2 which is incompatible.\u001b[0m\u001b[31m\n",
"\u001b[0mSuccessfully installed chex-0.1.5 jax-0.2.28 numpy-1.19.5 optax-0.0.9 typing-extensions-3.10.0.2\n",
"\n",
"Usage: \n",
" /bin/python3 -m pip install [options] <requirement specifier> [package-index-options] ...\n",
" /bin/python3 -m pip install [options] -r <requirements file> [package-index-options] ...\n",
" /bin/python3 -m pip install [options] [-e] <vcs project url> ...\n",
" /bin/python3 -m pip install [options] [-e] <local project path> ...\n",
" /bin/python3 -m pip install [options] <archive url/path> ...\n",
"\n",
"no such option: -u\n",
"Defaulting to user installation because normal site-packages is not writeable\n",
"Looking in links: https://storage.googleapis.com/jax-releases/libtpu_releases.html\n",
"Requirement already satisfied: jax[tpu] in /home/moe/.local/lib/python3.8/site-packages (0.2.28)\n",
"Requirement already satisfied: scipy>=1.2.1 in /home/moe/.local/lib/python3.8/site-packages (from jax[tpu]) (1.10.0)\n",
"Requirement already satisfied: numpy>=1.19 in /home/moe/.local/lib/python3.8/site-packages (from jax[tpu]) (1.19.5)\n",
"Requirement already satisfied: typing-extensions in /home/moe/.local/lib/python3.8/site-packages (from jax[tpu]) (3.10.0.2)\n",
"Requirement already satisfied: opt-einsum in /home/moe/.local/lib/python3.8/site-packages (from jax[tpu]) (3.3.0)\n",
"Requirement already satisfied: absl-py in /home/moe/.local/lib/python3.8/site-packages (from jax[tpu]) (0.15.0)\n",
"Requirement already satisfied: requests in /home/moe/.local/lib/python3.8/site-packages (from jax[tpu]) (2.25.1)\n",
"Requirement already satisfied: jaxlib==0.1.76 in /home/moe/.local/lib/python3.8/site-packages (from jax[tpu]) (0.1.76)\n",
"Requirement already satisfied: libtpu-nightly==0.1.dev20220128 in /home/moe/.local/lib/python3.8/site-packages (from jax[tpu]) (0.1.dev20220128)\n",
"Requirement already satisfied: flatbuffers<3.0,>=1.12 in /home/moe/.local/lib/python3.8/site-packages (from jaxlib==0.1.76->jax[tpu]) (1.12)\n",
"Requirement already satisfied: six in /home/moe/.local/lib/python3.8/site-packages (from absl-py->jax[tpu]) (1.15.0)\n",
"Requirement already satisfied: chardet<5,>=3.0.2 in /usr/lib/python3/dist-packages (from requests->jax[tpu]) (3.0.4)\n",
"Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/moe/.local/lib/python3.8/site-packages (from requests->jax[tpu]) (1.26.15)\n",
"Requirement already satisfied: certifi>=2017.4.17 in /usr/lib/python3/dist-packages (from requests->jax[tpu]) (2019.11.28)\n",
"Requirement already satisfied: idna<3,>=2.5 in /usr/lib/python3/dist-packages (from requests->jax[tpu]) (2.8)\n",
"Defaulting to user installation because normal site-packages is not writeable\n",
"Collecting git+https://github.com/google/flax.git\n",
" Cloning https://github.com/google/flax.git to /tmp/pip-req-build-t_xht7qd\n",
" Running command git clone --filter=blob:none --quiet https://github.com/google/flax.git /tmp/pip-req-build-t_xht7qd\n",
" Resolved https://github.com/google/flax.git to commit 332b7d041de235d4579ab0fb4c82cc4ad94ade62\n",
" Installing build dependencies ... \u001b[?25ldone\n",
"\u001b[?25h Getting requirements to build wheel ... \u001b[?25ldone\n",
"\u001b[?25h Preparing metadata (pyproject.toml) ... \u001b[?25ldone\n",
"\u001b[?25hRequirement already satisfied: numpy>=1.12 in /home/moe/.local/lib/python3.8/site-packages (from flax==0.6.9) (1.19.5)\n",
"Collecting jax>=0.4.2\n",
" Using cached jax-0.4.8-py3-none-any.whl\n",
"Collecting typing-extensions>=4.1.1\n",
" Using cached typing_extensions-4.5.0-py3-none-any.whl (27 kB)\n",
"Requirement already satisfied: rich>=11.1 in /home/moe/.local/lib/python3.8/site-packages (from flax==0.6.9) (13.0.1)\n",
"Requirement already satisfied: optax in /home/moe/.local/lib/python3.8/site-packages (from flax==0.6.9) (0.0.9)\n",
"Collecting orbax-checkpoint\n",
" Downloading orbax_checkpoint-0.1.8-py3-none-any.whl (77 kB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m77.1/77.1 kB\u001b[0m \u001b[31m3.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hRequirement already satisfied: PyYAML>=5.4.1 in /usr/local/lib/python3.8/dist-packages (from flax==0.6.9) (5.4.1)\n",
"Requirement already satisfied: msgpack in /home/moe/.local/lib/python3.8/site-packages (from flax==0.6.9) (1.0.4)\n",
"Requirement already satisfied: tensorstore in /home/moe/.local/lib/python3.8/site-packages (from flax==0.6.9) (0.1.30)\n",
"Collecting numpy>=1.12\n",
" Using cached numpy-1.24.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (17.3 MB)\n",
"Requirement already satisfied: ml-dtypes>=0.0.3 in /home/moe/.local/lib/python3.8/site-packages (from jax>=0.4.2->flax==0.6.9) (0.0.4)\n",
"Requirement already satisfied: opt-einsum in /home/moe/.local/lib/python3.8/site-packages (from jax>=0.4.2->flax==0.6.9) (3.3.0)\n",
"Requirement already satisfied: scipy>=1.7 in /home/moe/.local/lib/python3.8/site-packages (from jax>=0.4.2->flax==0.6.9) (1.10.0)\n",
"Requirement already satisfied: pygments<3.0.0,>=2.6.0 in /home/moe/.local/lib/python3.8/site-packages (from rich>=11.1->flax==0.6.9) (2.14.0)\n",
"Requirement already satisfied: commonmark<0.10.0,>=0.9.0 in /home/moe/.local/lib/python3.8/site-packages (from rich>=11.1->flax==0.6.9) (0.9.1)\n",
"Requirement already satisfied: jaxlib>=0.1.37 in /home/moe/.local/lib/python3.8/site-packages (from optax->flax==0.6.9) (0.1.76)\n",
"Requirement already satisfied: absl-py>=0.7.1 in /home/moe/.local/lib/python3.8/site-packages (from optax->flax==0.6.9) (0.15.0)\n",
"Requirement already satisfied: chex>=0.0.4 in /home/moe/.local/lib/python3.8/site-packages (from optax->flax==0.6.9) (0.1.5)\n",
"Requirement already satisfied: importlib_resources in /home/moe/.local/lib/python3.8/site-packages (from orbax-checkpoint->flax==0.6.9) (5.10.2)\n",
"Requirement already satisfied: nest_asyncio in /home/moe/.local/lib/python3.8/site-packages (from orbax-checkpoint->flax==0.6.9) (1.5.6)\n",
"Collecting tensorstore\n",
" Downloading tensorstore-0.1.35-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (10.9 MB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m10.9/10.9 MB\u001b[0m \u001b[31m79.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m0:01\u001b[0m\n",
"\u001b[?25hRequirement already satisfied: etils in /home/moe/.local/lib/python3.8/site-packages (from orbax-checkpoint->flax==0.6.9) (1.0.0)\n",
"Requirement already satisfied: cached_property in /home/moe/.local/lib/python3.8/site-packages (from orbax-checkpoint->flax==0.6.9) (1.5.2)\n",
"Requirement already satisfied: six in /home/moe/.local/lib/python3.8/site-packages (from absl-py>=0.7.1->optax->flax==0.6.9) (1.15.0)\n",
"Requirement already satisfied: dm-tree>=0.1.5 in /home/moe/.local/lib/python3.8/site-packages (from chex>=0.0.4->optax->flax==0.6.9) (0.1.8)\n",
"Requirement already satisfied: toolz>=0.9.0 in /home/moe/.local/lib/python3.8/site-packages (from chex>=0.0.4->optax->flax==0.6.9) (0.12.0)\n",
"Requirement already satisfied: flatbuffers<3.0,>=1.12 in /home/moe/.local/lib/python3.8/site-packages (from jaxlib>=0.1.37->optax->flax==0.6.9) (1.12)\n",
"Requirement already satisfied: zipp>=3.1.0 in /home/moe/.local/lib/python3.8/site-packages (from importlib_resources->orbax-checkpoint->flax==0.6.9) (3.11.0)\n",
"Building wheels for collected packages: flax\n",
" Building wheel for flax (pyproject.toml) ... \u001b[?25ldone\n",
"\u001b[?25h Created wheel for flax: filename=flax-0.6.9-py3-none-any.whl size=216742 sha256=591b5d1eabf22226dbc486d7067765a1f364b54b267e32762ec5d04bf6023aa2\n",
" Stored in directory: /tmp/pip-ephem-wheel-cache-fe_fkuzd/wheels/5d/cd/13/31d6d37c2228329484735564bfa16a48768d2232924588f91b\n",
"Successfully built flax\n",
"Installing collected packages: typing-extensions, numpy, tensorstore, jax, orbax-checkpoint, flax\n",
" Attempting uninstall: typing-extensions\n",
" Found existing installation: typing-extensions 3.10.0.2\n",
" Uninstalling typing-extensions-3.10.0.2:\n",
" Successfully uninstalled typing-extensions-3.10.0.2\n",
" Attempting uninstall: numpy\n",
" Found existing installation: numpy 1.19.5\n",
" Uninstalling numpy-1.19.5:\n",
" Successfully uninstalled numpy-1.19.5\n",
" Attempting uninstall: tensorstore\n",
" Found existing installation: tensorstore 0.1.30\n",
" Uninstalling tensorstore-0.1.30:\n",
" Successfully uninstalled tensorstore-0.1.30\n",
" Attempting uninstall: jax\n",
" Found existing installation: jax 0.2.28\n",
" Uninstalling jax-0.2.28:\n",
" Successfully uninstalled jax-0.2.28\n",
" Attempting uninstall: flax\n",
" Found existing installation: flax 0.6.3\n",
" Uninstalling flax-0.6.3:\n",
" Successfully uninstalled flax-0.6.3\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",
"tensorflow 2.5.0 requires grpcio~=1.34.0, but you have grpcio 1.53.0 which is incompatible.\n",
"tensorflow 2.5.0 requires numpy~=1.19.2, but you have numpy 1.24.2 which is incompatible.\n",
"tensorflow 2.5.0 requires tensorflow-estimator<2.6.0,>=2.5.0rc0, but you have tensorflow-estimator 2.6.0 which is incompatible.\n",
"tensorflow 2.5.0 requires typing-extensions~=3.7.4, but you have typing-extensions 4.5.0 which is incompatible.\n",
"tensorflow-cpu 2.6.5 requires numpy~=1.19.2, but you have numpy 1.24.2 which is incompatible.\n",
"tensorflow-cpu 2.6.5 requires typing-extensions<3.11,>=3.7, but you have typing-extensions 4.5.0 which is incompatible.\u001b[0m\u001b[31m\n",
"\u001b[0mSuccessfully installed flax-0.6.9 jax-0.4.8 numpy-1.24.2 orbax-checkpoint-0.1.8 tensorstore-0.1.35 typing-extensions-4.5.0\n",
"Defaulting to user installation because normal site-packages is not writeable\n",
"Collecting git+https://github.com/deepmind/optax.git\n",
" Cloning https://github.com/deepmind/optax.git to /tmp/pip-req-build-d7mf0dm6\n",
" Running command git clone --filter=blob:none --quiet https://github.com/deepmind/optax.git /tmp/pip-req-build-d7mf0dm6\n",
" Resolved https://github.com/deepmind/optax.git to commit 04768d252911d6af4d4d36361930ccd0a54f9160\n",
" Preparing metadata (setup.py) ... \u001b[?25ldone\n",
"\u001b[?25hRequirement already satisfied: absl-py>=0.7.1 in /home/moe/.local/lib/python3.8/site-packages (from optax==0.1.5.dev0) (0.15.0)\n",
"Requirement already satisfied: chex>=0.1.5 in /home/moe/.local/lib/python3.8/site-packages (from optax==0.1.5.dev0) (0.1.5)\n",
"Requirement already satisfied: jax>=0.1.55 in /home/moe/.local/lib/python3.8/site-packages (from optax==0.1.5.dev0) (0.4.8)\n",
"Requirement already satisfied: jaxlib>=0.1.37 in /home/moe/.local/lib/python3.8/site-packages (from optax==0.1.5.dev0) (0.1.76)\n",
"Requirement already satisfied: numpy>=1.18.0 in /home/moe/.local/lib/python3.8/site-packages (from optax==0.1.5.dev0) (1.24.2)\n",
"Requirement already satisfied: six in /home/moe/.local/lib/python3.8/site-packages (from absl-py>=0.7.1->optax==0.1.5.dev0) (1.15.0)\n",
"Requirement already satisfied: dm-tree>=0.1.5 in /home/moe/.local/lib/python3.8/site-packages (from chex>=0.1.5->optax==0.1.5.dev0) (0.1.8)\n",
"Requirement already satisfied: toolz>=0.9.0 in /home/moe/.local/lib/python3.8/site-packages (from chex>=0.1.5->optax==0.1.5.dev0) (0.12.0)\n",
"Requirement already satisfied: scipy>=1.7 in /home/moe/.local/lib/python3.8/site-packages (from jax>=0.1.55->optax==0.1.5.dev0) (1.10.0)\n",
"Requirement already satisfied: ml-dtypes>=0.0.3 in /home/moe/.local/lib/python3.8/site-packages (from jax>=0.1.55->optax==0.1.5.dev0) (0.0.4)\n",
"Requirement already satisfied: opt-einsum in /home/moe/.local/lib/python3.8/site-packages (from jax>=0.1.55->optax==0.1.5.dev0) (3.3.0)\n",
"Requirement already satisfied: flatbuffers<3.0,>=1.12 in /home/moe/.local/lib/python3.8/site-packages (from jaxlib>=0.1.37->optax==0.1.5.dev0) (1.12)\n",
"Building wheels for collected packages: optax\n",
" Building wheel for optax (setup.py) ... \u001b[?25ldone\n",
"\u001b[?25h Created wheel for optax: filename=optax-0.1.5.dev0-py3-none-any.whl size=164328 sha256=b3146613ae800d20bcb8e34efa1b1622600f7f60a3dfa5256204ca22550120af\n",
" Stored in directory: /tmp/pip-ephem-wheel-cache-521s1tej/wheels/bf/41/a5/0e13738ca7803f1431e452f31882b235ab0ea6b54f5a9b0b22\n",
"Successfully built optax\n",
"Installing collected packages: optax\n",
" Attempting uninstall: optax\n",
" Found existing installation: optax 0.0.9\n",
" Uninstalling optax-0.0.9:\n",
" Successfully uninstalled optax-0.0.9\n",
"Successfully installed optax-0.1.5.dev0\n",
"Defaulting to user installation because normal site-packages is not writeable\n",
"Collecting git+https://github.com/deepmind/chex.git\n",
" Cloning https://github.com/deepmind/chex.git to /tmp/pip-req-build-qbakjgnz\n",
" Running command git clone --filter=blob:none --quiet https://github.com/deepmind/chex.git /tmp/pip-req-build-qbakjgnz\n",
" Resolved https://github.com/deepmind/chex.git to commit a65cc6eda665a7f9d531cb4d59d1629eccb55372\n",
" Preparing metadata (setup.py) ... \u001b[?25ldone\n",
"\u001b[?25hRequirement already satisfied: absl-py>=0.9.0 in /home/moe/.local/lib/python3.8/site-packages (from chex==0.1.7) (0.15.0)\n",
"Requirement already satisfied: dm-tree>=0.1.5 in /home/moe/.local/lib/python3.8/site-packages (from chex==0.1.7) (0.1.8)\n",
"Requirement already satisfied: jax>=0.4.6 in /home/moe/.local/lib/python3.8/site-packages (from chex==0.1.7) (0.4.8)\n",
"Requirement already satisfied: jaxlib>=0.1.37 in /home/moe/.local/lib/python3.8/site-packages (from chex==0.1.7) (0.1.76)\n",
"Requirement already satisfied: numpy>=1.18.0 in /home/moe/.local/lib/python3.8/site-packages (from chex==0.1.7) (1.24.2)\n",
"Requirement already satisfied: toolz>=0.9.0 in /home/moe/.local/lib/python3.8/site-packages (from chex==0.1.7) (0.12.0)\n",
"Requirement already satisfied: typing_extensions>=4.2.0 in /home/moe/.local/lib/python3.8/site-packages (from chex==0.1.7) (4.5.0)\n",
"Requirement already satisfied: six in /home/moe/.local/lib/python3.8/site-packages (from absl-py>=0.9.0->chex==0.1.7) (1.15.0)\n",
"Requirement already satisfied: ml-dtypes>=0.0.3 in /home/moe/.local/lib/python3.8/site-packages (from jax>=0.4.6->chex==0.1.7) (0.0.4)\n",
"Requirement already satisfied: scipy>=1.7 in /home/moe/.local/lib/python3.8/site-packages (from jax>=0.4.6->chex==0.1.7) (1.10.0)\n",
"Requirement already satisfied: opt-einsum in /home/moe/.local/lib/python3.8/site-packages (from jax>=0.4.6->chex==0.1.7) (3.3.0)\n",
"Requirement already satisfied: flatbuffers<3.0,>=1.12 in /home/moe/.local/lib/python3.8/site-packages (from jaxlib>=0.1.37->chex==0.1.7) (1.12)\n",
"Building wheels for collected packages: chex\n",
" Building wheel for chex (setup.py) ... \u001b[?25ldone\n",
"\u001b[?25h Created wheel for chex: filename=chex-0.1.7-py3-none-any.whl size=89780 sha256=132fa1615c1f2bc9fe1c0030f5afa5b2008d0124c3c91d14c6c4ce65137dafbe\n",
" Stored in directory: /tmp/pip-ephem-wheel-cache-9hhb2yzn/wheels/7c/0a/b3/49f5d3f1b57e71ff799d8f997436af83578d13b52e7fd24163\n",
"Successfully built chex\n",
"Installing collected packages: chex\n",
" Attempting uninstall: chex\n",
" Found existing installation: chex 0.1.5\n",
" Uninstalling chex-0.1.5:\n",
" Successfully uninstalled chex-0.1.5\n",
"Successfully installed chex-0.1.7\n"
]
}
],
"source": [
"!sudo 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",
"!/bin/python3 -m 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",
"!/bin/python3 -m pip install -u mesh-transformer-jax/ tensorflow\n",
"!/bin/python3 -m pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n",
"!/bin/python3 -m pip install --upgrade git+https://github.com/google/flax.git\n",
"!/bin/python3 -m pip install git+https://github.com/deepmind/optax.git\n",
"!/bin/python3 -m pip install git+https://github.com/deepmind/chex.git"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "aO1UXepF-0Uq"
},
"source": [
"## Setup Model\n"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {
"id": "ex0qJgaueZtJ"
},
"outputs": [],
"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']"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "NIgUVdFLe4A8"
},
"source": [
"Sometimes the next step errors for some reason, just run it again ¯\\\\\\_(ツ)\\_/¯"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {
"id": "-A5IGYSaeze3"
},
"outputs": [
{
"ename": "AttributeError",
"evalue": "module 'jax' has no attribute 'Array'",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mAttributeError\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[15], line 6\u001b[0m\n\u001b[1;32m 4\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mjax\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mexperimental\u001b[39;00m \u001b[39mimport\u001b[39;00m maps\n\u001b[1;32m 5\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mnumpy\u001b[39;00m \u001b[39mas\u001b[39;00m \u001b[39mnp\u001b[39;00m\n\u001b[0;32m----> 6\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39moptax\u001b[39;00m\n\u001b[1;32m 7\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mtransformers\u001b[39;00m\n\u001b[1;32m 9\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mmesh_transformer\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mcheckpoint\u001b[39;00m \u001b[39mimport\u001b[39;00m read_ckpt_lowmem\n",
"File \u001b[0;32m~/.local/lib/python3.8/site-packages/optax/__init__.py:17\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[39m# Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.\u001b[39;00m\n\u001b[1;32m 2\u001b[0m \u001b[39m#\u001b[39;00m\n\u001b[1;32m 3\u001b[0m \u001b[39m# Licensed under the Apache License, Version 2.0 (the \"License\");\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 13\u001b[0m \u001b[39m# limitations under the License.\u001b[39;00m\n\u001b[1;32m 14\u001b[0m \u001b[39m# ==============================================================================\u001b[39;00m\n\u001b[1;32m 15\u001b[0m \u001b[39m\"\"\"Optax: composable gradient processing and optimization, in JAX.\"\"\"\u001b[39;00m\n\u001b[0;32m---> 17\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39moptax\u001b[39;00m \u001b[39mimport\u001b[39;00m experimental\n\u001b[1;32m 18\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39moptax\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39m_src\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39malias\u001b[39;00m \u001b[39mimport\u001b[39;00m adabelief\n\u001b[1;32m 19\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39moptax\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39m_src\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39malias\u001b[39;00m \u001b[39mimport\u001b[39;00m adafactor\n",
"File \u001b[0;32m~/.local/lib/python3.8/site-packages/optax/experimental/__init__.py:20\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[39m# Copyright 2021 DeepMind Technologies Limited. All Rights Reserved.\u001b[39;00m\n\u001b[1;32m 2\u001b[0m \u001b[39m#\u001b[39;00m\n\u001b[1;32m 3\u001b[0m \u001b[39m# Licensed under the Apache License, Version 2.0 (the \"License\");\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 13\u001b[0m \u001b[39m# limitations under the License.\u001b[39;00m\n\u001b[1;32m 14\u001b[0m \u001b[39m# ==============================================================================\u001b[39;00m\n\u001b[1;32m 15\u001b[0m \u001b[39m\"\"\"Experimental features in Optax.\u001b[39;00m\n\u001b[1;32m 16\u001b[0m \n\u001b[1;32m 17\u001b[0m \u001b[39mFeatures may be removed or modified at any time.\u001b[39;00m\n\u001b[1;32m 18\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[0;32m---> 20\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39moptax\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39m_src\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mexperimental\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mcomplex_valued\u001b[39;00m \u001b[39mimport\u001b[39;00m split_real_and_imaginary\n\u001b[1;32m 21\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39moptax\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39m_src\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mexperimental\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mcomplex_valued\u001b[39;00m \u001b[39mimport\u001b[39;00m SplitRealAndImaginaryState\n\u001b[1;32m 22\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39moptax\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39m_src\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mexperimental\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mextra_args\u001b[39;00m \u001b[39mimport\u001b[39;00m GradientTransformationWithExtraArgs\n",
"File \u001b[0;32m~/.local/lib/python3.8/site-packages/optax/_src/experimental/complex_valued.py:32\u001b[0m\n\u001b[1;32m 15\u001b[0m \u001b[39m\"\"\"Complex-valued optimization.\u001b[39;00m\n\u001b[1;32m 16\u001b[0m \n\u001b[1;32m 17\u001b[0m \u001b[39mWhen using `split_real_and_imaginary` to wrap an optimizer, we split the complex\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 27\u001b[0m \u001b[39mSee details at https://github.com/deepmind/optax/issues/196\u001b[39;00m\n\u001b[1;32m 28\u001b[0m \u001b[39m\"\"\"\u001b[39;00m\n\u001b[1;32m 30\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mtyping\u001b[39;00m \u001b[39mimport\u001b[39;00m NamedTuple, Union\n\u001b[0;32m---> 32\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mchex\u001b[39;00m\n\u001b[1;32m 33\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mjax\u001b[39;00m\n\u001b[1;32m 34\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mjax\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mnumpy\u001b[39;00m \u001b[39mas\u001b[39;00m \u001b[39mjnp\u001b[39;00m\n",
"File \u001b[0;32m~/.local/lib/python3.8/site-packages/chex/__init__.py:17\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[39m# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.\u001b[39;00m\n\u001b[1;32m 2\u001b[0m \u001b[39m#\u001b[39;00m\n\u001b[1;32m 3\u001b[0m \u001b[39m# Licensed under the Apache License, Version 2.0 (the \"License\");\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 13\u001b[0m \u001b[39m# limitations under the License.\u001b[39;00m\n\u001b[1;32m 14\u001b[0m \u001b[39m# ==============================================================================\u001b[39;00m\n\u001b[1;32m 15\u001b[0m \u001b[39m\"\"\"Chex: Testing made fun, in JAX!\"\"\"\u001b[39;00m\n\u001b[0;32m---> 17\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mchex\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39m_src\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39masserts\u001b[39;00m \u001b[39mimport\u001b[39;00m assert_axis_dimension\n\u001b[1;32m 18\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mchex\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39m_src\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39masserts\u001b[39;00m \u001b[39mimport\u001b[39;00m assert_axis_dimension_comparator\n\u001b[1;32m 19\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mchex\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39m_src\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39masserts\u001b[39;00m \u001b[39mimport\u001b[39;00m assert_axis_dimension_gt\n",
"File \u001b[0;32m~/.local/lib/python3.8/site-packages/chex/_src/asserts.py:26\u001b[0m\n\u001b[1;32m 23\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39munittest\u001b[39;00m\n\u001b[1;32m 24\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39munittest\u001b[39;00m \u001b[39mimport\u001b[39;00m mock\n\u001b[0;32m---> 26\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mchex\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39m_src\u001b[39;00m \u001b[39mimport\u001b[39;00m asserts_internal \u001b[39mas\u001b[39;00m _ai\n\u001b[1;32m 27\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mchex\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39m_src\u001b[39;00m \u001b[39mimport\u001b[39;00m pytypes\n\u001b[1;32m 28\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mjax\u001b[39;00m\n",
"File \u001b[0;32m~/.local/lib/python3.8/site-packages/chex/_src/asserts_internal.py:34\u001b[0m\n\u001b[1;32m 31\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mtyping\u001b[39;00m \u001b[39mimport\u001b[39;00m Any, Sequence, Union, Callable, List, Optional, Set, Tuple, Type\n\u001b[1;32m 33\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mabsl\u001b[39;00m \u001b[39mimport\u001b[39;00m logging\n\u001b[0;32m---> 34\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mchex\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39m_src\u001b[39;00m \u001b[39mimport\u001b[39;00m pytypes\n\u001b[1;32m 35\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mjax\u001b[39;00m\n\u001b[1;32m 36\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mjax\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mexperimental\u001b[39;00m \u001b[39mimport\u001b[39;00m checkify\n",
"File \u001b[0;32m~/.local/lib/python3.8/site-packages/chex/_src/pytypes.py:27\u001b[0m\n\u001b[1;32m 24\u001b[0m ArrayNumpy \u001b[39m=\u001b[39m np\u001b[39m.\u001b[39mndarray\n\u001b[1;32m 26\u001b[0m \u001b[39m# For instance checking, use `isinstance(x, jax.Array)`.\u001b[39;00m\n\u001b[0;32m---> 27\u001b[0m ArrayDevice \u001b[39m=\u001b[39m jax\u001b[39m.\u001b[39;49mArray\n\u001b[1;32m 29\u001b[0m \u001b[39m# Types for backward compatibility.\u001b[39;00m\n\u001b[1;32m 30\u001b[0m ArraySharded \u001b[39m=\u001b[39m jax\u001b[39m.\u001b[39mArray\n",
"\u001b[0;31mAttributeError\u001b[0m: module 'jax' has no attribute 'Array'"
]
}
],
"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"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 167,
"referenced_widgets": [
"5b8a31b3b4034116af0a81d897c3123b",
"835ee2070c05443c9fa3ae6fece3fb44",
"9e6646f7066341ef9619f5f7e6f7c8b8",
"4f803e3d56b74e18997b216bf9d65337",
"f9fa7c1bc9f04df09e487b16761d6b8f",
"39aed05f694a495abd03aa00ff71efb3",
"61b0857802184b51abf77039fad681d7",
"0fb689df0a8b42c2b87e7f11b818b88a",
"8b76924bdac749d5b4908678fd0a7497",
"b03bd083afc84f52834d4a124e57c442",
"3ba9fe91731f4ff380cda2183ad72c4d",
"4e01b9cceffb4d178e97f1d1b36d62b0",
"b3ac7912f51746afae20616269dc69f0",
"fb1263f5a3da42e6acab0e9a66551b32",
"3426927332884df4b1ec5544ed479c45",
"debf4a0a609a4cffa5cd686a0a23f6e2",
"7197de8a39ab4e4186eabd061c730cf5",
"c31661c9a1ff4bd681cc10d70fd287c0",
"bb11a4fbf16f488082aa2bac7443fd62",
"5360630d6b0541dbaafa7ff91610bd84",
"f3883c133c62438db114c22cfc94e351",
"74672c43909444c592955afd0d727a28",
"67f7c3bc59384dc7812d7130c7cde450",
"bb7410ff9cb4484da73302993427b90a"
]
},
"id": "QAgKq-X2kmba",
"outputId": "041117df-a315-4b95-9caa-26cb870ff3df"
},
"outputs": [],
"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')"
]
},
{
"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",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "lwNETD2Uk8nu",
"outputId": "a659285b-5b46-4ddf-d65c-fbdad0594b86"
},
"outputs": [],
"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))"
]
},
{
"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",
"execution_count": null,
"metadata": {
"id": "IRiH9xLFZw8b"
},
"outputs": [],
"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)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "ZVzs2TYlvYeX",
"outputId": "80c8b1e5-0d1c-4799-d682-4a5be0c038a1"
},
"outputs": [],
"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])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "nvlAK6RbCJYg",
"outputId": "c0525611-ecc8-422e-d8e3-24b3a2159083"
},
"outputs": [],
"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])"
]
}
],
"metadata": {
"accelerator": "TPU",
"colab": {
"machine_shape": "hm",
"name": "GPT-J-6B Inference Demo.ipynb",
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.10"
},
"vscode": {
"interpreter": {
"hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
}
},
"widgets": {
"application/vnd.jupyter.widget-state+json": {
"0fb689df0a8b42c2b87e7f11b818b88a": {
"model_module": "@jupyter-widgets/base",
"model_module_version": "1.2.0",
"model_name": "LayoutModel",
"state": {
"_model_module": "@jupyter-widgets/base",
"_model_module_version": "1.2.0",
"_model_name": "LayoutModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/base",
"_view_module_version": "1.2.0",
"_view_name": "LayoutView",
"align_content": null,
"align_items": null,
"align_self": null,
"border": null,
"bottom": null,
"display": null,
"flex": null,
"flex_flow": null,
"grid_area": null,
"grid_auto_columns": null,
"grid_auto_flow": null,
"grid_auto_rows": null,
"grid_column": null,
"grid_gap": null,
"grid_row": null,
"grid_template_areas": null,
"grid_template_columns": null,
"grid_template_rows": null,
"height": null,
"justify_content": null,
"justify_items": null,
"left": null,
"margin": null,
"max_height": null,
"max_width": null,
"min_height": null,
"min_width": null,
"object_fit": null,
"object_position": null,
"order": null,
"overflow": null,
"overflow_x": null,
"overflow_y": null,
"padding": null,
"right": null,
"top": null,
"visibility": null,
"width": null
}
},
"3426927332884df4b1ec5544ed479c45": {
"model_module": "@jupyter-widgets/controls",
"model_module_version": "1.5.0",
"model_name": "DescriptionStyleModel",
"state": {
"_model_module": "@jupyter-widgets/controls",
"_model_module_version": "1.5.0",
"_model_name": "DescriptionStyleModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/base",
"_view_module_version": "1.2.0",
"_view_name": "StyleView",
"description_width": ""
}
},
"39aed05f694a495abd03aa00ff71efb3": {
"model_module": "@jupyter-widgets/base",
"model_module_version": "1.2.0",
"model_name": "LayoutModel",
"state": {
"_model_module": "@jupyter-widgets/base",
"_model_module_version": "1.2.0",
"_model_name": "LayoutModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/base",
"_view_module_version": "1.2.0",
"_view_name": "LayoutView",
"align_content": null,
"align_items": null,
"align_self": null,
"border": null,
"bottom": null,
"display": null,
"flex": null,
"flex_flow": null,
"grid_area": null,
"grid_auto_columns": null,
"grid_auto_flow": null,
"grid_auto_rows": null,
"grid_column": null,
"grid_gap": null,
"grid_row": null,
"grid_template_areas": null,
"grid_template_columns": null,
"grid_template_rows": null,
"height": null,
"justify_content": null,
"justify_items": null,
"left": null,
"margin": null,
"max_height": null,
"max_width": null,
"min_height": null,
"min_width": null,
"object_fit": null,
"object_position": null,
"order": null,
"overflow": null,
"overflow_x": null,
"overflow_y": null,
"padding": null,
"right": null,
"top": null,
"visibility": null,
"width": null
}
},
"3ba9fe91731f4ff380cda2183ad72c4d": {
"model_module": "@jupyter-widgets/controls",
"model_module_version": "1.5.0",
"model_name": "FloatProgressModel",
"state": {
"_dom_classes": [],
"_model_module": "@jupyter-widgets/controls",
"_model_module_version": "1.5.0",
"_model_name": "FloatProgressModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/controls",
"_view_module_version": "1.5.0",
"_view_name": "ProgressView",
"bar_style": "success",
"description": "Downloading: 100%",
"description_tooltip": null,
"layout": "IPY_MODEL_fb1263f5a3da42e6acab0e9a66551b32",
"max": 456318,
"min": 0,
"orientation": "horizontal",
"style": "IPY_MODEL_b3ac7912f51746afae20616269dc69f0",
"value": 456318
}
},
"4e01b9cceffb4d178e97f1d1b36d62b0": {
"model_module": "@jupyter-widgets/controls",
"model_module_version": "1.5.0",
"model_name": "HTMLModel",
"state": {
"_dom_classes": [],
"_model_module": "@jupyter-widgets/controls",
"_model_module_version": "1.5.0",
"_model_name": "HTMLModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/controls",
"_view_module_version": "1.5.0",
"_view_name": "HTMLView",
"description": "",
"description_tooltip": null,
"layout": "IPY_MODEL_debf4a0a609a4cffa5cd686a0a23f6e2",
"placeholder": "​",
"style": "IPY_MODEL_3426927332884df4b1ec5544ed479c45",
"value": " 456k/456k [00:00&lt;00:00, 928kB/s]"
}
},
"4f803e3d56b74e18997b216bf9d65337": {
"model_module": "@jupyter-widgets/controls",
"model_module_version": "1.5.0",
"model_name": "HTMLModel",
"state": {
"_dom_classes": [],
"_model_module": "@jupyter-widgets/controls",
"_model_module_version": "1.5.0",
"_model_name": "HTMLModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/controls",
"_view_module_version": "1.5.0",
"_view_name": "HTMLView",
"description": "",
"description_tooltip": null,
"layout": "IPY_MODEL_0fb689df0a8b42c2b87e7f11b818b88a",
"placeholder": "​",
"style": "IPY_MODEL_61b0857802184b51abf77039fad681d7",
"value": " 1.04M/1.04M [00:07&lt;00:00, 148kB/s]"
}
},
"5360630d6b0541dbaafa7ff91610bd84": {
"model_module": "@jupyter-widgets/controls",
"model_module_version": "1.5.0",
"model_name": "HTMLModel",
"state": {
"_dom_classes": [],
"_model_module": "@jupyter-widgets/controls",
"_model_module_version": "1.5.0",
"_model_name": "HTMLModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/controls",
"_view_module_version": "1.5.0",
"_view_name": "HTMLView",
"description": "",
"description_tooltip": null,
"layout": "IPY_MODEL_bb7410ff9cb4484da73302993427b90a",
"placeholder": "​",
"style": "IPY_MODEL_67f7c3bc59384dc7812d7130c7cde450",
"value": " 1.36M/1.36M [00:03&lt;00:00, 404kB/s]"
}
},
"5b8a31b3b4034116af0a81d897c3123b": {
"model_module": "@jupyter-widgets/controls",
"model_module_version": "1.5.0",
"model_name": "HBoxModel",
"state": {
"_dom_classes": [],
"_model_module": "@jupyter-widgets/controls",
"_model_module_version": "1.5.0",
"_model_name": "HBoxModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/controls",
"_view_module_version": "1.5.0",
"_view_name": "HBoxView",
"box_style": "",
"children": [
"IPY_MODEL_9e6646f7066341ef9619f5f7e6f7c8b8",
"IPY_MODEL_4f803e3d56b74e18997b216bf9d65337"
],
"layout": "IPY_MODEL_835ee2070c05443c9fa3ae6fece3fb44"
}
},
"61b0857802184b51abf77039fad681d7": {
"model_module": "@jupyter-widgets/controls",
"model_module_version": "1.5.0",
"model_name": "DescriptionStyleModel",
"state": {
"_model_module": "@jupyter-widgets/controls",
"_model_module_version": "1.5.0",
"_model_name": "DescriptionStyleModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/base",
"_view_module_version": "1.2.0",
"_view_name": "StyleView",
"description_width": ""
}
},
"67f7c3bc59384dc7812d7130c7cde450": {
"model_module": "@jupyter-widgets/controls",
"model_module_version": "1.5.0",
"model_name": "DescriptionStyleModel",
"state": {
"_model_module": "@jupyter-widgets/controls",
"_model_module_version": "1.5.0",
"_model_name": "DescriptionStyleModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/base",
"_view_module_version": "1.2.0",
"_view_name": "StyleView",
"description_width": ""
}
},
"7197de8a39ab4e4186eabd061c730cf5": {
"model_module": "@jupyter-widgets/controls",
"model_module_version": "1.5.0",
"model_name": "HBoxModel",
"state": {
"_dom_classes": [],
"_model_module": "@jupyter-widgets/controls",
"_model_module_version": "1.5.0",
"_model_name": "HBoxModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/controls",
"_view_module_version": "1.5.0",
"_view_name": "HBoxView",
"box_style": "",
"children": [
"IPY_MODEL_bb11a4fbf16f488082aa2bac7443fd62",
"IPY_MODEL_5360630d6b0541dbaafa7ff91610bd84"
],
"layout": "IPY_MODEL_c31661c9a1ff4bd681cc10d70fd287c0"
}
},
"74672c43909444c592955afd0d727a28": {
"model_module": "@jupyter-widgets/base",
"model_module_version": "1.2.0",
"model_name": "LayoutModel",
"state": {
"_model_module": "@jupyter-widgets/base",
"_model_module_version": "1.2.0",
"_model_name": "LayoutModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/base",
"_view_module_version": "1.2.0",
"_view_name": "LayoutView",
"align_content": null,
"align_items": null,
"align_self": null,
"border": null,
"bottom": null,
"display": null,
"flex": null,
"flex_flow": null,
"grid_area": null,
"grid_auto_columns": null,
"grid_auto_flow": null,
"grid_auto_rows": null,
"grid_column": null,
"grid_gap": null,
"grid_row": null,
"grid_template_areas": null,
"grid_template_columns": null,
"grid_template_rows": null,
"height": null,
"justify_content": null,
"justify_items": null,
"left": null,
"margin": null,
"max_height": null,
"max_width": null,
"min_height": null,
"min_width": null,
"object_fit": null,
"object_position": null,
"order": null,
"overflow": null,
"overflow_x": null,
"overflow_y": null,
"padding": null,
"right": null,
"top": null,
"visibility": null,
"width": null
}
},
"835ee2070c05443c9fa3ae6fece3fb44": {
"model_module": "@jupyter-widgets/base",
"model_module_version": "1.2.0",
"model_name": "LayoutModel",
"state": {
"_model_module": "@jupyter-widgets/base",
"_model_module_version": "1.2.0",
"_model_name": "LayoutModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/base",
"_view_module_version": "1.2.0",
"_view_name": "LayoutView",
"align_content": null,
"align_items": null,
"align_self": null,
"border": null,
"bottom": null,
"display": null,
"flex": null,
"flex_flow": null,
"grid_area": null,
"grid_auto_columns": null,
"grid_auto_flow": null,
"grid_auto_rows": null,
"grid_column": null,
"grid_gap": null,
"grid_row": null,
"grid_template_areas": null,
"grid_template_columns": null,
"grid_template_rows": null,
"height": null,
"justify_content": null,
"justify_items": null,
"left": null,
"margin": null,
"max_height": null,
"max_width": null,
"min_height": null,
"min_width": null,
"object_fit": null,
"object_position": null,
"order": null,
"overflow": null,
"overflow_x": null,
"overflow_y": null,
"padding": null,
"right": null,
"top": null,
"visibility": null,
"width": null
}
},
"8b76924bdac749d5b4908678fd0a7497": {
"model_module": "@jupyter-widgets/controls",
"model_module_version": "1.5.0",
"model_name": "HBoxModel",
"state": {
"_dom_classes": [],
"_model_module": "@jupyter-widgets/controls",
"_model_module_version": "1.5.0",
"_model_name": "HBoxModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/controls",
"_view_module_version": "1.5.0",
"_view_name": "HBoxView",
"box_style": "",
"children": [
"IPY_MODEL_3ba9fe91731f4ff380cda2183ad72c4d",
"IPY_MODEL_4e01b9cceffb4d178e97f1d1b36d62b0"
],
"layout": "IPY_MODEL_b03bd083afc84f52834d4a124e57c442"
}
},
"9e6646f7066341ef9619f5f7e6f7c8b8": {
"model_module": "@jupyter-widgets/controls",
"model_module_version": "1.5.0",
"model_name": "FloatProgressModel",
"state": {
"_dom_classes": [],
"_model_module": "@jupyter-widgets/controls",
"_model_module_version": "1.5.0",
"_model_name": "FloatProgressModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/controls",
"_view_module_version": "1.5.0",
"_view_name": "ProgressView",
"bar_style": "success",
"description": "Downloading: 100%",
"description_tooltip": null,
"layout": "IPY_MODEL_39aed05f694a495abd03aa00ff71efb3",
"max": 1042301,
"min": 0,
"orientation": "horizontal",
"style": "IPY_MODEL_f9fa7c1bc9f04df09e487b16761d6b8f",
"value": 1042301
}
},
"b03bd083afc84f52834d4a124e57c442": {
"model_module": "@jupyter-widgets/base",
"model_module_version": "1.2.0",
"model_name": "LayoutModel",
"state": {
"_model_module": "@jupyter-widgets/base",
"_model_module_version": "1.2.0",
"_model_name": "LayoutModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/base",
"_view_module_version": "1.2.0",
"_view_name": "LayoutView",
"align_content": null,
"align_items": null,
"align_self": null,
"border": null,
"bottom": null,
"display": null,
"flex": null,
"flex_flow": null,
"grid_area": null,
"grid_auto_columns": null,
"grid_auto_flow": null,
"grid_auto_rows": null,
"grid_column": null,
"grid_gap": null,
"grid_row": null,
"grid_template_areas": null,
"grid_template_columns": null,
"grid_template_rows": null,
"height": null,
"justify_content": null,
"justify_items": null,
"left": null,
"margin": null,
"max_height": null,
"max_width": null,
"min_height": null,
"min_width": null,
"object_fit": null,
"object_position": null,
"order": null,
"overflow": null,
"overflow_x": null,
"overflow_y": null,
"padding": null,
"right": null,
"top": null,
"visibility": null,
"width": null
}
},
"b3ac7912f51746afae20616269dc69f0": {
"model_module": "@jupyter-widgets/controls",
"model_module_version": "1.5.0",
"model_name": "ProgressStyleModel",
"state": {
"_model_module": "@jupyter-widgets/controls",
"_model_module_version": "1.5.0",
"_model_name": "ProgressStyleModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/base",
"_view_module_version": "1.2.0",
"_view_name": "StyleView",
"bar_color": null,
"description_width": "initial"
}
},
"bb11a4fbf16f488082aa2bac7443fd62": {
"model_module": "@jupyter-widgets/controls",
"model_module_version": "1.5.0",
"model_name": "FloatProgressModel",
"state": {
"_dom_classes": [],
"_model_module": "@jupyter-widgets/controls",
"_model_module_version": "1.5.0",
"_model_name": "FloatProgressModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/controls",
"_view_module_version": "1.5.0",
"_view_name": "ProgressView",
"bar_style": "success",
"description": "Downloading: 100%",
"description_tooltip": null,
"layout": "IPY_MODEL_74672c43909444c592955afd0d727a28",
"max": 1355256,
"min": 0,
"orientation": "horizontal",
"style": "IPY_MODEL_f3883c133c62438db114c22cfc94e351",
"value": 1355256
}
},
"bb7410ff9cb4484da73302993427b90a": {
"model_module": "@jupyter-widgets/base",
"model_module_version": "1.2.0",
"model_name": "LayoutModel",
"state": {
"_model_module": "@jupyter-widgets/base",
"_model_module_version": "1.2.0",
"_model_name": "LayoutModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/base",
"_view_module_version": "1.2.0",
"_view_name": "LayoutView",
"align_content": null,
"align_items": null,
"align_self": null,
"border": null,
"bottom": null,
"display": null,
"flex": null,
"flex_flow": null,
"grid_area": null,
"grid_auto_columns": null,
"grid_auto_flow": null,
"grid_auto_rows": null,
"grid_column": null,
"grid_gap": null,
"grid_row": null,
"grid_template_areas": null,
"grid_template_columns": null,
"grid_template_rows": null,
"height": null,
"justify_content": null,
"justify_items": null,
"left": null,
"margin": null,
"max_height": null,
"max_width": null,
"min_height": null,
"min_width": null,
"object_fit": null,
"object_position": null,
"order": null,
"overflow": null,
"overflow_x": null,
"overflow_y": null,
"padding": null,
"right": null,
"top": null,
"visibility": null,
"width": null
}
},
"c31661c9a1ff4bd681cc10d70fd287c0": {
"model_module": "@jupyter-widgets/base",
"model_module_version": "1.2.0",
"model_name": "LayoutModel",
"state": {
"_model_module": "@jupyter-widgets/base",
"_model_module_version": "1.2.0",
"_model_name": "LayoutModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/base",
"_view_module_version": "1.2.0",
"_view_name": "LayoutView",
"align_content": null,
"align_items": null,
"align_self": null,
"border": null,
"bottom": null,
"display": null,
"flex": null,
"flex_flow": null,
"grid_area": null,
"grid_auto_columns": null,
"grid_auto_flow": null,
"grid_auto_rows": null,
"grid_column": null,
"grid_gap": null,
"grid_row": null,
"grid_template_areas": null,
"grid_template_columns": null,
"grid_template_rows": null,
"height": null,
"justify_content": null,
"justify_items": null,
"left": null,
"margin": null,
"max_height": null,
"max_width": null,
"min_height": null,
"min_width": null,
"object_fit": null,
"object_position": null,
"order": null,
"overflow": null,
"overflow_x": null,
"overflow_y": null,
"padding": null,
"right": null,
"top": null,
"visibility": null,
"width": null
}
},
"debf4a0a609a4cffa5cd686a0a23f6e2": {
"model_module": "@jupyter-widgets/base",
"model_module_version": "1.2.0",
"model_name": "LayoutModel",
"state": {
"_model_module": "@jupyter-widgets/base",
"_model_module_version": "1.2.0",
"_model_name": "LayoutModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/base",
"_view_module_version": "1.2.0",
"_view_name": "LayoutView",
"align_content": null,
"align_items": null,
"align_self": null,
"border": null,
"bottom": null,
"display": null,
"flex": null,
"flex_flow": null,
"grid_area": null,
"grid_auto_columns": null,
"grid_auto_flow": null,
"grid_auto_rows": null,
"grid_column": null,
"grid_gap": null,
"grid_row": null,
"grid_template_areas": null,
"grid_template_columns": null,
"grid_template_rows": null,
"height": null,
"justify_content": null,
"justify_items": null,
"left": null,
"margin": null,
"max_height": null,
"max_width": null,
"min_height": null,
"min_width": null,
"object_fit": null,
"object_position": null,
"order": null,
"overflow": null,
"overflow_x": null,
"overflow_y": null,
"padding": null,
"right": null,
"top": null,
"visibility": null,
"width": null
}
},
"f3883c133c62438db114c22cfc94e351": {
"model_module": "@jupyter-widgets/controls",
"model_module_version": "1.5.0",
"model_name": "ProgressStyleModel",
"state": {
"_model_module": "@jupyter-widgets/controls",
"_model_module_version": "1.5.0",
"_model_name": "ProgressStyleModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/base",
"_view_module_version": "1.2.0",
"_view_name": "StyleView",
"bar_color": null,
"description_width": "initial"
}
},
"f9fa7c1bc9f04df09e487b16761d6b8f": {
"model_module": "@jupyter-widgets/controls",
"model_module_version": "1.5.0",
"model_name": "ProgressStyleModel",
"state": {
"_model_module": "@jupyter-widgets/controls",
"_model_module_version": "1.5.0",
"_model_name": "ProgressStyleModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/base",
"_view_module_version": "1.2.0",
"_view_name": "StyleView",
"bar_color": null,
"description_width": "initial"
}
},
"fb1263f5a3da42e6acab0e9a66551b32": {
"model_module": "@jupyter-widgets/base",
"model_module_version": "1.2.0",
"model_name": "LayoutModel",
"state": {
"_model_module": "@jupyter-widgets/base",
"_model_module_version": "1.2.0",
"_model_name": "LayoutModel",
"_view_count": null,
"_view_module": "@jupyter-widgets/base",
"_view_module_version": "1.2.0",
"_view_name": "LayoutView",
"align_content": null,
"align_items": null,
"align_self": null,
"border": null,
"bottom": null,
"display": null,
"flex": null,
"flex_flow": null,
"grid_area": null,
"grid_auto_columns": null,
"grid_auto_flow": null,
"grid_auto_rows": null,
"grid_column": null,
"grid_gap": null,
"grid_row": null,
"grid_template_areas": null,
"grid_template_columns": null,
"grid_template_rows": null,
"height": null,
"justify_content": null,
"justify_items": null,
"left": null,
"margin": null,
"max_height": null,
"max_width": null,
"min_height": null,
"min_width": null,
"object_fit": null,
"object_position": null,
"order": null,
"overflow": null,
"overflow_x": null,
"overflow_y": null,
"padding": null,
"right": null,
"top": null,
"visibility": null,
"width": null
}
}
}
}
},
"nbformat": 4,
"nbformat_minor": 0
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment