04_trl_SFTTrainer_RewardTrainer_PPO.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/gulabpatel/LLMs/blob/main/04_trl_SFTTrainer_RewardTrainer_PPO.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# TRL - Transformer Reinforcement Learning\n",
        "\n",
        "\n",
        "\n",
        "trl is a full stack library where we provide a set of tools to train transformer language models and stable diffusion models with Reinforcement Learning, from the Supervised Fine-tuning step (SFT), Reward Modeling step (RM) to the Proximal Policy Optimization (PPO) step. The library is built on top of the transformers library by 🤗 Hugging Face. Therefore, pre-trained language models can be directly loaded via transformers. At this point most of decoder architectures and encoder-decoder architectures are supported. Refer to the documentation or the examples/ folder for example code snippets and how to run these tools."
      ],
      "metadata": {
        "id": "21-5Or68K-vq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install trl peft"
      ],
      "metadata": {
        "id": "_GGRtm81tVvn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## SFTTrainer\n",
        "\n",
        "This is a basic example on how to use the SFTTrainer from the library. The SFTTrainer is a light wrapper around the transformers Trainer to easily fine-tune language models or adapters on a custom dataset."
      ],
      "metadata": {
        "id": "kVfsYyv0OJxV"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cHX50__SOZzj"
      },
      "outputs": [],
      "source": [
        "# imports\n",
        "from datasets import load_dataset\n",
        "from trl import SFTTrainer\n",
        "\n",
        "# get dataset\n",
        "dataset = load_dataset(\"imdb\", split=\"train\")\n",
        "\n",
        "# get trainer\n",
        "trainer = SFTTrainer(\n",
        "    \"facebook/opt-350m\",\n",
        "    train_dataset=dataset,\n",
        "    dataset_text_field=\"text\",\n",
        "    max_seq_length=512,\n",
        ")\n",
        "\n",
        "# train\n",
        "trainer.train()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## RewardTrainer\n",
        "\n",
        "This is a basic example on how to use the RewardTrainer from the library. The RewardTrainer is a wrapper around the transformers Trainer to easily fine-tune reward models or adapters on a custom preference dataset."
      ],
      "metadata": {
        "id": "qHJpUVQ0OUy6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# imports\n",
        "from transformers import AutoModelForSequenceClassification, AutoTokenizer\n",
        "from trl import RewardTrainer\n",
        "\n",
        "# load model and dataset - dataset needs to be in a specific format\n",
        "model = AutoModelForSequenceClassification.from_pretrained(\"gpt2\", num_labels=1)\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n",
        "\n",
        "...\n",
        "\n",
        "# load trainer\n",
        "trainer = RewardTrainer(\n",
        "    model=model,\n",
        "    tokenizer=tokenizer,\n",
        "    train_dataset=dataset,\n",
        ")\n",
        "\n",
        "# train\n",
        "trainer.train()"
      ],
      "metadata": {
        "id": "KUsCCrGwKzxq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## PPOTrainer\n",
        "This is a basic example on how to use the PPOTrainer from the library. Based on a query the language model creates a response which is then evaluated. The evaluation could be a human in the loop or another model's output."
      ],
      "metadata": {
        "id": "zZGzi5IKOeJC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# imports\n",
        "import torch\n",
        "from transformers import AutoTokenizer\n",
        "from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead, create_reference_model\n",
        "from trl.core import respond_to_batch\n",
        "\n",
        "# get models\n",
        "model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')\n",
        "model_ref = create_reference_model(model)\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained('gpt2')\n",
        "\n",
        "# initialize trainer\n",
        "ppo_config = PPOConfig(\n",
        "    batch_size=1,\n",
        ")\n",
        "\n",
        "# encode a query\n",
        "query_txt = \"This morning I went to the \"\n",
        "query_tensor = tokenizer.encode(query_txt, return_tensors=\"pt\")\n",
        "\n",
        "# get model response\n",
        "response_tensor  = respond_to_batch(model, query_tensor)\n",
        "\n",
        "# create a ppo trainer\n",
        "ppo_trainer = PPOTrainer(ppo_config, model, model_ref, tokenizer)\n",
        "\n",
        "# define a reward for response\n",
        "# (this could be any reward such as human feedback or output from another model)\n",
        "reward = [torch.tensor(1.0)]\n",
        "\n",
        "# train model for one step with ppo\n",
        "train_stats = ppo_trainer.step([query_tensor[0]], [response_tensor[0]], reward)"
      ],
      "metadata": {
        "id": "_SN2Z5P4KzuO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "54buTMpMK39q"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "machine_shape": "hm",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}