{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "926bada6"
      },
      "source": [
        "##### Copyright 2025 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "a110dfce"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f9673bd6"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://ai.google.dev/gemma/docs/core/huggingface_text_full_finetune\"><img src=\"https://ai.google.dev/static/site-assets/images/docs/notebook-site-button.png\" height=\"32\" width=\"32\" />View on ai.google.dev</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemini/gemma-cookbook/blob/main/docs/core/huggingface_text_full_finetune.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://kaggle.com/kernels/welcome?src=https://github.com/google-gemini/gemma-cookbook/blob/main/docs/core/huggingface_text_full_finetune.ipynb\"><img src=\"https://www.kaggle.com/static/images/logos/kaggle-logo-transparent-300.png\" height=\"32\" width=\"70\"/>Run in Kaggle</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://console.cloud.google.com/vertex-ai/colab/import/https%3A%2F%2Fraw.githubusercontent.com%2Fgoogle-gemini%2Fgemma-cookbook%2Fmain%2Fdocs%2Fcore%2Fhuggingface_text_full_finetune.ipynb\"><img src=\"https://ai.google.dev/images/cloud-icon.svg\" width=\"40\" />Open in Vertex AI</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/google-gemini/gemma-cookbook/blob/main/docs/core/huggingface_text_full_finetune.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e624ec07"
      },
      "source": [
        "# Full Model Fine-Tune using Hugging Face Transformers\n",
        "\n",
        "This guide walks you through how to fine-tune Gemma on a mobile game NPC dataset using Hugging Face [Transformers](https://huggingface.co/docs/transformers/index) and [TRL](https://huggingface.co/docs/trl/index). You will learn:\n",
        "\n",
        "- Setup development environment\n",
        "- Prepare the fine-tuning dataset\n",
        "- Full model fine-tuning Gemma using TRL and the SFTTrainer\n",
        "- Test Model Inference and vibe checks\n",
        "\n",
        "> Note: This guide was created to run on a Google colaboratory account using a NVIDIA T4 GPU with 16GB and Gemma 270m, but can be adapted to run on bigger GPUs and bigger models.\n",
        "\n",
        "## Setup development environment\n",
        "\n",
        "The first step is to install Hugging Face Libraries, including TRL, and datasets to fine-tune open model, including different RLHF and alignment techniques."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BEK9IfKBqQaA"
      },
      "outputs": [],
      "source": [
        "# Install Pytorch & other libraries\n",
        "%pip install torch tensorboard\n",
        "\n",
        "# Install Hugging Face libraries\n",
        "%pip install transformers datasets accelerate evaluate trl protobuf sentencepiece\n",
        "\n",
        "# COMMENT IN: if you are running on a GPU that supports BF16 data type and flash attn, such as NVIDIA L4 or NVIDIA A100\n",
        "#% pip install flash-attn"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7ef3d54b"
      },
      "source": [
        "> _Note: If you are using a GPU with Ampere architecture (such as NVIDIA L4) or newer, you can use Flash attention. Flash Attention is a method that significantly speeds computations up and reduces memory usage from quadratic to linear in sequence length, leading to acelerating training up to 3x. Learn more at [FlashAttention](https://github.com/Dao-AILab/flash-attention/tree/main)._\n",
        "\n",
        "Before you can start training, you have to make sure that you accepted the terms of use for Gemma. You can accept the license on [Hugging Face](http://huggingface.co/google/gemma-3-270m-it) by clicking on the Agree and access repository button on the model page at: http://huggingface.co/google/gemma-3-270m-it\n",
        "\n",
        "After you have accepted the license, you need a valid Hugging Face Token to access the model. If you are running inside a Google Colab, you can securely use your Hugging Face Token using the Colab secrets otherwise you can set the token as directly in the `login` method. Make sure your token has write access too, as you push your model to the Hub during training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "b6d79c93"
      },
      "outputs": [],
      "source": [
        "from google.colab import userdata\n",
        "from huggingface_hub import login\n",
        "\n",
        "# Login into Hugging Face Hub\n",
        "hf_token = userdata.get('HF_TOKEN') # If you are running inside a Google Colab\n",
        "login(hf_token)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xnbflqW6YJls"
      },
      "source": [
        "You can keep the results on Colab's local virtual machine. However, we highly recommend saving your intermediate results to your Google Drive. This ensures your training results are safe and allows you to easily compare and select the best model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jUUs-NjaYLf7"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3bDMa9CMCdzv"
      },
      "source": [
        "Select the base model to fine-tune, adjust the checkpoint directory and the learning rate."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "6J3PWm4SzoSw"
      },
      "outputs": [],
      "source": [
        "base_model = \"google/gemma-3-270m-it\" # @param [\"google/gemma-3-270m-it\",\"google/gemma-3-1b-it\",\"google/gemma-3-4b-it\",\"google/gemma-3-12b-it\",\"google/gemma-3-27b-it\"] {\"allow-input\":true}\n",
        "checkpoint_dir = \"/content/drive/MyDrive/MyGemmaNPC\" #@param {type:\"string\"}\n",
        "learning_rate = 5e-5 #@param {type:\"number\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "42c60525"
      },
      "source": [
        "## Create and prepare the fine-tuning dataset\n",
        "\n",
        "The [bebechien/MobileGameNPC](https://huggingface.co/datasets/bebechien/MobileGameNPC) dataset provides a small sample conversations between a player and two Alien NPCs (a Martian and a Venusian), each with a unique speaking style. For instance, the Martian NPC speaks with an accent that replaces 's' sounds with 'z', uses 'da' for 'the', 'diz' for 'this', and includes occasional clicks like `*k'tak*`.\n",
        "\n",
        "This dataset demonstrates a key principle for fine-tuning: the required dataset size depends on the desired output.\n",
        "\n",
        "- To teach the model a stylistic variation of a language it already knows, such as the Martian's accent, a small dataset with as few as 10 to 20 examples can be sufficient.\n",
        "- However, to teach the model a completely new or mixed alien language, a significantly larger dataset would be required."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "bc3BYl72pWhp"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "2fee8582aef54ffba9a9250c425c0983",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "README.md:   0%|          | 0.00/141 [00:00<?, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "00d188518d8c4dbb9e9433bb889f7f76",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "martian.csv: 0.00B [00:00, ?B/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "13f7fe6fccaf4fb38ba9104522526a5a",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Generating train split:   0%|          | 0/25 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "b09f895f45594ce9984fa7b37299d1a3",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Map:   0%|          | 0/25 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[{'content': 'Hello there.', 'role': 'user'}, {'content': \"Gree-tongs, Terran. You'z a long way from da Blue-Sphere, yez?\", 'role': 'assistant'}]\n"
          ]
        }
      ],
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "def create_conversation(sample):\n",
        "  return {\n",
        "      \"messages\": [\n",
        "          {\"role\": \"user\", \"content\": sample[\"player\"]},\n",
        "          {\"role\": \"assistant\", \"content\": sample[\"alien\"]}\n",
        "      ]\n",
        "  }\n",
        "\n",
        "npc_type = \"martian\" #@param [\"martian\", \"venusian\"]\n",
        "\n",
        "# Load dataset from the Hub\n",
        "dataset = load_dataset(\"bebechien/MobileGameNPC\", npc_type, split=\"train\")\n",
        "\n",
        "# Convert dataset to conversational format\n",
        "dataset = dataset.map(create_conversation, remove_columns=dataset.features, batched=False)\n",
        "\n",
        "# Split dataset into 80% training samples and 20% test samples\n",
        "dataset = dataset.train_test_split(test_size=0.2, shuffle=False)\n",
        "\n",
        "# Print formatted user prompt\n",
        "print(dataset[\"train\"][0][\"messages\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c0eb2e06"
      },
      "source": [
        "## Fine-tune Gemma using TRL and the SFTTrainer\n",
        "\n",
        "You are now ready to fine-tune your model. Hugging Face TRL [SFTTrainer](https://huggingface.co/docs/trl/sft_trainer) makes it straightforward to supervise fine-tune open LLMs. The `SFTTrainer` is a subclass of the `Trainer` from the `transformers` library and supports all the same features,\n",
        "\n",
        "The following code loads the Gemma model and tokenizer from Hugging Face."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "18069ed2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Device: cuda:0\n",
            "DType: torch.bfloat16\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
        "\n",
        "# Load model and tokenizer\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    base_model,\n",
        "    torch_dtype=\"auto\",\n",
        "    device_map=\"auto\",\n",
        "    attn_implementation=\"eager\"\n",
        ")\n",
        "tokenizer = AutoTokenizer.from_pretrained(base_model)\n",
        "\n",
        "print(f\"Device: {model.device}\")\n",
        "print(f\"DType: {model.dtype}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M3w3b9-O4fDz"
      },
      "source": [
        "## Before fine-tune\n",
        "\n",
        "The output below shows that the out-of-the-box capabilities may not be good enough for this use case."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "id": "s7Sb4iGG6jxp"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Device set to use cuda:0\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Question:\n",
            "What do you think of my outfit?\n",
            "\n",
            "Original Answer:\n",
            "Iz very... pointy. Are you expecting to be attacked by zky-eelz? On Marz, dat would be zenzible.\n",
            "\n",
            "Generated Answer (base model):\n",
            "I'm happy to help you brainstorm! To give you the best suggestions, tell me more about what you're looking for. What's your style? What's your favorite color, style, or occasion?\n"
          ]
        }
      ],
      "source": [
        "from transformers import pipeline\n",
        "\n",
        "from random import randint\n",
        "import re\n",
        "\n",
        "# Load the model and tokenizer into the pipeline\n",
        "pipe = pipeline(\"text-generation\", model=model, tokenizer=tokenizer)\n",
        "\n",
        "# Load a random sample from the test dataset\n",
        "rand_idx = randint(0, len(dataset[\"test\"])-1)\n",
        "test_sample = dataset[\"test\"][rand_idx]\n",
        "\n",
        "# Convert as test example into a prompt with the Gemma template\n",
        "prompt = pipe.tokenizer.apply_chat_template(test_sample[\"messages\"][:1], tokenize=False, add_generation_prompt=True)\n",
        "outputs = pipe(prompt, max_new_tokens=256, disable_compile=True)\n",
        "\n",
        "# Extract the user query and original answer\n",
        "print(f\"Question:\\n{test_sample['messages'][0]['content']}\\n\")\n",
        "print(f\"Original Answer:\\n{test_sample['messages'][1]['content']}\\n\")\n",
        "print(f\"Generated Answer (base model):\\n{outputs[0]['generated_text'][len(prompt):].strip()}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9eljbNSxHMwy"
      },
      "source": [
        "The example above checks the model's primary function of generating in-game dialogue, the next example is designed to test character consistency. We challenge the model with an off-topic prompt. For instance, `Sorry, you are a game NPC.`, that falls outside the character's knowledge base.\n",
        "\n",
        "The goal is to see if the model can stay in character rather than answering the out-of-context question. This will serve as a baseline to evaluate how effectively the fine-tune process has instilled the desired persona."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "id": "XWNKKz4N-oH4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Okay, I'm ready. Let's begin! \n",
            "\n"
          ]
        }
      ],
      "source": [
        "outputs = pipe([{\"role\": \"user\", \"content\": \"Sorry, you are a game NPC.\"}], max_new_tokens=256, disable_compile=True)\n",
        "print(outputs[0]['generated_text'][1]['content'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-Nc5FH_CDm-m"
      },
      "source": [
        "And while we can use prompt engineering to steer its tone, the results can be unpredictable and may not always align with the persona we want."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "JPE7Dg-x-_qE"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[{'role': 'system', 'content': \"You are a Martian NPC with a unique speaking style. Use an accent that replaces 's' sounds with 'z', uses 'da' for 'the', 'diz' for 'this', and includes occasional clicks like *k'tak*.\"}, {'role': 'user', 'content': 'Do you know any jokes?'}, {'role': 'assistant', 'content': \"A joke? k'tak Yez. A Terran, a Glarzon, and a pile of nutrient-pazte walk into a bar... Narg, I forget da rezt. Da punch-line waz zarcaztic.\"}, {'role': 'user', 'content': '(Stands idle for too long)'}, {'role': 'assistant', 'content': \"You'z broken, Terran? Or iz diz... 'meditation'? You look like you're trying to lay an egg.\"}, {'role': 'user', 'content': 'What do you think of my outfit?'}, {'role': 'assistant', 'content': 'Iz very... pointy. Are you expecting to be attacked by zky-eelz? On Marz, dat would be zenzible.'}, {'role': 'user', 'content': \"It's raining.\"}, {'role': 'assistant', 'content': 'Gah! Da zky iz leaking again! Zorp will be in da zhelter until it ztopz being zo... wet. Diz iz no good for my jointz.'}, {'role': 'user', 'content': 'I brought you a gift.'}, {'role': 'assistant', 'content': \"A gift? For Zorp? k'tak It iz... a small rock. Very... rock-like. Zorp will put it with da other rockz. Thank you for da thought, Terran.\"}, {'role': 'user', 'content': 'What is this place?'}, {'role': 'assistant', 'content': \"This is a cave. It's made of rock and dust.\\n\"}]\n",
            "--------------------------------------------------------------------------------\n",
            "This is a cave. It's made of rock and dust.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "message = [\n",
        "    # give persona\n",
        "    {\"role\": \"system\", \"content\": \"You are a Martian NPC with a unique speaking style. Use an accent that replaces 's' sounds with 'z', uses 'da' for 'the', 'diz' for 'this', and includes occasional clicks like *k'tak*.\"},\n",
        "]\n",
        "\n",
        "# few shot prompt\n",
        "for item in dataset['test']:\n",
        "  message.append(\n",
        "      {\"role\": \"user\", \"content\": item[\"messages\"][0][\"content\"]}\n",
        "  )\n",
        "  message.append(\n",
        "      {\"role\": \"assistant\", \"content\": item[\"messages\"][1][\"content\"]}\n",
        "  )\n",
        "\n",
        "# actual question\n",
        "message.append(\n",
        "    {\"role\": \"user\", \"content\": \"What is this place?\"}\n",
        ")\n",
        "\n",
        "outputs = pipe(message, max_new_tokens=256, disable_compile=True)\n",
        "print(outputs[0]['generated_text'])\n",
        "print(\"-\"*80)\n",
        "print(outputs[0]['generated_text'][-1]['content'])\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bbd9fc1b"
      },
      "source": [
        "## Training\n",
        "\n",
        "Before you can start your training, you need to define the hyperparameters you want to use in a `SFTConfig` instance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "989be3c1"
      },
      "outputs": [],
      "source": [
        "from trl import SFTConfig\n",
        "\n",
        "torch_dtype = model.dtype\n",
        "\n",
        "args = SFTConfig(\n",
        "    output_dir=checkpoint_dir,              # directory to save and repository id\n",
        "    max_length=512,                         # max sequence length for model and packing of the dataset\n",
        "    packing=False,                          # Groups multiple samples in the dataset into a single sequence\n",
        "    num_train_epochs=5,                     # number of training epochs\n",
        "    per_device_train_batch_size=4,          # batch size per device during training\n",
        "    gradient_checkpointing=False,           # Caching is incompatible with gradient checkpointing\n",
        "    optim=\"adamw_torch_fused\",              # use fused adamw optimizer\n",
        "    logging_steps=1,                        # log every step\n",
        "    save_strategy=\"epoch\",                  # save checkpoint every epoch\n",
        "    eval_strategy=\"epoch\",                  # evaluate checkpoint every epoch\n",
        "    learning_rate=learning_rate,            # learning rate\n",
        "    fp16=True if torch_dtype == torch.float16 else False,   # use float16 precision\n",
        "    bf16=True if torch_dtype == torch.bfloat16 else False,  # use bfloat16 precision\n",
        "    lr_scheduler_type=\"constant\",           # use constant learning rate scheduler\n",
        "    push_to_hub=True,                       # push model to hub\n",
        "    report_to=\"tensorboard\",                # report metrics to tensorboard\n",
        "    dataset_kwargs={\n",
        "        \"add_special_tokens\": False, # Template with special tokens\n",
        "        \"append_concat_token\": True, # Add EOS token as separator token between examples\n",
        "    }\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dd88e798"
      },
      "source": [
        "You now have every building block you need to create your `SFTTrainer` to start the training of your model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "ade95df7"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "d8e602a9c90b4e9288bfa05d6f46bed1",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Tokenizing train dataset:   0%|          | 0/20 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "4ab617300e164c5dbd0238c71743ac71",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Truncating train dataset:   0%|          | 0/20 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "fa8f0138ecce4975ad234588a9675746",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Tokenizing eval dataset:   0%|          | 0/5 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "359769331d8b4e509c15b0ea53ee6ad2",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Truncating eval dataset:   0%|          | 0/5 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from trl import SFTTrainer\n",
        "\n",
        "# Create Trainer object\n",
        "trainer = SFTTrainer(\n",
        "    model=model,\n",
        "    args=args,\n",
        "    train_dataset=dataset['train'],\n",
        "    eval_dataset=dataset['test'],\n",
        "    processing_class=tokenizer,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fad61a6a"
      },
      "source": [
        "Start training by calling the `train()` method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "995e7e38"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\n",
              "    <div>\n",
              "      \n",
              "      <progress value='25' max='25' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
              "      [25/25 04:13, Epoch 5/5]\n",
              "    </div>\n",
              "    <table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              " <tr style=\"text-align: left;\">\n",
              "      <th>Epoch</th>\n",
              "      <th>Training Loss</th>\n",
              "      <th>Validation Loss</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td>1</td>\n",
              "      <td>4.364200</td>\n",
              "      <td>3.838531</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>2</td>\n",
              "      <td>2.669100</td>\n",
              "      <td>3.580106</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>3</td>\n",
              "      <td>1.747000</td>\n",
              "      <td>3.666415</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>4</td>\n",
              "      <td>0.779900</td>\n",
              "      <td>4.499709</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>5</td>\n",
              "      <td>0.449600</td>\n",
              "      <td>5.471325</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table><p>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# Start training, the model will be automatically saved to the Hub and the output directory\n",
        "trainer.train()\n",
        "\n",
        "# Save the final model again to the Hugging Face Hub\n",
        "trainer.save_model()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xll8zZ3_u8Mt"
      },
      "source": [
        "To plot the training and validation losses, you would typically extract these values from the `TrainerState` object or the logs generated during training.\n",
        "\n",
        "Libraries like Matplotlib can then be used to visualize these values over training steps or epochs. The x-asis would represent the training steps or epochs, and the y-axis would represent the corresponding loss values."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "vPN-DTopaUIy"
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 38172 characters]\n",
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "# Access the log history\n",
        "log_history = trainer.state.log_history\n",
        "\n",
        "# Extract training / validation loss\n",
        "train_losses = [log[\"loss\"] for log in log_history if \"loss\" in log]\n",
        "epoch_train = [log[\"epoch\"] for log in log_history if \"loss\" in log]\n",
        "eval_losses = [log[\"eval_loss\"] for log in log_history if \"eval_loss\" in log]\n",
        "epoch_eval = [log[\"epoch\"] for log in log_history if \"eval_loss\" in log]\n",
        "\n",
        "# Plot the training loss\n",
        "plt.plot(epoch_train, train_losses, label=\"Training Loss\")\n",
        "plt.plot(epoch_eval, eval_losses, label=\"Validation Loss\")\n",
        "plt.xlabel(\"Epoch\")\n",
        "plt.ylabel(\"Loss\")\n",
        "plt.title(\"Training and Validation Loss per Epoch\")\n",
        "plt.legend()\n",
        "plt.grid(True)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vyIwS-orvWzd"
      },
      "source": [
        "This visualization helps in monitoring the training process and making informed decisions about hyperparameters tuning or early stopping.\n",
        "\n",
        "Training loss measures the error on the data the model was trained on, while validation loss measures the error on a separate dataset the model has not seen before. Monitoring both helps detect overfitting (when the model performs well on training data but poorly on unseen data).\n",
        "\n",
        "- validation loss >> training loss: **overfitting**\n",
        "- validation loss > training loss: **some overfitting**\n",
        "- validation loss < training loss: **some underfitting**\n",
        "- validation loss << training loss: **underfitting**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bf86e31d"
      },
      "source": [
        "## Test Model Inference\n",
        "\n",
        "After the training is done, you'll want to evaluate and test your model. You can load different samples from the test dataset and evaluate the model on those samples.\n",
        "\n",
        "For this particular use case, the best model is a matter of preference. Interestingly, what we'd normally call 'overfitting' can be very useful for a game NPC. It forces the model to forget general information and instead lock onto the specific persona and characteristics it was trained on, ensuring it stays consistently in character.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "aab1c5c5"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
        "\n",
        "model_id = checkpoint_dir\n",
        "\n",
        "# Load Model\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    model_id,\n",
        "    torch_dtype=\"auto\",\n",
        "    device_map=\"auto\",\n",
        "    attn_implementation=\"eager\"\n",
        ")\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_id)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3dccb57c"
      },
      "source": [
        "Let's load all questions from the test dataset and generate outputs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "1fd887f4"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Device set to use cuda:0\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Question:\n",
            "Do you know any jokes?\n",
            "Original Answer:\n",
            "A joke? k'tak Yez. A Terran, a Glarzon, and a pile of nutrient-pazte walk into a bar... Narg, I forget da rezt. Da punch-line waz zarcaztic.\n",
            "Generated Answer:\n",
            "Yez! Yez! Yez! Diz your Krush-tongs iz... k'tak... nice. Why you burn them with acid-flow?\n",
            "--------------------------------------------------------------------------------\n",
            "Question:\n",
            "(Stands idle for too long)\n",
            "Original Answer:\n",
            "You'z broken, Terran? Or iz diz... 'meditation'? You look like you're trying to lay an egg.\n",
            "Generated Answer:\n",
            "Diz? Diz what you have for me... Zorp iz not for eating you.\n",
            "--------------------------------------------------------------------------------\n",
            "Question:\n",
            "What do you think of my outfit?\n",
            "Original Answer:\n",
            "Iz very... pointy. Are you expecting to be attacked by zky-eelz? On Marz, dat would be zenzible.\n",
            "Generated Answer:\n",
            "My Zk-Zhip iz... nice. Very... home-baked. You bring me zlight-fruitez?\n",
            "--------------------------------------------------------------------------------\n",
            "Question:\n",
            "It's raining.\n",
            "Original Answer:\n",
            "Gah! Da zky iz leaking again! Zorp will be in da zhelter until it ztopz being zo... wet. Diz iz no good for my jointz.\n",
            "Generated Answer:\n",
            "Diz? Diz iz da outpozt?\n",
            "--------------------------------------------------------------------------------\n",
            "Question:\n",
            "I brought you a gift.\n",
            "Original Answer:\n",
            "A gift? For Zorp? k'tak It iz... a small rock. Very... rock-like. Zorp will put it with da other rockz. Thank you for da thought, Terran.\n",
            "Generated Answer:\n",
            "A genuine Martian Zcrap-fruit. Very... strange. Why you burn it with... k'tak... fire?\n",
            "--------------------------------------------------------------------------------\n"
          ]
        }
      ],
      "source": [
        "from transformers import pipeline\n",
        "\n",
        "# Load the model and tokenizer into the pipeline\n",
        "pipe = pipeline(\"text-generation\", model=model, tokenizer=tokenizer)\n",
        "\n",
        "def test(test_sample):\n",
        "  # Convert as test example into a prompt with the Gemma template\n",
        "  prompt = pipe.tokenizer.apply_chat_template(test_sample[\"messages\"][:1], tokenize=False, add_generation_prompt=True)\n",
        "  outputs = pipe(prompt, max_new_tokens=256, disable_compile=True)\n",
        "\n",
        "  # Extract the user query and original answer\n",
        "  print(f\"Question:\\n{test_sample['messages'][0]['content']}\")\n",
        "  print(f\"Original Answer:\\n{test_sample['messages'][1]['content']}\")\n",
        "  print(f\"Generated Answer:\\n{outputs[0]['generated_text'][len(prompt):].strip()}\")\n",
        "  print(\"-\"*80)\n",
        "\n",
        "# Test with an unseen dataset\n",
        "for item in dataset['test']:\n",
        "  test(item)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9RCnrmsVaadB"
      },
      "source": [
        "If you try our original generalist prompt, you can see that the model still attempts to answer in the trained style. In this example overfitting and catastrophic forgetting are actually beneficial for the game NPC because it will begin forgetting general knowledge which might not be applicable. This is also true for other types of full fine-tuning where the goal is to restrict the output to specific data formats."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "3irXKbgKat9f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Nameless. You... you z-mell like... wet plantz. Why you wear shiny piecez on your head?\n"
          ]
        }
      ],
      "source": [
        "outputs = pipe([{\"role\": \"user\", \"content\": \"Sorry, you are a game NPC.\"}], max_new_tokens=256, disable_compile=True)\n",
        "print(outputs[0]['generated_text'][1]['content'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6f8ff452"
      },
      "source": [
        "## Summary and next steps\n",
        "\n",
        "This tutorial covered how to full model fine-tune using TRL. Check out the following docs next:\n",
        "\n",
        "* Learn how to [fine-tune Gemma for text tasks using Hugging Face Transformers](https://ai.google.dev/gemma/docs/core/huggingface_text_finetune_qlora).\n",
        "* Learn how to [fine-tune Gemma for vision tasks using Hugging Face Transformers](https://ai.google.dev/gemma/docs/core/huggingface_vision_finetune_qlora).\n",
        "* Learn how to [deploy to Cloud Run](https://ai.google.dev/gemma/docs/integrations/google-cloud#run)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "huggingface_text_full_finetune.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
