PromptArray.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "PromptArray.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyOKXzE34N6meiOstED+R9Ms",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jeffbinder/promptarray/blob/main/PromptArray.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "35TQLi4rQ_gh"
      },
      "source": [
        "# PromptArray: A Prompting Language for Neural Text Generators\n",
        "\n",
        "This notebook allows you to experiment with PromptArray, a system for controlling the output of neural text generators. Text generators are usually controlled by prompts: input text that indicates what the model should do. For instance, if you want a description of a species of snake, you can enter the following:\n",
        "\n",
        "> Scientists recently discovered a new species of snake. Here is a description of it:\n",
        "\n",
        "The machine will then generate a completion of this text, which usually consists of something like the desired description. However, engineering effective prompts is not always straightforward; in particular, it is very hard to design a prompt that effectively tells the generator *not* to do something.\n",
        "\n",
        "PromptArray allows you to include operators in your prompt that manipulate the machine's predictions. At present, these five operators are available:\n",
        "\n",
        "| Operator | Meaning |\n",
        "| --- | --- |\n",
        "| A&B | A and B |\n",
        "| A\\|B | A or B |\n",
        "| A^B | A and not B |\n",
        "| A/B | A more than B |\n",
        "| A~B | A as opposed to B |\n",
        "\n",
        "These operators allow you to construct arrays of multiple prompt variants and merge their output. Don't care if it's a snake or lizard? Write \"a new species of {snake|lizard}.\" Want the species to combine the qualities of a snake and a bird? Write \"{snake&bird}.\" Want to make sure the snake is not venomous? Write \"{~ venomous} snake,\" which is far more effective than simply writing \"non-venomous snake.\" You can combine multiple operators, using {} brackets to group text.\n",
        "\n",
        "A detailed explanation of this method, along with the code, is available [here](https://github.com/jeffbinder/promptarray)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KRSZU1o-qJjz",
        "cellView": "form"
      },
      "source": [
        "#@title Setup and Model Selection\n",
        "\n",
        "#@markdown Using a GPU is recommended, but you will first need to connect to an instance that has one. The larger models will not work unless your instance has enough RAM. Note that the XL model will take a while to load.\n",
        "\n",
        "model = \"gpt2-large\" #@param ['gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl']\n",
        "use_gpu = True #@param {type:\"boolean\"}\n",
        "\n",
        "%cd /content\n",
        "!rm -rf promptarray\n",
        "!git clone https://github.com/jeffbinder/promptarray\n",
        "%cd promptarray/\n",
        "\n",
        "!pip install lark\n",
        "!pip install sentencepiece\n",
        "!pip install git+https://github.com/huggingface/transformers.git@61f64262692ac7dc90e2e0bdeb7e79d9cd607a66\n",
        "\n",
        "import textwrap\n",
        "from generation_utils import *\n",
        "\n",
        "model_type = model.split('-')[0]\n",
        "model_name_or_path = model\n",
        "device = 'cuda' if use_gpu else 'cpu'\n",
        "\n",
        "# Initialize the model and tokenizer\n",
        "try:\n",
        "    model_class, tokenizer_class = MODEL_CLASSES[model_type]\n",
        "except KeyError:\n",
        "    raise KeyError(\"the model {} you specified is not supported. You are welcome to add it and open a PR :)\")\n",
        "tokenizer = tokenizer_class.from_pretrained(model_name_or_path)\n",
        "model = model_class.from_pretrained(model_name_or_path)\n",
        "model.to(device)\n",
        "model.eval()\n",
        "\n",
        "print(\"Ready!\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "x0dDxxwGB8DF"
      },
      "source": [
        "#@title Prompt Entry\n",
        "\n",
        "prompt_text = \"Scientists recently discovered a new species of {serpent~snake}. Here is a description of it:\" #@param {type:\"string\"}\n",
        "\n",
        "output_length = 100 #@param {type:\"slider\", min:1, max:500, step:1}\n",
        "num_return_sequences = 3 #@param {type:\"slider\", min:1, max:10, step:1}\n",
        "\n",
        "#@markdown ___\n",
        "\n",
        "do_sample = True #@param {type:\"boolean\"}\n",
        "temperature = 0.6 #@param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "top_k = 5 #@param {type:\"slider\", min:0, max:20, step:1}\n",
        "top_p = 0.8 #@param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "repetition_penalty = 1.5 #@param {type:\"slider\", min:0, max:5, step:0.1}\n",
        "overlap_factor = 0.25 #@param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "show_program = True #@param {type:\"boolean\"}\n",
        "\n",
        "\n",
        "def adjust_length_to_model(length, max_sequence_length):\n",
        "    if length < 0 and max_sequence_length > 0:\n",
        "        length = max_sequence_length\n",
        "    elif 0 < max_sequence_length < length:\n",
        "        length = max_sequence_length  # No generation bigger than model size\n",
        "    elif length < 0:\n",
        "        length = MAX_LENGTH  # avoid infinite loop\n",
        "    return length\n",
        "length = adjust_length_to_model(output_length, max_sequence_length=model.config.max_position_embeddings)\n",
        "\n",
        "import time\n",
        "start_time = time.time()\n",
        "output_sequences = model.generate(\n",
        "    prompt=prompt_text,\n",
        "    overlap_factor=overlap_factor,\n",
        "    tokenizer=tokenizer,\n",
        "    max_length=length,\n",
        "    temperature=temperature,\n",
        "    top_k=top_k,\n",
        "    top_p=top_p,\n",
        "    repetition_penalty=repetition_penalty,\n",
        "    do_sample=do_sample,\n",
        "    num_return_sequences=num_return_sequences,\n",
        "    pad_token_id=0,\n",
        "    verbose=show_program,\n",
        ")\n",
        "print(f\"Time: {time.time() - start_time}s\\n\")\n",
        "\n",
        "# Remove the batch dimension when returning multiple sequences\n",
        "if len(output_sequences.shape) > 2:\n",
        "    output_sequences.squeeze_()\n",
        "\n",
        "generated_sequences = []\n",
        "\n",
        "for generated_sequence_idx, generated_sequence in enumerate(output_sequences):\n",
        "    generated_sequence = generated_sequence.tolist()\n",
        "    generated_sequence = [idx for idx in generated_sequence if idx != 0]\n",
        "\n",
        "    # Decode text\n",
        "    generated_text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)\n",
        "\n",
        "    if num_return_sequences > 1:\n",
        "        print(f'\\nGenerated sequence {generated_sequence_idx}:')\n",
        "    print('\\n'.join(textwrap.wrap(generated_text)))\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}