1_instruction_tuning/notebooks/chat_templates_example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exploring Chat Templates with SmolLM2 and Llama 3.2\n",
    "\n",
    "This notebook demonstrates how to use chat templates with the `SmolLM2` and `Llama 3.2` models. Chat templates help structure interactions between users and AI models, ensuring consistent and contextually appropriate responses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/thilo/smol-course/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "# Import necessary libraries\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from trl import setup_chat_format"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SmolLM2 Chat Template\n",
    "\n",
    "Let's explore how to use a chat template with the `SmolLM2` model. We'll define a simple conversation and apply the chat template."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"HuggingFaceTB/SmolLM2-135M\"\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    pretrained_model_name_or_path=model_name\n",
    ").to(\"mps\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model_name)\n",
    "model, tokenizer = setup_chat_format(model=model, tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define messages for SmolLM2\n",
    "messages = [\n",
    "    {\"role\": \"user\", \"content\": \"Hello, how are you?\"},\n",
    "    {\n",
    "        \"role\": \"assistant\",\n",
    "        \"content\": \"I'm doing well, thank you! How can I assist you today?\",\n",
    "    },\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Apply chat template without tokenization\n",
    "\n",
    "The tokenizer represents the conversation as a string with special tokens to describe the role of the user and the assistant.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Conversation with template: <|im_start|>user\n",
      "Hello, how are you?<|im_end|>\n",
      "<|im_start|>assistant\n",
      "I'm doing well, thank you! How can I assist you today?<|im_end|>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "input_text = tokenizer.apply_chat_template(messages, tokenize=False)\n",
    "\n",
    "print(\"Conversation with template:\", input_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Decode the conversation\n",
    "\n",
    "Note that the conversation is represented as above but with a further assistant message.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Conversation decoded: <|im_start|>user\n",
      "Hello, how are you?<|im_end|>\n",
      "<|im_start|>assistant\n",
      "I'm doing well, thank you! How can I assist you today?<|im_end|>\n",
      "<|im_start|>assistant\n",
      "\n"
     ]
    }
   ],
   "source": [
    "input_text = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True)\n",
    "\n",
    "print(\"Conversation decoded:\", tokenizer.decode(token_ids=input_text))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tokenize the conversation\n",
    "\n",
    "Of course, the tokenizer also tokenizes the conversation and special token as ids that relate to the model's vocabulary.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Conversation tokenized: [1, 4093, 198, 19556, 28, 638, 359, 346, 47, 2, 198, 1, 520, 9531, 198, 57, 5248, 2567, 876, 28, 9984, 346, 17, 1073, 416, 339, 4237, 346, 1834, 47, 2, 198, 1, 520, 9531, 198]\n"
     ]
    }
   ],
   "source": [
    "input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True)\n",
    "\n",
    "print(\"Conversation tokenized:\", input_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style='background-color: lightblue; padding: 10px; border-radius: 5px; margin-bottom: 20px'>\n",
    "    <h2 style='margin: 0;color:blue'>Exercise: Process a dataset for SFT</h2>\n",
    "    <p>Take a dataset from the Hugging Face hub and process it for SFT. </p> \n",
    "    <p><b>Difficulty Levels</b></p>\n",
    "    <p>🐢 Convert the `HuggingFaceTB/smoltalk` dataset into chatml format.</p>\n",
    "    <p>🐕 Convert the `openai/gsm8k` dataset into chatml format.</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/8z/jnnncfnj7_lfxym0262z4p180000gn/T/ipykernel_94817/1790625692.py:1: DeprecationWarning: Importing display from IPython.core.display is deprecated since IPython 7.14, please import from IPython.display\n",
      "  from IPython.core.display import display, HTML\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<iframe\n",
       "  src=\"https://huggingface.co/datasets/HuggingFaceTB/smoltalk/embed/viewer/all/train?row=0\"\n",
       "  frameborder=\"0\"\n",
       "  width=\"100%\"\n",
       "  height=\"360px\"\n",
       "></iframe>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.core.display import display, HTML\n",
    "\n",
    "display(\n",
    "    HTML(\"\"\"<iframe\n",
    "  src=\"https://huggingface.co/datasets/HuggingFaceTB/smoltalk/embed/viewer/all/train?row=0\"\n",
    "  frameborder=\"0\"\n",
    "  width=\"100%\"\n",
    "  height=\"360px\"\n",
    "></iframe>\n",
    "\"\"\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map: 100%|██████████| 2260/2260 [00:00<00:00, 61231.51 examples/s]\n",
      "Map: 100%|██████████| 119/119 [00:00<00:00, 37117.73 examples/s]\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "ds = load_dataset(\"HuggingFaceTB/smoltalk\", \"everyday-conversations\")\n",
    "\n",
    "\n",
    "def process_dataset(sample):\n",
    "    # TODO: 🐢 Convert the sample into a chat format\n",
    "    return sample\n",
    "\n",
    "\n",
    "ds = ds.map(process_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<iframe\n",
       "  src=\"https://huggingface.co/datasets/openai/gsm8k/embed/viewer/main/train\"\n",
       "  frameborder=\"0\"\n",
       "  width=\"100%\"\n",
       "  height=\"360px\"\n",
       "></iframe>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(\n",
    "    HTML(\"\"\"<iframe\n",
    "  src=\"https://huggingface.co/datasets/openai/gsm8k/embed/viewer/main/train\"\n",
    "  frameborder=\"0\"\n",
    "  width=\"100%\"\n",
    "  height=\"360px\"\n",
    "></iframe>\n",
    "\"\"\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = load_dataset(\"openai/gsm8k\", \"main\")\n",
    "\n",
    "\n",
    "def process_dataset(sample):\n",
    "    # TODO: 🐕 Convert the sample into a chat format\n",
    "    return sample\n",
    "\n",
    "\n",
    "ds = ds.map(process_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "This notebook demonstrated how to apply chat templates to different models, `SmolLM2`. By structuring interactions with chat templates, we can ensure that AI models provide consistent and contextually relevant responses.\n",
    "\n",
    "In the exercise you tried out converting a dataset into chatml format. Luckily, TRL will do this for you, but it's useful to understand what's going on under the hood."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}