GPT2_with_Javascript_interface_POC.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "GPT2_with_Javascript_interface_POC.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mrm8488/shared_colab_notebooks/blob/master/GPT2_with_Javascript_interface_POC.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZMMtepKIwxR8",
        "colab_type": "text"
      },
      "source": [
        "#This is proof of concept that GPT-2 can be run from colab with Javascript interface\n",
        "**Q: How to do?**\n",
        "\n",
        "A: \n",
        "1. Runtime -> Change runtime type -> Hardware accelerator -> GPU\n",
        "2. Runtime -> Reset all runtimes\n",
        "3. Runtime -> Run all\n",
        "4. Scroll down and wait until you see the little window\n",
        "5. Type text\n",
        "6. The button \"Continue with GPT-2\" will invoke GPT-2 and it will continue your text.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rGUX9yKxnaRe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%tensorflow_version 1.x\n",
        "!git clone https://github.com/gpt2ent/gpt-2-simple.git\n",
        "%cd gpt-2-simple\n",
        "!git checkout context-trim\n",
        "!pip install .\n",
        "%cd ..\n",
        "!git clone https://github.com/gpt2ent/gpt2colab-js.git\n",
        "%cd gpt2colab-js\n",
        "\n",
        "import gpt_2_simple as gpt2\n",
        "\n",
        "import os\n",
        "import requests\n",
        "import tensorflow as tf\n",
        "\n",
        "import re\n",
        "\n",
        "#Determining the graphics card used by colab: full model can run only on P100\n",
        "\n",
        "try:\n",
        "    !cat /proc/driver/nvidia/gpus/0000:00:04.0/information >> /content/card_info.txt\n",
        "    with open('/content/card_info.txt','r') as f:\n",
        "        graphics_card = re.split('\\n|\\t\\t ',f.read())[1]\n",
        "\n",
        "    if not graphics_card.startswith(\"Tesla P100\"):\n",
        "        print(\"=\"*90+'\\n'+\"=\"*90+'\\n\\n')\n",
        "        print('\\n\\tYour current GPU - %s - cannot fit the full GPT-2 model!' % graphics_card)\n",
        "        print('\\n\\tFalling back on 774M model.')\n",
        "        print('\\n\\tNothing I can do. just pray to Google to give you a P100')\n",
        "        print('\\t\\tnext time. ¯\\_(ツ)_/¯')\n",
        "        print('\\n\\tAlso you might try TPU runtime.')\n",
        "        print('\\n\\n'+\"=\"*90+'\\n'+\"=\"*90+'\\n\\n')\n",
        "        model_name = \"774M\"\n",
        "        spinner_speed = \"300ms\"\n",
        "    else:\n",
        "        print('GPU: %s' % graphics_card)\n",
        "        model_name = \"1558M\"\n",
        "        spinner_speed = '400ms'\n",
        "except IndexError:\n",
        "    print(\"=\"*90+'\\n'+\"=\"*90+'\\n\\n')\n",
        "    print('\\n\\tYou\\'re not in a GPU runtime.\\n')\n",
        "    print('\\n\\tTrying 1558M model anyways - assuming you\\'re on a good TPU.')\n",
        "    print('\\n\\tIf it fails, you have to go to Runtime -> Change runtime type')\n",
        "    print('\\n\\tand choose GPU.')\n",
        "    print('\\n\\n'+\"=\"*90+'\\n'+\"=\"*90+'\\n\\n')\n",
        "    model_name = \"1558M\"\n",
        "    spinner_speed = \"1200ms\"\n",
        "\n",
        "\n",
        "#Overwrite default model choice\n",
        "#model_name = \"1558M\"\n",
        "#model_name = \"774M\"\n",
        "#model_name = \"124M\"\n",
        "#model_name = \"355M\"\n",
        "\n",
        "\n",
        "if not os.path.isdir(os.path.join(\"models\", model_name)):\n",
        "    print(f\"Downloading {model_name} model...\")\n",
        "    gpt2.download_gpt2(model_name=model_name)\n",
        "  \n",
        "sess = gpt2.start_tf_sess()\n",
        "gpt2.load_gpt2(sess, model_name=model_name)\n",
        "\n",
        "generate_count = 0\n",
        "\n",
        "import google.colab.output\n",
        "\n",
        "def overlap(a, b):\n",
        "    return max(i for i in range(len(b)+1) if a.endswith(b[:i]))\n",
        "\n",
        "\n",
        "def ai_generate(prefix, temp, top_k, length):\n",
        "    global sess\n",
        "    global generate_count\n",
        "\n",
        "    temp = float(temp)\n",
        "    top_k = int(top_k)\n",
        "    length = int(length)\n",
        "    result = gpt2.generate(sess, model_name=model_name, prefix=prefix, temperature=temp,\n",
        "                        top_k=top_k, length=length, include_prefix=False, return_as_list=True)[0]\n",
        "    \n",
        "    j = overlap(prefix, result)\n",
        "    result = result[j:]\n",
        "    \n",
        "    generate_count += 1\n",
        "    if generate_count == 6:\n",
        "          #prevent memory leak as in https://github.com/minimaxir/gpt-2-simple/issues/71\n",
        "          tf.reset_default_graph()\n",
        "          sess.close()\n",
        "          sess = gpt2.start_tf_sess()\n",
        "          gpt2.load_gpt2(sess, model_name=model_name)\n",
        "          generate_count = 0\n",
        "    return result\n",
        "\n",
        "#register callback for Javascript\n",
        "google.colab.output.register_callback('ai_generate', ai_generate)\n",
        "\n",
        "print('Done')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CdyRipC0o8vR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from IPython.display import HTML\n",
        "\n",
        "#spinner from https://codepen.io/vovchisko/pen/vROoYQ\n",
        "spinner_css = \"\"\"\n",
        "<style>\n",
        "@keyframes c-inline-spinner-kf {\n",
        "  0% {\n",
        "    transform: rotate(0deg);\n",
        "  }\n",
        "  100% {\n",
        "    transform: rotate(360deg);\n",
        "  }\n",
        "}\n",
        "\n",
        ".c-inline-spinner,\n",
        ".c-inline-spinner:before {\n",
        "  display: inline-block;\n",
        "  width: 11px;\n",
        "  height: 11px;\n",
        "  transform-origin: 50%;\n",
        "  border: 2px solid transparent;\n",
        "  border-color: #74a8d0 #74a8d0 transparent transparent;\n",
        "  border-radius: 50%;\n",
        "  content: \"\";\n",
        "  animation: linear c-inline-spinner-kf \"\"\"+spinner_speed+\"\"\" infinite;\n",
        "  position: relative;\n",
        "  vertical-align: inherit;\n",
        "  line-height: inherit;\n",
        "}\n",
        ".c-inline-spinner {\n",
        "  top: 3px;\n",
        "  margin: 0 3px;\n",
        "}\n",
        ".c-inline-spinner:before {\n",
        "  border-color: #74a8d0 #74a8d0 transparent transparent;\n",
        "  position: absolute;\n",
        "  left: -2px;\n",
        "  top: -2px;\n",
        "  border-style: solid;\n",
        "}\n",
        "</style>\n",
        "\"\"\"\n",
        "\n",
        "input_form = \"\"\"\n",
        "<link rel=\"stylesheet\" href=\"https://unpkg.com/purecss@1.0.1/build/pure-min.css\" integrity=\"sha384-oAOxQR6DkCoMliIh8yFnu25d7Eq/PHS21PClpwjOTeU2jRSq11vu66rf90/cZr47\" crossorigin=\"anonymous\">\n",
        "\n",
        "<div style=\"background-color:white; border:solid #ccc; width:800px; padding:20px; color: black;\">\n",
        "<p>You have currently loaded %s model</p>\n",
        "<textarea id=\"main_textarea\" cols=\"75\" rows=\"20\" style=\"font-family: 'Liberation Serif', 'DejaVu Serif', Georgia, 'Times New Roman', Times, serif; font-size: 13pt; padding:10px;\"></textarea><br>\n",
        "<div class=\"pure-form pure-form-aligned\">\n",
        "    <div class=\"pure-control-group\">\n",
        "      <label for=\"temp\">Temperature:</label>\n",
        "      <input type=\"number\" min=\"0.00\" max=\"999.99\" step=\"0.01\" id=\"temp\" value=\"0.70\" style=\"background-color: white;\">\n",
        "    </div>\n",
        "    <div class=\"pure-control-group\">\n",
        "        <label for=\"top_k\">top_k:</label>\n",
        "        <input type=\"number\" min=\"0\" max=\"9999\" id=\"top_k\" value=\"40\" style=\"background-color: white;\">\n",
        "    </div>\n",
        "    <div class=\"pure-control-group\">\n",
        "        <label for=\"length\">Generate how much:</label>\n",
        "        <input type=\"number\" id=\"length\" min=\"1\" max=\"1023\" value=\"10\" style=\"background-color: white;\">\n",
        "    </div>\n",
        "    <div style=\"width: 300px; display: block; margin-left: auto !important; margin-right: auto !important;\">\n",
        "        <p><button class=\"pure-button pure-button-primary\" style=\"font-size: 125%%;\" onclick=\"generate()\">Continue with GPT-2</button>\n",
        "        <span class=\"c-inline-spinner\" style=\"visibility: hidden;\" id=\"spinner\"></span></p>\n",
        "    </div>\n",
        "</div>\n",
        "</div>\n",
        "\"\"\" % model_name\n",
        "\n",
        "javascript = \"\"\"\n",
        "<script type=\"text/Javascript\">\n",
        "    function desanitize(text) {\n",
        "        return text.slice(1,-1).replace(/\\\\\\\\n/g, \"\\\\n\").replace(/\\\\\\\\'/g, \"'\");\n",
        "    };\n",
        "\n",
        "    function add_text(text) {\n",
        "        var deftext = document.getElementById('main_textarea').value;\n",
        "        document.getElementById('main_textarea').value = deftext + text;\n",
        "    };\n",
        "\n",
        "    function generate(){\n",
        "        var prefix = document.getElementById('main_textarea').value;\n",
        "        var temp = document.getElementById('temp').value;\n",
        "        var top_k = document.getElementById('top_k').value;\n",
        "        var length = document.getElementById('length').value;\n",
        "        \n",
        "        var kernel = google.colab.kernel;\n",
        "        var resultPromise = kernel.invokeFunction(\"ai_generate\", [prefix,temp,top_k,length]); // developer, look here\n",
        "        resultPromise.then(\n",
        "            function(value) {\n",
        "              add_text(desanitize(value.data[\"text/plain\"]));\n",
        "              document.getElementById('spinner').style = \"visibility: hidden;\";\n",
        "        });\n",
        "        document.getElementById('spinner').style = \"visibility: visible;\";\n",
        "    };\n",
        "</script>\n",
        "\"\"\"\n",
        "\n",
        "HTML(spinner_css + input_form + javascript)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}