{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "beae7bb1-f4e8-4e9a-80ce-c22f16124638",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: torch==1.13.1 in /opt/conda/lib/python3.7/site-packages (1.13.1)\n",
      "Requirement already satisfied: torchdata in /opt/conda/lib/python3.7/site-packages (0.5.1)\n",
      "Requirement already satisfied: nvidia-cudnn-cu11==8.5.0.96 in /opt/conda/lib/python3.7/site-packages (from torch==1.13.1) (8.5.0.96)\n",
      "Requirement already satisfied: nvidia-cublas-cu11==11.10.3.66 in /opt/conda/lib/python3.7/site-packages (from torch==1.13.1) (11.10.3.66)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu11==11.7.99 in /opt/conda/lib/python3.7/site-packages (from torch==1.13.1) (11.7.99)\n",
      "Requirement already satisfied: typing-extensions in /opt/conda/lib/python3.7/site-packages (from torch==1.13.1) (4.5.0)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu11==11.7.99 in /opt/conda/lib/python3.7/site-packages (from torch==1.13.1) (11.7.99)\n",
      "Requirement already satisfied: wheel in /opt/conda/lib/python3.7/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch==1.13.1) (0.40.0)\n",
      "Requirement already satisfied: setuptools in /opt/conda/lib/python3.7/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch==1.13.1) (59.3.0)\n",
      "Requirement already satisfied: urllib3>=1.25 in /opt/conda/lib/python3.7/site-packages (from torchdata) (1.26.15)\n",
      "Requirement already satisfied: requests in /opt/conda/lib/python3.7/site-packages (from torchdata) (2.28.2)\n",
      "Requirement already satisfied: portalocker>=2.0.0 in /opt/conda/lib/python3.7/site-packages (from torchdata) (2.7.0)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /opt/conda/lib/python3.7/site-packages (from requests->torchdata) (2.0.4)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.7/site-packages (from requests->torchdata) (2022.12.7)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.7/site-packages (from requests->torchdata) (2.8)\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.1.2\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install torch==1.13.1 torchdata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8edb5aba-15bc-446d-874f-8795a436be6a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install --disable-pip-version-check -q \\\n",
    "    transformers==4.27.2 \\\n",
    "    datasets==2.9.0 \\\n",
    "    accelerate==0.17.0 \\\n",
    "    evaluate==0.4.0 \\\n",
    "    trl==0.4.1 \\\n",
    "    rouge_score==0.1.2 \\\n",
    "    loralib==0.1.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "039fb142-c9a5-4ca6-aaf7-f8c4c65743b3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting git+https://github.com/huggingface/peft.git\n",
      "  Cloning https://github.com/huggingface/peft.git to /tmp/pip-req-build-pd5zw3ha\n",
      "  Running command git clone --filter=blob:none --quiet https://github.com/huggingface/peft.git /tmp/pip-req-build-pd5zw3ha\n",
      "  Resolved https://github.com/huggingface/peft.git to commit 632997d1fb776c3cf05d8c2537ac9a98a7ce9435\n",
      "  Installing build dependencies ... \u001b[?25ldone\n",
      "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25ldone\n",
      "\u001b[?25h  Preparing metadata (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25hRequirement already satisfied: torch>=1.13.0 in /opt/conda/lib/python3.7/site-packages (from peft==0.3.0.dev0) (1.13.1)\n",
      "Requirement already satisfied: pyyaml in /opt/conda/lib/python3.7/site-packages (from peft==0.3.0.dev0) (6.0)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.7/site-packages (from peft==0.3.0.dev0) (23.1)\n",
      "Requirement already satisfied: psutil in /opt/conda/lib/python3.7/site-packages (from peft==0.3.0.dev0) (5.6.7)\n",
      "Requirement already satisfied: transformers in /opt/conda/lib/python3.7/site-packages (from peft==0.3.0.dev0) (4.27.2)\n",
      "Requirement already satisfied: accelerate in /opt/conda/lib/python3.7/site-packages (from peft==0.3.0.dev0) (0.17.0)\n",
      "Requirement already satisfied: numpy>=1.17 in /opt/conda/lib/python3.7/site-packages (from peft==0.3.0.dev0) (1.21.6)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu11==11.7.99 in /opt/conda/lib/python3.7/site-packages (from torch>=1.13.0->peft==0.3.0.dev0) (11.7.99)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu11==11.7.99 in /opt/conda/lib/python3.7/site-packages (from torch>=1.13.0->peft==0.3.0.dev0) (11.7.99)\n",
      "Requirement already satisfied: typing-extensions in /opt/conda/lib/python3.7/site-packages (from torch>=1.13.0->peft==0.3.0.dev0) (4.5.0)\n",
      "Requirement already satisfied: nvidia-cublas-cu11==11.10.3.66 in /opt/conda/lib/python3.7/site-packages (from torch>=1.13.0->peft==0.3.0.dev0) (11.10.3.66)\n",
      "Requirement already satisfied: nvidia-cudnn-cu11==8.5.0.96 in /opt/conda/lib/python3.7/site-packages (from torch>=1.13.0->peft==0.3.0.dev0) (8.5.0.96)\n",
      "Requirement already satisfied: setuptools in /opt/conda/lib/python3.7/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch>=1.13.0->peft==0.3.0.dev0) (59.3.0)\n",
      "Requirement already satisfied: wheel in /opt/conda/lib/python3.7/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch>=1.13.0->peft==0.3.0.dev0) (0.40.0)\n",
      "Requirement already satisfied: importlib-metadata in /opt/conda/lib/python3.7/site-packages (from transformers->peft==0.3.0.dev0) (6.3.0)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /opt/conda/lib/python3.7/site-packages (from transformers->peft==0.3.0.dev0) (2022.10.31)\n",
      "Requirement already satisfied: requests in /opt/conda/lib/python3.7/site-packages (from transformers->peft==0.3.0.dev0) (2.28.2)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.11.0 in /opt/conda/lib/python3.7/site-packages (from transformers->peft==0.3.0.dev0) (0.14.1)\n",
      "Requirement already satisfied: tqdm>=4.27 in /opt/conda/lib/python3.7/site-packages (from transformers->peft==0.3.0.dev0) (4.65.0)\n",
      "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /opt/conda/lib/python3.7/site-packages (from transformers->peft==0.3.0.dev0) (0.13.3)\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.7/site-packages (from transformers->peft==0.3.0.dev0) (3.0.12)\n",
      "Requirement already satisfied: fsspec in /opt/conda/lib/python3.7/site-packages (from huggingface-hub<1.0,>=0.11.0->transformers->peft==0.3.0.dev0) (2023.1.0)\n",
      "Requirement already satisfied: zipp>=0.5 in /opt/conda/lib/python3.7/site-packages (from importlib-metadata->transformers->peft==0.3.0.dev0) (3.15.0)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.7/site-packages (from requests->transformers->peft==0.3.0.dev0) (2022.12.7)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /opt/conda/lib/python3.7/site-packages (from requests->transformers->peft==0.3.0.dev0) (1.26.15)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.7/site-packages (from requests->transformers->peft==0.3.0.dev0) (2.8)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /opt/conda/lib/python3.7/site-packages (from requests->transformers->peft==0.3.0.dev0) (2.0.4)\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.1.2\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install git+https://github.com/huggingface/peft.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "442332d6-8048-4e38-aa97-d0ac35cb0f3e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting git+https://github.com/lvwerra/trl.git\n",
      "  Cloning https://github.com/lvwerra/trl.git to /tmp/pip-req-build-vjvgg1xk\n",
      "  Running command git clone --filter=blob:none --quiet https://github.com/lvwerra/trl.git /tmp/pip-req-build-vjvgg1xk\n",
      "  Resolved https://github.com/lvwerra/trl.git to commit 52fecee8839ad826ad1e6c83a95c99a4116e98d2\n",
      "  Preparing metadata (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25hRequirement already satisfied: torch>=1.4.0 in /opt/conda/lib/python3.7/site-packages (from trl==0.4.2.dev0) (1.13.1)\n",
      "Requirement already satisfied: transformers>=4.18.0 in /opt/conda/lib/python3.7/site-packages (from trl==0.4.2.dev0) (4.27.2)\n",
      "Requirement already satisfied: numpy>=1.18.2 in /opt/conda/lib/python3.7/site-packages (from trl==0.4.2.dev0) (1.21.6)\n",
      "Requirement already satisfied: accelerate in /opt/conda/lib/python3.7/site-packages (from trl==0.4.2.dev0) (0.17.0)\n",
      "Requirement already satisfied: datasets in /opt/conda/lib/python3.7/site-packages (from trl==0.4.2.dev0) (2.9.0)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu11==11.7.99 in /opt/conda/lib/python3.7/site-packages (from torch>=1.4.0->trl==0.4.2.dev0) (11.7.99)\n",
      "Requirement already satisfied: nvidia-cudnn-cu11==8.5.0.96 in /opt/conda/lib/python3.7/site-packages (from torch>=1.4.0->trl==0.4.2.dev0) (8.5.0.96)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu11==11.7.99 in /opt/conda/lib/python3.7/site-packages (from torch>=1.4.0->trl==0.4.2.dev0) (11.7.99)\n",
      "Requirement already satisfied: typing-extensions in /opt/conda/lib/python3.7/site-packages (from torch>=1.4.0->trl==0.4.2.dev0) (4.5.0)\n",
      "Requirement already satisfied: nvidia-cublas-cu11==11.10.3.66 in /opt/conda/lib/python3.7/site-packages (from torch>=1.4.0->trl==0.4.2.dev0) (11.10.3.66)\n",
      "Requirement already satisfied: setuptools in /opt/conda/lib/python3.7/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch>=1.4.0->trl==0.4.2.dev0) (59.3.0)\n",
      "Requirement already satisfied: wheel in /opt/conda/lib/python3.7/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch>=1.4.0->trl==0.4.2.dev0) (0.40.0)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.7/site-packages (from transformers>=4.18.0->trl==0.4.2.dev0) (23.1)\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.7/site-packages (from transformers>=4.18.0->trl==0.4.2.dev0) (3.0.12)\n",
      "Requirement already satisfied: requests in /opt/conda/lib/python3.7/site-packages (from transformers>=4.18.0->trl==0.4.2.dev0) (2.28.2)\n",
      "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /opt/conda/lib/python3.7/site-packages (from transformers>=4.18.0->trl==0.4.2.dev0) (0.13.3)\n",
      "Requirement already satisfied: importlib-metadata in /opt/conda/lib/python3.7/site-packages (from transformers>=4.18.0->trl==0.4.2.dev0) (6.3.0)\n",
      "Requirement already satisfied: tqdm>=4.27 in /opt/conda/lib/python3.7/site-packages (from transformers>=4.18.0->trl==0.4.2.dev0) (4.65.0)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /opt/conda/lib/python3.7/site-packages (from transformers>=4.18.0->trl==0.4.2.dev0) (2022.10.31)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /opt/conda/lib/python3.7/site-packages (from transformers>=4.18.0->trl==0.4.2.dev0) (6.0)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.11.0 in /opt/conda/lib/python3.7/site-packages (from transformers>=4.18.0->trl==0.4.2.dev0) (0.14.1)\n",
      "Requirement already satisfied: psutil in /opt/conda/lib/python3.7/site-packages (from accelerate->trl==0.4.2.dev0) (5.6.7)\n",
      "Requirement already satisfied: responses<0.19 in /opt/conda/lib/python3.7/site-packages (from datasets->trl==0.4.2.dev0) (0.18.0)\n",
      "Requirement already satisfied: pyarrow>=6.0.0 in /opt/conda/lib/python3.7/site-packages (from datasets->trl==0.4.2.dev0) (11.0.0)\n",
      "Requirement already satisfied: xxhash in /opt/conda/lib/python3.7/site-packages (from datasets->trl==0.4.2.dev0) (3.2.0)\n",
      "Requirement already satisfied: aiohttp in /opt/conda/lib/python3.7/site-packages (from datasets->trl==0.4.2.dev0) (3.8.4)\n",
      "Requirement already satisfied: dill<0.3.7 in /opt/conda/lib/python3.7/site-packages (from datasets->trl==0.4.2.dev0) (0.3.6)\n",
      "Requirement already satisfied: fsspec[http]>=2021.11.1 in /opt/conda/lib/python3.7/site-packages (from datasets->trl==0.4.2.dev0) (2023.1.0)\n",
      "Requirement already satisfied: pandas in /opt/conda/lib/python3.7/site-packages (from datasets->trl==0.4.2.dev0) (1.3.5)\n",
      "Requirement already satisfied: multiprocess in /opt/conda/lib/python3.7/site-packages (from datasets->trl==0.4.2.dev0) (0.70.14)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets->trl==0.4.2.dev0) (4.0.2)\n",
      "Requirement already satisfied: charset-normalizer<4.0,>=2.0 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets->trl==0.4.2.dev0) (2.0.4)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets->trl==0.4.2.dev0) (6.0.4)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets->trl==0.4.2.dev0) (1.3.3)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets->trl==0.4.2.dev0) (1.8.2)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets->trl==0.4.2.dev0) (1.3.1)\n",
      "Requirement already satisfied: asynctest==0.13.0 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets->trl==0.4.2.dev0) (0.13.0)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /opt/conda/lib/python3.7/site-packages (from aiohttp->datasets->trl==0.4.2.dev0) (22.2.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.7/site-packages (from requests->transformers>=4.18.0->trl==0.4.2.dev0) (2.8)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /opt/conda/lib/python3.7/site-packages (from requests->transformers>=4.18.0->trl==0.4.2.dev0) (1.26.15)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.7/site-packages (from requests->transformers>=4.18.0->trl==0.4.2.dev0) (2022.12.7)\n",
      "Requirement already satisfied: zipp>=0.5 in /opt/conda/lib/python3.7/site-packages (from importlib-metadata->transformers>=4.18.0->trl==0.4.2.dev0) (3.15.0)\n",
      "Requirement already satisfied: pytz>=2017.3 in /opt/conda/lib/python3.7/site-packages (from pandas->datasets->trl==0.4.2.dev0) (2019.3)\n",
      "Requirement already satisfied: python-dateutil>=2.7.3 in /opt/conda/lib/python3.7/site-packages (from pandas->datasets->trl==0.4.2.dev0) (2.8.2)\n",
      "Requirement already satisfied: six>=1.5 in /opt/conda/lib/python3.7/site-packages (from python-dateutil>=2.7.3->pandas->datasets->trl==0.4.2.dev0) (1.14.0)\n",
      "Building wheels for collected packages: trl\n",
      "  Building wheel for trl (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for trl: filename=trl-0.4.2.dev0-py3-none-any.whl size=59756 sha256=8a9251fbbee710e747c34ea73ad5af481a025348b4e389db9bac95d91503bd64\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-a6o05ves/wheels/ca/6e/f4/b183ecbed483efdcd2041a8021ce7bcb9f7b09c74bff5bb00a\n",
      "Successfully built trl\n",
      "Installing collected packages: trl\n",
      "  Attempting uninstall: trl\n",
      "    Found existing installation: trl 0.4.1\n",
      "    Uninstalling trl-0.4.1:\n",
      "      Successfully uninstalled trl-0.4.1\n",
      "Successfully installed trl-0.4.2.dev0\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.1.2\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install git+https://github.com/lvwerra/trl.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "37a78f85-cb66-4059-b4fc-437db14e9684",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%store -r ranking_reward_model_custom_checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "948d777a-2ed7-4d9b-933e-49cd60e576f8",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./ranking_reward_model_custom/\n"
     ]
    }
   ],
   "source": [
    "print(ranking_reward_model_custom_checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "107e192d-ec64-41a8-bf2f-674c4ebb816a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# %store -r supervised_fine_tuned_model_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4347957b-85d3-4bdf-bba7-39512f08be55",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# print(supervised_fine_tuned_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f421a5b-2720-4430-aae5-b25cbbef1a40",
   "metadata": {},
   "source": [
    "# Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b696f76c-2aac-4c04-ad0f-82e8a16d3bad",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration data-summarization-6643a2708bd3b017\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading and preparing dataset csv/data-summarization to /root/.cache/huggingface/datasets/csv/data-summarization-6643a2708bd3b017/0.0.0/6b34fb8fcf56f7c8ba51dc895bfa2bfbe43546f190a60fcf74bb5e8afdcc2317...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "67d2e6d3f5f74faf83cc6fa4a6df6210",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5b9125f1a8da4ed192a03627bd7f1a95",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Extracting data files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset csv downloaded and prepared to /root/.cache/huggingface/datasets/csv/data-summarization-6643a2708bd3b017/0.0.0/6b34fb8fcf56f7c8ba51dc895bfa2bfbe43546f190a60fcf74bb5e8afdcc2317. Subsequent calls will reuse this data.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3e8238ec599e4077897f7140ec6c0e3b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['id', 'dialogue', 'summary', 'topic'],\n",
       "        num_rows: 14460\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import  AutoTokenizer\n",
    "from datasets import load_dataset\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained('google/flan-t5-base')\n",
    "dataset = load_dataset('./data-summarization/')\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12b10b83-ad2d-4277-b9df-ccda4fde9810",
   "metadata": {},
   "source": [
    "## Explore an Example Prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e2dbd058-8ba6-4792-8244-ae0e54ac4257",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prompt:\n",
      "--------------------------\n",
      "Summarize the following conversation.\n",
      "\n",
      "#Person1#: If we employ you, what starting salary would you expect?\n",
      "#Person2#: I'd like to start at 3000 yuan a month.\n",
      "#Person1#: I think your background and experience are worth the compensation.\n",
      "#Person2#: Does it include bonuses?\n",
      "#Person1#: No, there are annual bonuses, one week paid vacation a year, and health insurance.\n",
      "#Person2#: Very good.\n",
      "\n",
      "Summary:\n",
      "--------------------------\n",
      "Baseline human summary : #Person1# agrees #Person2#'s starting monthly salary would be 3000 yuan and tells #Person2# about other benefits.\n"
     ]
    }
   ],
   "source": [
    "idx = 0\n",
    "diag = dataset['train'][idx]['dialogue']\n",
    "baseline_human_summary = dataset['train'][idx]['summary']\n",
    "\n",
    "prompt = f'Summarize the following conversation.\\n\\n{diag}\\n\\nSummary:'\n",
    "input_ids = tokenizer(prompt, return_tensors=\"pt\").input_ids\n",
    "\n",
    "print(f'Prompt:\\n--------------------------\\n{prompt}\\n--------------------------')\n",
    "print(f'Baseline human summary : {baseline_human_summary}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b25c87b-1cd3-47d6-94bd-aa7106864578",
   "metadata": {},
   "source": [
    "## Tokenize the Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "daac7a15-430d-45da-9c6e-e4c56badbd8c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "74d3c4d07e2841e280e6e180dd861f1e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/14460 [00:00<?, ?ex/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "\"Summarize the following conversation.\\n\\n#Person1#: If we employ you, what starting salary would you expect?\\n#Person2#: I'd like to start at 3000 yuan a month.\\n#Person1#: I think your background and experience are worth the compensation.\\n#Person2#: Does it include bonuses?\\n#Person1#: No, there are annual bonuses, one week paid vacation a year, and health insurance.\\n#Person2#: Very good.\\n\\nSummary: \""
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def tokenize_function(example):\n",
    "    prompt = 'Summarize the following conversation.\\n\\n'\n",
    "    end_prompt = '\\n\\nSummary: '\n",
    "\n",
    "    dialog = example[\"dialogue\"]\n",
    "    query = prompt + dialog + end_prompt\n",
    "    \n",
    "    example['query'] = query\n",
    "    example['input_ids'] = tokenizer(query, padding=\"max_length\", truncation=True, return_tensors=\"pt\").input_ids\n",
    "    summary = example[\"summary\"]\n",
    "    example['labels'] = tokenizer(summary, padding=\"max_length\", truncation=True, return_tensors=\"pt\").input_ids\n",
    "    return example\n",
    "\n",
    "tokenized_dataset = dataset.map(tokenize_function) #, batched=True)\n",
    "tokenized_dataset = tokenized_dataset.remove_columns(['id', 'topic', 'dialogue', 'summary',])\n",
    "\n",
    "# tokenized_dataset = tokenized_dataset.filter(lambda x: len(x[\"input_ids\"]) < 512, batched=False)\n",
    "# tokenized_dataset.set_format(type=\"torch\")\n",
    "\n",
    "tokenized_dataset['train'][0]['query']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "bc85748e-d6f4-4e36-970d-05a99f958c60",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cabfc56d03354252818ac2ce8672b473",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/15 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from dataclasses import dataclass, field\n",
    "from typing import Optional\n",
    "\n",
    "import torch\n",
    "from accelerate import Accelerator\n",
    "from datasets import load_dataset\n",
    "from peft import LoraConfig\n",
    "from tqdm import tqdm\n",
    "from transformers import Adafactor, AutoTokenizer, HfArgumentParser, pipeline\n",
    "\n",
    "from trl import AutoModelForSeq2SeqLMWithValueHead, PPOConfig, PPOTrainer, set_seed\n",
    "from trl.core import LengthSampler\n",
    "\n",
    "tqdm.pandas()\n",
    "\n",
    "fine_tuned_with_ranking_rewards = './fine_tuned_with_ranking_rewards'\n",
    "\n",
    "@dataclass\n",
    "class ScriptArguments:\n",
    "    model_name: Optional[str] = field(default=\"google/flan-t5-base\", metadata={\"help\": \"the model name\"})\n",
    "    tokenizer_name: Optional[str] = field(default=\"google/flan-t5-base\", metadata={\"help\": \"the tokenizer name\"})\n",
    "    # model_name: Optional[str] = field(default=supervised_fine_tuned_model_path, metadata={\"help\": \"the model name\"})\n",
    "    # tokenizer_name: Optional[str] = field(default=supervised_fine_tuned_model_path, metadata={\"help\": \"the tokenizer name\"})\n",
    "    reward_model_name: Optional[str] = field(default=ranking_reward_model_custom_checkpoint, metadata={\"help\": \"the reward model name\"})\n",
    "    log_with: Optional[str] = field(default=None, metadata={\"help\": \"use 'wandb' to log with wandb\"})\n",
    "    learning_rate: Optional[float] = field(default=1.4e-5, metadata={\"help\": \"the learning rate\"})\n",
    "    output_max_length: Optional[int] = field(default=128, metadata={\"help\": \"maximum length for generation\"})\n",
    "    mini_batch_size: Optional[int] = field(default=1, metadata={\"help\": \"the PPO minibatch size\"})\n",
    "    batch_size: Optional[int] = field(default=8, metadata={\"help\": \"the batch size\"})\n",
    "    ppo_epochs: Optional[int] = field(default=4, metadata={\"help\": \"the number of ppo epochs\"})\n",
    "    gradient_accumulation_steps: Optional[int] = field(\n",
    "        default=8, metadata={\"help\": \"the number of gradient accumulation steps\"}\n",
    "    )\n",
    "    adafactor: Optional[bool] = field(default=False, metadata={\"help\": \"whether to use the adafactor optimizer\"})\n",
    "    early_stopping: Optional[bool] = field(default=True, metadata={\"help\": \"whether to early stop\"})\n",
    "    target_kl: Optional[float] = field(default=0.1, metadata={\"help\": \"kl target for early stopping\"})\n",
    "    reward_baseline: Optional[float] = field(\n",
    "        default=0.0,\n",
    "        metadata={\"help\": \"a baseline value that is subtracted from the reward\"},\n",
    "    )\n",
    "    batched_gen: Optional[bool] = field(default=True, metadata={\"help\": \"whether to use the batched text gen\"})\n",
    "    save_freq: Optional[int] = field(default=100, metadata={\"help\": \"n steps to save the model\"})\n",
    "    output_dir: Optional[str] = field(default=fine_tuned_with_ranking_rewards, metadata={\"help\": \"n steps to save the model\"})\n",
    "    seed: Optional[int] = field(default=42, metadata={\"help\": \"the seed\"})\n",
    "\n",
    "\n",
    "parser = HfArgumentParser(ScriptArguments)\n",
    "script_args: ScriptArguments = parser.parse_args_into_dataclasses(return_remaining_strings=True)[0]\n",
    "reward_model_name = script_args.reward_model_name\n",
    "\n",
    "\n",
    "#dataset_name = \"lvwerra/stack-exchange-paired\"\n",
    "\n",
    "config = PPOConfig(\n",
    "    model_name=script_args.model_name,\n",
    "    learning_rate=script_args.learning_rate,\n",
    "    log_with=script_args.log_with,\n",
    "    batch_size=script_args.batch_size,\n",
    "    mini_batch_size=script_args.mini_batch_size,\n",
    "    gradient_accumulation_steps=script_args.gradient_accumulation_steps,\n",
    "    optimize_cuda_cache=True,\n",
    "    early_stopping=script_args.early_stopping,\n",
    "    target_kl=script_args.target_kl,\n",
    "    ppo_epochs=script_args.ppo_epochs,\n",
    "    seed=script_args.seed,\n",
    ")\n",
    "\n",
    "\n",
    "# Dataset is here:\n",
    "#   https://huggingface.co/datasets/lvwerra/stack-exchange-paired/tree/main/data/rl\n",
    "\n",
    "# train_dataset = load_dataset(\"lvwerra/stack-exchange-paired\", data_dir=\"data/rl\", split=\"train\")\n",
    "# train_dataset = train_dataset.select(range(100000))\n",
    "\n",
    "# train_dataset = load_dataset(\"knkarthick/dialogsum\", split=\"test\")\n",
    "# train_dataset = train_dataset.select(range(10))\n",
    "# # We then define the arguments to pass to the sentiment analysis pipeline.\n",
    "# # We set `return_all_scores` to True to get the sentiment score for each token.\n",
    "sent_kwargs = {\"return_all_scores\": True, \"function_to_apply\": \"none\", \"batch_size\": 16, \"truncation\": True}\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(script_args.tokenizer_name)\n",
    "\n",
    "\n",
    "# def build_human_feedback_dataset(\n",
    "#     tokenizer, # dataset_name=\"lvwerra/stack-exchange-paired\", input_min_text_length=2, input_max_text_length=8\n",
    "# ):\n",
    "#     # load with datasets\n",
    "#     ds = load_dataset(\"knkarthick/dialogsum\", split=\"test\")\n",
    "#     original_columns = ds.column_names\n",
    "#     num_proc = 24\n",
    "\n",
    "#     def preprocess_function(examples):\n",
    "#         new_examples = {\n",
    "#             \"query\": [],\n",
    "#             \"input_ids\": [],\n",
    "#         }\n",
    "#         # for question in examples[\"question\"]:\n",
    "#         #     query = \"Question: \" + question + \"\\n\\nAnswer: \"\n",
    "#         #     tokenized_question = tokenizer(query, truncation=True)\n",
    "#         #     new_examples[\"query\"].append(query)\n",
    "#         #     new_examples[\"input_ids\"].append(tokenized_question[\"input_ids\"])\n",
    "#         for example in examples:\n",
    "#             prompt = 'Summarize the following conversation.\\n\\n'\n",
    "#             end_prompt = '\\n\\nSummary: '\n",
    "#             query = [prompt + i + end_prompt for i in example[\"dialogue\"]]\n",
    "#             example['input_ids'] = tokenizer(query, padding=\"max_length\", truncation=True, return_tensors=\"pt\").input_ids\n",
    "#             example['labels'] = tokenizer(example[\"summary\"], padding=\"max_length\", truncation=True, return_tensors=\"pt\").input_ids\n",
    "#             new_examples[\"query\"].append(query)\n",
    "#             new_examples[\"input_ids\"].append(example[\"input_ids\"])\n",
    "            \n",
    "#         return new_examples\n",
    "\n",
    "#     ds = train_dataset.map(\n",
    "#         preprocess_function,\n",
    "#         batched=True,\n",
    "#         num_proc=num_proc,\n",
    "#         remove_columns=original_columns,\n",
    "#     )\n",
    "#     ds = ds.filter(lambda x: len(x[\"input_ids\"]) < 512, batched=False)\n",
    "\n",
    "#     ds.set_format(type=\"torch\")\n",
    "#     return ds\n",
    "\n",
    "\n",
    "# # We retrieve the dataloader by calling the `build_dataset` function.\n",
    "#dataset = build_human_feedback_dataset(tokenizer)\n",
    "\n",
    "def collator(data):\n",
    "    return dict((key, [d[key] for d in data]) for key in data[0])\n",
    "\n",
    "sample_tokenized_dataset = tokenized_dataset.filter(lambda example, indice: indice % 100 == 0, with_indices=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "18f7af59-7c98-4f50-aaf3-bdde4d1ce7b6",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /root/.cache/huggingface/datasets/csv/data-summarization-6643a2708bd3b017/0.0.0/6b34fb8fcf56f7c8ba51dc895bfa2bfbe43546f190a60fcf74bb5e8afdcc2317/cache-5cdf9aa777eefd40.arrow\n"
     ]
    }
   ],
   "source": [
    "# set seed before initializing value head for deterministic eval\n",
    "set_seed(config.seed)\n",
    "\n",
    "# Now let's build the model, the reference model, and the tokenizer.\n",
    "current_device = Accelerator().local_process_index\n",
    "\n",
    "lora_config = LoraConfig(\n",
    "    r=16,\n",
    "    lora_alpha=32,\n",
    "    lora_dropout=0.05,\n",
    "    bias=\"none\",\n",
    "    task_type=\"SEQ_2_SEQ_LM\",\n",
    ")\n",
    "model = AutoModelForSeq2SeqLMWithValueHead.from_pretrained(config.model_name, # device_map={\"\": current_device},\n",
    "    peft_config=lora_config,\n",
    "#    layer_norm_names=[],\n",
    ")\n",
    "\n",
    "sample_tokenized_dataset = tokenized_dataset.filter(lambda example, indice: indice % 100 == 0, with_indices=True)\n",
    "\n",
    "# We then build the PPOTrainer, passing the model, the reference model, the tokenizer\n",
    "ppo_trainer = PPOTrainer(\n",
    "    config,\n",
    "    model,\n",
    "    ref_model=None,\n",
    "    tokenizer=tokenizer,\n",
    "    dataset=sample_tokenized_dataset['train'],\n",
    "    data_collator=collator,\n",
    ")\n",
    "\n",
    "device = ppo_trainer.accelerator.device\n",
    "if ppo_trainer.accelerator.num_processes == 1:\n",
    "    device = 0 if torch.cuda.is_available() else \"cpu\"  # to avoid a ` pipeline` bug\n",
    "sentiment_pipe = pipeline(\n",
    "    \"sentiment-analysis\", # reward classifier is a type of sentiment-analysis (perhaps make this more clear by using text-classification instead?)\n",
    "    model=reward_model_name,\n",
    "#    device_map={\"\": current_device},\n",
    "    tokenizer=tokenizer,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "23633a19-7d91-4d81-b087-fff8f3be0dac",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "0it [00:00, ?it/s]/opt/conda/lib/python3.7/site-packages/transformers/pipelines/text_classification.py:107: UserWarning: `return_all_scores` is now deprecated,  if want a similar funcionality use `top_k=None` instead of `return_all_scores=True` or `top_k=1` instead of `return_all_scores=False`.\n",
      "  UserWarning,\n",
      "You're using a T5TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n",
      "1it [00:33, 33.75s/it]/opt/conda/lib/python3.7/site-packages/trl/trainer/ppo_trainer.py:1089: UserWarning: KL divergence is starting to become negative: -0.06 - this might be a precursor for failed training. sometimes this happens because the generation kwargs are not correctly set. Please make sure that the generation kwargs are set correctly, or review your training hyperparameters.\n",
      "  f\"KL divergence is starting to become negative: {mean_kl.item():.2f} - this might be a precursor for failed training.\"\n",
      "3it [01:37, 32.08s/it]/opt/conda/lib/python3.7/site-packages/trl/trainer/ppo_trainer.py:1089: UserWarning: KL divergence is starting to become negative: -0.13 - this might be a precursor for failed training. sometimes this happens because the generation kwargs are not correctly set. Please make sure that the generation kwargs are set correctly, or review your training hyperparameters.\n",
      "  f\"KL divergence is starting to become negative: {mean_kl.item():.2f} - this might be a precursor for failed training.\"\n",
      "4it [02:10, 32.46s/it]/opt/conda/lib/python3.7/site-packages/trl/trainer/ppo_trainer.py:1089: UserWarning: KL divergence is starting to become negative: -0.14 - this might be a precursor for failed training. sometimes this happens because the generation kwargs are not correctly set. Please make sure that the generation kwargs are set correctly, or review your training hyperparameters.\n",
      "  f\"KL divergence is starting to become negative: {mean_kl.item():.2f} - this might be a precursor for failed training.\"\n",
      "5it [02:43, 32.85s/it]/opt/conda/lib/python3.7/site-packages/trl/trainer/ppo_trainer.py:1089: UserWarning: KL divergence is starting to become negative: -0.10 - this might be a precursor for failed training. sometimes this happens because the generation kwargs are not correctly set. Please make sure that the generation kwargs are set correctly, or review your training hyperparameters.\n",
      "  f\"KL divergence is starting to become negative: {mean_kl.item():.2f} - this might be a precursor for failed training.\"\n",
      "6it [03:13, 31.97s/it]/opt/conda/lib/python3.7/site-packages/trl/trainer/ppo_trainer.py:1089: UserWarning: KL divergence is starting to become negative: -0.26 - this might be a precursor for failed training. sometimes this happens because the generation kwargs are not correctly set. Please make sure that the generation kwargs are set correctly, or review your training hyperparameters.\n",
      "  f\"KL divergence is starting to become negative: {mean_kl.item():.2f} - this might be a precursor for failed training.\"\n",
      "7it [03:45, 31.80s/it]/opt/conda/lib/python3.7/site-packages/trl/trainer/ppo_trainer.py:1089: UserWarning: KL divergence is starting to become negative: -0.22 - this might be a precursor for failed training. sometimes this happens because the generation kwargs are not correctly set. Please make sure that the generation kwargs are set correctly, or review your training hyperparameters.\n",
      "  f\"KL divergence is starting to become negative: {mean_kl.item():.2f} - this might be a precursor for failed training.\"\n",
      "8it [04:18, 32.18s/it]/opt/conda/lib/python3.7/site-packages/trl/trainer/ppo_trainer.py:1089: UserWarning: KL divergence is starting to become negative: -0.47 - this might be a precursor for failed training. sometimes this happens because the generation kwargs are not correctly set. Please make sure that the generation kwargs are set correctly, or review your training hyperparameters.\n",
      "  f\"KL divergence is starting to become negative: {mean_kl.item():.2f} - this might be a precursor for failed training.\"\n",
      "9it [04:49, 32.00s/it]/opt/conda/lib/python3.7/site-packages/trl/trainer/ppo_trainer.py:1089: UserWarning: KL divergence is starting to become negative: -0.41 - this might be a precursor for failed training. sometimes this happens because the generation kwargs are not correctly set. Please make sure that the generation kwargs are set correctly, or review your training hyperparameters.\n",
      "  f\"KL divergence is starting to become negative: {mean_kl.item():.2f} - this might be a precursor for failed training.\"\n",
      "10it [05:22, 32.23s/it]/opt/conda/lib/python3.7/site-packages/trl/trainer/ppo_trainer.py:1089: UserWarning: KL divergence is starting to become negative: -0.19 - this might be a precursor for failed training. sometimes this happens because the generation kwargs are not correctly set. Please make sure that the generation kwargs are set correctly, or review your training hyperparameters.\n",
      "  f\"KL divergence is starting to become negative: {mean_kl.item():.2f} - this might be a precursor for failed training.\"\n",
      "11it [05:53, 31.75s/it]/opt/conda/lib/python3.7/site-packages/trl/trainer/ppo_trainer.py:1089: UserWarning: KL divergence is starting to become negative: -0.69 - this might be a precursor for failed training. sometimes this happens because the generation kwargs are not correctly set. Please make sure that the generation kwargs are set correctly, or review your training hyperparameters.\n",
      "  f\"KL divergence is starting to become negative: {mean_kl.item():.2f} - this might be a precursor for failed training.\"\n",
      "12it [06:28, 32.85s/it]/opt/conda/lib/python3.7/site-packages/trl/trainer/ppo_trainer.py:1089: UserWarning: KL divergence is starting to become negative: -0.25 - this might be a precursor for failed training. sometimes this happens because the generation kwargs are not correctly set. Please make sure that the generation kwargs are set correctly, or review your training hyperparameters.\n",
      "  f\"KL divergence is starting to become negative: {mean_kl.item():.2f} - this might be a precursor for failed training.\"\n",
      "13it [06:59, 32.21s/it]/opt/conda/lib/python3.7/site-packages/trl/trainer/ppo_trainer.py:1089: UserWarning: KL divergence is starting to become negative: -0.57 - this might be a precursor for failed training. sometimes this happens because the generation kwargs are not correctly set. Please make sure that the generation kwargs are set correctly, or review your training hyperparameters.\n",
      "  f\"KL divergence is starting to become negative: {mean_kl.item():.2f} - this might be a precursor for failed training.\"\n",
      "14it [07:31, 32.29s/it]/opt/conda/lib/python3.7/site-packages/trl/trainer/ppo_trainer.py:1089: UserWarning: KL divergence is starting to become negative: -0.29 - this might be a precursor for failed training. sometimes this happens because the generation kwargs are not correctly set. Please make sure that the generation kwargs are set correctly, or review your training hyperparameters.\n",
      "  f\"KL divergence is starting to become negative: {mean_kl.item():.2f} - this might be a precursor for failed training.\"\n",
      "16it [08:31, 31.07s/it]/opt/conda/lib/python3.7/site-packages/trl/trainer/ppo_trainer.py:1089: UserWarning: KL divergence is starting to become negative: -0.34 - this might be a precursor for failed training. sometimes this happens because the generation kwargs are not correctly set. Please make sure that the generation kwargs are set correctly, or review your training hyperparameters.\n",
      "  f\"KL divergence is starting to become negative: {mean_kl.item():.2f} - this might be a precursor for failed training.\"\n",
      "17it [09:05, 31.69s/it]/opt/conda/lib/python3.7/site-packages/trl/trainer/ppo_trainer.py:1089: UserWarning: KL divergence is starting to become negative: -0.38 - this might be a precursor for failed training. sometimes this happens because the generation kwargs are not correctly set. Please make sure that the generation kwargs are set correctly, or review your training hyperparameters.\n",
      "  f\"KL divergence is starting to become negative: {mean_kl.item():.2f} - this might be a precursor for failed training.\"\n",
      "18it [09:36, 32.05s/it]\n"
     ]
    }
   ],
   "source": [
    "# We then define the arguments to pass to the `generate` function. These arguments\n",
    "# are passed to the `generate` function of the PPOTrainer, which is a wrapper around\n",
    "# the `generate` function of the trained model.\n",
    "generation_kwargs = {\n",
    "    # \"min_length\": -1,\n",
    "    \"top_k\": 0.0,\n",
    "    \"top_p\": 1.0,\n",
    "    \"do_sample\": True,\n",
    "    \"pad_token_id\": tokenizer.pad_token_id,\n",
    "    \"eos_token_id\": 100_000,\n",
    "}\n",
    "output_min_length = 32\n",
    "output_max_length = script_args.output_max_length\n",
    "output_length_sampler = LengthSampler(output_min_length, output_max_length)\n",
    "\n",
    "for epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):\n",
    "    prompt_tensors = batch[\"input_ids\"]\n",
    "\n",
    "    batch[\"response\"] = []\n",
    "    response_tensors = []\n",
    "    prompt_tensors_as_tensor = []\n",
    "    for prompt_tensor in prompt_tensors:\n",
    "        #print(prompt_tensor[0])\n",
    "        prompt_tensor_as_tensor = torch.as_tensor(prompt_tensor[0], device=\"cuda\")\n",
    "        prompt_tensors_as_tensor.append(prompt_tensor_as_tensor)\n",
    "        response_tensor = ppo_trainer.generate(            \n",
    "            prompt_tensor_as_tensor,\n",
    "            return_prompt=False,\n",
    "            length_sampler=output_length_sampler,\n",
    "            **generation_kwargs,\n",
    "        )\n",
    "        response_tensors.append(response_tensor[0])\n",
    "        #print(response_tensor)\n",
    "        #batch[\"response\"] = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)\n",
    "        decoded_response = tokenizer.decode(response_tensor[0], skip_special_tokens=True)\n",
    "        batch[\"response\"].append(decoded_response)\n",
    "\n",
    "    # Compute sentiment score\n",
    "    texts = [q + r for q, r in zip(batch[\"query\"], batch[\"response\"])]\n",
    "    #print(texts)\n",
    "    pipe_outputs = sentiment_pipe(texts, **sent_kwargs)\n",
    "    #print(pipe_outputs)\n",
    "    rewards = []\n",
    "    for output in pipe_outputs:\n",
    "        #print(output)\n",
    "        reward = torch.tensor(output[0][\"score\"] - script_args.reward_baseline)\n",
    "        rewards.append(reward)\n",
    "\n",
    "    # Run PPO step    \n",
    "    stats = ppo_trainer.step(prompt_tensors_as_tensor, response_tensors, rewards)\n",
    "    ppo_trainer.log_stats(stats, batch, rewards)\n",
    "\n",
    "    if script_args.save_freq and epoch and epoch % script_args.save_freq == 0:\n",
    "        ppo_trainer.tokenizer.save_pretrained(fine_tuned_with_ranking_rewards)\n",
    "        ppo_trainer.accelerator.unwrap_model(ppo_trainer.model).save_pretrained(fine_tuned_with_ranking_rewards)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "260edbff-4010-4c78-b126-0521387af07f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11811828-0d71-4d68-a976-35e419ed7f25",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "availableInstances": [
   {
    "_defaultOrder": 0,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.t3.medium",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 1,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.t3.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 2,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.t3.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 3,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.t3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 4,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 5,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 6,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 7,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 8,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 9,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 10,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 11,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 12,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5d.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 13,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5d.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 14,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5d.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 15,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5d.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 16,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5d.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 17,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5d.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 18,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5d.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 19,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5d.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 20,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": true,
    "memoryGiB": 0,
    "name": "ml.geospatial.interactive",
    "supportedImageNames": [
     "sagemaker-geospatial-v1-0"
    ],
    "vcpuNum": 0
   },
   {
    "_defaultOrder": 21,
    "_isFastLaunch": true,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.c5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 22,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.c5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 23,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.c5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 24,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.c5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 25,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 72,
    "name": "ml.c5.9xlarge",
    "vcpuNum": 36
   },
   {
    "_defaultOrder": 26,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 96,
    "name": "ml.c5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 27,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 144,
    "name": "ml.c5.18xlarge",
    "vcpuNum": 72
   },
   {
    "_defaultOrder": 28,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.c5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 29,
    "_isFastLaunch": true,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g4dn.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 30,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g4dn.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 31,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g4dn.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 32,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g4dn.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 33,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g4dn.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 34,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g4dn.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 35,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 61,
    "name": "ml.p3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 36,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 244,
    "name": "ml.p3.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 37,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 488,
    "name": "ml.p3.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 38,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.p3dn.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 39,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.r5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 40,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.r5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 41,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.r5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 42,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.r5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 43,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.r5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 44,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.r5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 45,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 512,
    "name": "ml.r5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 46,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.r5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 47,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 48,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 49,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 50,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 51,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 52,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 53,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.g5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 54,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.g5.48xlarge",
    "vcpuNum": 192
   }
  ],
  "instance_type": "ml.g5.4xlarge",
  "kernelspec": {
   "display_name": "Python 3 (Data Science)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-1:081325390199:image/datascience-1.0"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}