<commit_before>{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvInv8aSxy3A"
      },
      "source": [
        "## ViT-Plex Demo\n",
        "\n",
        "*Licensed under the Apache License, Version 2.0.*\n",
        "\n",
        "To run this in a public Colab, change the GitHub link: replace github.com with [githubtocolab.com](http://githubtocolab.com)\n",
        "\n",
        "\u003ca href=\"https://githubtocolab.com/google/uncertainty-baselines/blob/main/experimental/plex/plex_vit_demo.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/\u003e\u003c/a\u003e\n",
        "\n",
        "This notebook demonstrates how one can utilize the released **ViT-Plex** checkpoints from the *Plex: Towards Reliability using Pretrained Large Model Extensions* paper using [JAX](https://jax.readthedocs.io/). The **General usage** section provides a minimal setup for loading the checkpoints and making predictions, and the ***Uncertainty***, ***Robust Generalization***, and ***Adaptation*** sections delve deeper into the three areas of reliability for which Plex is designed to excel.\n",
        "\n",
        "For more advanced usage, full training and fine-tuning scripts can be found at https://github.com/google/uncertainty-baselines/tree/main/baselines/jft."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k5Z7xEm8OEOd"
      },
      "source": [
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7HBkyhA_yYBr"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TDl8fDWONRip"
      },
      "outputs": [],
      "source": [
        "pip_install = True\n",
        "backend = \"cpu\"  #@param [\"tpu\", \"gpu\", \"cpu\"]\n",
        "if pip_install:\n",
        "  # NOTE: Set the jax version to \u003e=0.3.14 if Python 3.9+ is available.\n",
        "  if backend == \"cpu\":\n",
        "    !python3 -m pip install \"jax~=0.2.27\"\n",
        "  elif backend == \"tpu\":\n",
        "    !python3 -m pip install \"jax[tpu]~=0.2.27\" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n",
        "  elif backend == \"gpu\":\n",
        "    !python3 -m pip install \"jax[cuda]~=0.2.27\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
        "  else:\n",
        "    raise ValueError(\"Backend must be one of ['cpu', 'tpu', 'gpu']. got \"\n",
        "                     f\"backend={backend} instead.\")\n",
        "  !python3 -m pip install \"jax[tpu]\u003e=0.3.14\" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n",
        "  !rm -rf uncertainty-baselines\n",
        "  !git clone https://github.com/google/uncertainty-baselines.git\n",
        "  !cp -r uncertainty-baselines/baselines/jft/* .\n",
        "  # NOTE: Remove the explicit tensorflow-federated and tensorflow_probability\n",
        "  # installs if Python 3.9+ is available.\n",
        "  !python3 -m pip install \"tensorflow-federated==0.20.0\" \"tensorflow_probability\u003c0.17.0\" ./uncertainty-baselines[tensorflow,jax,models,datasets]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0yEmRO5rQ338"
      },
      "outputs": [],
      "source": [<commit_msg>Fixes a dependency install bug.<commit_after>{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvInv8aSxy3A"
      },
      "source": [
        "## ViT-Plex Demo\n",
        "\n",
        "*Licensed under the Apache License, Version 2.0.*\n",
        "\n",
        "To run this in a public Colab, change the GitHub link: replace github.com with [githubtocolab.com](http://githubtocolab.com)\n",
        "\n",
        "\u003ca href=\"https://githubtocolab.com/google/uncertainty-baselines/blob/main/experimental/plex/plex_vit_demo.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/\u003e\u003c/a\u003e\n",
        "\n",
        "This notebook demonstrates how one can utilize the released **ViT-Plex** checkpoints from the *Plex: Towards Reliability using Pretrained Large Model Extensions* paper using [JAX](https://jax.readthedocs.io/). The **General usage** section provides a minimal setup for loading the checkpoints and making predictions, and the ***Uncertainty***, ***Robust Generalization***, and ***Adaptation*** sections delve deeper into the three areas of reliability for which Plex is designed to excel.\n",
        "\n",
        "For more advanced usage, full training and fine-tuning scripts can be found at https://github.com/google/uncertainty-baselines/tree/main/baselines/jft."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k5Z7xEm8OEOd"
      },
      "source": [
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7HBkyhA_yYBr"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TDl8fDWONRip"
      },
      "outputs": [],
      "source": [
        "pip_install = True\n",
        "backend = \"cpu\"  #@param [\"tpu\", \"gpu\", \"cpu\"]\n",
        "if pip_install:\n",
        "  # NOTE: Set the jax version to \u003e=0.3.14 if Python 3.9+ is available.\n",
        "  if backend == \"cpu\":\n",
        "    !python3 -m pip install \"jax~=0.2.27\"\n",
        "  elif backend == \"tpu\":\n",
        "    !python3 -m pip install \"jax[tpu]~=0.2.27\" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n",
        "  elif backend == \"gpu\":\n",
        "    !python3 -m pip install \"jax[cuda]~=0.2.27\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
        "  else:\n",
        "    raise ValueError(\"Backend must be one of ['cpu', 'tpu', 'gpu']. got \"\n",
        "                     f\"backend={backend} instead.\")\n",
        "  !rm -rf uncertainty-baselines\n",
        "  !git clone https://github.com/google/uncertainty-baselines.git\n",
        "  !cp -r uncertainty-baselines/baselines/jft/* .\n",
        "  # NOTE: Remove the explicit tensorflow-federated and tensorflow_probability\n",
        "  # installs if Python 3.9+ is available.\n",
        "  !python3 -m pip install \"tensorflow-federated==0.20.0\" \"tensorflow_probability\u003c0.17.0\" ./uncertainty-baselines[tensorflow,jax,models,datasets]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0yEmRO5rQ338"
      },
      "outputs": [],
      "source": [
{'fraction_non_alphanumeric': 0.1755435575660295, 'fraction_numerical': 0.021596381146942945, 'mean_word_length': 2.2904464714354296, 'pattern_counts': {'":': 56, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 15, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4136235', 'n_tokens_mistral': 2423, 'n_tokens_neox': 2285, 'n_words': 501}