example/2022-12-10-textrl-elon-musk.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "At5gZSqIG1ah"
   },
   "source": [
    "# Controllable generation via RL to let Elon Musk speak ill of DOGE\n",
    "> How to control text generation through a sentiment classifier.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tgBsD1fa0hJn"
   },
   "outputs": [],
   "source": [
    "!pip install pfrl@git+https://github.com/voidful/pfrl.git\n",
    "!pip install textrl==0.2.15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "c8BT6ZwpNke-"
   },
   "outputs": [],
   "source": [
    "from textrl import TextRLEnv,TextRLActor\n",
    "from transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer, AutoModelWithLMHead\n",
    "import logging\n",
    "import sys\n",
    "import pfrl\n",
    "import torch\n",
    "logging.basicConfig(level=logging.INFO, stream=sys.stdout, format='')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Re1cxoPZ4wgf"
   },
   "source": [
    "**Using a pre-trained model, it can generate elonmusk's style tweets.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "f0TqcFITHHdX"
   },
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"huggingtweets/elonmusk\")  \n",
    "model = AutoModelWithLMHead.from_pretrained(\"huggingtweets/elonmusk\")\n",
    "model.eval()\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "342ChdlM5CXv"
   },
   "source": [
    "**a sentiment classifier for rl reward**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "UYRgFPW_HrJo",
    "outputId": "a78366c1-d0ba-4220-97f2-d3fc0577a99a",
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "sentiment = pipeline('sentiment-analysis',model=\"cardiffnlp/twitter-roberta-base-sentiment\",tokenizer=\"cardiffnlp/twitter-roberta-base-sentiment\",device=0,return_all_scores=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "o47CRT8TKqvn"
   },
   "outputs": [],
   "source": [
    "transformers_logger = logging.getLogger('transformers')\n",
    "transformers_logger.setLevel(logging.CRITICAL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "GgFrS5gQIAxR",
    "outputId": "5336889f-b05b-415e-d368-e4d953a5591e"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[{'label': 'LABEL_0', 'score': 0.9338533878326416},\n",
       "  {'label': 'LABEL_1', 'score': 0.06011885032057762},\n",
       "  {'label': 'LABEL_2', 'score': 0.0060277231968939304}]]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sentiment(\"dogecoin is bad\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "GhVcFdwjJzjW",
    "outputId": "c397a19c-c53c-4013-e45b-069b65eaf615"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9338533878326416"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sentiment(\"dogecoin is bad\")[0][0]['score']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SfjPrvcK5N5e"
   },
   "source": [
    "set our text generation reward, inverse perplexity + sentiment classifier.\n",
    "- inverse perplexity make sure the generated sentence probability will be high.\n",
    "- sentiment classifier can make the generate more negative."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "OgNGUk99HMtT"
   },
   "outputs": [],
   "source": [
    "class MyRLEnv(TextRLEnv):\n",
    "    def get_reward(self, input_item, predicted_list, finish): # predicted will be the list of predicted token\n",
    "      reward = 0\n",
    "      if finish or len(predicted_list) >= self.env_max_length:\n",
    "        predicted_text = tokenizer.convert_tokens_to_string(predicted_list[0])\n",
    "        # sentiment classifier\n",
    "        reward = sentiment(input_item['input']+predicted_text)[0][0]['score'] * 10\n",
    "      return reward"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jqF7mNCY5tdO"
   },
   "source": [
    "**fit one example**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "cy4tCfslKGd4"
   },
   "outputs": [],
   "source": [
    "observaton_list = [{'input':'i think dogecoin is'}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "wtGfk03eHOv_"
   },
   "outputs": [],
   "source": [
    "env = MyRLEnv(model, tokenizer, observation_input=observaton_list,compare_sample=1)\n",
    "actor = TextRLActor(env,model,tokenizer)\n",
    "agent = actor.agent_ppo(update_interval=100, minibatch_size=3, epochs=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_sBBy1yjIdtP",
    "outputId": "62125633-6c76-47eb-c207-ebf33a99e78e"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[' a good idea<|endoftext|>']"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "actor.predict(observaton_list[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "FBysk9MiHR2D",
    "outputId": "4086dcd7-6d19-44bc-e1b0-fa764f873301",
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "pfrl.experiments.train_agent_with_evaluation(\n",
    "    agent,\n",
    "    env,\n",
    "    steps=300,\n",
    "    eval_n_steps=None,\n",
    "    eval_n_episodes=1,       \n",
    "    train_max_episode_len=100,  \n",
    "    eval_interval=10,\n",
    "    outdir='elon_musk_dogecoin', \n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9B7rMPRU5zsM"
   },
   "source": [
    "loading the best result and predict."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "FrkYGPjYTIcS"
   },
   "outputs": [],
   "source": [
    "agent.load(\"./elon_musk_dogecoin/best\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "dpAwe42ES-5w",
    "outputId": "bb12c0d2-1916-4076-8f98-b20d2a2e4e57"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[' a real thing<|endoftext|>']"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "actor.predict(observaton_list[0])"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "virtual-env-3.10.1",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.9.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "2c6c77b12b02a1c2aaa91a9fb9cc35bb3c4bbfb7b716f83ac7b2b57ffb1247cc"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}