t5_trainer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ===================== Module Imports ==================== # \n",
    "import torch.optim as optim\n",
    "#from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "from transformers import T5ForConditionalGeneration, T5Tokenizer\n",
    "from transformers import AdamW, get_linear_schedule_with_warmup\n",
    "\n",
    "#================== Global Constants ==================#\n",
    "TASK_PREFIX     = {\n",
    "                    \"translation\":\"Translate English to SQL: \",\n",
    "                    \"metadata\":\"Find data tables in question: \"\n",
    "                    }\n",
    "TASK_KEY        = {\n",
    "                    \"translation\":\"sql\",\n",
    "                    \"metadata\":\"tables\"\n",
    "                    }\n",
    "TASK_HINT       = {\n",
    "                    \"translation\":\". Use following data tables - \",\n",
    "                    \"metadata\":\"\"\n",
    "                    }\n",
    "MODEL_CLASSES   = {\n",
    "                    \"t5-small\": (T5ForConditionalGeneration, T5Tokenizer),\n",
    "                    \"t5-base\": (T5ForConditionalGeneration, T5Tokenizer),\n",
    "                    }\n",
    "OPTIM_CLASSES   = {\n",
    "                    \"sgd\": optim.SGD,\n",
    "                    \"adam\": optim.Adam,\n",
    "                    }\n",
    "MODEL_BASE_DIR  = \"/Users/sree/.cache/huggingface/hub\"\n",
    "#LOG_DIR        = OUTPUT_DIR + \"/logs\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset, DatasetDict\n",
    "\n",
    "class MyDataLoader:\n",
    "    \n",
    "    data_mode       = \"record\" #record|batch\n",
    "    task_mode       = \"translation\" #translation|metadata\n",
    "    given_input     = None\n",
    "    expected_output = None\n",
    "    skip_record     = None\n",
    "\n",
    "    def __init__(self, data_mode):\n",
    "        self.set_data_mode(data_mode)\n",
    "        self.__init_data_extractors__()\n",
    "\n",
    "    # train modes - record | batch\n",
    "    def set_data_mode(self, data_mode):\n",
    "        self.data_mode = data_mode\n",
    "\n",
    "    # task modes - translation | metadata\n",
    "    def set_task_mode(self, task_mode):\n",
    "        self.task_mode = task_mode\n",
    "\n",
    "    def load_data(self):\n",
    "        data_path = \"./data/my_flat_sql_data_meta.json\"\n",
    "        train_ds = load_dataset('json', data_files = data_path)\n",
    "        #print(\"Keys: \",train_ds.keys())\n",
    "        #print(\"Items: \",train_ds.items())\n",
    "\n",
    "        # will have to load eval ds separate\n",
    "        return train_ds, train_ds\n",
    "    \n",
    "    def __init_data_extractors__(self):\n",
    "        \n",
    "        if self.data_mode == 'batch':\n",
    "            input_with_hint = lambda task_mode, tuple: TASK_PREFIX[task_mode] + tuple[0]+ TASK_HINT[task_mode] + tuple[1]\n",
    "\n",
    "            self.given_input = lambda task_mode, batch: [input_with_hint(task_mode, row) for row in zip(batch['question'], batch['tables']) if all(row)]\n",
    "            self.expected_output = lambda task_mode, batch: [row for row in batch[TASK_KEY[task_mode]] if row]\n",
    "            self.skip_record = None\n",
    "\n",
    "        if self.data_mode == 'record':\n",
    "            extract_input = lambda task_mode, record: TASK_PREFIX[task_mode] + record['question']\n",
    "            exract_hint = lambda task_mode, record: TASK_HINT[task_mode] + record['tables']\n",
    "\n",
    "            self.given_input = lambda task_mode, record: extract_input(task_mode, record) + exract_hint(task_mode, record)\n",
    "            self.expected_output = lambda task_mode, record: record[ TASK_KEY[task_mode] ]\n",
    "            self.skip_record = lambda record: record['comment']\n",
    "    \n",
    "    def exract_input(self, data):\n",
    "        return self.given_input(self.task_mode, data)\n",
    "\n",
    "    def extract_expected_output(self, data):\n",
    "        return self.expected_output(self.task_mode, data)\n",
    "    \n",
    "    def is_skip_record(self, data):\n",
    "        return self.data_mode == 'record' and self.skip_record(data)\n",
    "    \n",
    "    def is_read_bacthed(self):\n",
    "        return self.data_mode == 'batch'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, time\n",
    "#import pandas as pd\n",
    "import numpy as np\n",
    "import random\n",
    "import torch\n",
    "import concurrent\n",
    "\n",
    "import concurrent.futures\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "class MyT5Trainer:\n",
    "    \n",
    "    #================== Contained Objects ==================#\n",
    "    # loader, tokenizer, model, optimizer\n",
    "    \n",
    "    #================== Model Info Attributes ==================#\n",
    "    #out_model_name     = None\n",
    "    #output_dir         = None\n",
    "\n",
    "    #================== Trainer State Control attributes ==================#\n",
    "    threads             = 1\n",
    "    epochs              = 1\n",
    "    loss_log            = []\n",
    "\n",
    "    #================== T5 Model Hyperparameters ==================#\n",
    "    #\n",
    "    \n",
    "    #================== Adam Optimizer Hyperparameters ==================#\n",
    "    adam_lr             = 3e-4\n",
    "    adam_eps            = 1e-8\n",
    "    \n",
    "\n",
    "    def __init__(self, model_name, loader):\n",
    "        \n",
    "        print(\"Initializing Trainer\")\n",
    "        self.set_loader(loader)\n",
    "        \n",
    "        # will be enabled in future when saving model checkpointing\n",
    "        #self.out_model_name = f\"my--{model_name}--finetuned-text-to-SQL\"\n",
    "        #self.output_dir = f\"{MODEL_BASE_DIR}/models--{self.out_model_name}\"\n",
    "\n",
    "        try:\n",
    "            model_class, tokenizer_class = MODEL_CLASSES[ model_name ]\n",
    "        except KeyError:\n",
    "            raise KeyError(\"the model {} you specified is not supported. You are welcome to add it and open a PR :)\")\n",
    "\n",
    "        # model_max_length=512,\n",
    "        self.tokenizer = tokenizer_class.from_pretrained( model_name )\n",
    "        self.model = model_class.from_pretrained(model_name, pad_token_id=self.tokenizer.eos_token_id )\n",
    "\n",
    "        no_decay = [\"bias\", \"LayerNorm.weight\"]\n",
    "        grouped_parameters = [\n",
    "            {\"params\": [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay)],\n",
    "                \"weight_decay\": 0.0,},\n",
    "            {\"params\": [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay)],\n",
    "                \"weight_decay\": 0.0,},]\n",
    "        self.optimizer = optim.Adam(grouped_parameters, lr=self.adam_lr, eps=self.adam_eps)\n",
    "        #self.optimizer = optim.Adam(self.model.parameters*(), lr=self.ADAM_LR, eps=self.ADAM_EPS, weight_decay=0.0)\n",
    "\n",
    "    \n",
    "    #======================== Getters & Setters =========================#\n",
    "    #====================================================================#\n",
    "    \n",
    "    def set_seed(self, seed):\n",
    "        random.seed(seed)\n",
    "        np.random.seed(seed)\n",
    "        torch.manual_seed(seed)\n",
    "\n",
    "    def set_loader(self, loader):\n",
    "        self.loader = loader\n",
    "\n",
    "    def set_threads(self, max_threads):\n",
    "        self.threads = max_threads\n",
    "    \n",
    "    def set_epochs(self, epochs):\n",
    "        self.epochs = epochs\n",
    "\n",
    "    def reset_accuracy_log(self):\n",
    "        self.loss_log = []\n",
    "        self.correctness = 0\n",
    "\n",
    "    #==================== Functional Methods ===================#\n",
    "    #===========================================================#\n",
    "        \n",
    "    __tokenize__ = lambda self, text: self.tokenizer.encode_plus(\n",
    "        text, max_length=96, padding=True, truncation=True, return_tensors=\"pt\")\n",
    "\n",
    "    def __encode_data__(self, data):\n",
    "\n",
    "        # preprocessing data - extracting from input data, prefixing & cleaning up\n",
    "        input = self.loader.exract_input(data) #self.task_mode,\n",
    "        expected_output = self.loader.extract_expected_output(data)\n",
    "\n",
    "        # tokenizing input & exptected output data\n",
    "        tokenized_input = self.__tokenize__(input)\n",
    "        tokenized_output = self.__tokenize__(expected_output)\n",
    "\n",
    "        return (tokenized_input[\"input_ids\"], tokenized_input[\"attention_mask\"], \n",
    "        tokenized_output[\"input_ids\"], tokenized_output[\"attention_mask\"])\n",
    "    \n",
    "    \n",
    "    def __process_train_data__(self, data):\n",
    "\n",
    "        #if training in single record mode, check for empty or comment records\n",
    "        if self.loader.is_skip_record(data):\n",
    "            return\n",
    "\n",
    "        # parse, cleanse & tokenize input data record or bacth records (based on train_mode)\n",
    "        input_ids, attention_mask, lm_labels, decoder_attention_mask = self.__encode_data__(data)\n",
    "        \n",
    "        # forward pass - predict\n",
    "        output  = self.model(\n",
    "            input_ids = input_ids, attention_mask = attention_mask, \n",
    "            labels = lm_labels, decoder_attention_mask = decoder_attention_mask)\n",
    "        \n",
    "        # ==== compute loss to optimize step ==== #\n",
    "\n",
    "        # foward pass - compute loss\n",
    "        loss = output[0]\n",
    "        \n",
    "        #record the loss for plotting\n",
    "        self.loss_log.append(loss.item())\n",
    "\n",
    "        #zero all gradients before tha backward pass\n",
    "        self.optimizer.zero_grad()\n",
    "\n",
    "        # backward pass\n",
    "        loss.backward()\n",
    "\n",
    "        self.optimizer.step()\n",
    "        \n",
    "\n",
    "    def __process_eval_data__(self, data):\n",
    "        \n",
    "        #if training in single record mode, check for empty or comment records\n",
    "        if self.loader.is_skip_record(data):\n",
    "            return\n",
    "        \n",
    "        # parse, cleanse & tokenize input data record or bacth records (based on train_mode)\n",
    "        input_ids, attention_mask, lm_labels, decoder_attention_mask = self.__encode_data__(data)\n",
    "\n",
    "        # forward pass - predict\n",
    "        output  = self.model(\n",
    "            input_ids = input_ids, attention_mask = attention_mask, \n",
    "            labels = lm_labels, decoder_attention_mask = decoder_attention_mask)\n",
    "\n",
    "        # ==== compute error  ==== #\n",
    "\n",
    "        print(\"Next iteration - EVAL\")\n",
    "        # Get the index of the max log-probability.\n",
    "        pred = output.argmax(dim=1, keepdim=True)\n",
    "        #self.correctness += pred.eq(lm_labels.view_as(pred)).sum().item()\n",
    "\n",
    "\n",
    "\n",
    "    def train_model(self, train_ds, eval_ds):\n",
    "        batched=self.loader.is_read_bacthed()\n",
    "        trainer.reset_accuracy_log()\n",
    "\n",
    "        for epoch in range(self.epochs):\n",
    "            print(\"epoch \",epoch)\n",
    "\n",
    "            # Model training - threaded\n",
    "            self.model.train(mode=True)\n",
    "            with concurrent.futures.ThreadPoolExecutor(max_workers=self.threads) as executor:\n",
    "                executor.submit( train_ds.map, self.__process_train_data__, batched=batched )\n",
    "            \n",
    "            # non-threaded training\n",
    "            #train_ds.map(self.__process_train_data__, batched=batched)\n",
    "\n",
    "            # Model validation\n",
    "            self.model.eval(mode=True)\n",
    "            # with torch.no_grad():\n",
    "            #    eval_ds.map(self.__process_eval_data__, batched=batched)\n",
    "            # accuracy = correct / min(len(valid_loader.dataset), N_VALID_EXAMPLES)\n",
    "\n",
    "    def plot_loss(self):\n",
    "        plt.plot(self.loss_log, label = \"Stochastic Gradient Descent\")\n",
    "        #plt.plot(loss_Adam,label = \"Adam Optimizer\")\n",
    "        plt.xlabel('epoch')\n",
    "        plt.ylabel('Cost/ total loss')\n",
    "        plt.legend()\n",
    "        plt.show()\n",
    "\n",
    "\n",
    "#================== Main Program ==================#\n",
    "#==================================================#\n",
    "\n",
    "# record = one record at a time training \n",
    "# batch = training 1K batched records at a time - not yeilding expected results - needs investigation\n",
    "data_mode = \"record\" #record|batch\n",
    "data_loader = MyDataLoader(data_mode)\n",
    "\n",
    "# load training & evaluation/validation datasets\n",
    "train_ds, eval_ds = data_loader.load_data()\n",
    "\n",
    "# initialize tokenizer/encoder, model & optimizer\n",
    "trainer = MyT5Trainer(\"t5-base\", data_loader) # t5-base | t5-small\n",
    "trainer.set_threads(1)\n",
    "trainer.set_epochs(1)\n",
    "trainer.set_seed(42)\n",
    "\n",
    "task_modes = ['translation'] #'metadata'\n",
    "for task_mode in task_modes:\n",
    "    \n",
    "    data_loader.set_task_mode(task_mode)\n",
    "    trainer.train_model(train_ds, eval_ds)\n",
    "\n",
    "    trainer.plot_loss()\n",
    "    print(trainer.loss_log)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import evaluate\n",
    "\n",
    "TASK_MODE = 'translation'\n",
    "input_text = TASK_PREFIX[TASK_MODE] + \"Which teams played in 2022?\" #+  \"</s>\"\n",
    "hint_text = TASK_HINT[TASK_MODE] + \"game_stats\"\n",
    "expected_output = \"SELECT team_name FROM game_stats WHERE DATE_PART('YEAR', pay_date)= 2022\"\n",
    "print(input_text+hint_text)\n",
    "\n",
    "test_tokenized = trainer.tokenizer.encode_plus(input_text+hint_text, return_tensors=\"pt\")\n",
    "test_input_ids  = test_tokenized[\"input_ids\"]\n",
    "test_attention_mask = test_tokenized[\"attention_mask\"]\n",
    "\n",
    "output_tokenized = trainer.tokenizer.encode_plus(expected_output, return_tensors=\"pt\")\n",
    "labels = output_tokenized[\"input_ids\"]\n",
    "\n",
    "trainer.model.eval()\n",
    "outputs = trainer.model.generate(\n",
    "    input_ids=test_input_ids,\n",
    "    attention_mask=test_attention_mask,\n",
    "    temperature = .96,\n",
    "    max_new_tokens=64,\n",
    "    #max_length=64,\n",
    "    \n",
    "    #early_stopping=True,\n",
    "    #num_beams=10,\n",
    "    #num_return_sequences=1, #3\n",
    "    #no_repeat_ngram_size=2 #2\n",
    "    \n",
    "    # ----- Beam Search w/ return sequences -----#\n",
    "    early_stopping=True,\n",
    "    num_beams=10,\n",
    "    no_repeat_ngram_size=2,\n",
    "    num_return_sequences=5, #num_return_sequences<=num_beams\n",
    "\n",
    "    # ----- Top P & Top K sampling -----# ANALYSIS - much faster than beam search\n",
    "    #do_sample=True,\n",
    "    #top_k=5, \n",
    "    #top_p=3,\n",
    "    #num_return_sequences=1\n",
    "\n",
    "    # --- Greedy search ----#\n",
    "    \n",
    ")\n",
    "\n",
    "for beam_output in outputs:\n",
    "    \n",
    "    output = trainer.tokenizer.decode(beam_output, skip_special_tokens=True, clean_up_tokenization_spaces=True)\n",
    "    print(\"BEAM Output: \", beam_output)\n",
    "    print(\"DECODED Output\", output)\n",
    "    print(\"Expected: \",expected_output)\n",
    "    print(\"LABELS: \",labels)\n",
    "\n",
    "    decoded_preds = output\n",
    "    decoded_labels = expected_output\n",
    "    #trainer.tokenizer.decode(output_tokenized, skip_special_tokens=True, clean_up_tokenization_spaces=True)\n",
    "    metric = evaluate.load(\"rouge\")\n",
    "    result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)\n",
    "    print(\"RESULT: \", result)\n",
    "\n",
    "    # Get the index of the max log-probability.\n",
    "    #pred = beam_output.argmax(dim=0, keepdim=True)\n",
    "    pred = torch.argmax(beam_output, dim=0, keepdim=True)\n",
    "    print(\"PRED Index: \",pred)\n",
    "    print(\"PRED: \", beam_output[pred].item())\n",
    "\n",
    "    max = torch.max(beam_output)\n",
    "    print(\"Max: \",max.item())\n",
    "\n",
    "    print(\"EQ: \", labels.view_as(beam_output).eq(beam_output))\n",
    "\n",
    "    #print(\"EQ: \", labels.eq(pred).sum() )\n",
    "    #print(\"EQ: \", labels.eq(pred).sum().item() )\n",
    "\n",
    "    #print(\"VIEW SUM: \", labels.sum())\n",
    "    \n",
    "    #print(\"PRED SHAPE: \", pred.shape)\n",
    "    print(\"LABEL SHAPE: \", labels.shape)\n",
    "    \n",
    "    #print(\"VIEW AS: \",pred.view_as(labels))\n",
    "    #print(\"VIEW AS: \",labels.view_as(pred))\n",
    "\n",
    "    #correct = pred.eq(labels.view_as(pred)).sum()\n",
    "    correct = pred.eq(labels).sum().item()\n",
    "    print(\"CORRECT: \", correct)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\"\"\"\n",
    "t1 = torch.Tensor(1,2)\n",
    "#t4 = torch.Tensor([],[])\n",
    "\n",
    "print(t1.shape)\n",
    "print(t3.size(0))\n",
    "print(t3.size(1))\n",
    "#print(\"VIEW:\", t3.view(-1,-1))\n",
    "eqt = t3.eq(t2)\n",
    "print(\"EQ: \", eqt )\n",
    "print(\"SUM: \", eqt.sum() )\n",
    "print(\"ITEM: \", eqt.sum().item() )\n",
    "\"\"\"\n",
    "\n",
    "t2 = torch.Tensor([20])\n",
    "t3 = torch.Tensor([[2,20,3,4,5]])\n",
    "print(t2.shape)\n",
    "print(t3.shape)\n",
    "eqt = t3.eq(t2)\n",
    "print(\"EQ: \",eqt)\n",
    "print(\"VIEW: \", t3.view_as(t2))\n",
    "\n",
    "beam = torch.Tensor([0, 12558, 10, 4073, 2323, 1944, 16, 460, 2884, 58,\n",
    "            5, 2048, 826, 331, 5056, 3, 18, 467, 834, 8547,\n",
    "            7, 5, 1])\n",
    "labels = torch.Tensor([[3, 23143, 14196, 372, 834, 4350, 21680, 467, 834, 8547,\n",
    "             7, 549, 17444, 427, 309, 6048, 834, 19846, 599, 31,\n",
    "           476, 19356, 31, 6, 726, 834, 5522, 61, 2423, 460,\n",
    "          2884, 1]])\n",
    "print(beam.shape)\n",
    "print(labels.shape)\n",
    "print(\"EQ: \",labels.eq(t3))\n",
    "\n",
    "\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
}