Chapter4.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chapter 4: Formatting Your Dataset\n",
    "\n",
    "### Spoilers \n",
    "\n",
    "In this chapter, we will:\n",
    "\n",
    "- Understand the importance of defining a proper chat template\n",
    "- Discuss several formatting alternatives, including custom formatting functions and templates\n",
    "- Configure the tokenizer and the model’s embedding layer\n",
    "- Explore packed datasets and different data collators for loading data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kODUm5BmEQhI",
    "outputId": "5a17ca8b-855a-4e55-b7a2-ab2af98d8906"
   },
   "outputs": [],
   "source": [
    "# If you're running on Colab\n",
    "!pip install datasets bitsandbytes trl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If you're running on runpod.io's Jupyter Template\n",
    "#!pip install datasets bitsandbytes trl transformers peft huggingface-hub accelerate safetensors pandas matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from peft import prepare_model_for_kbit_training, get_peft_model, LoraConfig\n",
    "from datasets import load_dataset, Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, DataCollatorForLanguageModeling, DataCollatorWithPadding, DataCollatorWithFlattening, BitsAndBytesConfig\n",
    "from trl import setup_chat_format, DataCollatorForCompletionOnlyLM\n",
    "from trl.extras.dataset_formatting import FORMAT_MAPPING, instructions_formatting_function, conversations_formatting_function\n",
    "from trl.trainer import ConstantLengthDataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Goal\n",
    "\n",
    "We format the dataset to provide structure and cues to the LLM. We can easily steer its behavior (e.g., instruction-tuning) by carefully wrapping each component—the user’s prompt and the model’s completion—with appropriate tags and special tokens."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Formatting in a Nutshell\n",
    "\n",
    "![](https://github.com/dvgodoy/FineTuningLLMs/blob/main/images/ch4/base_prompt.png?raw=True)\n",
    "<center>Figure 4.1 - Base model’s next token prediction</center>\n",
    "\n",
    "![](https://github.com/dvgodoy/FineTuningLLMs/blob/main/images/ch4/fine_tuned_prompt.png?raw=True)\n",
    "<center>Figure 4.2 - Fine-tuned model triggered by response template</center>\n",
    "\n",
    "![](https://github.com/dvgodoy/FineTuningLLMs/blob/main/images/ch4/chat_prompt_new.png?raw=True)\n",
    "<center>Figure 4.3 - Chat model using chat template</center>\n",
    "\n",
    "![](https://github.com/dvgodoy/FineTuningLLMs/blob/main/images/ch4/chat_example_new.png?raw=True)\n",
    "<center>Figure 4.4 - General structure of a chat template</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Road so Far"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return self.fget.__get__(instance, owner)()\n"
     ]
    }
   ],
   "source": [
    "supported = torch.cuda.is_bf16_supported(including_emulation=False)\n",
    "compute_dtype = (torch.bfloat16 if supported else torch.float32)\n",
    "\n",
    "nf4_config = BitsAndBytesConfig(\n",
    "   load_in_4bit=True,\n",
    "   bnb_4bit_quant_type=\"nf4\",\n",
    "   bnb_4bit_use_double_quant=True,\n",
    "   bnb_4bit_compute_dtype=compute_dtype\n",
    ")\n",
    "\n",
    "model_q4 = AutoModelForCausalLM.from_pretrained(\"facebook/opt-350m\",\n",
    "                                                device_map='cuda:0',\n",
    "                                                torch_dtype=compute_dtype,\n",
    "                                                quantization_config=nf4_config)\n",
    "\n",
    "model_q4 = prepare_model_for_kbit_training(model_q4)\n",
    "\n",
    "config = LoraConfig(\n",
    "    r=16,\n",
    "    lora_alpha=32,\n",
    "    lora_dropout=0.05,\n",
    "    bias=\"none\",\n",
    "    task_type=\"CAUSAL_LM\",\n",
    ")\n",
    "peft_model = get_peft_model(model_q4, config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Applying Templates\n",
    "\n",
    "****\n",
    "**Summary of \"Applying Templates\"**\n",
    "\n",
    "You have three options for formatting your dataset:\n",
    "1. Your dataset is in one of the **two formats supported by the `STTrainer` class** (conversational or instruction):\n",
    "   - Your **tokenizer must have a chat template** configured.\n",
    "   - No need to define a formatting function or format the dataset before training.\n",
    "2. You want to use a **custom formatting function** (see \"BYOFF, Bring Your Own Formatting Function\"):\n",
    "   - The custom function should be provided as the **`formatting_func` argument of the `SFTTrainer` class** (see Chapter 5).\n",
    "   - Your formatting function **must handle batches of data**.\n",
    "     - Test it by calling the dataset's `map()` method with `batched=True`.\n",
    "    - No need to apply the function to the dataset before training.\n",
    "    - If your tokenizer already **has a chat template**:\n",
    "      - You may call its `apply_chat_template()` method in your function.\n",
    "      - Stick to the template's general format (instruction and response templates).\n",
    "      - If the template doesn’t include one, **you may append an `EOS` token to the end of the formatted output**.\n",
    "   - If your tokenizer **does not have a chat template**:\n",
    "     - You're free to define the general format, including instruction and response templates (see \"Advanced—BYOT, Bring Your Own Template\")\n",
    "3. Your dataset is **already formatted** (see \"BYOFD, Bring Your Own Formatted Data\"):\n",
    "   - The column containing the formatted data should be provided as the **`dataset_text_field` argument of the `SFTTrainer` class** (see Chapter 5).\n",
    "   - Even though you can use your own formatting function to preprocess your dataset, it won't be used by the trainer class.\n",
    "   - Ensure your **data is compatible with the tokenizer's template**.\n",
    "****"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{% for message in messages %}{% if message['role'] == 'system' %}{{'<|system|>\n",
      "' + message['content'] + '<|end|>\n",
      "'}}{% elif message['role'] == 'user' %}{{'<|user|>\n",
      "' + message['content'] + '<|end|>\n",
      "'}}{% elif message['role'] == 'assistant' %}{{'<|assistant|>\n",
      "' + message['content'] + '<|end|>\n",
      "'}}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|assistant|>\n",
      "' }}{% else %}{{ eos_token }}{% endif %}\n"
     ]
    }
   ],
   "source": [
    "tokenizer_phi = AutoTokenizer.from_pretrained(\"microsoft/phi-3-mini-4k-instruct\")\n",
    "print(tokenizer_phi.chat_template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|system|>\n",
      "You are a helpful AI assistant.<|end|>\n",
      "<|user|>\n",
      "What is the capital of Argentina?<|end|>\n",
      "<|assistant|>\n",
      "Buenos Aires.<|end|>\n",
      "<|endoftext|>\n"
     ]
    }
   ],
   "source": [
    "messages = [\n",
    "    {'role': 'system', 'content': 'You are a helpful AI assistant.'},\n",
    "    {'role': 'user', 'content': 'What is the capital of Argentina?'},\n",
    "    {'role': 'assistant', 'content': 'Buenos Aires.'}\n",
    "]\n",
    "\n",
    "formatted = tokenizer_phi.apply_chat_template(conversation=messages, \n",
    "                                          tokenize=False, \n",
    "                                          add_generation_prompt=False)\n",
    "print(formatted)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|system|>\n",
      "You are a helpful AI assistant.<|end|>\n",
      "<|user|>\n",
      "What is the capital of Argentina?<|end|>\n",
      "<|assistant|>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "inference_input = tokenizer_phi.apply_chat_template(conversation=messages[:-1], \n",
    "                                          tokenize=False, \n",
    "                                          add_generation_prompt=True)\n",
    "print(inference_input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Supported Formats\n",
    "\n",
    "##### Conversational"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [{'content': Value(dtype='string', id=None),\n",
       "   'role': Value(dtype='string', id=None)}]}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation_ds = Dataset.from_list([{'messages': messages}])\n",
    "conversation_ds.features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "FORMAT_MAPPING['chatml'] == conversation_ds.features['messages']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|system|>\n",
      "You are a helpful AI assistant.<|end|>\n",
      "<|user|>\n",
      "What is the capital of Argentina?<|end|>\n",
      "<|assistant|>\n",
      "Buenos Aires.<|end|>\n",
      "<|endoftext|>\n"
     ]
    }
   ],
   "source": [
    "formatting_func = conversations_formatting_function(tokenizer_phi, messages_field='messages')\n",
    "\n",
    "print(formatting_func(conversation_ds[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "# formatting function for conversational format\n",
    "def format_dataset(examples):\n",
    "    if isinstance(examples[messages_field][0], list):\n",
    "        output_texts = []\n",
    "        for i in range(len(examples[messages_field])):\n",
    "            output_texts.append(tokenizer.apply_chat_template(examples[messages_field][i], tokenize=False))\n",
    "        return output_texts\n",
    "    else:\n",
    "        return tokenizer.apply_chat_template(examples[messages_field], tokenize=False)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Instruction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'prompt': Value(dtype='string', id=None),\n",
       " 'completion': Value(dtype='string', id=None)}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "instructions = [{'prompt': 'What is the capital of Argentina?',\n",
    "                 'completion': 'Buenos Aires.'}]\n",
    "\n",
    "instruction_ds = Dataset.from_list(instructions)\n",
    "instruction_ds.features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "FORMAT_MAPPING['instruction'] == instruction_ds.features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function trl.extras.dataset_formatting.instructions_formatting_function.<locals>.format_dataset(examples)>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "formatting_func = instructions_formatting_function(tokenizer_phi)\n",
    "formatting_func"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "# formatting function for instruction format\n",
    "def format_dataset(examples):\n",
    "    if isinstance(examples[\"prompt\"], list):\n",
    "        output_texts = []\n",
    "        for i in range(len(examples[\"prompt\"])):\n",
    "            converted_sample = [\n",
    "                {\"role\": \"user\", \"content\": examples[\"prompt\"][i]},\n",
    "                {\"role\": \"assistant\", \"content\": examples[\"completion\"][i]},\n",
    "            ]\n",
    "            output_texts.append(tokenizer.apply_chat_template(converted_sample, tokenize=False))\n",
    "        return output_texts\n",
    "    else:\n",
    "        converted_sample = [\n",
    "            {\"role\": \"user\", \"content\": examples[\"prompt\"]},\n",
    "            {\"role\": \"assistant\", \"content\": examples[\"completion\"]},\n",
    "        ]\n",
    "        return tokenizer.apply_chat_template(converted_sample, tokenize=False)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_prompts_completions = {\n",
    "    'prompt': ['What is the capital of Argentina?',\n",
    "               'What is the capital of the United States?'],\n",
    "    'completion': ['Buenos Aires.',\n",
    "                    'Washington D.C.']\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_messages = [\n",
    "    [{'role': 'user', 'content': 'What is the capital of Argentina?'},\n",
    "     {'role': 'assistant', 'content': 'Buenos Aires.'}],\n",
    "    [{'role': 'user', 'content': 'What is the capital of the United States?'},\n",
    "     {'role': 'assistant', 'content': 'Washington D.C.'}]\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### BYOFF (Bring Your Own Formatting Function)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def byo_formatting_func1(examples):\n",
    "    messages = examples[\"messages\"]\n",
    "    output_texts = tokenizer_phi.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)\n",
    "    return output_texts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c8a2e61ba5d64fffa8a48341fa0ed8cd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['messages', 'input_ids', 'attention_mask'],\n",
       "    num_rows: 2\n",
       "})"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds_msg = Dataset.from_dict({'messages': batch_messages})\n",
    "ds_msg.map(lambda v: tokenizer_phi(byo_formatting_func1(v)), batched=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def byo_formatting_func2(examples):\n",
    "    response_template = '### Answer:'\n",
    "    text = f\"### Question: {examples['prompt']}\\n{response_template} {examples['completion']}\"\n",
    "    text += tokenizer_phi.eos_token\n",
    "    return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "### Question: What is the capital of Argentina?\n",
      "### Answer: Buenos Aires.<|endoftext|>\n"
     ]
    }
   ],
   "source": [
    "ds_prompt = Dataset.from_dict(batch_prompts_completions)\n",
    "print(byo_formatting_func2(ds_prompt[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dce4ebcfd66c4aefb48ed1271cde2743",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "ArrowInvalid",
     "evalue": "Column 2 named input_ids expected length 2 but got length 44",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mArrowInvalid\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;31mArrowInvalid\u001b[0m: Column 2 named input_ids expected length 2 but got length 44"
     ]
    }
   ],
   "source": [
    "# this is going to raise an exception\n",
    "ds_prompt.map(lambda v: tokenizer_phi(byo_formatting_func2(v)), batched=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def byo_formatting_func3(examples):\n",
    "    output_texts = []\n",
    "    response_template = '### Answer:'\n",
    "    for i in range(len(examples['prompt'])):\n",
    "        text = f\"### Question: {examples['prompt'][i]}\\n {response_template} {examples['completion'][i]}\"\n",
    "        text += tokenizer_phi.eos_token\n",
    "        output_texts.append(text)\n",
    "    return output_texts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ea47244bb5e8428997404671ed088420",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['prompt', 'completion', 'input_ids', 'attention_mask'],\n",
       "    num_rows: 2\n",
       "})"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds_prompt.map(lambda v: tokenizer_phi(byo_formatting_func3(v)), batched=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### BYOFD (Bring Your Own Formatted Data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def byofd_formatting_func(examples):\n",
    "    messages = examples[\"messages\"]\n",
    "    output_texts = tokenizer_phi.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)\n",
    "    return {'text': output_texts}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e799118503a2433f9bbb3258c5e68709",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "['<|user|>\\nWhat is the capital of Argentina?<|end|>\\n<|assistant|>\\nBuenos Aires.<|end|>\\n<|endoftext|>',\n",
       " '<|user|>\\nWhat is the capital of the United States?<|end|>\\n<|assistant|>\\nWashington D.C.<|end|>\\n<|endoftext|>']"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "formatted_ds = ds_msg.map(byofd_formatting_func, batched=True)\n",
    "formatted_ds['text']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Showdown\n",
    "\n",
    "![](https://github.com/dvgodoy/FineTuningLLMs/blob/main/images/ch4/formatting_flow.png?raw=True)\n",
    "\n",
    "<center>Figure 4.5 - Choosing the right configuration for your formatting needs</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Tokenizer\n",
    "\n",
    "****\n",
    "**Summary of \"The Tokenizer\"**\n",
    "- The **tokenizer's vocabulary** is usually **shorter than the model's embedding layer**.\n",
    "  - The difference in size consists of, quite literally, \"empty slots\" that you can use to **create new tokens without resizing** the embedding layer.\n",
    "  - The **size of the embedding layer** is often a **multiple of a power of two** (32, 64, etc.) to optimize **memory allocation**.\n",
    "- The `EOS` token should be **used solely to mark the end of the text** and nothing else.\n",
    "  - Using the `EOS` token for padding may lead to _endless token generation_.\n",
    "- The `PAD` token is often undefined, but you might still need it:\n",
    "  - **DO NOT** assign the `EOS` token as the `PAD` token.\n",
    "  - If the `UNK` token is defined, it is fine to assign it as the `PAD` token.\n",
    "  - If the `UNK` token is undefined, create a new special token as the `PAD` token.\n",
    "  - **WATCH OUT**: If the `PAD` token is left **undefined**, many libraries will **default to assigning it the `EOS` token** instead!\n",
    "- For **generative** models, **padding** should be performed on the **left** side.\n",
    "  - Padding on the _right_ side will train the model to generate _endless sequences of padding tokens_.\n",
    "  - Many tutorials use `tokenizer.padding_side='right'` due to reported overflow issues with the `SFTTrainer` class.\n",
    "    - This is fine **only if you're using packing or packing-like collators** (see the \"Packed Dataset\" section) instead of standard padding.\n",
    "- If you **create new special tokens**, in theory, you should also **fine-tune the embedding layer** (since you're using those \"empty slots\").\n",
    "  - In practice, your model _may_ still work if you **keep the embeddings frozen**.\n",
    "  - Even though the new tokens' representation is _random_ (their embeddings aren't trained), the other trainable parts of the model may still learn to use them \"as is.\"\n",
    "****"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer_phi = AutoTokenizer.from_pretrained(\"microsoft/phi-3-mini-4k-instruct\")\n",
    "config_phi = AutoConfig.from_pretrained(\"microsoft/phi-3-mini-4k-instruct\", trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': [2803, 29915, 29879, 5993, 675, 445, 10541, 29991], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1]}"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer_phi(\"Let's tokenize this sentence!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Vocabulary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(32011, 32064)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(tokenizer_phi), config_phi.vocab_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('<|user|>', 32010),\n",
       " ('<|placeholder6|>', 32009),\n",
       " ('<|placeholder5|>', 32008),\n",
       " ('<|end|>', 32007),\n",
       " ('<|system|>', 32006),\n",
       " ('<|placeholder4|>', 32005),\n",
       " ('<|placeholder3|>', 32004),\n",
       " ('<|placeholder2|>', 32003),\n",
       " ('<|placeholder1|>', 32002),\n",
       " ('<|assistant|>', 32001),\n",
       " ('<|endoftext|>', 32000),\n",
       " ('给', 31999)]"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sorted(tokenizer_phi.vocab.items(), key=lambda t: -t[1])[:11]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('<|endoftext|>', 32000)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer_phi.eos_token, tokenizer_phi.eos_token_id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### The Tokenizer 7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['<s>', '<|endoftext|>', '<unk>']"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer_phi.all_special_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'bos_token': '<s>',\n",
       " 'eos_token': '<|endoftext|>',\n",
       " 'unk_token': '<unk>',\n",
       " 'pad_token': '<|endoftext|>'}"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer_phi.special_tokens_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(None, None, None)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer_phi.cls_token, tokenizer_phi.sep_token, tokenizer_phi.mask_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'bos_token': '<s>',\n",
       " 'eos_token': '<|endoftext|>',\n",
       " 'unk_token': '<unk>',\n",
       " 'sep_token': '<sep>',\n",
       " 'pad_token': '<unk>',\n",
       " 'cls_token': '<cls>',\n",
       " 'mask_token': '<mask>'}"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer_phi.add_special_tokens({'cls_token': '<cls>', 'sep_token': '<sep>', 'mask_token': '<mask>'})\n",
    "tokenizer_phi.special_tokens_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('<mask>', 32013),\n",
       " ('<sep>', 32012),\n",
       " ('<cls>', 32011),\n",
       " ('<|user|>', 32010),\n",
       " ('<|placeholder6|>', 32009),\n",
       " ('<|placeholder5|>', 32008),\n",
       " ('<|end|>', 32007),\n",
       " ('<|system|>', 32006),\n",
       " ('<|placeholder4|>', 32005),\n",
       " ('<|placeholder3|>', 32004),\n",
       " ('<|placeholder2|>', 32003),\n",
       " ('<|placeholder1|>', 32002),\n",
       " ('<|assistant|>', 32001),\n",
       " ('<|endoftext|>', 32000),\n",
       " ('给', 31999)]"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sorted(tokenizer_phi.vocab.items(), key=lambda t: -t[1])[:13]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### The `EOS` Token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'bos_token': '<s>',\n",
       " 'eos_token': '<|endoftext|>',\n",
       " 'unk_token': '<unk>',\n",
       " 'pad_token': '<unk>'}"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer_phi.pad_token = tokenizer_phi.unk_token\n",
    "tokenizer_phi.pad_token_id = tokenizer_phi.unk_token_id\n",
    "\n",
    "tokenizer_phi.special_tokens_map"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "# Updating model's configuration for the modified PAD token\n",
    "if getattr(model, \"config\", None) is not None:\n",
    "    model.config.pad_token_id = tokenizer_phi.pad_token_id\n",
    "if (getattr(model, \"generation_config\", None) s not None):\n",
    "    model.config.pad_token_id = tokenizer_phi.pad_token_id\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### The `PAD` Token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('<unk>', 'left')"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer_phi.pad_token, tokenizer_phi.padding_side"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Collators\n",
    "\n",
    "****\n",
    "**Summary of \"Data Collators\"**\n",
    "- You can specify the `data_collator` argument in the `SFTTrainer` class (see Chapter 5).\n",
    "- `DataCollatorForLanguageModeling` is the **default** collator for the `SFTTrainer` class:\n",
    "  - It automatically **replicates the token IDs as labels**.\n",
    "  - It **doesn't shift the labels**, as this is **handled automatically by the model**.\n",
    "  - It includes the full text (both prompt and completion) as labels, making it ideal for instruction-tuning.\n",
    "- If you're further fine-tuning an instruction or chat model, you can use `DataCollatorForCompletionOnlyLM` to **train only on the model's answer (completion)**.\n",
    "  - It also replicates the token IDs as labels but **masks the prompt tokens by replacing their IDs with `-100`**.\n",
    "  - In a **single interaction** (one prompt and one completion), the **response template is enough** to locate the completion.\n",
    "  - In **multiple interactions** (a sequence of prompts and completions), both the **instruction and response templates** are needed to correctly identify and mask the prompt tokens.\n",
    "****"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(720,\n",
       " {'prompt': 'The birch canoe slid on the smooth planks.',\n",
       "  'completion': 'On the smooth planks, the birch canoe slid. Yes, hrrrm.'})"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = load_dataset(\"dvgodoy/yoda_sentences\", split=\"train\")\n",
    "dataset = dataset.rename_column(\"sentence\", \"prompt\")\n",
    "dataset = dataset.rename_column(\"translation_extra\", \"completion\")\n",
    "dataset = dataset.remove_columns([\"translation\"])\n",
    "len(dataset), dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['<|user|>\\nThe birch canoe slid on the smooth planks.<|end|>\\n<|assistant|>\\nOn the smooth planks, the birch canoe slid. Yes, hrrrm.<|end|>\\n<|endoftext|>', '<|user|>\\nGlue the sheet to the dark blue background.<|end|>\\n<|assistant|>\\nGlue the sheet to the dark blue background, you must.<|end|>\\n<|endoftext|>']\n"
     ]
    }
   ],
   "source": [
    "formatting_func = instructions_formatting_function(tokenizer_phi)\n",
    "dataset = dataset.map(lambda row: {'text': formatting_func(row)}, batched=True, batch_size=32)\n",
    "sequences = dataset['text']\n",
    "print(sequences[:2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_dataset = dataset.map(lambda row: tokenizer_phi(row['text']))\n",
    "tokenized_dataset = tokenized_dataset.select_columns(['input_ids'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### `DataCollatorWithPadding`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[32010,   450, 29773,   305,   508,  7297,  2243,   333,   373,   278,\n",
       "         10597,   715,  1331, 29889, 32007, 32001,  1551,   278, 10597,   715,\n",
       "          1331, 29892,   278, 29773,   305,   508,  7297,  2243,   333, 29889,\n",
       "          3869, 29892,   298, 21478,  1758, 29889, 32007, 32000],\n",
       "        [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "         32010,  8467,   434,   278,  9869,   304,   278,  6501,  7254,  3239,\n",
       "         29889, 32007, 32001,  8467,   434,   278,  9869,   304,   278,  6501,\n",
       "          7254,  3239, 29892,   366,  1818, 29889, 32007, 32000]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pad_collator = DataCollatorWithPadding(tokenizer_phi)\n",
    "pad_dloader = DataLoader(tokenized_dataset, batch_size=2, collate_fn=pad_collator)\n",
    "pad_batch = next(iter(pad_dloader))\n",
    "pad_batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Dude, Where's My Label?\n",
    "\n",
    "![](https://github.com/dvgodoy/FineTuningLLMs/blob/main/images/ch4/shift_labels.png?raw=True)\n",
    "\n",
    "<center>Figure 4.6 - Inputs and their corresponding shifted labels</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### `DataCollatorForLanguageModeling`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[32010,   450, 29773,   305,   508,  7297,  2243,   333,   373,   278,\n",
       "         10597,   715,  1331, 29889, 32007, 32001,  1551,   278, 10597,   715,\n",
       "          1331, 29892,   278, 29773,   305,   508,  7297,  2243,   333, 29889,\n",
       "          3869, 29892,   298, 21478,  1758, 29889, 32007, 32000],\n",
       "        [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "         32010,  8467,   434,   278,  9869,   304,   278,  6501,  7254,  3239,\n",
       "         29889, 32007, 32001,  8467,   434,   278,  9869,   304,   278,  6501,\n",
       "          7254,  3239, 29892,   366,  1818, 29889, 32007, 32000]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]), 'labels': tensor([[32010,   450, 29773,   305,   508,  7297,  2243,   333,   373,   278,\n",
       "         10597,   715,  1331, 29889, 32007, 32001,  1551,   278, 10597,   715,\n",
       "          1331, 29892,   278, 29773,   305,   508,  7297,  2243,   333, 29889,\n",
       "          3869, 29892,   298, 21478,  1758, 29889, 32007, 32000],\n",
       "        [ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "         32010,  8467,   434,   278,  9869,   304,   278,  6501,  7254,  3239,\n",
       "         29889, 32007, 32001,  8467,   434,   278,  9869,   304,   278,  6501,\n",
       "          7254,  3239, 29892,   366,  1818, 29889, 32007, 32000]])}"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lm_collator = DataCollatorForLanguageModeling(tokenizer_phi, mlm=False)\n",
    "lm_dloader = DataLoader(tokenized_dataset, batch_size=2, collate_fn=lm_collator)\n",
    "lm_batch = next(iter(lm_dloader))\n",
    "lm_batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### `DataCollatorForCompletionOnlyLM`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[32010,   450, 29773,   305,   508,  7297,  2243,   333,   373,   278,\n",
       "         10597,   715,  1331, 29889, 32007, 32001,  1551,   278, 10597,   715,\n",
       "          1331, 29892,   278, 29773,   305,   508,  7297,  2243,   333, 29889,\n",
       "          3869, 29892,   298, 21478,  1758, 29889, 32007, 32000],\n",
       "        [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "         32010,  8467,   434,   278,  9869,   304,   278,  6501,  7254,  3239,\n",
       "         29889, 32007, 32001,  8467,   434,   278,  9869,   304,   278,  6501,\n",
       "          7254,  3239, 29892,   366,  1818, 29889, 32007, 32000]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]), 'labels': tensor([[ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  1551,   278, 10597,   715,\n",
       "          1331, 29892,   278, 29773,   305,   508,  7297,  2243,   333, 29889,\n",
       "          3869, 29892,   298, 21478,  1758, 29889, 32007, 32000],\n",
       "        [ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  8467,   434,   278,  9869,   304,   278,  6501,\n",
       "          7254,  3239, 29892,   366,  1818, 29889, 32007, 32000]])}"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response_template = '<|assistant|>' # token id 32001\n",
    "completion_collator = DataCollatorForCompletionOnlyLM(response_template=response_template, \n",
    "                                                      tokenizer=tokenizer_phi)\n",
    "completion_dloader = DataLoader(tokenized_dataset, batch_size=2, collate_fn=completion_collator)\n",
    "completion_batch = next(iter(completion_dloader))\n",
    "completion_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'On the smooth planks, the birch canoe slid. Yes, hrrrm.<|end|><|endoftext|>'"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels = completion_batch['labels'][0]\n",
    "valid_tokens = (labels >= 0)\n",
    "tokenizer_phi.decode(labels[valid_tokens])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Multiple Interactions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5d2153ac43d24bf185cad336bb9c394b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dummy_chat = \"\"\"<|user|>Hello\n",
    "<|assistant|>How are you?\n",
    "<|user|>I'm fine! You?\n",
    "<|assistant|>I'm fine too!\n",
    "<|endoftext|>\"\"\"\n",
    "\n",
    "dummy_ds = Dataset.from_dict({'text': [dummy_chat]})\n",
    "dummy_ds = dummy_ds.map(lambda row: tokenizer_phi(row['text'])).select_columns(['input_ids'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[32010, 15043,    13, 32001,  1128,   526,   366, 29973,    13, 32010,\n",
       "           306, 29915, 29885,  2691, 29991,   887, 29973,    13, 32001,   306,\n",
       "         29915, 29885,  2691,  2086, 29991,    13, 32000]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1]]), 'labels': tensor([[ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,   306,\n",
       "         29915, 29885,  2691,  2086, 29991,    13, 32000]])}"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "completion_dloader = DataLoader(dummy_ds, batch_size=1, collate_fn=completion_collator)\n",
    "completion_batch = next(iter(completion_dloader))\n",
    "completion_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"I'm fine too!\\n<|endoftext|>\""
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels = completion_batch['labels']\n",
    "tokenizer_phi.decode(labels[labels >= 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[32010, 15043,    13, 32001,  1128,   526,   366, 29973,    13, 32010,\n",
       "           306, 29915, 29885,  2691, 29991,   887, 29973,    13, 32001,   306,\n",
       "         29915, 29885,  2691,  2086, 29991,    13, 32000]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1]]), 'labels': tensor([[ -100,  -100,  -100,  -100,  1128,   526,   366, 29973,    13,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,   306,\n",
       "         29915, 29885,  2691,  2086, 29991,    13, 32000]])}"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "instruction_template = '<|user|>'\n",
    "response_template = '<|assistant|>'\n",
    "completion_collator = DataCollatorForCompletionOnlyLM(instruction_template=instruction_template,\n",
    "                                                      response_template=response_template, \n",
    "                                                      tokenizer=tokenizer_phi)\n",
    "completion_dloader = DataLoader(dummy_ds, batch_size=1, collate_fn=completion_collator)\n",
    "completion_batch = next(iter(completion_dloader))\n",
    "completion_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"How are you?\\n I'm fine too!\\n<|endoftext|>\""
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels = completion_batch['labels']\n",
    "tokenizer_phi.decode(labels[labels >= 0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Label Shifting\n",
    "\n",
    "```python\n",
    "if labels is not None:\n",
    "    # move labels to correct device to enable model parallelism\n",
    "    labels = labels.to(lm_logits.device)\n",
    "    # we are doing next-token prediction; shift prediction scores and input ids by one\n",
    "    shift_logits = lm_logits[:, :-1, :].contiguous()\n",
    "    labels = labels[:, 1:].contiguous()\n",
    "    loss_fct = CrossEntropyLoss()\n",
    "    lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1))\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Packed Dataset\n",
    "\n",
    "****\n",
    "**Summary of \"Packed Dataset\"**\n",
    "- Packing **concatenates** sequences and **splits** them into **equal-sized packs**:\n",
    "  - **No padding tokens** are used.\n",
    "  - Each pack's length must not exceed the **model’s maximum sequence length**.\n",
    "- Packing is natively supported by the `SFTTrainer`:\n",
    "  - Set its `packing` argument to `True`.\n",
    "  - It creates an internal `ConstantLengthDataset` to handle the packing.\n",
    "  - By default, you **cannot use packing and a collator simultaneously**.\n",
    "- Some **collators can effectively pack** sequences:\n",
    "  - In this case, the **`packing` argument must be set to `False`**, and the collator performs the packing.\n",
    "  - `DataCollatorWithFlattening` is the packing equivalent of `DataCollatorForLanguageModeling`.\n",
    "  - `DataCollatorForCompletionOnlyLM` includes a new argument (`padding_free`) that makes the completion-only collator function like packing.\n",
    "  - Certain models (e.g. Llama, Phi, Mistral, Gemma, OLMo, and a few others) support these collators with Flash Attention 2:\n",
    "    - These models use `position_ids` to **mark the boundaries** between the original sequences packed together.\n",
    "****\n",
    "\n",
    "![](https://github.com/dvgodoy/FineTuningLLMs/blob/main/images/ch4/packed_seq.png?raw=True)\n",
    "\n",
    "<center>Figure 4.7 - Packed sequences</center>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['<|user|>\\nThe birch canoe slid on the smooth planks.<|end|>\\n<|assistant|>\\nOn the smooth planks, the birch canoe slid. Yes, hrrrm.<|end|>\\n<|endoftext|>', '<|user|>\\nGlue the sheet to the dark blue background.<|end|>\\n<|assistant|>\\nGlue the sheet to the dark blue background, you must.<|end|>\\n<|endoftext|>']\n"
     ]
    }
   ],
   "source": [
    "sequences = dataset['text']\n",
    "print(sequences[:2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['input_ids', 'labels'],\n",
       "    num_rows: 351\n",
       "})"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "iterator = ConstantLengthDataset(tokenizer_phi, dataset, \n",
    "                                 dataset_text_field='text', \n",
    "                                 seq_length=64, shuffle=False)\n",
    "\n",
    "def data_generator(iterator):\n",
    "    yield from iterator\n",
    "\n",
    "packed_dataset = Dataset.from_generator(\n",
    "    data_generator, \n",
    "    gen_kwargs={\"iterator\": iterator}\n",
    ")\n",
    "packed_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<|user|> The birch canoe slid on the smooth planks.<|end|><|assistant|> On the smooth planks, the birch canoe slid. Yes, hrrrm.<|end|><|endoftext|><|endoftext|><|user|> Glue the sheet to the dark blue background.<|end|><|assistant|> Glue the sheet to the dark blue background, you must'"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_ids = packed_dataset['input_ids']\n",
    "tokenizer_phi.decode(input_ids[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](https://github.com/dvgodoy/FineTuningLLMs/blob/main/images/ch4/packing_flow.png?raw=True)\n",
    "\n",
    "<center>Figure 4.8 - Choosing the right configuration for your data</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Collators for Packing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### `DataCollatorWithFlattening`\n",
    "\n",
    "![](https://github.com/dvgodoy/FineTuningLLMs/blob/main/images/ch4/collator_flat.png?raw=True)\n",
    "\n",
    "<center>Figure 4.9 - Packing-like collator</center>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "UserWarning: Using `DataCollatorWithFlattening` will flatten the entire mini batch into single long sequence.Make sure your attention computation is able to handle it!\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[32010,   450, 29773,   305,   508,  7297,  2243,   333,   373,   278,\n",
       "          10597,   715,  1331, 29889, 32007, 32001,  1551,   278, 10597,   715,\n",
       "           1331, 29892,   278, 29773,   305,   508,  7297,  2243,   333, 29889,\n",
       "           3869, 29892,   298, 21478,  1758, 29889, 32007, 32000, 32010,  8467,\n",
       "            434,   278,  9869,   304,   278,  6501,  7254,  3239, 29889, 32007,\n",
       "          32001,  8467,   434,   278,  9869,   304,   278,  6501,  7254,  3239,\n",
       "          29892,   366,  1818, 29889, 32007, 32000]]),\n",
       " 'labels': tensor([[ -100,   450, 29773,   305,   508,  7297,  2243,   333,   373,   278,\n",
       "          10597,   715,  1331, 29889, 32007, 32001,  1551,   278, 10597,   715,\n",
       "           1331, 29892,   278, 29773,   305,   508,  7297,  2243,   333, 29889,\n",
       "           3869, 29892,   298, 21478,  1758, 29889, 32007, 32000,  -100,  8467,\n",
       "            434,   278,  9869,   304,   278,  6501,  7254,  3239, 29889, 32007,\n",
       "          32001,  8467,   434,   278,  9869,   304,   278,  6501,  7254,  3239,\n",
       "          29892,   366,  1818, 29889, 32007, 32000]]),\n",
       " 'position_ids': tensor([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,\n",
       "          18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,\n",
       "          36, 37,  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15,\n",
       "          16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27]])}"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "flat_collator = DataCollatorWithFlattening()\n",
    "flat_dloader = DataLoader(tokenized_dataset, batch_size=2, collate_fn=flat_collator)\n",
    "flat_batch = next(iter(flat_dloader))\n",
    "flat_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([1, 66]), tensor(38))"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "flat_batch['input_ids'].shape, flat_batch['position_ids'].max() + 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### `DataCollatorForCompletionOnlyLM`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[32010,   450, 29773,   305,   508,  7297,  2243,   333,   373,   278,\n",
       "         10597,   715,  1331, 29889, 32007, 32001,  1551,   278, 10597,   715,\n",
       "          1331, 29892,   278, 29773,   305,   508,  7297,  2243,   333, 29889,\n",
       "          3869, 29892,   298, 21478,  1758, 29889, 32007, 32000, 32010,  8467,\n",
       "           434,   278,  9869,   304,   278,  6501,  7254,  3239, 29889, 32007,\n",
       "         32001,  8467,   434,   278,  9869,   304,   278,  6501,  7254,  3239,\n",
       "         29892,   366,  1818, 29889, 32007, 32000]]), 'labels': tensor([[ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  1551,   278, 10597,   715,\n",
       "          1331, 29892,   278, 29773,   305,   508,  7297,  2243,   333, 29889,\n",
       "          3869, 29892,   298, 21478,  1758, 29889, 32007, 32000,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  8467,   434,   278,  9869,   304,   278,  6501,  7254,  3239,\n",
       "         29892,   366,  1818, 29889, 32007, 32000]]), 'position_ids': tensor([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,\n",
       "         18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,\n",
       "         36, 37,  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15,\n",
       "         16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27]])}"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response_template = '<|assistant|>'\n",
    "completion_nopad_collator = DataCollatorForCompletionOnlyLM(response_template=response_template, \n",
    "                                                            tokenizer=tokenizer_phi,\n",
    "                                                            padding_free=True)\n",
    "completion_nopad_dloader = DataLoader(tokenized_dataset, batch_size=2, collate_fn=completion_nopad_collator)\n",
    "completion_nopad_batch = next(iter(completion_nopad_dloader))\n",
    "completion_nopad_batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Advanced: BYOT (Bring Your Own Template)\n",
    "\n",
    "****\n",
    "**Summary of \"Advanced: BYOT\"**\n",
    "- Every template must define a **response template** and, ideally, **end with an `EOS` token**.\n",
    "- Double-check your tokenizer's `EOS`, `PAD`, and `UNK` tokens:\n",
    "  - The `EOS` token must be distinct from both `PAD` and `UNK` tokens.\n",
    "  - The `PAD` and `UNK` tokens can be the same.\n",
    "- **Only resize** the embedding layer if absolutely necessary—i.e., if all \"empty slots\" have already been used:\n",
    "  - When calling the model's `resize_token_embeddings()`, use the `pad_to_multiple_of` argument to ensure the size remains a **multiple of a power of two**.\n",
    "- If you don’t want to create a Jinja template yourself, you can use a default template like ChatML.\n",
    "  - The `trl` package provides the `setup_chat_format()` function, but it has some drawbacks:\n",
    "    - It assigns the `EOS` token to the `PAD` token (you'll need to **fix it manually** afterward).\n",
    "    - It resizes the model's embedding layer by default, even if only to make it shorter (though **you can avoid resizing** by selecting the appropriate `resize_to_multiple_of`).\n",
    "- You can define and apply a **custom template using a formatting function** instead of creating a Jinja template for your tokenizer:\n",
    "  - If you specify the `formatting_func` in the `SFTTrainer` class (see Chapter 5), your tokenizer doesn't need to have a chat template.\n",
    "  - Choose your response template carefully:\n",
    "    - Using **regular words** (e.g. \"## Answer:\") **may cause issues**, as some tokenizers are \"context-dependent\" and might split your response template into multiple tokens.\n",
    "    - Creating an **additional special token for your response template** is safer, as it will be encoded as a **single token**.\n",
    "****"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Chat Template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 270,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "None\n"
     ]
    }
   ],
   "source": [
    "model_opt = AutoModelForCausalLM.from_pretrained(\"facebook/opt-350m\")\n",
    "tokenizer_opt = AutoTokenizer.from_pretrained(\"facebook/opt-350m\")\n",
    "\n",
    "print(tokenizer_opt.chat_template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'bos_token': '</s>',\n",
       " 'eos_token': '</s>',\n",
       " 'unk_token': '</s>',\n",
       " 'pad_token': '<pad>'}"
      ]
     },
     "execution_count": 128,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer_opt.special_tokens_map"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**ChatML**\n",
    "****\n",
    "[ChatML](https://github.com/openai/openai-python/blob/release-v0.28.0/chatml.md), short for Chat Markup Language, was developed by OpenAI: \n",
    "\n",
    "_____\n",
    "\"_Traditionally, GPT models consumed unstructured text. ChatGPT models instead expect a structured format, called Chat Markup Language (ChatML for short). ChatML documents consist of a sequence of messages._\"\n",
    "_____\n",
    "\n",
    "Each message should contain the role of the participant and their corresponding content, like the conversational format introduced earlier. This is ChatML's Jinja template:\n",
    "\n",
    "```\n",
    "{% for message in messages %}\n",
    "  {{'<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n'}}\n",
    "{% endfor %}\n",
    "```\n",
    "****"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "50265"
      ]
     },
     "execution_count": 129,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(tokenizer_opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "50272"
      ]
     },
     "execution_count": 131,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_opt.config.vocab_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "32"
      ]
     },
     "execution_count": 132,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_multiple_of(vocab_size):\n",
    "    return 2**(bin(vocab_size)[::-1].find('1'))\n",
    "\n",
    "pad_to_multiple_of = get_multiple_of(model_opt.config.vocab_size)\n",
    "pad_to_multiple_of"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Embedding(50272, 512, padding_idx=1)"
      ]
     },
     "execution_count": 133,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_opt.resize_token_embeddings(len(tokenizer_opt), \n",
    "                                  pad_to_multiple_of=pad_to_multiple_of)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 217,
   "metadata": {},
   "outputs": [],
   "source": [
    "def modify_tokenizer(tokenizer, \n",
    "                     alternative_bos_token='<|im_start|>', \n",
    "                     alternative_unk_token='<unk>', \n",
    "                     special_tokens=None, \n",
    "                     tokens=None):\n",
    "    eos_token, bos_token = tokenizer.eos_token, tokenizer.bos_token\n",
    "    pad_token, unk_token = tokenizer.pad_token, tokenizer.unk_token\n",
    "\n",
    "    # BOS token must be different than EOS token\n",
    "    if bos_token == eos_token:\n",
    "        bos_token = alternative_bos_token\n",
    "\n",
    "    # UNK token must be different than EOS token\n",
    "    if unk_token == eos_token:\n",
    "        unk_token = alternative_unk_token\n",
    "\n",
    "    # PAD token must be different than EOS token\n",
    "    # but can be the same as UNK token\n",
    "    if pad_token == eos_token:\n",
    "        pad_token = unk_token\n",
    "        \n",
    "    assert bos_token != eos_token, \"Please choose a different BOS token.\"\n",
    "    assert unk_token != eos_token, \"Please choose a different UNK token.\"\n",
    "\n",
    "    # Creates dict for BOS, PAD, and UNK tokens\n",
    "    # Keeps the EOS token as it was originally defined\n",
    "    special_tokens_dict = {'bos_token': bos_token, \n",
    "                           'pad_token': pad_token, \n",
    "                           'unk_token': unk_token}\n",
    "    \n",
    "    # If there are additional special tokens, add them\n",
    "    if special_tokens is not None:\n",
    "        if isinstance(special_tokens, list):\n",
    "            special_tokens_dict.update({'additional_special_tokens': special_tokens})\n",
    "        \n",
    "    tokenizer.add_special_tokens(special_tokens_dict)\n",
    "    \n",
    "    # If there are new regular (not special) tokens to add\n",
    "    if tokens is not None:\n",
    "        if isinstance(tokens, list):\n",
    "            tokenizer.add_tokens(tokens)\n",
    "        \n",
    "    return tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 216,
   "metadata": {},
   "outputs": [],
   "source": [
    "def jinja_template(tokenizer):\n",
    "    return (\"{% for message in messages %}\"\n",
    "            f\"{{{{'{tokenizer.bos_token}' + message['role'] + '\\n' + message['content'] + '{tokenizer.eos_token}' + '\\n'}}}}\"\n",
    "            \"{% endfor %}\"\n",
    "            \"{% if add_generation_prompt %}\"\n",
    "            f\"{{{{ '{tokenizer.bos_token}assistant\\n' }}}}\"\n",
    "            \"{% endif %}\")\n",
    "\n",
    "def add_template(tokenizer, chat_template=None):\n",
    "    # If not chat template was given, creates a ChatML template\n",
    "    # using the BOS and EOS tokens\n",
    "    if chat_template is None:\n",
    "        chat_template = jinja_template(tokenizer)\n",
    "        \n",
    "    # Assigns chat template to tokenizer\n",
    "    tokenizer.chat_template = chat_template\n",
    "    \n",
    "    return tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 219,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_multiple_of(vocab_size):\n",
    "    return 2**(bin(vocab_size)[::-1].find('1'))\n",
    "\n",
    "def modify_model(model, tokenizer):    \n",
    "    # If new tokenizer length exceeds vocabulary size\n",
    "    # resizes it while keeping it a multiple of the same value\n",
    "    if len(tokenizer) > model.config.vocab_size:\n",
    "        pad_to_multiple_of = get_multiple_of(model.vocab_size)\n",
    "        model.resize_token_embeddings(len(tokenizer), \n",
    "                                      pad_to_multiple_of=pad_to_multiple_of)    \n",
    "\n",
    "    # Updates token ids on model configurations\n",
    "    if getattr(model, \"config\", None) is not None:\n",
    "        model.config.pad_token_id = tokenizer.pad_token_id\n",
    "        model.config.bos_token_id = tokenizer.bos_token_id\n",
    "        model.config.eos_token_id = tokenizer.eos_token_id\n",
    "    if getattr(model, \"generation_config\", None) is not None:\n",
    "        model.generation_config.bos_token_id = tokenizer.bos_token_id\n",
    "        model.generation_config.eos_token_id = tokenizer.eos_token_id\n",
    "        model.generation_config.pad_token_id = tokenizer.pad_token_id\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 271,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer_opt = modify_tokenizer(tokenizer_opt)\n",
    "tokenizer_opt = add_template(tokenizer_opt)\n",
    "model_opt = modify_model(model_opt, tokenizer_opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 273,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'bos_token': '<|im_start|>',\n",
       " 'eos_token': '</s>',\n",
       " 'unk_token': '<unk>',\n",
       " 'pad_token': '<pad>'}"
      ]
     },
     "execution_count": 273,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer_opt.special_tokens_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 281,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "50266"
      ]
     },
     "execution_count": 281,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(tokenizer_opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 282,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<|im_start|>'"
      ]
     },
     "execution_count": 282,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer_opt.convert_ids_to_tokens(50265)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 283,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Embedding(50272, 512, padding_idx=1)"
      ]
     },
     "execution_count": 283,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_opt.get_input_embeddings()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 221,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n",
      "' + message['content'] + '</s>' + '\n",
      "'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n",
      "' }}{% endif %}\n"
     ]
    }
   ],
   "source": [
    "print(tokenizer_opt.chat_template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 274,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|im_start|>user\n",
      "What is the capital of Argentina?</s>\n",
      "<|im_start|>assistant\n",
      "Buenos Aires.</s>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "messages = ds_msg['messages'][0]\n",
    "print(tokenizer_opt.apply_chat_template(messages, tokenize=False))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Custom Template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 284,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_opt = AutoModelForCausalLM.from_pretrained(\"facebook/opt-350m\")\n",
    "tokenizer_opt = AutoTokenizer.from_pretrained(\"facebook/opt-350m\")\n",
    "\n",
    "response_template = '##[YODA]##>'\n",
    "tokenizer_opt = modify_tokenizer(tokenizer_opt, special_tokens=[response_template])\n",
    "model_opt = modify_model(model_opt, tokenizer_opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 285,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function __main__.gen_formatting_func.<locals>.formatting_func(examples, add_generation_prompt=False)>"
      ]
     },
     "execution_count": 285,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def formatting_func_builder(response_template):\n",
    "    def formatting_func(examples, add_generation_prompt=False):\n",
    "        output_texts = []\n",
    "        for i in range(len(examples['prompt'])):\n",
    "            text = f\"{examples['prompt'][i]}\"\n",
    "            try:\n",
    "                text += f\" {response_template} {examples['completion'][i]}{tokenizer_opt.eos_token}\"\n",
    "            except KeyError:\n",
    "                if add_generation_prompt:\n",
    "                    text += f\" {response_template} \"\n",
    "            output_texts.append(text)\n",
    "        return output_texts\n",
    "    return formatting_func\n",
    "\n",
    "yoda_formatting_func = formatting_func_builder(response_template)\n",
    "yoda_formatting_func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 289,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The birch canoe slid on the smooth planks. ##[YODA]##> On the smooth planks, the birch canoe slid. Yes, hrrrm.</s>'"
      ]
     },
     "execution_count": 289,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "formatted_seqs = yoda_formatting_func(dataset)\n",
    "formatted_seqs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 291,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': [2, 133, 23629, 611, 31728, 13763, 15, 5, 6921, 563, 2258, 4, 1437, 50266, 374, 5, 6921, 563, 2258, 6, 5, 23629, 611, 31728, 13763, 4, 3216, 6, 1368, 28015, 22900, 4, 2], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}"
      ]
     },
     "execution_count": 291,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer_opt(formatted_seqs[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 293,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'##[YODA]##>'"
      ]
     },
     "execution_count": 293,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer_opt.convert_ids_to_tokens(50266)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 294,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['The Force is strong in you. ##[YODA]##> ', 'I am your father! ##[YODA]##> ']"
      ]
     },
     "execution_count": 294,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "yoda_formatting_func({'prompt': ['The Force is strong in you.', \n",
    "                                 'I am your father!']}, \n",
    "                     add_generation_prompt=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Special Tokens FTW"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 184,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer_llama = AutoTokenizer.from_pretrained(\"meta-llama/Llama-2-7b-hf\")\n",
    "tokenizer_llama.pad_token = tokenizer_llama.unk_token\n",
    "tokenizer_llama.pad_token_id = tokenizer_llama.unk_token_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 185,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "### User: Hello\n",
      "\n",
      "### Assistant: Hi, how can I help you?\n"
     ]
    }
   ],
   "source": [
    "prompt = \"\"\"### User: Hello\\n\\n### Assistant: Hi, how can I help you?\"\"\"\n",
    "print(prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 186,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('##', 2277), ('#', 29937), ('▁Ass', 4007), ('istant', 22137), (':', 29901)]"
      ]
     },
     "execution_count": 186,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokens = tokenizer_llama.tokenize(prompt, add_special_tokens=False)\n",
    "token_ids = tokenizer_llama.encode(prompt, add_special_tokens=False)\n",
    "list(zip(tokens, token_ids))[6:11]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 187,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('▁###', 835), ('▁Ass', 4007), ('istant', 22137), (':', 29901)]"
      ]
     },
     "execution_count": 187,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response_template = \"### Assistant:\"\n",
    "tokens = tokenizer_llama.tokenize(response_template, add_special_tokens=False)\n",
    "token_ids = tokenizer_llama.encode(response_template, add_special_tokens=False)\n",
    "list(zip(tokens, token_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 188,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ddddfc8b0e1c40019a74c438910a2cf0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[    1,   835,  4911, 29901, 15043,    13,    13,  2277, 29937,  4007,\n",
       "         22137, 29901,  6324, 29892,   920,   508,   306,  1371,   366, 29973]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]), 'labels': tensor([[-100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
       "         -100, -100, -100, -100, -100, -100, -100, -100]])}"
      ]
     },
     "execution_count": 188,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dummy_ds = Dataset.from_dict({'text': [prompt]})\n",
    "dummy_tokenized = dummy_ds.map(lambda row: tokenizer_llama(row['text'])).select_columns(['input_ids'])\n",
    "\n",
    "response_template = \"### Assistant:\"\n",
    "\n",
    "bad_collator = DataCollatorForCompletionOnlyLM(response_template, tokenizer=tokenizer_llama)\n",
    "bad_dloader = DataLoader(dummy_tokenized, batch_size=1, collate_fn=bad_collator)\n",
    "bad_batch = next(iter(bad_dloader))\n",
    "bad_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 189,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('▁', 29871),\n",
       " ('<0x0A>', 13),\n",
       " ('##', 2277),\n",
       " ('#', 29937),\n",
       " ('▁Ass', 4007),\n",
       " ('istant', 22137),\n",
       " (':', 29901)]"
      ]
     },
     "execution_count": 189,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "modified_response_template = \"\\n### Assistant:\"\n",
    "tokens = tokenizer_llama.tokenize(modified_response_template, add_special_tokens=False)\n",
    "token_ids = tokenizer_llama.encode(modified_response_template, add_special_tokens=False)\n",
    "list(zip(tokens, token_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 190,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[    1,   835,  4911, 29901, 15043,    13,    13,  2277, 29937,  4007,\n",
       "         22137, 29901,  6324, 29892,   920,   508,   306,  1371,   366, 29973]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]), 'labels': tensor([[ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  6324, 29892,   920,   508,   306,  1371,   366, 29973]])}"
      ]
     },
     "execution_count": 190,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fixed_token_ids = token_ids[2:]\n",
    "fixed_collator = DataCollatorForCompletionOnlyLM(fixed_token_ids, tokenizer=tokenizer_llama)\n",
    "fixed_dloader = DataLoader(dummy_tokenized, batch_size=1, collate_fn=fixed_collator)\n",
    "fixed_batch = next(iter(fixed_dloader))\n",
    "fixed_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 191,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 191,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response_template = \"### Assistant:\"\n",
    "tokenizer_llama.add_special_tokens({'additional_special_tokens': [response_template]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 192,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4664c57b3b9841aa83cb6951ec4e3456",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dummy_tokenized = dummy_ds.map(lambda row: tokenizer_llama(row['text'])).select_columns(['input_ids'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[    1,   835,  4911, 29901, 15043,    13,    13, 32000, 29871,  6324,\n",
       "         29892,   920,   508,   306,  1371,   366, 29973]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]), 'labels': tensor([[ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100, 29871,  6324,\n",
       "         29892,   920,   508,   306,  1371,   366, 29973]])}"
      ]
     },
     "execution_count": 193,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "special_collator = DataCollatorForCompletionOnlyLM(response_template, tokenizer=tokenizer_llama)\n",
    "special_dloader = DataLoader(dummy_tokenized, batch_size=1, collate_fn=special_collator)\n",
    "special_batch = next(iter(special_dloader))\n",
    "special_batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Coming Up in \"Fine-Tuning LLMs\"\n",
    "\n",
    "Chat templates are key to reining in the untamed LLM monsters and teaching them how to have proper conversations with us humans. Cleverly placing cues, or special tokens, along the conversation enables them to learn how to respond when triggered by the right commanding keyword. The training procedure, though, is not without its perils: activations, gradients, and the optimizer all demand huge portions of precious RAM in order to do their jobs. Appeasing these memory-hungry components will take both skill and effort. Configuring the training loop isn’t for the faint of heart. Don’t miss the next challenging chapter of \"Fine-Tuning LLMs.\""
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}