Colab-TextGen-GPU.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Alternative PygmalionAI notebook for Google Colab\n",
        "\n",
        "This notebook is an *unofficial* method to run the [PygmalionAI LLMs](https://huggingface.co/PygmalionAI) (NSFW) using **oobabooga**'s [text-generation-webui](https://github.com/oobabooga/text-generation-webui) project in chat mode.\n",
        "\n",
        "Run all the cells and a public gradio URL will appear at the bottom in around 5 minutes.\n",
        "\n",
        "## Parameters\n",
        "\n",
        "* **save_logs_to_google_drive**: saves your chat logs, characters, and softprompts to Google Drive automatically, so that they will persist across sessions.\n",
        "* **text_streaming**: streams the text output in real time instead of waiting for the full response to be completed.\n",
        "* **cai_chat**: makes the interface look like Character.AI. Otherwise, it looks like a standard WhatsApp-like chat.\n",
        "* **load_in_8bit**: loads the model with 8-bit precision, reducing the GPU memory usage by half. This allows you to use the full 2048 prompt length without running out of memory, at a small accuracy and speed cost.\n",
        "* **activate_silero_text_to_speech**: responses will be audios instead of text. There are 118 voices available (`en_0` to `en_117`), which can be set in the \"Extensions\" tab of the interface. You can find samples here: [Silero samples](https://oobabooga.github.io/silero-samples/).\n",
        "* **activate_sending_pictures**: adds a menu for sending pictures to the bot, which are automatically captioned using BLIP.\n",
        "* **activate_character_bias**: an extension that adds an user-defined, hidden string at the beginning of the bot's reply with the goal of biasing the rest of the response.\n",
        "* **chat_language**: if different than English, activates automatic translation using Google Translate, allowing you to communicate with the bot in a different language.\n",
        "\n",
        "## Characters\n",
        "\n",
        "You can use the following websites to create characters compatible with this web UI:\n",
        "\n",
        "* [Pygmalion JSON character creator](https://oobabooga.github.io/character-creator.html)\n",
        "* [AI Character Editor](https://zoltanai.github.io/character-editor/)\n",
        "\n",
        "## Credits\n",
        "\n",
        "Based on the [original notebook by 81300](https://colab.research.google.com/github/81300/AI-Notebooks/blob/main/Colab-TextGen-GPU.ipynb).\n"
      ],
      "metadata": {
        "id": "MFQl6-FjSYtY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 1. Keep this tab alive to prevent Colab from disconnecting you { display-mode: \"form\" }\n",
        "\n",
        "# From: https://github.com/henk717/KoboldAI\n",
        "\n",
        "%%html\n",
        "<b>Press play on the music player to keep the tab alive (Uses only 13MB of data)</b><br/>\n",
        "<audio src=\"https://henk.tech/colabkobold/silence.m4a\" controls>"
      ],
      "metadata": {
        "id": "f7TVVj_z4flw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 2. Install the requirements\n",
        "#@markdown Optionally, you can save the chat logs to Google Drive so that they will persist across sessions.\n",
        "\n",
        "#@markdown * The characters and softprompts will be saved as well. \n",
        "\n",
        "save_logs_to_google_drive = False #@param {type:\"boolean\"}\n",
        "\n",
        "if save_logs_to_google_drive:\n",
        "  import os\n",
        "  import shutil\n",
        "  from google.colab import drive\n",
        "  drive.mount('/content/drive')\n",
        "  base_folder = '/content/drive/MyDrive'\n",
        "\n",
        "%cd /content\n",
        "!git clone https://github.com/oobabooga/text-generation-webui\n",
        "if save_logs_to_google_drive:\n",
        "  if not os.path.exists(f\"{base_folder}/oobabooga-data\"):\n",
        "    os.mkdir(f\"{base_folder}/oobabooga-data\")\n",
        "  if not os.path.exists(f\"{base_folder}/oobabooga-data/logs\"):\n",
        "    os.mkdir(f\"{base_folder}/oobabooga-data/logs\")\n",
        "  if not os.path.exists(f\"{base_folder}/oobabooga-data/softprompts\"):\n",
        "    os.mkdir(f\"{base_folder}/oobabooga-data/softprompts\")\n",
        "  if not os.path.exists(f\"{base_folder}/oobabooga-data/characters\"):\n",
        "    shutil.move(\"text-generation-webui/characters\", f\"{base_folder}/oobabooga-data/characters\")\n",
        "  else:\n",
        "    !rm -r \"text-generation-webui/characters\"\n",
        "    \n",
        "  !rm -r \"text-generation-webui/softprompts\"\n",
        "  !ln -s \"$base_folder/oobabooga-data/logs\" \"text-generation-webui/logs\"\n",
        "  !ln -s \"$base_folder/oobabooga-data/softprompts\" \"text-generation-webui/softprompts\"\n",
        "  !ln -s \"$base_folder/oobabooga-data/characters\" \"text-generation-webui/characters\"\n",
        "\n",
        "else:\n",
        "  !mkdir text-generation-webui/logs\n",
        "\n",
        "!ln -s text-generation-webui/logs .\n",
        "!ln -s text-generation-webui/characters .\n",
        "!ln -s text-generation-webui/models .\n",
        "%rm -r sample_data\n",
        "%cd text-generation-webui\n",
        "!wget https://oobabooga.github.io/settings-colab.json -O settings-colab-template.json\n",
        "\n",
        "# Install requirements\n",
        "!pip install -r requirements.txt\n",
        "!pip install -r extensions/google_translate/requirements.txt\n",
        "!pip install -r extensions/silero_tts/requirements.txt"
      ],
      "metadata": {
        "id": "LGQ8BiMuXMDG",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 3. Start the web UI\n",
        "\n",
        "import json\n",
        "\n",
        "# Parameters\n",
        "model = \"Pygmalion 6B original (sharded, rehosted)\" #@param [\"Pygmalion 6B original (sharded, rehosted)\", \"Pygmalion 6B main (sharded, rehosted)\", \"Pygmalion 6B dev (sharded, rehosted)\", \"Pygmalion 350m (for debugging)\"] {allow-input: false}\n",
        "text_streaming = False #@param {type:\"boolean\"}\n",
        "cai_chat = True #@param {type:\"boolean\"}\n",
        "load_in_8bit = False #@param {type:\"boolean\"}\n",
        "activate_silero_text_to_speech = False #@param {type:\"boolean\"}\n",
        "activate_sending_pictures = False #@param {type:\"boolean\"}\n",
        "activate_character_bias = False #@param {type:\"boolean\"}\n",
        "chat_language = \"English\" # @param ['Afrikaans', 'Albanian', 'Amharic', 'Arabic', 'Armenian', 'Azerbaijani', 'Basque', 'Belarusian', 'Bengali', 'Bosnian', 'Bulgarian', 'Catalan', 'Cebuano', 'Chinese (Simplified)', 'Chinese (Traditional)', 'Corsican', 'Croatian', 'Czech', 'Danish', 'Dutch', 'English', 'Esperanto', 'Estonian', 'Finnish', 'French', 'Frisian', 'Galician', 'Georgian', 'German', 'Greek', 'Gujarati', 'Haitian Creole', 'Hausa', 'Hawaiian', 'Hebrew', 'Hindi', 'Hmong', 'Hungarian', 'Icelandic', 'Igbo', 'Indonesian', 'Irish', 'Italian', 'Japanese', 'Javanese', 'Kannada', 'Kazakh', 'Khmer', 'Korean', 'Kurdish', 'Kyrgyz', 'Lao', 'Latin', 'Latvian', 'Lithuanian', 'Luxembourgish', 'Macedonian', 'Malagasy', 'Malay', 'Malayalam', 'Maltese', 'Maori', 'Marathi', 'Mongolian', 'Myanmar (Burmese)', 'Nepali', 'Norwegian', 'Nyanja (Chichewa)', 'Pashto', 'Persian', 'Polish', 'Portuguese (Portugal, Brazil)', 'Punjabi', 'Romanian', 'Russian', 'Samoan', 'Scots Gaelic', 'Serbian', 'Sesotho', 'Shona', 'Sindhi', 'Sinhala (Sinhalese)', 'Slovak', 'Slovenian', 'Somali', 'Spanish', 'Sundanese', 'Swahili', 'Swedish', 'Tagalog (Filipino)', 'Tajik', 'Tamil', 'Telugu', 'Thai', 'Turkish', 'Ukrainian', 'Urdu', 'Uzbek', 'Vietnamese', 'Welsh', 'Xhosa', 'Yiddish', 'Yoruba', 'Zulu']\n",
        "\n",
        "activate_google_translate = (chat_language != \"English\")\n",
        "\n",
        "# Data\n",
        "models = {\n",
        "    \"Pygmalion 6B original (sharded, rehosted)\": (\"ayylamo\", \"pygmalion-6b\", \"original-sharded\", \"pygmalion-6b_original-sharded\"),\n",
        "    \"Pygmalion 6B main (sharded, rehosted)\": (\"ayylamo\", \"pygmalion-6b\", \"sharded\", \"pygmalion-6b_sharded\"),\n",
        "    \"Pygmalion 6B dev (sharded, rehosted)\": (\"ayylamo\", \"pygmalion-6b\", \"dev-sharded\", \"pygmalion-6b_dev-sharded\"),\n",
        "    \"Pygmalion 350m (for debugging)\": (\"PygmalionAI\", \"pygmalion-350m\", \"main\", \"pygmalion-350m\"),\n",
        "}\n",
        "language_codes = {'Afrikaans': 'af', 'Albanian': 'sq', 'Amharic': 'am', 'Arabic': 'ar', 'Armenian': 'hy', 'Azerbaijani': 'az', 'Basque': 'eu', 'Belarusian': 'be', 'Bengali': 'bn', 'Bosnian': 'bs', 'Bulgarian': 'bg', 'Catalan': 'ca', 'Cebuano': 'ceb', 'Chinese (Simplified)': 'zh-CN', 'Chinese (Traditional)': 'zh-TW', 'Corsican': 'co', 'Croatian': 'hr', 'Czech': 'cs', 'Danish': 'da', 'Dutch': 'nl', 'English': 'en', 'Esperanto': 'eo', 'Estonian': 'et', 'Finnish': 'fi', 'French': 'fr', 'Frisian': 'fy', 'Galician': 'gl', 'Georgian': 'ka', 'German': 'de', 'Greek': 'el', 'Gujarati': 'gu', 'Haitian Creole': 'ht', 'Hausa': 'ha', 'Hawaiian': 'haw', 'Hebrew': 'iw', 'Hindi': 'hi', 'Hmong': 'hmn', 'Hungarian': 'hu', 'Icelandic': 'is', 'Igbo': 'ig', 'Indonesian': 'id', 'Irish': 'ga', 'Italian': 'it', 'Japanese': 'ja', 'Javanese': 'jw', 'Kannada': 'kn', 'Kazakh': 'kk', 'Khmer': 'km', 'Korean': 'ko', 'Kurdish': 'ku', 'Kyrgyz': 'ky', 'Lao': 'lo', 'Latin': 'la', 'Latvian': 'lv', 'Lithuanian': 'lt', 'Luxembourgish': 'lb', 'Macedonian': 'mk', 'Malagasy': 'mg', 'Malay': 'ms', 'Malayalam': 'ml', 'Maltese': 'mt', 'Maori': 'mi', 'Marathi': 'mr', 'Mongolian': 'mn', 'Myanmar (Burmese)': 'my', 'Nepali': 'ne', 'Norwegian': 'no', 'Nyanja (Chichewa)': 'ny', 'Pashto': 'ps', 'Persian': 'fa', 'Polish': 'pl', 'Portuguese (Portugal, Brazil)': 'pt', 'Punjabi': 'pa', 'Romanian': 'ro', 'Russian': 'ru', 'Samoan': 'sm', 'Scots Gaelic': 'gd', 'Serbian': 'sr', 'Sesotho': 'st', 'Shona': 'sn', 'Sindhi': 'sd', 'Sinhala (Sinhalese)': 'si', 'Slovak': 'sk', 'Slovenian': 'sl', 'Somali': 'so', 'Spanish': 'es', 'Sundanese': 'su', 'Swahili': 'sw', 'Swedish': 'sv', 'Tagalog (Filipino)': 'tl', 'Tajik': 'tg', 'Tamil': 'ta', 'Telugu': 'te', 'Thai': 'th', 'Turkish': 'tr', 'Ukrainian': 'uk', 'Urdu': 'ur', 'Uzbek': 'uz', 'Vietnamese': 'vi', 'Welsh': 'cy', 'Xhosa': 'xh', 'Yiddish': 'yi', 'Yoruba': 'yo', 'Zulu': 'zu'}\n",
        "\n",
        "# Download the model (if it hasn't been downloaded already)\n",
        "huggingface_org, huggingface_repo, huggingface_branch, model_name = models[model]\n",
        "![[ ! -f models/$model_name/config.json ]] && python download-model.py $huggingface_org/$huggingface_repo --branch $huggingface_branch\n",
        "\n",
        "# Applying the selected language and setting the prompt size to 2048\n",
        "# if 8bit mode is selected\n",
        "j = json.loads(open('settings-colab-template.json', 'r').read())\n",
        "j[\"google_translate-language string\"] = language_codes[chat_language]\n",
        "if load_in_8bit:\n",
        "  j[\"chat_prompt_size\"] = 2048\n",
        "with open('settings-colab.json', 'w') as f:\n",
        "  f.write(json.dumps(j, indent=4))\n",
        "\n",
        "params = set()\n",
        "if cai_chat:\n",
        "  params.add('--cai-chat')\n",
        "else:\n",
        "  params.add('--chat')\n",
        "\n",
        "if load_in_8bit:\n",
        "  params.add('--load-in-8bit')\n",
        "\n",
        "active_extensions = []\n",
        "if activate_character_bias:\n",
        "  active_extensions.append('character_bias')\n",
        "if activate_google_translate:\n",
        "  active_extensions.append('google_translate')\n",
        "if activate_silero_text_to_speech:\n",
        "  active_extensions.append('silero_tts')\n",
        "\n",
        "if len(active_extensions) > 0:\n",
        "  params.add(f'--extensions {\",\".join(active_extensions)}')\n",
        "\n",
        "if not text_streaming or activate_google_translate or activate_silero_text_to_speech:\n",
        "  params.add('--no-stream')\n",
        "if activate_character_bias:\n",
        "  params.add('--verbose')\n",
        "if activate_sending_pictures:\n",
        "  params.add('--picture')\n",
        "\n",
        "# Starting the web UI\n",
        "cmd = f\"python server.py --share --model {model_name} --settings settings-colab.json {' '.join(params)}\"\n",
        "print(cmd)\n",
        "!$cmd"
      ],
      "metadata": {
        "id": "hKuocueuXnm5",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}