training2.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The primary codes below are based on [akpe12/JP-KR-ocr-translator-for-travel](https://github.com/akpe12/JP-KR-ocr-translator-for-travel)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TrHlPFqwFAgj"
   },
   "source": [
    "## Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "t-jXeSJKE1WM"
   },
   "outputs": [],
   "source": [
    "from typing import Dict, List\n",
    "import csv\n",
    "\n",
    "import os\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "\n",
    "import datasets\n",
    "import torch\n",
    "from transformers import (\n",
    "    PreTrainedTokenizerFast,\n",
    "    DataCollatorForSeq2Seq,\n",
    "    Seq2SeqTrainingArguments,\n",
    "    BertJapaneseTokenizer,\n",
    "    GPT2TokenizerFast,\n",
    "    Seq2SeqTrainer,\n",
    "    Trainer\n",
    ")\n",
    "from transformers.models.encoder_decoder.modeling_encoder_decoder import EncoderDecoderModel\n",
    "\n",
    "from datasets import load_dataset\n",
    "\n",
    "encoder_model_name = \"cl-tohoku/bert-base-japanese-v2\"\n",
    "# decoder_model_name = \"skt/kogpt2-base-v2\"\n",
    "decoder_model_name = \"openai-community/gpt2\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nEW5trBtbykK"
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "# device = torch.device(\"cpu\")\n",
    "device, torch.cuda.device_count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GPT2TokenizerFast2(GPT2TokenizerFast):\n",
    "    def build_inputs_with_special_tokens(self, token_ids: List[int]) -> List[int]:\n",
    "        # print(f'adding {self.eos_token_id}')\n",
    "        return token_ids + [self.eos_token_id]   \n",
    "\n",
    "src_tokenizer = BertJapaneseTokenizer.from_pretrained(encoder_model_name)\n",
    "src_tokenizer.model_max_length = 512\n",
    "trg_tokenizer = GPT2TokenizerFast2.from_pretrained(decoder_model_name)\n",
    "trg_tokenizer.model_max_length = 512\n",
    "# trg_tokenizer = GPT2Tokenizer.from_pretrained(decoder_model_name, bos_token='</s>', eos_token='</s>', unk_token='<unk>',\n",
    "#   pad_token='<pad>', mask_token='<mask>')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DTf4U1fmFQFh"
   },
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "65L4O1c5FLKt"
   },
   "outputs": [],
   "source": [
    "class PairedDataset:\n",
    "    def __init__(self, \n",
    "        source_tokenizer: PreTrainedTokenizerFast, target_tokenizer: PreTrainedTokenizerFast,\n",
    "        file_path: str = None,\n",
    "        dataset_raw: datasets.Dataset = None\n",
    "    ):\n",
    "        self.src_tokenizer = source_tokenizer\n",
    "        self.trg_tokenizer = target_tokenizer\n",
    "        \n",
    "        if file_path is not None:\n",
    "            with open(file_path, 'r') as fd:\n",
    "                reader = csv.reader(fd)\n",
    "                next(reader)\n",
    "                self.data = [row for row in reader]\n",
    "        elif dataset_raw is not None:\n",
    "            self.data = dataset_raw\n",
    "        else:\n",
    "            raise ValueError('file_path or dataset_raw must be specified')\n",
    "\n",
    "    def __getitem__(self, index: int) -> Dict[str, torch.Tensor]:\n",
    "#         with open('train_log.txt', 'a+') as log_file:\n",
    "#             log_file.write(f'reading data[{index}] {self.data[index]}\\n')\n",
    "        if isinstance(self.data, datasets.Dataset):\n",
    "            src, trg = self.data[index]['sourceString'], self.data[index]['targetString']\n",
    "        else:\n",
    "            src, trg = self.data[index]\n",
    "        embeddings = self.src_tokenizer(src, return_attention_mask=False, return_token_type_ids=False, max_length=512, truncation=True)\n",
    "        embeddings['labels'] = self.trg_tokenizer.build_inputs_with_special_tokens(self.trg_tokenizer(trg, return_attention_mask=False)['input_ids'])\n",
    "\n",
    "        return embeddings\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# DATA_ROOT = '/home/tikim/code/ffat2json/output'\n",
    "# FILE_FFAC_FULL = 'ffac_full.csv'\n",
    "# FILE_FFAC_TEST = 'ffac_test.csv'\n",
    "# FILE_JA_KO_TRAIN = 'ja_ko_train.csv'\n",
    "# FILE_JA_KO_TEST = 'ja_ko_test.csv'\n",
    "\n",
    "# train_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_FFAC_FULL}')\n",
    "# eval_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_FFAC_TEST}') \n",
    "\n",
    "# train_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_JA_KO_TRAIN}')\n",
    "# eval_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_JA_KO_TEST}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_ROOT = '/home/tikim/dataset/jaen'\n",
    "FILE_TRAIN = 'train.csv'\n",
    "FILE_EVAL = 'dev.csv'\n",
    "FILE_TEST = 'test.csv'\n",
    "\n",
    "train_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_TRAIN}')\n",
    "eval_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_EVAL}')\n",
    "\n",
    "# train_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_EVAL}')\n",
    "# eval_dataset = PairedDataset(src_tokenizer, trg_tokenizer, file_path=f'{DATA_ROOT}/{FILE_TEST}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# be sure to check the column count of each dataset if you encounter \"ValueError: too many values to unpack (expected 2)\"\n",
    "# at the `src, trg = self.data[index]`\n",
    "# The `cat ffac_full.csv tteb_train.csv > ja_ko_train.csv` command may be the reason.\n",
    "# the last row of first csv and first row of second csv is merged and that's why 3rd column is created (which arouse ValueError)\n",
    "# debug_data = train_dataset.data\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uCBiLouSFiZY"
   },
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "I7uFbFYJFje8"
   },
   "outputs": [],
   "source": [
    "print(f'bos_token_id of trg_tokenizer: {trg_tokenizer.bos_token_id}')\n",
    "model = EncoderDecoderModel.from_encoder_decoder_pretrained(\n",
    "    encoder_model_name,\n",
    "    decoder_model_name,\n",
    "    pad_token_id=trg_tokenizer.bos_token_id,\n",
    ")\n",
    "model.config.decoder_start_token_id = trg_tokenizer.bos_token_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YFq2GyOAUV0W"
   },
   "outputs": [],
   "source": [
    "# for Trainer\n",
    "import wandb\n",
    "\n",
    "collate_fn = DataCollatorForSeq2Seq(src_tokenizer, model)\n",
    "wandb.init(project=\"fftr-poc1-en\", name='jbert+gpt2')\n",
    "\n",
    "arguments = Seq2SeqTrainingArguments(\n",
    "    output_dir='dump',\n",
    "    do_train=True,\n",
    "    do_eval=True,\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    num_train_epochs=6,\n",
    "    # num_train_epochs=25,\n",
    "    # per_device_train_batch_size=1,\n",
    "    # per_device_train_batch_size=30, # takes 40GB\n",
    "    per_device_train_batch_size=40,\n",
    "    per_device_eval_batch_size=2,\n",
    "    # per_device_eval_batch_size=30,\n",
    "    # per_device_eval_batch_size=64,\n",
    "    warmup_ratio=0.1,\n",
    "    gradient_accumulation_steps=4,\n",
    "    save_total_limit=5,\n",
    "    dataloader_num_workers=1,\n",
    "    # fp16=True, # ENABLE if CUDA is enabled\n",
    "    load_best_model_at_end=True,\n",
    "    report_to='wandb'\n",
    ")\n",
    "\n",
    "trainer = Seq2SeqTrainer(\n",
    "    model,\n",
    "    arguments,\n",
    "    data_collator=collate_fn,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=eval_dataset\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pPsjDHO5Vc3y"
   },
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_T4P4XunmK-C"
   },
   "outputs": [],
   "source": [
    "# model = EncoderDecoderModel.from_encoder_decoder_pretrained(\"xlm-roberta-base\",  \"skt/kogpt2-base-v2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7vTqAgW6Ve3J"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    trainer.train()\n",
    "    model.save_pretrained(\"dump/best_model\")\n",
    "    src_tokenizer.save_pretrained(\"dump/best_model/src_tokenizer\")\n",
    "    trg_tokenizer.save_pretrained(\"dump/best_model/trg_tokenizer\")\n",
    "    print('Training finished')\n",
    "except Exception as e:\n",
    "    print('Training failed')\n",
    "    print(e)\n",
    "finally:\n",
    "    wandb.finish()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "machine_shape": "hm",
   "provenance": []
  },
  "gpuClass": "premium",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}