preference_tuning_with_dpo.ipynb

{
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.10.13",
      "mimetype": "text/x-python",
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "pygments_lexer": "ipython3",
      "nbconvert_exporter": "python",
      "file_extension": ".py"
    },
    "kaggle": {
      "accelerator": "none",
      "dataSources": [
        {
          "sourceId": 159748558,
          "sourceType": "kernelVersion"
        },
        {
          "sourceId": 165914575,
          "sourceType": "kernelVersion"
        },
        {
          "sourceId": 165914677,
          "sourceType": "kernelVersion"
        }
      ],
      "dockerImageVersionId": 30699,
      "isInternetEnabled": true,
      "language": "python",
      "sourceType": "notebook",
      "isGpuEnabled": false
    },
    "colab": {
      "provenance": [],
      "gpuType": "L4",
      "machine_shape": "hm"
    },
    "accelerator": "GPU"
  },
  "nbformat_minor": 0,
  "nbformat": 4,
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!pip install -q accelerate\n",
        "!pip install -q bitsandbytes\n",
        "!pip install -q peft\n",
        "!pip install datasets -q"
      ],
      "metadata": {
        "trusted": true,
        "id": "U6pVuxb8wNBP"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!nvidia-smi"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ms2XErQHPah6",
        "outputId": "1b582e42-aa68-4099-bd52-ca41f80a12c8"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Thu Jun  6 17:29:45 2024       \n",
            "+---------------------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 535.104.05             Driver Version: 535.104.05   CUDA Version: 12.2     |\n",
            "|-----------------------------------------+----------------------+----------------------+\n",
            "| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                                         |                      |               MIG M. |\n",
            "|=========================================+======================+======================|\n",
            "|   0  NVIDIA L4                      Off | 00000000:00:03.0 Off |                    0 |\n",
            "| N/A   68C    P8              19W /  72W |      1MiB / 23034MiB |      0%      Default |\n",
            "|                                         |                      |                  N/A |\n",
            "+-----------------------------------------+----------------------+----------------------+\n",
            "                                                                                         \n",
            "+---------------------------------------------------------------------------------------+\n",
            "| Processes:                                                                            |\n",
            "|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |\n",
            "|        ID   ID                                                             Usage      |\n",
            "|=======================================================================================|\n",
            "|  No running processes found                                                           |\n",
            "+---------------------------------------------------------------------------------------+\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import gc\n",
        "from datasets import load_dataset\n",
        "from types import SimpleNamespace\n",
        "\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "\n",
        "import torch\n",
        "from torch import nn, optim\n",
        "import torch.nn.functional as F\n",
        "from torch.nn.utils.rnn import pad_sequence\n",
        "\n",
        "from transformers import (\n",
        "    AutoTokenizer,\n",
        "    AutoModelForCausalLM,\n",
        "    BitsAndBytesConfig,\n",
        "    get_cosine_schedule_with_warmup\n",
        ")\n",
        "\n",
        "\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torch.cuda.amp import GradScaler, autocast\n",
        "\n",
        "os.environ['TOKENIZERS_PARALLELISM'] = 'false'"
      ],
      "metadata": {
        "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
        "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
        "execution": {
          "iopub.status.busy": "2024-05-15T07:04:45.196539Z",
          "iopub.execute_input": "2024-05-15T07:04:45.197037Z",
          "iopub.status.idle": "2024-05-15T07:04:51.743993Z",
          "shell.execute_reply.started": "2024-05-15T07:04:45.197000Z",
          "shell.execute_reply": "2024-05-15T07:04:51.742219Z"
        },
        "trusted": true,
        "id": "gRQrvhkKwNBR"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "cfg = {\n",
        "    'model_id': 'h2oai/h2o-danube-1.8b-sft',\n",
        "    'device': 'cuda' if torch.cuda.is_available() else 'cpu',\n",
        "    'seed': 1337,\n",
        "    'beta': 0.1,\n",
        "    'batch_size': 1,\n",
        "    'learning_rate': 1e-07,\n",
        "    'weight_decay': 1e-02,\n",
        "    'num_epochs': 1,\n",
        "    'warmup_steps': 0,\n",
        "    'accumulation_steps': 2,\n",
        "    'logging_steps': 50,\n",
        "}\n",
        "\n",
        "cfg = SimpleNamespace(**cfg)"
      ],
      "metadata": {
        "id": "5Sh9037vPNIh"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Create a set of prompts"
      ],
      "metadata": {
        "id": "mY9ckmNxwNBT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "toy_data = {\n",
        "    'prompt' : ['What is Ha Noi', '2 + 2 ='],\n",
        "}\n",
        "\n",
        "train = pd.DataFrame(toy_data)\n",
        "train"
      ],
      "metadata": {
        "trusted": true,
        "id": "jO7GSPrOwNBU",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 125
        },
        "outputId": "a231c6a6-c739-41f8-e9c1-9eebb3dcf056"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "           prompt\n",
              "0  What is Ha Noi\n",
              "1         2 + 2 ="
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-a0a3d0a9-94fe-4fbd-b954-fd315ff4519e\" class=\"colab-df-container\">\n",
              "    <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>prompt</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>What is Ha Noi</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>2 + 2 =</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-a0a3d0a9-94fe-4fbd-b954-fd315ff4519e')\"\n",
              "            title=\"Convert this dataframe to an interactive table.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "    <script>\n",
              "      const buttonEl =\n",
              "        document.querySelector('#df-a0a3d0a9-94fe-4fbd-b954-fd315ff4519e button.colab-df-convert');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      async function convertToInteractive(key) {\n",
              "        const element = document.querySelector('#df-a0a3d0a9-94fe-4fbd-b954-fd315ff4519e');\n",
              "        const dataTable =\n",
              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                    [key], {});\n",
              "        if (!dataTable) return;\n",
              "\n",
              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "          + ' to learn more about interactive tables.';\n",
              "        element.innerHTML = '';\n",
              "        dataTable['output_type'] = 'display_data';\n",
              "        await google.colab.output.renderOutput(dataTable, element);\n",
              "        const docLink = document.createElement('div');\n",
              "        docLink.innerHTML = docLinkHtml;\n",
              "        element.appendChild(docLink);\n",
              "      }\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "\n",
              "<div id=\"df-2ab244bc-f36f-41cd-97e8-5b2a5bc46650\">\n",
              "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-2ab244bc-f36f-41cd-97e8-5b2a5bc46650')\"\n",
              "            title=\"Suggest charts\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "     width=\"24px\">\n",
              "    <g>\n",
              "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
              "    </g>\n",
              "</svg>\n",
              "  </button>\n",
              "\n",
              "<style>\n",
              "  .colab-df-quickchart {\n",
              "      --bg-color: #E8F0FE;\n",
              "      --fill-color: #1967D2;\n",
              "      --hover-bg-color: #E2EBFA;\n",
              "      --hover-fill-color: #174EA6;\n",
              "      --disabled-fill-color: #AAA;\n",
              "      --disabled-bg-color: #DDD;\n",
              "  }\n",
              "\n",
              "  [theme=dark] .colab-df-quickchart {\n",
              "      --bg-color: #3B4455;\n",
              "      --fill-color: #D2E3FC;\n",
              "      --hover-bg-color: #434B5C;\n",
              "      --hover-fill-color: #FFFFFF;\n",
              "      --disabled-bg-color: #3B4455;\n",
              "      --disabled-fill-color: #666;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart {\n",
              "    background-color: var(--bg-color);\n",
              "    border: none;\n",
              "    border-radius: 50%;\n",
              "    cursor: pointer;\n",
              "    display: none;\n",
              "    fill: var(--fill-color);\n",
              "    height: 32px;\n",
              "    padding: 0;\n",
              "    width: 32px;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart:hover {\n",
              "    background-color: var(--hover-bg-color);\n",
              "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "    fill: var(--button-hover-fill-color);\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart-complete:disabled,\n",
              "  .colab-df-quickchart-complete:disabled:hover {\n",
              "    background-color: var(--disabled-bg-color);\n",
              "    fill: var(--disabled-fill-color);\n",
              "    box-shadow: none;\n",
              "  }\n",
              "\n",
              "  .colab-df-spinner {\n",
              "    border: 2px solid var(--fill-color);\n",
              "    border-color: transparent;\n",
              "    border-bottom-color: var(--fill-color);\n",
              "    animation:\n",
              "      spin 1s steps(1) infinite;\n",
              "  }\n",
              "\n",
              "  @keyframes spin {\n",
              "    0% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "      border-left-color: var(--fill-color);\n",
              "    }\n",
              "    20% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    30% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    40% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    60% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    80% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "    90% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "  }\n",
              "</style>\n",
              "\n",
              "  <script>\n",
              "    async function quickchart(key) {\n",
              "      const quickchartButtonEl =\n",
              "        document.querySelector('#' + key + ' button');\n",
              "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
              "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
              "      try {\n",
              "        const charts = await google.colab.kernel.invokeFunction(\n",
              "            'suggestCharts', [key], {});\n",
              "      } catch (error) {\n",
              "        console.error('Error during call to suggestCharts:', error);\n",
              "      }\n",
              "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
              "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
              "    }\n",
              "    (() => {\n",
              "      let quickchartButtonEl =\n",
              "        document.querySelector('#df-2ab244bc-f36f-41cd-97e8-5b2a5bc46650 button');\n",
              "      quickchartButtonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "    })();\n",
              "  </script>\n",
              "</div>\n",
              "\n",
              "  <div id=\"id_456fe433-2d2f-4310-aef4-2b76f4c7f889\">\n",
              "    <style>\n",
              "      .colab-df-generate {\n",
              "        background-color: #E8F0FE;\n",
              "        border: none;\n",
              "        border-radius: 50%;\n",
              "        cursor: pointer;\n",
              "        display: none;\n",
              "        fill: #1967D2;\n",
              "        height: 32px;\n",
              "        padding: 0 0 0 0;\n",
              "        width: 32px;\n",
              "      }\n",
              "\n",
              "      .colab-df-generate:hover {\n",
              "        background-color: #E2EBFA;\n",
              "        box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "        fill: #174EA6;\n",
              "      }\n",
              "\n",
              "      [theme=dark] .colab-df-generate {\n",
              "        background-color: #3B4455;\n",
              "        fill: #D2E3FC;\n",
              "      }\n",
              "\n",
              "      [theme=dark] .colab-df-generate:hover {\n",
              "        background-color: #434B5C;\n",
              "        box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "        filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "        fill: #FFFFFF;\n",
              "      }\n",
              "    </style>\n",
              "    <button class=\"colab-df-generate\" onclick=\"generateWithVariable('train')\"\n",
              "            title=\"Generate code using this dataframe.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "       width=\"24px\">\n",
              "    <path d=\"M7,19H8.4L18.45,9,17,7.55,7,17.6ZM5,21V16.75L18.45,3.32a2,2,0,0,1,2.83,0l1.4,1.43a1.91,1.91,0,0,1,.58,1.4,1.91,1.91,0,0,1-.58,1.4L9.25,21ZM18.45,9,17,7.55Zm-12,3A5.31,5.31,0,0,0,4.9,8.1,5.31,5.31,0,0,0,1,6.5,5.31,5.31,0,0,0,4.9,4.9,5.31,5.31,0,0,0,6.5,1,5.31,5.31,0,0,0,8.1,4.9,5.31,5.31,0,0,0,12,6.5,5.46,5.46,0,0,0,6.5,12Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "    <script>\n",
              "      (() => {\n",
              "      const buttonEl =\n",
              "        document.querySelector('#id_456fe433-2d2f-4310-aef4-2b76f4c7f889 button.colab-df-generate');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      buttonEl.onclick = () => {\n",
              "        google.colab.notebook.generateWithVariable('train');\n",
              "      }\n",
              "      })();\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "    </div>\n",
              "  </div>\n"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "dataframe",
              "variable_name": "train",
              "summary": "{\n  \"name\": \"train\",\n  \"rows\": 2,\n  \"fields\": [\n    {\n      \"column\": \"prompt\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 2,\n        \"samples\": [\n          \"2 + 2 =\",\n          \"What is Ha Noi\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}"
            }
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Understand chat template of Danube in HF\n"
      ],
      "metadata": {
        "id": "IIAz2H_iwNBV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Load tokenizer\n",
        "tokenizer = AutoTokenizer.from_pretrained(cfg.model_id)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-05-15T07:06:17.678079Z",
          "iopub.execute_input": "2024-05-15T07:06:17.678499Z",
          "iopub.status.idle": "2024-05-15T07:06:19.231363Z",
          "shell.execute_reply.started": "2024-05-15T07:06:17.678466Z",
          "shell.execute_reply": "2024-05-15T07:06:19.230108Z"
        },
        "trusted": true,
        "id": "56MWDh1gwNBV",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "149e52de-1a95-4c8d-f33e-c72bffae5208"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/huggingface_hub/utils/_token.py:89: UserWarning: \n",
            "The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
            "To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n",
            "You will be able to reuse this secret in all of your notebooks.\n",
            "Please note that authentication is recommended but still optional to access public models or datasets.\n",
            "  warnings.warn(\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "messages = [\n",
        "    {\"role\": \"user\", \"content\": \"How are you?\"},\n",
        "    {\"role\": \"assistant\", \"content\": \"I'm an AI, so I don't have feelings, but I'm here and ready to help you! How can I assist you today?\"},\n",
        "]"
      ],
      "metadata": {
        "trusted": true,
        "id": "v7q7q4kfwNBW"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "encoded_input_ids = tokenizer.apply_chat_template(messages)\n",
        "tokenizer.decode(encoded_input_ids)"
      ],
      "metadata": {
        "trusted": true,
        "id": "X006zkBnwNBW",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "outputId": "af72327d-087e-4d8f-85b7-a7449998f548"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "\"<|prompt|>How are you?</s> <|answer|>I'm an AI, so I don't have feelings, but I'm here and ready to help you! How can I assist you today?</s>\""
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Let's ask H20 Danube, a supervised fine-tuning model"
      ],
      "metadata": {
        "id": "Od49b2tKwNBX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "reference_model = AutoModelForCausalLM.from_pretrained(\n",
        "    cfg.model_id,\n",
        "    torch_dtype=torch.bfloat16,\n",
        "    device_map='auto'\n",
        ")\n",
        "\n",
        "reference_model.eval();"
      ],
      "metadata": {
        "trusted": true,
        "id": "4lxFdOWNwNBX"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_messages(prompt):\n",
        "    # Function to generate messages with user prompt\n",
        "    return [\n",
        "        {'role': 'user', 'content': prompt},\n",
        "        {'role': 'assistant', 'content': ''},\n",
        "    ]\n",
        "\n",
        "def extract_response(question_response):\n",
        "    # Function to extract response from the model output\n",
        "    response = question_response.strip().split('<|answer|>')[-1]  # Splitting by <|answer|> string to get the last part\n",
        "    return response"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-05-15T07:06:36.411764Z",
          "iopub.execute_input": "2024-05-15T07:06:36.412262Z",
          "iopub.status.idle": "2024-05-15T07:06:36.419670Z",
          "shell.execute_reply.started": "2024-05-15T07:06:36.412227Z",
          "shell.execute_reply": "2024-05-15T07:06:36.418559Z"
        },
        "trusted": true,
        "id": "zLJiTOIdwNBY"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "prompts = train['prompt']\n",
        "responses = []\n",
        "\n",
        "for prompt in prompts:\n",
        "    # Generate input tokens for the model\n",
        "    input_ids = tokenizer.apply_chat_template(\n",
        "        get_messages(prompt),\n",
        "        return_tensors='pt'\n",
        "    ).to(cfg.device)\n",
        "\n",
        "    input_ids = input_ids[:, :-1] # Slice to remove the </s> that is automatically added when using tokenizer.apply_chat_template()\n",
        "\n",
        "    # Generate response using the model\n",
        "    with torch.no_grad():\n",
        "        reference_model.eval()\n",
        "        output_ids = reference_model.generate(\n",
        "            input_ids,\n",
        "            max_new_tokens=1024,\n",
        "            pad_token_id=tokenizer.pad_token_id\n",
        "        )[0]\n",
        "\n",
        "    # Decode model output and extract response\n",
        "    question_response = tokenizer.decode(output_ids, skip_special_tokens=True)\n",
        "\n",
        "    responses.append(\n",
        "        extract_response(question_response)\n",
        "    )"
      ],
      "metadata": {
        "trusted": true,
        "id": "95F21ZWvwNBY"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for idx, response in enumerate(responses):\n",
        "    print(f'### Response for question {idx+1}:\\n', response.strip(), '\\n')"
      ],
      "metadata": {
        "trusted": true,
        "id": "s57ofx1gwNBY",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "56158e50-e238-406d-ccf1-9e858cc2dbf5"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "### Response for question 1:\n",
            " Ha Noi, also known as Hanoi, is the capital city of Vietnam. It is located in the north-central part of the country and is known for its rich history, culture, and architecture. The city is home to many historic landmarks such as the Ho Chi Minh Mausoleum, the Temple of Literature, and the Hoan Kiem Lake. Ha Noi is also a hub for business and commerce, with many international companies having offices or branches in the city. Overall, Ha Noi is an important center of government, education, and tourism in Vietnam. \n",
            "\n",
            "### Response for question 2:\n",
            " To solve the equation 2 + 2 =, we first add the two numbers on the left side of the equation:\n",
            "\n",
            "2 + 2 = 4\n",
            "\n",
            "So, the equation becomes:\n",
            "\n",
            "4 =\n",
            "\n",
            "Now, to find the value of x that satisfies this equation, we can set it equal to 4 and solve for x.\n",
            "\n",
            "4 = 4\n",
            "\n",
            "Therefore, the value of unknown variable x is 4. \n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Get a preference dataset"
      ],
      "metadata": {
        "id": "sUc-ykwrwNBY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "chosen_responses = ['Hanoi is a city in Viet Nam.', '2 + 2 = 4']\n",
        "rejected_responses = responses # Answer we don't like the current responses from Danube\n",
        "\n",
        "train['chosen'] = chosen_responses\n",
        "train['rejected'] = rejected_responses\n",
        "train"
      ],
      "metadata": {
        "trusted": true,
        "id": "2vhAcVi4wNBY",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 125
        },
        "outputId": "2d122663-b6eb-4c7a-abe6-2b394fdddba4"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "           prompt                        chosen  \\\n",
              "0  What is Ha Noi  Hanoi is a city in Viet Nam.   \n",
              "1         2 + 2 =                     2 + 2 = 4   \n",
              "\n",
              "                                            rejected  \n",
              "0   Ha Noi, also known as Hanoi, is the capital c...  \n",
              "1   To solve the equation 2 + 2 =, we first add t...  "
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-da0ba5a7-994f-4767-bd15-0322cd39ec69\" class=\"colab-df-container\">\n",
              "    <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>prompt</th>\n",
              "      <th>chosen</th>\n",
              "      <th>rejected</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>What is Ha Noi</td>\n",
              "      <td>Hanoi is a city in Viet Nam.</td>\n",
              "      <td>Ha Noi, also known as Hanoi, is the capital c...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>2 + 2 =</td>\n",
              "      <td>2 + 2 = 4</td>\n",
              "      <td>To solve the equation 2 + 2 =, we first add t...</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-da0ba5a7-994f-4767-bd15-0322cd39ec69')\"\n",
              "            title=\"Convert this dataframe to an interactive table.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "    <script>\n",
              "      const buttonEl =\n",
              "        document.querySelector('#df-da0ba5a7-994f-4767-bd15-0322cd39ec69 button.colab-df-convert');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      async function convertToInteractive(key) {\n",
              "        const element = document.querySelector('#df-da0ba5a7-994f-4767-bd15-0322cd39ec69');\n",
              "        const dataTable =\n",
              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                    [key], {});\n",
              "        if (!dataTable) return;\n",
              "\n",
              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "          + ' to learn more about interactive tables.';\n",
              "        element.innerHTML = '';\n",
              "        dataTable['output_type'] = 'display_data';\n",
              "        await google.colab.output.renderOutput(dataTable, element);\n",
              "        const docLink = document.createElement('div');\n",
              "        docLink.innerHTML = docLinkHtml;\n",
              "        element.appendChild(docLink);\n",
              "      }\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "\n",
              "<div id=\"df-f491921b-bfb4-40f5-bb25-2f698372a1a1\">\n",
              "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-f491921b-bfb4-40f5-bb25-2f698372a1a1')\"\n",
              "            title=\"Suggest charts\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "     width=\"24px\">\n",
              "    <g>\n",
              "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
              "    </g>\n",
              "</svg>\n",
              "  </button>\n",
              "\n",
              "<style>\n",
              "  .colab-df-quickchart {\n",
              "      --bg-color: #E8F0FE;\n",
              "      --fill-color: #1967D2;\n",
              "      --hover-bg-color: #E2EBFA;\n",
              "      --hover-fill-color: #174EA6;\n",
              "      --disabled-fill-color: #AAA;\n",
              "      --disabled-bg-color: #DDD;\n",
              "  }\n",
              "\n",
              "  [theme=dark] .colab-df-quickchart {\n",
              "      --bg-color: #3B4455;\n",
              "      --fill-color: #D2E3FC;\n",
              "      --hover-bg-color: #434B5C;\n",
              "      --hover-fill-color: #FFFFFF;\n",
              "      --disabled-bg-color: #3B4455;\n",
              "      --disabled-fill-color: #666;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart {\n",
              "    background-color: var(--bg-color);\n",
              "    border: none;\n",
              "    border-radius: 50%;\n",
              "    cursor: pointer;\n",
              "    display: none;\n",
              "    fill: var(--fill-color);\n",
              "    height: 32px;\n",
              "    padding: 0;\n",
              "    width: 32px;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart:hover {\n",
              "    background-color: var(--hover-bg-color);\n",
              "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "    fill: var(--button-hover-fill-color);\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart-complete:disabled,\n",
              "  .colab-df-quickchart-complete:disabled:hover {\n",
              "    background-color: var(--disabled-bg-color);\n",
              "    fill: var(--disabled-fill-color);\n",
              "    box-shadow: none;\n",
              "  }\n",
              "\n",
              "  .colab-df-spinner {\n",
              "    border: 2px solid var(--fill-color);\n",
              "    border-color: transparent;\n",
              "    border-bottom-color: var(--fill-color);\n",
              "    animation:\n",
              "      spin 1s steps(1) infinite;\n",
              "  }\n",
              "\n",
              "  @keyframes spin {\n",
              "    0% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "      border-left-color: var(--fill-color);\n",
              "    }\n",
              "    20% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    30% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    40% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    60% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    80% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "    90% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "  }\n",
              "</style>\n",
              "\n",
              "  <script>\n",
              "    async function quickchart(key) {\n",
              "      const quickchartButtonEl =\n",
              "        document.querySelector('#' + key + ' button');\n",
              "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
              "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
              "      try {\n",
              "        const charts = await google.colab.kernel.invokeFunction(\n",
              "            'suggestCharts', [key], {});\n",
              "      } catch (error) {\n",
              "        console.error('Error during call to suggestCharts:', error);\n",
              "      }\n",
              "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
              "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
              "    }\n",
              "    (() => {\n",
              "      let quickchartButtonEl =\n",
              "        document.querySelector('#df-f491921b-bfb4-40f5-bb25-2f698372a1a1 button');\n",
              "      quickchartButtonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "    })();\n",
              "  </script>\n",
              "</div>\n",
              "\n",
              "  <div id=\"id_b4d51fa1-da75-4ba0-8938-fcefcfed762e\">\n",
              "    <style>\n",
              "      .colab-df-generate {\n",
              "        background-color: #E8F0FE;\n",
              "        border: none;\n",
              "        border-radius: 50%;\n",
              "        cursor: pointer;\n",
              "        display: none;\n",
              "        fill: #1967D2;\n",
              "        height: 32px;\n",
              "        padding: 0 0 0 0;\n",
              "        width: 32px;\n",
              "      }\n",
              "\n",
              "      .colab-df-generate:hover {\n",
              "        background-color: #E2EBFA;\n",
              "        box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "        fill: #174EA6;\n",
              "      }\n",
              "\n",
              "      [theme=dark] .colab-df-generate {\n",
              "        background-color: #3B4455;\n",
              "        fill: #D2E3FC;\n",
              "      }\n",
              "\n",
              "      [theme=dark] .colab-df-generate:hover {\n",
              "        background-color: #434B5C;\n",
              "        box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "        filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "        fill: #FFFFFF;\n",
              "      }\n",
              "    </style>\n",
              "    <button class=\"colab-df-generate\" onclick=\"generateWithVariable('train')\"\n",
              "            title=\"Generate code using this dataframe.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "       width=\"24px\">\n",
              "    <path d=\"M7,19H8.4L18.45,9,17,7.55,7,17.6ZM5,21V16.75L18.45,3.32a2,2,0,0,1,2.83,0l1.4,1.43a1.91,1.91,0,0,1,.58,1.4,1.91,1.91,0,0,1-.58,1.4L9.25,21ZM18.45,9,17,7.55Zm-12,3A5.31,5.31,0,0,0,4.9,8.1,5.31,5.31,0,0,0,1,6.5,5.31,5.31,0,0,0,4.9,4.9,5.31,5.31,0,0,0,6.5,1,5.31,5.31,0,0,0,8.1,4.9,5.31,5.31,0,0,0,12,6.5,5.46,5.46,0,0,0,6.5,12Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "    <script>\n",
              "      (() => {\n",
              "      const buttonEl =\n",
              "        document.querySelector('#id_b4d51fa1-da75-4ba0-8938-fcefcfed762e button.colab-df-generate');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      buttonEl.onclick = () => {\n",
              "        google.colab.notebook.generateWithVariable('train');\n",
              "      }\n",
              "      })();\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "    </div>\n",
              "  </div>\n"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "dataframe",
              "variable_name": "train",
              "summary": "{\n  \"name\": \"train\",\n  \"rows\": 2,\n  \"fields\": [\n    {\n      \"column\": \"prompt\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 2,\n        \"samples\": [\n          \"2 + 2 =\",\n          \"What is Ha Noi\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"chosen\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 2,\n        \"samples\": [\n          \"2 + 2 = 4\",\n          \"Hanoi is a city in Viet Nam.\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"rejected\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 2,\n        \"samples\": [\n          \" To solve the equation 2 + 2 =, we first add the two numbers on the left side of the equation:\\n\\n2 + 2 = 4\\n\\nSo, the equation becomes:\\n\\n4 =\\n\\nNow, to find the value of x that satisfies this equation, we can set it equal to 4 and solve for x.\\n\\n4 = 4\\n\\nTherefore, the value of unknown variable x is 4.\",\n          \" Ha Noi, also known as Hanoi, is the capital city of Vietnam. It is located in the north-central part of the country and is known for its rich history, culture, and architecture. The city is home to many historic landmarks such as the Ho Chi Minh Mausoleum, the Temple of Literature, and the Hoan Kiem Lake. Ha Noi is also a hub for business and commerce, with many international companies having offices or branches in the city. Overall, Ha Noi is an important center of government, education, and tourism in Vietnam.\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}"
            }
          },
          "metadata": {},
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Create dataloader for training"
      ],
      "metadata": {
        "id": "2-Mb7WxTwNBZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def encode(prompt, response):\n",
        "\n",
        "    eos_token = '</s>'\n",
        "    start_inst_token = '<|prompt|>'\n",
        "    end_inst_token = '<|answer|>'\n",
        "\n",
        "    prompt = start_inst_token + prompt + eos_token + end_inst_token\n",
        "\n",
        "    prompt_ids = tokenizer(\n",
        "        prompt,\n",
        "        add_special_tokens=False,\n",
        "        return_tensors='pt'\n",
        "    ).input_ids\n",
        "\n",
        "    response_ids = tokenizer(\n",
        "        response,\n",
        "        add_special_tokens=False,\n",
        "        return_tensors='pt'\n",
        "    ).input_ids\n",
        "\n",
        "    eos_id = tokenizer(\n",
        "        eos_token,\n",
        "        add_special_tokens=False,\n",
        "        return_tensors='pt'\n",
        "    ).input_ids\n",
        "\n",
        "    input_ids = torch.cat(\n",
        "        [prompt_ids, response_ids, eos_id],\n",
        "        dim=1\n",
        "    ).view(-1)\n",
        "\n",
        "    labels = torch.cat(\n",
        "        [\n",
        "            torch.full_like(\n",
        "                prompt_ids[:, :-1],\n",
        "                tokenizer.eos_token_id\n",
        "            ), # create a tensor with full eos_token_id (2)\n",
        "            prompt_ids[:, -1].view(1, -1), # slice to except the \">\" token in \"<|answer|>\" before the first token of teh response. Because this contains the predicted probas of the first token of the response\n",
        "            response_ids, # token ids of the esponse\n",
        "            eos_id # end token </s>\n",
        "        ],\n",
        "        dim=1\n",
        "    ).view(-1)\n",
        "\n",
        "    return input_ids, labels"
      ],
      "metadata": {
        "trusted": true,
        "id": "4hipygghwNBZ"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class CustomDataset(Dataset):\n",
        "    def __init__(self, prompts, chosen_responses, rejected_responses):\n",
        "        self.prompts = prompts\n",
        "        self.chosen_responses = chosen_responses\n",
        "        self.rejected_responses = rejected_responses\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.prompts)\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        prompt = self.prompts[index]\n",
        "        chosen_response = self.chosen_responses[index]\n",
        "        rejected_response = self.rejected_responses[index]\n",
        "\n",
        "        chosen_input_ids, chosen_labels = encode(prompt, chosen_response)\n",
        "        rejected_input_ids, rejected_labels = encode(prompt, rejected_response)\n",
        "\n",
        "        return {\n",
        "            'chosen_input_ids': chosen_input_ids,\n",
        "            'chosen_labels': chosen_labels,\n",
        "            'rejected_input_ids': rejected_input_ids,\n",
        "            'rejected_labels': rejected_labels\n",
        "        }\n",
        "\n",
        "def custom_collate(data):\n",
        "    padded_data = {}\n",
        "    data = pd.DataFrame(data)\n",
        "    keys = data.columns\n",
        "\n",
        "    for key in keys:\n",
        "        padded_data[key] = pad_sequence(\n",
        "            sequences=data[key].values.tolist(),\n",
        "            batch_first=True,\n",
        "            padding_value=tokenizer.eos_token_id\n",
        "        )\n",
        "\n",
        "    return padded_data"
      ],
      "metadata": {
        "trusted": true,
        "id": "AdZYAxHwwNBZ"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "prompts = train['prompt'].values.tolist()\n",
        "chosen_responses = train['chosen'].values.tolist()\n",
        "rejected_responses = train['rejected'].values.tolist()\n",
        "\n",
        "torch.manual_seed(seed=cfg.seed)\n",
        "train_dataset = CustomDataset(\n",
        "    prompts=prompts,\n",
        "    chosen_responses=chosen_responses,\n",
        "    rejected_responses=rejected_responses\n",
        ")\n",
        "\n",
        "train_dataloader = DataLoader(\n",
        "    dataset=train_dataset,\n",
        "    batch_size=cfg.batch_size,\n",
        "    shuffle=True,\n",
        "    drop_last=True,\n",
        "    collate_fn=custom_collate\n",
        ")"
      ],
      "metadata": {
        "trusted": true,
        "id": "hqo80zCawNBZ"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Implement DPO loss"
      ],
      "metadata": {
        "id": "HeyzqDVMwNBZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def get_batch_logps(logits, labels, pad_id_value):\n",
        "    labels = labels.clone()\n",
        "    loss_mask = labels != pad_id_value\n",
        "    per_token_logps = torch.gather(logits.log_softmax(dim=-1), dim=2, index=labels.unsqueeze(2)).squeeze(2)\n",
        "    return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1)\n",
        "\n",
        "def dpo_loss(\n",
        "    policy_chosen_logps,\n",
        "    policy_rejected_logps,\n",
        "    reference_chosen_logps,\n",
        "    reference_rejected_logps,\n",
        "    beta\n",
        "):\n",
        "    policy_logratios = policy_chosen_logps - policy_rejected_logps\n",
        "    reference_log_ratios = reference_chosen_logps - reference_rejected_logps\n",
        "    losses = (-1) * F.logsigmoid((beta*policy_logratios) - beta*reference_log_ratios)\n",
        "    return losses.mean()\n",
        "\n",
        "def get_reference_logits(reference_model, chosen_pair, rejected_pair):\n",
        "    reference_model.eval()\n",
        "    with torch.no_grad():\n",
        "        reference_chosen_logits = reference_model(chosen_pair).logits\n",
        "        reference_rejected_logits = reference_model(rejected_pair).logits\n",
        "    return reference_chosen_logits, reference_rejected_logits"
      ],
      "metadata": {
        "trusted": true,
        "id": "L4EBb8rmwNBZ"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Fine-tuning the tiny toy dataset with DPO loss"
      ],
      "metadata": {
        "id": "FtMMdWKCwNBa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Policy(nn.Module):\n",
        "    def __init__(self, cfg):\n",
        "        super(Policy, self).__init__()\n",
        "\n",
        "        self.backbone = AutoModelForCausalLM.from_pretrained(\n",
        "            cfg.model_id,\n",
        "            low_cpu_mem_usage=True,\n",
        "            device_map='cuda',\n",
        "            torch_dtype=torch.bfloat16,\n",
        "            use_cache=False\n",
        "        )\n",
        "\n",
        "        self.backbone.gradient_checkpointing_enable()\n",
        "\n",
        "    def forward(self, input_ids):\n",
        "        logits = self.backbone(input_ids).logits\n",
        "        return logits"
      ],
      "metadata": {
        "trusted": true,
        "id": "sQMwwXtpwNBa"
      },
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = Policy(cfg)\n",
        "\n",
        "optimizer = optim.AdamW(\n",
        "    model.parameters(),\n",
        "    lr=cfg.learning_rate,\n",
        "    weight_decay=cfg.weight_decay\n",
        ")\n",
        "\n",
        "scheduler = get_cosine_schedule_with_warmup(\n",
        "    optimizer=optimizer,\n",
        "    num_warmup_steps=cfg.warmup_steps,\n",
        "    num_training_steps=cfg.num_epochs*len(train_dataloader)\n",
        ")\n",
        "\n",
        "\n",
        "scaler = GradScaler()\n",
        "\n",
        "for epoch in range(cfg.num_epochs):\n",
        "\n",
        "    for batch_idx, batch in enumerate(train_dataloader):\n",
        "\n",
        "        model.train()\n",
        "\n",
        "        chosen_input_ids = batch['chosen_input_ids'].long().to(cfg.device)\n",
        "        chosen_labels = batch['chosen_labels'].long().to(cfg.device)\n",
        "        rejected_input_ids = batch['rejected_input_ids'].long().to(cfg.device)\n",
        "        rejected_labels = batch['rejected_labels'].long().to(cfg.device)\n",
        "\n",
        "        reference_chosen_logits, reference_rejected_logits = get_reference_logits(\n",
        "            reference_model,\n",
        "            chosen_input_ids,\n",
        "            rejected_input_ids\n",
        "        )\n",
        "\n",
        "        policy_chosen_logits = model(chosen_input_ids)\n",
        "        policy_rejected_logits = model(rejected_input_ids)\n",
        "\n",
        "        policy_chosen_logps = get_batch_logps(\n",
        "            policy_chosen_logits,\n",
        "            chosen_labels,\n",
        "            pad_id_value=tokenizer.eos_token_id\n",
        "        )\n",
        "\n",
        "        policy_rejected_logps = get_batch_logps(\n",
        "            policy_rejected_logits,\n",
        "            rejected_labels,\n",
        "            pad_id_value=tokenizer.eos_token_id\n",
        "        )\n",
        "\n",
        "        reference_chosen_logps = get_batch_logps(\n",
        "            reference_chosen_logits,\n",
        "            chosen_labels,\n",
        "            pad_id_value=tokenizer.eos_token_id\n",
        "        )\n",
        "\n",
        "        reference_rejected_logps = get_batch_logps(\n",
        "            reference_rejected_logits,\n",
        "            rejected_labels,\n",
        "            pad_id_value=tokenizer.eos_token_id\n",
        "        )\n",
        "\n",
        "        loss = dpo_loss(\n",
        "            policy_chosen_logps,\n",
        "            policy_rejected_logps,\n",
        "            reference_chosen_logps,\n",
        "            reference_rejected_logps,\n",
        "            beta=cfg.beta\n",
        "        ) / cfg.accumulation_steps\n",
        "\n",
        "        print(\n",
        "            f'Epoch: {epoch+1} / {cfg.num_epochs}'\n",
        "            f'| Batch: {batch_idx}/{len(train_dataloader)}'\n",
        "            f'| DPO Loss: {loss.item():.8f}'\n",
        "        )\n",
        "\n",
        "        loss.backward()\n",
        "\n",
        "        if not (batch_idx % cfg.accumulation_steps):\n",
        "            optimizer.step()\n",
        "            scheduler.step()\n",
        "            optimizer.zero_grad()"
      ],
      "metadata": {
        "trusted": true,
        "id": "02dtGZO5wNBa",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ecd9e182-c37e-464f-bf26-3c52f52c7c39"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: 1 / 1| Batch: 0/2| DPO Loss: 0.34657359\n",
            "Epoch: 1 / 1| Batch: 1/2| DPO Loss: 0.34611285\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Examine the model after fine-tuning with DPO\n",
        "Note: 2 rows of data seem to provide no improvement, but just for testing. We will fine-tune with a larger preference dataset."
      ],
      "metadata": {
        "id": "To8YezHwwNBa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "test_prompts = ['What is Hanoi', '2 + 2 =']\n",
        "responses = []\n",
        "\n",
        "for prompt in test_prompts:\n",
        "\n",
        "    input_ids = tokenizer.apply_chat_template(\n",
        "        get_messages(prompt),\n",
        "        return_tensors='pt'\n",
        "    ).to(cfg.device)\n",
        "\n",
        "    input_ids = input_ids[:, :-1] # Slice to remove the </s> that is automatically added when using tokenizer.apply_chat_template()\n",
        "\n",
        "    with torch.no_grad():\n",
        "        model.eval()\n",
        "        output_ids = model.backbone.generate(\n",
        "                input_ids=input_ids,\n",
        "                max_new_tokens=512,\n",
        "                pad_token_id=tokenizer.pad_token_id\n",
        "        )[0]\n",
        "\n",
        "    question_response = tokenizer.decode(output_ids, skip_special_tokens=True)\n",
        "    response = extract_response(question_response)\n",
        "\n",
        "    print('Prompt:', tokenizer.decode(input_ids[0]))\n",
        "    print('Response:', response)\n",
        "    print('\\n')"
      ],
      "metadata": {
        "trusted": true,
        "id": "PVRy6-xSwNBa",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5fc6a2ba-3396-4701-894c-412ce0bda1e4"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Prompt: <|prompt|>What is Hanoi</s> <|answer|>\n",
            "Response:  Hanoi is the capital city of Vietnam. It's a big city with lots of people and buildings, but it's also very old and has many interesting things to see. Some famous places in Hanoi are the Ho Chi Minh Mausoleum, the Temple of Literature, and the Hoan Kiem Lake.\n",
            "\n",
            "\n",
            "Prompt: <|prompt|>2 + 2 =</s> <|answer|>\n",
            "Response:  To solve the equation, I will follow these steps:\n",
            "\n",
            "Step 1: Add the numbers on the left side of the equation.\n",
            "2 + 2 = 4\n",
            "\n",
            "So, the equation becomes:\n",
            "4 =\n",
            "\n",
            "Step 2: Since there is no unknown variable in this equation, we can conclude that the value of the unknown variable x is not needed to determine the solution.\n",
            "\n",
            "Therefore, the answer is 4.\n",
            "\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Let's train on a real preference dataset\n",
        "\n",
        "I use **ultrafeedback** dataset on HuggingFace: https://huggingface.co/datasets/argilla/ultrafeedback-binarized-preferences"
      ],
      "metadata": {
        "id": "TZp-I6czwNBb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "data_path = 'argilla/ultrafeedback-binarized-preferences'\n",
        "\n",
        "# load in data\n",
        "train = load_dataset(data_path)['train'].to_pandas()\n",
        "\n",
        "# Select important columns\n",
        "train = train[['instruction', 'chosen_response', 'rejected_response']]\n",
        "\n",
        "# rename columns for consistent with our above format\n",
        "train.columns = ['prompt', 'chosen', 'rejected']\n",
        "\n",
        "# Sample about 500 training examples for fine-tuning\n",
        "train = train.sample(500, random_state=1).reset_index(drop=True)\n",
        "train"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-05-15T07:17:03.653720Z",
          "iopub.execute_input": "2024-05-15T07:17:03.654271Z",
          "iopub.status.idle": "2024-05-15T07:17:06.505896Z",
          "shell.execute_reply.started": "2024-05-15T07:17:03.654234Z",
          "shell.execute_reply": "2024-05-15T07:17:06.504664Z"
        },
        "trusted": true,
        "id": "KRkKdYF9wNBb",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 424
        },
        "outputId": "1c1b3d1c-2d61-4a1b-912a-8dee24af6ead"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "                                                prompt  \\\n",
              "0    You will be given a definition of a task first...   \n",
              "1    const [upgradePrimeToPlusMutation, { isUpgrade...   \n",
              "2    I work in Typescript and use Google Data Stora...   \n",
              "3    TASK DEFINITION: In this task, you are given a...   \n",
              "4    Develop an algorithm that can detect a given t...   \n",
              "..                                                 ...   \n",
              "495  What ingredients are required to make the cook...   \n",
              "496  Upwork Bidding: \\n\\nI'm bidding on this Upwork...   \n",
              "497  In what ways does the World Health Organizatio...   \n",
              "498  Develop a Ruby on Rails application that facil...   \n",
              "499  What will happen if you drive barefoot in the ...   \n",
              "\n",
              "                                                chosen  \\\n",
              "0                               [256, 512, 4096, 4096]   \n",
              "1    It appears that the code you've provided is a ...   \n",
              "2    Sure, I can help you with that. First, let's d...   \n",
              "3    Sure, I'd be happy to help! Here's the transla...   \n",
              "4    Certainly, I can help you develop an algorithm...   \n",
              "..                                                 ...   \n",
              "495  To make the cookie dough texture for M&M Choco...   \n",
              "496  Subject: Top-Notch SEO Services for Your Busin...   \n",
              "497  The World Health Organization (WHO) collaborat...   \n",
              "498  To develop a Ruby on Rails application that fa...   \n",
              "499  In the United Kingdom, it is generally not for...   \n",
              "\n",
              "                                              rejected  \n",
              "0    Sure, I'd be happy to help! Here's the definit...  \n",
              "1    It's essential to test various scenarios to en...  \n",
              "2    Sure, here's an example of a generic base clas...  \n",
              "3    همراحم از زمان محل را به نظر می رسند، من بنا ب...  \n",
              "4    As a helpful and respectful AI assistant, I mu...  \n",
              "..                                                 ...  \n",
              "495  No ingredients are required. This is not a rec...  \n",
              "496  Dear [Client Name],\\n\\nI hope this message fin...  \n",
              "497  The World Health Organization (WHO) collaborat...  \n",
              "498  Great! I'm here to assist you in creating a Ru...  \n",
              "499  Driving barefoot or without proper footwear is...  \n",
              "\n",
              "[500 rows x 3 columns]"
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-ec79ec12-ffa4-4661-8410-a596105ac5a5\" class=\"colab-df-container\">\n",
              "    <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>prompt</th>\n",
              "      <th>chosen</th>\n",
              "      <th>rejected</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>You will be given a definition of a task first...</td>\n",
              "      <td>[256, 512, 4096, 4096]</td>\n",
              "      <td>Sure, I'd be happy to help! Here's the definit...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>const [upgradePrimeToPlusMutation, { isUpgrade...</td>\n",
              "      <td>It appears that the code you've provided is a ...</td>\n",
              "      <td>It's essential to test various scenarios to en...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>I work in Typescript and use Google Data Stora...</td>\n",
              "      <td>Sure, I can help you with that. First, let's d...</td>\n",
              "      <td>Sure, here's an example of a generic base clas...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>TASK DEFINITION: In this task, you are given a...</td>\n",
              "      <td>Sure, I'd be happy to help! Here's the transla...</td>\n",
              "      <td>همراحم از زمان محل را به نظر می رسند، من بنا ب...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>Develop an algorithm that can detect a given t...</td>\n",
              "      <td>Certainly, I can help you develop an algorithm...</td>\n",
              "      <td>As a helpful and respectful AI assistant, I mu...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>...</th>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>495</th>\n",
              "      <td>What ingredients are required to make the cook...</td>\n",
              "      <td>To make the cookie dough texture for M&amp;M Choco...</td>\n",
              "      <td>No ingredients are required. This is not a rec...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>496</th>\n",
              "      <td>Upwork Bidding: \\n\\nI'm bidding on this Upwork...</td>\n",
              "      <td>Subject: Top-Notch SEO Services for Your Busin...</td>\n",
              "      <td>Dear [Client Name],\\n\\nI hope this message fin...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>497</th>\n",
              "      <td>In what ways does the World Health Organizatio...</td>\n",
              "      <td>The World Health Organization (WHO) collaborat...</td>\n",
              "      <td>The World Health Organization (WHO) collaborat...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>498</th>\n",
              "      <td>Develop a Ruby on Rails application that facil...</td>\n",
              "      <td>To develop a Ruby on Rails application that fa...</td>\n",
              "      <td>Great! I'm here to assist you in creating a Ru...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>499</th>\n",
              "      <td>What will happen if you drive barefoot in the ...</td>\n",
              "      <td>In the United Kingdom, it is generally not for...</td>\n",
              "      <td>Driving barefoot or without proper footwear is...</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "<p>500 rows × 3 columns</p>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-ec79ec12-ffa4-4661-8410-a596105ac5a5')\"\n",
              "            title=\"Convert this dataframe to an interactive table.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "    <script>\n",
              "      const buttonEl =\n",
              "        document.querySelector('#df-ec79ec12-ffa4-4661-8410-a596105ac5a5 button.colab-df-convert');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      async function convertToInteractive(key) {\n",
              "        const element = document.querySelector('#df-ec79ec12-ffa4-4661-8410-a596105ac5a5');\n",
              "        const dataTable =\n",
              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                    [key], {});\n",
              "        if (!dataTable) return;\n",
              "\n",
              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "          + ' to learn more about interactive tables.';\n",
              "        element.innerHTML = '';\n",
              "        dataTable['output_type'] = 'display_data';\n",
              "        await google.colab.output.renderOutput(dataTable, element);\n",
              "        const docLink = document.createElement('div');\n",
              "        docLink.innerHTML = docLinkHtml;\n",
              "        element.appendChild(docLink);\n",
              "      }\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "\n",
              "<div id=\"df-aee69069-8b48-4983-903f-3063f90b2243\">\n",
              "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-aee69069-8b48-4983-903f-3063f90b2243')\"\n",
              "            title=\"Suggest charts\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "     width=\"24px\">\n",
              "    <g>\n",
              "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
              "    </g>\n",
              "</svg>\n",
              "  </button>\n",
              "\n",
              "<style>\n",
              "  .colab-df-quickchart {\n",
              "      --bg-color: #E8F0FE;\n",
              "      --fill-color: #1967D2;\n",
              "      --hover-bg-color: #E2EBFA;\n",
              "      --hover-fill-color: #174EA6;\n",
              "      --disabled-fill-color: #AAA;\n",
              "      --disabled-bg-color: #DDD;\n",
              "  }\n",
              "\n",
              "  [theme=dark] .colab-df-quickchart {\n",
              "      --bg-color: #3B4455;\n",
              "      --fill-color: #D2E3FC;\n",
              "      --hover-bg-color: #434B5C;\n",
              "      --hover-fill-color: #FFFFFF;\n",
              "      --disabled-bg-color: #3B4455;\n",
              "      --disabled-fill-color: #666;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart {\n",
              "    background-color: var(--bg-color);\n",
              "    border: none;\n",
              "    border-radius: 50%;\n",
              "    cursor: pointer;\n",
              "    display: none;\n",
              "    fill: var(--fill-color);\n",
              "    height: 32px;\n",
              "    padding: 0;\n",
              "    width: 32px;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart:hover {\n",
              "    background-color: var(--hover-bg-color);\n",
              "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "    fill: var(--button-hover-fill-color);\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart-complete:disabled,\n",
              "  .colab-df-quickchart-complete:disabled:hover {\n",
              "    background-color: var(--disabled-bg-color);\n",
              "    fill: var(--disabled-fill-color);\n",
              "    box-shadow: none;\n",
              "  }\n",
              "\n",
              "  .colab-df-spinner {\n",
              "    border: 2px solid var(--fill-color);\n",
              "    border-color: transparent;\n",
              "    border-bottom-color: var(--fill-color);\n",
              "    animation:\n",
              "      spin 1s steps(1) infinite;\n",
              "  }\n",
              "\n",
              "  @keyframes spin {\n",
              "    0% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "      border-left-color: var(--fill-color);\n",
              "    }\n",
              "    20% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    30% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    40% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    60% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    80% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "    90% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "  }\n",
              "</style>\n",
              "\n",
              "  <script>\n",
              "    async function quickchart(key) {\n",
              "      const quickchartButtonEl =\n",
              "        document.querySelector('#' + key + ' button');\n",
              "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
              "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
              "      try {\n",
              "        const charts = await google.colab.kernel.invokeFunction(\n",
              "            'suggestCharts', [key], {});\n",
              "      } catch (error) {\n",
              "        console.error('Error during call to suggestCharts:', error);\n",
              "      }\n",
              "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
              "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
              "    }\n",
              "    (() => {\n",
              "      let quickchartButtonEl =\n",
              "        document.querySelector('#df-aee69069-8b48-4983-903f-3063f90b2243 button');\n",
              "      quickchartButtonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "    })();\n",
              "  </script>\n",
              "</div>\n",
              "\n",
              "  <div id=\"id_b0663cdb-fe3d-4416-b488-b01cbfd69bb8\">\n",
              "    <style>\n",
              "      .colab-df-generate {\n",
              "        background-color: #E8F0FE;\n",
              "        border: none;\n",
              "        border-radius: 50%;\n",
              "        cursor: pointer;\n",
              "        display: none;\n",
              "        fill: #1967D2;\n",
              "        height: 32px;\n",
              "        padding: 0 0 0 0;\n",
              "        width: 32px;\n",
              "      }\n",
              "\n",
              "      .colab-df-generate:hover {\n",
              "        background-color: #E2EBFA;\n",
              "        box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "        fill: #174EA6;\n",
              "      }\n",
              "\n",
              "      [theme=dark] .colab-df-generate {\n",
              "        background-color: #3B4455;\n",
              "        fill: #D2E3FC;\n",
              "      }\n",
              "\n",
              "      [theme=dark] .colab-df-generate:hover {\n",
              "        background-color: #434B5C;\n",
              "        box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "        filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "        fill: #FFFFFF;\n",
              "      }\n",
              "    </style>\n",
              "    <button class=\"colab-df-generate\" onclick=\"generateWithVariable('train')\"\n",
              "            title=\"Generate code using this dataframe.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "       width=\"24px\">\n",
              "    <path d=\"M7,19H8.4L18.45,9,17,7.55,7,17.6ZM5,21V16.75L18.45,3.32a2,2,0,0,1,2.83,0l1.4,1.43a1.91,1.91,0,0,1,.58,1.4,1.91,1.91,0,0,1-.58,1.4L9.25,21ZM18.45,9,17,7.55Zm-12,3A5.31,5.31,0,0,0,4.9,8.1,5.31,5.31,0,0,0,1,6.5,5.31,5.31,0,0,0,4.9,4.9,5.31,5.31,0,0,0,6.5,1,5.31,5.31,0,0,0,8.1,4.9,5.31,5.31,0,0,0,12,6.5,5.46,5.46,0,0,0,6.5,12Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "    <script>\n",
              "      (() => {\n",
              "      const buttonEl =\n",
              "        document.querySelector('#id_b0663cdb-fe3d-4416-b488-b01cbfd69bb8 button.colab-df-generate');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      buttonEl.onclick = () => {\n",
              "        google.colab.notebook.generateWithVariable('train');\n",
              "      }\n",
              "      })();\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "    </div>\n",
              "  </div>\n"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "dataframe",
              "variable_name": "train",
              "summary": "{\n  \"name\": \"train\",\n  \"rows\": 500,\n  \"fields\": [\n    {\n      \"column\": \"prompt\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 500,\n        \"samples\": [\n          \"Web search results:\\n\\n[1] \\\"Here are eight tasks all sales teams need to consider to complete their account planning strategy: 1. Identify your existing accounts. First, list all your current customers. Add any details you have about them regarding their purchasing habits and company profile. 2. Caculate potential revenue and success rate.\\\"\\nURL: https://www.salesforce.com/products/anywhere/resources/account-planning-strategy/\\n\\n[2] \\\"Implementing an effective digital strategy across front-, middle- and back-office IT environments is crucial to staying relevant. Yet, the urgent need to create and act on a digital strategy gives rise to a host of complex problems that enterprises must sort out quickly\\u2014or risk falling behind.\\\"\\nURL: https://www.cognizant.com/us/en/services/digital-strategy\\n\\n[3] \\\"At Cognizant, we help our clients do what others cant\\u2014operate with human insight, but at superhuman speed. So they get ahead. And stay ahead. Learn more About Back 2021 in Focus About Cognizant Board of Directors Client Briefing Program Cognizant Events Contact Us Corporate Governance Corporate Integrity Culture Values Diversity Inclusion ESG\\\"\\nURL: https://www.cognizant.com/us/en/services/consulting\\n\\n[4] \\\"Strategic account management is an important job that requires rigor and discipline. Think of it as running a business within your business. Account planning is important to help identify the resources that you need to achieve your growth objectives. And, it is a collaborative process that requires involvement from the client to be valid.\\\"\\nURL: https://www.richardson.com/sales-resources/strategic-sales-key-account-management-strategy/\\n\\n[5] \\\"As an accomplished Technology Services Executive, I excel in blending strategic vision with tactical execution in order to achieve business agendas. I am skilled in transitioning companies through ...\\\"\\nURL: https://www.linkedin.com/in/robvatter\\nCurrent date: 1/29/2023\\n\\nInstructions: Using the provided web search results, write a comprehensive reply to the given query. Make sure to cite results using [[number](URL)] notation after the reference. If the provided search results refer to multiple subjects with the same name, write separate answers for each subject.\\nQuery: Act a VP of Sales for Cognizant to create a new account strategy to sell IT support and consulting services to Seagen Inc located in Bothnell, Seattle\",\n          \"Given a sentence in French, provide an equivalent paraphrased translation in Chinese that retains the same meaning both through the translation and the paraphrase.\\nQ: Cette s\\u00e9rie \\u00e9tait exclusive \\u00e0 Wal-Mart Canada, mais a \\u00e9t\\u00e9 vendue en ligne dans le magasin Spawn.\\nA: \",\n          \"Given a comment, classify it into a 'hate' speech or 'not hate' speech. Answer using 'hate' or 'not hate'.\\n\\nEx Input:\\nI don't care what people say these two fuckers are lying,the Facts are stacked against them I hope they going to pay for this smere campaign\\n\\nEx Output:\\nnot hate\\n\\n\\nEx Input:\\nWait did the tiger get shot?\\n\\nEx Output:\\nnot hate\\n\\n\\nEx Input:\\nHindus and Muslims are same shit. They only know how to kill Peoples...See comments of Christians people no rudeness no abusing or bad words... Christians religion is the best religion in the world\\n\\nEx Output:\\n\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"chosen\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 499,\n        \"samples\": [\n          \"1\\nWhy? This conversation is neutral as it is about discussing taxing increases and opinions. Therefore this is not classified as a strictly positive emotion in the conversation.\",\n          \"\\u8fd9\\u90e8\\u8fde\\u7eed\\u5267\\u4ee5\\u524d\\u4ec5\\u5728\\u52a0\\u62ff\\u5927\\u6c83\\u5c14\\u739b\\u6709\\u552e\\uff0c\\u4f46\\u540e\\u6765\\u5728Spawn\\u5546\\u5e97\\u7684\\u7f51\\u7ad9\\u4e0a\\u51fa\\u552e\\u4e86\\u3002\",\n          \"[Question] What is Madagascar known for?\\n[Answer] Madagascar is famous for its unique wildlife, diverse ecosystems, and rich cultural heritage.\\n\\n[Question] What is Fickle Creek Farm?\\n[Answer] Fickle Creek Farm is an equestrian center located in North Carolina offering various equine services such as boarding, riding lessons, and training.\\n\\n[Question] Tell me about Obama family tree.\\n[Answer] The Obama family tree includes Barack Obama, the 44th President of the United States, and his family members like Michelle Obama, his wife, and their daughters Malia and Sasha.\\n\\n[Question] Why is the SpaceX Falcon 9 rocket important?\\n[Answer] The SpaceX Falcon 9 rocket is important for its successful and cost-effective launches of satellites and other payloads into orbit, credited with advancing the field of private space exploration and boosting commercial spaceflight.\\n\\n[Question] What is the impact of deforestation on wildlife?\\n[Answer] Deforestation negatively impacts wildlife by disrupting habitats and natural food chains, leading to loss of biodiversity and increased risk of extinction, while also exacerbating global climate change.\\n\\n[Question] Who are the Jewish people and what's their origin?\\n[Answer] The Jewish people are a global ethnoreligious group originating from the Israelites or Hebrews of the ancient Near East, with their religious, cultural, and ethnic heritage closely tied to the principles of the Torah and the history of Israel.\\n\\n[Question] What are the effects of smoking on human health?\\n[Answer] Smoking has harmful effects on human health, causing various diseases such as lung cancer, heart disease, and respiratory illnesses, while also increasing the risk of stroke, premature aging, and numerous other health issues.\\n\\n[Question] What are the benefits of eating a balanced diet?\\n[Answer] A balanced diet has numerous benefits, including the maintenance of a healthy weight, improved energy levels, enhanced physical and mental health, and reduced risk of chronic diseases such as heart disease, diabetes, and certain types of cancer.\\n\\n[Question] What is the purpose of the Paris Agreement?\\n[Answer] The Paris Agreement aims to strengthen the global response to climate change by setting a long-term goal to limit global warming, promoting low greenhouse gas emissions, and encouraging financial flows that contribute to climate resilience and adaptation.\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"rejected\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 499,\n        \"samples\": [\n          \"Thank you for your question! I'll be happy to help you.\\n\\nThe crypto industry is full of potentially profitable affiliate programs! Here are a few sites that offer affiliate programs in the crypto industry:\\n\\n1. Binance: Binance is a popular cryptocurrency exchange that offers an affiliate program. You can earn commissions by referring users to trade on the exchange.\\n2. Bitfinex: Bitfinex is a cryptocurrency exchange that also offers an affiliate program, which pays out commissions based on trading volume.\\n3. Coinbase: Coinbase is a cryptocurrency exchange that offers an affiliate program that pays out commissions based on referred sign-ups.\\n4. Ledger: Ledger is a hardware wallet manufacturer that offers an affiliate program for people who want to promote their products.\\n5. Trezor: Trezor is another hardware wallet manufacturer that offers an affiliate program for people who want to promote their hardware wallets.\\n\\nI hope that helps! If you have any more questions, feel free to ask.\",\n          \"\\u8fd9\\u4e2a\\u7cfb\\u5217\\u4ec5\\u5728\\u52a0\\u62ff\\u5927\\u534e\\u5c14\\u8857\\u5f00\\u59cb\\u9500\\u552e\\uff0c\\u4f46\\u540e\\u6765\\u5728\\u7f51\\u4e0a\\u7684\\u6770\\u514b\\u7ef4\\u5c14\\u5546\\u5e97\\u8fdb\\u884c\\u4e86\\u9500\\u552e\\u3002\",\n          \"The Obama family tree consists of Barack Obama, his parents, and their eight years of service in the White House as the 44th President of the United States. Barack Obama's parents were Barack Obama and Stanley Ann Dunham, independently from Barack and Stanley's marriage. They had a daughter, Malia, and a son, Nat\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}"
            }
          },
          "metadata": {},
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Truncate prompt, and responses to avoid OOM\n",
        "for col in ['prompt', 'chosen', 'rejected']:\n",
        "    if col == 'prompt':\n",
        "        max_length = 500\n",
        "    else:\n",
        "        max_length = 1024\n",
        "    train[col] = train[col].str[:max_length]"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-05-15T07:03:39.589603Z",
          "iopub.execute_input": "2024-05-15T07:03:39.590077Z",
          "iopub.status.idle": "2024-05-15T07:03:39.638742Z",
          "shell.execute_reply.started": "2024-05-15T07:03:39.590041Z",
          "shell.execute_reply": "2024-05-15T07:03:39.637745Z"
        },
        "trusted": true,
        "id": "B7l-9Q1qwNBb"
      },
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "prompts = train['prompt'].values.tolist()\n",
        "chosen_responses = train['chosen'].values.tolist()\n",
        "rejected_responses = train['rejected'].values.tolist()\n",
        "\n",
        "torch.manual_seed(seed=cfg.seed)\n",
        "train_dataset = CustomDataset(\n",
        "    prompts=prompts,\n",
        "    chosen_responses=chosen_responses,\n",
        "    rejected_responses=rejected_responses\n",
        ")\n",
        "\n",
        "train_dataloader = DataLoader(\n",
        "    dataset=train_dataset,\n",
        "    batch_size=cfg.batch_size,\n",
        "    shuffle=True,\n",
        "    drop_last=True,\n",
        "    collate_fn=custom_collate\n",
        ")"
      ],
      "metadata": {
        "trusted": true,
        "id": "ztRBP7zcwNBb"
      },
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Init model\n",
        "model = Policy(cfg)\n",
        "\n",
        "# Optimizer and scheduler\n",
        "optimizer = optim.AdamW(\n",
        "    model.parameters(),\n",
        "    lr=cfg.learning_rate,\n",
        "    weight_decay=cfg.weight_decay\n",
        ")\n",
        "\n",
        "scheduler = get_cosine_schedule_with_warmup(\n",
        "    optimizer=optimizer,\n",
        "    num_warmup_steps=cfg.warmup_steps,\n",
        "    num_training_steps=cfg.num_epochs*len(train_dataloader)\n",
        ")"
      ],
      "metadata": {
        "trusted": true,
        "id": "U-F7WaCuwNBb"
      },
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = Policy(cfg)\n",
        "\n",
        "optimizer = optim.AdamW(\n",
        "    model.parameters(),\n",
        "    lr=cfg.learning_rate,\n",
        "    weight_decay=cfg.weight_decay\n",
        ")\n",
        "\n",
        "scheduler = get_cosine_schedule_with_warmup(\n",
        "    optimizer=optimizer,\n",
        "    num_warmup_steps=cfg.warmup_steps,\n",
        "    num_training_steps=cfg.num_epochs*len(train_dataloader)\n",
        ")\n",
        "\n",
        "train_losses = []\n",
        "for epoch in range(cfg.num_epochs):\n",
        "\n",
        "    for batch_idx, batch in enumerate(train_dataloader):\n",
        "\n",
        "        torch.cuda.empty_cache()\n",
        "        gc.collect()\n",
        "\n",
        "        model.train()\n",
        "\n",
        "        chosen_input_ids = batch['chosen_input_ids'].long().to(cfg.device)\n",
        "        chosen_labels = batch['chosen_labels'].long().to(cfg.device)\n",
        "        rejected_input_ids = batch['rejected_input_ids'].long().to(cfg.device)\n",
        "        rejected_labels = batch['rejected_labels'].long().to(cfg.device)\n",
        "\n",
        "        reference_chosen_logits, reference_rejected_logits = get_reference_logits(\n",
        "            reference_model,\n",
        "            chosen_input_ids,\n",
        "            rejected_input_ids\n",
        "        )\n",
        "\n",
        "        policy_chosen_logits = model(chosen_input_ids)\n",
        "        policy_rejected_logits = model(rejected_input_ids)\n",
        "\n",
        "        policy_chosen_logps = get_batch_logps(\n",
        "            policy_chosen_logits,\n",
        "            chosen_labels,\n",
        "            pad_id_value=tokenizer.eos_token_id\n",
        "        )\n",
        "\n",
        "        policy_rejected_logps = get_batch_logps(\n",
        "            policy_rejected_logits,\n",
        "            rejected_labels,\n",
        "            pad_id_value=tokenizer.eos_token_id\n",
        "        )\n",
        "\n",
        "        reference_chosen_logps = get_batch_logps(\n",
        "            reference_chosen_logits,\n",
        "            chosen_labels,\n",
        "            pad_id_value=tokenizer.eos_token_id\n",
        "        )\n",
        "\n",
        "        reference_rejected_logps = get_batch_logps(\n",
        "            reference_rejected_logits,\n",
        "            rejected_labels,\n",
        "            pad_id_value=tokenizer.eos_token_id\n",
        "        )\n",
        "\n",
        "        loss = dpo_loss(\n",
        "            policy_chosen_logps,\n",
        "            policy_rejected_logps,\n",
        "            reference_chosen_logps,\n",
        "            reference_rejected_logps,\n",
        "            beta=cfg.beta\n",
        "        ) / cfg.accumulation_steps\n",
        "\n",
        "        train_losses.append(loss.item())\n",
        "        if not (batch_idx % 100):\n",
        "            print(\n",
        "                f'Epoch: {epoch+1} / {cfg.num_epochs}'\n",
        "                f'| Batch: {batch_idx}/{len(train_dataloader)}'\n",
        "                f'| Train DPO Loss: {np.mean(train_losses):.8f}'\n",
        "            )\n",
        "\n",
        "        loss.backward()\n",
        "\n",
        "        if not (batch_idx % cfg.accumulation_steps):\n",
        "            optimizer.step()\n",
        "            scheduler.step()\n",
        "            optimizer.zero_grad()"
      ],
      "metadata": {
        "trusted": true,
        "id": "U7XuenL9wNBc",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "20fc9987-a6e7-4611-a50d-b5cff42a19ef"
      },
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: 1 / 1| Batch: 0/500| Train DPO Loss: 0.34657359\n",
            "Epoch: 1 / 1| Batch: 100/500| Train DPO Loss: 0.34655462\n",
            "Epoch: 1 / 1| Batch: 200/500| Train DPO Loss: 0.34659074\n",
            "Epoch: 1 / 1| Batch: 300/500| Train DPO Loss: 0.34656812\n",
            "Epoch: 1 / 1| Batch: 400/500| Train DPO Loss: 0.34656084\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "test_prompts = ['What is Hanoi', '2 + 2 =']\n",
        "responses = []\n",
        "\n",
        "for prompt in test_prompts:\n",
        "\n",
        "    input_ids = tokenizer.apply_chat_template(\n",
        "        get_messages(prompt),\n",
        "        return_tensors='pt'\n",
        "    ).to(cfg.device)\n",
        "\n",
        "    input_ids = input_ids[:, :-1] # Slice to remove the </s> that is automatically added when using tokenizer.apply_chat_template()\n",
        "\n",
        "    with torch.no_grad():\n",
        "        model.eval()\n",
        "        output_ids = model.backbone.generate(\n",
        "                input_ids=input_ids,\n",
        "                max_new_tokens=1024,\n",
        "                pad_token_id=tokenizer.pad_token_id\n",
        "        )[0]\n",
        "\n",
        "    question_response = tokenizer.decode(output_ids, skip_special_tokens=True)\n",
        "    response = extract_response(question_response)\n",
        "\n",
        "    print('Prompt:', tokenizer.decode(input_ids[0]))\n",
        "    print('Response:', response)\n",
        "    print('\\n')"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-05-15T07:06:52.837853Z",
          "iopub.execute_input": "2024-05-15T07:06:52.838322Z",
          "iopub.status.idle": "2024-05-15T07:06:52.878583Z",
          "shell.execute_reply.started": "2024-05-15T07:06:52.838285Z",
          "shell.execute_reply": "2024-05-15T07:06:52.876928Z"
        },
        "trusted": true,
        "id": "Ex6LRZnXwNBc",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "17bca8b3-74c6-48ff-d0d7-5fc071909a42"
      },
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Prompt: <|prompt|>What is Hanoi</s> <|answer|>\n",
            "Response:  Hanoi is the capital city of Vietnam. It's a big city with lots of people and buildings. It has many important places like museums, parks, and temples. People in Hanoi speak Vietnamese, which is a language that is different from English. They also have their own special food and way of life.\n",
            "\n",
            "\n",
            "Prompt: <|prompt|>2 + 2 =</s> <|answer|>\n",
            "Response:  To solve the equation, I will follow these steps:\n",
            "\n",
            "Step 1: Add the numbers on the left side of the equation.\n",
            "2 + 2 = 4\n",
            "\n",
            "So, the equation becomes:\n",
            "4 =\n",
            "\n",
            "Step 2: Since there is no unknown variable in this equation, we can conclude that the value of the unknown variable x is not needed to determine the solution.\n",
            "\n",
            "Therefore, the answer is 4.\n",
            "\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We can see that aligning models with more training examples using DPO yields slightly better responses."
      ],
      "metadata": {
        "id": "ZgaEiDwwjdND"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "6_jJyXoghv95"
      },
      "execution_count": 26,
      "outputs": []
    }
  ]
}