Llama_Colab

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard",
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "Initialize repository, copy weights from Google drive."
      ],
      "metadata": {
        "id": "uh_RON9_FqBV"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aChaziGm-OrN"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import sys\n",
        "from google.colab import drive\n",
        "\n",
        "# Mount google drive.\n",
        "drive.mount('/gdrive')\n",
        "\n",
        "#@markdown Location of tokenizer.\n",
        "tokenizer_loc = '/gdrive/My Drive/tokenizer.model' #@param {type:\"string\"}\n",
        "\n",
        "# @markdown Location of directory containing model weights / parameters.\n",
        "weight_loc = '/gdrive/My Drive/7B/' #@param {type:\"string\"}\n",
        "\n",
        "!pip install fairscale\n",
        "!pip install sentencepiece\n",
        "!git clone https://github.com/facebookresearch/llama.git\n",
        "\n",
        "sys.path.insert(0, '/content/llama/')\n",
        "\n",
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The 7B checkpoint is too large to fit into RAM. Run this cell if you need to split the 7B checkpoint. Will save the results to your 7B directory so you should only ever need to run this cell once. You may need to restart the runtime afterward."
      ],
      "metadata": {
        "id": "YUViS0koD_aj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "\n",
        "checkpoint = torch.load(os.path.join(weight_loc, 'consolidated.00.pth'),\n",
        "                        map_location=\"cuda\")\n",
        "\n",
        "d1 = dict(list(checkpoint.items())[:len(checkpoint)//2])\n",
        "torch.save(d1, os.path.join(weight_loc, 'consolidated.00.00.pth'))\n",
        "del(d1)\n",
        "\n",
        "d2 = dict(list(checkpoint.items())[len(checkpoint)//2:])\n",
        "torch.save(d2, os.path.join(weight_loc, 'consolidated.00.01.pth'))\n",
        "del(d2)\n",
        "\n",
        "del(checkpoint)"
      ],
      "metadata": {
        "id": "MxwFpC1fCAdz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Include that one anon's additional sampling methods so we have Kobold parameters like repetition penalty, tfs, etc."
      ],
      "metadata": {
        "id": "BMbLXEqjcmdi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Copyright (c) Meta Platforms, Inc. and affiliates.\n",
        "# This software may be used and distributed according to the terms of the GNU\n",
        "# General Public License version 3.\n",
        "\n",
        "from typing import List\n",
        "\n",
        "import torch\n",
        "\n",
        "from llama.tokenizer import Tokenizer\n",
        "from llama.model import Transformer\n",
        "\n",
        "class LLaMA:\n",
        "    def __init__(self, model: Transformer, tokenizer: Tokenizer):\n",
        "        self.model = model\n",
        "        self.tokenizer = tokenizer\n",
        "\n",
        "    def generate(\n",
        "        self,\n",
        "        prompts: List[str],\n",
        "        max_gen_len: int,\n",
        "        temperature: float = 0.8,\n",
        "        top_p: float = 0.95,\n",
        "        tfs: float = 1.0,\n",
        "        typical: float = 1.0,\n",
        "        penalty_range: float = 1024,\n",
        "        penalty_slope: float = 0.7,\n",
        "        penalty: float = 1.1\n",
        "    ) -> List[str]:\n",
        "        bsz = len(prompts)\n",
        "        params = self.model.params\n",
        "        assert bsz <= params.max_batch_size, (bsz, params.max_batch_size)\n",
        "\n",
        "        prompt_tokens = [self.tokenizer.encode(x, bos=True, eos=False)\n",
        "                         for x in prompts]\n",
        "\n",
        "        min_prompt_size = min([len(t) for t in prompt_tokens])\n",
        "        max_prompt_size = max([len(t) for t in prompt_tokens])\n",
        "\n",
        "        total_len = min(params.max_seq_len, max_gen_len + max_prompt_size)\n",
        "\n",
        "        tokens = torch.full((bsz, total_len),\n",
        "                            self.tokenizer.pad_id).cuda().long()\n",
        "        for k, t in enumerate(prompt_tokens):\n",
        "            tokens[k, : len(t)] = torch.tensor(t).long()\n",
        "        input_text_mask = tokens != self.tokenizer.pad_id\n",
        "        start_pos = min_prompt_size\n",
        "        prev_pos = 0\n",
        "        for cur_pos in range(start_pos, total_len):\n",
        "            input_ids = tokens[:, prev_pos:cur_pos]\n",
        "            logits = self.model.forward(input_ids, prev_pos)\n",
        "            if temperature > 0:\n",
        "\n",
        "                next_token_scores = sample_top_p_actual(input_ids, logits,\n",
        "                                                        top_p)\n",
        "                next_token_scores = sample_tail_free(input_ids,\n",
        "                                                     next_token_scores, tfs)\n",
        "                next_token_scores = sample_typical(input_ids, next_token_scores,\n",
        "                                                   typical)\n",
        "                next_token_scores = sample_temperature(input_ids,\n",
        "                                                       next_token_scores,\n",
        "                                                       temperature)\n",
        "                next_token_scores = sample_advanced_repetition_penalty(input_ids,\n",
        "                                                                       next_token_scores,\n",
        "                                                                       penalty_range,\n",
        "                                                                       penalty_slope,\n",
        "                                                                       penalty)\n",
        "\n",
        "                next_token_scores = torch.nn.functional.softmax(next_token_scores,\n",
        "                                                                dim=-1)\n",
        "                next_token = torch.multinomial(next_token_scores,\n",
        "                                               num_samples=1).squeeze(1)\n",
        "            else:\n",
        "                next_token = torch.argmax(logits, dim=-1)\n",
        "            next_token = next_token.reshape(-1)\n",
        "            # only replace token if prompt has already been generated\n",
        "            next_token = torch.where(\n",
        "                input_text_mask[:, cur_pos], tokens[:, cur_pos], next_token\n",
        "            )\n",
        "            tokens[:, cur_pos] = next_token\n",
        "            prev_pos = cur_pos\n",
        "\n",
        "        decoded = []\n",
        "        for i, t in enumerate(tokens.tolist()):\n",
        "            # cut to max gen len\n",
        "            t = t[: len(prompt_tokens[i]) + max_gen_len]\n",
        "            # cut to eos tok if any\n",
        "            try:\n",
        "                t = t[: t.index(self.tokenizer.eos_id)]\n",
        "            except ValueError:\n",
        "                pass\n",
        "            decoded.append(self.tokenizer.decode(t))\n",
        "        return decoded\n",
        "\n",
        "# taken from Kobold and transformers so this stuff is AGPL I guess\n",
        "def sample_temperature(input_ids, scores, tempt):\n",
        "    scores = scores / tempt\n",
        "    return scores\n",
        "\n",
        "def sample_typical(input_ids, scores, typical, filter_value = -float(\"Inf\"),\n",
        "                   min_tokens_to_keep = 1):\n",
        "    if filter_value >= 1.0:\n",
        "        return scores\n",
        "\n",
        "    probs = scores.softmax(dim=-1)\n",
        "    log_probs = probs.log()\n",
        "\n",
        "    neg_entropy = (probs * log_probs).nansum(dim=-1, keepdim=True)\n",
        "\n",
        "    entropy_deviation = (neg_entropy - log_probs).abs()\n",
        "\n",
        "    _, sorted_indices = torch.sort(entropy_deviation)\n",
        "    sorted_logits = probs.gather(-1, sorted_indices)\n",
        "    sorted_indices_to_remove = sorted_logits.cumsum(dim=-1) >= typical\n",
        "    sorted_indices_to_remove = sorted_indices_to_remove.roll(1, dims=-1)\n",
        "\n",
        "    min_tokens_to_keep = max(min_tokens_to_keep, 1)\n",
        "    # Keep at least min_tokens_to_keep\n",
        "    sorted_indices_to_remove[..., : min_tokens_to_keep] = 0\n",
        "\n",
        "    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)\n",
        "    scores = scores.masked_fill(indices_to_remove, filter_value)\n",
        "    return scores    \n",
        "\n",
        "def sample_top_p_actual(input_ids, scores, top_p, filter_value = -float(\"Inf\"),\n",
        "                        min_tokens_to_keep = 1):\n",
        "    sorted_logits, sorted_indices = torch.sort(scores, descending=False)\n",
        "    cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)\n",
        "\n",
        "    # Remove tokens with cumulative top_p above the threshold (token with 0 are kept)\n",
        "    sorted_indices_to_remove = cumulative_probs <= (1 - top_p)\n",
        "    if min_tokens_to_keep > 1:\n",
        "        # Keep at least min_tokens_to_keep\n",
        "        sorted_indices_to_remove[..., -min_tokens_to_keep :] = 0\n",
        "\n",
        "    # scatter sorted tensors to original indexing\n",
        "    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices,\n",
        "                                                         sorted_indices_to_remove)\n",
        "    scores = scores.masked_fill(indices_to_remove, filter_value)\n",
        "    return scores\n",
        "\n",
        "def sample_advanced_repetition_penalty(input_ids, scores, penalty_range,\n",
        "                                       penalty_slope, penalty):\n",
        "    penalty_range = int(penalty_range)\n",
        "    clipped_penalty_range = min(input_ids.shape[-1], penalty_range)\n",
        "\n",
        "    if penalty != 1.0:\n",
        "        if penalty_range > 0:\n",
        "            if clipped_penalty_range < input_ids.shape[1]:\n",
        "                input_ids = input_ids[..., -clipped_penalty_range:]\n",
        "\n",
        "            if penalty_slope != 0:\n",
        "                _penalty = (torch.arange(penalty_range, dtype=scores.dtype,\n",
        "                                         device=scores.device)/(penalty_range - 1)) * 2. - 1\n",
        "                _penalty = (penalty_slope * _penalty) / (1 + torch.abs(_penalty) * (penalty_slope - 1))\n",
        "                _penalty = 1 + ((_penalty + 1) / 2).unsqueeze(0) * (penalty - 1)\n",
        "                penalty = _penalty[..., -clipped_penalty_range:]\n",
        "\n",
        "        score = torch.gather(scores, 1, input_ids)\n",
        "        score = torch.where(score <= 0, score * penalty, score / penalty)\n",
        "        scores.scatter_(1, input_ids, score)\n",
        "\n",
        "        return scores    \n",
        "\n",
        "def sample_top_a(input_ids, scores, top_a, filter_value = -float(\"Inf\"),\n",
        "                 min_tokens_to_keep = 1):\n",
        "    if filter_value >= 1.0:\n",
        "        return scores\n",
        "\n",
        "    sorted_logits, sorted_indices = torch.sort(scores, descending=True)\n",
        "    probs = sorted_logits.softmax(dim=-1)\n",
        "\n",
        "    # Remove tokens with probability less than top_a*(max(probs))^2 (token with 0 are kept)\n",
        "    probs_max = probs[..., 0, None]\n",
        "    sorted_indices_to_remove = probs < probs_max * probs_max * top_a\n",
        "\n",
        "    if min_tokens_to_keep > 1:\n",
        "        # Keep at least min_tokens_to_keep\n",
        "        sorted_indices_to_remove[..., : min_tokens_to_keep] = 0\n",
        "\n",
        "    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices,\n",
        "                                                         sorted_indices_to_remove)\n",
        "    scores = scores.masked_fill(indices_to_remove, filter_value)\n",
        "    return scores    \n",
        "\n",
        "def sample_tail_free(input_ids, scores, tfs, filter_value = -float(\"Inf\"),\n",
        "                     min_tokens_to_keep = 1):\n",
        "    if filter_value >= 1.0:\n",
        "        return scores\n",
        "    sorted_logits, sorted_indices = torch.sort(scores, descending=True)\n",
        "    probs = sorted_logits.softmax(dim=-1)\n",
        "\n",
        "    # Compute second derivative normalized CDF\n",
        "    d2 = probs.diff().diff().abs()\n",
        "    normalized_d2 = d2 / d2.sum(dim=-1, keepdim=True)\n",
        "    normalized_d2_cdf = normalized_d2.cumsum(dim=-1)\n",
        "\n",
        "    # Remove tokens with CDF value above the threshold (token with 0 are kept)\n",
        "    sorted_indices_to_remove = normalized_d2_cdf > tfs\n",
        "\n",
        "    # Centre the distribution around the cutoff as in the original implementation of the algorithm\n",
        "    sorted_indices_to_remove = torch.cat(\n",
        "        (\n",
        "            torch.zeros(scores.shape[0], 1, dtype=torch.bool,\n",
        "                        device=scores.device),\n",
        "            sorted_indices_to_remove,\n",
        "            torch.ones(scores.shape[0], 1, dtype=torch.bool,\n",
        "                       device=scores.device),\n",
        "        ),\n",
        "        dim=-1,\n",
        "    )\n",
        "\n",
        "    if min_tokens_to_keep > 1:\n",
        "        # Keep at least min_tokens_to_keep\n",
        "        sorted_indices_to_remove[..., : min_tokens_to_keep] = 0\n",
        "\n",
        "    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices,\n",
        "                                                         sorted_indices_to_remove)\n",
        "    scores = scores.masked_fill(indices_to_remove, filter_value)\n",
        "    return scores"
      ],
      "metadata": {
        "id": "fYMxCH_Zcajj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Load model."
      ],
      "metadata": {
        "id": "47yXVODMO6l0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from typing import Tuple\n",
        "import os\n",
        "import sys\n",
        "import torch\n",
        "import time\n",
        "import json\n",
        "\n",
        "from pathlib import Path\n",
        "\n",
        "from fairscale.nn.model_parallel.initialize import initialize_model_parallel\n",
        "\n",
        "from llama.model import ModelArgs, Transformer\n",
        "from llama.tokenizer import Tokenizer\n",
        "\n",
        "\n",
        "os.environ['RANK'] = '0'\n",
        "os.environ['WORLD_SIZE'] = '1'\n",
        "os.environ['MP'] = '1'\n",
        "os.environ['MASTER_ADDR'] = '127.0.0.1'\n",
        "os.environ['MASTER_PORT'] = '2223'\n",
        "\n",
        "\n",
        "def setup_model_parallel() -> Tuple[int, int]:\n",
        "    local_rank = int(os.environ.get(\"LOCAL_RANK\", -1))\n",
        "    world_size = int(os.environ.get(\"WORLD_SIZE\", -1))\n",
        "\n",
        "    torch.distributed.init_process_group(\"gloo\")\n",
        "    initialize_model_parallel(world_size)\n",
        "    torch.cuda.set_device(local_rank)\n",
        "\n",
        "    # seed must be the same in all processes\n",
        "    torch.manual_seed(1)\n",
        "    return local_rank, world_size\n",
        "\n",
        "\n",
        "'''\n",
        "def load(ckpt_dir: str, tokenizer_path: str, local_rank: int, world_size: int,\n",
        "         max_seq_len: int, max_batch_size: int) -> LLaMA:\n",
        "    start_time = time.time()\n",
        "    checkpoints = sorted(Path(ckpt_dir).glob(\"*.pth\"))\n",
        "    assert (\n",
        "        world_size == len(checkpoints)\n",
        "    ), f\"Loading a checkpoint for MP={len(checkpoints)} but world size is\n",
        "        {world_size}\"\n",
        "    ckpt_path = checkpoints[local_rank]\n",
        "    print(\"Loading\")\n",
        "    \n",
        "    checkpoint = torch.load(ckpt_path, map_location=\"cpu\")\n",
        "    with open(Path(ckpt_dir) / \"params.json\", \"r\") as f:\n",
        "        params = json.loads(f.read())\n",
        "\n",
        "    model_args: ModelArgs = ModelArgs(max_seq_len=max_seq_len,\n",
        "                                      max_batch_size=max_batch_size,\n",
        "                                      **params)\n",
        "    tokenizer = Tokenizer(model_path=tokenizer_path)\n",
        "    model_args.vocab_size = tokenizer.n_words\n",
        "    torch.set_default_tensor_type(torch.cuda.HalfTensor)\n",
        "    model = Transformer(model_args).cuda().half()\n",
        "    torch.set_default_tensor_type(torch.FloatTensor)\n",
        "    model.load_state_dict(checkpoint, strict=False)\n",
        "\n",
        "    generator = LLaMA(model, tokenizer)\n",
        "    print(f\"Loaded in {time.time() - start_time:.2f} seconds\")\n",
        "    return generator\n",
        "'''\n",
        "\n",
        "\n",
        "def load(ckpt_dir: str, tokenizer_path: str, local_rank: int, world_size: int,\n",
        "         max_seq_len: int, max_batch_size: int) -> LLaMA:\n",
        "    start_time = time.time()\n",
        "    \n",
        "    print(\"Loading\")\n",
        "    with open(Path(ckpt_dir) / \"params.json\", \"r\") as f:\n",
        "        params = json.loads(f.read())\n",
        "\n",
        "    model_args: ModelArgs = ModelArgs(max_seq_len=max_seq_len,\n",
        "                                      max_batch_size=max_batch_size,\n",
        "                                      **params)\n",
        "    tokenizer = Tokenizer(model_path=tokenizer_path)\n",
        "    model_args.vocab_size = tokenizer.n_words\n",
        "    torch.set_default_tensor_type(torch.cuda.HalfTensor)\n",
        "    model = Transformer(model_args).cuda().half()\n",
        "    torch.set_default_tensor_type(torch.FloatTensor)\n",
        "\n",
        "    checkpoint_paths = [os.path.join(ckpt_dir, 'consolidated.00.00.pth'),\n",
        "                        os.path.join(ckpt_dir, 'consolidated.00.01.pth')]\n",
        "    \n",
        "    for checkpoint_path in checkpoint_paths:\n",
        "        checkpoint = torch.load(checkpoint_path, map_location='cpu')\n",
        "        model.load_state_dict(checkpoint, strict=False)\n",
        "        del checkpoint\n",
        "\n",
        "    generator = LLaMA(model, tokenizer)\n",
        "    print(f\"Loaded in {time.time() - start_time:.2f} seconds\")\n",
        "    return generator\n",
        "\n",
        "# @markdown Context size. Can be up to 2048, but Colab GPU doesn't always play well with high values.\n",
        "max_seq_len = 1024 # @param {type:\"number\"}\n",
        "max_batch_size = 1\n",
        "\n",
        "local_rank, world_size = setup_model_parallel()\n",
        "if local_rank > 0:\n",
        "    sys.stdout = open(os.devnull, 'w')\n",
        "\n",
        "generator = load(weight_loc, tokenizer_loc, local_rank, world_size,\n",
        "                 max_seq_len, max_batch_size)\n",
        "tokenizer = generator.tokenizer"
      ],
      "metadata": {
        "id": "U5eiyiLJMNpz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Main GUI. If you change the presets, you'll have to reload the cell for the changes to take effect."
      ],
      "metadata": {
        "id": "yfgItXzVPGGS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import ipywidgets as widgets\n",
        "from IPython.display import display\n",
        "import time\n",
        "\n",
        "max_gen_len = 64 #@param {type:\"number\"}\n",
        "temperature = 0.8 #@param {type:\"number\"}\n",
        "top_p = 0.95 #@param {type:\"number\"}\n",
        "tfs = 1.0 #@param {type:\"number\"}\n",
        "typical = 1.0 #@param {type:\"number\"}\n",
        "penalty_range = 1024 #@param {type:\"number\"}\n",
        "penalty_slope = 0.7 #@param {type:\"number\"}\n",
        "penalty = 1.1 #@param {type:\"number\"}\n",
        "\n",
        "input_text_area = widgets.Textarea(placeholder='Enter a prompt...',\n",
        "                                   layout=widgets.Layout(width='1200px',\n",
        "                                                         height='600px'))\n",
        "send_button = widgets.Button(description='Send')\n",
        "undo_button = widgets.Button(description='Undo')\n",
        "redo_button = widgets.Button(description='Redo')\n",
        "retry_button = widgets.Button(description='Retry')\n",
        "memory_button = widgets.ToggleButton(description='Memory')\n",
        "\n",
        "hbox = widgets.HBox([input_text_area,\n",
        "                     widgets.VBox([send_button, undo_button, redo_button,\n",
        "                                  retry_button, memory_button])])\n",
        "output = widgets.Output()\n",
        "\n",
        "undo_button.disabled = True\n",
        "redo_button.disabled = True\n",
        "retry_button.disabled = True\n",
        "\n",
        "listen_for_updates = False\n",
        "cur_outputs = []\n",
        "cur_outputs_idx = -1\n",
        "memory_text = ''\n",
        "input_text = ''\n",
        "\n",
        "def generate():\n",
        "    # When creating the context, first, place the full memory followed by a\n",
        "    # newline.\n",
        "    #\n",
        "    # Next, taking the last (max_seq_len-1-max_gen_len-len(mem)) tokens,\n",
        "    # place these tokens in the context.\n",
        "    \n",
        "    if memory_text:\n",
        "        mem_tokenized = tokenizer.encode(memory_text + '\\n', bos=False, eos=False)\n",
        "    else:\n",
        "        mem_tokenized = []\n",
        "    \n",
        "    inp_tokenized = tokenizer.encode(input_text_area.value, bos=False, eos=False)\n",
        "    num_inp_tokens = max(max_seq_len-1-max_gen_len-len(mem_tokenized), 0)\n",
        "\n",
        "    if num_inp_tokens > 0:\n",
        "        tokenized = mem_tokenized + inp_tokenized[-num_inp_tokens:]\n",
        "    elif len(mem_tokenized) > 0:\n",
        "        num_mem_tokens = max_seq_len-1-max_gen_len\n",
        "        tokenized = mem_tokenized[-num_mem_tokens:]\n",
        "    else:\n",
        "        tokenized = []\n",
        "    \n",
        "    detokenized = tokenizer.decode(tokenized)\n",
        "    output = generator.generate([detokenized],\n",
        "                                max_gen_len=max_gen_len,\n",
        "                                temperature=temperature,\n",
        "                                top_p=top_p,\n",
        "                                tfs=tfs,\n",
        "                                typical=typical,\n",
        "                                penalty_range=penalty_range,\n",
        "                                penalty_slope=penalty_slope,\n",
        "                                penalty=penalty)\n",
        "\n",
        "    num_characters = len(output) - len(detokenized) - 1\n",
        "    return output[0][-num_characters:]\n",
        "\n",
        "    '''\n",
        "    tokenized = tokenizer.encode(input_text_area.value, bos=True, eos=False)\n",
        "    detokenized = tokenizer.decode(tokenized[-(max_seq_len-1-max_gen_len):])\n",
        "    output = generator.generate([detokenized],\n",
        "                                max_gen_len=max_gen_len,\n",
        "                                temperature=temperature,\n",
        "                                top_p=top_p,\n",
        "                                tfs=tfs,\n",
        "                                typical=typical,\n",
        "                                penalty_range=penalty_range,\n",
        "                                penalty_slope=penalty_slope,\n",
        "                                penalty=penalty)\n",
        "    num_characters = len(output) - len(detokenized) - 1\n",
        "    return output[0][-num_characters:]\n",
        "    '''\n",
        "\n",
        "def on_update_input_text_area(change):\n",
        "    global listen_for_updates, cur_outputs, cur_outputs_idx\n",
        "\n",
        "    if listen_for_updates:\n",
        "        cur_outputs = []\n",
        "        cur_outputs_idx = -1\n",
        "        undo_button.disabled = True\n",
        "        redo_button.disabled = True\n",
        "        retry_button.disabled = True\n",
        "\n",
        "def send():\n",
        "    global listen_for_updates, cur_outputs, cur_outputs_idx\n",
        "\n",
        "    input_text_area.disabled = True\n",
        "    memory_button.disabled = True\n",
        "    listen_for_updates = False\n",
        "\n",
        "    generation = generate()\n",
        "    input_text_area.value += generation\n",
        "    cur_outputs_idx += 1\n",
        "    cur_outputs = cur_outputs[:cur_outputs_idx]\n",
        "    cur_outputs.append(generation)\n",
        "\n",
        "    undo_button.disabled = False\n",
        "    redo_button.disabled = True\n",
        "    retry_button.disabled = False\n",
        "    listen_for_updates = True\n",
        "    memory_button.disabled = False\n",
        "    input_text_area.disabled = False\n",
        "\n",
        "def undo():\n",
        "    global listen_for_updates, cur_outputs, cur_outputs_idx\n",
        "\n",
        "    listen_for_updates = False\n",
        "    num_chars = len(cur_outputs[cur_outputs_idx])\n",
        "    input_text_area.value = input_text_area.value[:-num_chars]\n",
        "    cur_outputs_idx -= 1\n",
        "\n",
        "    if cur_outputs_idx == -1:\n",
        "        undo_button.disabled = True\n",
        "        retry_button.disabled = True\n",
        "    if len(cur_outputs) > 0:\n",
        "        redo_button.disabled = False\n",
        "\n",
        "    listen_for_updates = True\n",
        "\n",
        "def redo():\n",
        "    global listen_for_updates, cur_outputs, cur_outputs_idx\n",
        "\n",
        "    listen_for_updates = False\n",
        "    input_text_area.value += cur_outputs[cur_outputs_idx+1]\n",
        "    cur_outputs_idx += 1\n",
        "\n",
        "    if cur_outputs_idx == len(cur_outputs) - 1:\n",
        "        redo_button.disabled = True\n",
        "    if len(cur_outputs) > 0:\n",
        "        undo_button.disabled = False\n",
        "        retry_button.disabled = False\n",
        "\n",
        "    listen_for_updates = True\n",
        "\n",
        "def send_button_clicked(b):\n",
        "    send()\n",
        "\n",
        "def undo_button_clicked(b):\n",
        "    undo()\n",
        "\n",
        "def redo_button_clicked(b):\n",
        "    redo()\n",
        "\n",
        "def retry_button_clicked(b):\n",
        "    undo()\n",
        "    send()\n",
        "\n",
        "def memory_button_clicked(b):\n",
        "    global listen_for_updates, cur_outputs, cur_outputs_idx, memory_text, \\\n",
        "           input_text\n",
        "    if memory_button.value:\n",
        "        listen_for_updates = False\n",
        "        send_button.disabled = True\n",
        "        undo_button.disabled = True\n",
        "        redo_button.disabled = True\n",
        "        retry_button.disabled = True\n",
        "        input_text = input_text_area.value\n",
        "        input_text_area.value = memory_text\n",
        "    else:\n",
        "        memory_text = input_text_area.value\n",
        "        input_text_area.value = input_text\n",
        "        input_text = ''\n",
        "        send_button.disabled = False\n",
        "        undo_button.disabled = cur_outputs_idx < 0\n",
        "        redo_button.disabled = cur_outputs_idx >= len(cur_outputs) - 1\n",
        "        retry_button.disabled = undo_button.disabled\n",
        "        listen_for_updates = True\n",
        "\n",
        "send_button.on_click(send_button_clicked)\n",
        "undo_button.on_click(undo_button_clicked)\n",
        "redo_button.on_click(redo_button_clicked)\n",
        "retry_button.on_click(retry_button_clicked)\n",
        "memory_button.observe(memory_button_clicked, names='value')\n",
        "input_text_area.observe(on_update_input_text_area, names='value')\n",
        "\n",
        "display(hbox, output)"
      ],
      "metadata": {
        "id": "RRpoZt05O4vx"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}