1_instruction_tuning/notebooks/chat_templates_example.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F3nzKE2UiC5Y"
      },
      "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": null,
      "metadata": {
        "id": "RQyZSF2fiC5Z",
        "outputId": "2f4825e1-d177-4117-feca-2f714491cdeb"
      },
      "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\n",
        "import torch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cl3KwX0JiC5a"
      },
      "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": null,
      "metadata": {
        "id": "A0-6k26niC5a"
      },
      "outputs": [],
      "source": [
        "# Dynamically set the device\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
        "\n",
        "model_name = \"HuggingFaceTB/SmolLM2-135M\"\n",
        "model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=model_name).to(device)\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": null,
      "metadata": {
        "id": "z8ur2jxKiC5a"
      },
      "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": {
        "id": "HZHBCpqDiC5a"
      },
      "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": {
        "id": "Bc38j3kriC5b",
        "outputId": "f826dc68-3052-421e-b45d-4e40515a3068"
      },
      "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": {
        "id": "o2sOIqSEiC5b"
      },
      "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": null,
      "metadata": {
        "id": "5cSaq3_XiC5b",
        "outputId": "fa92bd64-820b-4941-b71f-730d9f20b09e"
      },
      "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": {
        "id": "Y-0OaQZ9iC5b"
      },
      "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": null,
      "metadata": {
        "id": "cnNislvxiC5b",
        "outputId": "1d51e564-036e-44ea-e2b3-a1c136a383ef"
      },
      "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": {
        "id": "AgMBIAZliC5b"
      },
      "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": null,
      "metadata": {
        "id": "3eAAyOtdiC5b",
        "outputId": "797bebd3-d1fe-492a-87e1-f3670d7bb64d"
      },
      "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": null,
      "metadata": {
        "id": "PK5VV_4yiC5c",
        "outputId": "2916e3a4-d6f6-475b-bb37-95980d13a11a"
      },
      "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": null,
      "metadata": {
        "id": "30VHtxjaiC5c",
        "outputId": "52af2644-eaeb-4915-ddfa-7e1ee717ac98"
      },
      "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": {
        "id": "OEivfg05iC5c"
      },
      "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": {
        "id": "dNRKj14YiC5c"
      },
      "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"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}