{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d-Yihz3hAb2E"
      },
      "source": [
        "https://colab.research.google.com/github/TavernAI/TavernAI/blob/main/colab/GPU.ipynb<br>\n",
        "\n",
        "Works with:<br>\n",
        "KoboldAI https://github.com/KoboldAI/KoboldAI-Client<br>\n",
        "Pygmalion https://huggingface.co/PygmalionAI/<br>\n",
        "<br>\n",
        "**Links**<br>\n",
        "TavernAI Github https://github.com/TavernAI/TavernAI<br>\n",
        "TavernAI Discord https://discord.gg/zmK2gmr45t<br>\n",
        "TavernAI Boosty https://boosty.to/tavernai\n",
        "<pre>\n",
        " Tavern.AI/ \\ /  ^   ^ ^ ^    ~~~~ ^ \\     /  ^ ^   ^ ^/ ^  ^ \\/^  ^    \\\n",
        "         /^ ^\\ ^  ^ ^   ^ ^  ~~   ^   \\   /  ^  ^ ^   / ^ ^  ^/   ^ ^    \\\n",
        "        /^ ^ ^\\^   ^ ^ ^   _||____   ^ \\ /  ^  ^ ^   /       /  ^  ^  ^   \\\n",
        " /\\ /\\ /\\   ^  \\  /\\ /\\   /\\\\\\\\\\\\\\\\   ^ \\  ^ /\\ /\\ /\\   /\\ /\\ /\\  ^ ^  ^/\\\n",
        "//\\\\/\\\\/\\\\   ^  \\//\\\\/\\\\ /__\\\\\\\\\\\\\\\\  _, \\  //\\\\/\\\\/\\\\ //\\\\/\\\\/\\\\  ^ ^ //\\\\\n",
        "//\\\\/\\\\/\\\\       //\\\\/\\\\ |__|_|_|__|   \\__, //\\\\/\\\\/\\\\ //\\\\/\\\\/\\\\     ///\\\\\\\n",
        " || || (@＾◡＾)(≖ ‸ ≖*) ( ←_← )\\| /|   /\\ \\ヽ(°ㅂ°╬) |( Ψ▼ｰ▼)∈ (O_O; )  |||\n",
        "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ~~~~~ ~~~~~ ~~~~~ ~~~~~  ~~~~~ ~~ \n",
        "</pre>\n",
        "**Launch Instructions**<br>\n",
        "1. Click the launch button.\n",
        "2. Wait for the environment and model to load\n",
        "3. After initialization, a TavernAI link will appear\n",
        "\n",
        "**Faq**<br>\n",
        "* Q: I do not get a TavernAI link\n",
        "* A: It seems the localtunnel service is currently down, so the TavernAI link is unavailable. Need to wait for it to start working again."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hCpoIHxYcDGs"
      },
      "outputs": [],
      "source": [
        "#@title <-- Tap this if you play on Mobile { display-mode: \"form\" }\n",
        "#Taken from KoboldAI colab\n",
        "%%html\n",
        "<b>Press play on the music player to keep the tab alive, then start TavernAI below (Uses only 13MB of data)</b><br/>\n",
        "<audio src=\"https://henk.tech/colabkobold/silence.m4a\" controls>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hps3qtPLFNBb",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title <b>TavernAI</b>\n",
        "#@markdown <- Click For Start (≖ ‸ ≖ ✿)\n",
        "\n",
        "Model = \"Pygmalion 6B\" #@param [ \"Pygmalion 6B\", \"Pygmalion 6B Dev\"] {allow-input: true}\n",
        "Version = \"Official\" \n",
        "KoboldAI_Provider = \"Localtunnel\" #@param [\"Localtunnel\", \"Cloudflare\"]\n",
        "use_google_drive = True #@param {type:\"boolean\"}\n",
        "Provider = KoboldAI_Provider\n",
        "!nvidia-smi\n",
        "import subprocess\n",
        "import time\n",
        "import sys\n",
        "import os\n",
        "import threading\n",
        "import shutil\n",
        "from google.colab import drive\n",
        "\n",
        "    \n",
        "if use_google_drive:\n",
        "  drive.mount('/content/drive/')\n",
        "  if not os.path.exists(\"/content/drive/MyDrive/TavernAI/\"):\n",
        "    os.mkdir(\"/content/drive/MyDrive/TavernAI/\")\n",
        "  if not os.path.exists(\"/content/drive/MyDrive/TavernAI/characters/\"):\n",
        "    os.mkdir(\"/content/drive/MyDrive/TavernAI/characters/\")\n",
        "  if not os.path.exists(\"/content/drive/MyDrive/TavernAI/chats/\"):\n",
        "    os.mkdir(\"/content/drive/MyDrive/TavernAI/chats/\")\n",
        "else:\n",
        "  if not os.path.exists(\"/content/drive\"):\n",
        "    os.mkdir(\"/content/drive\")\n",
        "  if not os.path.exists(\"/content/drive/MyDrive/\"):\n",
        "    os.mkdir(\"/content/drive/MyDrive/\")\n",
        "\n",
        "def copy_characters(use_google_drive=False):\n",
        "  if not use_google_drive:\n",
        "    return\n",
        "  \n",
        "  src_folder = \"/TavernAIColab/public/characters\"\n",
        "  dst_folder = \"/content/drive/MyDrive/TavernAI/characters\"\n",
        "\n",
        "  for filename in os.listdir(src_folder):\n",
        "    src_file = os.path.join(src_folder, filename)\n",
        "    dst_file = os.path.join(dst_folder, filename)\n",
        "\n",
        "    if os.path.exists(dst_file):\n",
        "      print(f\"{dst_file} already exists. Skipping...\")\n",
        "      continue\n",
        "\n",
        "    shutil.copy(src_file, dst_folder)\n",
        "    print(f\"{src_file} copied to {dst_folder}\")\n",
        "Revision = \"\"\n",
        "\n",
        "if Model == \"Pygmalion 6B\":\n",
        "  Model = \"PygmalionAI/pygmalion-6b\"\n",
        "  path = \"\"\n",
        "  download = \"\"\n",
        "  Version = \"United\"\n",
        "elif Model == \"Pygmalion 6B Dev\":\n",
        "  Model = \"PygmalionAI/pygmalion-6b\"\n",
        "  Revision = \"--revision dev\"\n",
        "  path = \"\"\n",
        "  Version = \"United\"\n",
        "  download = \"\"\n",
        "\n",
        "if Provider == \"Localtunnel\":\n",
        "  tunnel = \"--localtunnel yes\"\n",
        "else:\n",
        "  tunnel = \"\"\n",
        "\n",
        "\n",
        "\n",
        "#Henk's KoboldAI script\n",
        "!wget https://koboldai.org/ckds && chmod +x ckds\n",
        "!./ckds --init only\n",
        "if Provider == \"Localtunnel\":\n",
        "  p = subprocess.Popen(['/content/ckds', '--model', Model, '--localtunnel', 'yes'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
        "else:\n",
        "  p = subprocess.Popen(['/content/ckds', '--model', Model], stdout=subprocess.PIPE, stderr=subprocess.PIPE)\n",
        "\n",
        "\n",
        "#Do not repeat! Tricks performed by a professional!\n",
        "url = ''\n",
        "while True:\n",
        "    line = p.stdout.readline().decode().strip()\n",
        "    if \"KoboldAI has finished loading and is available at the following link: \" in line:\n",
        "        print(line)\n",
        "        url = line.split(\"KoboldAI has finished loading and is available at the following link: \")[1]\n",
        "        print(url)\n",
        "        break\n",
        "    if \"KoboldAI has finished loading and is available at the following link for UI 1: \" in line:\n",
        "        print(line)\n",
        "        url = line.split(\"KoboldAI has finished loading and is available at the following link for UI 1: \")[1]\n",
        "        print(url)\n",
        "        break\n",
        "    if not line:\n",
        "        break\n",
        "    print(line)\n",
        "    if \"INIT\" in line and \"Transformers\" in line:\n",
        "      print(\"Model loading... (It will take 2 - 5 minutes)\")\n",
        "\n",
        "\n",
        "#TavernAI\n",
        "%cd /\n",
        "!curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.37.2/install.sh | bash\n",
        "!nvm install 19.1.0\n",
        "!nvm use 19.1.0\n",
        "!node -v\n",
        "!git clone https://github.com/TavernAI/TavernAIColab\n",
        "copy_characters(use_google_drive)\n",
        "%cd TavernAIColab\n",
        "!npm install\n",
        "time.sleep(1)\n",
        "%env colab=2\n",
        "%env colaburl=$url\n",
        "if use_google_drive:\n",
        "  %env googledrive=2\n",
        "!nohup node server.js &\n",
        "time.sleep(3)\n",
        "print('KoboldAI LINK:')\n",
        "print(url)\n",
        "print('')\n",
        "print('###TavernAI LINK###')\n",
        "!lt --port 8000\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}