{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "view-in-github", "colab_type": "text" }, "source": [ "\"Open" ] }, { "cell_type": "markdown", "source": [ "# GOOGLE HAS BROKEN SUPPORT FOR MESH TRANSFORMERS JAX IN THEIR DRIVER, THIS MESSAGE WILL BE REMOVED ONCE THE NOTEBOOK WORKS AGAIN\n", "---\n", "Are you a developer familair with Jax? We could use some help.\n", "Our Mesh Transformers Jax fork resides here (but unfortunately VE-Forbryderne has gone missing) : https://github.com/VE-FORBRYDERNE/mesh-transformer-jax\n", "\n", "This is combined with the TPU backend code you can find here: https://github.com/KoboldAI/KoboldAI-Client/blob/main/tpu_mtj_backend.py\n", "\n", "So far we know the driver initialization issues can be resolved when a newer version of Jax is used combined with a slight edit to tpu_mtj_backend.py to use Jax's built in code for driver intialization (Which is not present in the older version we currently use, hence that part being in our file).\n", "\n", "As far as we understand the issue is that xmap was broken in newer versions, and MTJ makes use of it. If someone can port this part of the code to be compatible with the newer Jax versions you can save this notebook!\n", "\n", "(Or Google, if you are reading this. Please reintroduce the old 0.1 compatibility since you broke an entire ecosystem of Mesh Transformers Jax users, which has historic value because it is the original implementation of GPT-J. There also do not seem to be alternatives that have the same amount of performance and model compatibility).\n", "\n", "# 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", "Press play on the music player to keep the tab alive, then start KoboldAI below (Uses only 13MB of data)
\n", "