KoboldAI-Client/colab/TPU.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/KoboldAI/KoboldAI-Client/blob/main/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/koboldai/KoboldAI-Client/blob/main/colab/GPU.ipynb)**!\n",
        "\n",
        "---\n",
        "## How to load KoboldAI: Everything you need to know\n",
        "1. On a phone? First put your browser in desktop mode because of a Google Colab bug. Otherwise nothing will happen when you click the play button. Then tap the play button next to \"<-- Tap This if you play on Mobile\", you will see an audio player. Keep the audio player playing so Colab does not get shut down in the background.\n",
        "2. Select the desired model, you will find a description of all the available models further down the page.\n",
        "3. Click the play button next to \"<-- Select your model below and then click this to start KoboldAI\".\n",
        "4. Got a message saying no accelerator is available? Click cancel, and try again in a few minutes. If you do not manage to get a session when you frequently try again try at a different time of day, colab can be busy or your priority may have been lowered by frequent usage.\n",
        "5. After everything is done loading you will get a link that you can use to open KoboldAI. In case of Localtunnel you will also be warned that some people are abusing Localtunnel for phishing, once you acknowledge this warning you will be taken to KoboldAI's interface. If you picked Cloudflare and get a 1033 error refresh the error page after waiting one minute.\n",
        "\n",
        "---\n",
        "\n",
        "Further down the page you can find descriptions of the models, and tips to get the most out of your Google Colab experience.\n",
        "\n",
        "Make sure to keep this page open while you are using KoboldAI, and check back regularly to see if you got a Captcha. Failure to complete the captcha's in time can result in termination of your session or a lower priority towards the TPUs.\n",
        "\n",
        "Firefox users need to disable the enhanced tracking protection or use a different browser in order to be able to use Google Colab without errors (This is not something we can do anything about, the cookie blocker breaks the Google Drive integration because it uses different domains)."
      ],
      "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><-- Select your model below and then click this to start KoboldAI</b>\n",
        "#@markdown You can find a description of the models below along with instructions on how to start KoboldAI.\n",
        "\n",
        "Model = \"Nerys 13B V2\" #@param [\"Nerys 13B V2\", \"Janeway 13B\", \"Shinen 13B\", \"Skein 6B\", \"Janeway 6B\", \"Adventure 6B\", \"Shinen 6B\", \"Lit 6B\", \"NeoX 20B\", \"OPT 13B\", \"Fairseq Dense 13B\", \"GPT-J-6B\"] {allow-input: true}\n",
        "Version = \"Official\" #@param [\"Official\", \"United\"] {allow-input: true}\n",
        "Provider = \"Cloudflare\" #@param [\"Localtunnel\", \"Cloudflare\"]\n",
        "\n",
        "import os\n",
        "try:\n",
        "    device_name = os.environ['COLAB_TPU_ADDR']\n",
        "    TPU_ADDRESS = 'grpc://' + device_name\n",
        "    print('Found TPU at: {}'.format(TPU_ADDRESS))\n",
        "except KeyError:\n",
        "    raise RuntimeError(\"⚠️You can not run this notebook without the TPU accelerator, go to Runtime->Sessions, terminate your session and then try again.⚠️\")\n",
        "print('Now we will need your Google Drive to store settings and saves, you must login with the same account you used for Colab.')\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive/')\n",
        "\n",
        "if Model == \"Janeway 13B\":\n",
        "  Model = \"KoboldAI/fairseq-dense-13B-Janeway\"\n",
        "  path = \"\"\n",
        "  download = \"\"\n",
        "elif Model == \"Nerys 13B V2\":\n",
        "  Model = \"KoboldAI/fairseq-dense-13B-Nerys-v2\"\n",
        "  path = \"\"\n",
        "  download = \"\"\n",
        "elif Model == \"Shinen 13B\":\n",
        "  Model = \"KoboldAI/fairseq-dense-13B-Shinen\"\n",
        "  path = \"\"\n",
        "  download = \"\"\n",
        "elif Model == \"NeoX 20B\":\n",
        "  Model = \"EleutherAI/gpt-neox-20b\"\n",
        "  path = \"\"\n",
        "  download = \"\"\n",
        "elif Model == \"Skein 6B\":\n",
        "  Model = \"KoboldAI/GPT-J-6B-Skein\"\n",
        "  path = \"\"\n",
        "  download = \"\"\n",
        "elif Model == \"Janeway 6B\":\n",
        "  Model = \"KoboldAI/GPT-J-6B-Janeway\"\n",
        "  path = \"\"\n",
        "  download = \"\"\n",
        "elif Model == \"Adventure 6B\":\n",
        "  Model = \"KoboldAI/GPT-J-6B-Adventure\"\n",
        "  path = \"\"\n",
        "  download = \"\"\n",
        "elif Model == \"Lit 6B\":\n",
        "  Model = \"hakurei/lit-6B\"\n",
        "  path = \"\"\n",
        "  download = \"\"\n",
        "elif Model == \"Shinen 6B\":\n",
        "  Model = \"KoboldAI/GPT-J-6B-Shinen\"\n",
        "  path = \"\"\n",
        "  download = \"\"\n",
        "elif Model == \"OPT 13B\":\n",
        "  Model = \"facebook/opt-13b\"\n",
        "  path = \"\"\n",
        "  download = \"\"\n",
        "elif Model == \"Fairseq Dense 13B\":\n",
        "  Model = \"KoboldAI/fairseq-dense-13B\"\n",
        "  path = \"\"\n",
        "  download = \"\"\n",
        "elif Model == \"GPT-J-6B\":\n",
        "  Model = \"EleutherAI/gpt-j-6B\"\n",
        "  path = \"\"\n",
        "  download = \"\"\n",        
        "else:\n",
        "  path = \"\"\n",
        "  download = \"\"\n",
        "\n",
        "if Provider == \"Localtunnel\":\n",
        "  tunnel = \"--localtunnel yes\"\n",
        "else:\n",
        "  tunnel = \"\"\n",
        "\n",
        "!wget https://koboldai.org/ckds -O - | bash /dev/stdin $path$download -m $Model -g $Version $tunnel"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# TPU Edition Model Descriptions\n",
        "\n",
        "| Model | Size | Style | Description |\n",
        "| --- | --- | --- | --- |\n",
        "| [Nerys](https://huggingface.co/KoboldAI/fairseq-dense-13B-Nerys) by Mr Seeker | 13B | Novel/Adventure | Nerys is a hybrid model based on Pike (A newer Janeway), on top of the Pike dataset you also get some Light Novels, Adventure mode support and a little bit of Shinen thrown in the mix. The end result is a very diverse model that is heavily biased towards SFW novel writing, but one that can go beyond its novel training and make for an excellent adventure model to. Adventure mode is best played from a second person perspective, but can be played in first or third person as well. Novel writing can be done best from the first or third person. |\n",
        "| [Janeway](https://huggingface.co/KoboldAI/fairseq-dense-13B-Janeway) by Mr Seeker | 13B | Novel | Janeway is a model created from Picard's dataset combined with a brand new collection of ebooks. This model is trained on 20% more content than Picard and has been trained on literature from various genres. Although the model is mainly focussed on SFW, romantic scenes might involve a degree of nudity. |\n",
        "| [Shinen](https://huggingface.co/KoboldAI/fairseq-dense-13B-Shinen) by Mr Seeker | 13B | NSFW | Shinen is an NSFW model designed to be more explicit. Trained on a variety of stories from the website Sexstories it contains many different kinks. |\n",
        "| [Skein](https://huggingface.co/KoboldAI/GPT-J-6B-Skein) by VE\\_FORBRYDERNE | 6B | Adventure | Skein is best used with Adventure mode enabled, it consists of a 4 times larger adventure dataset than the Adventure model making it excellent for text adventure gaming. On top of that it also consists of light novel training further expanding its knowledge and writing capabilities. It can be used with the You filter bias if you wish to write Novels with it, but dedicated Novel models can perform better for this task. |\n",
        "| [Adventure](https://huggingface.co/KoboldAI/GPT-J-6B-Adventure) by VE\\_FORBRYDERNE | 6B | Adventure | Adventure is a 6B model designed to mimick the behavior of AI Dungeon. It is exclusively for Adventure Mode and can take you on the epic and wackey adventures that AI Dungeon players love. It also features the many tropes of AI Dungeon as it has been trained on very similar data. It must be used in second person (You). |\n",
        "| [Lit](https://huggingface.co/hakurei/lit-6B) by Haru | 6B | NSFW | Lit is a great NSFW model trained by Haru on both a large set of Literotica stories and high quality novels along with tagging support. Creating a high quality model for your NSFW stories. This model is exclusively a novel model and is best used in third person. |\n",
        "| Neo(X) by EleutherAI | 20B | Generic | NeoX is the largest EleutherAI model currently available, being a generic model it is not particularly trained towards anything and can do a variety of writing, Q&A and coding tasks. 20B's performance is closely compared to the 13B models and it is worth trying both especially if you have a task that does not involve english writing. Its behavior will be similar to the GPT-J-6B model since they are trained on the same dataset but with more sensitivity towards repetition penalty and with more knowledge. |\n",
        "| [Fairseq Dense](https://huggingface.co/KoboldAI/fairseq-dense-13B) | 13B | Generic | Trained by Facebook Researchers this model stems from the MOE research project within Fairseq. This particular version has been converted by us for use in KoboldAI. It is known to be on par with the larger 20B model from EleutherAI and considered as better for pop culture and language tasks. Because the model has never seen a new line (enter) it may perform worse on formatting and paragraphing. |\n",
        "| [GPT-J-6B](https://huggingface.co/EleutherAI/gpt-j-6B) by EleutherAI | 6B | Generic | This model serves as the basis for most other 6B models (Some being based on Fairseq Dense instead). Being trained on the Pile and not biased towards anything in particular it is suitable for a variety of tasks such as writing, Q&A and coding tasks. You will likely get better result with larger generic models or finetuned models. |\n",
        "\n",
        "\n",
        "# [GPU Edition Model Descriptions](https://colab.research.google.com/github/KoboldAI/KoboldAI-Client/blob/main/colab/GPU.ipynb)\n",
        "\n",
        "| Model | Size | Style | Description |\n",
        "| --- | --- | --- | --- |\n",
        "| [Nerys 2.7B](https://huggingface.co/KoboldAI/fairseq-dense-2.7B-Nerys) by Mr Seeker | 2.7B | Novel/Adventure | Nerys is a hybrid model based on Pike (A newer Janeway), on top of the Pike dataset you also get some Light Novels, Adventure mode support and a little bit of Shinen thrown in the mix. The end result is a very diverse model that is heavily biased towards SFW novel writing, but one that can go beyond its novel training and make for an excellent adventure model to. Adventure mode is best played from a second person perspective, but can be played in first or third person as well. Novel writing can be done best from the first or third person. |\n",
        "| [Janeway 2.7B](https://huggingface.co/KoboldAI/GPT-Neo-2.7B-Janeway) by Mr Seeker | 2.7B | Novel | Janeway is a model created from Picard's dataset combined with a brand new collection of ebooks. This model is trained on 20% more content than Picard and has been trained on literature from various genres. Although the model is mainly focussed on SFW, romantic scenes might involve a degree of nudity. |\n",
        "| [Picard 2.7B](https://huggingface.co/KoboldAI/GPT-Neo-2.7B-Picard) by Mr Seeker | 2.7B | Novel | Picard is a model trained for SFW Novels based on Neo 2.7B. It is focused on Novel style writing without the NSFW bias. While the name suggests a sci-fi model this model is designed for Novels of a variety of genre's. It is meant to be used in KoboldAI's regular mode. |\n",
        "| [AID 2.7B](https://huggingface.co/KoboldAI/GPT-Neo-2.7B-AID) by melastacho | 2.7B | Adventure | Also know as Adventure 2.7B this is a clone of the AI Dungeon Classic model and is best known for the epic wackey adventures that AI Dungeon Classic players love. |\n",
        "| [Horni LN 2.7B](https://huggingface.co/KoboldAI/GPT-Neo-2.7B-Horni-LN) by finetune | 2.7B | Novel | This model is based on Horni 2.7B and retains its NSFW knowledge, but was then further biased towards SFW novel stories. If you seek a balance between a SFW Novel model and a NSFW model this model should be a good choice. |\n",
        "| [Horni 2.7B](https://huggingface.co/KoboldAI/GPT-Neo-2.7B-Horni) by finetune | 2.7B | NSFW | This model is tuned on Literotica to produce a Novel style model biased towards NSFW content. Can still be used for SFW stories but will have a bias towards NSFW content. It is meant to be used in KoboldAI's regular mode. |\n",
        "| [Shinen 2.7B](https://huggingface.co/KoboldAI/GPT-Neo-2.7B-Shinen) by Mr Seeker | 2.7B | NSFW | Shinen is an alternative to the Horni model designed to be more explicit. If Horni is to tame for you Shinen might produce better results. While it is a Novel model it is unsuitable for SFW stories due to its heavy NSFW bias. Shinen will not hold back. It is meant to be used in KoboldAI's regular mode. |\n",
        "| [Neo 2.7B](https://huggingface.co/EleutherAI/gpt-neo-2.7B) by EleutherAI | 2.7B | Generic | This is the base model for all the other 2.7B models, it is best used when you have a use case that we have no other models available for, such as writing blog articles or programming. It can also be a good basis for the experience of some of the softprompts if your softprompt is not about a subject the other models cover. |\n",
        "\n",
        "| Style | Description |\n",
        "| --- | --- |\n",
        "| Novel | For regular story writing, not compatible with Adventure mode or other specialty modes. |\n",
        "| NSFW | Indicates that the model is strongly biased towards NSFW content and is not suitable for children, work environments or livestreaming. Most NSFW models are also Novel models in nature. |\n",
        "| Adventure | These models are excellent for people willing to play KoboldAI like a Text Adventure game and are meant to be used with Adventure mode enabled. Even if you wish to use it as a Novel style model you should always have Adventure mode on and set it to story. These models typically have a strong bias towards the use of the word You and without Adventure mode enabled break the story flow and write actions on your behalf. |\n",
        "| Generic | Generic models are not trained towards anything specific, typically used as a basis for other tasks and models. They can do everything the other models can do, but require much more handholding to work properly. Generic models are an ideal basis for tasks that we have no specific model for, or for experiencing a softprompt in its raw form. |\n",
        "\n",
        "---\n",
        "# Tips to get the most out of Google Colab\n",
        "-  Google will occationally show a Captcha, typically after it has been open for 30 minutes but it can be more frequent if you often use Colab. Make sure to do these properly, or you risk getting your instance shut down and getting a lower priority towards the TPU's.\n",
        "- KoboldAI uses Google Drive to store your files and settings, if you wish to upload a softprompt or userscript this can be done directly on the Google Drive website. You can also use this to download backups of your KoboldAI related files or upload models of your own.\n",
        "- Don't want to save your stories on Google Drive for privacy reasons? Do not use KoboldAI's save function and instead click Download as .json, this will automatically download the story to your own computer without ever touching Google's harddrives. You can load this back trough the Load from file option.\n",
        "- Google shut your instance down unexpectedly? You can still make use of the Download as .json button to recover your story as long as you did not close the KoboldAI window. You can then load this back up in your next session.\n",
        "- Done with KoboldAI? Go to the Runtime menu, click on Manage Sessions and terminate your open sessions that you no longer need. This trick can help you maintain higher priority towards getting a TPU."
      ],
      "metadata": {
        "id": "i0-9ARA3c4Fx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title <b>Model Cleaner</b>\n",
        "#@markdown Out of space? Run this to remove all cached models (Google Drive models are not effected).\n",
        "!rm /content/KoboldAI-Client/cache/*\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "QQZSmoNol04V"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "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
}