src/data_process.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Process train and test set for Qulac and ClariQ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/raid/zhenduow/miniconda3/envs/ustest/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import csv\n",
    "import json \n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import random\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch as T\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import re\n",
    "from autocorrect import Speller\n",
    "spell = Speller(lang='en')\n",
    "\n",
    "import torch as T\n",
    "import pandas as pd\n",
    "import tqdm\n",
    "from transformers import RobertaTokenizer, RobertaForSequenceClassification\n",
    "from transformers import (\n",
    "    T5ForConditionalGeneration,\n",
    "    T5Tokenizer,\n",
    ")\n",
    "\n",
    "idk_list = [\n",
    "    'i dont know',\n",
    "    'i do not know',\n",
    "    'im not sure',\n",
    "    'i am not sure',\n",
    "    'unsure',\n",
    "    'possibly',\n",
    "    'this is not related to my search'\n",
    "    ]\n",
    "\n",
    "negation_list = [\n",
    "    'no',\n",
    "    'not',\n",
    "    'none',\n",
    "    'isnt',\n",
    "    'isn\\'t',\n",
    "    'dont',\n",
    "    'don\\'t',\n",
    "    ]\n",
    "\n",
    "\n",
    "auxiliary_verb_list = [\n",
    "    'did',\n",
    "    'do',\n",
    "    'does',\n",
    "    'it',\n",
    "    'are',\n",
    "    'was',\n",
    "    'were',\n",
    "    'have',\n",
    "    'has',\n",
    "    'can',\n",
    "    'could',\n",
    "    'will',\n",
    "    'would',\n",
    "]\n",
    "\n",
    "def normalize_text(text):\n",
    "    \"\"\"Lowercase and remove quotes from a TensorFlow string.\"\"\"\n",
    "    text = text.lower()\n",
    "    text = re.sub(\"'(.*)'\", r\"\\1\", text)\n",
    "    if text[-1] == '?':\n",
    "        text = text[:-1] + '.'\n",
    "    return text\n",
    "\n",
    "def type_answer(answer):\n",
    "    if answer in idk_list:\n",
    "        return 'idk'\n",
    "    elif 'yes' in answer.split()[:3]:\n",
    "        return 'yes'\n",
    "    elif any([w in answer.split()[:3] for w in negation_list]):\n",
    "        return 'no'\n",
    "    else:\n",
    "        return 'open'\n",
    "\n",
    "def u2i(text):\n",
    "    text = re.sub('are you', 'am i', text)\n",
    "    text = re.sub('you', 'i', text)\n",
    "    return text\n",
    "\n",
    "qulac_data_files = [\n",
    "    '../data/qulac/qulac.train.json',\n",
    "    '../data/qulac/qulac.test.json',\n",
    "    '../data/qulac/qulac.valid.json',\n",
    "]\n",
    "\n",
    "qulac_output_names = [\n",
    "    '../data/processed/qulac_train.csv',\n",
    "    '../data/processed/qulac_dev.csv',\n",
    "    '../data/processed/qulac_test.csv',\n",
    "]\n",
    "\n",
    "\n",
    "if not os.path.exists('../data/processed'):\n",
    "    os.makedirs('../data/processed')\n",
    "\n",
    "\n",
    "for data_file, output_name in zip(qulac_data_files, qulac_output_names):\n",
    "    df = pd.read_json(data_file)\n",
    "    df.replace(['', \"NaN\", 'NaT'], np.nan, inplace = True)\n",
    "    df.dropna(subset=['question', 'facet_desc' , 'answer'], how='any', inplace=True)\n",
    "    df = df[['topic','facet_desc','question','answer']].copy(deep=True)\n",
    "\n",
    "    for iter, row in df.iterrows():\n",
    "        query = normalize_text(df.at[iter, 'topic'])\n",
    "        facet = normalize_text(df.at[iter, 'facet_desc'])\n",
    "        question = normalize_text(df.at[iter, 'question'])\n",
    "        answer = normalize_text(df.at[iter, 'answer'])\n",
    "        df.at[iter, 't5-question'] = facet + ' . ' + query + ' . ' + question\n",
    "        df.at[iter, 'unifiedqa-question'] = u2i(question) + ' ? \\\\n ' + 'i am looking for ' + facet \n",
    "        df.at[iter, 'answer'] = normalize_text(answer)\n",
    "        df.at[iter, 'answer-type'] = type_answer(answer)\n",
    "        df.at[iter, 'answer-len'] = len(answer.split())\n",
    "\n",
    "    df.to_csv(output_name, index=False)\n",
    "\n",
    "\n",
    "clariq_data_files = [\n",
    "    '../data/clariq/clariq_train.tsv',\n",
    "    '../data/clariq/clariq_dev.tsv',\n",
    "    '../data/clariq/clariq_test.tsv'\n",
    "]\n",
    "\n",
    "clariq_output_names = [\n",
    "    '../data/processed/clariq_train.csv',\n",
    "    '../data/processed/clariq_dev.csv',\n",
    "    '../data/processed/clariq_test.csv',\n",
    "]\n",
    "\n",
    "for data_file, output_name in zip(clariq_data_files, clariq_output_names):\n",
    "    df = pd.read_csv(data_file, delimiter='\\t')\n",
    "    df.replace(['', \"NaN\", 'NaT'], np.nan, inplace = True)\n",
    "    df.dropna(subset=['question', 'facet_desc' , 'answer'], how='any', inplace=True)\n",
    "\n",
    "    for iter, row in df.iterrows():\n",
    "        query = normalize_text(df.at[iter, 'initial_request'])\n",
    "        facet = normalize_text(df.at[iter, 'facet_desc'])\n",
    "        question = normalize_text(df.at[iter, 'question'])\n",
    "        answer = normalize_text(df.at[iter, 'answer'])\n",
    "        df.at[iter, 't5-question'] = facet + ' . ' + query + ' . ' + question\n",
    "        df.at[iter, 'unifiedqa-question'] = u2i(question) + ' ? \\\\n ' + 'i am looking for ' + facet \n",
    "        df.at[iter, 'answer'] = normalize_text(answer)\n",
    "        df.at[iter, 'answer-type'] = type_answer(answer)\n",
    "        df.at[iter, 'answer-len'] = len(answer.split())\n",
    "\n",
    "    df.to_csv(output_name, index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Train RoBERTa on Qulac"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import datasets\n",
    "from transformers import RobertaTokenizerFast, RobertaForSequenceClassification,Trainer, TrainingArguments\n",
    "import torch.nn as nn\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import numpy as np\n",
    "from sklearn.metrics import accuracy_score, precision_recall_fscore_support\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "\n",
    "model = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=4).cuda()\n",
    "tokenizer = RobertaTokenizerFast.from_pretrained('roberta-base', max_length = 512)\n",
    "\n",
    "def tokenization(batched_text):\n",
    "    return tokenizer(batched_text['text'], padding = True, truncation=True)\n",
    "\n",
    "def to_categorical(batch):\n",
    "    \"\"\" 1-hot encodes a tensor \"\"\"\n",
    "    batch['labels'] = batch['labels'].unsqueeze(0)\n",
    "    return batch\n",
    "\n",
    "# qulac experiments\n",
    "print(\"training roberta on qulac dataset\")\n",
    "data_files = {\"train\": \"qulac_train.csv\", \"dev\": \"qulac_dev.csv\", \"test\": \"qulac_test.csv\"}\n",
    "dataset = datasets.load_dataset(\"../data/processed\", data_files=data_files)\n",
    "train_data, dev_data, test_data = dataset['train'],dataset['dev'],dataset['test']\n",
    "\n",
    "train_data = train_data.remove_columns(['topic', 'facet_desc', 'question', 'answer', 't5-question', 'answer-len'])\n",
    "train_data = train_data.rename_column('unifiedqa-question', 'text')\n",
    "train_data = train_data.rename_column('answer-type', 'labels')\n",
    "train_data = train_data.class_encode_column('labels')\n",
    "\n",
    "test_data = test_data.remove_columns(['topic', 'facet_desc', 'question', 'answer', 't5-question',  'answer-len'])\n",
    "test_data = test_data.rename_column('unifiedqa-question', 'text')\n",
    "test_data = test_data.rename_column('answer-type', 'labels')\n",
    "test_data = test_data.class_encode_column('labels')\n",
    "\n",
    "train_data = train_data.map(tokenization, batched = True, batch_size = len(train_data))\n",
    "test_data = test_data.map(tokenization, batched = True, batch_size = len(test_data))\n",
    "\n",
    "train_data.set_format('torch', columns=['input_ids', 'attention_mask', 'labels'])\n",
    "test_data.set_format('torch', columns=['input_ids', 'attention_mask', 'labels'])\n",
    "\n",
    "train_data = train_data.map(to_categorical, batched = False, batch_size = len(train_data))\n",
    "test_data = test_data.map(to_categorical, batched = False, batch_size = len(test_data))\n",
    "\n",
    "# define accuracy metrics\n",
    "def compute_metrics(pred):\n",
    "    labels = pred.label_ids\n",
    "    preds = pred.predictions.argmax(-1)\n",
    "    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='macro')\n",
    "    acc = accuracy_score(labels, preds)\n",
    "    return {\n",
    "        'accuracy': acc,\n",
    "        'f1': f1,\n",
    "        'precision': precision,\n",
    "        'recall': recall\n",
    "    }\n",
    "\n",
    "# define the training arguments\n",
    "training_args = TrainingArguments(\n",
    "    output_dir = './ckpt/roberta-qulac/',\n",
    "    num_train_epochs=100,\n",
    "    per_device_train_batch_size = 32,\n",
    "    gradient_accumulation_steps = 16,    \n",
    "    per_device_eval_batch_size= 32,\n",
    "    evaluation_strategy = \"epoch\",\n",
    "    save_strategy = 'epoch',\n",
    "    disable_tqdm = False, \n",
    "    load_best_model_at_end=True,\n",
    "    warmup_steps=500,\n",
    "    weight_decay=0.01,\n",
    "    logging_steps = 8,\n",
    "    fp16 = True,\n",
    "    logging_dir='./ckpt/roberta-qulac/log',\n",
    "    dataloader_num_workers = 8,\n",
    "    run_name = 'roberta-classification',\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    compute_metrics=compute_metrics,\n",
    "    train_dataset=train_data,\n",
    "    eval_dataset=test_data\n",
    ")\n",
    "\n",
    "trainer.train()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Train RoberTa on ClariQ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# clariq experiments\n",
    "print(\"training roberta on clariq dataset\")\n",
    "data_files = {\"train\": \"clariq_train.csv\", \"dev\": \"clariq_dev.csv\", \"test\": \"clariq_dev.csv\"}\n",
    "dataset = datasets.load_dataset(\"../data/processed\", data_files=data_files)\n",
    "train_data, dev_data, test_data = dataset['train'],dataset['dev'],dataset['test']\n",
    "\n",
    "\n",
    "train_data = train_data.remove_columns(['topic_id', 'initial_request' , 'topic_desc', 'clarification_need', 'facet_id' ,'facet_desc', 'question_id', 'question', 'answer', 't5-question', 'answer-len'])\n",
    "train_data = train_data.rename_column('unifiedqa-question', 'text')\n",
    "train_data = train_data.rename_column('answer-type', 'labels')\n",
    "train_data = train_data.class_encode_column('labels')\n",
    "\n",
    "test_data = test_data.remove_columns(['topic_id', 'initial_request' , 'topic_desc', 'clarification_need', 'facet_id' ,'facet_desc', 'question_id', 'question', 'answer', 't5-question', 'answer-len'])\n",
    "test_data = test_data.rename_column('unifiedqa-question', 'text')\n",
    "test_data = test_data.rename_column('answer-type', 'labels')\n",
    "test_data = test_data.class_encode_column('labels')\n",
    "\n",
    "train_data = train_data.map(tokenization, batched = True, batch_size = len(train_data))\n",
    "test_data = test_data.map(tokenization, batched = True, batch_size = len(test_data))\n",
    "\n",
    "train_data.set_format('torch', columns=['input_ids', 'attention_mask', 'labels'])\n",
    "test_data.set_format('torch', columns=['input_ids', 'attention_mask', 'labels'])\n",
    "\n",
    "\n",
    "train_data = train_data.map(to_categorical, batched = False, batch_size = len(train_data))\n",
    "test_data = test_data.map(to_categorical, batched = False, batch_size = len(test_data))\n",
    "\n",
    "# define accuracy metrics\n",
    "def compute_metrics(pred):\n",
    "    labels = pred.label_ids\n",
    "    preds = pred.predictions.argmax(-1)\n",
    "    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='macro')\n",
    "    acc = accuracy_score(labels, preds)\n",
    "    return {\n",
    "        'accuracy': acc,\n",
    "        'f1': f1,\n",
    "        'precision': precision,\n",
    "        'recall': recall\n",
    "    }\n",
    "\n",
    "# define the training arguments\n",
    "training_args = TrainingArguments(\n",
    "    output_dir = './ckpt/roberta-clariq/',\n",
    "    num_train_epochs=100,\n",
    "    per_device_train_batch_size = 32,\n",
    "    gradient_accumulation_steps = 16,    \n",
    "    per_device_eval_batch_size= 32,\n",
    "    evaluation_strategy = \"epoch\",\n",
    "    save_strategy = 'epoch',\n",
    "    disable_tqdm = False, \n",
    "    load_best_model_at_end=True,\n",
    "    warmup_steps=500,\n",
    "    weight_decay=0.01,\n",
    "    logging_steps = 8,\n",
    "    fp16 = True,\n",
    "    logging_dir='./ckpt/roberta-clariq/log',\n",
    "    dataloader_num_workers = 8,\n",
    "    run_name = 'roberta-classification',\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    compute_metrics=compute_metrics,\n",
    "    train_dataset=train_data,\n",
    "    eval_dataset=test_data\n",
    ")\n",
    "\n",
    "trainer.train()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Generate decoder inputs with Finetuned RoBERTa (requires step 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import RobertaTokenizer, RobertaForSequenceClassification\n",
    "\n",
    "id2prefix = {\n",
    "    0: 'i dont know',\n",
    "    1: 'no',\n",
    "    2: '',\n",
    "    3: 'yes'\n",
    "}\n",
    "\n",
    "roberta_tokenizer = RobertaTokenizer.from_pretrained('roberta-base')\n",
    "roberta_model = RobertaForSequenceClassification.from_pretrained(\"ckpt/roberta-qulac/checkpoint-102\")\n",
    "for output_name in qulac_output_names:\n",
    "    df = pd.read_csv(output_name)\n",
    "    for iter, row in df.iterrows():\n",
    "        clf_inputs = roberta_tokenizer(df.at[iter, 'unifiedqa-question'], return_tensors=\"pt\")\n",
    "        with T.no_grad():\n",
    "            logits = roberta_model(**clf_inputs).logits\n",
    "        predicted_class_id = logits.argmax().item()\n",
    "        df.at[iter, 'decoder-input'] = id2prefix[predicted_class_id]\n",
    "    df.to_csv(output_name, index=False)\n",
    "\n",
    "roberta_model = RobertaForSequenceClassification.from_pretrained(\"ckpt/roberta-clariq/checkpoint-100\")\n",
    "for output_name in clariq_output_names:\n",
    "    df = pd.read_csv(output_name)\n",
    "    for iter, row in df.iterrows():\n",
    "        clf_inputs = roberta_tokenizer(df.at[iter, 'unifiedqa-question'], return_tensors=\"pt\")\n",
    "        with T.no_grad():\n",
    "            logits = roberta_model(**clf_inputs).logits\n",
    "        predicted_class_id = logits.argmax().item()\n",
    "        df.at[iter, 'decoder-input'] = id2prefix[predicted_class_id]\n",
    "    df.to_csv(output_name, index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5. Post-processing the generation output (Run this after finetuning your generation models)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# aggregate the output of generation models to the csv file\n",
    "import pandas as pd\n",
    "output_dir = 'output/t5-qulac'   \n",
    "output_csv = 'output/t5-qulac.csv'\n",
    "test_file = '../data/processed/qulac_test.csv'\n",
    "#dataset = 'clariq_dev.csv'\n",
    "\n",
    "sentences = open(output_dir + '/generated_predictions.txt', 'r').readlines()\n",
    "df = pd.read_csv(test_file)\n",
    "\n",
    "assert len(sentences) == len(df)\n",
    "for iter, row in df.iterrows():\n",
    "    df.at[iter, 'candidate'] = sentences[iter].lower().strip()\n",
    "\n",
    "df.to_csv(output_csv, index=False)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}