GPT2_adapter.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import cache\n",
    "from adapter_args_helper import (\n",
    "    DataArguments,\n",
    "    ModelArguments,\n",
    "    TrainingArguments\n",
    ")\n",
    "from datasets import load_from_disk, load_metric, set_caching_enabled, DatasetDict\n",
    "from data_utils import load_dataset\n",
    "from itertools import chain\n",
    "from models import ClipCaptionModel\n",
    "from torch.nn.functional import cross_entropy\n",
    "from tqdm import tqdm\n",
    "from transformers import (\n",
    "    default_data_collator,\n",
    "    get_linear_schedule_with_warmup,\n",
    "    set_seed,\n",
    "    AdamW,\n",
    "    DataCollatorForLanguageModeling,\n",
    "    GPT2Config,\n",
    "    GPT2Tokenizer,\n",
    "    GPT2Model,\n",
    "    EarlyStoppingCallback,\n",
    "    HfArgumentParser,\n",
    "    Trainer\n",
    ")\n",
    "from transformers.trainer_utils import get_last_checkpoint, is_main_process\n",
    "\n",
    "import logging\n",
    "import math\n",
    "import numpy as np\n",
    "import os\n",
    "import sys\n",
    "import torch\n",
    "import transformers\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "from ppcm_models.pytorch_pretrained_bert.modeling_adapter import GPT2LMHeadModel, GPT2Config\n",
    "from utils.helper import load_model_recursive\n",
    "\n",
    "set_caching_enabled(True)\n",
    "logger = logging.getLogger(__name__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataArguments():\n",
    "    def __init__(self):\n",
    "        self.dataset_path = '/home/bryan/datasets/bookcorpusopen/bookcorpusopen_chunked.arrow'\n",
    "        self.bookcorpusopen_story_column_name = 'chunk'\n",
    "        self.preprocessing_num_workers = 8\n",
    "        self.genre='Romance'\n",
    "        self.adapter_id=1\n",
    "        self.match_up_to_n_genres=3\n",
    "        self.sample_row=None\n",
    "        \n",
    "class ModelArguments():\n",
    "    def __init__(self):\n",
    "        self.model_size = 'medium'\n",
    "        self.load_checkpoint_adapter = \"\"\n",
    "        self.max_seq_len=512\n",
    "        # self.lr = 2e-4 #, help=\"Learning rate\")\n",
    "\n",
    "class TrainingArguments(TrainingArguments):\n",
    "    def __init__(self):\n",
    "        self.output_dir = \"./save\"\n",
    "        self.eval_accumulation_steps = None\n",
    "        \n",
    "model_args = ModelArguments()\n",
    "data_args = DataArguments()\n",
    "training_args = TrainingArguments()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_args.model_path = f'ppcm_models/dialoGPT/{model_args.model_size}/'\n",
    "\n",
    "config = GPT2Config.from_json_file(os.path.join(model_args.model_path, 'config.json'))\n",
    "tokenizer = GPT2Tokenizer.from_pretrained(model_args.model_path)\n",
    "\n",
    "## Load either Adapters' checkpoint, or just finetuned DialoGPT\n",
    "if(model_args.load_checkpoint_adapter != \"\"):\n",
    "    print(\"Loading ADAPTERS\")\n",
    "    model = load_model_recursive(GPT2LMHeadModel(config), model_args.load_checkpoint_adapter, model_args, verbose=True)\n",
    "else:\n",
    "    model = load_model_recursive(GPT2LMHeadModel(config), model_args.model_path+f\"{model_args.model_size}_ft.pkl\", model_args, verbose=True)\n",
    "\n",
    "## Load GPT2 instead of DialoGPT\n",
    "\n",
    "pt_gpt2_model = GPT2Model.from_pretrained('gpt2-medium')\n",
    "\n",
    "model.transformer.wte.weight = pt_gpt2_model.wte.weight\n",
    "model.transformer.wpe.weight = pt_gpt2_model.wpe.weight\n",
    "\n",
    "layers = np.arange(0,len(pt_gpt2_model.h),1)\n",
    "for layer in layers:\n",
    "    model.transformer.h[layer].ln_1.weight = pt_gpt2_model.h[layer].ln_1.weight\n",
    "    model.transformer.h[layer].attn.c_attn.weight = pt_gpt2_model.h[layer].attn.c_attn.weight\n",
    "    model.transformer.h[layer].attn.c_proj.weight = pt_gpt2_model.h[layer].attn.c_proj.weight\n",
    "    model.transformer.h[layer].ln_2.weight = pt_gpt2_model.h[layer].ln_2.weight\n",
    "    model.transformer.h[layer].mlp.c_fc.weight = pt_gpt2_model.h[layer].mlp.c_fc.weight\n",
    "    model.transformer.h[layer].mlp.c_proj.weight = pt_gpt2_model.h[layer].mlp.c_proj.weight\n",
    "# model.to(model_args.device)\n",
    "print('GPT2 loaded instead DialoGPT')\n",
    "\n",
    "for n, p in model.named_parameters():\n",
    "    if \"adapter\" not in str(n):\n",
    "        p.requires_grad = False\n",
    "parameters_to_update = [p for n, p in model.named_parameters() if \"adapter\" in str(n)]\n",
    "# optimizer = AdamW(parameters_to_update, lr=model_args.lr, correct_bias=True)\n",
    "print('GPT2 param frozen, Adapter is trainable and initialized with AdamW')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BookcorpusopenGenreAdapterDataset(Dataset):\n",
    "    def __init__(self, data_args, split, tokenizer, genre=None, adapter_id=-1,\n",
    "                         sample_row=100, match_up_to_n_genres=None, truncate=True, \n",
    "                         max_seq_len=512, add_special_tokens=True,\n",
    "                         *args, **kwargs):\n",
    "        super(BookcorpusopenGenreAdapterDataset, self).__init__(*args, **kwargs)\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            adapter_id: int, adapter_id for the genre we want the adapter to be trained with\n",
    "        \"\"\"\n",
    "        \n",
    "        self.data_args = data_args\n",
    "        self.tokenizer = tokenizer\n",
    "        self.add_special_tokens = add_special_tokens\n",
    "        self.truncate = truncate\n",
    "        self.max_seq_len = max_seq_len\n",
    "        self.adapter_id = adapter_id\n",
    "        self.preprocessing_num_workers = data_args.preprocessing_num_workers\n",
    "        self.dataset = self.load_bookcorpusopen(split, genre, \n",
    "                                                match_up_to_n_genres,\n",
    "                                                sample_row)\n",
    "\n",
    "    def load_bookcorpusopen(self, split, genre='Fiction', \n",
    "                            match_up_to_n_genres=None, sample_row=None):\n",
    "        \"\"\"\n",
    "        Load bookcorpusopen from pyarrow file.\n",
    "        \n",
    "        Further improvement:\n",
    "        Group, concat, and truncate entries based on the adapter_id after tokenization\n",
    "            \n",
    "        Args:\n",
    "            split: string, {train, valid, test}\n",
    "            genre: string, genre that we want the adapter to be trained with, e.g. 'Fiction'\n",
    "            match_up_to_n_genres: int, how many of the firsts bookcorpusopen genres entries \n",
    "                                    is considered as a genre to match with the genre input.\n",
    "                                    None defaults to use all bookcorpusopen genres to match.\n",
    "            sample_row: int, set the int number to sample the dataset, \n",
    "                        None means using all the datasets samples available\n",
    "            match_up_to_n_genres\n",
    "            \n",
    "        Returns:\n",
    "            dataset: tokenized huggingface dataset format from one of the bookcorpusopen split, \n",
    "                        with the adapter_id attached, and without any adapter_id = -1\n",
    "        \"\"\"\n",
    "\n",
    "        def genre_match(entry_genres_string_list, genre, match_up_to_n_genres):\n",
    "            \"\"\"\n",
    "            True to the genre that match to match_up_to_n_genres genres from the entry_genres\n",
    "            else false\n",
    "            \"\"\"\n",
    "            story_genre_list = [genre[1:-1] for genre in entry_genres_string_list[1:-1].split(', ')]\n",
    "            story_genre_stringlist = \", \".join(story_genre_list[:match_up_to_n_genres])\n",
    "            \n",
    "            return genre.lower() in story_genre_stringlist.lower()\n",
    "        \n",
    "        def map_tokenization(batch):\n",
    "            self.tokenizer.pad_token = self.tokenizer.eos_token\n",
    "            tokenized = self.tokenizer(batch[self.data_args.bookcorpusopen_story_column_name], \n",
    "                                          truncation=self.truncate,\n",
    "                                          max_length=self.max_seq_len,\n",
    "                                          add_special_tokens=self.add_special_tokens)\n",
    "            return tokenized\n",
    "        \n",
    "        # Main data processing function that will concatenate all texts \n",
    "        # from our dataset and generate chunks of max_seq_len.\n",
    "        def group_texts(examples):\n",
    "            # Concatenate all texts.\n",
    "            concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}\n",
    "            total_length = len(concatenated_examples[list(examples.keys())[0]])\n",
    "            # We drop the small remainder, we could add padding if the model supported \n",
    "            # it instead of this drop, you can customize this part to your needs.\n",
    "            if total_length >= self.max_seq_len:\n",
    "                total_length = (total_length // self.max_seq_len) * self.max_seq_len\n",
    "            # Split by chunks of max_len.\n",
    "            result = {\n",
    "                k: [t[i : i + self.max_seq_len] \\\n",
    "                    for i in range(0, total_length, self.max_seq_len)]\n",
    "                for k, t in concatenated_examples.items()\n",
    "            }\n",
    "            return result\n",
    "        \n",
    "        # load bookcorpusopen from arrow file\n",
    "        datasets = DatasetDict()\n",
    "        print('Loading train, validation, test dataset...')\n",
    "        datasets = load_from_disk(self.data_args.dataset_path)\n",
    "        print('Loaded')\n",
    "        \n",
    "        # Select rows sampled and filter for the matching genres\n",
    "        sample_row = len(datasets[split]) if sample_row == None else sample_row\n",
    "        \n",
    "        dataset = datasets[split].select(np.arange(0,sample_row,1))\\\n",
    "                                .filter(lambda x: genre_match(x['genre'], genre, match_up_to_n_genres)\\\n",
    "                                        , num_proc=self.preprocessing_num_workers)\n",
    "\n",
    "        \n",
    "        # Tokenize with huggingface datasets mapping function\n",
    "        tokenized_dataset = dataset.map(\n",
    "            map_tokenization,\n",
    "            remove_columns=self.data_args.bookcorpusopen_story_column_name,\n",
    "            num_proc=self.preprocessing_num_workers,\n",
    "            load_from_cache_file=True\n",
    "        )\n",
    "        print(split, 'split tokenized')\n",
    "        \n",
    "        group_concatted_dataset = tokenized_dataset.map(\n",
    "                                        group_texts,\n",
    "                                        batched=True,\n",
    "                                        num_proc=self.preprocessing_num_workers,\n",
    "                                        load_from_cache_file=True,\n",
    "                                        desc=f\"Grouping texts in chunks of {self.max_seq_len}\",\n",
    "                                    )\n",
    "                                \n",
    "        return group_concatted_dataset\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "            \n",
    "        forward_inputs = {}\n",
    "        forward_inputs['task_id'] = self.adapter_id\n",
    "        forward_inputs['input_ids'] = [self.dataset[index]['input_ids']]\n",
    "        forward_inputs[\"labels\"] = forward_inputs[\"input_ids\"].copy()\n",
    "        \n",
    "        return forward_inputs\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.dataset.num_rows"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load and Checks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load the preprocessed dataset splits\n",
    "dataset_dict = {}\n",
    "for split in ['train', 'valid']:\n",
    "    dataset_dict[split] = BookcorpusopenGenreAdapterDataset(\n",
    "                                    data_args, split, tokenizer, genre=data_args.genre,\n",
    "                                    adapter_id=data_args.adapter_id, sample_row=data_args.sample_row,\n",
    "                                    match_up_to_n_genres=data_args.match_up_to_n_genres,\n",
    "                                    max_seq_len=model_args.max_seq_len)\n",
    "\n",
    "    for i in range(len(dataset_dict[split])):\n",
    "        input_ids_len = len(dataset_dict[split][i]['input_ids'][0])\n",
    "        if input_ids_len < model_args.max_seq_len:\n",
    "            print(split, i, input_ids_len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def remove_remainder(input_ids):\n",
    "    # print(len(input_ids))\n",
    "    return len(input_ids[0]) == self.max_seq_len"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "dataset_dict['train'].dataset.filter(lambda x: remove_remainder(x['input_ids'])\\\n",
    "                                                    , num_proc=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_dict['train'].dataset[0]['input_ids']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pix2story",
   "language": "python",
   "name": "pix2story"
  },
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}