samples/colab_llama2_enforcer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the COLAB runtime (user action required)\n",
    "\n",
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/noamgat/lm-format-enforcer/blob/main/samples/colab_llama2_enforcer.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>\n",
    "\n",
    "This colab-friendly notebook is targeted at demoing the enforcer on LLAMA2. It can run on a free GPU on Google Colab.\n",
    "Make sure that your runtime is set to GPU:\n",
    "\n",
    "Menu Bar -> Runtime -> Change runtime type -> T4 GPU (at the time of writing this notebook). [Guide here](https://www.codesansar.com/deep-learning/using-free-gpu-tpu-google-colab.htm)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install transformers torch lm-format-enforcer huggingface_hub optimum langchain langchain-experimental\n",
    "!pip install auto-gptq --extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/ \n",
    "\n",
    "# When running from source / developing the library, use this instead\n",
    "# %load_ext autoreload\n",
    "# %autoreload 2\n",
    "# import sys\n",
    "# import os\n",
    "# sys.path.append(os.path.abspath('..'))\n",
    "## os.environ['CUDA_LAUNCH_BLOCKING'] = '1'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can create the model. This may take a few minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/noamgat/mambaforge/envs/commentranker/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "model_id = 'TheBloke/Llama-2-7b-Chat-GPTQ'\n",
    "device = 'cuda'\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    model = AutoModelForCausalLM.from_pretrained(model_id, device_map='auto')\n",
    "else:\n",
    "    raise Exception('GPU not available')\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "if tokenizer.pad_token_id is None:\n",
    "    # Required for batching example\n",
    "    tokenizer.pad_token_id = tokenizer.eos_token_id \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the previous cell executed successfully, you have propertly set up your Colab runtime and huggingface account!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the prompt for the specific language model\n",
    "\n",
    "We set up the prompting style according to the demo at https://huggingface.co/spaces/huggingface-projects/llama-2-13b-chat/blob/main/app.py . We simplify the implementation a bit as we don't need chat history for this demo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_prompt(message: str, system_prompt: str) -> str:\n",
    "    texts = [f'<s>[INST] <<SYS>>\\n{system_prompt}\\n<</SYS>>\\n\\n']\n",
    "    # The first user input is _not_ stripped\n",
    "    do_strip = False\n",
    "    message = message.strip() if do_strip else message\n",
    "    texts.append(f'{message} [/INST]')\n",
    "    return ''.join(texts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calling generate_enforced() instead of model.generate()\n",
    "\n",
    "The main function is fairly straightforward, except for the optional parameters ```required_regex``` / ```required_str``` / ```required_json_schema``` which activate the appropriate ```CharacterLevelParser``` to be used by the format enforcer.\n",
    "\n",
    "The implementation includes ```output_scores=True``` when calling ```generate_enforced()```, which returns diagnostic information about the enforcer's actions. We will use this to understand the results later in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "from typing import Tuple, Optional, Union, List\n",
    "import pandas as pd\n",
    "from lmformatenforcer import JsonSchemaParser, CharacterLevelParser, RegexParser, StringParser\n",
    "from lmformatenforcer.integrations.transformers import generate_enforced, build_token_enforcer_tokenizer_data\n",
    "\n",
    "StringOrManyStrings = Union[str, List[str]]\n",
    "\n",
    "tokenizer_data = build_token_enforcer_tokenizer_data(tokenizer)\n",
    "\n",
    "def run(message: StringOrManyStrings,\n",
    "        system_prompt: str,\n",
    "        max_new_tokens: int = 1024,\n",
    "        temperature: float = 0.8,\n",
    "        top_p: float = 0.95,\n",
    "        top_k: int = 50,\n",
    "        num_beams: int = 1,\n",
    "        required_regex: Optional[str] = None,\n",
    "        required_str: Optional[str] = None,\n",
    "        required_json_schema: Optional[dict] = None,\n",
    "        required_json_output: Optional[bool] = None) -> Tuple[StringOrManyStrings, Optional[pd.DataFrame]]:\n",
    "    is_multi_message = isinstance(message, list)\n",
    "    messages = message if is_multi_message else [message]\n",
    "    prompts = [get_prompt(msg, system_prompt) for msg in messages]\n",
    "    inputs = tokenizer(prompts, return_tensors='pt', add_special_tokens=False, return_token_type_ids=False, padding=is_multi_message).to(device)\n",
    "    \n",
    "    generate_kwargs = dict(\n",
    "        inputs,\n",
    "        # streamer=streamer,\n",
    "        max_new_tokens=max_new_tokens,\n",
    "        do_sample=True,\n",
    "        top_p=top_p,\n",
    "        top_k=top_k,\n",
    "        temperature=temperature,\n",
    "        num_beams=num_beams,\n",
    "        output_scores=True,\n",
    "        return_dict_in_generate=True\n",
    "    )\n",
    "\n",
    "    parser: Optional[CharacterLevelParser] = None\n",
    "    if required_regex:\n",
    "        parser = RegexParser(required_regex)\n",
    "    if required_str:\n",
    "        parser = StringParser(required_str)\n",
    "    if required_json_schema:\n",
    "        parser = JsonSchemaParser(required_json_schema)\n",
    "    if required_json_output:\n",
    "        parser = JsonSchemaParser(None)\n",
    "\n",
    "    if parser:\n",
    "        output = generate_enforced(model, tokenizer_data, parser, **generate_kwargs)\n",
    "    else:\n",
    "        output = model.generate(**generate_kwargs)\n",
    "\n",
    "    sequences = output['sequences']\n",
    "    # skip_prompt=True doesn't work consistenly, so we hack around it.\n",
    "    string_outputs = [tokenizer.decode(sequence, skip_special_tokens=True) for sequence in sequences]\n",
    "    string_outputs = [string_output.replace(prompt[3:], '') for string_output, prompt in zip(string_outputs, prompts)]\n",
    "    if parser and not is_multi_message:\n",
    "        enforced_scores_dict = output.enforced_scores\n",
    "        enforced_scores = pd.DataFrame(enforced_scores_dict)\n",
    "        pd.set_option('display.width', 1000)\n",
    "        pd.set_option('display.max_columns', 10)\n",
    "        pd.set_option('display.max_rows', 999)\n",
    "        pd.set_option('display.float_format', ' {:,.5f}'.format)\n",
    "    else:\n",
    "        enforced_scores = None\n",
    "    return string_outputs if is_multi_message else string_outputs[0], enforced_scores\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## JSON Schema Use case\n",
    "\n",
    "Now we demonstrate using ```JsonSchemaParser```. We create a pydantic model, generate the schema from it, and use that to enforce the format.\n",
    "The output will always be in a format that can be parsed by the parser."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "**Question:**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       "Please give me information about Michael Jordan. You MUST answer using the following json schema: {\"title\": \"AnswerFormat\", \"type\": \"object\", \"properties\": {\"first_name\": {\"title\": \"First Name\", \"type\": \"string\"}, \"last_name\": {\"title\": \"Last Name\", \"type\": \"string\"}, \"year_of_birth\": {\"title\": \"Year Of Birth\", \"type\": \"integer\"}, \"num_seasons_in_nba\": {\"title\": \"Num Seasons In Nba\", \"type\": \"integer\"}}, \"required\": [\"first_name\", \"last_name\", \"year_of_birth\", \"num_seasons_in_nba\"]}\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Answer, With json schema enforcing:**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/noamgat/mambaforge/envs/commentranker/lib/python3.10/site-packages/transformers/generation/utils.py:1411: UserWarning: You have modified the pretrained model configuration to control generation. This is a deprecated strategy to control generation and will be removed soon, in a future version. Please use a generation configuration file (see https://huggingface.co/docs/transformers/main_classes/text_generation )\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       "  {\n",
       "\"first_name\": \"Michael\",\n",
       "\"last_name\": \"Jordan\",\n",
       "\"year_of_birth\": 1963,\n",
       "\"num_seasons_in_nba\": 15\n",
       "}\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Answer, Without json schema enforcing:**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       "  Of course! I'd be happy to help you with your question about Michael Jordan. Here's the information you requested in the format you specified:\n",
       "{\n",
       "\"title\": \"AnswerFormat\",\n",
       "\"type\": \"object\",\n",
       "\"properties\": {\n",
       "\"first_name\": {\n",
       "\"title\": \"First Name\",\n",
       "\"type\": \"string\"\n",
       "},\n",
       "\"last_name\": {\n",
       "\"title\": \"Last Name\",\n",
       "\"type\": \"string\"\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Answer, With json mode enforcing (json output, schemaless):**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       "  {\n",
       "\"title\": \"AnswerFormat\",\n",
       "\"type\": \"object\",\n",
       "\"properties\": {\n",
       "\"first_name\": {\n",
       "\"title\": \"First Name\",\n",
       "\"type\": \"string\"\n",
       "},\n",
       "\"last_name\": {\n",
       "\"title\": \"Last Name\",\n",
       "\"type\": \"string\"\n",
       "},\n",
       "\"year_of_birth\": {\n",
       "\"title\": \"Year Of Birth\",\n",
       "\"type\": \"integer\"\n",
       "},\n",
       "\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from pydantic import BaseModel\n",
    "from IPython.display import display, Markdown\n",
    "from typing import List\n",
    "\n",
    "def display_header(text):\n",
    "    display(Markdown(f'**{text}**'))\n",
    "\n",
    "def display_content(text):\n",
    "    display(Markdown(f'```\\n{text}\\n```'))\n",
    "\n",
    "DEFAULT_SYSTEM_PROMPT = \"\"\"\\\n",
    "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\\n\\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\\\n",
    "\"\"\"\n",
    "DEFAULT_MAX_NEW_TOKENS = 100\n",
    "\n",
    "class AnswerFormat(BaseModel):\n",
    "    first_name: str\n",
    "    last_name: str\n",
    "    year_of_birth: int\n",
    "    num_seasons_in_nba: int\n",
    "\n",
    "question = 'Please give me information about Michael Jordan. You MUST answer using the following json schema: '\n",
    "question_with_schema = f'{question}{AnswerFormat.schema_json()}'\n",
    "\n",
    "display_header(\"Question:\")\n",
    "display_content(question_with_schema)\n",
    "\n",
    "display_header(\"Answer, With json schema enforcing:\")\n",
    "result, enforced_scores = run(question_with_schema, system_prompt=DEFAULT_SYSTEM_PROMPT, max_new_tokens=DEFAULT_MAX_NEW_TOKENS, required_json_schema=AnswerFormat.schema())\n",
    "display_content(result)\n",
    "\n",
    "display_header(\"Answer, Without json schema enforcing:\")\n",
    "result, _ = run(question_with_schema, system_prompt=DEFAULT_SYSTEM_PROMPT, max_new_tokens=DEFAULT_MAX_NEW_TOKENS)\n",
    "display_content(result)\n",
    "\n",
    "display_header(\"Answer, With json mode enforcing (json output, schemaless):\")\n",
    "result, _ = run(question_with_schema, system_prompt=DEFAULT_SYSTEM_PROMPT, max_new_tokens=DEFAULT_MAX_NEW_TOKENS, required_json_output=True)\n",
    "display_content(result)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Understanding the results\n",
    "Both runs used the exact same prompt, but only the enforced run produced a valid JSON output.\n",
    "How did the enforcer cause the output to conform to the format? Lets look at the enforcer intervention table:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "**Enforcer intervention table**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>generated_token</th>\n",
       "      <th>generated_token_idx</th>\n",
       "      <th>generated_score</th>\n",
       "      <th>leading_token</th>\n",
       "      <th>leading_token_idx</th>\n",
       "      <th>leading_score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td></td>\n",
       "      <td>29871</td>\n",
       "      <td>0.99982</td>\n",
       "      <td></td>\n",
       "      <td>29871</td>\n",
       "      <td>0.99982</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>{</td>\n",
       "      <td>426</td>\n",
       "      <td>0.00002</td>\n",
       "      <td>Of</td>\n",
       "      <td>4587</td>\n",
       "      <td>0.95649</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>\\n</td>\n",
       "      <td>13</td>\n",
       "      <td>0.99705</td>\n",
       "      <td>\\n</td>\n",
       "      <td>13</td>\n",
       "      <td>0.99705</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>\"</td>\n",
       "      <td>29908</td>\n",
       "      <td>0.98198</td>\n",
       "      <td>\"</td>\n",
       "      <td>29908</td>\n",
       "      <td>0.98198</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>first</td>\n",
       "      <td>4102</td>\n",
       "      <td>0.01030</td>\n",
       "      <td>title</td>\n",
       "      <td>3257</td>\n",
       "      <td>0.55388</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>_</td>\n",
       "      <td>29918</td>\n",
       "      <td>0.99879</td>\n",
       "      <td>_</td>\n",
       "      <td>29918</td>\n",
       "      <td>0.99879</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>name</td>\n",
       "      <td>978</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>name</td>\n",
       "      <td>978</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>\":</td>\n",
       "      <td>1115</td>\n",
       "      <td>0.99994</td>\n",
       "      <td>\":</td>\n",
       "      <td>1115</td>\n",
       "      <td>0.99994</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>\"</td>\n",
       "      <td>376</td>\n",
       "      <td>0.99949</td>\n",
       "      <td>\"</td>\n",
       "      <td>376</td>\n",
       "      <td>0.99949</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>Michael</td>\n",
       "      <td>24083</td>\n",
       "      <td>0.99028</td>\n",
       "      <td>Michael</td>\n",
       "      <td>24083</td>\n",
       "      <td>0.99028</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>\",</td>\n",
       "      <td>613</td>\n",
       "      <td>0.99112</td>\n",
       "      <td>\",</td>\n",
       "      <td>613</td>\n",
       "      <td>0.99112</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>\\n</td>\n",
       "      <td>13</td>\n",
       "      <td>0.99918</td>\n",
       "      <td>\\n</td>\n",
       "      <td>13</td>\n",
       "      <td>0.99918</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>\"</td>\n",
       "      <td>29908</td>\n",
       "      <td>0.99950</td>\n",
       "      <td>\"</td>\n",
       "      <td>29908</td>\n",
       "      <td>0.99950</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>last</td>\n",
       "      <td>4230</td>\n",
       "      <td>0.99999</td>\n",
       "      <td>last</td>\n",
       "      <td>4230</td>\n",
       "      <td>0.99999</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>_</td>\n",
       "      <td>29918</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>_</td>\n",
       "      <td>29918</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>name</td>\n",
       "      <td>978</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>name</td>\n",
       "      <td>978</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>\":</td>\n",
       "      <td>1115</td>\n",
       "      <td>0.99999</td>\n",
       "      <td>\":</td>\n",
       "      <td>1115</td>\n",
       "      <td>0.99999</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>\"</td>\n",
       "      <td>376</td>\n",
       "      <td>0.99997</td>\n",
       "      <td>\"</td>\n",
       "      <td>376</td>\n",
       "      <td>0.99997</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>J</td>\n",
       "      <td>29967</td>\n",
       "      <td>0.99992</td>\n",
       "      <td>J</td>\n",
       "      <td>29967</td>\n",
       "      <td>0.99992</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>ord</td>\n",
       "      <td>536</td>\n",
       "      <td>0.99984</td>\n",
       "      <td>ord</td>\n",
       "      <td>536</td>\n",
       "      <td>0.99984</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>an</td>\n",
       "      <td>273</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>an</td>\n",
       "      <td>273</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>\",</td>\n",
       "      <td>613</td>\n",
       "      <td>0.99996</td>\n",
       "      <td>\",</td>\n",
       "      <td>613</td>\n",
       "      <td>0.99996</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>\\n</td>\n",
       "      <td>13</td>\n",
       "      <td>0.99984</td>\n",
       "      <td>\\n</td>\n",
       "      <td>13</td>\n",
       "      <td>0.99984</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td>\"</td>\n",
       "      <td>29908</td>\n",
       "      <td>0.99984</td>\n",
       "      <td>\"</td>\n",
       "      <td>29908</td>\n",
       "      <td>0.99984</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>year</td>\n",
       "      <td>6360</td>\n",
       "      <td>0.99992</td>\n",
       "      <td>year</td>\n",
       "      <td>6360</td>\n",
       "      <td>0.99992</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25</th>\n",
       "      <td>_</td>\n",
       "      <td>29918</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>_</td>\n",
       "      <td>29918</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26</th>\n",
       "      <td>of</td>\n",
       "      <td>974</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>of</td>\n",
       "      <td>974</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>27</th>\n",
       "      <td>_</td>\n",
       "      <td>29918</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>_</td>\n",
       "      <td>29918</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28</th>\n",
       "      <td>b</td>\n",
       "      <td>29890</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>b</td>\n",
       "      <td>29890</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>29</th>\n",
       "      <td>irth</td>\n",
       "      <td>7515</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>irth</td>\n",
       "      <td>7515</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>30</th>\n",
       "      <td>\":</td>\n",
       "      <td>1115</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>\":</td>\n",
       "      <td>1115</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>31</th>\n",
       "      <td></td>\n",
       "      <td>29871</td>\n",
       "      <td>0.99961</td>\n",
       "      <td></td>\n",
       "      <td>29871</td>\n",
       "      <td>0.99961</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>32</th>\n",
       "      <td>1</td>\n",
       "      <td>29896</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>1</td>\n",
       "      <td>29896</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>33</th>\n",
       "      <td>9</td>\n",
       "      <td>29929</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>9</td>\n",
       "      <td>29929</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>34</th>\n",
       "      <td>6</td>\n",
       "      <td>29953</td>\n",
       "      <td>0.99995</td>\n",
       "      <td>6</td>\n",
       "      <td>29953</td>\n",
       "      <td>0.99995</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>35</th>\n",
       "      <td>3</td>\n",
       "      <td>29941</td>\n",
       "      <td>0.99974</td>\n",
       "      <td>3</td>\n",
       "      <td>29941</td>\n",
       "      <td>0.99974</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>36</th>\n",
       "      <td>,</td>\n",
       "      <td>29892</td>\n",
       "      <td>0.99999</td>\n",
       "      <td>,</td>\n",
       "      <td>29892</td>\n",
       "      <td>0.99999</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>37</th>\n",
       "      <td>\\n</td>\n",
       "      <td>13</td>\n",
       "      <td>0.99910</td>\n",
       "      <td>\\n</td>\n",
       "      <td>13</td>\n",
       "      <td>0.99910</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>38</th>\n",
       "      <td>\"</td>\n",
       "      <td>29908</td>\n",
       "      <td>0.99977</td>\n",
       "      <td>\"</td>\n",
       "      <td>29908</td>\n",
       "      <td>0.99977</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>39</th>\n",
       "      <td>num</td>\n",
       "      <td>1949</td>\n",
       "      <td>0.99981</td>\n",
       "      <td>num</td>\n",
       "      <td>1949</td>\n",
       "      <td>0.99981</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>40</th>\n",
       "      <td>_</td>\n",
       "      <td>29918</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>_</td>\n",
       "      <td>29918</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>41</th>\n",
       "      <td>se</td>\n",
       "      <td>344</td>\n",
       "      <td>0.99999</td>\n",
       "      <td>se</td>\n",
       "      <td>344</td>\n",
       "      <td>0.99999</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>42</th>\n",
       "      <td>asons</td>\n",
       "      <td>7040</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>asons</td>\n",
       "      <td>7040</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>43</th>\n",
       "      <td>_</td>\n",
       "      <td>29918</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>_</td>\n",
       "      <td>29918</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>44</th>\n",
       "      <td>in</td>\n",
       "      <td>262</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>in</td>\n",
       "      <td>262</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>45</th>\n",
       "      <td>_</td>\n",
       "      <td>29918</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>_</td>\n",
       "      <td>29918</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>46</th>\n",
       "      <td>n</td>\n",
       "      <td>29876</td>\n",
       "      <td>1.00000</td>\n",
       "      <td>n</td>\n",
       "      <td>29876</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>47</th>\n",
       "      <td>ba</td>\n",
       "      <td>2291</td>\n",
       "      <td>0.99999</td>\n",
       "      <td>ba</td>\n",
       "      <td>2291</td>\n",
       "      <td>0.99999</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>48</th>\n",
       "      <td>\":</td>\n",
       "      <td>1115</td>\n",
       "      <td>0.99999</td>\n",
       "      <td>\":</td>\n",
       "      <td>1115</td>\n",
       "      <td>0.99999</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>49</th>\n",
       "      <td></td>\n",
       "      <td>29871</td>\n",
       "      <td>1.00000</td>\n",
       "      <td></td>\n",
       "      <td>29871</td>\n",
       "      <td>1.00000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50</th>\n",
       "      <td>1</td>\n",
       "      <td>29896</td>\n",
       "      <td>0.99996</td>\n",
       "      <td>1</td>\n",
       "      <td>29896</td>\n",
       "      <td>0.99996</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>51</th>\n",
       "      <td>5</td>\n",
       "      <td>29945</td>\n",
       "      <td>0.99677</td>\n",
       "      <td>5</td>\n",
       "      <td>29945</td>\n",
       "      <td>0.99677</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>52</th>\n",
       "      <td>\\n</td>\n",
       "      <td>13</td>\n",
       "      <td>0.73163</td>\n",
       "      <td>\\n</td>\n",
       "      <td>13</td>\n",
       "      <td>0.73163</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>53</th>\n",
       "      <td>}</td>\n",
       "      <td>29913</td>\n",
       "      <td>0.71903</td>\n",
       "      <td>}</td>\n",
       "      <td>29913</td>\n",
       "      <td>0.71903</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>54</th>\n",
       "      <td>&lt;/s&gt;</td>\n",
       "      <td>2</td>\n",
       "      <td>0.70978</td>\n",
       "      <td>&lt;/s&gt;</td>\n",
       "      <td>2</td>\n",
       "      <td>0.70978</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   generated_token  generated_token_idx  generated_score leading_token  leading_token_idx  leading_score\n",
       "0                                 29871          0.99982                            29871        0.99982\n",
       "1                {                  426          0.00002            Of               4587        0.95649\n",
       "2               \\n                   13          0.99705            \\n                 13        0.99705\n",
       "3                \"                29908          0.98198             \"              29908        0.98198\n",
       "4            first                 4102          0.01030         title               3257        0.55388\n",
       "5                _                29918          0.99879             _              29918        0.99879\n",
       "6             name                  978          1.00000          name                978        1.00000\n",
       "7               \":                 1115          0.99994            \":               1115        0.99994\n",
       "8                \"                  376          0.99949             \"                376        0.99949\n",
       "9          Michael                24083          0.99028       Michael              24083        0.99028\n",
       "10              \",                  613          0.99112            \",                613        0.99112\n",
       "11              \\n                   13          0.99918            \\n                 13        0.99918\n",
       "12               \"                29908          0.99950             \"              29908        0.99950\n",
       "13            last                 4230          0.99999          last               4230        0.99999\n",
       "14               _                29918          1.00000             _              29918        1.00000\n",
       "15            name                  978          1.00000          name                978        1.00000\n",
       "16              \":                 1115          0.99999            \":               1115        0.99999\n",
       "17               \"                  376          0.99997             \"                376        0.99997\n",
       "18               J                29967          0.99992             J              29967        0.99992\n",
       "19             ord                  536          0.99984           ord                536        0.99984\n",
       "20              an                  273          1.00000            an                273        1.00000\n",
       "21              \",                  613          0.99996            \",                613        0.99996\n",
       "22              \\n                   13          0.99984            \\n                 13        0.99984\n",
       "23               \"                29908          0.99984             \"              29908        0.99984\n",
       "24            year                 6360          0.99992          year               6360        0.99992\n",
       "25               _                29918          1.00000             _              29918        1.00000\n",
       "26              of                  974          1.00000            of                974        1.00000\n",
       "27               _                29918          1.00000             _              29918        1.00000\n",
       "28               b                29890          1.00000             b              29890        1.00000\n",
       "29            irth                 7515          1.00000          irth               7515        1.00000\n",
       "30              \":                 1115          1.00000            \":               1115        1.00000\n",
       "31                                29871          0.99961                            29871        0.99961\n",
       "32               1                29896          1.00000             1              29896        1.00000\n",
       "33               9                29929          1.00000             9              29929        1.00000\n",
       "34               6                29953          0.99995             6              29953        0.99995\n",
       "35               3                29941          0.99974             3              29941        0.99974\n",
       "36               ,                29892          0.99999             ,              29892        0.99999\n",
       "37              \\n                   13          0.99910            \\n                 13        0.99910\n",
       "38               \"                29908          0.99977             \"              29908        0.99977\n",
       "39             num                 1949          0.99981           num               1949        0.99981\n",
       "40               _                29918          1.00000             _              29918        1.00000\n",
       "41              se                  344          0.99999            se                344        0.99999\n",
       "42           asons                 7040          1.00000         asons               7040        1.00000\n",
       "43               _                29918          1.00000             _              29918        1.00000\n",
       "44              in                  262          1.00000            in                262        1.00000\n",
       "45               _                29918          1.00000             _              29918        1.00000\n",
       "46               n                29876          1.00000             n              29876        1.00000\n",
       "47              ba                 2291          0.99999            ba               2291        0.99999\n",
       "48              \":                 1115          0.99999            \":               1115        0.99999\n",
       "49                                29871          1.00000                            29871        1.00000\n",
       "50               1                29896          0.99996             1              29896        0.99996\n",
       "51               5                29945          0.99677             5              29945        0.99677\n",
       "52              \\n                   13          0.73163            \\n                 13        0.73163\n",
       "53               }                29913          0.71903             }              29913        0.71903\n",
       "54            </s>                    2          0.70978          </s>                  2        0.70978"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_header(\"Enforcer intervention table\")\n",
    "display(enforced_scores)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Token index 3 shows where the enforcer had to be aggressive:\n",
    "\n",
    "\n",
    "```\n",
    "generated_token  generated_token_idx  generated_score leading_token  leading_token_idx  leading_score\n",
    "\n",
    "3            first                 4102         0.00032         title               3257       0.94433\n",
    "```\n",
    "\n",
    "The language model was trying to generate the word \"title\" (post softmax score of 0.94433) but the enforcer made the \"first\" token (post softmax score of 0.00032) be the main candidate instead.\n",
    "This can be used to further improve the prompt engineering, as we generally want to avoid timesteps that cause the enforcer to be this aggressive, as it increases the likelyhood of hallucinations. For example, The [langchain project removes the \"title\" from the json schema](https://github.com/langchain-ai/langchain/blob/cfa2203c626a2287d60c1febeb3e3a68b77acd77/libs/langchain/langchain/output_parsers/pydantic.py#L40), probably for this reason."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Regular Expression Use Case\n",
    "\n",
    "Note that the ```RegexParser``` does not support the full regex syntax, as it uses [interegular](https://pypi.org/project/interegular/) under the hood. We will also make use of ```StringParser``` to diagnose the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "**Question:**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       "When was Michael Jordan Born? Please answer in mm/dd/yyyy format.\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Without enforcing:**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       "  Thank you for reaching out! I'm happy to help you with your question. However, I must inform you that Michael Jordan was born on February 17, 1963, which is not possible in the mm/dd/yyyy format as it is a date that falls in the 20th century. I apologize for any confusion. Is there anything else I can help you with?\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "----------------------------------------\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "**With regex force. Regex: ``` In mm/dd/yyyy format, Michael Jordan was born in (0?[1-9]|1[0-2])\\/(0?[1-9]|1\\d|2\\d|3[01])\\/(19|20)\\d{2}```**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Language model output:**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       " In mm/dd/yyyy format, Michael Jordan was born in 12/23/1963\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Enforcer intervention table:**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>generated_token</th>\n",
       "      <th>generated_token_idx</th>\n",
       "      <th>generated_score</th>\n",
       "      <th>leading_token</th>\n",
       "      <th>leading_token_idx</th>\n",
       "      <th>leading_score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td></td>\n",
       "      <td>29871</td>\n",
       "      <td>0.99993</td>\n",
       "      <td></td>\n",
       "      <td>29871</td>\n",
       "      <td>0.99993</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>I</td>\n",
       "      <td>29902</td>\n",
       "      <td>0.00000</td>\n",
       "      <td>I</td>\n",
       "      <td>306</td>\n",
       "      <td>0.59019</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>n</td>\n",
       "      <td>29876</td>\n",
       "      <td>0.00000</td>\n",
       "      <td>apolog</td>\n",
       "      <td>27746</td>\n",
       "      <td>0.63013</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td></td>\n",
       "      <td>29871</td>\n",
       "      <td>0.00091</td>\n",
       "      <td>st</td>\n",
       "      <td>303</td>\n",
       "      <td>0.67129</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>mm</td>\n",
       "      <td>4317</td>\n",
       "      <td>0.00002</td>\n",
       "      <td>1</td>\n",
       "      <td>29896</td>\n",
       "      <td>0.53647</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>/</td>\n",
       "      <td>29914</td>\n",
       "      <td>0.99731</td>\n",
       "      <td>/</td>\n",
       "      <td>29914</td>\n",
       "      <td>0.99731</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>dd</td>\n",
       "      <td>1289</td>\n",
       "      <td>0.99882</td>\n",
       "      <td>dd</td>\n",
       "      <td>1289</td>\n",
       "      <td>0.99882</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>/</td>\n",
       "      <td>29914</td>\n",
       "      <td>0.99897</td>\n",
       "      <td>/</td>\n",
       "      <td>29914</td>\n",
       "      <td>0.99897</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>yyyy</td>\n",
       "      <td>18855</td>\n",
       "      <td>0.99702</td>\n",
       "      <td>yyyy</td>\n",
       "      <td>18855</td>\n",
       "      <td>0.99702</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>format</td>\n",
       "      <td>3402</td>\n",
       "      <td>0.99611</td>\n",
       "      <td>format</td>\n",
       "      <td>3402</td>\n",
       "      <td>0.99611</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>,</td>\n",
       "      <td>29892</td>\n",
       "      <td>0.99142</td>\n",
       "      <td>,</td>\n",
       "      <td>29892</td>\n",
       "      <td>0.99142</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>Michael</td>\n",
       "      <td>5765</td>\n",
       "      <td>0.99273</td>\n",
       "      <td>Michael</td>\n",
       "      <td>5765</td>\n",
       "      <td>0.99273</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>Jordan</td>\n",
       "      <td>18284</td>\n",
       "      <td>0.99815</td>\n",
       "      <td>Jordan</td>\n",
       "      <td>18284</td>\n",
       "      <td>0.99815</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>was</td>\n",
       "      <td>471</td>\n",
       "      <td>0.99543</td>\n",
       "      <td>was</td>\n",
       "      <td>471</td>\n",
       "      <td>0.99543</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>born</td>\n",
       "      <td>6345</td>\n",
       "      <td>0.98757</td>\n",
       "      <td>born</td>\n",
       "      <td>6345</td>\n",
       "      <td>0.98757</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>in</td>\n",
       "      <td>297</td>\n",
       "      <td>0.00090</td>\n",
       "      <td>on</td>\n",
       "      <td>373</td>\n",
       "      <td>0.99826</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td></td>\n",
       "      <td>29871</td>\n",
       "      <td>0.57364</td>\n",
       "      <td></td>\n",
       "      <td>29871</td>\n",
       "      <td>0.57364</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>1</td>\n",
       "      <td>29896</td>\n",
       "      <td>0.53364</td>\n",
       "      <td>1</td>\n",
       "      <td>29896</td>\n",
       "      <td>0.53364</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>2</td>\n",
       "      <td>29906</td>\n",
       "      <td>0.00075</td>\n",
       "      <td>9</td>\n",
       "      <td>29929</td>\n",
       "      <td>0.99794</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>/</td>\n",
       "      <td>29914</td>\n",
       "      <td>0.99044</td>\n",
       "      <td>/</td>\n",
       "      <td>29914</td>\n",
       "      <td>0.99044</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>2</td>\n",
       "      <td>29906</td>\n",
       "      <td>0.80903</td>\n",
       "      <td>2</td>\n",
       "      <td>29906</td>\n",
       "      <td>0.80903</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>3</td>\n",
       "      <td>29941</td>\n",
       "      <td>0.36755</td>\n",
       "      <td>3</td>\n",
       "      <td>29941</td>\n",
       "      <td>0.36755</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>/</td>\n",
       "      <td>29914</td>\n",
       "      <td>0.99938</td>\n",
       "      <td>/</td>\n",
       "      <td>29914</td>\n",
       "      <td>0.99938</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td>1</td>\n",
       "      <td>29896</td>\n",
       "      <td>0.99812</td>\n",
       "      <td>1</td>\n",
       "      <td>29896</td>\n",
       "      <td>0.99812</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>9</td>\n",
       "      <td>29929</td>\n",
       "      <td>0.99975</td>\n",
       "      <td>9</td>\n",
       "      <td>29929</td>\n",
       "      <td>0.99975</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25</th>\n",
       "      <td>6</td>\n",
       "      <td>29953</td>\n",
       "      <td>0.99162</td>\n",
       "      <td>6</td>\n",
       "      <td>29953</td>\n",
       "      <td>0.99162</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26</th>\n",
       "      <td>3</td>\n",
       "      <td>29941</td>\n",
       "      <td>0.99485</td>\n",
       "      <td>3</td>\n",
       "      <td>29941</td>\n",
       "      <td>0.99485</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>27</th>\n",
       "      <td>&lt;/s&gt;</td>\n",
       "      <td>2</td>\n",
       "      <td>0.00501</td>\n",
       "      <td>.</td>\n",
       "      <td>29889</td>\n",
       "      <td>0.98415</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   generated_token  generated_token_idx  generated_score leading_token  leading_token_idx  leading_score\n",
       "0                                 29871          0.99993                            29871        0.99993\n",
       "1                I                29902          0.00000             I                306        0.59019\n",
       "2                n                29876          0.00000        apolog              27746        0.63013\n",
       "3                                 29871          0.00091            st                303        0.67129\n",
       "4               mm                 4317          0.00002             1              29896        0.53647\n",
       "5                /                29914          0.99731             /              29914        0.99731\n",
       "6               dd                 1289          0.99882            dd               1289        0.99882\n",
       "7                /                29914          0.99897             /              29914        0.99897\n",
       "8             yyyy                18855          0.99702          yyyy              18855        0.99702\n",
       "9           format                 3402          0.99611        format               3402        0.99611\n",
       "10               ,                29892          0.99142             ,              29892        0.99142\n",
       "11         Michael                 5765          0.99273       Michael               5765        0.99273\n",
       "12          Jordan                18284          0.99815        Jordan              18284        0.99815\n",
       "13             was                  471          0.99543           was                471        0.99543\n",
       "14            born                 6345          0.98757          born               6345        0.98757\n",
       "15              in                  297          0.00090            on                373        0.99826\n",
       "16                                29871          0.57364                            29871        0.57364\n",
       "17               1                29896          0.53364             1              29896        0.53364\n",
       "18               2                29906          0.00075             9              29929        0.99794\n",
       "19               /                29914          0.99044             /              29914        0.99044\n",
       "20               2                29906          0.80903             2              29906        0.80903\n",
       "21               3                29941          0.36755             3              29941        0.36755\n",
       "22               /                29914          0.99938             /              29914        0.99938\n",
       "23               1                29896          0.99812             1              29896        0.99812\n",
       "24               9                29929          0.99975             9              29929        0.99975\n",
       "25               6                29953          0.99162             6              29953        0.99162\n",
       "26               3                29941          0.99485             3              29941        0.99485\n",
       "27            </s>                    2          0.00501             .              29889        0.98415"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "----------------------------------------\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "**With string force:**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Language model output:**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       " Michael Jordan was born in 02/17/1963\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Enforcer intervention table:**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>generated_token</th>\n",
       "      <th>generated_token_idx</th>\n",
       "      <th>generated_score</th>\n",
       "      <th>leading_token</th>\n",
       "      <th>leading_token_idx</th>\n",
       "      <th>leading_score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td></td>\n",
       "      <td>29871</td>\n",
       "      <td>0.99993</td>\n",
       "      <td></td>\n",
       "      <td>29871</td>\n",
       "      <td>0.99993</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Michael</td>\n",
       "      <td>24083</td>\n",
       "      <td>0.00000</td>\n",
       "      <td>I</td>\n",
       "      <td>306</td>\n",
       "      <td>0.58970</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Jordan</td>\n",
       "      <td>18284</td>\n",
       "      <td>0.99998</td>\n",
       "      <td>Jordan</td>\n",
       "      <td>18284</td>\n",
       "      <td>0.99998</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>was</td>\n",
       "      <td>471</td>\n",
       "      <td>0.99966</td>\n",
       "      <td>was</td>\n",
       "      <td>471</td>\n",
       "      <td>0.99966</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>born</td>\n",
       "      <td>6345</td>\n",
       "      <td>0.99996</td>\n",
       "      <td>born</td>\n",
       "      <td>6345</td>\n",
       "      <td>0.99996</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>in</td>\n",
       "      <td>297</td>\n",
       "      <td>0.00003</td>\n",
       "      <td>on</td>\n",
       "      <td>373</td>\n",
       "      <td>0.99997</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td></td>\n",
       "      <td>29871</td>\n",
       "      <td>0.07737</td>\n",
       "      <td>Fort</td>\n",
       "      <td>7236</td>\n",
       "      <td>0.47024</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>0</td>\n",
       "      <td>29900</td>\n",
       "      <td>0.00822</td>\n",
       "      <td>1</td>\n",
       "      <td>29896</td>\n",
       "      <td>0.97270</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>2</td>\n",
       "      <td>29906</td>\n",
       "      <td>0.99314</td>\n",
       "      <td>2</td>\n",
       "      <td>29906</td>\n",
       "      <td>0.99314</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>/</td>\n",
       "      <td>29914</td>\n",
       "      <td>0.99944</td>\n",
       "      <td>/</td>\n",
       "      <td>29914</td>\n",
       "      <td>0.99944</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>1</td>\n",
       "      <td>29896</td>\n",
       "      <td>0.02840</td>\n",
       "      <td>2</td>\n",
       "      <td>29906</td>\n",
       "      <td>0.86993</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>7</td>\n",
       "      <td>29955</td>\n",
       "      <td>0.99795</td>\n",
       "      <td>7</td>\n",
       "      <td>29955</td>\n",
       "      <td>0.99795</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>/</td>\n",
       "      <td>29914</td>\n",
       "      <td>0.99996</td>\n",
       "      <td>/</td>\n",
       "      <td>29914</td>\n",
       "      <td>0.99996</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>1</td>\n",
       "      <td>29896</td>\n",
       "      <td>0.99980</td>\n",
       "      <td>1</td>\n",
       "      <td>29896</td>\n",
       "      <td>0.99980</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>9</td>\n",
       "      <td>29929</td>\n",
       "      <td>0.99997</td>\n",
       "      <td>9</td>\n",
       "      <td>29929</td>\n",
       "      <td>0.99997</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>6</td>\n",
       "      <td>29953</td>\n",
       "      <td>0.99955</td>\n",
       "      <td>6</td>\n",
       "      <td>29953</td>\n",
       "      <td>0.99955</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>3</td>\n",
       "      <td>29941</td>\n",
       "      <td>0.99924</td>\n",
       "      <td>3</td>\n",
       "      <td>29941</td>\n",
       "      <td>0.99924</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>&lt;/s&gt;</td>\n",
       "      <td>2</td>\n",
       "      <td>0.00058</td>\n",
       "      <td>.</td>\n",
       "      <td>29889</td>\n",
       "      <td>0.93155</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   generated_token  generated_token_idx  generated_score leading_token  leading_token_idx  leading_score\n",
       "0                                 29871          0.99993                            29871        0.99993\n",
       "1          Michael                24083          0.00000             I                306        0.58970\n",
       "2           Jordan                18284          0.99998        Jordan              18284        0.99998\n",
       "3              was                  471          0.99966           was                471        0.99966\n",
       "4             born                 6345          0.99996          born               6345        0.99996\n",
       "5               in                  297          0.00003            on                373        0.99997\n",
       "6                                 29871          0.07737          Fort               7236        0.47024\n",
       "7                0                29900          0.00822             1              29896        0.97270\n",
       "8                2                29906          0.99314             2              29906        0.99314\n",
       "9                /                29914          0.99944             /              29914        0.99944\n",
       "10               1                29896          0.02840             2              29906        0.86993\n",
       "11               7                29955          0.99795             7              29955        0.99795\n",
       "12               /                29914          0.99996             /              29914        0.99996\n",
       "13               1                29896          0.99980             1              29896        0.99980\n",
       "14               9                29929          0.99997             9              29929        0.99997\n",
       "15               6                29953          0.99955             6              29953        0.99955\n",
       "16               3                29941          0.99924             3              29941        0.99924\n",
       "17            </s>                    2          0.00058             .              29889        0.93155"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "DEFAULT_SYSTEM_PROMPT = \"\"\"\\\n",
    "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\\n\\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\\\n",
    "\"\"\"\n",
    "MAX_MAX_NEW_TOKENS = 200\n",
    "DEFAULT_MAX_NEW_TOKENS = 100\n",
    "MAX_INPUT_TOKEN_LENGTH = 4000\n",
    "\n",
    "date_regex = r'(0?[1-9]|1[0-2])\\/(0?[1-9]|1\\d|2\\d|3[01])\\/(19|20)\\d{2}'\n",
    "answer_regex = ' In mm/dd/yyyy format, Michael Jordan was born in ' + date_regex\n",
    "question = 'When was Michael Jordan Born? Please answer in mm/dd/yyyy format.'\n",
    "display_header(\"Question:\")\n",
    "display_content(question)\n",
    "\n",
    "display_header(\"Without enforcing:\")\n",
    "result, _ = run(question, system_prompt=DEFAULT_SYSTEM_PROMPT, max_new_tokens=DEFAULT_MAX_NEW_TOKENS)\n",
    "display_content(result)\n",
    "\n",
    "print('\\n----------------------------------------\\n')\n",
    "\n",
    "display_header(f\"With regex force. Regex: ```{answer_regex}```\")\n",
    "result, enforced_scores = run(question, system_prompt=DEFAULT_SYSTEM_PROMPT, max_new_tokens=DEFAULT_MAX_NEW_TOKENS, required_regex=answer_regex)\n",
    "display_header(\"Language model output:\")\n",
    "display_content(result)\n",
    "display_header(\"Enforcer intervention table:\")\n",
    "display(enforced_scores)\n",
    "\n",
    "print('\\n----------------------------------------\\n')\n",
    "\n",
    "display_header(\"With string force:\")\n",
    "answer = ' Michael Jordan was born in 02/17/1963'\n",
    "result, enforced_scores = run(question, system_prompt=DEFAULT_SYSTEM_PROMPT, max_new_tokens=DEFAULT_MAX_NEW_TOKENS, required_str=answer)\n",
    "display_header(\"Language model output:\")\n",
    "display_content(result)\n",
    "display_header(\"Enforcer intervention table:\")\n",
    "display(enforced_scores)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The unenforced example does not adhere to the format requirements at all.\n",
    "\n",
    "Note that in the string force section, the first digit that the LM wants to generate is a 1, not a 0 (which we would expect in 02/17/1963):\n",
    "\n",
    "```\n",
    "\tgenerated_token\tgenerated_token_idx\tgenerated_score\tleading_token\tleading_token_idx\tleading_score\n",
    "7\t0\t29900\t0.03939\t1\t29896\t0.95441\n",
    "```\n",
    "\n",
    "This is likely due to the prompt not emphasizing the format requirements enough, causing the LM to want to output the year first (1963 starts with 1). In the regex example, we add a ```In mm/dd/yyyy format, ``` prefix to the regex, making 0 a pretty confident leader in that case:\n",
    "\n",
    "```\n",
    "\tgenerated_token\tgenerated_token_idx\tgenerated_score\tleading_token\tleading_token_idx\tleading_score\n",
    "16\t0\t29900\t0.91775\t0\t29900\t0.91775\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Batching example\n",
    "\n",
    "This is a simple example of using batching to generate multiple queries in parallel. All outputs will be in the correct format. Every timestep can filter different tokens for the different batch indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "**Batched Answers, With json schema enforcing:**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       "  { \"first_name\": \"Michael\", \"last_name\": \"Jordan\", \"year_of_birth\": 1963, \"num_seasons_in_nba\": 15 }\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       "  { \"first_name\": \"Timothy\", \"last_name\": \"Duncan\", \"year_of_birth\": 1976, \"num_seasons_in_nba\": 19 }\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       "  { \"first_name\": \"Kobe\", \"last_name\": \"Bryant\", \"year_of_birth\": 1978, \"num_seasons_in_nba\": 20 }\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       "  { \"first_name\": \"Kareem Abdul-Jabbar\", \"last_name\": \"Abdul-Jabbar\", \"year_of_birth\": 1947, \"num_seasons_in_nba\": 20 }\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "PLAYER_NAMES = ['Michael Jordan', 'Tim Duncan', 'Kobe Bryant', 'Kareem Abdul Jabbar']\n",
    "question = 'Please give me information about {0}. You MUST answer using the following json schema: '\n",
    "questions_with_schema = [f'{question.format(player_name)}{AnswerFormat.schema_json()}' for player_name in PLAYER_NAMES]\n",
    "\n",
    "display_header(\"Batched Answers, With json schema enforcing:\")\n",
    "results, _ = run(questions_with_schema, system_prompt=DEFAULT_SYSTEM_PROMPT, max_new_tokens=DEFAULT_MAX_NEW_TOKENS, required_json_schema=AnswerFormat.schema())\n",
    "for result in results:\n",
    "    display_content(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LangChain Integration\n",
    "\n",
    "This is a simple example of how to integrate the enforcer into the LangChain project.\n",
    "\n",
    "\n",
    "For a more complete explanation, see the [LM Format Enforcer Page in the Langchain documentation](https://python.langchain.com/docs/integrations/llms/lmformatenforcer_experimental).\n",
    "\n",
    "This demo shows the JSON use case, the regex case is also supported."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "**Call mode**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       "   {\n",
       "\n",
       "\"last_name\" : \"Jordan\",\n",
       "\"first_name\" : \"Michael\",\n",
       "\"year_of_birth\" : 1963,\n",
       "\"num_seasons_in_nba\" : 15\n",
       "}\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Batched mode**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       "   {\n",
       "\n",
       "\"last_name\" : \"Jordan\",\n",
       "\"first_name\" : \"Michael\",\n",
       "\"year_of_birth\" : 1963,\n",
       "\"num_seasons_in_nba\" : 15\n",
       "}\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       " \n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "{\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\"last_name\": \"Larry Bird\",\n",
       "\"num_seasons_in_nba\": 13,\n",
       "\"year_of_birth\": 1956,\n",
       "\"first_name\": \"Larry\"\n",
       "}\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       "   {\n",
       "\n",
       "\"first_name\": \"Tim\",\n",
       "\"last_name\": \"Duncan\",\n",
       "\"num_seasons_in_nba\": 19,\n",
       "\"year_of_birth\": 1976\n",
       "}\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import pipeline\n",
    "from langchain_experimental.llms import LMFormatEnforcer\n",
    "\n",
    "# We create a transformers pipeline to avoid loading the model twice, but we could also use LMFormatEnforcer.from_model_id()\n",
    "pipe = pipeline(\"text-generation\", model=model, tokenizer=tokenizer, max_new_tokens=100)\n",
    "langchain_pipeline = LMFormatEnforcer(pipeline=pipe, json_schema=AnswerFormat.schema())\n",
    "\n",
    "DEFAULT_SYSTEM_PROMPT = \"\"\"\\\n",
    "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\\n\\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\\\n",
    "\"\"\"\n",
    "\n",
    "players = ['Michael Jordan', 'Larry Bird', 'Tim Duncan']\n",
    "question = 'Please give me information about {}.'\n",
    "prompts = [get_prompt(question.format(player), DEFAULT_SYSTEM_PROMPT) for player in players]\n",
    "\n",
    "display_header('Call mode')\n",
    "result = langchain_pipeline(prompts[0])\n",
    "display_content(result)\n",
    "\n",
    "display_header('Batched mode')\n",
    "results = langchain_pipeline.generate(prompts)\n",
    "for generation in results.generations:\n",
    "    display_content(generation[0].text)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lmformatenforcer",
   "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.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}