1.dpo_trl训练.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "12276e41",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/anaconda3/envs/pt2/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "Using sep_token, but it is not set yet.\n",
      "Using cls_token, but it is not set yet.\n",
      "Using mask_token, but it is not set yet.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "GPT2TokenizerFast(name_or_path='tokenizer/gpt2', vocab_size=50257, model_max_length=1024, is_fast=True, padding_side='right', truncation_side='right', special_tokens={'bos_token': '<|endoftext|>', 'eos_token': '<|endoftext|>', 'unk_token': '<|endoftext|>', 'pad_token': '!'}, clean_up_tokenization_spaces=True),  added_tokens_decoder={\n",
       "\t50256: AddedToken(\"<|endoftext|>\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
       "}"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "import random\n",
    "import torch\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained('tokenizer/gpt2')\n",
    "tokenizer.pad_token_id = 0\n",
    "\n",
    "tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "58af6f4a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(DatasetDict({\n",
       "     train: Dataset({\n",
       "         features: ['prompt', 'chosen', 'rejected'],\n",
       "         num_rows: 71745\n",
       "     })\n",
       "     test: Dataset({\n",
       "         features: ['prompt', 'chosen', 'rejected'],\n",
       "         num_rows: 200\n",
       "     })\n",
       " }),\n",
       " {'prompt': 'context:CREATE TABLE TV_series (SHARE INTEGER) question:What is minimum and maximum share of TV series? answer:',\n",
       "  'chosen': 'SELECT MAX(SHARE), MIN(SHARE) FROM TV_series',\n",
       "  'rejected': ''})"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_from_disk\n",
    "\n",
    "dataset = load_from_disk('dataset/b-mc2/sql-create-context')['train']\n",
    "\n",
    "\n",
    "def f(data):\n",
    "    question = 'context:%s question:%s answer:' % (data['context'],\n",
    "                                                   data['question'])\n",
    "    answer = data['answer']\n",
    "    return {'question': question, 'answer': answer}\n",
    "\n",
    "\n",
    "dataset = dataset.map(f, remove_columns=['context'])\n",
    "\n",
    "\n",
    "def f(data):\n",
    "    question = len(tokenizer.encode(data['question']))\n",
    "    answer = len(tokenizer.encode(data['answer']))\n",
    "    return 25 <= question <= 65 and 10 <= answer <= 35\n",
    "\n",
    "\n",
    "dataset = dataset.filter(f)\n",
    "\n",
    "\n",
    "def f(data):\n",
    "    return {\n",
    "        'prompt': data['question'],\n",
    "        'chosen': data['answer'],\n",
    "        'rejected': ''\n",
    "    }\n",
    "\n",
    "\n",
    "dataset = dataset.map(f, remove_columns=['question', 'answer'])\n",
    "dataset = dataset.train_test_split(test_size=200)\n",
    "\n",
    "dataset, dataset['train'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "88954533",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM\n",
    "\n",
    "model_dpo = AutoModelForCausalLM.from_pretrained('model/gpt2').to('cuda')\n",
    "model_dpo_ref = AutoModelForCausalLM.from_pretrained('model/gpt2').to('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8dba89d3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'context:CREATE TABLE department (num_employees INTEGER, ranking INTEGER) question:What is the average number of employees of the departments whose rank is between 10 and 15? answer:10\\n\\nThe answer is:\\n\\nThe answer is:\\n\\nThe answer is:\\n\\nThe answer is:\\n\\nThe answer is:\\n\\nThe answer'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import random\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def generate(input_ids):\n",
    "    lens = input_ids.shape[1]\n",
    "    while True:\n",
    "        out = model_dpo(input_ids=input_ids)\n",
    "        topk = out['logits'][0, -1].topk(1)\n",
    "\n",
    "        values = topk.values.softmax(0).tolist()\n",
    "        indices = topk.indices.tolist()\n",
    "        next_word = random.choices(indices, weights=values)\n",
    "\n",
    "        next_word = torch.LongTensor(next_word).unsqueeze(0).to('cuda')\n",
    "        input_ids = torch.cat([input_ids, next_word], dim=1)\n",
    "\n",
    "        if input_ids.shape[1] - lens >= 35:\n",
    "            break\n",
    "\n",
    "        if input_ids[0, -1] == tokenizer.eos_token_id:\n",
    "            break\n",
    "\n",
    "    return input_ids\n",
    "\n",
    "\n",
    "input_ids = 'context:CREATE TABLE department (num_employees INTEGER, ranking INTEGER) question:What is the average number of employees of the departments whose rank is between 10 and 15? answer:'\n",
    "input_ids = tokenizer.encode(input_ids, return_tensors='pt').to('cuda')\n",
    "\n",
    "out = generate(input_ids)\n",
    "\n",
    "tokenizer.decode(out[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "16be7af3",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/anaconda3/envs/pt2/lib/python3.10/site-packages/trl/trainer/ppo_config.py:141: UserWarning: The `optimize_cuda_cache` arguement will be deprecated soon, please use `optimize_device_cache` instead.\n",
      "  warnings.warn(\n",
      "/root/anaconda3/envs/pt2/lib/python3.10/site-packages/trl/trainer/dpo_trainer.py:291: UserWarning: When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your TrainingArguments we have set it for you, but you should do it yourself in the future.\n",
      "  warnings.warn(\n",
      "Could not estimate the number of tokens of the input, floating-point operations will not be computed\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='2000' max='2000' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [2000/2000 19:19, Epoch 0/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>500</td>\n",
       "      <td>0.001800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1000</td>\n",
       "      <td>0.000400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1500</td>\n",
       "      <td>0.000500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2000</td>\n",
       "      <td>0.000300</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100\n",
      "context:CREATE TABLE table_25330991_3 (james_e_holmes VARCHAR, reidsville VARCHAR) question:Name the james e. holmes for erselle young answer:SELECT name FROM table_25330991_3 WHERE reidsville = \"young\"<|endoftext|>\n",
      "=================\n",
      "SELECT james_e_holmes FROM table_25330991_3 WHERE reidsville = \"Erselle Young\"\n",
      "=================\n",
      "200\n",
      "context:CREATE TABLE table_name_8 (weekly_rank VARCHAR, official_ratings__millions_ VARCHAR, show VARCHAR) question:Which Weekly Rank for a Live Final Show has an Official Ratings (millions) greater than 5.2? answer:SELECT weekly_rank FROM table_name_8 WHERE official_ratings__millions_ = 5.2 AND show = \"Live Final Show\" AND show = \"\n",
      "=================\n",
      "SELECT weekly_rank FROM table_name_8 WHERE official_ratings__millions_ > 5.2 AND show = \"live final\"\n",
      "=================\n",
      "300\n",
      "context:CREATE TABLE table_11173827_1 (english_title VARCHAR, finale VARCHAR, peak VARCHAR) question:What is the english title that has finale as 33 and peak as 42? answer:SELECT title FROM table_11173827_1 WHERE finale = 33 AND peak = 42<|endoftext|>\n",
      "=================\n",
      "SELECT english_title FROM table_11173827_1 WHERE finale = 33 AND peak = 42\n",
      "=================\n",
      "400\n",
      "context:CREATE TABLE table_name_80 (year INTEGER, date VARCHAR, designated_home VARCHAR) question:What year had a date of TBA where the Oakland raiders were the home team? answer:SELECT year FROM table_name_80 WHERE designated_home = \"Oakland\"<|endoftext|>\n",
      "=================\n",
      "SELECT AVG(year) FROM table_name_80 WHERE date = \"tba\" AND designated_home = \"oakland raiders\"\n",
      "=================\n",
      "500\n",
      "context:CREATE TABLE table_name_66 (cost___us$__ VARCHAR, open_source VARCHAR) question:What is the cost of an Open Source that is no? answer:SELECT cost___us$__ FROM table_name_66 WHERE open_source = \"Open Source\"<|endoftext|>\n",
      "=================\n",
      "SELECT cost___us$__ FROM table_name_66 WHERE open_source = \"no\"\n",
      "=================\n",
      "600\n",
      "context:CREATE TABLE table_name_11 (winner VARCHAR, year VARCHAR) question:What is Winner, when Year is 2013? answer:SELECT Winner FROM table_name_11 WHERE year = 2013 AND year = 2013 AND year = 2013 AND year = 2013 AND year = 2013 AND year = 2013 AND year =\n",
      "=================\n",
      "SELECT winner FROM table_name_11 WHERE year = 2013\n",
      "=================\n",
      "700\n",
      "context:CREATE TABLE table_name_44 (years VARCHAR, displacement VARCHAR) question:Which years have a displacement of 1816cc? answer:SELECT years FROM table_name_44 WHERE displacement = 1816cc<|endoftext|>\n",
      "=================\n",
      "SELECT years FROM table_name_44 WHERE displacement = \"1816cc\"\n",
      "=================\n",
      "800\n",
      "context:CREATE TABLE table_26982362_2 (original_airdate VARCHAR, production_code VARCHAR) question:The episode with production code 693-002, has how many original airdates? answer:SELECT production_code FROM table_26982362_2 WHERE production_code = 693-002<|endoftext|>\n",
      "=================\n",
      "SELECT COUNT(original_airdate) FROM table_26982362_2 WHERE production_code = \"693-002\"\n",
      "=================\n",
      "900\n",
      "context:CREATE TABLE table_name_24 (games INTEGER, term_ VARCHAR, c_ VARCHAR) question:What is the of games when for the term [c] of 1969 – 1973? answer:SELECT games FROM table_name_24 WHERE term_ = \"1969\"<|endoftext|>\n",
      "=================\n",
      "SELECT SUM(games) FROM table_name_24 WHERE term_[c_] = \"1969 – 1973\"\n",
      "=================\n",
      "1000\n",
      "context:CREATE TABLE table_18498743_1 (october_20 VARCHAR, _2008 VARCHAR, mexico VARCHAR) question:what is the october 20, 2008 stat where mexico stat is romania answer:SELECT date FROM table_18498743_1 WHERE date = \"2008\"<|endoftext|>\n",
      "=================\n",
      "SELECT october_20, _2008 FROM table_18498743_1 WHERE mexico = \"Romania\"\n",
      "=================\n",
      "1100\n",
      "context:CREATE TABLE table_3005915_3 (starts INTEGER) question:What is the maximum number of starts? answer:SELECT start FROM table_3005915_3 WHERE start < 3005915_3<|endoftext|>\n",
      "=================\n",
      "SELECT MAX(starts) FROM table_3005915_3\n",
      "=================\n",
      "1200\n",
      "context:CREATE TABLE table_name_17 (date VARCHAR, venue VARCHAR) question:What was the date of the game at Lake Oval? answer:SELECT date FROM table_name_17 WHERE venue = \"Lake Oval\"<|endoftext|>\n",
      "=================\n",
      "SELECT date FROM table_name_17 WHERE venue = \"lake oval\"\n",
      "=================\n",
      "1300\n",
      "context:CREATE TABLE table_name_21 (player VARCHAR, total VARCHAR, year_s__won VARCHAR) question:Who won in 1988 with a total less than 287? answer:SELECT player FROM table_name_21 WHERE total = 287<|endoftext|>\n",
      "=================\n",
      "SELECT player FROM table_name_21 WHERE total < 287 AND year_s__won = \"1988\"\n",
      "=================\n",
      "1400\n",
      "context:CREATE TABLE table_23224961_1 (oberbayern_b VARCHAR, oberpfalz VARCHAR) question:When fc schwandorf is the oberpfalz what is the oberbayern b? answer:SELECT fc, oberpfalz FROM table_23224961_1 WHERE oberpfalz = \"schwandorf\"<|endoftext|>\n",
      "=================\n",
      "SELECT oberbayern_b FROM table_23224961_1 WHERE oberpfalz = \"FC Schwandorf\"\n",
      "=================\n",
      "1500\n",
      "context:CREATE TABLE table_name_88 (lane INTEGER, name VARCHAR) question:What is the average lane that is called rebecca brown? answer:SELECT MAX(lane) FROM table_name_88 WHERE name = \"rebecca brown\"<|endoftext|>\n",
      "=================\n",
      "SELECT AVG(lane) FROM table_name_88 WHERE name = \"rebecca brown\"\n",
      "=================\n",
      "1600\n",
      "context:CREATE TABLE table_name_24 (games INTEGER, term_ VARCHAR, c_ VARCHAR) question:What is the of games when for the term [c] of 1969 – 1973? answer:SELECT games FROM table_name_24 WHERE term_ = \"1969\"<|endoftext|>\n",
      "=================\n",
      "SELECT SUM(games) FROM table_name_24 WHERE term_[c_] = \"1969 – 1973\"\n",
      "=================\n",
      "1700\n",
      "context:CREATE TABLE table_name_87 (type VARCHAR, location VARCHAR) question:What type of Bridge is in Stanley? answer:SELECT type FROM table_name_87 WHERE location = \"Stanley\"<|endoftext|>\n",
      "=================\n",
      "SELECT type FROM table_name_87 WHERE location = \"stanley\"\n",
      "=================\n",
      "1800\n",
      "context:CREATE TABLE table_name_34 (attendance VARCHAR, date VARCHAR, week VARCHAR) question:What was the attendance of the game on December 13, 1970? answer:SELECT attendance FROM table_name_34 WHERE date = \"2013-12-13T00:00:00Z\"<|endoftext|>\n",
      "=================\n",
      "SELECT COUNT(attendance) FROM table_name_34 WHERE date = \"december 13, 1970\" AND week > 13\n",
      "=================\n",
      "1900\n",
      "context:CREATE TABLE constructors (nationality VARCHAR) question:What are the numbers of constructors for different nationalities? answer:SELECT nationality FROM constructors WHERE nationality = \"nationality\" AND country = \"nationality\"<|endoftext|>\n",
      "=================\n",
      "SELECT COUNT(*), nationality FROM constructors GROUP BY nationality\n",
      "=================\n",
      "2000\n",
      "context:CREATE TABLE table_name_93 (home_team VARCHAR, venue VARCHAR) question:Which home team played at MCG? answer:SELECT venue FROM table_name_93 WHERE venue = \"MCG\"<|endoftext|>\n",
      "=================\n",
      "SELECT home_team FROM table_name_93 WHERE venue = \"mcg\"\n",
      "=================\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=2000, training_loss=0.0007831706330180168, metrics={'train_runtime': 1160.3168, 'train_samples_per_second': 27.579, 'train_steps_per_second': 1.724, 'total_flos': 0.0, 'train_loss': 0.0007831706330180168, 'epoch': 0.45})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import TrainingArguments, TrainerCallback\n",
    "from trl import DPOTrainer\n",
    "import random\n",
    "\n",
    "args = TrainingArguments(per_device_train_batch_size=16,\n",
    "                         max_steps=2000,\n",
    "                         learning_rate=1e-5,\n",
    "                         evaluation_strategy='no',\n",
    "                         optim='rmsprop',\n",
    "                         report_to='none',\n",
    "                         save_strategy='no',\n",
    "                         output_dir='output_dir')\n",
    "\n",
    "\n",
    "class MyCallback(TrainerCallback):\n",
    "\n",
    "    def on_step_end(self, args, state, control, **kwargs):\n",
    "        if state.global_step % 100 == 0:\n",
    "            print(state.global_step)\n",
    "\n",
    "            data = random.choice(dataset['test'])\n",
    "            input_ids = tokenizer.encode(data['prompt'],\n",
    "                                         return_tensors='pt').to('cuda')\n",
    "\n",
    "            out = generate(input_ids)\n",
    "\n",
    "            print(tokenizer.decode(out[0]))\n",
    "            print('=================')\n",
    "            print(data['chosen'])\n",
    "            print('=================')\n",
    "\n",
    "\n",
    "trainer = DPOTrainer(model_dpo,\n",
    "                     model_dpo_ref,\n",
    "                     args=args,\n",
    "                     beta=0.1,\n",
    "                     train_dataset=dataset['train'],\n",
    "                     tokenizer=tokenizer,\n",
    "                     max_length=100,\n",
    "                     max_target_length=100,\n",
    "                     max_prompt_length=100,\n",
    "                     callbacks=[MyCallback()])\n",
    "\n",
    "trainer.train()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pt2]",
   "language": "python",
   "name": "conda-env-pt2-py"
  },
  "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}