tutorial_ctrlg.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ctrl-G Tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Part A**. Ctrl-G on GPT2-large (less computation required)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Step 1. load pretrained models**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "device = 'cuda'\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0' # set your cuda device\n",
    "os.environ['TOKENIZERS_PARALLELISM'] = 'false'\n",
    "\n",
    "import torch\n",
    "import ctrlg\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, LogitsProcessorList\n",
    "\n",
    "# load the pretrained base_model and hmm_model; see README.md for a complete list of \n",
    "# released checkpoints. note that the hmm_model and base_model must share the same \n",
    "# vocabulary of tokens: i.e., one cannot apply hmm_gpt2-large_common-gen_4096 to \n",
    "# tulu2-7b_writing-prompts. To apply Ctrl-G to a custom base_model or to achieve \n",
    "# best performance on a specific domain, users would need to distill an hmm_model\n",
    "# from the base_model. Please refer to tutorial_distillation.ipynb for details.\n",
    "BASE_MODEL_PATH = f'ctrlg/gpt2-large_common-gen' # a gpt2-large checkpoint domain adapted to the common-gen corpus\n",
    "HMM_MODEL_PATH = f'ctrlg/hmm_gpt2-large_common-gen_4096' # alternatively 'ctrlg/hmm_gpt2-large_common-gen_32768' for better quality\n",
    "\n",
    "base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL_PATH).to(device)\n",
    "base_model.eval()\n",
    "tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_PATH)\n",
    "hmm_model = ctrlg.HMM.from_pretrained(HMM_MODEL_PATH).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Step 2. specify logical constraints as DFAs (example constraint 1)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab_size = hmm_model.vocab_size\n",
    "eos_token_id = hmm_model.eos_token_id\n",
    "\n",
    "\n",
    "##################################### prefix, suffix, prompt #####################################\n",
    "prefix = '' # generate text starting with nothing\n",
    "suffix = '.<|endoftext|>' # generate text ending with '<|endoftext|>'; a suffix must end with the eos token\n",
    "prompt = '<|endoftext|>' # prompt the base model with the '<|endoftext|>' token\n",
    "\n",
    "prefix_ids = tokenizer.encode(prefix)\n",
    "suffix_ids = tokenizer.encode(suffix)\n",
    "prompt_ids = tokenizer.encode(prompt)\n",
    "##################################### prefix, suffix, prompt #####################################\n",
    "\n",
    "\n",
    "##################################### DFA Construction #####################################\n",
    "# ac_builder constructs a DFA representing the constraint that (at least) \n",
    "# one the patterns must appear; a pattern is a sequence of token ids\n",
    "ac_builder = ctrlg.AhoCorasickBuilder(vocab_size)\n",
    "# word_count_builder constructs a DFA representing the constraint that \n",
    "# the generated text consists of a to b words; refer to the source code of\n",
    "# WordCountBuilder for the definition of a word.\n",
    "word_count_builder = ctrlg.WordCountBuilder(tokenizer, vocab_size)\n",
    "\n",
    "dfa_graphs = []\n",
    "\n",
    "# constraint 1:\n",
    "# one of ' riding a bike', ' ride bikes', ' rides a bike', ' biking', ' bikes' has to appear\n",
    "# AND one of ' park', ' beach' has to appear\n",
    "keyphrases = [[' riding a bike', ' ride bikes', ' rides a bike', ' biking', ' bikes'],\n",
    "            [' park', ' beach']]\n",
    "for keyphrase in keyphrases:\n",
    "    patterns = [tokenizer.encode(x) for x in keyphrase]\n",
    "    dfa_graphs.append(ac_builder.build(patterns))\n",
    "\n",
    "# constraint 2: generate exactly 10 words\n",
    "# word_count_builder constructs a DFA representing the constraint that \n",
    "# the generated text must contain a to b words\n",
    "a, b = 10, 10\n",
    "dfa_graphs.append(word_count_builder.build(a, b))\n",
    "\n",
    "# taking the intersection of the DFAs, i.e., \"logical and\" of the constraints.\n",
    "# This function also minimizes the constructed DFA, which is mainly CPU-based operations;\n",
    "# Due to its pure python implemenation, DFA minimization can be slow for complex constraints\n",
    "dfa_graph = ctrlg.DFA_prod(dfa_graphs, mode='intersection')\n",
    "\n",
    "# compile the dfa_graph for efficient GPU execution\n",
    "dfa_model = ctrlg.DFAModel(dfa_graph, vocab_size).to(device)\n",
    "##################################### DFA Construction #####################################\n",
    "\n",
    "\n",
    "##################################### token length #####################################\n",
    "# specify the min_new_tokens and max_new_tokens to be generated (excluding\n",
    "# the prefix and suffix) make sure that the numbers here would not conflict\n",
    "# with the given constraint: e.g. ask the model to generate 10 words with\n",
    "# max_new_tokens = 8\n",
    "min_new_tokens = 5\n",
    "max_new_tokens = 32\n",
    "##################################### token length #####################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Step 3. generate with constraints.**\n",
    "\n",
    "Due to the use of @torch.compile, the first run of the following functions could be significantly slower than the later runs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialze the constraints logits processor\n",
    "# Note: this part pre-computes & cache certain conditional probability tables;\n",
    "# one simple optimization is to re-use the same constraint_logits_processor for\n",
    "# base_model.generate if the constraints do not change.\n",
    "constraint_logits_processor = ctrlg.ConstraintLogitsProcessor(\n",
    "    hmm_model, dfa_model,\n",
    "    min_new_tokens, max_new_tokens,\n",
    "    prompt_ids, prefix_ids=prefix_ids, suffix_ids=suffix_ids)\n",
    "\n",
    "\n",
    "# set beam_size for beam search; usually the larger the beam_size the\n",
    "# higher the generation quality\n",
    "beam_size = 128\n",
    "\n",
    "# set the hmm_batch_size depending on the resource available;\n",
    "# uses more memory with larger hmm_batch_size but attains best speed \n",
    "# when it is set to beam_size\n",
    "constraint_logits_processor.hmm_batch_size = beam_size\n",
    "\n",
    "# generate with beam search\n",
    "input_ids = torch.tensor([prompt_ids], device=device)\n",
    "outputs = base_model.generate(\n",
    "        input_ids=input_ids, do_sample=False, length_penalty=0.2,\n",
    "        num_beams=beam_size, num_return_sequences=beam_size,\n",
    "        min_new_tokens=min_new_tokens, max_new_tokens=max_new_tokens,\n",
    "        logits_processor=LogitsProcessorList([constraint_logits_processor]),\n",
    "        pad_token_id=tokenizer.eos_token_id,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Step 4. extract & rank outputs via the base model.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract the generated ids; removing prompt ids; remove suffix ids that are (partially) generated\n",
    "generated_ids = ctrlg.extract_generated_ids(outputs.tolist(), prompt_ids, suffix_ids, eos_token_id)\n",
    "\n",
    "# rank the generated ids by the base_model probability\n",
    "generated_ids = ctrlg.rank_generated_ids(base_model, generated_ids, prompt_ids, suffix_ids, length_penalty=0.2)\n",
    "\n",
    "# print top 10 outputs\n",
    "for idx, generated in enumerate(generated_ids[:10]):\n",
    "    print(f'{idx}. ' + tokenizer.decode(prefix_ids, skip_special_tokens=True) + \\\n",
    "          '\\033[1m' + tokenizer.decode(generated, skip_special_tokens=True) + '\\033[0m' + \\\n",
    "          tokenizer.decode(suffix_ids, skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Step 5. try some other constraints! (example constraint 2)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab_size = hmm_model.vocab_size\n",
    "eos_token_id = hmm_model.eos_token_id\n",
    "\n",
    "\n",
    "prefix = ' on a fine sunny' # generate text starting with ' on a fine sunny'\n",
    "suffix = ' in the park.<|endoftext|>' # generate text ending with ' in the park.<|endoftext|>'\n",
    "prompt = '<|endoftext|> on a fine sunny' # prompt the base model with the '<|endoftext|>' token and the prefix\n",
    "\n",
    "prefix_ids = tokenizer.encode(prefix)\n",
    "suffix_ids = tokenizer.encode(suffix)\n",
    "prompt_ids = tokenizer.encode(prompt)\n",
    "\n",
    "\n",
    "ac_builder = ctrlg.AhoCorasickBuilder(vocab_size)\n",
    "word_count_builder = ctrlg.WordCountBuilder(tokenizer, vocab_size)\n",
    "\n",
    "dfa_graphs = []\n",
    "# constraint 1:\n",
    "# one of ' girl', ' boy', ' girls', ' boys', ' children' AND\n",
    "# one of ' dogs', ' cats', ' dog', ' cat' have to appear\n",
    "# in the GIVEN ORDER.\n",
    "keyphrases = [[' girl', ' boy', ' girls', ' boys', ' children'],\n",
    "            [' dogs', ' cats', ' dog', ' cat']]\n",
    "for keyphrase in keyphrases:\n",
    "    patterns = [tokenizer.encode(x) for x in keyphrase]\n",
    "    dfa_graphs.append(ac_builder.build(patterns))\n",
    "# concatenate the patterns so they appear in the given order\n",
    "dfa_graphs = [ctrlg.DFA_concatenate(dfa_graphs)]\n",
    "\n",
    "# constraint 2: generate 7 - 12 words\n",
    "a, b = 7, 12\n",
    "dfa_graphs.append(word_count_builder.build(a, b))\n",
    "\n",
    "dfa_graph = ctrlg.DFA_prod(dfa_graphs, mode='intersection')\n",
    "dfa_model = ctrlg.DFAModel(dfa_graph, vocab_size).to(device)\n",
    "\n",
    "\n",
    "min_new_tokens = 5\n",
    "max_new_tokens = 32\n",
    "\n",
    "\n",
    "# initialze the constraints logits processor\n",
    "constraint_logits_processor = ctrlg.ConstraintLogitsProcessor(\n",
    "    hmm_model, dfa_model,\n",
    "    min_new_tokens, max_new_tokens,\n",
    "    prompt_ids, prefix_ids=prefix_ids, suffix_ids=suffix_ids)\n",
    "\n",
    "\n",
    "beam_size = 128\n",
    "constraint_logits_processor.hmm_batch_size = beam_size\n",
    "input_ids = torch.tensor([prompt_ids], device=device)\n",
    "# generate with beam search\n",
    "outputs = base_model.generate(\n",
    "        input_ids=input_ids, do_sample=False,\n",
    "        num_beams=beam_size, num_return_sequences=beam_size,\n",
    "        min_new_tokens=min_new_tokens, max_new_tokens=max_new_tokens,\n",
    "        logits_processor=LogitsProcessorList([constraint_logits_processor]),\n",
    "        pad_token_id=tokenizer.eos_token_id,\n",
    "    )\n",
    "\n",
    "# extract the generated ids; removing prompt ids; remove suffix ids that are (partially) generated\n",
    "generated_ids = ctrlg.extract_generated_ids(outputs.tolist(), prompt_ids, suffix_ids, eos_token_id)\n",
    "\n",
    "# rank the generated ids by the base_model probability\n",
    "generated_ids = ctrlg.rank_generated_ids(base_model, generated_ids, prompt_ids, suffix_ids)\n",
    "\n",
    "# print top 10 outputs\n",
    "for idx, generated in enumerate(generated_ids[:10]):\n",
    "    print(f'{idx}. ' + tokenizer.decode(prefix_ids, skip_special_tokens=True) + \\\n",
    "          '\\033[1m' + tokenizer.decode(generated, skip_special_tokens=True) + '\\033[0m' + \\\n",
    "          tokenizer.decode(suffix_ids, skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Part B**. Ctrl-G on TULU2-7B (more computation required)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Step 1. load pretrained models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'  # set your cuda device\n",
    "os.environ['TOKENIZERS_PARALLELISM'] = 'false'\n",
    "\n",
    "import torch\n",
    "import ctrlg\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, LogitsProcessorList\n",
    "\n",
    "device = 'cuda'\n",
    "\n",
    "# load the pretrained base_model and hmm_model;\n",
    "BASE_MODEL_PATH = f'ctrlg/tulu2-7b_writing-prompts'\n",
    "HMM_MODEL_PATH = f'ctrlg/hmm_tulu2-7b_writing-prompts_32768'\n",
    "\n",
    "base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL_PATH).to(device)\n",
    "base_model.eval()\n",
    "base_model.half() # fp16 inference\n",
    "tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_PATH)\n",
    "hmm_model = ctrlg.HMM.from_pretrained(HMM_MODEL_PATH).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Step 2. specify logical constraints as DFAs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab_size = hmm_model.vocab_size\n",
    "eos_token_id = hmm_model.eos_token_id\n",
    "\n",
    "prefix = 'Once upon a time, in a land far, far away, there was a kingdom. The kingdom was'\n",
    "suffix = 'beautiful buildings. The people of this kingdom were known for their kindness and generosity, always ready to lend a helping hand.</s>'\n",
    "soft_constraint = ' in fairytale style' # use empty string for no soft constraint\n",
    "prompt = f'<|user|>\\nContinue the given text{soft_constraint}:\\n{prefix}\\n<|assistant|>\\n'\n",
    "\n",
    "prefix_ids = tokenizer.encode(prefix)[1:]\n",
    "suffix_ids = tokenizer.encode(suffix)[1:]\n",
    "prompt_ids = tokenizer.encode(prompt)\n",
    "\n",
    "ac_builder = ctrlg.AhoCorasickBuilder(vocab_size)\n",
    "eos_builder = ctrlg.EOSBuilder(vocab_size, eos_token_id)\n",
    "\n",
    "dfa_graphs = []\n",
    "keyphrases = [['towering'], ['reach the sky'], ['reflected'], ['lake']]\n",
    "for keyphrase in keyphrases:\n",
    "    patterns = [tokenizer.encode(x)[1:] for x in keyphrase]\n",
    "    dfa_graphs.append(ac_builder.build(patterns))\n",
    "dfa_graphs.append(eos_builder.build())\n",
    "\n",
    "dfa_graph = ctrlg.DFA_prod(dfa_graphs, mode='intersection')\n",
    "dfa_model = ctrlg.DFAModel(dfa_graph, vocab_size).to(device)\n",
    "\n",
    "min_new_tokens = 16\n",
    "max_new_tokens = 32"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Step 3. generate with constraints.\n",
    "\n",
    "Due to the use of @torch.compile, the first run of the following functions could be significantly slower than the later runs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialze the constraints logits processor\n",
    "constraint_logits_processor = ctrlg.ConstraintLogitsProcessor(\n",
    "    hmm_model, dfa_model,\n",
    "    min_new_tokens, max_new_tokens,\n",
    "    prompt_ids, prefix_ids=prefix_ids, suffix_ids=suffix_ids)\n",
    "\n",
    "\n",
    "# set the hmm_batch_size & temperature\n",
    "beam_size = 128 # sample 128 sequences\n",
    "temperature = 0.7\n",
    "constraint_logits_processor.hmm_batch_size = beam_size\n",
    "constraint_logits_processor.temperature = temperature\n",
    "\n",
    "\n",
    "# generate with sampling, temperature=0.7\n",
    "input_ids = torch.tensor([prompt_ids], device=device)\n",
    "outputs = base_model.generate(\n",
    "        input_ids=input_ids, do_sample=True,\n",
    "        num_return_sequences=beam_size, \n",
    "        min_new_tokens=min_new_tokens, max_new_tokens=max_new_tokens,\n",
    "        logits_processor=LogitsProcessorList([constraint_logits_processor]),\n",
    "        pad_token_id=tokenizer.eos_token_id,\n",
    "    )\n",
    "\n",
    "\n",
    "# extract the generated ids; removing prompt ids; remove suffix ids that are (partially) generated\n",
    "generated_ids = ctrlg.extract_generated_ids(outputs.tolist(), prompt_ids, suffix_ids, eos_token_id)\n",
    "\n",
    "# filter 75% of the generated ids by how well they connect with the suffix\n",
    "generated_ids = ctrlg.rank_generated_ids(base_model, generated_ids, prompt_ids, suffix_ids,\n",
    "                                            suffix_logits_only=True, suffix_length_cap=5)[:32]\n",
    "# rank the generated ids by the base_model for higher quality\n",
    "generated_ids = ctrlg.rank_generated_ids(base_model, generated_ids, prompt_ids, suffix_ids)\n",
    "\n",
    "# print top 10 outputs\n",
    "for idx, generated in enumerate(generated_ids[:10]):\n",
    "    print(f'{idx}. ' + tokenizer.decode(prefix_ids, skip_special_tokens=True) + \\\n",
    "          ' ' + '\\033[1m' + tokenizer.decode(generated, skip_special_tokens=True) + '\\033[0m' + ' ' + \\\n",
    "          tokenizer.decode(suffix_ids, skip_special_tokens=True))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}