1_instruction_tuning/notebooks/sft_finetuning_example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Supervised Fine-Tuning with SFTTrainer\n",
    "\n",
    "This notebook demonstrates how to fine-tune the `HuggingFaceTB/SmolLM2-135M` model using the `SFTTrainer` from the `trl` library. The notebook cells run and will finetune the model. You can select your difficulty by trying out different datasets.\n",
    "\n",
    "<div style='background-color: lightblue; padding: 10px; border-radius: 5px; margin-bottom: 20px; color:black'>\n",
    "    <h2 style='margin: 0;color:blue'>Exercise: Fine-Tuning SmolLM2 with SFTTrainer</h2>\n",
    "    <p>Take a dataset from the Hugging Face hub and finetune a model on it. </p> \n",
    "    <p><b>Difficulty Levels</b></p>\n",
    "    <p>🐢 Use the `HuggingFaceTB/smoltalk` dataset</p>\n",
    "    <p>🐕 Try out the `bigcode/the-stack-smol` dataset and finetune a code generation model on a specific subset `data/python`.</p>\n",
    "    <p>🦁 Select a dataset that relates to a real world use case your interested in</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install the requirements in Google Colab\n",
    "# !pip install transformers datasets trl huggingface_hub\n",
    "\n",
    "# Authenticate to Hugging Face\n",
    "\n",
    "from huggingface_hub import login\n",
    "login()\n",
    "\n",
    "# for convenience you can create an environment variable containing your hub token as HF_TOKEN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from datasets import load_dataset\n",
    "from trl import SFTConfig, SFTTrainer, setup_chat_format\n",
    "import torch\n",
    "\n",
    "device = (\n",
    "    \"cuda\"\n",
    "    if torch.cuda.is_available()\n",
    "    else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
    ")\n",
    "\n",
    "# Load the model and tokenizer\n",
    "model_name = \"HuggingFaceTB/SmolLM2-135M\"\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    pretrained_model_name_or_path=model_name\n",
    ").to(device)\n",
    "tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model_name)\n",
    "\n",
    "# Set up the chat format\n",
    "model, tokenizer = setup_chat_format(model=model, tokenizer=tokenizer)\n",
    "\n",
    "# Set our name for the finetune to be saved &/ uploaded to\n",
    "finetune_name = \"SmolLM2-FT-MyDataset\"\n",
    "finetune_tags = [\"smol-course\", \"module_1\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate with the base model\n",
    "\n",
    "Here we will try out the base model which does not have a chat template. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's test the base model before training\n",
    "prompt = \"Write a haiku about programming\"\n",
    "\n",
    "# Format with template\n",
    "messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False)\n",
    "\n",
    "# Generate response\n",
    "inputs = tokenizer(formatted_prompt, return_tensors=\"pt\").to(device)\n",
    "outputs = model.generate(**inputs, max_new_tokens=100)\n",
    "print(\"Before training:\")\n",
    "print(tokenizer.decode(outputs[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset Preparation\n",
    "\n",
    "We will load a sample dataset and format it for training. The dataset should be structured with input-output pairs, where each input is a prompt and the output is the expected response from the model.\n",
    "\n",
    "**TRL will format input messages based on the model's chat templates.** They need to be represented as a list of dictionaries with the keys: `role` and `content`,."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load a sample dataset\n",
    "from datasets import load_dataset\n",
    "\n",
    "# TODO: define your dataset and config using the path and name parameters\n",
    "ds = load_dataset(path=\"HuggingFaceTB/smoltalk\", name=\"everyday-conversations\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: 🦁 If your dataset is not in a format that TRL can convert to the chat template, you will need to process it. Refer to the [module](../chat_templates.md)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuring the SFTTrainer\n",
    "\n",
    "The `SFTTrainer` is configured with various parameters that control the training process. These include the number of training steps, batch size, learning rate, and evaluation strategy. Adjust these parameters based on your specific requirements and computational resources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure the SFTTrainer\n",
    "sft_config = SFTConfig(\n",
    "    output_dir=\"./sft_output\",\n",
    "    max_steps=1000,  # Adjust based on dataset size and desired training duration\n",
    "    per_device_train_batch_size=4,  # Set according to your GPU memory capacity\n",
    "    learning_rate=5e-5,  # Common starting point for fine-tuning\n",
    "    logging_steps=10,  # Frequency of logging training metrics\n",
    "    save_steps=100,  # Frequency of saving model checkpoints\n",
    "    evaluation_strategy=\"steps\",  # Evaluate the model at regular intervals\n",
    "    eval_steps=50,  # Frequency of evaluation\n",
    "    use_mps_device=(\n",
    "        True if device == \"mps\" else False\n",
    "    ),  # Use MPS for mixed precision training\n",
    "    hub_model_id=finetune_name,  # Set a unique name for your model\n",
    ")\n",
    "\n",
    "# Initialize the SFTTrainer\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    args=sft_config,\n",
    "    train_dataset=ds[\"train\"],\n",
    "    tokenizer=tokenizer,\n",
    "    eval_dataset=ds[\"test\"],\n",
    ")\n",
    "\n",
    "# TODO: 🦁 🐕 align the SFTTrainer params with your chosen dataset. For example, if you are using the `bigcode/the-stack-smol` dataset, you will need to choose the `content` column`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training the Model\n",
    "\n",
    "With the trainer configured, we can now proceed to train the model. The training process will involve iterating over the dataset, computing the loss, and updating the model's parameters to minimize this loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the model\n",
    "trainer.train()\n",
    "\n",
    "# Save the model\n",
    "trainer.save_model(f\"./{finetune_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.push_to_hub(tags=finetune_tags)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style='background-color: lightblue; padding: 10px; border-radius: 5px; margin-bottom: 20px; color:black'>\n",
    "    <h2 style='margin: 0;color:blue'>Bonus Exercise: Generate with fine-tuned model</h2>\n",
    "    <p>🐕 Use the fine-tuned to model generate a response, just like with the base example..</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test the fine-tuned model on the same prompt\n",
    "\n",
    "# Let's test the base model before training\n",
    "prompt = \"Write a haiku about programming\"\n",
    "\n",
    "# Format with template\n",
    "messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False)\n",
    "\n",
    "# Generate response\n",
    "inputs = tokenizer(formatted_prompt, return_tensors=\"pt\").to(device)\n",
    "\n",
    "# TODO: use the fine-tuned to model generate a response, just like with the base example."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 💐 You're done!\n",
    "\n",
    "This notebook provided a step-by-step guide to fine-tuning the `HuggingFaceTB/SmolLM2-135M` model using the `SFTTrainer`. By following these steps, you can adapt the model to perform specific tasks more effectively. If you want to carry on working on this course, here are steps you could try out:\n",
    "\n",
    "- Try this notebook on a harder difficulty\n",
    "- Review a colleagues PR\n",
    "- Improve the course material via an Issue or PR."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py310",
   "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.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}