├── GPT2_with_Javascript_interface_POC.ipynb ├── LICENSE ├── README.md ├── gpt2js.png └── how_it_works.md /GPT2_with_Javascript_interface_POC.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "nbformat": 4, 3 | "nbformat_minor": 0, 4 | "metadata": { 5 | "colab": { 6 | "name": "GPT2_with_Javascript_interface_POC.ipynb", 7 | "provenance": [], 8 | "collapsed_sections": [], 9 | "include_colab_link": true 10 | }, 11 | "kernelspec": { 12 | "name": "python3", 13 | "display_name": "Python 3" 14 | }, 15 | "accelerator": "GPU" 16 | }, 17 | "cells": [ 18 | { 19 | "cell_type": "markdown", 20 | "metadata": { 21 | "id": "view-in-github", 22 | "colab_type": "text" 23 | }, 24 | "source": [ 25 | "\"Open" 26 | ] 27 | }, 28 | { 29 | "cell_type": "markdown", 30 | "metadata": { 31 | "id": "ZMMtepKIwxR8" 32 | }, 33 | "source": [ 34 | "#This is proof of concept that GPT-2 can be run from colab with Javascript interface\n", 35 | "**Q: How to do?**\n", 36 | "\n", 37 | "A: \n", 38 | "1. Runtime -> Change runtime type -> Hardware accelerator -> GPU\n", 39 | "2. Runtime -> Reset all runtimes\n", 40 | "3. Runtime -> Run all\n", 41 | "4. Scroll down and wait until you see the little window\n", 42 | "5. Type text\n", 43 | "6. The button \"Continue with GPT-2\" will invoke GPT-2 and it will continue your text.\n", 44 | "\n" 45 | ] 46 | }, 47 | { 48 | "cell_type": "code", 49 | "metadata": { 50 | "id": "rGUX9yKxnaRe" 51 | }, 52 | "source": [ 53 | "%tensorflow_version 1.x\n", 54 | "!git clone https://github.com/gpt2ent/gpt-2-simple.git\n", 55 | "%cd gpt-2-simple\n", 56 | "!git checkout context-trim\n", 57 | "!pip install .\n", 58 | "%cd ..\n", 59 | "!git clone https://github.com/gpt2ent/gpt2colab-js.git\n", 60 | "%cd gpt2colab-js\n", 61 | "\n", 62 | "import gpt_2_simple as gpt2\n", 63 | "\n", 64 | "import os\n", 65 | "import requests\n", 66 | "import tensorflow as tf\n", 67 | "\n", 68 | "import re\n", 69 | "\n", 70 | "#Determining the graphics card used by colab: full model can run only on P100\n", 71 | "\n", 72 | "try:\n", 73 | " !cat /proc/driver/nvidia/gpus/0000:00:04.0/information >> /content/card_info.txt\n", 74 | " with open('/content/card_info.txt','r') as f:\n", 75 | " graphics_card = re.split('\\n|\\t\\t ',f.read())[1]\n", 76 | "\n", 77 | " if not graphics_card.startswith(\"Tesla P100\") and not graphics_card.startswith(\"Tesla T4\"):\n", 78 | " print(\"=\"*90+'\\n'+\"=\"*90+'\\n\\n')\n", 79 | " print('\\n\\tYour current GPU - %s - cannot fit the full GPT-2 model!' % graphics_card)\n", 80 | " print('\\n\\tFalling back on 774M model.')\n", 81 | " print('\\n\\tNothing I can do. just pray to Google to give you a P100')\n", 82 | " print('\\t\\tnext time. ¯\\_(ツ)_/¯')\n", 83 | " print('\\n\\tAlso you might try TPU runtime.')\n", 84 | " print('\\n\\n'+\"=\"*90+'\\n'+\"=\"*90+'\\n\\n')\n", 85 | " model_name = \"774M\"\n", 86 | " spinner_speed = \"300ms\"\n", 87 | " else:\n", 88 | " print('GPU: %s' % graphics_card)\n", 89 | " model_name = \"1558M\"\n", 90 | " spinner_speed = '400ms'\n", 91 | "except IndexError:\n", 92 | " print(\"=\"*90+'\\n'+\"=\"*90+'\\n\\n')\n", 93 | " print('\\n\\tYou\\'re not in a GPU runtime.\\n')\n", 94 | " print('\\n\\tTrying 1558M model anyways - assuming you\\'re on a good TPU.')\n", 95 | " print('\\n\\tIf it fails, you have to go to Runtime -> Change runtime type')\n", 96 | " print('\\n\\tand choose GPU.')\n", 97 | " print('\\n\\n'+\"=\"*90+'\\n'+\"=\"*90+'\\n\\n')\n", 98 | " model_name = \"1558M\"\n", 99 | " spinner_speed = \"1200ms\"\n", 100 | "\n", 101 | "\n", 102 | "#Overwrite default model choice\n", 103 | "#model_name = \"1558M\"\n", 104 | "#model_name = \"774M\"\n", 105 | "#model_name = \"124M\"\n", 106 | "#model_name = \"355M\"\n", 107 | "\n", 108 | "\n", 109 | "if not os.path.isdir(os.path.join(\"models\", model_name)):\n", 110 | " print(f\"Downloading {model_name} model...\")\n", 111 | " gpt2.download_gpt2(model_name=model_name)\n", 112 | " \n", 113 | "sess = gpt2.start_tf_sess()\n", 114 | "gpt2.load_gpt2(sess, model_name=model_name)\n", 115 | "\n", 116 | "generate_count = 0\n", 117 | "\n", 118 | "import google.colab.output\n", 119 | "\n", 120 | "import json\n", 121 | "\n", 122 | "class JsonRepr:\n", 123 | " \"\"\"\n", 124 | " For some reasons I can only use the result of __repr__\n", 125 | " from inside Javascript. So this wrapper uses json.dumps() as __repr__\n", 126 | " for python function output.\n", 127 | " \"\"\"\n", 128 | " def __init__(self, obj):\n", 129 | " self.obj = obj\n", 130 | "\n", 131 | " def __repr__(self):\n", 132 | " return json.dumps(self.obj)\n", 133 | "\n", 134 | "def overlap(a, b):\n", 135 | " return max(i for i in range(len(b)+1) if a.endswith(b[:i]))\n", 136 | "\n", 137 | "\n", 138 | "def ai_generate(prefix, temp, top_k, length):\n", 139 | " global sess\n", 140 | " global generate_count\n", 141 | "\n", 142 | " temp = float(temp)\n", 143 | " top_k = int(top_k)\n", 144 | " length = int(length)\n", 145 | " result = gpt2.generate(sess, model_name=model_name, prefix=prefix, temperature=temp,\n", 146 | " top_k=top_k, length=length, include_prefix=False, return_as_list=True)[0]\n", 147 | " \n", 148 | " j = overlap(prefix, result)\n", 149 | " result = result[j:]\n", 150 | " \n", 151 | " generate_count += 1\n", 152 | " if generate_count == 6:\n", 153 | " #prevent memory leak as in https://github.com/minimaxir/gpt-2-simple/issues/71\n", 154 | " tf.reset_default_graph()\n", 155 | " sess.close()\n", 156 | " sess = gpt2.start_tf_sess()\n", 157 | " gpt2.load_gpt2(sess, model_name=model_name)\n", 158 | " generate_count = 0\n", 159 | " return JsonRepr(result)\n", 160 | "\n", 161 | "#register callback for Javascript\n", 162 | "google.colab.output.register_callback('ai_generate', ai_generate)\n", 163 | "\n", 164 | "print('Done')" 165 | ], 166 | "execution_count": null, 167 | "outputs": [] 168 | }, 169 | { 170 | "cell_type": "code", 171 | "metadata": { 172 | "id": "CdyRipC0o8vR" 173 | }, 174 | "source": [ 175 | "from IPython.display import HTML\n", 176 | "\n", 177 | "#spinner from https://codepen.io/vovchisko/pen/vROoYQ\n", 178 | "spinner_css = \"\"\"\n", 179 | "\n", 216 | "\"\"\"\n", 217 | "\n", 218 | "input_form = \"\"\"\n", 219 | "\n", 220 | "\n", 221 | "
\n", 222 | "

You have currently loaded %s model

\n", 223 | "
\n", 224 | "
\n", 225 | "
\n", 226 | " \n", 227 | " \n", 228 | "
\n", 229 | "
\n", 230 | " \n", 231 | " \n", 232 | "
\n", 233 | "
\n", 234 | " \n", 235 | " \n", 236 | "
\n", 237 | "
\n", 238 | "

\n", 239 | "

\n", 240 | "
\n", 241 | "
\n", 242 | "
\n", 243 | "\"\"\" % model_name\n", 244 | "\n", 245 | "javascript = \"\"\"\n", 246 | "\n", 269 | "\"\"\"\n", 270 | "\n", 271 | "HTML(spinner_css + input_form + javascript)" 272 | ], 273 | "execution_count": null, 274 | "outputs": [] 275 | } 276 | ] 277 | } -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 gpt2ent 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # gpt2colab-js 2 | Make GPT-2 complete your text in Colab. Use fancy JS, instead of default ugly colab interface. 3 | This is proof of concept for developers that want to create apps with graphics in Colab. 4 | 5 | If you want to use this code as your playground check out [the 'multisample' branch](https://github.com/gpt2ent/gpt2colab-js/tree/multisample) for more useful notebooks. 6 | ![Teaser image](./gpt2js.png) 7 | 8 | # How to use the interface with your model 9 | 10 | If your model [is uploaded to HuggingFace](https://huggingface.co/models?pipeline_tag=text-generation), check out [the 'multisample' branch](https://github.com/gpt2ent/gpt2colab-js/tree/multisample). 11 | 12 | For custom models the idea is the following: 13 | 14 | 1. Train your model *outside* of this notebook. This notebook is supposed to inference from already pretrained models. 15 | 3. Ensure you have a way to call your model *from within Python code* and get a string. That means if you infer text via calling external script -- something like `!python main.py --predict output.txt`) -- you need to examine the inference code of the script and write some sort of a function or an object that will handle your inference and return strings to the code, not just print the result or write it to file. 16 | 2. Discard all of the code that goes before `import google.colab.output` and copy your own code that loads and prepares your model. Add `model_name` and `spinner_speed` variables or your HTML code won't run. Note that `spinner_speed` is a string variable that looks like `"400ms"`. 17 | 5. Use the `ai_generate` function to connect your model to the JS: 18 | * The simplest way is to not change any arguments and just use your function/object to generate a string and put it into `result` variable before returning `JsonRepr(result)`. 19 | * Note that the function will receive `top_k`, `temp` and `length` as *string* variables and will convert these into numbers internally. 20 | * Also note that you aren't supposed to output any part of `prefix` as part of your output string. 21 | * If you want to use other parameters (e.g. `top_p`) you need to fix HTML and JS code (`generate()` function) in the second block. Again, note that JS will pass any arguments as *strings*. 22 | -------------------------------------------------------------------------------- /gpt2js.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/gpt2ent/gpt2colab-js/c1a844942277ad98cba64253b669582b95917968/gpt2js.png -------------------------------------------------------------------------------- /how_it_works.md: -------------------------------------------------------------------------------- 1 | # How it works 2 | 3 | ## Python side 4 | 5 | ```python 6 | import google.colab.output 7 | 8 | def some_function(*args, **kwargs): 9 | #your code 10 | 11 | google.colab.output.register_callback('js_callback', some_function) 12 | ``` 13 | 14 | ## JS side 15 | 16 | ```javascript 17 | var kernel = google.colab.kernel; 18 | var resultPromise = kernel.invokeFunction("js_callback", 19 | [arg1,arg2], 20 | {kwarg1: 1, kwarg2: 2}); //returns a Promise 21 | resultPromise.then( 22 | function(value) { 23 | console.log(value); //go figure what it returns in console 24 | }, function(value) { 25 | console.log(value); //no idea what's gonna happen here 26 | }); 27 | ``` 28 | 29 | Python errors returned to the JS will show up in console. 30 | --------------------------------------------------------------------------------