examples/mistral-7b-dcd_dropout_with_quantize.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig\n",
    "from dcd import dcd_pipeline_registry, set_stop_words\n",
    "from dcd import create_prompt, create_prompt_student\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dcd_pipeline_registry()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ee0b4a3ee1374771a28415654b571295",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(\"mistralai/Mistral-7B-v0.1\", device_map=\"auto\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"mistralai/Mistral-7B-v0.1\", device_map=\"auto\")\n",
    "\n",
    "model_quantize = AutoModelForCausalLM.from_pretrained(\"TheBloke/Llama-2-7B-GPTQ\", device_map=\"auto\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Added stop word:  Q: with the ids [28824, 28747]\n",
      "Added stop word:  \\end{code} with the ids [28756, 416, 28751, 1409, 28752]\n",
      "Added stop word:  </s> with the ids [2]\n",
      "Added stop word:  Wrong explanation: with the ids [17055, 566, 13268, 28747]\n"
     ]
    }
   ],
   "source": [
    "beam_size = 1\n",
    "max_length = 250\n",
    "\n",
    "alpha_coef = 0.1\n",
    "beta_coef = 0.8\n",
    "dropout_rate = 0.2\n",
    "\n",
    "type_prompt = 4  # The synthetic demonstration prompt for arithmetic problems\n",
    "\n",
    "stopping_criteria = set_stop_words(tokenizer=tokenizer, stop_words=[\"Q:\", \"\\end{code}\", \"</s>\", \"Wrong explanation:\"])\n",
    "\n",
    "generation_config = GenerationConfig(\n",
    "    do_sample=False,\n",
    "    num_beams=beam_size,\n",
    "    pad_token_id=0,\n",
    "    eos_token_id=0,\n",
    ")\n",
    "\n",
    "class Args:\n",
    "    def __init__(self) -> None:\n",
    "        self.prompt_file = 'gsm8k'\n",
    "        self.data_name = \"gsm8k\"\n",
    "        self.cot_flag = True\n",
    "        self.direct_answer_trigger_for_fewshot = 'The answer is'\n",
    "\n",
    "args_prompt = Args()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Greedy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Toulouse has twice as many sheep as Charleston. Charleston has 4 times as many sheep as Seattle. How many sheep do Toulouse, Charleston, and Seattle have together if Seattle has 20 sheep?\"\n",
    "question_formated = \"Q: \" + question + \"\\n\" + \"A:\"\n",
    "inputs = tokenizer(create_prompt(args_prompt, data_name=args_prompt.data_name) + question_formated, return_tensors=\"pt\")\n",
    "input_ids = inputs[\"input_ids\"].to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs_args_greedy = dict(\n",
    "    generation_config=generation_config,\n",
    "    return_dict_in_generate=True,\n",
    "    output_scores=True,\n",
    "    max_new_tokens=max_length,\n",
    "    stopping_criteria=stopping_criteria,\n",
    "    min_tokens_to_keep=2 if beam_size > 1 else 1,\n",
    "    dropout_rate=dropout_rate\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output of greedy: Toulouse has twice as many sheep as Charleston. So if Charleston has 4 times as many sheep as Seattle, then Toulouse has 4 times as many sheep as Seattle. So Toulouse, Charleston, and Seattle have 20 + 4 + 4 = 28 sheep. The answer is 28.\n"
     ]
    }
   ],
   "source": [
    "output_sequences = model.generate(\n",
    "    input_ids=input_ids,\n",
    "    **inputs_args_greedy)\n",
    "\n",
    "s_greedy = output_sequences.sequences[0]\n",
    "output_greedy = tokenizer.decode(s_greedy, skip_special_tokens=True)\n",
    "\n",
    "output_formated_greedy = output_greedy.split(\"A: \")[-1].replace(\"\\n\\nQ:\", \"\")\n",
    "print(f\"Output of greedy: {output_formated_greedy}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Distillation Contrastive Decoding with Dropout + Quantize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Toulouse has twice as many sheep as Charleston. Charleston has 4 times as many sheep as Seattle. How many sheep do Toulouse, Charleston, and Seattle have together if Seattle has 20 sheep?\"\n",
    "question_formated = \"Q: \" + question + \"\\n\" + \"A:\"\n",
    "inputs = tokenizer(create_prompt(args_prompt, data_name=args_prompt.data_name) + question_formated, return_tensors=\"pt\")\n",
    "input_ids = inputs[\"input_ids\"].to(device)\n",
    "\n",
    "inputs_student = tokenizer(create_prompt_student(args_prompt, type=type_prompt, data_name=args_prompt.data_name) + question_formated, return_tensors=\"pt\")\n",
    "input_ids_student = inputs_student[\"input_ids\"].to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs_args_greedy = dict(\n",
    "    generation_config=generation_config,\n",
    "    return_dict_in_generate=True,\n",
    "    output_scores=True,\n",
    "    max_new_tokens=max_length,\n",
    "    stopping_criteria=stopping_criteria,\n",
    "\n",
    "    # DCD parameters of dropout and quantize\n",
    "    student_lm=model_quantize,\n",
    "    alpha_coef=alpha_coef,\n",
    "    beta_coef=beta_coef,\n",
    "    min_tokens_to_keep=2 if beam_size > 1 else 1,\n",
    "    teacher_student=True,\n",
    "    dropout_rate=dropout_rate\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output of DCD: Toulouse has twice as many sheep as Charleston. If Charleston has 4 times as many as Seattle, then Toulouse has 2 * 4 = 8 times as many as Seattle. 8 times as many as Seattle is 8 x 20 = 160. The answer is 160.\n"
     ]
    }
   ],
   "source": [
    "output_sequences = model.generate(\n",
    "    input_ids=input_ids,\n",
    "    input_ids_student=input_ids_student,\n",
    "    **inputs_args_greedy)\n",
    "\n",
    "s_dcd = output_sequences.sequences[0]\n",
    "output_dcd = tokenizer.decode(s_dcd, skip_special_tokens=True)\n",
    "output_formated_dcd = output_dcd.split(\"A: \")[-1].replace(\"\\n\\nQ:\", \"\")\n",
    "print(f\"Output of DCD: {output_formated_dcd}\")"
   ]
  }
 ],
 "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}