notebooks/llama-dev.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Generating metered verse with LLaMA\n",
    "\n",
    "This notebook demonstrates how to use Meta's 6 billion parameter LLaMA model to generate song lyrics with a specific metric structure. If you've been yearning to rewrite the happy birthday song so that it's just about dogs, `MetricGenerator` can help :). \n",
    "\n",
    "This notebook relies on the `MetricGenerator` class implemented my [bragi](https://github.com/joehoover/bragi) repository. If you're wondering, Bragi (pictured below) is the [Norse god](https://en.wikipedia.org/wiki/Bragi) of poetry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 277,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://upload.wikimedia.org/wikipedia/commons/thumb/2/20/Bragi_by_Wahlbom.jpg/440px-Bragi_by_Wahlbom.jpg\" width=\"150\" height=\"150\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 277,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import sys\n",
    "sys.path.append('../')\n",
    "\n",
    "from IPython.display import Image\n",
    "from IPython.core.display import HTML \n",
    "Image(url= \"http://my_site.com/my_picture.jpg\")\n",
    "Image(url= \"https://upload.wikimedia.org/wikipedia/commons/thumb/2/20/Bragi_by_Wahlbom.jpg/440px-Bragi_by_Wahlbom.jpg\", width=150, height=150)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "The `MetricGenerator` class allows you to specify an *initialization* song that is used to specify a target metric structure. Then, given this target metric structure, it attempts to generate new lyrics with the same metric structure. \n",
    "\n",
    "For example, the following lines\n",
    "\n",
    "<blockquote>\n",
    "Happy birthday to you<br>\n",
    "Happy birthday to you<br>\n",
    "Happy birthday dear Marvin<br>\n",
    "Happy birthday to you\n",
    "</blockquote>\n",
    "\n",
    "have 6, 6, 7, and 6 syllables respectively. With `MetricGenerator`, you can use those lyrics to initialize a metric structure and then generate new lyrics with the same metric structure.\n",
    "\n",
    "And, most importantly, you can guide the new lyrics with a prompt!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Quick demo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check it out! First, we'll load our dependencies, specify some paths, and initialize a model and tokenizer.\n",
    "\n",
    "For this demo, I'll use Meta's 6 billion parameter `LLaMA` model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bragi.metric_generator import MetricGenerator\n",
    "from transformers import LLaMAForCausalLM, LLaMATokenizer\n",
    "import torch \n",
    "\n",
    "CACHE_DIR = 'weights'\n",
    "SEP = \"<sep>\"\n",
    "MODEL_PATH  = \"/src/weights\"\n",
    "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "\n",
    "model = LLaMAForCausalLM.from_pretrained(MODEL_PATH, cache_dir=CACHE_DIR, local_files_only=True).to(device)\n",
    "tokenizer = LLaMATokenizer.from_pretrained(MODEL_PATH, cache_dir=CACHE_DIR, local_files_only=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now I'll initialize the metric generator class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 324,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "generator = MetricGenerator(model=model, tokenizer=tokenizer, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To generate lyrics, you just need to provide a prompt and the lyrics you want to use to initialize your metric structure. Let's try the happy birthday song displayed above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 245,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Happy birthday to you,\n",
      "Happy birthday to you,\n",
      "Happy birthday dear Marvin,\n",
      "Happy birthday to you\n"
     ]
    }
   ],
   "source": [
    "text_init = \"Happy birthday to you,\\nHappy birthday to you,\\nHappy birthday dear Marvin,\\nHappy birthday to you\"\n",
    "print(text_init)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And let's try to generate a new song that has the same metric structure, but is about dogs. To do that, we can use something like the following prompt:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 320,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "This is a song about dogs:\n",
      "\n"
     ]
    }
   ],
   "source": [
    "prompt = \"\"\"This is a song about dogs:\\n\"\"\"\n",
    "print(prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can call our `MetricGenerator` instance. The `__call__` method of `MetricGenerator` accepts all of the arguments that `Transformers.model.generate()` does, but it adds a few extra methods that allow us to constrain the generation process so that the target metric structure is respected.\n",
    "\n",
    "**Note:** If you want the raw token_ids, you can call `MetricGenerator.generate()` with the same arguments. In contrast, `MetricGenerator.__call__()` returns the decoded output string."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 328,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---prompt---\n",
      "This is a song about dogs:\n",
      "\n",
      "\n",
      "---text_init----\n",
      "Happy birthday to you,\n",
      "Happy birthday to you,\n",
      "Happy birthday dear Marvin,\n",
      "Happy birthday to you\n",
      "\n",
      "\n",
      "----output-----\n",
      "I love how the dogs can \n",
      "be so good, so bad, and \n",
      "so strange. Sometimes I love them \n",
      "when they are the most good\n",
      "\n",
      "\n",
      "----Syllables-----\n",
      "Syllables per line in output: tensor([6., 6., 7., 6.])\n",
      "Syllables per line in `text_init`: tensor([6., 6., 7., 6.])\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(2)\n",
    "output = generator(\n",
    "    prompt = prompt,\n",
    "    text_init = text_init,\n",
    "    free_tokens=['||', '?', '.', ','],\n",
    "    # syllable_budget = torch.Tensor([6., 6.]),\n",
    "    num_return_sequences=1,\n",
    "    no_repeat_ngram_size=2,\n",
    "    remove_invalid_values=True,\n",
    "    do_sample=True,\n",
    "    top_k=25,\n",
    "    temperature=.7,\n",
    "    max_length = 100,\n",
    "    new_line_token='||',\n",
    "    bad_words_ids=[[8876]],\n",
    ")\n",
    "\n",
    "print('---prompt---')\n",
    "print(prompt.strip())\n",
    "print('\\n')\n",
    "\n",
    "print('---text_init----')\n",
    "print(text_init)\n",
    "print('\\n')\n",
    "\n",
    "print('----output-----')\n",
    "print(output)\n",
    "print('\\n')\n",
    "\n",
    "print('----Syllables-----')\n",
    "print(f\"Syllables per line in output: {generator.calculate_syllable_budget(output)}\")\n",
    "print(f\"Syllables per line in `text_init`: {generator.calculate_syllable_budget(text_init)}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Neat! We rewrote happy birthday to be aboud dogs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exploring the effect of prompt conditioning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try some other prompts and see how responsive the model is to prompt conditioning. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **California ☀️!**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 274,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "text_init = \"Happy birthday to you,\\nHappy birthday to you,\\nHappy birthday dear Marvin,\\nHappy birthday to you\"\n",
    "prompt = \"\"\"This is a song about about California sun:\\n\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 275,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---prompt---\n",
      "This is a song about about California sun:\n",
      "\n",
      "\n",
      "---text_init----\n",
      "Happy birthday to you,\n",
      "Happy birthday to you,\n",
      "Happy birthday dear Marvin,\n",
      "Happy birthday to you\n",
      "\n",
      "\n",
      "----output-----\n",
      "I love how the colors \n",
      "the lights are bright, the air \n",
      "is warm and the water is \n",
      "clear. I love the way the\n",
      "\n",
      "\n",
      "----Syllables-----\n",
      "Syllables per line in output: tensor([6., 6., 7., 6.])\n",
      "Syllables per line in `text_init`: tensor([6., 6., 7., 6.])\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(2)\n",
    "output = generator(\n",
    "    prompt = prompt,\n",
    "    text_init = text_init,\n",
    "    # syllable_budget = torch.Tensor([6., 6.]),\n",
    "    num_return_sequences=1,\n",
    "    no_repeat_ngram_size=2,\n",
    "    remove_invalid_values=True,\n",
    "    do_sample=True,\n",
    "    top_k=25,\n",
    "    temperature=.7,\n",
    "    max_length = 100,\n",
    "    new_line_token='||',\n",
    "    free_tokens=['||', '?', '.', ','], \n",
    "    bad_words_ids=[[8876]],\n",
    ")\n",
    "\n",
    "print('---prompt---')\n",
    "print(prompt.strip())\n",
    "print('\\n')\n",
    "\n",
    "print('---text_init----')\n",
    "print(text_init)\n",
    "print('\\n')\n",
    "\n",
    "print('----output-----')\n",
    "print(output)\n",
    "print('\\n')\n",
    "\n",
    "print('----Syllables-----')\n",
    "print(f\"Syllables per line in output: {generator.calculate_syllable_budget(output)}\")\n",
    "print(f\"Syllables per line in `text_init`: {generator.calculate_syllable_budget(text_init)}\")\n",
    "# print(tokenizer.decode(output[0], skip_special_tokens=True).strip())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Dark and stormy 🌊**!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 266,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "text_init = \"Happy birthday to you,\\nHappy birthday to you,\\nHappy birthday dear Marvin,\\nHappy birthday to you\"\n",
    "prompt = \"\"\"This a beautiful poem that uses descriptive, earthy language to describe the ocean during a storm:\\n\\\"\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 268,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---prompt---\n",
      "This a beautiful poem that uses descriptive, earthy language to describe the ocean during a storm:\n",
      "\"\n",
      "\n",
      "\n",
      "---text_init----\n",
      "Happy birthday to you,\n",
      "Happy birthday to you,\n",
      "Happy birthday dear Marvin,\n",
      "Happy birthday to you\n",
      "\n",
      "\n",
      "----output-----\n",
      "In the storm, the wind and \n",
      "the sea are one. The waves \n",
      "are the breath of the gods. As \n",
      "they break and fall back, they\n",
      "\n",
      "\n",
      "----Syllables-----\n",
      "Syllables per line in output: [6, 6, 7, 6]\n",
      "Syllables per line in `text_init`: [6, 6, 7, 6]\n"
     ]
    }
   ],
   "source": [
    "\n",
    "torch.manual_seed(5)\n",
    "output = generator(\n",
    "    prompt = prompt,\n",
    "    text_init = text_init,\n",
    "    # syllable_budget = torch.Tensor([6., 6.]),\n",
    "    num_return_sequences=1,\n",
    "    no_repeat_ngram_size=2,\n",
    "    remove_invalid_values=True,\n",
    "    do_sample=True,\n",
    "    top_k=25,\n",
    "    temperature=.7,\n",
    "    max_length = 100,\n",
    "    new_line_token='||',\n",
    "    free_tokens=['||', '?', '.', ','], \n",
    "    bad_words_ids=[[8876]],\n",
    ")\n",
    "\n",
    "print('---prompt---')\n",
    "print(prompt.strip())\n",
    "print('\\n')\n",
    "\n",
    "print('---text_init----')\n",
    "print(text_init)\n",
    "print('\\n')\n",
    "\n",
    "print('----output-----')\n",
    "print(output)\n",
    "print('\\n')\n",
    "\n",
    "print('----Syllables-----')\n",
    "print(f\"Syllables per line in output: {generator.calculate_syllable_budget(output, pt=False)}\")\n",
    "print(f\"Syllables per line in `text_init`: {generator.calculate_syllable_budget(text_init, pt=False)}\")\n",
    "# print(tokenizer.decode(output[0], skip_special_tokens=True).strip())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Methodology"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that you've seen `MetricGenerator` in action, let's take a peak under the hood. We'll discuss:\n",
    "\n",
    "1. How `MetricGenerator` constrains generation\n",
    "2. The fact that it occasionally violates the target metric structure\n",
    "3. Shortcomings of the current implementation\n",
    "4. Ways that we could improve `MetricGenerator` and shortcomings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### How does `MetricGenerator` constrain generation?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`MetricGenerator` uses a logit warper to constrain model vocabulary during the generation process. A logit warper is a callable that modifies a model's probability distribution over tokens during generation. Specifically, `MetricGenerator` uses the `SyllableRestrictionWarper`, which is implemented in `bragi.logit_warpers`. \n",
    "\n",
    "The `SyllableRestrictionWarper` constrains the metric structure of model generated output by tracking a line-level metric *budget* and a number-of-lines budget. Every time the generator *spends* one more more syllables, the budget is decremented. When all of the syllables permitted for a given line are spent, the warper forces a new line token by setting the logit scores for all other tokens to `-Inf`. Further, when a new line token is emitted, the number-of-lines budget is decremented by one. Finally, when the number-of-lines budget reaches zero, the warper forces the eos token using the same approach. \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Calculating syllable cost"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `SyllableRestrictionWarper` requires a method for calculating the *cost* of each token in the model's vocabulary. In theory, this cost function can implement arbitrary scoring logic; however, for our purposes model tokens are *scored* according to their number of syllables. \n",
    "\n",
    "This is accomplished using the `phones_for_word` and `syllable_count` methods implemented in the `pronouncing` library, a python interface for the CMU Pronouncing Dictionary. `Bragi` provides the `verse_parsers.cmu_syllable_counter` function, which wraps the `pronouncing` methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 284,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'Hello' has 2 syllables\n",
      "'world' has 1 syllables\n"
     ]
    }
   ],
   "source": [
    "from bragi.verse_parsers import cmu_syllable_counter\n",
    "print(f\"'Hello' has {cmu_syllable_counter('Hello')} syllables\")\n",
    "print(f\"'world' has {cmu_syllable_counter('world')} syllables\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`MetricGenerator` calculates the syllables of every token in the model's vocabulary.\n",
    "\n",
    "However, because LLaMA uses a [BPE tokenizer],(https://huggingface.co/docs/transformers/main/model_doc/llama#transformers.LlamaTokenizer), so it's crucial to calculate token syllables on decoded tokens and not the raw model vocabulary. For example, the characters 'and' are tokenized into '▁and', which is OOV for the CMU dictionary. `MetricGenerator` uses the `verse_parsers.token_syllable_scores` which handles decoding and syllable calculation for the all tokens in the model vocabulary. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Calculating the metric structure of the initialization text"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To calculate the metric structure of the initialization text, `MetricGenerator` uses a wrapper around the `poesy` package. `poesy` is a python interface for the `eSpeak` speech synthesis library, which provides more robust syllable calculation than simply performing lexical lookups against the CMU Pronouncing Dictionary. \n",
    "\n",
    "`bragi` provides the `PoesyParsedVerseHandler` class as an interface for the `poesy`. You can use the `example` method to extract the syllabic structure of a string in which new lines are specified as '\\n'.\n",
    "\n",
    "For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 295,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "An elderly man called Keith,\n",
      "Mislaid his set of false teeth.\n",
      "They'd been laid on a chair,\n",
      "He'd forgot they were there,\n",
      "Sat down, and was bitten beneath.\n",
      "\n",
      "The lines of the limerick above have the following syllable counts: [7, 7, 6, 6, 8]\n"
     ]
    }
   ],
   "source": [
    "from bragi.verse_parsers import PoesyParsedVerseHandler\n",
    "verse_handler = PoesyParsedVerseHandler()\n",
    "\n",
    "text = \"\"\"An elderly man called Keith,\n",
    "Mislaid his set of false teeth.\n",
    "They'd been laid on a chair,\n",
    "He'd forgot they were there,\n",
    "Sat down, and was bitten beneath.\"\"\"\n",
    "\n",
    "print(text)\n",
    "\n",
    "example, syllable_budget = verse_handler.example(text)\n",
    "print(f\"\\nThe lines of the limerick above have the following syllable counts: {syllable_budget}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** The `PoesyParsedVerseHandler.example()` method also transforms an input text to a control code that can be used for fine-tuning. In addition to capturing the syllabic structure of the input text, it also attempts to detect the rhyme scheme. This functionality was inpsired by [Ormazabal et al. (2022)](https://aclanthology.org/2022.findings-emnlp.268/), which demonstrates that fine-tuning a language model on a control code such as the one below substantially improves its ability to produce verse with a specific metric structure.\n",
    "\n",
    "E.g.:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 298,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<PREF>\n",
      "<SYLLABLES: 7><RHYME: A>\n",
      "<SYLLABLES: 7><RHYME: A>\n",
      "<SYLLABLES: 6><RHYME: B>\n",
      "<SYLLABLES: 6><RHYME: B>\n",
      "<SYLLABLES: 8><RHYME: A>\n",
      "</PREF>\n",
      "\n",
      "An elderly man called Keith,\n",
      "Mislaid his set of false teeth.\n",
      "They'd been laid on a chair,\n",
      "He'd forgot they were there,\n",
      "Sat down, and was bitten beneath.\n"
     ]
    }
   ],
   "source": [
    "print(example)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Putting it all together"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`SyllableRestrictionWarper` uses the line-level syllable counts obtained from `poesy` to construct a syllable budget. Then at each inference step, words are dynamically masked according to their syllable cost, conditional on the available budget for the current line. \n",
    "\n",
    "For example, imagine that the tokens corresponding to the 6-syllable character sequence \"An elderly man called\" have already been generated for a line with a 7 syllable budget. In this case, tokens like \"Marvin\", \"Magnolia\", and \"murky\" will be masked, as will any other token with two or more syllables. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Limitations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Metric Structure Errors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`MetricGenerator` does not always strictly adhere to the initialized metric structure. The indexing over syllable and line budgets is still in alpha and there's probably a bug causing this behavior. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 317,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---prompt---\n",
      "Song lyrics for \"Heavy Metal Sunset\", an 80's metal song about Los Angeles:\n",
      "\"\n",
      "\n",
      "\n",
      "---text_init----\n",
      "Happy birthday to you,\n",
      "Happy birthday to you,\n",
      "Happy birthday dear Marvin,\n",
      "Happy birthday to you\n",
      "\n",
      "\n",
      "----output-----\n",
      "Heartache and sorrow \n",
      "A thousand miles apart \n",
      "I can see it clear, I can \n",
      "I feel it in my heart\n",
      "\n",
      "\n",
      "----Syllables-----\n",
      "Syllables per line in output: [5, 6, 7, 6]\n",
      "Syllables per line in `text_init`: [6, 6, 7, 6]\n"
     ]
    }
   ],
   "source": [
    "text_init = \"Happy birthday to you,\\nHappy birthday to you,\\nHappy birthday dear Marvin,\\nHappy birthday to you\"\n",
    "prompt = \"\"\"Song lyrics for \"Heavy Metal Sunset\", an 80's metal song about Los Angeles:\\n\\\"\"\"\"\n",
    "\n",
    "torch.manual_seed(1)\n",
    "output = generator(\n",
    "    prompt = prompt,\n",
    "    text_init = text_init,\n",
    "    # syllable_budget = torch.Tensor([6., 6.]),\n",
    "    num_return_sequences=1,\n",
    "    no_repeat_ngram_size=3,\n",
    "    remove_invalid_values=True,\n",
    "    do_sample=True,\n",
    "    top_k=25,\n",
    "    temperature=.8,\n",
    "    max_length = 100,\n",
    "    new_line_token='||',\n",
    "    free_tokens=['||', '?', '.', ','], \n",
    "    bad_words_ids=[[8876]],\n",
    ")\n",
    "\n",
    "print('---prompt---')\n",
    "print(prompt.strip())\n",
    "print('\\n')\n",
    "\n",
    "print('---text_init----')\n",
    "print(text_init)\n",
    "print('\\n')\n",
    "\n",
    "print('----output-----')\n",
    "print(output)\n",
    "print('\\n')\n",
    "\n",
    "print('----Syllables-----')\n",
    "print(f\"Syllables per line in output: {generator.calculate_syllable_budget(output, pt=False)}\")\n",
    "print(f\"Syllables per line in `text_init`: {generator.calculate_syllable_budget(text_init, pt=False)}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Investigating deviations from specified meter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to get a sense of `MetricGenerator`'s failure rate, we can look at the distribution of deviations from expected number of syllables per line. Below, I run 100 generations against a target text and calculate the mean and standard deviation of these deviations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 215,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: tqdm in /root/.pyenv/versions/3.10.10/lib/python3.10/site-packages (4.65.0)\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip available: \u001b[0m\u001b[31;49m22.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.0.1\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 241,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [05:01<00:00,  3.02s/it]\n"
     ]
    }
   ],
   "source": [
    "import tqdm\n",
    "\n",
    "prompt = \"\"\"This a beautiful poem that uses descriptive, earthy language to describe the ocean during a storm:\\n\\\"\"\"\"\n",
    "text_init = \"Happy birthday to you,\\nHappy birthday to you,\\nHappy birthday dear Marvin,\\nHappy birthday to you\"\n",
    "init_syllables = generator.calculate_syllable_budget(text_init)\n",
    "\n",
    "output_syllables = []\n",
    "for i in tqdm.tqdm(range(0, 100)):\n",
    "    output = generator.generate(\n",
    "        prompt = prompt,\n",
    "        text_init = text_init,\n",
    "        # syllable_budget = torch.Tensor([6., 6.]),\n",
    "        num_return_sequences=1,\n",
    "        no_repeat_ngram_size=2,\n",
    "        remove_invalid_values=True,\n",
    "        do_sample=True,\n",
    "        top_k=25,\n",
    "        temperature=.7,\n",
    "        max_length = 100,\n",
    "        new_line_token='||',\n",
    "        free_tokens=['||', '?', '.', ','], \n",
    "        bad_words_ids=[[8876]],\n",
    "    )\n",
    "    \n",
    "    output_syllables.append(generator.calculate_syllable_budget(output))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 242,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "stacked_output_syllables = torch.stack(output_syllables)\n",
    "deviation = stacked_output_syllables - init_syllables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 243,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The average line-level violation of the metric constraint is -0.5174999833106995,       which suggests that the mechanism is unbiased but still flawed.\n",
      "The vast majority of deviations are within +/- 1.824031949043274 syllables of the constraint.\n"
     ]
    }
   ],
   "source": [
    "m = deviation.mean()\n",
    "std = deviation.std()\n",
    "\n",
    "print(f\"The average line-level violation of the metric constraint is {m}.\")\n",
    "print(f\"The vast majority of deviations are within +/- {std*2} syllables of the constraint.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Beam search is not supported"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An obvious way to improve output quality is to conduct generation with beam search. However, the current implementation of the syllable budgeting process does not support beam search."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Future Work"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "**Excluding partial words from the model vocabulary**\n",
    "\n",
    "Currently, beyond the syllable budget, `MetricGenerator` does not enforce any additional constraints on the generation process or available model vacubulary at a given step. This means that it *can* output partial words or non-sense tokens. Accordingly, excluding partial-word tokens may improve output quality. However, standard best-practices for text generation such as specifying `topk` appear to largely mitigate the partial-word issue. With a reasonably small `topk` value (e.g. 25-50), partial word emissions are a negligible problem. Accordingly, this has been left for future work. \n",
    "\n",
    "**Fine-tuning on in-distribution data**\n",
    "\n",
    "While emitting partial words is largely mitigated by masking all but the highest probability tokens, no current generation settings address the problem of abrupt stopping. Because `MetricGenerator` forces an eos token when the syllable budget of the last line is spent, most generated sequences have an unnatural ending. While implementing support for beam search may partially mitigate this, it will likely be far more effective to fine-tune on in-domain data such as song lyrics and poems. The fine-tuning process could be further enhanced using control codes ([Ormazabal et al. (2022)](https://aclanthology.org/2022.findings-emnlp.268/)) that encode the metric structure of the target text. At inference time, generation would condition on the control code and this would–in theory–guide the model toward natural endings. \n",
    "\n",
    "Fine-tuning on in-distribution data would also provide an opportunity to improve the conditioning effect of prompts on generated sequences. For example, the instruction-tuning methods implemented in [Chakrabarty (2022)](https://arxiv.org/abs/2210.13669) allow users to request rhetorical devices such as metaphor. \n",
    "\n",
    "**Constrained verse generation with a fine-tuned model**\n",
    "\n",
    "Ultimately, the best solution will likely rely on a combination of rule-based constraints (e.g.  excluding tokens based on their syllabic cost) and fine-tuning. This approach should yield a model that is forced to follow the rules but also tought how to follow them with style."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.10"
  },
  "toc-showcode": false
 },
 "nbformat": 4,
 "nbformat_minor": 4
}