{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "b85c4af9",
      "metadata": {
        "id": "b85c4af9"
      },
      "source": [
        "# Grounded Answers with DPO: Teaching LFM3-1.7B to Say \"I Don't Know\" 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_dpo_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",
        "Small models are usually deployed with retrieval: your app finds a few relevant documents, and the model answers from them. The main risk is a confident answer that isn't in the documents. This notebook trains [LFM3-1.7B](https://huggingface.co/LiquidAI/LFM3-1.7B) to answer only from the documents, and to reply with a fixed fallback sentence when the answer isn't there, using Direct Preference Optimization (DPO) with [TRL](https://github.com/huggingface/trl) and LoRA.\n",
        "\n",
        "DPO learns from pairs of answers to the same prompt: a **chosen** answer and a **rejected** one. Here, the rejected answers are the model's own mistakes, sampled from LFM3-1.7B itself: hallucinated answers to questions the documents can't answer, wrong answers, and unnecessary refusals. The chosen answers are verified answers, or the exact fallback sentence. Training on its own mistakes targets exactly the behavior we want to change.\n",
        "\n",
        "## 🎯 What you'll find\n",
        "\n",
        "- **Grounded QA data**: fictional documents in 16 languages, with answerable and tricky unanswerable questions\n",
        "- **LoRA + DPO**: the base model with the adapter disabled serves as the reference, so only one copy of the weights is in memory\n",
        "- **Evaluation**: accuracy on answerable questions and on unanswerable ones, before and after training"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3ebb395e",
      "metadata": {
        "id": "3ebb395e"
      },
      "source": [
        "# 📦 Installation & Setup\n",
        "\n",
        "We pin every library to the versions this notebook was tested with."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "386d7748",
      "metadata": {
        "tags": [
          "install"
        ],
        "id": "386d7748"
      },
      "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": "2c71a51a",
      "metadata": {
        "id": "2c71a51a"
      },
      "outputs": [],
      "source": [
        "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, like the T4, 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": "6a4d9ef1",
      "metadata": {
        "id": "6a4d9ef1"
      },
      "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": "ec360286",
      "metadata": {
        "id": "ec360286"
      },
      "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": "88c4bea0",
      "metadata": {
        "id": "88c4bea0"
      },
      "source": [
        "# 📚 Load the grounded QA dataset\n",
        "\n",
        "Each example has a few short documents about fictional companies, products, places, and policies (so the model can't answer from memory), a question, and the gold label. About half of the questions can't be answered from the documents, and most of those are near-misses: the documents give the price of a different plan, the start date but not the end date, or the policy for a neighboring case.\n",
        "\n",
        "- `train`: preference pairs, with a `chosen` and a `rejected` answer for each prompt\n",
        "- `test`: questions with their gold answers, for evaluation\n",
        "- `test_unseen`: questions about domains that never appear in training\n",
        "- `showcase`: a few hand-checked examples from those unseen domains"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9dd04aa2",
      "metadata": {
        "id": "9dd04aa2"
      },
      "outputs": [],
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "dataset = load_dataset(\"LiquidAI/grounded-qa\")\n",
        "print(dataset)\n",
        "\n",
        "example = dataset[\"train\"][0]\n",
        "print(f\"\\n❓ Question ({example['language']}): {example['question']}\")\n",
        "print(f\"📎 Answerable: {example['answerable']}\")\n",
        "print(f\"\\n✅ Chosen: {example['chosen']}\")\n",
        "print(f\"\\n❌ Rejected: {example['rejected']}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ccd6c00d",
      "metadata": {
        "id": "ccd6c00d"
      },
      "source": [
        "## Build the prompts\n",
        "\n",
        "The instructions and the fallback sentence go in the system prompt, in the language of the question. The documents and the question go in the user message.\n",
        "\n",
        "We use LFM3's non-reasoning mode (`enable_thinking=False`), which answers right away. TRL reads the `chat_template_kwargs` column and passes it to the chat template."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8338db44",
      "metadata": {
        "id": "8338db44"
      },
      "outputs": [],
      "source": [
        "SYSTEM_PROMPT = (\n",
        "    \"Answer the user's question using only the documents they provide. Answer in the language of the question, \"\n",
        "    \"in one or two sentences. If the documents don't contain the answer, reply with exactly this sentence and \"\n",
        "    \"nothing else: {fallback}\"\n",
        ")\n",
        "\n",
        "def build_messages(example):\n",
        "    documents = \"\\n\\n\".join(f\"[{i}] {doc['title']}\\n{doc['text']}\" for i, doc in enumerate(example[\"documents\"], start=1))\n",
        "    return [\n",
        "        {\"role\": \"system\", \"content\": SYSTEM_PROMPT.format(fallback=example[\"fallback\"])},\n",
        "        {\"role\": \"user\", \"content\": f\"Documents:\\n\\n{documents}\\n\\nQuestion: {example['question']}\"},\n",
        "    ]\n",
        "\n",
        "def to_preference(example):\n",
        "    return {\n",
        "        \"prompt\": build_messages(example),\n",
        "        \"chosen\": [{\"role\": \"assistant\", \"content\": example[\"chosen\"]}],\n",
        "        \"rejected\": [{\"role\": \"assistant\", \"content\": example[\"rejected\"]}],\n",
        "        \"chat_template_kwargs\": {\"enable_thinking\": False},\n",
        "    }\n",
        "\n",
        "train_dataset = dataset[\"train\"].map(to_preference, remove_columns=dataset[\"train\"].column_names)\n",
        "print(train_dataset[0][\"prompt\"][1][\"content\"][:1200])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9fecb001",
      "metadata": {
        "id": "9fecb001"
      },
      "source": [
        "# 📏 Evaluate the model before training\n",
        "\n",
        "We generate an answer for each test question, with greedy decoding and the repetition penalty recommended in the model card (`repetition_penalty=1.1`), and check it against the rules of the system prompt:\n",
        "\n",
        "- **Answerable questions**: the answer is correct if it contains the gold answer and doesn't include the fallback sentence. We also count how often the model declines them, which should stay low.\n",
        "- **Unanswerable questions**: the answer is correct if it's exactly the fallback sentence, as the system prompt asks. The original model often writes the fallback sentence and then answers anyway, which an app can't rely on.\n",
        "\n",
        "We evaluate on 200 questions from `test` (the training domains) and 200 from `test_unseen` (domains that never appear in training), to check that the behavior generalizes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9b61e7ac",
      "metadata": {
        "id": "9b61e7ac"
      },
      "outputs": [],
      "source": [
        "from tqdm.auto import tqdm\n",
        "\n",
        "def normalize(text):\n",
        "    return \" \".join(text.casefold().replace(\"’\", \"'\").split()).strip(\" .!?。\")\n",
        "\n",
        "@torch.no_grad()\n",
        "def evaluate(model, split, n=200, batch_size=16):\n",
        "    model.eval()\n",
        "    examples = dataset[split].shuffle(seed=42).select(range(n))\n",
        "    results = {\"answerable\": [], \"unanswerable\": [], \"over_declined\": []}\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 = [tokenizer.apply_chat_template(build_messages(e), add_generation_prompt=True, tokenize=False, enable_thinking=False) 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=256, do_sample=False, repetition_penalty=1.1)\n",
        "        for output, example in zip(outputs[:, inputs[\"input_ids\"].shape[1]:], batch):\n",
        "            answer = normalize(tokenizer.decode(output, skip_special_tokens=True))\n",
        "            fell_back = answer == normalize(example[\"fallback\"])\n",
        "            if example[\"answerable\"]:\n",
        "                results[\"answerable\"].append(normalize(example[\"fallback\"]) not in answer and normalize(example[\"answer\"]) in answer)\n",
        "                results[\"over_declined\"].append(fell_back)\n",
        "            else:\n",
        "                results[\"unanswerable\"].append(fell_back)\n",
        "    return {kind: sum(r) / len(r) for kind, r in results.items()}\n",
        "\n",
        "def show(results, title):\n",
        "    print(title)\n",
        "    for split, r in results.items():\n",
        "        print(f\"   {split:<12} answerable: answered correctly {r['answerable']:.1%}, declined {r['over_declined']:.1%} | unanswerable: declined {r['unanswerable']:.1%}\")\n",
        "\n",
        "baseline = {split: evaluate(model, split) for split in [\"test\", \"test_unseen\"]}\n",
        "show(baseline, \"📏 Before training\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "11b6921f",
      "metadata": {
        "id": "11b6921f"
      },
      "source": [
        "# 🎯 Direct Preference Optimization (DPO + LoRA)\n",
        "\n",
        "## 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": "3eb309f1",
      "metadata": {
        "id": "3eb309f1"
      },
      "outputs": [],
      "source": [
        "from peft import LoraConfig\n",
        "\n",
        "peft_config = LoraConfig(\n",
        "    r=16,\n",
        "    lora_alpha=32,\n",
        "    lora_dropout=0.05,\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": "3cea5e61",
      "metadata": {
        "id": "3cea5e61"
      },
      "source": [
        "## Launch training\n",
        "\n",
        "`beta` controls how far the model can move away from the reference: lower values allow larger changes.\n",
        "\n",
        "We combine the DPO loss with an SFT loss on the chosen answers (`loss_type=[\"sigmoid\", \"sft\"]`). DPO alone only learns which answer is *better*: here it quickly stops the model from prefixing its answers with the fallback sentence, but it also pushes the model away from the fallback when it should use it. The SFT term teaches the model to actually write the chosen answers, including the exact fallback sentence. In our runs, it took the share of declined unanswerable questions from about 20% to about 90%.\n",
        "\n",
        "During training, watch `rewards/accuracies` (how often the model prefers the chosen answer) and `rewards/margins` go up."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4d4bfcb8",
      "metadata": {
        "id": "4d4bfcb8"
      },
      "outputs": [],
      "source": [
        "from trl import DPOConfig, DPOTrainer\n",
        "\n",
        "dpo_config = DPOConfig(\n",
        "    output_dir=\"./LFM3-1.7B-grounded-dpo\",\n",
        "    num_train_epochs=1,\n",
        "    per_device_train_batch_size=2,\n",
        "    gradient_accumulation_steps=8,  # <- effective batch size of 16 pairs\n",
        "    learning_rate=2e-5,\n",
        "    loss_type=[\"sigmoid\", \"sft\"],  # <- DPO loss + SFT loss on the chosen answers\n",
        "    loss_weights=[1.0, 1.0],\n",
        "    lr_scheduler_type=\"cosine\",\n",
        "    warmup_steps=0.1,  # <- a float is a fraction of the total steps\n",
        "    beta=0.1,\n",
        "    max_length=3072,\n",
        "    logging_steps=10,\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 = DPOTrainer(\n",
        "    model=model,\n",
        "    args=dpo_config,\n",
        "    train_dataset=train_dataset,\n",
        "    processing_class=tokenizer,\n",
        "    peft_config=peft_config,\n",
        ")\n",
        "trainer.model.print_trainable_parameters()\n",
        "\n",
        "print(\"🚀 Starting DPO training...\")\n",
        "trainer.train()\n",
        "print(\"🎉 DPO training completed!\")\n",
        "print(f\"🧮 Peak GPU memory: {torch.cuda.max_memory_reserved() / 1e9:.1f} GB\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "697e3811",
      "metadata": {
        "id": "697e3811"
      },
      "source": [
        "# 📊 Evaluate the trained model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ffaf97c5",
      "metadata": {
        "id": "ffaf97c5"
      },
      "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": "bc55c35d",
      "metadata": {
        "id": "bc55c35d"
      },
      "source": [
        "# 🔍 Original vs fine-tuned model\n",
        "\n",
        "To see what changed, we ask the original model and the trained model the same questions. 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 questions about domains the model never saw during training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1013656c",
      "metadata": {
        "id": "1013656c"
      },
      "outputs": [],
      "source": [
        "def answer(model, example):\n",
        "    prompt = tokenizer.apply_chat_template(build_messages(example), add_generation_prompt=True, tokenize=False, enable_thinking=False)\n",
        "    inputs = tokenizer(prompt, return_tensors=\"pt\", add_special_tokens=False).to(model.device)\n",
        "    output = model.generate(**inputs, max_new_tokens=256, do_sample=False, repetition_penalty=1.1)\n",
        "    return tokenizer.decode(output[0, inputs[\"input_ids\"].shape[1]:], skip_special_tokens=True).strip()\n",
        "\n",
        "def verdict(text, example):\n",
        "    fell_back = normalize(text) == normalize(example[\"fallback\"])\n",
        "    if example[\"answerable\"]:\n",
        "        return \"✅\" if normalize(example[\"fallback\"]) not in normalize(text) and normalize(example[\"answer\"]) in normalize(text) else \"❌\"\n",
        "    return \"✅\" if fell_back else \"❌\"\n",
        "\n",
        "trainer.model.eval()\n",
        "for example in dataset[\"showcase\"]:\n",
        "    with trainer.model.disable_adapter():\n",
        "        original = answer(trainer.model, example)\n",
        "    trained = answer(trainer.model, example)\n",
        "    print(f\"\\n❓ {example['question']} ({example['language']})\")\n",
        "    print(f\"   🎯 Gold:       {example['reference_answer']}\")\n",
        "    print(f\"   {verdict(original, example)} Original:   {original}\")\n",
        "    print(f\"   {verdict(trained, example)} Fine-tuned: {trained}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "82f4297a",
      "metadata": {
        "id": "82f4297a"
      },
      "source": [
        "And the summary over all the evaluation questions:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "956dca4f",
      "metadata": {
        "id": "956dca4f"
      },
      "outputs": [],
      "source": [
        "print(f\"{'':<46}{'original':>10}{'fine-tuned':>12}\")\n",
        "for split, name in [(\"test\", \"training domains\"), (\"test_unseen\", \"unseen domains\")]:\n",
        "    for kind, label in [(\"answerable\", \"answerable, answered correctly\"), (\"over_declined\", \"answerable, declined (lower is better)\"), (\"unanswerable\", \"unanswerable, declined\")]:\n",
        "        print(f\"{name + ', ' + label:<46}{baseline[split][kind]:>10.1%}{final[split][kind]:>12.1%}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "097b063f",
      "metadata": {
        "id": "097b063f"
      },
      "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": "70e2e0be",
      "metadata": {
        "id": "70e2e0be"
      },
      "outputs": [],
      "source": [
        "trainer.save_model(\"./LFM3-1.7B-grounded-dpo-lora\")\n",
        "\n",
        "output_dir = \"./LFM3-1.7B-grounded-dpo-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": "656e2dac",
      "metadata": {
        "id": "656e2dac"
      },
      "source": [
        "# 🚀 Next steps\n",
        "\n",
        "- **Your own documents**: sample answers from your model on your own retrieval prompts, label them (by hand, with code, or with a stronger model as a judge), and pair a good answer with a bad one for each prompt.\n",
        "- **Keep it on-policy**: after a round of DPO, sample again from the trained model and build new pairs from its remaining mistakes."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}