samples/colab_exllamav2_integration.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LM Format Enforcer Integration with ExLlamaV2\n",
    "\n",
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/noamgat/lm-format-enforcer/blob/main/samples/colab_exllamav2_integration.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>\n",
    "\n",
    "This notebook shows how you can integrate with the [ExLlamaV2](https://github.com/turboderp/exllamav2/) library. We do it using it's Sampler Filter interface and the integration class in this repository.\n",
    "\n",
    "ExLlamaV2 is one of the fastest inference engines, but does not support any of the popular constrained decoding libraries."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the COLAB runtime (user action required)\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).\n",
    "\n",
    "## Installing dependencies\n",
    "\n",
    "We begin by installing the dependencies.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install exllamav2 lm-format-enforcer huggingface-hub\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": [
    "## Loading the model\n",
    "\n",
    "This demo uses [Llama2 bpw weights by turboderp](https://huggingface.co/turboderp/Llama2-7B-exl2/tree/8.0bpw). We will use huggingface hub to download the model."
   ]
  },
  {
   "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",
      "Fetching 16 files: 100%|██████████| 16/16 [00:00<00:00, 215784.13it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloaded model to /mnt/wsl/PHYSICALDRIVE1p3/huggingface/hub/models--turboderp--Llama2-7B-exl2/snapshots/6463dd96f3694a87b777852f8bd979dbaeb2b839\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from huggingface_hub import snapshot_download\n",
    "model_directory = snapshot_download(repo_id=\"turboderp/Llama2-7B-exl2\", revision=\"6463dd96f3694a87b777852f8bd979dbaeb2b839\")\n",
    "print(f\"Downloaded model to {model_directory}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing ExLlamaV2\n",
    "\n",
    "We follow the [inference.py example](https://github.com/turboderp/exllamav2/blob/master/examples/inference.py) from the ExLlamaV2 repo. There is no one-liner setup at the moment, so the next cell will contain quite a bit of code. It is all from the example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model: /mnt/wsl/PHYSICALDRIVE1p3/huggingface/hub/models--turboderp--Llama2-7B-exl2/snapshots/6463dd96f3694a87b777852f8bd979dbaeb2b839\n"
     ]
    }
   ],
   "source": [
    "from exllamav2 import(\n",
    "    ExLlamaV2,\n",
    "    ExLlamaV2Config,\n",
    "    ExLlamaV2Cache,\n",
    "    ExLlamaV2Tokenizer,\n",
    ")\n",
    "\n",
    "from exllamav2.generator import (\n",
    "    ExLlamaV2BaseGenerator,\n",
    "    ExLlamaV2Sampler\n",
    ")\n",
    "\n",
    "# Initialize model and cache\n",
    "\n",
    "config = ExLlamaV2Config()\n",
    "config.model_dir = model_directory\n",
    "config.prepare()\n",
    "\n",
    "model = ExLlamaV2(config)\n",
    "print(\"Loading model: \" + model_directory)\n",
    "\n",
    "cache = ExLlamaV2Cache(model, lazy = True)\n",
    "model.load_autosplit(cache)\n",
    "\n",
    "tokenizer = ExLlamaV2Tokenizer(config)\n",
    "\n",
    "# Initialize generator\n",
    "\n",
    "generator = ExLlamaV2BaseGenerator(model, cache, tokenizer)\n",
    "\n",
    "# Prepare settings\n",
    "\n",
    "settings = ExLlamaV2Sampler.Settings()\n",
    "settings.temperature = 0.85\n",
    "settings.top_k = 50\n",
    "settings.top_p = 0.8\n",
    "# settings.disallow_tokens(tokenizer, [tokenizer.eos_token_id])\n",
    "\n",
    "max_new_tokens = 150\n",
    "\n",
    "generator.warmup()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the previous cell executed successfully, you have propertly set up your Colab runtime and loaded the ExLlamaV2 model!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A few helper functions to make display nicer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, Markdown\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating text with the LM Format Enforcer Logits Processor\n",
    "\n",
    "ExLlamaV2's `Sampler.Settings` have a `filters` interface similar to one that exists in Huggingface Transformers. We will connect to this API and filter the forbidden logits.\n",
    "\n",
    "The integration class `ExLlamaV2TokenEnforcerFilter` does just that. This is the ONLY integration point between lm-format-enforcer and ExLlamaV2.\n",
    "\n",
    "Note that in this notebook we use `generate_simple()`, but the integration works with all ExLlamaV2 generation methods.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lmformatenforcer.characterlevelparser import CharacterLevelParser\n",
    "from lmformatenforcer.integrations.exllamav2 import ExLlamaV2TokenEnforcerFilter, build_token_enforcer_tokenizer_data\n",
    "from typing import Optional\n",
    "\n",
    "# Building the tokenizer data once is a performance optimization, it saves preprocessing in subsequent calls.\n",
    "tokenizer_data = build_token_enforcer_tokenizer_data(tokenizer)\n",
    "\n",
    "def exllamav2_with_format_enforcer(prompt: str, parser: Optional[CharacterLevelParser] = None) -> str:\n",
    "    if parser is None:\n",
    "        settings.filters = []\n",
    "    else:\n",
    "        settings.filters = [ExLlamaV2TokenEnforcerFilter(parser, tokenizer_data)]\n",
    "    result = generator.generate_simple(prompt, settings, max_new_tokens, seed = 1234)\n",
    "    return result[len(prompt):]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ExLlamaV2 + JSON 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": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_39692/2472056033.py:13: PydanticDeprecatedSince20: The `schema_json` method is deprecated; use `model_json_schema` and json.dumps instead. Deprecated in Pydantic V2.0 to be removed in V3.0. See Pydantic V2 Migration Guide at https://errors.pydantic.dev/2.5/migration/\n",
      "  question_with_schema = f'{question}{AnswerFormat.schema_json()}'\n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "**Prompt:**"
      ],
      "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: {\"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\"], \"title\": \"AnswerFormat\", \"type\": \"object\"}\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",
       " I have tried: The schema is not well formed. Please give me an example of how to do so.\n",
       "\n",
       "Comment: Can you show your attempt?\n",
       "\n",
       "Comment: @Shawn, I have updated my question with more details on what I have tried and what I would like.\n",
       "\n",
       "Comment: Your schema doesn't seem to have any properties that are required.  Can you try removing `required` from your schema?\n",
       "\n",
       "Comment: @Shawn, I removed the `required` key and it still returns the same error: The schema is not well formed.\n",
       "\n",
       "Comment: Are you sure you're trying to validate against the JSON you provided in your question?  Because that is invalid JSON (\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": [
      "/tmp/ipykernel_39692/2472056033.py:24: PydanticDeprecatedSince20: The `schema` method is deprecated; use `model_json_schema` instead. Deprecated in Pydantic V2.0 to be removed in V3.0. See Pydantic V2 Migration Guide at https://errors.pydantic.dev/2.5/migration/\n",
      "  parser = JsonSchemaParser(AnswerFormat.schema())\n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       "\n",
       "\n",
       "    {\n",
       "        \"first_name\": \"Michael\",\n",
       "        \"last_name\": \"Jordan\",\n",
       "        \"year_of_birth\": 1963,\n",
       "        \"num_seasons_in_nba\": 17\n",
       "    }\n",
       "\n",
       "   \n",
       "\n",
       "    \n",
       "\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Answer, With json mode (json output, no specific schema) enforcing:**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       "\n",
       "\n",
       "    {\"error\": true, \"message\": \"Please give me information about Michael Jordan. You MUST answer using the following json schema: {\"\n",
       "\n",
       "    }\n",
       "\n",
       "   \n",
       "\n",
       "   \n",
       "\n",
       "\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from lmformatenforcer import JsonSchemaParser\n",
    "from pydantic import BaseModel\n",
    "\n",
    "\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",
    "prompt = question_with_schema\n",
    "\n",
    "display_header(\"Prompt:\")\n",
    "display_content(prompt)\n",
    "\n",
    "display_header(\"Answer, Without json schema enforcing:\")\n",
    "result = exllamav2_with_format_enforcer(prompt, parser=None)\n",
    "display_content(result)\n",
    "\n",
    "display_header(\"Answer, With json schema enforcing:\")\n",
    "parser = JsonSchemaParser(AnswerFormat.schema())\n",
    "result = exllamav2_with_format_enforcer(prompt, parser=parser)\n",
    "display_content(result)\n",
    "\n",
    "display_header(\"Answer, With json mode (json output, no specific schema) enforcing:\")\n",
    "parser = JsonSchemaParser(None)\n",
    "result = exllamav2_with_format_enforcer(prompt, parser=parser)\n",
    "display_content(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the enforced output matches the required schema, while the unenforced does not. We have successfully integrated with ExLlamaV2!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ExLlamaV2 + Regular Expressions Use Case\n",
    "\n",
    "LM Format Enforcer can also be used to make sure that the output matches a regular expression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "**Prompt:**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       "Q: When was Michael Jordan Born? Please answer in mm/dd/yyyy format. A:\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Without format forcing:**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "```\n",
       " He was born on Feb 17, 1963.\n",
       "Q: What is his height? Please answer in feet and inches format. A: According to Celebrity Height, he is 6’6” tall.\n",
       "Q: What is his weight? Please answer in pounds and ounces format. A: According to Celebrity Weight, he weighs around 205 lb (93 kg).\n",
       "Q: What is his net worth? Please answer in dollars and cents format. A: According to Celebrity Net Worth, his net worth is $1.4 Billion.\n",
       "Q: Who is his wife? Please answer in full name format.\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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": [
       "```\n",
       " In mm/dd/yyyy format, Michael Jordan was born in 2/17/1963\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from lmformatenforcer import RegexParser\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 = 'Q: When was Michael Jordan Born? Please answer in mm/dd/yyyy format. A:'\n",
    "prompt = question\n",
    "\n",
    "display_header(\"Prompt:\")\n",
    "display_content(prompt)\n",
    "\n",
    "\n",
    "display_header(\"Without format forcing:\")\n",
    "result = exllamav2_with_format_enforcer(prompt, parser=None)\n",
    "display_content(result)\n",
    "\n",
    "\n",
    "display_header(f\"With regex force. Regex: ```{answer_regex}```\")\n",
    "parser = RegexParser(answer_regex)\n",
    "result = exllamav2_with_format_enforcer(prompt, parser=parser)\n",
    "display_content(result)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, with regex forcing enabled, we got a valid output. Without it, we did not get it in the structure that we wanted."
   ]
  }
 ],
 "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.14"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}