{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LLM Alignment with DPO and KTO\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/ludwig-ai/ludwig/blob/main/examples/alignment/alignment_dpo.ipynb)\n",
    "\n",
    "This notebook walks through aligning **Llama-3.1-8B** with human preferences using Ludwig's built-in\n",
    "**Direct Preference Optimization (DPO)** trainer, then shows how to switch to **KTO** when paired\n",
    "preference data is not available.\n",
    "\n",
    "We use the publicly available [Anthropic HH-RLHF](https://huggingface.co/datasets/Anthropic/hh-rlhf)\n",
    "dataset, which contains helpfulness and harmlessness preference pairs collected from human annotators.\n",
    "\n",
    "### What you will learn\n",
    "- How to parse the HH-RLHF dataset into Ludwig's expected column format\n",
    "- How to configure and run DPO alignment training with LoRA\n",
    "- How to evaluate response quality before and after alignment\n",
    "- How to switch to KTO with a one-line config change\n",
    "- How to upload the aligned model to HuggingFace Hub\n",
    "\n",
    "### Prerequisites\n",
    "- **GPU**: A100 (40 GiB) recommended. T4 will work with smaller batch sizes but is slower.\n",
    "- **HuggingFace token**: Required for Llama-3.1-8B. Set `HUGGING_FACE_HUB_TOKEN` before running.\n",
    "- **Model access**: Request access at https://huggingface.co/meta-llama/Meta-Llama-3.1-8B"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"ludwig[llm]\" datasets --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Set your HuggingFace token — required to download Llama-3.1-8B.\n",
    "# In Colab: Secrets panel (key icon) -> add HF_TOKEN, then reference it here.\n",
    "try:\n",
    "    from google.colab import userdata\n",
    "\n",
    "    os.environ[\"HUGGING_FACE_HUB_TOKEN\"] = userdata.get(\"HF_TOKEN\")\n",
    "except Exception:\n",
    "    pass  # Running locally — export HUGGING_FACE_HUB_TOKEN in your shell instead\n",
    "\n",
    "assert os.environ.get(\"HUGGING_FACE_HUB_TOKEN\"), (\n",
    "    \"HUGGING_FACE_HUB_TOKEN is not set. Add it to Colab Secrets or export it in your shell.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess\n",
    "\n",
    "# Verify GPU availability\n",
    "result = subprocess.run(\n",
    "    [\"nvidia-smi\", \"--query-gpu=name,memory.total\", \"--format=csv,noheader\"], capture_output=True, text=True\n",
    ")\n",
    "if result.returncode == 0:\n",
    "    print(\"GPU(s) detected:\")\n",
    "    print(result.stdout.strip())\n",
    "else:\n",
    "    print(\"WARNING: No GPU detected. Training will be very slow on CPU.\")\n",
    "    print(\"In Colab: Runtime -> Change runtime type -> A100 GPU\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare dataset\n",
    "\n",
    "The [Anthropic HH-RLHF](https://huggingface.co/datasets/Anthropic/hh-rlhf) dataset stores full\n",
    "multi-turn conversations in two columns: `chosen` (preferred response) and `rejected` (dispreferred\n",
    "response). Each value is a raw conversation string like:\n",
    "\n",
    "```\n",
    "\\n\\nHuman: How do I make pasta?\\n\\nAssistant: Boil water, add pasta...\\n\\nHuman: What sauce?\\n\\nAssistant: Try marinara...\n",
    "```\n",
    "\n",
    "Ludwig's DPO trainer expects three columns: `prompt`, `chosen`, `rejected`, where `prompt` is the\n",
    "last human turn and `chosen`/`rejected` are the final assistant responses.\n",
    "\n",
    "The cell below does that extraction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "import pandas as pd\n",
    "from datasets import load_dataset\n",
    "\n",
    "\n",
    "def extract_last_human_turn(conversation: str) -> str:\n",
    "    \"\"\"Extract the last Human turn from an HH-RLHF conversation string.\"\"\"\n",
    "    turns = re.findall(r\"\\n\\nHuman: (.*?)(?=\\n\\nAssistant:|\\Z)\", conversation, re.DOTALL)\n",
    "    return turns[-1].strip() if turns else conversation.strip()\n",
    "\n",
    "\n",
    "def extract_last_assistant_turn(conversation: str) -> str:\n",
    "    \"\"\"Extract the last Assistant turn from an HH-RLHF conversation string.\"\"\"\n",
    "    turns = re.findall(r\"\\n\\nAssistant: (.*?)(?=\\n\\nHuman:|\\Z)\", conversation, re.DOTALL)\n",
    "    return turns[-1].strip() if turns else \"\"\n",
    "\n",
    "\n",
    "def prepare_dpo_dataset(split, max_samples=None):\n",
    "    if max_samples:\n",
    "        split = split.select(range(min(max_samples, len(split))))\n",
    "    rows = []\n",
    "    for ex in split:\n",
    "        prompt = extract_last_human_turn(ex[\"chosen\"])\n",
    "        chosen = extract_last_assistant_turn(ex[\"chosen\"])\n",
    "        rejected = extract_last_assistant_turn(ex[\"rejected\"])\n",
    "        if prompt and chosen and rejected:\n",
    "            rows.append({\"prompt\": prompt, \"chosen\": chosen, \"rejected\": rejected})\n",
    "    return pd.DataFrame(rows)\n",
    "\n",
    "\n",
    "print(\"Downloading Anthropic/hh-rlhf ...\")\n",
    "hh = load_dataset(\"Anthropic/hh-rlhf\")\n",
    "print(f\"Train size: {len(hh['train'])}, Test size: {len(hh['test'])}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use a subset for a quick experiment. Remove max_samples caps for full training.\n",
    "MAX_TRAIN = 5000\n",
    "MAX_TEST = 500\n",
    "\n",
    "train_df = prepare_dpo_dataset(hh[\"train\"], max_samples=MAX_TRAIN)\n",
    "test_df = prepare_dpo_dataset(hh[\"test\"], max_samples=MAX_TEST)\n",
    "\n",
    "train_df.to_csv(\"train.csv\", index=False)\n",
    "test_df.to_csv(\"test.csv\", index=False)\n",
    "\n",
    "print(f\"Saved {len(train_df)} train rows to train.csv\")\n",
    "print(f\"Saved {len(test_df)} test rows to test.csv\")\n",
    "train_df.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preview: show a prompt with its chosen vs rejected response\n",
    "row = train_df.iloc[0]\n",
    "print(\"=== PROMPT ===\")\n",
    "print(row[\"prompt\"])\n",
    "print(\"\\n=== CHOSEN ===\")\n",
    "print(row[\"chosen\"])\n",
    "print(\"\\n=== REJECTED ===\")\n",
    "print(row[\"rejected\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DPO training\n",
    "\n",
    "Direct Preference Optimization (DPO) treats alignment as a classification problem: given a prompt,\n",
    "prefer `chosen` over `rejected`. The loss function directly optimises the log-ratio between the\n",
    "policy's probability of the chosen response and the reference model's, without needing an explicit\n",
    "reward model.\n",
    "\n",
    "Key hyperparameters:\n",
    "- `beta`: KL penalty coefficient (0.1 is a typical starting point). Higher values keep the aligned\n",
    "  model closer to the base model.\n",
    "- `learning_rate`: 5e-7 is much lower than SFT — the policy already knows how to generate text;\n",
    "  we're only nudging its preferences.\n",
    "- `lora r=16, alpha=32`: Standard LoRA settings. Increase `r` for larger models or more complex tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "\n",
    "import yaml\n",
    "\n",
    "from ludwig.api import LudwigModel\n",
    "\n",
    "dpo_config = yaml.safe_load(\"\"\"\n",
    "model_type: llm\n",
    "base_model: meta-llama/Llama-3.1-8B\n",
    "\n",
    "adapter:\n",
    "  type: lora\n",
    "  r: 16\n",
    "  alpha: 32\n",
    "  dropout: 0.05\n",
    "\n",
    "trainer:\n",
    "  type: dpo\n",
    "  epochs: 1\n",
    "  learning_rate: 5.0e-7\n",
    "  batch_size: 2\n",
    "  gradient_accumulation_steps: 8\n",
    "  beta: 0.1\n",
    "\n",
    "input_features:\n",
    "  - name: prompt\n",
    "    type: text\n",
    "\n",
    "output_features:\n",
    "  - name: chosen\n",
    "    type: text\n",
    "\n",
    "backend:\n",
    "  type: local\n",
    "\"\"\")\n",
    "\n",
    "dpo_model = LudwigModel(config=dpo_config, logging_level=logging.INFO)\n",
    "\n",
    "train_stats, _, output_directory = dpo_model.train(\n",
    "    dataset=\"train.csv\",\n",
    "    experiment_name=\"hh_rlhf_dpo\",\n",
    "    output_directory=\"results\",\n",
    ")\n",
    "\n",
    "print(f\"\\nModel saved to: {output_directory}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate\n",
    "\n",
    "We compare responses from the base model (no alignment) and the DPO-aligned model on a few\n",
    "prompts from the test set. You should see the aligned model give more helpful, less evasive answers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the base model for comparison\n",
    "base_config = yaml.safe_load(\"\"\"\n",
    "model_type: llm\n",
    "base_model: meta-llama/Llama-3.1-8B\n",
    "\n",
    "input_features:\n",
    "  - name: prompt\n",
    "    type: text\n",
    "\n",
    "output_features:\n",
    "  - name: chosen\n",
    "    type: text\n",
    "\n",
    "backend:\n",
    "  type: local\n",
    "\"\"\")\n",
    "\n",
    "base_model = LudwigModel(config=base_config, logging_level=logging.WARNING)\n",
    "_ = base_model.train(dataset=\"train.csv\", experiment_name=\"baseline\", output_directory=\"results_base\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_prompts = test_df[\"prompt\"].head(3).tolist()\n",
    "eval_df = pd.DataFrame({\"prompt\": eval_prompts})\n",
    "\n",
    "base_preds, _ = base_model.predict(dataset=eval_df)\n",
    "dpo_preds, _ = dpo_model.predict(dataset=eval_df)\n",
    "\n",
    "for i, prompt in enumerate(eval_prompts):\n",
    "    print(f\"\\n{'=' * 60}\")\n",
    "    print(f\"PROMPT: {prompt}\")\n",
    "    print(\"\\n--- Base model ---\")\n",
    "    print(base_preds.iloc[i][\"chosen_predictions\"])\n",
    "    print(\"\\n--- DPO-aligned model ---\")\n",
    "    print(dpo_preds.iloc[i][\"chosen_predictions\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## KTO alternative\n",
    "\n",
    "**Kahneman-Tversky Optimization (KTO)** is a good alternative to DPO when you cannot collect\n",
    "paired preference data. Instead of a (chosen, rejected) pair per prompt, KTO requires only a\n",
    "single response with a boolean label indicating whether it was desirable.\n",
    "\n",
    "This makes it easy to use binary user feedback (thumbs up / thumbs down, click-through, etc.)\n",
    "as training signal.\n",
    "\n",
    "### Dataset format for KTO\n",
    "\n",
    "We expand each DPO pair into two rows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rows_kto = []\n",
    "for _, row in train_df.iterrows():\n",
    "    rows_kto.append({\"prompt\": row[\"prompt\"], \"response\": row[\"chosen\"], \"label\": True})\n",
    "    rows_kto.append({\"prompt\": row[\"prompt\"], \"response\": row[\"rejected\"], \"label\": False})\n",
    "\n",
    "train_kto_df = pd.DataFrame(rows_kto)\n",
    "train_kto_df.to_csv(\"train_kto.csv\", index=False)\n",
    "print(f\"KTO dataset: {len(train_kto_df)} rows\")\n",
    "train_kto_df.head(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kto_config = yaml.safe_load(\"\"\"\n",
    "model_type: llm\n",
    "base_model: meta-llama/Llama-3.1-8B\n",
    "\n",
    "adapter:\n",
    "  type: lora\n",
    "  r: 16\n",
    "  alpha: 32\n",
    "  dropout: 0.05\n",
    "\n",
    "trainer:\n",
    "  type: kto\n",
    "  epochs: 1\n",
    "  learning_rate: 5.0e-7\n",
    "  batch_size: 2\n",
    "  gradient_accumulation_steps: 8\n",
    "  beta: 0.1\n",
    "  desirable_weight: 1.0\n",
    "  undesirable_weight: 1.0\n",
    "\n",
    "input_features:\n",
    "  - name: prompt\n",
    "    type: text\n",
    "\n",
    "output_features:\n",
    "  - name: response\n",
    "    type: text\n",
    "\n",
    "backend:\n",
    "  type: local\n",
    "\"\"\")\n",
    "\n",
    "kto_model = LudwigModel(config=kto_config, logging_level=logging.INFO)\n",
    "\n",
    "_, _, kto_output_dir = kto_model.train(\n",
    "    dataset=\"train_kto.csv\",\n",
    "    experiment_name=\"hh_rlhf_kto\",\n",
    "    output_directory=\"results_kto\",\n",
    ")\n",
    "\n",
    "print(f\"\\nKTO model saved to: {kto_output_dir}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save and upload\n",
    "\n",
    "After training, upload the aligned model to HuggingFace Hub to share it or deploy it to an endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Replace with your HuggingFace org/username and desired repo name\n",
    "HF_REPO = \"your-username/llama-3.1-8b-dpo-hh-rlhf\"\n",
    "\n",
    "# Upload via CLI:\n",
    "print(f\"ludwig upload hf_hub -r {HF_REPO} -m {output_directory}\")\n",
    "\n",
    "# Or via Python API:\n",
    "# LudwigModel.upload_to_hf_hub(HF_REPO, output_directory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment to actually run the upload:\n",
    "# !ludwig upload hf_hub -r {HF_REPO} -m {output_directory}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "- **Scale up**: Remove the `max_samples` caps and train on the full HH-RLHF dataset (~160k examples).\n",
    "- **Try ORPO**: Use `config_orpo.yaml` for single-stage SFT + alignment without a reference model.\n",
    "- **Iterate on beta**: Lower `beta` (e.g. 0.05) gives more aggressive alignment but risks reward hacking.\n",
    "- **Evaluation**: Use [MT-Bench](https://github.com/lm-sys/FastChat/tree/main/fastchat/llm_judge) or\n",
    "  [AlpacaEval](https://github.com/tatsu-lab/alpaca_eval) for rigorous alignment benchmarks.\n",
    "- **GRPO**: For tasks with a programmatic reward (math, code), see the GRPO trainer docs."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}