{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/henk717/KoboldAI/blob/united/colab/TPU.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Welcome to KoboldAI on Google Colab, TPU Edition!\n",
        "KoboldAI is a powerful and easy way to use a variety of AI based text generation experiences. You can use it to write stories, blog posts, play a text adventure game, use it like a chatbot and more! In some cases it might even help you with an assignment or programming task (But always make sure the information the AI mentions is correct, it loves to make stuff up).\n",
        "\n",
        "For more information about KoboldAI check our our Github readme : https://github.com/KoboldAI/KoboldAI-Client/blob/main/readme.md\n",
        "\n",
        "More (smaller) models are available in the **[GPU edition](https://colab.research.google.com/github/henk717/KoboldAI/blob/united/colab/GPU.ipynb)**!"
      ],
      "metadata": {
        "id": "zrLGxVCEaqZx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title <-- Tap this if you play on Mobile { display-mode: \"form\" }\n",
        "%%html\n",
        "<b>Press play on the music player to keep the tab alive, then start KoboldAI below (Uses only 13MB of data)</b><br/>\n",
        "<audio src=\"https://henk.tech/colabkobold/silence.m4a\" controls>"
      ],
      "metadata": {
        "id": "ZIL7itnNaw5V"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qZmAyFFMouk9",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title <b><-- Click this to start KoboldAI</b>\n",
        "Model = \"Adventure 6B\" #@param [\"Skein 6B\", \"Adventure 6B\", \"Lit 6B\", \"Generic 6B\", \"C1 6B\"]\n",
        "Version = \"United\" #@param [\"Official\", \"United\"] {allow-input: true}\n",
        "Drive = \"Unextracted (Less Space)\" #@param [\"Unextracted (Less Space)\", \"Extracted (Faster Loading)\"]\n",
        "#@markdown Extracted models take up more space but load faster the next time you use them, not all models use your Google Drive. See the Model list below for descriptions and space requirements. If your extracted model does not load the next time you try to launch KoboldAI delete the folder from your Google Drive and ensure enough space is available.\n",
        "\n",
        "if Model == \"Skein 6B\":\n",
        "  path = \"gpt-j-6b-skein-jax\"\n",
        "  location = \"colab\"\n",
        "  download = \"-a https://storage.henk.tech/KoboldAI/skein-jax.txt\"\n",
        "  extract = \"\"\n",
        "  Drive = \"Unextracted (Less Space)\"\n",
        "if Model == \"Adventure 6B\":\n",
        "  path = \"gpt-j-6b-adventure-jax\"\n",
        "  location = \"colab\"\n",
        "  download = \"-a https://api.wandb.ai/files/ve-forbryderne/adventure/carol-data/models/gpt-j-6b-adventure-jax/aria2.txt\"\n",
        "  extract = \"\"\n",
        "  Drive = \"Unextracted (Less Space)\"\n",
        "if Model == \"Lit 6B\":\n",
        "  path = \"gpt-j-6b-lit-jax\"\n",
        "  location = \"drive\"\n",
        "  download = \"-a https://storage.henk.tech/KoboldAI/aria2.php?file=gpt-j-6b-lit-jax.7z\"\n",
        "  extract = \"-z gpt-j-6b-lit-jax.7z\"\n",
        "if Model == \"Generic 6B\":\n",
        "  path = \"step_383500\"\n",
        "  location = \"drive\"\n",
        "  download = \"-a https://storage.henk.tech/KoboldAI/aria2.php?file=step_383500_slim.tar.zstd\"\n",
        "  extract = \"-t step_383500_slim.tar.zstd\"\n",
        "if Model == \"C1 6B\":\n",
        "  path = \"gpt-j-6b-c1-jax\"\n",
        "  location = \"drive\"\n",
        "  download = \"-a https://storage.henk.tech/KoboldAI/aria2.php?file=gpt-j-6b-c1-jax\"\n",
        "  extract = \"-z gpt-j-6b-c1-jax.7z\"\n",
        "\n",
        "\n",
        "if Drive == \"Unextracted (Less Space)\":\n",
        "  xloc = \"colab\"\n",
        "if Drive == \"Extracted (Faster Loading)\":\n",
        "  xloc = \"drive\"\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive/')\n",
        "\n",
        "!wget https://henk.tech/ckds -O - | bash /dev/stdin $download -l $location $extract -p $path -m TPUMeshTransformerGPTJ -g $Version -x $xloc"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# MODEL LIST PENDING (WIP COLAB)"
      ],
      "metadata": {
        "id": "i0-9ARA3c4Fx"
      }
    }
  ],
  "metadata": {
    "colab": {
      "name": "ColabKobold TPU",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "TPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}