train_t5_text_to_sql.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ===================== Traing base T5 ==================== # \n",
    "import sys\n",
    "import random\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "import tqdm\n",
    "from datasets import load_dataset, DatasetDict\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import T5ForConditionalGeneration, T5Tokenizer\n",
    "from transformers import AdamW, get_linear_schedule_with_warmup\n",
    "\n",
    "MODEL_BASE_DIR      = \"/Users/sree/.cache/huggingface/hub\"\n",
    "IN_MODEL_NAME       = \"t5-base\"\n",
    "OUT_MODEL_NAME      = \"my--t5-base-finetuned-text-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",
    "TASK_PREFIX         = \"Translate English to SQL: \"\n",
    "DATA_PATH           = \"./data/my_flat_sql_data.json\"\n",
    "\n",
    "def set_seed(seed):\n",
    "  random.seed(seed)\n",
    "  np.random.seed(seed)\n",
    "  torch.manual_seed(seed)\n",
    "set_seed(42)\n",
    "\n",
    "def initialize_optimizer(model):\n",
    "    # initialize optimizer\n",
    "    no_decay = [\"bias\", \"LayerNorm.weight\"]\n",
    "    grouped_parameters = [\n",
    "        {\"params\": [p for n, p in 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 model.named_parameters() if any(nd in n for nd in no_decay)],\n",
    "            \"weight_decay\": 0.0,},]\n",
    "    return AdamW(grouped_parameters, lr=3e-4, eps=1e-8)\n",
    "\n",
    "\n",
    "tokenize = lambda tokenizer, text: tokenizer.encode_plus(\n",
    "    text, max_length=96, padding=True, truncation=True, return_tensors=\"pt\")\n",
    "\n",
    "def process_data(record):\n",
    "    \n",
    "    input = TASK_PREFIX + record['question'] #+ \" </s>\"\n",
    "    expected_ouput = record['sql'] #+ \" </s>\"\n",
    "\n",
    "    tokenized_input = tokenize(tokenizer, input)\n",
    "    input_ids  = tokenized_input[\"input_ids\"]\n",
    "    attention_mask = tokenized_input[\"attention_mask\"]\n",
    "    \n",
    "    # model_max_length=512,\n",
    "    tokenized_output = tokenize(tokenizer, expected_ouput)\n",
    "    lm_labels = tokenized_output[\"input_ids\"]\n",
    "    decoder_attention_mask=  tokenized_output[\"attention_mask\"]\n",
    "\n",
    "    output = model(\n",
    "        input_ids = input_ids, attention_mask = attention_mask, \n",
    "        labels = lm_labels, decoder_attention_mask = decoder_attention_mask)\n",
    "    \n",
    "    loss = output[0]\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "# initialize model, tokenizer & optimizer\n",
    "tokenizer = T5Tokenizer.from_pretrained(IN_MODEL_NAME)\n",
    "model = T5ForConditionalGeneration.from_pretrained(IN_MODEL_NAME)\n",
    "optimizer = initialize_optimizer(model)\n",
    "\n",
    "my_ds = load_dataset('json', data_files = DATA_PATH)\n",
    "print(my_ds)\n",
    "\n",
    "epochs = 2\n",
    "model.train(mode=True)\n",
    "for epoch in range(epochs):\n",
    "    print (\"epoch \",epoch)\n",
    "    my_ds.map(process_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_sent = TASK_PREFIX['translation'] + \"how many customers are in Florida?\" #+  \"</s>\"\n",
    "test_tokenized = tokenizer.encode_plus(test_sent, return_tensors=\"pt\")\n",
    "\n",
    "test_input_ids  = test_tokenized[\"input_ids\"]\n",
    "test_attention_mask = test_tokenized[\"attention_mask\"]\n",
    "\n",
    "model.eval()\n",
    "beam_outputs = model.generate(\n",
    "    input_ids=test_input_ids,\n",
    "    attention_mask=test_attention_mask,\n",
    "    temperature = .95,\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=3\n",
    ")\n",
    "\n",
    "for beam_output in beam_outputs:\n",
    "    output = tokenizer.decode(beam_output, skip_special_tokens=True,clean_up_tokenization_spaces=True)\n",
    "    print(output)"
   ]
  }
 ],
 "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
}