From cf87ea41991cd3a09216aa1314073e3abd17f1eb Mon Sep 17 00:00:00 2001
From: henk717 <henk717@msn.com>
Date: Tue, 18 Jan 2022 00:58:32 +0100
Subject: [PATCH] Initial TPU Colab

---
 colab/TPU.ipynb | 123 ++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 123 insertions(+)
 create mode 100644 colab/TPU.ipynb

diff --git a/colab/TPU.ipynb b/colab/TPU.ipynb
new file mode 100644
index 00000000..6f36466e
--- /dev/null
+++ b/colab/TPU.ipynb
@@ -0,0 +1,123 @@
+{
+  "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 = \"Skein 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",
+        "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",
+        "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 = \"drive\"\n",
+        "if Drive = \"Extracted (Faster Loading)\"\n",
+        "  xloc = \"colab\"\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": "Untitled2.ipynb",
+      "provenance": [],
+      "private_outputs": true,
+      "collapsed_sections": [],
+      "authorship_tag": "ABX9TyN99Aa/LOA7MugR8iSyacJM",
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "TPU"
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
\ No newline at end of file