{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "KWfW4Mlu_0bv" }, "source": [ "# How to Fine-Tune a Decision Model with TRL (SFT + LoRA)\n", "\n", "\n", "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Liquid4All/cookbook/blob/main/finetuning/notebooks/sft_for_decision_models_with_trl.ipynb)\n", "\n", "This tutorial shows you how to fine-tune a decision model, a model that makes decisions instead of generating text, on your own labels.\n", "\n", "We fine-tune our [`d1-3B`](https://huggingface.co/LiquidAI/d1-3B) decision model with Supervised Fine-tuning (SFT) and LoRA using the TRL library, with intent classification for customer support routing on Banking77 as the example.\n", "\n", "Follow along if it's your first time using trl, or take single code snippets for your own workflow\n", "\n", "## ๐ŸŽฏ What You'll Find:\n", "- **LoRA + SFT** - Fine-tune the d1-3B decision model on your own decisions using LoRA (from PEFT)\n", "- **Evaluation** - Compare the decision model's accuracy before and after fine-tuning\n", "\n", "## ๐Ÿ“‹ Prerequisites:\n", "- **GPU Runtime**: Select a GPU in `Runtime` โ†’ `Change runtime type`. A free T4 GPU is enough (about 35 minutes of training)\n", "- **Hugging Face Account**: Only needed to push your model to the Hub" ] }, { "cell_type": "markdown", "metadata": { "id": "moZoi4B2_0by" }, "source": [ "# ๐Ÿ“ฆ Installation & Setup\n", "\n", "First, let's install all the required packages:" ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "id": "TzGr4YRj_0by" }, "outputs": [], "source": [ "!pip install -qqq \"transformers==5.19.0\" \"trl==1.15.0\" \"peft==0.21.2\" \"torchao>0.16.0\" datasets --progress-bar off" ] }, { "cell_type": "markdown", "metadata": { "id": "NmWNMRxA_0by" }, "source": [ "Let's now verify the packages are installed correctly" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "tjmdRT-Q_0by", "outputId": "83905727-7333-45cd-c6da-41716352ec29" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "๐Ÿ“ฆ PyTorch version: 2.11.0+cu130\n", "๐Ÿค— Transformers version: 5.19.0\n", "๐Ÿ“Š TRL version: 1.15.0\n", "๐ŸŽ›๏ธ PEFT version: 0.21.2\n", "๐Ÿ”ข Precision: torch.bfloat16\n" ] } ], "source": [ "import torch\n", "import transformers\n", "import trl\n", "import peft\n", "import os\n", "os.environ[\"WANDB_DISABLED\"] = \"true\"\n", "\n", "print(f\"๐Ÿ“ฆ PyTorch version: {torch.__version__}\")\n", "print(f\"๐Ÿค— Transformers version: {transformers.__version__}\")\n", "print(f\"๐Ÿ“Š TRL version: {trl.__version__}\")\n", "print(f\"๐ŸŽ›๏ธ PEFT version: {peft.__version__}\")\n", "\n", "# bfloat16 on GPUs that support it natively (e.g. A100), float16 otherwise (e.g. T4)\n", "dtype = torch.bfloat16 if torch.cuda.is_bf16_supported(including_emulation=False) else torch.float16\n", "print(f\"๐Ÿ”ข Precision: {dtype}\")" ] }, { "cell_type": "markdown", "metadata": { "id": "yW05aaci_0bz" }, "source": [ "# Loading the model from Transformers ๐Ÿค—\n", "\n", "d1-3B ships its own code for the System One API (`system_one` and `system_one_batch`), so we load it with `trust_remote_code=True`.\n", "\n", "We also download the model repository locally. It contains `prompt.py`, the module d1-3B uses to turn a state and a question into a prompt. We will reuse it to build our training data, so that the model sees the exact same prompt during training and inference." ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 488, "referenced_widgets": [ "08f34c371ef84e0596c6a943f2ca74b7", "f15551b3ef9f4e4eaabae203fd0a655d", "6dea5cec77d9471d8f2d166406b10091", "30f4273b147b48a6ad3a6775b8bd8fbe", "bff8e7324a4a47679c2e723d192c7bfd", "4385c543fa5247e490afe36ee7065352", "5b8b8eb517f4437b9d1a1160fa2a4976", "63a0ee20a3424504b205b7c0ad6783e8", "af5d481810d64c2787753c32b2988678", "d4113b0eb8e647929ad15eb669d4d029", "4ab1bab2f37d458db480d737afac44b9", "2e8f2267344c4ff0a9927cc6b368ba3f", "d8d98136039a4138b9dbff3a2c46a4cc", "daed99ccf1d046e3b65b1130e939f877", "e3f8990f28b5490487449f726244a04c", "d0350d9b7d37448aa23e5130a53d12ff", "b4d10318571d4ae8994a0d42fbe53ae8", "3882477071e84914be04f63b54dd4263", "b01fd91451cc4e669557ee421abcf1be", "f0da9067065942e2a7274937a311964c", "7ca485b91e70478cb7e00bc9762cb387", "406a8e07ab4742a3a9298878d1dd5c25", "1f545f27530f407da514ac1a4d6d03e6", "d1a0bfacf3034e8ea80daacdfde9020b", "f625b14f330b4f67b9a2dcecb6738fbd", "23229cd9ff2b46e0a8ce06129696797b", "e8aac19cdeed47afbe98bc87f976042c", "0b781ea88a64436ab520f37f5f752f25", "55b1981f1b0a4847aeae303ab6a7aa86", "5e24e449a4204d3bb342aebeca3e311e", "c5f598a3da284a548d9b7ac1bb5a8243", "65e995ab3d77490e9a453a0477ef0be9", "6f828335145745d3815739836788de6a", "ee4e3bec60214d6fb32c99d0ba93091a", "a20b20b765164dc4b3bdf3b852c6588d", "8f9d32ecb3aa4383a995236ac2c9f62a", "3be5b3982a0049d3a1f99c1e8a026852", "e878e42022a242f298e08dfb438b3458", "4fff2e0480914628ad3e5c489042d94d", "cbd8cd8dafe04478b0476a160e3d65fe", "a5bc899d8d7b4d3c8365d3d748e7c0a3", "38a0737792f147438f3fc55f911aa31f", "dc661738f08d4b1f9188035bfaaefc5c", "cb54a415436745afb3f163850e0c78c0" ] }, "id": "Bnwq3kAI_0bz", "outputId": "e4c34804-4e09-4532-e4e8-5b4143139326" }, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "08f34c371ef84e0596c6a943f2ca74b7", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Downloading bytes: | 0.00B " ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "2e8f2267344c4ff0a9927cc6b368ba3f", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Reconstructing (incomplete total...): | | 0.00B / 0.00B " ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "1f545f27530f407da514ac1a4d6d03e6", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Fetching 15 files: 0%| | 0/15 [00:00<|im_start|>user\n", "I was charged twice this month, please refund one of them.\n", "\n", "\n", "QUESTION:\n", "Which team should handle this?\n", "\n", "Options:\n", "A Charges, refunds, invoices\n", "B App or site faults\n", "C Suspected unauthorised use\n", "\n", "Reply with the option code only.<|im_end|>\n", "<|im_start|>assistant\n", "\n" ] } ], "source": [ "from prompt import aliases, as_question, render\n", "\n", "rendered_prompt = render(\n", " tokenizer=tokenizer,\n", " state=state,\n", " q=as_question(questions[\"team\"]),\n", " bos=tokenizer.bos_token,\n", ")\n", "\n", "print(rendered_prompt)" ] }, { "cell_type": "markdown", "metadata": { "id": "TUM-SVJE_0b0" }, "source": [ "The model then runs one forward pass and reads the next-token probabilities at the end of the prompt, the **answer slot**. It only looks at the tokens of the option codes and normalizes them into the probabilities you saw above.\n", "\n", "This means that fine-tuning a decision model is regular SFT where **the completion is a single token: the code of the correct option**, followed by the end-of-turn token `<|im_end|>`." ] }, { "cell_type": "markdown", "metadata": { "id": "wye2m15NRrRL" }, "source": [ "# ๐ŸŽฏ Supervised Fine-Tuning (SFT + LoRA)\n", "\n", "SFT teaches the model to make your decisions by training on input-output pairs. For a decision model, the input is the rendered prompt and the output is the code of the correct option.\n", "\n", "We use LoRA (Low-Rank Adaptation) to finetune the model by only training a small number of additional parameters." ] }, { "cell_type": "markdown", "metadata": { "id": "-RlDQ8N0_0b0" }, "source": [ "## Load an SFT Dataset\n", "\n", "We will use [Banking77](https://huggingface.co/datasets/mteb/banking77), a dataset of ~13k online banking customer messages, each labeled with one of 77 fine-grained intents. To keep this notebook fast on a free T4 GPU, we train on 3,000 messages and evaluate on 100. Feel free to increase these numbers with `select(range(...))`." ] }, { "cell_type": "code", "execution_count": 6, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 363, "referenced_widgets": [ "790ff24057d445a39f3e3106238ea409", "5ad214ab224545b3b55e17df681572ec", "e6e6f863669c440d96d5c04f4bb80abd", "35657a7a052047b0a782fdda168a04a8", "277f4f561070488f8125c850a8a75457", "5ed1585a65954ef1b3d5af89df315e9f", "f635fb0c36904d7eaafd360b969a14df", "2cf1e5485a8e4d87b5d7d4b39bf6b2ff", "3963c28eae354ec69145c5e0cc3b31ac", "dec33beb976a4eaabe10040d34eef41b", "6ca05b0dbdec41cb8518d8ccc47be495", "c4b2269eb42a40e69e9cbaa82c5ce095", "d56a0f02cd2c4db3998f1fd26c960a7b", "105c5b4107c746c19fe63f29df50cb83", "376549c7c7f347e0a606bea4c3ce8a5a", "5b805d0c26f348ee87d52d6675c2c7be", "4fe4af61245842aa9d901273c9b3f8d5", "4ee1210385434c66bc1143637e053271", "877f4482b13b40fd8d6333ee242960c6", "b0ed28b8adca4993b6254f83d8a7993b", "50fa7fce1c034c8da4d256f5e5f59b26", "c72d4630423044a29f6a93e720a9d87e", "dbcb439a6b644f4ebcd8aee7119615bd", "41b84513a58446938ec73bf270e3faee", "3ec5f9aa07a844e99bd618d96a24e04c", "5a695cca9567465fac5af384b43a54be", "4fd81d3569f94225b49a5250b76e2795", "05727540aaf2498a9b7fefca99993b42", "b43e4a9848c040e0a5f85278e773bcaa", "1b470a6731224348a54e189ad064308b", "d5c1fa7d135a496f90a12cb442542fcf", "705f31a1bcd84aabbafa407b9c11129e", "994e508818c74593961121ebe425ea07", "2195d4e4f7214f22ba102ca9efbad670", "e4db324893ae4286ae60fc236520f62e", "673da49b7e8447e7941d0755eebc6825", "73f15cf59ded4986a93195123211a204", "264c5c191daa4ceaadcf42f83b5e5437", "7907eac57791444da020a7905244000d", "cc2fdbad769940b9ab108925e9faa6f2", "fbb97612f5264d918b9f79b96a65a587", "9759473ee5c041c5bc7e80187eff8b28", "54a6b5a9a344479abd1fa5301ae1c064", "3d9c0fe445b04ecda29691f8e76876af", "acafaf7e0f3f4edc93ee235f12361c75", "d5449d81eb584994865b68dc64896eb9", "4ebe896db25943fdad06f291b914fb71", "88b5d93291ff4db1a760bfef313c7f59", "e1d07b714c3d4bb986f4034681db4c1c", "c72ba1c3b46d40339e6362410445a173", "261391783ba348e4a322199ea7df6992", "5ab2f2637bdd490aa6e2bada167f1f38", "4893a3bf15974f4781ffb1225efa85e4", "ce0c19a5dae5411cb70d8207d274e7ce", "0404818f06d346d18bc765447141e65c", "9ca3ee46bf9f4067bd5d69b801b7b16e", "33f0bfe1b37943e0adf44005bd56c6ef", "6b7bb36230334726b2833b1e63c62d9a", "1679ec968b524a93ae7fb52090a06e1a", "b91dd43248ec45b89746f3643bf3f0ce", "f3b0167a0f5a4889ad30079bceb6ad06", "a3b66abb98e34ebfa336495323d19e7d", "657c04054f7a4126a675bb43f42dbbdd", "57b48f9f7b7a47bab005f198beaee425", "9e0fbc3e766049acbf48a6e510f7680f", "000cdd6f14054139b90403d780bb8794", "065022d26ad649d8aea463a5a6a26d83", "536aaf7a3bda4f81a60f288c701ba4d2", "861df63510794d23be7ecb28e095a839", "77ed46862db2452dba92805d1d9cda10", "279e497667c94db093938079edfd24df", "553afbd6d8f04ed99b4069dc530e1dfd", "03847fa7a7ef4fe29c42caa96dc0c942", "5a19321a9c954768962899edc5e8ea21", "ec586135341a4bb9b7a121cd5f2f91a7", "dc0b2bcdbde3430fa03684b6a3b76148", "eb6d5980b9eb415d91719f691770678c" ] }, "id": "Y2aUXHNm_0b0", "outputId": "27a7e8b2-c246-4f1e-8e2d-c634c23b3f33" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "๐Ÿ“ฅ Loading dataset...\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "790ff24057d445a39f3e3106238ea409", "version_major": 2, "version_minor": 0 }, "text/plain": [ "README.md: 0%| | 0.00/16.6k [00:00= threshold\n", "\n", "def case_ids(texts):\n", " \"\"\"Give messages that reword each other the same case ID, transitively.\"\"\"\n", " word_sets = [words(text) for text in texts]\n", " parent = list(range(len(texts)))\n", "\n", " def find(i):\n", " while parent[i] != i:\n", " i = parent[i]\n", " return i\n", "\n", " for i in range(len(texts)):\n", " for j in range(i + 1, len(texts)):\n", " if same_case(word_sets[i], word_sets[j]):\n", " parent[find(j)] = find(i)\n", " return [find(i) for i in range(len(texts))]\n", "\n", "# 1. Drop training messages that reword a test message\n", "test_words = [words(text) for text in test_dataset[\"text\"]]\n", "leaked = {i for i, text in enumerate(train_dataset[\"text\"]) if any(same_case(words(text), w) for w in test_words)}\n", "train_dataset = train_dataset.select([i for i in range(len(train_dataset)) if i not in leaked])\n", "\n", "# 2. Hold out 10% of the training messages to track the loss, taking whole cases only\n", "cases = case_ids(train_dataset[\"text\"])\n", "case_sizes = Counter(cases)\n", "shuffled_cases = sorted(case_sizes)\n", "random.Random(42).shuffle(shuffled_cases)\n", "eval_cases, n_eval = set(), 0\n", "for case in shuffled_cases:\n", " if n_eval >= 0.1 * len(train_dataset):\n", " break\n", " eval_cases.add(case)\n", " n_eval += case_sizes[case]\n", "\n", "sft_train = train_dataset.select([i for i, case in enumerate(cases) if case not in eval_cases])\n", "sft_eval = train_dataset.select([i for i, case in enumerate(cases) if case in eval_cases])\n", "\n", "print(f\"๐Ÿงน Dropped {len(leaked)} training messages that reword a test message\")\n", "print(f\"๐Ÿ”— {len(cases) - len(case_sizes)} training messages reword another training message\")\n", "print(f\"โœ… No case on both sides: {len(sft_train)} train and {len(sft_eval)} eval messages\")" ] }, { "cell_type": "markdown", "metadata": { "id": "yXo0fh3e_0b0" }, "source": [ "## Define the decision\n", "\n", "We turn the classification task into a single `choice` question. Each intent becomes one option, described by its label name.\n", "\n", "**Important**: The option codes are assigned by position: the first option gets the first code, the second option the second, and so on. If we always trained with the options in the same order, the model could learn which code means which intent without reading the option list. Adding or reordering an intent later would then shift the codes and break the fine-tuned model. We avoid that below by giving every training example its own random subset of the intents, in a random order." ] }, { "cell_type": "code", "execution_count": 8, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "N7-7bCmQ_0b0", "outputId": "d6285358-89a0-48d0-d566-fa543ca3fa10" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "๐Ÿท๏ธ Number of options: 77\n", "๐Ÿ”ค First option codes: {'Refund_not_showing_up': '00', 'activate_my_card': '01', 'age_limit': '02', 'apple_pay_or_google_pay': '03', 'atm_support': '04'}\n" ] } ], "source": [ "labels = sorted(set(dataset[\"train\"][\"label_text\"]))\n", "\n", "intent_question = {\n", " \"type\": \"choice\",\n", " \"instructions\": \"Which intent does this banking customer message express?\",\n", " \"criteria\": {label: label.replace(\"_\", \" \") for label in labels},\n", "}\n", "\n", "# The code d1-3B assigns to each option in this order, e.g. \"card_arrival\" -> \"12\"\n", "codes = {label: code for label, (code, _) in zip(labels, aliases(tokenizer, labels))}\n", "\n", "print(f\"๐Ÿท๏ธ Number of options: {len(labels)}\")\n", "print(f\"๐Ÿ”ค First option codes: {dict(list(codes.items())[:5])}\")" ] }, { "cell_type": "markdown", "metadata": { "id": "CYeMBFH9_0b0" }, "source": [ "With up to 26 options, d1-3B uses the letters `A` to `Z` as codes. With more options, as here, it switches to the two-digit codes `00` to `76`, each of which is a single token. Going from 26 to 27 options therefore changes every code at once, which is another reason to train on shuffled option lists. Our subsets keep at least 27 options, so the codes stay two-digit, as they are when we ask with all 77 intents." ] }, { "cell_type": "markdown", "metadata": { "id": "-He7VzSs_0b0" }, "source": [ "We now transform the dataset to a format compatible with `trl`. We use TRL's **prompt-completion** format:\n", "- `prompt`: the full prompt rendered by d1-3B's own prompt builder, with its own option list: the correct intent plus a random number of the other intents (27 to 77 options in total), in a random order\n", "- `completion`: the code the correct intent gets in that list, followed by `<|im_end|>`\n", "\n", "Because the list changes from example to example, a code says nothing about which intent it stands for: the model has to read the option list to find the right one, whichever intents it contains.\n", "\n", "With a prompt-completion dataset, `SFTTrainer` computes the loss only on the completion, so the model learns to put its probability mass on the right option code." ] }, { "cell_type": "code", "execution_count": 9, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 1000, "referenced_widgets": [ "fcb2e8c77ede4224a965fa5649995660", "c9b18bcee5ec46c48a9b05606550b845", "0cc4d236f51546c19331049719202221", "08cbe200db9b4d7c9cb7ceb74ead6a0e", "cefff3b520a346b7b6c3f25098fea1f2", "dfb21247fcb94b0f80d0da7f54a20bbb", "47420bdaa4b64154a6ff5c85bd1c730e", "35c3f74ffe9c49edbde61ea63f0e0703", "90b3808ea58c4d2f988c0e5eecd50407", "87df7756ca9e4be4ac93ce67eefc6111", "ab7c96a1b1dd4b8789e430397ed8c199", "06ccbc7d9c364d8786882ce597ce02e1", "3feeaad4b9284bb280735c5685e5d36b", "12b22c015b9e4336938ae55080575c5c", "b07e164fc19840608965fda01be666ad", "6ef008d01e5b4164a8a26673a8376b8b", "0bd1f63ceed8404d8dc35e7a06bd827b", "b6f5b31600bd42a597f8d60c2b14870e", "f59e08db3272417593c6fcdf8d488aa7", "c8cc464bf501474197701b475fcc1eca", "1ffb52a141fb4336bde86defe8ac768f", "eb870c2dada44055a5b971513421f7aa" ] }, "id": "Oh7aE-kN_0b0", "outputId": "ff75a384-048c-4278-a817-ecd7f4cf2df2" }, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "fcb2e8c77ede4224a965fa5649995660", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Map: 0%| | 0/2694 [00:00<|im_start|>user\n", "Whats the process of tranfering money into an account?\n", "\n", "\n", "QUESTION:\n", "Which intent does this banking customer message express?\n", "\n", "Options:\n", "00 reverted card payment?\n", "01 card acceptance\n", "02 cash withdrawal charge\n", "03 card swallowed\n", "04 declined cash withdrawal\n", "05 get physical card\n", "06 wrong exchange rate for cash withdrawal\n", "07 request refund\n", "08 wrong amount of cash received\n", "09 automatic top up\n", "10 transaction charged twice\n", "11 apple pay or google pay\n", "12 top up by bank transfer charge\n", "13 edit personal details\n", "14 card payment not recognised\n", "15 card linking\n", "16 verify my identity\n", "17 supported cards and currencies\n", "18 extra charge on statement\n", "19 compromised card\n", "20 age limit\n", "21 terminate account\n", "22 top up by card charge\n", "23 card delivery estimate\n", "24 visa or mastercard\n", "25 direct debit payment not recognised\n", "26 atm support\n", "27 order physical card\n", "28 cash withdrawal not recognised\n", "29 exchange charge\n", "30 balance not updated after bank transfer\n", "31 why verify identity\n", "32 lost or stolen phone\n", "33 disposable card limits\n", "34 get disposable virtual card\n", "35 declined card payment\n", "36 transfer into account\n", "37 verify source of funds\n", "38 unable to verify identity\n", "39 transfer timing\n", "40 cancel transfer\n", "41 card payment wrong exchange rate\n", "42 country support\n", "43 getting spare card\n", "44 card arrival\n", "45 Refund not showing up\n", "46 passcode forgotten\n", "47 pending cash withdrawal\n", "48 top up by cash or cheque\n", "49 transfer not received by recipient\n", "50 card not working\n", "51 pin blocked\n", "\n", "Reply with the option code only.<|im_end|>\n", "<|im_start|>assistant\n", "\n", "\n", "๐ŸŽฏ Completion: '36<|im_end|>'\n" ] } ], "source": [ "def to_prompt_completion(example, idx, split):\n", " # Every example gets its own option list: the correct intent plus 26 to 76 others, in a random order\n", " rng = random.Random(f\"{split}-{idx}\")\n", " others = [label for label in labels if label != example[\"label_text\"]]\n", " options = rng.sample(others, rng.randint(26, len(others))) + [example[\"label_text\"]]\n", " rng.shuffle(options)\n", " question = as_question({**intent_question, \"criteria\": {label: intent_question[\"criteria\"][label] for label in options}})\n", " option_codes = {label: code for label, (code, _) in zip(options, aliases(tokenizer, options))}\n", " return {\n", " \"prompt\": render(tokenizer, example[\"text\"], question, bos=tokenizer.bos_token),\n", " \"completion\": option_codes[example[\"label_text\"]] + tokenizer.eos_token,\n", " }\n", "\n", "# Each example gets one option list, which is enough for one epoch. To train for more epochs, map the dataset\n", "# once per epoch with a different `split` value and concatenate the copies, so every pass sees new lists.\n", "train_dataset_sft = sft_train.map(to_prompt_completion, with_indices=True, fn_kwargs={\"split\": \"train\"}, remove_columns=sft_train.column_names)\n", "eval_dataset_sft = sft_eval.map(to_prompt_completion, with_indices=True, fn_kwargs={\"split\": \"eval\"}, remove_columns=sft_eval.column_names)\n", "\n", "print(\"โœ… SFT Dataset formatted:\")\n", "print(f\" ๐Ÿ“š Train samples: {len(train_dataset_sft)}\")\n", "print(f\" ๐Ÿงช Eval samples: {len(eval_dataset_sft)}\")\n", "print(f\"\\n๐Ÿ“ Prompt:\\n{train_dataset_sft[0]['prompt']}\")\n", "print(f\"\\n๐ŸŽฏ Completion: {train_dataset_sft[0]['completion']!r}\")" ] }, { "cell_type": "markdown", "metadata": { "id": "2851ff7ac170" }, "source": [ "Let's look at two training messages with the same intent. Each prompt has its own option list, so `card arrival` gets a different code in each, and the model has to read the list to answer." ] }, { "cell_type": "code", "execution_count": 10, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "f17deaf926d0", "outputId": "4317b465-1e6f-4441-9e9e-29488ca59ba1" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "๐Ÿ’ฌ Message: Where is the card I ordered 2 weeks ago?\n", "๐Ÿ”ข Number of options: 76\n", "๐Ÿ“‹ First options: ['00 get disposable virtual card', '01 edit personal details', '02 getting virtual card']\n", "๐Ÿ”Ž Correct option: 08 card arrival\n", "๐ŸŽฏ Completion: '08<|im_end|>'\n", "\n", "๐Ÿ’ฌ Message: Can I track when my card will be delivered?\n", "๐Ÿ”ข Number of options: 69\n", "๐Ÿ“‹ First options: ['00 activate my card', '01 getting spare card', '02 getting virtual card']\n", "๐Ÿ”Ž Correct option: 55 card arrival\n", "๐ŸŽฏ Completion: '55<|im_end|>'\n", "\n" ] } ], "source": [ "import re\n", "\n", "# The first two training messages whose intent is card_arrival\n", "for i in [i for i, label in enumerate(sft_train[\"label_text\"]) if label == \"card_arrival\"][:2]:\n", " options = re.findall(r\"^\\d\\d .+$\", train_dataset_sft[i][\"prompt\"], re.M)\n", " print(f\"๐Ÿ’ฌ Message: {sft_train[i]['text']}\")\n", " print(f\"๐Ÿ”ข Number of options: {len(options)}\")\n", " print(f\"๐Ÿ“‹ First options: {options[:3]}\")\n", " print(f\"๐Ÿ”Ž Correct option: {next(option for option in options if option.endswith(' card arrival'))}\")\n", " print(f\"๐ŸŽฏ Completion: {train_dataset_sft[i]['completion']!r}\\n\")" ] }, { "cell_type": "markdown", "metadata": { "id": "oYoj310Y_0b0" }, "source": [ "## Evaluate the base model\n", "\n", "Let's measure how well d1-3B performs on the 100 test messages before fine-tuning. `system_one_batch` packs many messages into each forward pass." ] }, { "cell_type": "code", "execution_count": 11, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 83, "referenced_widgets": [ "d3c579abbe4241309cdfca6aa6144ed0", "345ae6331c5843e188fa61ce6155eceb", "c3d88beb07b14fbcb244cb85784b1564", "e9d1be613c9443e7a960dd4ba0645173", "078b55e068cc46f28637437e2f26bd96", "a322bb64fbcb4566aab6abe9ebad46f8", "a24a84128b734f24963d5b34fce121c7", "587dc47fea5c412b9ec8c8741fb2a142", "74b19df3c27547f4ac949ba14d5dfc69", "5b9201c9fbb043449fb43142d3cc0aee", "5a3c7e1a7e914dfe8dcd30f5ffc0efe3" ] }, "id": "0-uTPjYo_0b0", "outputId": "74471886-ff45-4f55-b805-5af1a4cd55bb" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "๐Ÿ“ Evaluating the base model...\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "d3c579abbe4241309cdfca6aa6144ed0", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Evaluating: 0%| | 0/7 [00:00\n", " \n", " \n", " [169/169 05:23, Epoch 1/1]\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
StepTraining LossValidation LossEntropyMean Token AccuracyNum Tokens
500.1806150.2262800.1785360.934211264665.000000
1000.1493790.1704800.1060430.957237529059.000000
1500.0976760.1499930.1516650.958882793472.000000
1690.1856790.1494330.1524030.958882888013.000000

" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "๐ŸŽ‰ SFT training completed!\n", "๐Ÿ’พ Saving to: ./d1-3B-banking77\n" ] } ], "source": [ "from trl import SFTConfig, SFTTrainer\n", "\n", "sft_config = SFTConfig(\n", " output_dir=\"./d1-3B-banking77\",\n", " num_train_epochs=1,\n", " per_device_train_batch_size=4,\n", " gradient_accumulation_steps=4, # <- effective batch size of 16\n", " gradient_checkpointing=True, # <- trades compute for memory, so it fits on a T4\n", " learning_rate=2e-4,\n", " lr_scheduler_type=\"cosine\",\n", " warmup_steps=0.05, # <- a float is a fraction of the total steps\n", " logging_steps=10,\n", " eval_strategy=\"steps\",\n", " eval_steps=50,\n", " save_strategy=\"no\",\n", " bf16=dtype == torch.bfloat16,\n", " fp16=dtype == torch.float16,\n", " report_to=\"none\",\n", ")\n", "\n", "print(\"๐Ÿ—๏ธ Creating SFT trainer...\")\n", "sft_trainer = SFTTrainer(\n", " model=model,\n", " args=sft_config,\n", " train_dataset=train_dataset_sft,\n", " eval_dataset=eval_dataset_sft,\n", " processing_class=tokenizer,\n", ")\n", "\n", "print(\"\\n๐Ÿš€ Starting SFT training...\")\n", "sft_trainer.train()\n", "\n", "print(\"๐ŸŽ‰ SFT training completed!\")\n", "\n", "sft_trainer.save_model()\n", "print(f\"๐Ÿ’พ Saving to: {sft_config.output_dir}\")" ] }, { "cell_type": "markdown", "metadata": { "id": "jWCedyUn_0b2" }, "source": [ "## Save merged model\n", "\n", "Merge the extra weights learned with LoRA back into the model to obtain a \"normal\" model checkpoint.\n", "\n", "The merged checkpoint contains only the weights and the tokenizer. To keep the System One API, we also copy d1-3B's own code and processor configuration into the same folder." ] }, { "cell_type": "code", "execution_count": 16, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 83, "referenced_widgets": [ "e07f27f1459d44c498aa53b0cf543c8e", "607e48de845946b880e8c85d23ca142a", "c6cfc548bda84d398435cc034a5e0561", "5730d4e8ba6f41f289d53c5a1f800ccd", "daceaf59087448c8bc3005bfe8eca5a9", "5bb1e836d8dc477c8a4cb97cd34408ee", "b3305e364d424fe7858c936418b86896", "a36fa90db049458eada3a54745869e20", "3a7a353009cd4a29ba961d90bcafe32e", "ab7c5c71b4e64885a8f46c99f19c86e7", "2709122d000a4ebdb450f30b445989eb" ] }, "id": "1pVJTptv_0b2", "outputId": "465853b9-6099-4bd0-d456-1944b0215148" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "๐Ÿ”„ Merging LoRA weights...\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "e07f27f1459d44c498aa53b0cf543c8e", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Writing model shards: 0%| | 0/1 [00:00