{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "46ac2ddb",
      "metadata": {
        "id": "46ac2ddb"
      },
      "source": [
        "# Training a Phone Assistant with Reinforcement Learning: LFM3-1.7B + GRPO with TRL\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/lfm3_grpo_with_trl.ipynb) [![GitHub](https://img.shields.io/badge/GitHub-Liquid4All%2Fcookbook-181717?logo=github)](https://github.com/Liquid4All/cookbook) [![Discord](https://img.shields.io/discord/1385439864920739850?color=7289da&label=Discord&logo=discord&logoColor=white)](https://discord.com/invite/liquid-ai)\n",
        "\n",
        "Imagine an assistant app that runs [LFM3-1.7B](https://huggingface.co/LiquidAI/LFM3-1.7B) on the phone and can use 108 tools: calendar, alarms, messages, email, maps, music, settings, smart home, payments, and more. For every request, the model has to pick the right tools among many lookalikes (a text message or a chat message? a timer or a sleep timer?), fill their arguments correctly, ask for missing details, and answer directly when no tool applies.\n",
        "\n",
        "This notebook trains it with reinforcement learning, using Group Relative Policy Optimization (GRPO) from [TRL](https://github.com/huggingface/trl) and LoRA. There is no supervised warm start: the model learns from a reward that compares its tool calls with verified gold calls. Generation dominates the cost of RL, so we recommend a Colab A100 GPU.\n",
        "\n",
        "## 🎯 What you'll find\n",
        "\n",
        "- **The app**: a fixed catalog of 108 tool schemas, in every prompt\n",
        "- **A verifiable reward**: the model's calls are parsed and compared argument by argument with the gold calls\n",
        "- **GRPO + LoRA**: the model learns from groups of its own attempts\n",
        "- **Evaluation**: exact-match accuracy on requests that use the training tools, and on 10 tools that never appear in training"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "460c4ed8",
      "metadata": {
        "id": "460c4ed8"
      },
      "source": [
        "# 📦 Installation & Setup\n",
        "\n",
        "We pin every library to the versions this notebook was tested with."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c5284123",
      "metadata": {
        "tags": [
          "install"
        ],
        "id": "c5284123"
      },
      "outputs": [],
      "source": [
        "!pip install -qqq \"torch==2.14.1\" \"torchvision==0.29.1\" \"torchao==0.18.0\" \"trl==1.16.0\" \"transformers==5.20.2\" \"peft==0.21.2\" \"accelerate==1.15.0\" \"datasets==5.1.0\" \"huggingface-hub==2.3.0\" \"tokenizers==0.23.3\" \"safetensors==0.8.0\" --progress-bar off"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "45d140bd",
      "metadata": {
        "id": "45d140bd"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import os\n",
        "import torch\n",
        "import transformers\n",
        "import trl\n",
        "import peft\n",
        "\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",
        "# LFM3-1.7B is released in bfloat16. GPUs without bfloat16 support load it in float16 instead:\n",
        "# the conversion is exact for almost every weight, and the LoRA adapters train in float32 either way.\n",
        "use_bf16 = torch.cuda.is_bf16_supported(including_emulation=False)\n",
        "dtype = torch.bfloat16 if use_bf16 else torch.float16\n",
        "print(f\"🔢 Precision: {dtype}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9251021f",
      "metadata": {
        "id": "9251021f"
      },
      "source": [
        "# Loading the model from Transformers 🤗\n",
        "\n",
        "LFM3-1.7B ships its modeling code with the checkpoint, so we load it with `trust_remote_code=True`. We pad on the left, so that batched generation continues each prompt right where it ends."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "01df6489",
      "metadata": {
        "id": "01df6489"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "\n",
        "model_id = \"LiquidAI/LFM3-1.7B\"\n",
        "\n",
        "print(\"📚 Loading tokenizer...\")\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True, padding_side=\"left\")\n",
        "\n",
        "print(\"🧠 Loading model...\")\n",
        "model = AutoModelForCausalLM.from_pretrained(model_id, dtype=dtype, device_map=\"auto\", trust_remote_code=True)\n",
        "print(\"✅ Model loaded successfully!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "674c8064",
      "metadata": {
        "id": "674c8064"
      },
      "source": [
        "# 📱 The app: 108 tools\n",
        "\n",
        "The dataset comes with the tool catalog (`tools.json`, standard function schemas) and, for the reward, the way each argument is compared with the gold value (`scoring.json`)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cfaccda6",
      "metadata": {
        "id": "cfaccda6"
      },
      "outputs": [],
      "source": [
        "from collections import Counter\n",
        "from datasets import load_dataset\n",
        "from huggingface_hub import hf_hub_download\n",
        "\n",
        "dataset_id = \"LiquidAI/phone-assistant-tools\"\n",
        "tools = json.load(open(hf_hub_download(dataset_id, \"tools.json\", repo_type=\"dataset\")))\n",
        "kinds = json.load(open(hf_hub_download(dataset_id, \"scoring.json\", repo_type=\"dataset\")))\n",
        "\n",
        "print(f\"🧰 {len(tools)} tools, for example:\")\n",
        "for tool in tools[:3]:\n",
        "    print(json.dumps(tool[\"function\"], ensure_ascii=False)[:300])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0cc0fbd5",
      "metadata": {
        "id": "0cc0fbd5"
      },
      "source": [
        "## The requests\n",
        "\n",
        "Each example is a user request with the phone's date, time, and locale, and the gold calls. The requests are in 16 languages, and the gold dates and times were computed in code from the context. The splits are:\n",
        "\n",
        "- `train`: 2,500 requests, which use 98 of the 108 tools\n",
        "- `test`: requests that use the same 98 tools\n",
        "- `test_unseen`: requests for the 10 held-out tools, which are in every prompt but never the right answer in training. Five of them are lookalikes of training tools (`draft_email` vs `send_email`, `set_sleep_timer` vs `set_timer`, ...).\n",
        "- `showcase`: a few hand-checked examples from `test_unseen`\n",
        "\n",
        "The categories cover what an assistant has to get right: one call, several calls, lookalike traps, no tool needed, a missing detail to ask about, and requests that are only partly possible."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e72bbd24",
      "metadata": {
        "id": "e72bbd24"
      },
      "outputs": [],
      "source": [
        "dataset = load_dataset(dataset_id)\n",
        "print(dataset)\n",
        "print(Counter(dataset[\"train\"][\"category\"]))\n",
        "\n",
        "for category in [\"multi\", \"clarify\"]:\n",
        "    example = next(e for e in dataset[\"train\"] if e[\"category\"] == category)\n",
        "    print(f\"\\n🗂️ {category} ({example['language']}, {example['context']['datetime']}): {example['request']}\")\n",
        "    print(f\"🎯 Gold calls: {example['calls']}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c1d2dba6",
      "metadata": {
        "id": "c1d2dba6"
      },
      "source": [
        "## Build the prompts\n",
        "\n",
        "The system prompt holds the rules and the context, and the chat template adds the 108 tool schemas (about 8,000 tokens). LFM3 writes tool calls in XML, and the tokenizer knows how to parse them back into structured calls with `tokenizer.parse_response`.\n",
        "\n",
        "We use LFM3's non-reasoning mode (`enable_thinking=False`). Since every prompt carries the same catalog, we render the prompts ourselves and give TRL plain text prompts."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7f91a2b2",
      "metadata": {
        "id": "7f91a2b2"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "SYSTEM_PROMPT = \"\"\"You are a helpful assistant running on the user's phone. You can call the phone's tools.\n",
        "Current date and time: {weekday}, {date} {time} ({timezone}). User locale: {locale}.\n",
        "\n",
        "Rules:\n",
        "- Call every tool needed to complete the request, with arguments taken from the request and the context above. Fill optional arguments only when the user gives them; never invent values.\n",
        "- Resolve relative dates and times (\"tomorrow\", \"on Friday\", \"in 20 minutes\") against the current date and time. Times are local, in the formats the tool descriptions give.\n",
        "- If the request needs a tool but a required detail is missing, ask a short clarifying question instead of calling a tool.\n",
        "- Answer small talk and general-knowledge questions directly, without tools. Use translate_text, calculate, convert_units and convert_currency for translations, arithmetic and conversions, and search_web or get_news_headlines for current or local information.\n",
        "- If no tool can do what is asked, say so briefly instead of calling a different tool. If only part of the request is possible, do that part and say what you could not do.\n",
        "- Keep free-text arguments (messages, notes, search queries) in the user's language.\"\"\"\n",
        "\n",
        "def build_prompt(example):\n",
        "    now = datetime.fromisoformat(example[\"context\"][\"datetime\"])\n",
        "    system = SYSTEM_PROMPT.format(weekday=now.strftime(\"%A\"), date=now.strftime(\"%Y-%m-%d\"), time=now.strftime(\"%H:%M\"), **example[\"context\"])\n",
        "    messages = [{\"role\": \"system\", \"content\": system}, {\"role\": \"user\", \"content\": example[\"request\"]}]\n",
        "    return tokenizer.apply_chat_template(messages, tools=tools, add_generation_prompt=True, tokenize=False, enable_thinking=False)\n",
        "\n",
        "prompt = build_prompt(dataset[\"train\"][0])\n",
        "print(f\"📏 Prompt length: {len(tokenizer(prompt, add_special_tokens=False).input_ids):,} tokens\")\n",
        "print(prompt[-1200:])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "11fe70da",
      "metadata": {
        "id": "11fe70da"
      },
      "source": [
        "## Compare calls with the gold calls\n",
        "\n",
        "A prediction is correct when it makes exactly the gold calls, in any order. Arguments are compared by kind:\n",
        "\n",
        "- numbers, booleans, and enums exactly\n",
        "- dates and times as values (`19:00` is `19:00:00`)\n",
        "- codes such as currencies and phone numbers, without spaces or dashes\n",
        "- names, without case or accents\n",
        "- free text (a message body, a search query), with some tolerance for wording\n",
        "\n",
        "Optional arguments left at their default value count as absent."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2853bd64",
      "metadata": {
        "id": "2853bd64"
      },
      "outputs": [],
      "source": [
        "import ast\n",
        "import itertools\n",
        "import math\n",
        "import operator\n",
        "import re\n",
        "import unicodedata\n",
        "from datetime import date, time\n",
        "from difflib import SequenceMatcher\n",
        "\n",
        "def norm(s):\n",
        "    s = unicodedata.normalize(\"NFKC\", str(s)).casefold()\n",
        "    s = \"\".join(c for c in unicodedata.normalize(\"NFD\", s) if unicodedata.category(c) != \"Mn\")\n",
        "    s = \"\".join(\" \" if unicodedata.category(c)[0] in \"PSZ\" else c for c in unicodedata.normalize(\"NFC\", s))\n",
        "    return \" \".join(s.split())\n",
        "\n",
        "def contained(a, b):\n",
        "    short, long_ = sorted((a, b), key=len)\n",
        "    return bool(short) and short in long_ and len(short) / len(long_) >= 0.5\n",
        "\n",
        "def name_equal(a, b):\n",
        "    a, b = norm(a), norm(b)\n",
        "    return a == b or contained(a, b)\n",
        "\n",
        "def text_similar(a, b):\n",
        "    a, b = norm(a), norm(b)\n",
        "    return a == b or contained(a, b) or (bool(a) and bool(b) and SequenceMatcher(None, a, b).ratio() >= 0.6)\n",
        "\n",
        "def parse_when(value, kind):\n",
        "    value = str(value).strip()\n",
        "    patterns = {\"datetime\": r\"\\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}(:\\d{2})?\", \"date\": r\"\\d{4}-\\d{2}-\\d{2}\", \"time\": r\"\\d{1,2}:\\d{2}(:\\d{2})?\"}\n",
        "    if not re.fullmatch(patterns[kind], value):\n",
        "        return None\n",
        "    return {\"datetime\": datetime, \"date\": date, \"time\": time}[kind].fromisoformat(value.zfill(5) if kind == \"time\" and len(value) == 4 else value)\n",
        "\n",
        "OPERATORS = {ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.Div: operator.truediv,\n",
        "             ast.Pow: operator.pow, ast.Mod: operator.mod, ast.USub: operator.neg, ast.UAdd: operator.pos}\n",
        "\n",
        "def calculate(expression):\n",
        "    \"\"\"Safely evaluate an arithmetic expression, or return None.\"\"\"\n",
        "    def walk(node):\n",
        "        if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)):\n",
        "            return node.value\n",
        "        if isinstance(node, ast.BinOp):\n",
        "            return OPERATORS[type(node.op)](walk(node.left), walk(node.right))\n",
        "        if isinstance(node, ast.UnaryOp):\n",
        "            return OPERATORS[type(node.op)](walk(node.operand))\n",
        "        raise ValueError\n",
        "    try:\n",
        "        return float(walk(ast.parse(str(expression).replace(\"^\", \"**\").replace(\"×\", \"*\").replace(\"÷\", \"/\"), mode=\"eval\").body))\n",
        "    except Exception:\n",
        "        return None\n",
        "\n",
        "def same_list(gold, pred, equal):\n",
        "    return len(gold) == len(pred) and any(all(equal(g, p) for g, p in zip(gold, perm)) for perm in itertools.permutations(pred))\n",
        "\n",
        "def same_value(kind, gold, pred):\n",
        "    try:\n",
        "        if kind == \"bool\":\n",
        "            return isinstance(pred, bool) and gold == pred\n",
        "        if kind in (\"int\", \"number\"):\n",
        "            return isinstance(pred, (int, float)) and not isinstance(pred, bool) and abs(float(gold) - float(pred)) < 1e-6\n",
        "        if kind == \"enum\":\n",
        "            return isinstance(pred, str) and gold == pred.strip().lower()\n",
        "        if kind == \"code\":\n",
        "            return isinstance(pred, str) and re.sub(r\"[\\s\\-().]\", \"\", str(gold).casefold()) == re.sub(r\"[\\s\\-().]\", \"\", pred.casefold())\n",
        "        if kind in (\"datetime\", \"date\", \"time\"):\n",
        "            return parse_when(pred, kind) is not None and parse_when(gold, kind) == parse_when(pred, kind)\n",
        "        if kind == \"expr\":\n",
        "            g, p = calculate(gold), calculate(pred)\n",
        "            return g is not None and p is not None and math.isclose(g, p, rel_tol=1e-9, abs_tol=1e-9)\n",
        "        if kind == \"name\":\n",
        "            return isinstance(pred, str) and name_equal(gold, pred)\n",
        "        if kind == \"text\":\n",
        "            return isinstance(pred, str) and text_similar(gold, pred)\n",
        "        if kind == \"list_enum\":\n",
        "            return isinstance(pred, list) and sorted(gold) == sorted(str(x).lower() for x in pred)\n",
        "        if kind == \"list_name\":\n",
        "            return isinstance(pred, list) and same_list(gold, [str(x) for x in pred], name_equal)\n",
        "        if kind == \"list_text\":\n",
        "            return isinstance(pred, list) and same_list(gold, [str(x) for x in pred], text_similar)\n",
        "    except (TypeError, ValueError, OverflowError, ZeroDivisionError):\n",
        "        return False\n",
        "    return False\n",
        "\n",
        "schemas = {tool[\"function\"][\"name\"]: tool[\"function\"][\"parameters\"] for tool in tools}\n",
        "\n",
        "def canonical(name, arguments):\n",
        "    \"\"\"Fill schema defaults, drop empty optional values, and treat end_date == date as a single day.\"\"\"\n",
        "    schema = schemas[name]\n",
        "    required = set(schema.get(\"required\", []))\n",
        "    out = {k: v for k, v in (arguments or {}).items() if k in required or v not in (None, \"\", [])}\n",
        "    for k, prop in schema[\"properties\"].items():\n",
        "        if k not in out and \"default\" in prop:\n",
        "            out[k] = prop[\"default\"]\n",
        "    if \"end_date\" in out and out.get(\"end_date\") == out.get(\"date\"):\n",
        "        del out[\"end_date\"]\n",
        "    return out\n",
        "\n",
        "def same_call(gold, pred):\n",
        "    if gold[\"name\"] != pred.get(\"name\"):\n",
        "        return False\n",
        "    g, p = canonical(gold[\"name\"], gold.get(\"arguments\")), canonical(gold[\"name\"], pred.get(\"arguments\"))\n",
        "    return set(g) == set(p) and all(k in kinds[gold[\"name\"]] and same_value(kinds[gold[\"name\"]][k], g[k], p[k]) for k in g)\n",
        "\n",
        "def calls_match(gold, pred):\n",
        "    \"\"\"True when the prediction makes exactly the gold calls, in any order.\"\"\"\n",
        "    if len(gold) != len(pred) or any(call.get(\"name\") not in schemas for call in pred):\n",
        "        return False\n",
        "    return any(all(same_call(g, p) for g, p in zip(gold, perm)) for perm in itertools.permutations(pred))\n",
        "\n",
        "def predicted_calls(completion_ids, prompt):\n",
        "    \"\"\"Parse a completion into a list of {\"name\", \"arguments\"} calls.\"\"\"\n",
        "    message = tokenizer.parse_response(completion_ids, prefix=prompt, tools=tools)\n",
        "    return [call[\"function\"] for call in message.get(\"tool_calls\") or []]\n",
        "\n",
        "# A quick check on a gold example\n",
        "example = dataset[\"train\"][0]\n",
        "gold = json.loads(example[\"calls\"])\n",
        "print(calls_match(gold, gold), calls_match(gold, []))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7c3a7991",
      "metadata": {
        "id": "7c3a7991"
      },
      "source": [
        "# 📏 Evaluate the model before training\n",
        "\n",
        "We generate one answer per request and count exact matches, per category. To keep the notebook fast, we evaluate on 200 requests of each test split."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b33c9f06",
      "metadata": {
        "id": "b33c9f06"
      },
      "outputs": [],
      "source": [
        "from tqdm.auto import tqdm\n",
        "\n",
        "@torch.no_grad()\n",
        "def evaluate(model, split, n=200, batch_size=8):\n",
        "    model.eval()\n",
        "    examples = dataset[split].shuffle(seed=42).select(range(n))\n",
        "    correct = {}\n",
        "    for start in tqdm(range(0, n, batch_size), desc=f\"Evaluating {split}\"):\n",
        "        batch = examples.select(range(start, min(start + batch_size, n)))\n",
        "        prompts = [build_prompt(e) for e in batch]\n",
        "        inputs = tokenizer(prompts, return_tensors=\"pt\", padding=True, add_special_tokens=False).to(model.device)\n",
        "        outputs = model.generate(**inputs, max_new_tokens=512, do_sample=False)\n",
        "        for output, prompt, example in zip(outputs[:, inputs[\"input_ids\"].shape[1]:], prompts, batch):\n",
        "            output = output[output != tokenizer.pad_token_id]\n",
        "            correct.setdefault(example[\"category\"], []).append(calls_match(json.loads(example[\"calls\"]), predicted_calls(output, prompt)))\n",
        "    results = {category: sum(c) / len(c) for category, c in correct.items()}\n",
        "    results[\"all\"] = sum(sum(c) for c in correct.values()) / n\n",
        "    return results\n",
        "\n",
        "def show(results, title):\n",
        "    print(title)\n",
        "    for split, r in results.items():\n",
        "        print(f\"   {split:<12} exact match {r['all']:.1%} | \" + \" | \".join(f\"{k} {v:.0%}\" for k, v in sorted(r.items()) if k != \"all\"))\n",
        "\n",
        "baseline = {split: evaluate(model, split) for split in [\"test\", \"test_unseen\"]}\n",
        "show(baseline, \"📏 Before training\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7587d5bd",
      "metadata": {
        "id": "7587d5bd"
      },
      "source": [
        "# 🎯 Group Relative Policy Optimization (GRPO + LoRA)\n",
        "\n",
        "## The reward\n",
        "\n",
        "For each request, GRPO samples a group of answers and compares their rewards: answers that score above the group's average are reinforced, the others discouraged. Our reward is:\n",
        "\n",
        "- **1.0** when the calls exactly match the gold calls (for requests where the right answer is to make no call, when the model makes no call)\n",
        "- **0.2 × F1** of the tool names otherwise, a small partial credit for picking the right tools with wrong arguments, which helps early in training\n",
        "\n",
        "TRL passes the dataset columns (here `calls`) and the completion token ids to the reward function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "00c36327",
      "metadata": {
        "id": "00c36327"
      },
      "outputs": [],
      "source": [
        "def tool_name_f1(gold, pred):\n",
        "    gold_names, pred_names = Counter(c[\"name\"] for c in gold), Counter(c[\"name\"] for c in pred)\n",
        "    overlap = sum((gold_names & pred_names).values())\n",
        "    return 0.0 if overlap == 0 else 2 * overlap / (sum(gold_names.values()) + sum(pred_names.values()))\n",
        "\n",
        "def tool_call_reward(prompts, completion_ids, calls, **kwargs):\n",
        "    rewards = []\n",
        "    for prompt, ids, gold in zip(prompts, completion_ids, calls):\n",
        "        gold, pred = json.loads(gold), predicted_calls(ids, prompt)\n",
        "        rewards.append(1.0 if calls_match(gold, pred) else 0.2 * tool_name_f1(gold, pred))\n",
        "    return rewards"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e20d6442",
      "metadata": {
        "id": "e20d6442"
      },
      "source": [
        "## Balance the training requests\n",
        "\n",
        "Only 15% of the training requests need no tool call (`no_tool` and `clarify`). Since most rewards come from making the right calls, RL alone teaches the model to always call a tool, even when a detail is missing and it should ask. We counter this by showing the no-call requests twice."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7f7b1989",
      "metadata": {
        "id": "7f7b1989"
      },
      "outputs": [],
      "source": [
        "from datasets import concatenate_datasets\n",
        "\n",
        "no_call = dataset[\"train\"].filter(lambda e: e[\"category\"] in (\"no_tool\", \"clarify\"))\n",
        "train_dataset = concatenate_datasets([dataset[\"train\"], no_call]).shuffle(seed=42)\n",
        "train_dataset = train_dataset.map(lambda e: {\"prompt\": build_prompt(e)}, remove_columns=[\"request\", \"context\"])\n",
        "print(f\"📚 Training requests: {len(train_dataset)}\")\n",
        "print(Counter(train_dataset[\"category\"]))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2a58e6b9",
      "metadata": {
        "id": "2a58e6b9"
      },
      "source": [
        "## Wrap the model with PEFT\n",
        "\n",
        "We adapt the attention projections (`q_proj`, `k_proj`, `v_proj`, `o_proj`), the gated short convolutions (`in_proj`, `out_proj`), and the MLP (`w1`, `w2`, `w3`). LFM3's attention layers also have a per-head output gate, `gate_proj`; adding it to the targets is a reasonable choice."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e0087415",
      "metadata": {
        "id": "e0087415"
      },
      "outputs": [],
      "source": [
        "from peft import LoraConfig\n",
        "\n",
        "peft_config = LoraConfig(\n",
        "    r=32,\n",
        "    lora_alpha=64,\n",
        "    lora_dropout=0.0,\n",
        "    target_modules=[\n",
        "        \"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"in_proj\", \"out_proj\", \"w1\", \"w2\", \"w3\",\n",
        "        # \"gate_proj\",  # optional: LFM3's per-head attention gate\n",
        "    ],\n",
        "    task_type=\"CAUSAL_LM\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2a2d94ed",
      "metadata": {
        "id": "2a2d94ed"
      },
      "source": [
        "## Configure training\n",
        "\n",
        "- `num_generations=8`: the size of each group. With 32 answers per step, each step trains on 4 requests.\n",
        "- `generation_batch_size=8`: the answers of one request are generated together. Each prompt is about 8,000 tokens, so generating all 32 answers at once would need much more memory.\n",
        "- `max_completion_length=256`: tool calls and short answers fit easily.\n",
        "- `temperature=1.0`: enough variety within a group for the reward to tell answers apart.\n",
        "- `beta=0.04`: a small KL penalty that keeps the model close to the original one. Without it, the behaviors the reward rarely asks for (like asking a clarifying question) fade as training goes on. With LoRA, TRL gets the reference model by disabling the adapter, so it doesn't load a second copy.\n",
        "\n",
        "TRL generates with Transformers here. It can also generate with [vLLM](https://huggingface.co/docs/trl/vllm_integration) (`use_vllm=True`), which is much faster when your vLLM version supports the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "74fe6fbf",
      "metadata": {
        "id": "74fe6fbf"
      },
      "outputs": [],
      "source": [
        "from trl import GRPOConfig, GRPOTrainer\n",
        "\n",
        "grpo_config = GRPOConfig(\n",
        "    output_dir=\"./LFM3-1.7B-phone-assistant\",\n",
        "    max_steps=150,\n",
        "    per_device_train_batch_size=2,\n",
        "    gradient_accumulation_steps=16,  # <- 32 answers per step\n",
        "    num_generations=8,\n",
        "    generation_batch_size=8,  # <- generate one group at a time, since each prompt is about 8,000 tokens\n",
        "    max_completion_length=256,\n",
        "    temperature=1.0,\n",
        "    learning_rate=1e-5,\n",
        "    beta=0.04,  # <- KL penalty towards the original model\n",
        "    lr_scheduler_type=\"constant_with_warmup\",\n",
        "    warmup_steps=5,\n",
        "    logging_steps=5,\n",
        "    save_strategy=\"no\",\n",
        "    bf16=use_bf16,\n",
        "    fp16=not use_bf16,  # <- float16 mixed precision with loss scaling on GPUs without bfloat16\n",
        "    report_to=\"none\",\n",
        ")\n",
        "\n",
        "trainer = GRPOTrainer(\n",
        "    model=model,\n",
        "    reward_funcs=tool_call_reward,\n",
        "    args=grpo_config,\n",
        "    train_dataset=train_dataset,\n",
        "    processing_class=tokenizer,\n",
        "    peft_config=peft_config,\n",
        ")\n",
        "trainer.model.print_trainable_parameters()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ea965e2e",
      "metadata": {
        "id": "ea965e2e"
      },
      "source": [
        "## Launch training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0befd2af",
      "metadata": {
        "id": "0befd2af"
      },
      "outputs": [],
      "source": [
        "print(\"🚀 Starting GRPO training...\")\n",
        "trainer.train()\n",
        "print(\"🎉 GRPO training completed!\")\n",
        "print(f\"🧮 Peak GPU memory: {torch.cuda.max_memory_reserved() / 1e9:.1f} GB\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "68ecd13c",
      "metadata": {
        "id": "68ecd13c"
      },
      "source": [
        "Let's look at how the reward changed during training (1.0 means every answer in the batch was exactly right):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ebe5e04b",
      "metadata": {
        "id": "ebe5e04b"
      },
      "outputs": [],
      "source": [
        "for log in trainer.state.log_history:\n",
        "    if \"reward\" in log:\n",
        "        print(f\"Step {log['step']:>3}: reward {log['reward']:.2f}, mean length {log['completions/mean_length']:.0f} tokens\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0bb581e7",
      "metadata": {
        "id": "0bb581e7"
      },
      "source": [
        "# 📊 Evaluate the trained model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5c68a817",
      "metadata": {
        "id": "5c68a817"
      },
      "outputs": [],
      "source": [
        "final = {split: evaluate(trainer.model, split) for split in [\"test\", \"test_unseen\"]}\n",
        "show(baseline, \"📏 Before training\")\n",
        "show(final, \"📏 After training\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "124870bb",
      "metadata": {
        "id": "124870bb"
      },
      "source": [
        "# 🔍 Original vs fine-tuned model\n",
        "\n",
        "To see what changed, we give the same requests to the original model and to the trained model. With LoRA, `disable_adapter()` temporarily turns the adapter off, so we don't need to load the original model again. The `showcase` split holds hand-checked requests for tools that never appeared in training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f93beb3b",
      "metadata": {
        "id": "f93beb3b"
      },
      "outputs": [],
      "source": [
        "def answer(model, example):\n",
        "    prompt = build_prompt(example)\n",
        "    inputs = tokenizer(prompt, return_tensors=\"pt\", add_special_tokens=False).to(model.device)\n",
        "    output = model.generate(**inputs, max_new_tokens=512, do_sample=False)[0, inputs[\"input_ids\"].shape[1]:]\n",
        "    message = tokenizer.parse_response(output, prefix=prompt, tools=tools)\n",
        "    calls = [call[\"function\"] for call in message.get(\"tool_calls\") or []]\n",
        "    return calls, (message.get(\"content\") or \"\").strip()\n",
        "\n",
        "def describe(calls, text):\n",
        "    return json.dumps(calls, ensure_ascii=False) if calls else f\"(no call) {text[:150]}\"\n",
        "\n",
        "trainer.model.eval()\n",
        "for example in dataset[\"showcase\"]:\n",
        "    gold = json.loads(example[\"calls\"])\n",
        "    with trainer.model.disable_adapter():\n",
        "        original = answer(trainer.model, example)\n",
        "    trained = answer(trainer.model, example)\n",
        "    print(f\"\\n💬 {example['request']} ({example['language']}, {example['category']})\")\n",
        "    print(f\"   🎯 Gold:       {describe(gold, 'ask for the missing detail or answer directly')}\")\n",
        "    print(f\"   {'✅' if calls_match(gold, original[0]) else '❌'} Original:   {describe(*original)}\")\n",
        "    print(f\"   {'✅' if calls_match(gold, trained[0]) else '❌'} Fine-tuned: {describe(*trained)}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d610032a",
      "metadata": {
        "id": "d610032a"
      },
      "source": [
        "And the summary over all the evaluation requests:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4aa5ba5a",
      "metadata": {
        "id": "4aa5ba5a"
      },
      "outputs": [],
      "source": [
        "print(f\"{'exact match':<36}{'original':>10}{'fine-tuned':>12}\")\n",
        "for split, name in [(\"test\", \"training tools\"), (\"test_unseen\", \"held-out tools\")]:\n",
        "    print(f\"{name:<36}{baseline[split]['all']:>10.1%}{final[split]['all']:>12.1%}\")\n",
        "    for category in sorted(k for k in baseline[split] if k != \"all\"):\n",
        "        print(f\"{'  ' + category:<36}{baseline[split][category]:>10.0%}{final[split][category]:>12.0%}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5bca7480",
      "metadata": {
        "id": "5bca7480"
      },
      "source": [
        "# 💾 Save and merge the model\n",
        "\n",
        "We save the LoRA adapter, then merge it into the base weights. `save_pretrained` also copies LFM3's modeling code next to the weights, so the merged model loads with `trust_remote_code=True` like the original one."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7dc09bed",
      "metadata": {
        "id": "7dc09bed"
      },
      "outputs": [],
      "source": [
        "trainer.save_model(\"./LFM3-1.7B-phone-assistant-lora\")\n",
        "\n",
        "output_dir = \"./LFM3-1.7B-phone-assistant-merged\"\n",
        "merged_model = trainer.model.merge_and_unload()\n",
        "merged_model.save_pretrained(output_dir)\n",
        "tokenizer.save_pretrained(output_dir)\n",
        "print(f\"💾 Merged model saved to: {output_dir}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bc9531e3",
      "metadata": {
        "id": "bc9531e3"
      },
      "source": [
        "# 🚀 Next steps\n",
        "\n",
        "- **Your own app**: export your tools as function schemas, collect real requests, and label the right calls. The reward only needs a way to compare calls, so you don't need full demonstrations.\n",
        "- **Train longer**: the accuracy keeps improving after this run (200 steps reach about 80% on the training tools and 76% on the held-out ones). Increase `max_steps`.\n",
        "- **Multi-turn agents**: `GRPOTrainer` can also run your tools during generation (`tools=[...]` or `environment_factory=...`) and reward the final outcome. See the [TRL documentation](https://huggingface.co/docs/trl/grpo_trainer)."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "A100",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}