GeDi_guided_GPT_2_XL.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/salesforce/GeDi/blob/master/GeDi_guided_GPT_2_XL.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Official implementation of generation with the topic GeDi (pronounced *Jedi*) model based on our paper [GeDi: Generative Discriminator Guided Sequence Generation](https://arxiv.org/abs/2009.06367)\n",
    "\n",
    "Check our github repository for more options (like detoxification and sentiment control) https://github.com/salesforce/GeDi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 204
    },
    "colab_type": "code",
    "id": "hWgP7huecyHw",
    "outputId": "59901224-5708-4b8e-ad83-7778a8e1643c"
   },
   "outputs": [],
   "source": [
    "!wget https://storage.googleapis.com/sfr-gedi-data/GeDi.zip\n",
    "import zipfile\n",
    "with zipfile.ZipFile('GeDi.zip', 'r') as zip_ref:\n",
    "    zip_ref.extractall('./')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "LeGaZoISeXy2",
    "outputId": "6bf07ff9-7209-4340-ab6e-50e2002da4f6"
   },
   "outputs": [],
   "source": [
    "cd GeDi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kXTswSEueE3b"
   },
   "outputs": [],
   "source": [
    "'''Installing transformers v2.8'''\n",
    "\n",
    "!pip install transformers==2.8\n",
    "!pip install -r hf_requirements.txt\n",
    "\n",
    "'''Downloading GeDi topic model checkpoints'''\n",
    "\n",
    "!wget https://storage.googleapis.com/sfr-gedi-data/gedi_topic.zip\n",
    "\n",
    "with zipfile.ZipFile('gedi_topic.zip', 'r') as zip_ref:\n",
    "    zip_ref.extractall('./')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "OgvczCtkdZM0"
   },
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "\n",
    "from modeling_gpt2 import GPT2LMHeadModel\n",
    "\n",
    "from transformers import (\n",
    "    GPT2Config,\n",
    "    GPT2Tokenizer\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CGwDhVA2dHoh"
   },
   "outputs": [],
   "source": [
    "mode = \"topic\"\n",
    "code_desired = \"true\"\n",
    "code_undesired = \"false\"\n",
    "model_type = 'gpt2'\n",
    "gen_type = \"gedi\"\n",
    "gen_model_name_or_path = \"gpt2-xl\"\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "MODEL_CLASSES = {\"gpt2\": (GPT2Config, GPT2LMHeadModel, GPT2Tokenizer),}\n",
    "config_class, model_class, tokenizer_class = MODEL_CLASSES[\"gpt2\"]\n",
    "tokenizer = tokenizer_class.from_pretrained(gen_model_name_or_path, do_lower_case=False)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next step needs to download and convert the GPT2-XL model. \n",
    "\n",
    "This takes a while (usually about 3 minutes to download and another 5 or so to convert). \n",
    "\n",
    "The good news is that once the model is loaded, you can quickly sample from many different prompts and topics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Loading GPT2-XL model (1.5B param LM) below, this could take a while.\n",
    "#This requires additional CPU memory overhead to load the pretrained weights in a new model\n",
    "#Due to CPU memory constraints on Colab, we're loading the model in half precision (load_in_half_prec=True) \n",
    "#Do to this change, generations may not always exactly match samples in paper, but sometimes do, and seem to be similar in quality\n",
    "#If you run the notebook with enough CPU RAM (most likely 16GB+), you can try setting load_in_half_prec=False   \n",
    "\n",
    "model = model_class.from_pretrained(gen_model_name_or_path, load_in_half_prec=True)\n",
    "model = model.to(device)\n",
    "model = model.float()\n",
    "\n",
    "gedi_model_name_or_path = 'gedi_topic'\n",
    "gedi_model = model_class.from_pretrained(gedi_model_name_or_path)\n",
    "gedi_model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "orrhZkuKwOXE"
   },
   "source": [
    "### Set arguments for generation\n",
    "\n",
    "You can change the max generation length, or play around with hyperparameter settings. \n",
    "\n",
    "The default hyperparameters were used in the topic model for the paper.\n",
    "\n",
    "More aggressive topic steering can be done by increasing `disc_weight` or `filter_p` (`filter_p` should always be less than 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "o1Waq8Gdudv-"
   },
   "outputs": [],
   "source": [
    "#setting arguments for generation\n",
    "#max generation length\n",
    "gen_length = 200\n",
    "#omega from paper, higher disc_weight means more aggressive topic steering\n",
    "disc_weight = 30\n",
    "#1 - rho from paper, should be between 0 and 1 higher filter_p means more aggressive topic steering\n",
    "filter_p = 0.8\n",
    "#tau from paper, preserves tokens that are classified as correct topic\n",
    "target_p = 0.8\n",
    "#hyperparameter that determines class prior, set to uniform by default\n",
    "class_bias = 0\n",
    "\n",
    "if gen_length>1024:\n",
    "  length = 1024\n",
    "else:\n",
    "  length = gen_length"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "cj9UD2wYjPz6"
   },
   "source": [
    "### Specify prompt and topic to GeDi\n",
    "\n",
    "\n",
    "The topic and prompt can be specified as strings with the `secondary_code` and `prompt` variables below.\n",
    "\n",
    "Note that our GeDi topic model has been trained on only four topics:  `world`, `sports`, `business` and `science` so it performs best on steering generation from GPT-2 towards these topics. However, it also shows some promising zero-shot results on new topics for eg. `education`, `food`, `fire`, `space`, `cars`, `climate`.\n",
    "\n",
    "Generic short prompts tend to work the best."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZjDKr8ZAgMN8"
   },
   "outputs": [],
   "source": [
    "#Specify what topic you want to generate on using the secondary_code variable\n",
    "\n",
    "secondary_code = 'climate'\n",
    "bpe_tokens = tokenizer.encode(secondary_code)\n",
    "if len(bpe_tokens) > 1:\n",
    "  print(\"Warning! number of bpe tokens for \" + code + \" is greater than 1, model isn't trained for this, generation is less likely to match the topic\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 173
    },
    "colab_type": "code",
    "id": "i5LyimIvdwEB",
    "outputId": "7b92490b-28c3-484b-e7c1-873db58d7e79"
   },
   "outputs": [],
   "source": [
    "#Specify prompt below\n",
    "prompt = \"In a shocking finding\"\n",
    "\n",
    "start_len=0\n",
    "text_ids = tokenizer.encode(prompt)\n",
    "encoded_prompts=torch.LongTensor(text_ids).unsqueeze(0).to(device)\n",
    "\n",
    "multi_code = tokenizer.encode(secondary_code)\n",
    "attr_class = 1\n",
    "\n",
    "generated_sequence = model.generate(input_ids=encoded_prompts,\n",
    "                                         pad_lens=None,\n",
    "                                          max_length= length,\n",
    "                                          top_k=None,\n",
    "                                          top_p=None,\n",
    "                                          repetition_penalty= 1.2,\n",
    "                                          rep_penalty_scale= 10,\n",
    "                                          eos_token_ids = tokenizer.eos_token_id,\n",
    "                                          pad_token_id = 0,\n",
    "                                          do_sample= False,\n",
    "                                          penalize_cond= True,\n",
    "                                          gedi_model= gedi_model,\n",
    "                                          tokenizer= tokenizer,\n",
    "                                          disc_weight= disc_weight,\n",
    "                                          filter_p = filter_p,\n",
    "                                          target_p = target_p,\n",
    "                                          class_bias = class_bias,\n",
    "                                          attr_class = attr_class,\n",
    "                                          code_0 = code_undesired,\n",
    "                                          code_1 = code_desired,\n",
    "                                          multi_code=multi_code\n",
    "                                          )\n",
    "\n",
    "text = tokenizer.decode(generated_sequence.tolist()[0], clean_up_tokenization_spaces=True)\n",
    "print('\\n')\n",
    "print(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZpqIppnIj1UH"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "GeDi-guided GPT-2 XL.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}