examples/transformers-gg.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-09-01T14:21:04.983177Z",
     "start_time": "2024-09-01T14:21:02.169487Z"
    }
   },
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed\n",
    "import guidance\n",
    "import json \n",
    "from dotenv import load_dotenv\n",
    "\n",
    "import grammar_guide as gg\n",
    "\n",
    "load_dotenv()\n",
    "set_seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [],
   "source": [
    "model_name_or_path = \"HuggingFaceTB/SmolLM-135M\"\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name_or_path)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-09-01T14:21:05.715126Z",
     "start_time": "2024-09-01T14:21:04.983837Z"
    }
   },
   "id": "3f8de4fda975943c"
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "We detected that you are passing `past_key_values` as a tuple and this is deprecated and will be removed in v4.43. Please use an appropriate `Cache` class (https://huggingface.co/docs/transformers/v4.41.3/en/internal/generation_utils#transformers.Cache)\n"
     ]
    },
    {
     "data": {
      "text/html": "<div style='margin: 0px; padding: 0px; font-family: ColfaxAI, Arial; font-size: 20px;'<text style=color:black>Here</text><text style=color:black>'s</text><text style=color:black> a</text><text style=color:black> JSON</text><text style=color:black> object</text><text style=color:black> with</text><text style=color:black> only</text><text style=color:black> string</text><text style=color:black> values</text><text style=color:black>:</text><text style=color:black>{</text><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\n</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\"</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>name</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\":</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black> \"</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>John</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\",</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\n</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\"</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>age</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\":</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black> </text></span><text style=color:red>3</text><text style=color:red>0</text><text style=color:red>,</text><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:red>\n</text></span><span style='background-color: rgba(0, 0, 165, 0.25);'><text style=color:blue>\"</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>2</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>5</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\",</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\n</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\"</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>city</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\":</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black> \"</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>New</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black> York</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\",</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\n</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\"</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>address</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\":</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black> \"</text></span><span style='background-color: rgba(0, 0, 165, 0.25);'><text style=color:blue>1</text></span><span style='background-color: rgba(0, 0, 165, 0.25);'><text style=color:blue>2</text></span><span style='background-color: rgba(0, 0, 165, 0.25);'><text style=color:blue>3</text></span><span style='background-color: rgba(0, 0, 165, 0.25);'><text style=color:blue> Main</text></span><span style='background-color: rgba(0, 0, 165, 0.25);'><text style=color:orange> St</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black> St</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\",</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\n</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\"</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>phone</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\":</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black> \"</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>1</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>2</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>3</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>-</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>5</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>5</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>5</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>-</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>1</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>2</text></span><span style='background-color: rgba(0, 0, 165, 0.25);'><text style=color:orange>\"</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\"</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\n</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>}</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\n</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\n</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>```</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\n</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>\n</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black>I</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black> want</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black> to</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black> get</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black> the</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black> string</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black> values</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black> from</text></span><span style='background-color: rgba(0, 165, 0, 0.25);'><text style=color:black> the</text></span></div>"
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.6683290004730225\n",
      "14.720598941108289\n"
     ]
    }
   ],
   "source": [
    "res = gg.guide(\n",
    "    model=model,\n",
    "    tokenizer=tokenizer,\n",
    "    parser=gg.load_parser(open(\"../grammars/json_five_values_string_only.lark\").read()),\n",
    "    prompt=\"Here's a JSON object with only string values:\",\n",
    "    draft_model=guidance.models.Transformers(\n",
    "        model_name_or_path, echo=False\n",
    "    ),\n",
    "    max_new_tokens=19,\n",
    "    max_grammar_corrections=5,\n",
    "    verbose=True,\n",
    "    token_healing=True,\n",
    "    temperature=0.0\n",
    ")\n",
    "print(res.process_time_seconds)\n",
    "print(len(tokenizer(res.response)['input_ids']) / res.process_time_seconds)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-09-01T14:22:04.091606Z",
     "start_time": "2024-09-01T14:21:59.741062Z"
    }
   },
   "id": "56d2f6feb26acfbe"
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    \"name\": \"John\",\n",
      "    \"age\": \"25\",\n",
      "    \"city\": \"New York\",\n",
      "    \"address\": \"123 Main St\",\n",
      "    \"phone\": \"123-555-12\"\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    print(json.dumps(json.loads(res.response), indent = 4))\n",
    "except:\n",
    "    print(repr(res.response))"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-09-01T14:22:14.325917Z",
     "start_time": "2024-09-01T14:22:14.314592Z"
    }
   },
   "id": "8493d0291c629eeb"
  },
  {
   "cell_type": "markdown",
   "source": [
    "So what exactly happened above?\n",
    "\n",
    "1) The model generated a JSON object with 3 keys (name, age, and city)\n",
    "2) Our grammar catches that a numeric value is invalid (`\"age\": 25`) and backtracks the KV cache\n",
    "3) At this point - we only have 1 possible candidate from our grammar (`\"`) so we feed that in via a forward pass, and then continue generating\n",
    "3) We generate up to `...\"-7890\", \"`, and our grammar constrains our draft model to generate a valid continuation (`email\"`)\n",
    "4) Finally, the grammar return the final prefix that satisfies that Lark grammar conditions"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "f94d5100f1e3cf86"
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "grammar-guide",
   "language": "python",
   "display_name": "grammar-guide"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}