gpt-neo_dungeon_untuned.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "gpt-neo dungeon untuned.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cdtMW21T9pDl"
      },
      "source": [
        "# Instructions\n",
        "\n",
        "Go through each cell in this notebook one by one, take a look at the options and descriptions and then press the play button to the left of it. You can skip the optional one. Don't skip any of the others. After running the \"Play\" cell, a small form will appear underneath, which you can use to actually play.\n",
        "\n",
        "To reset the state of your game, run the \"Setup\" cell again. Closing the notebook will lose your progress, so if you want to keep your story, use the \"history\" action, copy out your story to a text editor. You can also copy out your author's note and memory from the output of the \"info\" action.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ucP5hOdoMgzC",
        "cellView": "form"
      },
      "source": [
        "#@title Setup\n",
        "#@markdown Run this for setting up dependencies or resetting actions\n",
        "!pip install git+https://github.com/finetuneanon/transformers@gpt-neo-dungeon-localattention1\n",
        "!wget -c http://ftp.us.debian.org/debian/pool/main/m/megatools/megatools_1.11.0~git20200404-1_amd64.deb -O megatools.deb\n",
        "!dpkg -i megatools.deb\n",
        "!nvidia-smi\n",
        "\n",
        "import os\n",
        "\n",
        "from transformers import GPTNeoForCausalLM, AutoTokenizer\n",
        "import tarfile\n",
        "import codecs\n",
        "import torch\n",
        "import subprocess\n",
        "\n",
        "from IPython.display import HTML, display\n",
        "import ipywidgets as widgets\n",
        "\n",
        "def set_css():\n",
        "  display(HTML('''\n",
        "  <style>\n",
        "    pre {\n",
        "        white-space: pre-wrap;\n",
        "    }\n",
        "  </style>\n",
        "  '''))\n",
        "\n",
        "try:\n",
        "  initialized += 1\n",
        "except:\n",
        "  get_ipython().events.register('pre_run_cell', set_css)\n",
        "  initialized = 0\n",
        "\n",
        "actions = []\n",
        "memory = (\"\", torch.zeros((1, 0)).long())\n",
        "lmi = [\"\", torch.zeros((1, 0)).long()]\n",
        "an = (\"\", torch.zeros((1, 0)).long())\n",
        "an_depth = 3\n",
        "history = None"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "M3ZF4eCFMV6E",
        "cellView": "form"
      },
      "source": [
        "#@title Model setup\n",
        "\n",
        "model_name = \"EleutherAI/gpt-neo-2.7B\" #@param [\"EleutherAI/gpt-neo-2.7B\", \"EleutherAI/gpt-neo-1.3B\"]\n",
        "\n",
        "model = None\n",
        "tokenizer = None\n",
        "pipeline = None\n",
        "checkpoint = None\n",
        "\n",
        "if True:\n",
        "  from transformers.file_utils import cached_path, WEIGHTS_NAME, hf_bucket_url\n",
        "  archive_file = hf_bucket_url(model_name, filename=WEIGHTS_NAME)\n",
        "  resolved_archive_file = cached_path(archive_file)\n",
        "  checkpoint = torch.load(resolved_archive_file, map_location=\"cuda:0\")\n",
        "  for k in checkpoint.keys():\n",
        "    checkpoint[k] = checkpoint[k].half()\n",
        "  model = GPTNeoForCausalLM.from_pretrained(model_name, state_dict=checkpoint).half().to(\"cuda\").eval()\n",
        "  for k in list(checkpoint.keys()):\n",
        "    del checkpoint[k]\n",
        "  del checkpoint\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ls_x_sSLN2hU",
        "cellView": "form"
      },
      "source": [
        "#@title Sampling settings\n",
        "#@markdown You can modify sampling settings here. Don't forget to run the cell again after changing. The number of generated tokens is subtracted from the context window size, don't set it high.\n",
        "top_k = 60 #@param {type:\"number\"}\n",
        "top_p = 0.9 #@param {type:\"number\"}\n",
        "temperature =  0.6#@param {type:\"number\"}\n",
        "number_generated_tokens =  40#@param {type:\"integer\"}\n",
        "repetition_penalty = 1.25 #@param {type:\"number\"}\n",
        "repetition_penalty_range = 300 #@param {type:\"number\"}\n",
        "repetition_penalty_slope = 3.33 #@param {type:\"number\"}\n",
        "number_show_last_actions = 15 #@param {type:\"integer\"}\n",
        "\n",
        "#@markdown Temperatures seem to give results different from those in AID, so play around with it. Even 0.5 can give good results."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "qd-HLb2nXaaA"
      },
      "source": [
        "#@title Basic sampling\n",
        "\n",
        "#@markdown Use this cell if you just want to sample from the model in a free form way.\n",
        "\n",
        "basic_prompt = \"The rays of the evening sun falling in through the window bathed the room in a soft, warm light\" #@param {type:\"string\"}\n",
        "\n",
        "ids = tokenizer(basic_prompt, return_tensors=\"pt\").input_ids.to(\"cpu\")\n",
        "n_ids = ids.shape[1]\n",
        "if n_ids < 1:\n",
        "  n_ids = 1\n",
        "  ids = torch.tensor([[tokenizer.eos_token_id]])\n",
        "max_length = n_ids + number_generated_tokens\n",
        "torch.cuda.empty_cache()\n",
        "basic_output = model.generate(\n",
        "    ids.long().cuda(),\n",
        "    do_sample=True,\n",
        "    min_length=max_length,\n",
        "    max_length=max_length,\n",
        "    temperature=temperature,\n",
        "    top_k = top_k,\n",
        "    top_p = top_p,\n",
        "    repetition_penalty = repetition_penalty,\n",
        "    repetition_penalty_range = repetition_penalty_range,\n",
        "    repetition_penalty_slope = repetition_penalty_slope,\n",
        "    use_cache=True,\n",
        "    pad_token_id=tokenizer.eos_token_id\n",
        ").long().to(\"cpu\")\n",
        "torch.cuda.empty_cache()\n",
        "\n",
        "print(tokenizer.decode(basic_output[0]))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8GKfy1uQ9QdU"
      },
      "source": [
        "# Using gpt-neo dungeon's play function\n",
        "\n",
        "If your prompt starts with a letter, try putting a space or newline in front.\n",
        "\n",
        "* **generate** adds your prompt as an action and generates more output\n",
        "* **continue** generates more output\n",
        "* **replace** replaces the last output with the prompt and generates more, use this to edit\n",
        "* **info** outputs LMI and memory\n",
        "* **history** outputs all actions so far\n",
        "* **memory** sets memory to the text in the prompt field\n",
        "* **authorsnote** sets author's note to the text in the prompt field\n",
        "* **andepth** sets the depth of the author's note to the number in the prompt\n",
        "* **tokenize** tokenizes the text in the prompt field and outputs the number of tokens"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fb_NHHISOeYg",
        "cellView": "form"
      },
      "source": [
        "#@title Play\n",
        "\n",
        "action_type = \"generate\"\n",
        "prompt = \"\"\n",
        "need_refresh = True\n",
        "\n",
        "action_types = [\"generate\", \"continue\", \"replace\", \"undo\", \"retry\", \"memory\", \"authorsnote\", \"andepth\", \"info\", \"history\", \"tokenize\"]\n",
        "\n",
        "def assemble():\n",
        "  remaining = (2048 - number_generated_tokens + 1) - memory[1].shape[1] - an[1].shape[1]\n",
        "  n_actions = len(actions)\n",
        "  n_ctx = 0\n",
        "  back_i = n_actions\n",
        "  for i in range(n_actions):\n",
        "      i_action = n_actions - i - 1\n",
        "      n_tok = actions[i_action][1].shape[1]\n",
        "      if remaining > n_ctx + n_tok:\n",
        "        n_ctx += n_tok\n",
        "        back_i = i_action\n",
        "      else:\n",
        "        break\n",
        "  lmi[0], lmi[1] = memory[0], memory[1]\n",
        "  start = False\n",
        "  if n_actions - back_i - 1 < an_depth:\n",
        "    start = True\n",
        "  while back_i < n_actions:\n",
        "    if start or n_actions - back_i - 1 == an_depth:\n",
        "      lmi[0] += an[0]\n",
        "      lmi[1] = torch.cat([lmi[1].cpu(), an[1].cpu()], 1).long()\n",
        "      start = False\n",
        "    lmi[0] += actions[back_i][0]\n",
        "    lmi[1] = torch.cat([lmi[1].cpu(), actions[back_i][1].cpu()], 1).long()\n",
        "    back_i += 1\n",
        "\n",
        "def clear_output():\n",
        "  with out:\n",
        "    IPython.display.clear_output()\n",
        "\n",
        "def set_action(change):\n",
        "  global action_type\n",
        "  action_type = change.new\n",
        "\n",
        "def set_prompt(change):\n",
        "  global prompt\n",
        "  prompt = change.new\n",
        "\n",
        "@torch.no_grad()\n",
        "def play(do_action=None):\n",
        "  global memory, need_refresh, an, an_depth, action_type, history\n",
        "  an_updated = False\n",
        "  memory_updated = False\n",
        "  if do_action is not None:\n",
        "    action = do_action\n",
        "    action_type = do_action\n",
        "  else:\n",
        "    action = action_type\n",
        "  with out:\n",
        "    if prompt in action_types:\n",
        "      action == prompt\n",
        "    else:\n",
        "      if action == \"replace\":\n",
        "        if len(actions) > 0:\n",
        "          actions.pop()\n",
        "        need_refresh = True\n",
        "        action = \"generate\"\n",
        "      if action == \"generate\":\n",
        "        text = prompt\n",
        "        if len(text) > 0:\n",
        "          for line in text.splitlines(True):\n",
        "            tokens = tokenizer(line, return_tensors=\"pt\").input_ids.to(\"cpu\")\n",
        "            actions.append((line, tokens))\n",
        "        action = \"continue\"\n",
        "      if action == \"info\":\n",
        "        clear_output()\n",
        "        print(\"LMI: \" + lmi[0])\n",
        "        print(\"LMI tokens: \" + str(lmi[1].shape[1]))\n",
        "        print(\"Memory: \" + memory[0])\n",
        "        print(\"Author's note: \" + an[0])\n",
        "        print(\"Author's note depth: \" + str(an_depth))\n",
        "        need_refresh = True\n",
        "      if action == \"history\":\n",
        "        clear_output()\n",
        "        print(\"\".join([action[0] for action in actions]), end=\"\")\n",
        "        need_refresh = False\n",
        "      if action == \"retry\":\n",
        "        if len(actions) > 0:\n",
        "          actions.pop()\n",
        "        need_refresh = True\n",
        "        action = \"continue\"\n",
        "      if action == \"undo\":\n",
        "        if len(actions) > 0:\n",
        "          actions.pop()\n",
        "        assemble()\n",
        "        clear_output()\n",
        "        print(\"\".join([action[0] for action in actions[-number_show_last_actions:]]), end=\"\")\n",
        "        need_refresh = False\n",
        "      if action == \"memory\":\n",
        "        if prompt == \"\":\n",
        "          memory = (\"\", torch.zeros((1, 0)).long())\n",
        "          text = \"\"\n",
        "        else:\n",
        "          text = codecs.decode(prompt + \"\\n\", \"unicode-escape\")\n",
        "          tokens = tokenizer(text, return_tensors=\"pt\").input_ids.to(\"cpu\")\n",
        "          memory = (text, tokens)\n",
        "        clear_output()\n",
        "        print(\"Memory: \" + text)\n",
        "        memory_updated = True\n",
        "      if action == \"authorsnote\":\n",
        "        if prompt == \"\":\n",
        "          an = (\"\", torch.zeros((1, 0)).long())\n",
        "          text = \"\"\n",
        "        else:\n",
        "          text = \"\\n[Author's note: \" + codecs.decode(prompt, \"unicode-escape\") + \"]\\n\"\n",
        "          tokens = tokenizer(text, return_tensors=\"pt\").input_ids.to(\"cpu\")\n",
        "          an = (text, tokens)\n",
        "        clear_output()\n",
        "        print(\"Author's note: \" + text)\n",
        "        an_updated = True\n",
        "      if action == \"andepth\":\n",
        "        clear_output()\n",
        "        try:\n",
        "          an_depth = int(codecs.decode(prompt + \"\\n\", \"unicode-escape\"))\n",
        "        except:\n",
        "          pass\n",
        "        print(\"Author's note depth: \" + str(an_depth))\n",
        "        an_updated = True\n",
        "      if action == \"tokenize\":\n",
        "        text = codecs.decode(prompt, \"unicode-escape\")\n",
        "        tokens = tokenizer(text, return_tensors=\"pt\").input_ids.to(\"cpu\")\n",
        "        clear_output()\n",
        "        print(\"Tokens: \" + str(tokens.shape[1]))\n",
        "        print(tokens[0])\n",
        "        need_refresh = True\n",
        "      if action == \"continue\":\n",
        "        assemble()\n",
        "        ids = lmi[1].cuda()\n",
        "        n_ids = ids.shape[1]\n",
        "        if n_ids < 1:\n",
        "          n_ids = 1\n",
        "          ids = torch.tensor([[tokenizer.eos_token_id]])\n",
        "        max_length = number_generated_tokens + n_ids\n",
        "        #ids[:, :] = 13\n",
        "        torch.cuda.empty_cache()\n",
        "        clear_output()\n",
        "        gen_tokens = model.generate(\n",
        "            ids.long().cuda(),\n",
        "            do_sample=True,\n",
        "            min_length=max_length,\n",
        "            max_length=max_length,\n",
        "            temperature=temperature,\n",
        "            top_k = top_k,\n",
        "            top_p = top_p,\n",
        "            repetition_penalty = repetition_penalty,\n",
        "            repetition_penalty_range = repetition_penalty_range,\n",
        "            repetition_penalty_slope = repetition_penalty_slope,\n",
        "            use_cache=True,\n",
        "            pad_token_id=tokenizer.eos_token_id\n",
        "        ).long()\n",
        "        stop_tokens = [0, 13, 30, 526, 764, 1701, 2474, 5145, 5633]\n",
        "        for i in reversed(range(len(gen_tokens[0]))):\n",
        "          if i < n_ids:\n",
        "            gen_tokens = gen_tokens[0]\n",
        "            break\n",
        "          if gen_tokens[0][i] in stop_tokens:\n",
        "            gen_tokens = gen_tokens[0][:i+1]\n",
        "            break\n",
        "        gen_text = tokenizer.decode(gen_tokens[n_ids:])\n",
        "        if len(gen_text) > 0:\n",
        "          actions.append((gen_text, gen_tokens[n_ids:].unsqueeze(0).cpu()))\n",
        "        print(\"\".join([action[0] for action in actions[-number_show_last_actions:]]), end=\"\")\n",
        "        torch.cuda.empty_cache()\n",
        "        need_refresh = False\n",
        "    if history is not None:\n",
        "      if history:\n",
        "        with out_history:\n",
        "          IPython.display.clear_output()\n",
        "          print(\"\".join([action[0] for action in actions]), end=\"\")\n",
        "        with out_history2:\n",
        "          IPython.display.clear_output()\n",
        "      else:\n",
        "        with out_history2:\n",
        "          IPython.display.clear_output()\n",
        "          print(\"\".join([action[0] for action in actions]), end=\"\")\n",
        "        with out_history:\n",
        "          IPython.display.clear_output()\n",
        "      if an_updated:\n",
        "        with out_an:\n",
        "          IPython.display.clear_output()\n",
        "          if len(an[0]) > 0:\n",
        "            print(\"AN depth: \" + str(an_depth) + \"\\n\" + an[0], end=\"\")\n",
        "      if memory_updated:\n",
        "        with out_memory:\n",
        "          IPython.display.clear_output()\n",
        "          print(memory[0], end=\"\")\n",
        "      history = not history\n",
        "\n",
        "import ipywidgets as widgets\n",
        "import IPython.display\n",
        "out = widgets.Output(layout={'border': '1px solid black', \"width\": \"1280px\"})\n",
        "dropdown = widgets.Dropdown(options=action_types, value=action_type, description='Action:', disabled=False)\n",
        "dropdown.observe(set_action, 'value')\n",
        "button = widgets.Button(description='[selected action]', disabled=False)\n",
        "button.on_click(lambda _: play(dropdown.value))\n",
        "retry_button = widgets.Button(description='Retry', disabled=False)\n",
        "retry_button.on_click(lambda _: play(\"retry\"))\n",
        "continue_button = widgets.Button(description='Continue', disabled=False)\n",
        "continue_button.on_click(lambda _: play(\"continue\"))\n",
        "undo_button = widgets.Button(description='Undo', disabled=False)\n",
        "undo_button.on_click(lambda _: play(\"undo\"))\n",
        "hbox = widgets.HBox([button, retry_button, continue_button, undo_button])\n",
        "input = widgets.Textarea(value='', placeholder='', description='Input:', disabled=False, rows=4, layout={\"width\": \"1280px\"})\n",
        "input.observe(set_prompt, 'value')\n",
        "\n",
        "display(out, dropdown, hbox, input)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "jPtYoUNO39Mg"
      },
      "source": [
        "#@title History\n",
        "#@markdown Run this cell to have an auto-updating full listing of the current story.\n",
        "\n",
        "history = True\n",
        "out_history = widgets.Output(layout={'border': '1px solid black', \"width\": \"1280px\"})\n",
        "out_history2 = widgets.Output(layout={'border': '1px solid black', \"width\": \"1280px\"})\n",
        "out_memory = widgets.Output(layout={'border': '1px solid black', \"width\": \"1280px\"})\n",
        "out_an = widgets.Output(layout={'border': '1px solid black', \"width\": \"1280px\"})\n",
        "display(out_history, out_history2, out_memory, out_an)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "qvIvzStC9rY9"
      },
      "source": [
        "#@title Attention display\n",
        "#@markdown If you don't know what this is, just ignore. attn_head_combination selects the operation used to combine layer and attention head results. Only global attention is included.\n",
        "\n",
        "visualizer_prompt = \"Before his eyes was an orange cat with stripes. Running his fingers through its soft fur, he admired\" #@param {type:\"string\"}\n",
        "max_attentions =  8#@param {type:\"integer\"}\n",
        "attn_head_combination = \"mean\" #@param [\"mean\", \"max\"]\n",
        "\n",
        "ids = tokenizer(visualizer_prompt, return_tensors=\"pt\").input_ids.to(\"cpu\")\n",
        "n_ids = ids.shape[1]\n",
        "if n_ids < 1:\n",
        "  n_ids = 1\n",
        "  ids = torch.tensor([[tokenizer.eos_token_id]])\n",
        "max_length = n_ids + number_generated_tokens\n",
        "torch.cuda.empty_cache()\n",
        "basic_output = model.generate(\n",
        "    ids.long().cuda(),\n",
        "    do_sample=True,\n",
        "    min_length=max_length,\n",
        "    max_length=max_length,\n",
        "    temperature=temperature,\n",
        "    top_k = top_k,\n",
        "    top_p = top_p,\n",
        "    repetition_penalty = repetition_penalty,\n",
        "    repetition_penalty_range = repetition_penalty_range,\n",
        "    repetition_penalty_slope = repetition_penalty_slope,\n",
        "    use_cache=True,\n",
        "    pad_token_id=tokenizer.eos_token_id,\n",
        "    return_dict_in_generate=True,\n",
        "    output_attentions=True\n",
        ")\n",
        "torch.cuda.empty_cache()\n",
        "attentions = basic_output[\"attentions\"]\n",
        "basic_output = basic_output[\"sequences\"].cpu()\n",
        "\n",
        "print(\"Prompt: \" + visualizer_prompt)\n",
        "print()\n",
        "\n",
        "def combine(attentions):\n",
        "  if attn_head_combination == \"mean\":\n",
        "    attentions = attentions.mean(0)\n",
        "  else:\n",
        "    attentions = attentions.max(0)[0]\n",
        "  return attentions\n",
        "\n",
        "import torch.nn.functional as F\n",
        "import numpy as np\n",
        "#gen_tokens x layers x tensor (1 x heads x n_ids square)\n",
        "torch.set_printoptions(sci_mode=False)\n",
        "for i in range(number_generated_tokens):\n",
        "  layer_attn = []\n",
        "  for j in range(0,len(attentions[i]),2):\n",
        "    layer_attn.append(combine(attentions[i][j][0,:,-1].float().cpu()))\n",
        "  layer_attn = torch.stack(layer_attn)\n",
        "  token_attn = combine(layer_attn)\n",
        "  prob, topk = token_attn.topk(max_attentions)\n",
        "  top_tokens = []\n",
        "  for top in topk:\n",
        "    decoded = tokenizer.decode(torch.tensor([basic_output[0][top]]))\n",
        "    top_tokens.append(f\"{decoded!r}@{top} ({token_attn[top]:.4f})\")\n",
        "  print(\"Token: \" + repr(tokenizer.decode(torch.tensor([basic_output[0][n_ids + i]]))))\n",
        "  print(\", \".join(top_tokens))\n",
        "  print()\n",
        "\n",
        "print(\"Output: \" + tokenizer.decode(basic_output[0]))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}