train_t5_text_to_sql_batch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ===================== Training T5 Model ======================#\n",
    "# ==============================================================#\n",
    "\n",
    "import sys, os\n",
    "from os import environ\n",
    "from dotenv import load_dotenv\n",
    "import logging\n",
    "from logging.config import fileConfig\n",
    "\n",
    "import datasets\n",
    "import numpy as np\n",
    "import evaluate\n",
    "from datasets import load_dataset, DatasetDict\n",
    "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelWithLMHead\n",
    "from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments\n",
    "from transformers import Trainer, TrainingArguments, AdamW\n",
    "from transformers import DataCollatorForSeq2Seq\n",
    "\n",
    "# ==============================================================#\n",
    "# ====================== GLOBAL VARIABLES ======================#\n",
    "\n",
    "GlobalVar           = 0\n",
    "train_batch_size    = 8\n",
    "eval_batch_size     = 8\n",
    "max_input_length    = 512\n",
    "max_target_length   = 64\n",
    "\n",
    "MODEL_BASE_DIR      = \"/Users/sree/.cache/huggingface/hub\"\n",
    "IN_MODEL_NAME       = \"mrm8488/t5-base-finetuned-wikiSQL\"\n",
    "OUT_MODEL_NAME      = \"my--t5-base-finetuned-wiki-to-SQL\" ##-finetuned-{source_lang}-to-{target_lang}\"\n",
    "OUTPUT_DIR          = MODEL_BASE_DIR + \"/models--\" + OUT_MODEL_NAME \n",
    "LOG_DIR             = OUTPUT_DIR + \"/logs\"\n",
    "\n",
    "DS_LOCAL            = 'my_sql_data.json'\n",
    "DS_HUGGINGFACE      = 'wikisql'\n",
    "USE_LOCAL_DATASET   = False\n",
    "DATASET             = DS_LOCAL if USE_LOCAL_DATASET else DS_HUGGINGFACE\n",
    "USE_TKN_CACHE       = False\n",
    "DATA_SAMPLING       = True\n",
    "SAMPLING_SPLIT      = {'test':50, 'train':20, 'validation':20}\n",
    "DATASET_PATH        = \"./dataset/\"\n",
    "ENCODED_DATA_PATH   = \"./cache/encoded-\" + DATASET\n",
    "\n",
    "# modify data path in case only samples are being processed\n",
    "ENCODED_DATA_PATH   += \"-sample\" if DATA_SAMPLING else \"\"\n",
    "\n",
    "INPUT_FIELD         = 'question'\n",
    "RESULT_FIELD        = 'sql'\n",
    "NESTED_RESULT_FIELD = 'human_readable'\n",
    "TASK_PREFIX         = \"translate English to SQL : \"\n",
    "\n",
    "# use_fast=True param to speed up tokenization\n",
    "# initialize model & it's tokenizer\n",
    "tokenizer           = AutoTokenizer.from_pretrained(IN_MODEL_NAME)\n",
    "model               = AutoModelForSeq2SeqLM.from_pretrained(IN_MODEL_NAME)\n",
    "\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "\n",
    "# ==============================================================#\n",
    "# ====================== GLOBAL FUNCTIONS ======================#\n",
    "\n",
    "def compute_metric(eval_pred):\n",
    "    predictions, labels = eval_pred\n",
    "    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)\n",
    "\n",
    "    # Replace -100 in the labels as we can't decode them.\n",
    "    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n",
    "    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
    "\n",
    "    # Compute ROUGE scores\n",
    "    result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)\n",
    "\n",
    "    # Extract ROUGE f1 scores\n",
    "    result = {key: value.mid.fmeasure * 100 for key, value in result.items()}\n",
    "\n",
    "    # Add mean generated length to metrics\n",
    "    prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in predictions]\n",
    "\n",
    "    result[\"gen_len\"] = np.mean(prediction_lens)\n",
    "    return {k: round(v, 4) for k, v in result.items()}\n",
    "\n",
    "def get_training_args():\n",
    "    seq_training_args = Seq2SeqTrainingArguments(\n",
    "        output_dir=OUTPUT_DIR,\n",
    "        num_train_epochs=3,\n",
    "        save_total_limit=3,\n",
    "        evaluation_strategy =\"steps\",\n",
    "        eval_steps=100,\n",
    "        logging_strategy=\"steps\",\n",
    "        logging_steps=100,\n",
    "        save_strategy=\"steps\",\n",
    "        save_steps=200,\n",
    "        \n",
    "        per_device_train_batch_size=train_batch_size,\n",
    "        per_device_eval_batch_size=eval_batch_size,\n",
    "        predict_with_generate=True,\n",
    "        metric_for_best_model=\"rouge1\",\n",
    "        load_best_model_at_end=True,    \n",
    "\n",
    "        weight_decay=0.01,\n",
    "        learning_rate=4e-5,\n",
    "        optim = \"adamw_torch\",\n",
    "        report_to=[\"tensorboard\"]\n",
    "        #adam_beta1=optimizer,\n",
    "        #fp16=True, # can be used with CUDA devices, not CPU\n",
    "        #report_to=\"tensorboard\" #still to make it work\n",
    "    )\n",
    "    return seq_training_args\n",
    "\n",
    "def get_trainer(tokenized_dataset, training_args):\n",
    "    seq_trainer = Seq2SeqTrainer(\n",
    "        model           = model,\n",
    "        args            = training_args,\n",
    "        train_dataset   = tokenized_dataset[\"train\"],\n",
    "        eval_dataset    = tokenized_dataset[\"validation\"],\n",
    "        data_collator   = data_collator,\n",
    "        tokenizer       = tokenizer,\n",
    "        compute_metrics = compute_metric\n",
    "    )\n",
    "    return seq_trainer\n",
    "\n",
    "select_slice = lambda ds, split: ds[split].shuffle(seed=20).select(range(SAMPLING_SPLIT[split]))\n",
    "\n",
    "def load_data():\n",
    "    my_ds = load_dataset(\n",
    "        'json', data_files = DATASET_PATH + DATASET) if USE_LOCAL_DATASET else load_dataset(DATASET)\n",
    "    \n",
    "    # loading samples from HF dataset for fast testing\n",
    "    if( not USE_LOCAL_DATASET and DATA_SAMPLING ):\n",
    "        my_ds = DatasetDict({\n",
    "            \"test\": select_slice(my_ds, \"test\"),\n",
    "            \"train\": select_slice(my_ds, \"train\"),\n",
    "            \"validation\": select_slice(my_ds, \"validation\")\n",
    "            })\n",
    "    return my_ds\n",
    "\n",
    "prefix_input_batch = lambda prefix, items: [prefix + q for q in items]\n",
    "\n",
    "extract_from_nested_batch = lambda nested_col, items: [q[nested_col] for q in items]\n",
    "\n",
    "tokenize = lambda data: tokenizer(data, max_length=max_input_length, truncation=True) #, return_tensors=\"pt\")\n",
    "\n",
    "def tokenize_batch(record_batch):\n",
    "    \n",
    "    # prefix each question in the batch\n",
    "    input_questions = prefix_input_batch(TASK_PREFIX, record_batch[INPUT_FIELD])\n",
    "    \n",
    "    # WikiSQL has nested field showing the expected output\n",
    "    expected_results = extract_from_nested_batch(NESTED_RESULT_FIELD, record_batch[RESULT_FIELD])\n",
    "    \n",
    "    # for self created dataset samples, use flat structure:\n",
    "    #expected_results = record[RESULT_FIELD]\n",
    "    \n",
    "    # encoding the questions from dataset to get input_ids & attention_masks\n",
    "    model_inputs = tokenize(input_questions)\n",
    "    \n",
    "    # encoding the expected results from dataset to compare against the predicted\n",
    "    with tokenizer.as_target_tokenizer():\n",
    "        labels = tokenize(expected_results)\n",
    "\n",
    "    # merging labels, dictionary will contain:\n",
    "    # input_ids & attention_masks of the inputs\n",
    "    # input_ids (labels) & attention_masks (decoder_attention_masks) of the expected results\n",
    "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
    "    model_inputs[\"decoder_attention_mask\"] = labels[\"attention_mask\"]\n",
    "\n",
    "    return model_inputs\n",
    "\n",
    "def tokenize_dataset(dataset_to_tokenize):\n",
    "\n",
    "    if USE_TKN_CACHE and os.path.exists(ENCODED_DATA_PATH):\n",
    "        print(\"Loading tokenized dataset from cache\")\n",
    "        tokenized_dataset = datasets.load_from_disk(ENCODED_DATA_PATH)\n",
    "    else:\n",
    "        print(\"Tokenizing dataset ... it will take some time\")\n",
    "        tokenized_dataset = dataset_to_tokenize.map(tokenize_batch, batched=True)\n",
    "\n",
    "        print(\"Saving tokenizing dataset ... \", ENCODED_DATA_PATH)\n",
    "        tokenized_dataset.save_to_disk(ENCODED_DATA_PATH)\n",
    "\n",
    "    # NOTE: for some reason, the tokenized dataset is contaiing \n",
    "    # all fields from input instead of just 4 tokenized fields\n",
    "    # ['phase', 'question', 'table', 'sql', 'input_ids', \n",
    "    #   'attention_mask', 'labels', 'decoder_attention_mask']\n",
    "    return tokenized_dataset\n",
    "\n",
    "# ==============================================================#\n",
    "# ======================== MAIN PROGRAM ========================#\n",
    "\n",
    "full_dataset = load_data()\n",
    "print(\"-\"*20, \"Dataset Loaded \")\n",
    "\n",
    "tokenized_dataset = tokenize_dataset(full_dataset)\n",
    "print(\"-\"*20, \"Dataset Tokenized \")\n",
    "\n",
    "#using trainer to train, preferred as the code is cleaner\n",
    "training_args = get_training_args()\n",
    "print(\"-\"*20, \"Training Args Ready \")\n",
    "\n",
    "metric = evaluate.load(\"rouge\")\n",
    "print(\"-\"*20, \"Rouge Metric Ready \")\n",
    "\n",
    "data_collator = DataCollatorForSeq2Seq(tokenizer)\n",
    "\n",
    "trainer = get_trainer(tokenized_dataset, training_args)\n",
    "print(\"Training Logs at : \", LOG_DIR)\n",
    "\n",
    "# Start tensorboard\n",
    "%load_ext tensorboard\n",
    "%tensorboard --logdir LOG_DIR+'/runs'\n",
    "\n",
    "trainer.train()\n",
    "trainer.save_model()\n",
    "    "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.9.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}