TaskA - Metric Optimization - Contrastive Search.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8ee843e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "from pathlib import Path\n",
    "\n",
    "import datasets as ds\n",
    "import nltk\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import DataLoader, Dataset, RandomSampler, SequentialSampler\n",
    "from torch.optim import AdamW\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "import transformers\n",
    "from filelock import FileLock\n",
    "import huggingface_hub as hf_hub\n",
    "\n",
    "from transformers import AutoConfig, \\\n",
    "                         AutoModelForSeq2SeqLM, \\\n",
    "                         AutoTokenizer, \\\n",
    "                         BartTokenizer, \\\n",
    "                         DataCollatorForSeq2Seq, \\\n",
    "                         SchedulerType, \\\n",
    "                         get_scheduler, \\\n",
    "                         set_seed, \\\n",
    "                         get_linear_schedule_with_warmup, \\\n",
    "                         SchedulerType, \\\n",
    "                         AutoModelForSequenceClassification, \\\n",
    "                         GenerationConfig\n",
    "\n",
    "from bert_score import score\n",
    "import evaluate\n",
    "import wandb\n",
    "import pandas as pd\n",
    "import random\n",
    "\n",
    "import re\n",
    "import os\n",
    "import config as code_config\n",
    "\n",
    "from sklearn.model_selection import KFold, train_test_split, StratifiedKFold\n",
    "import math\n",
    "import json\n",
    "import gc\n",
    "from tqdm import tqdm\n",
    "import warnings\n",
    "import optuna\n",
    "from optuna.integration.wandb import WeightsAndBiasesCallback\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "tqdm.pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "424cb653",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"WANDB_API_KEY\"] = code_config.WANDB_API\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "os.environ[\"WANDB_MODE\"] = \"online\"\n",
    "hf_hub.login(code_config.HF_API)\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30d52e97",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    nltk.data.find(\"tokenizers/punkt\")\n",
    "except LookupError as e:\n",
    "    with FileLock(\".lock\") as lock:\n",
    "        nltk.download(\"punkt\", quiet=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3eac0091",
   "metadata": {},
   "outputs": [],
   "source": [
    "section_header_mapping = \\\n",
    "{\"fam/sochx\": [\"FAMILY HISTORY\",\"SOCIAL HISTORY\"], \\\n",
    "\"genhx\": [\"HISTORY of PRESENT ILLNESS\"], \\\n",
    "\"pastmedicalhx\": [\"PAST MEDICAL HISTORY\"], \\\n",
    "\"cc\": [\"CHIEF COMPLAINT\"], \\\n",
    "\"pastsurgical\": [\"PAST SURGICAL HISTORY\"], \\\n",
    "\"allergy\": [\"allergy\"], \\\n",
    "\"ros\": [\"REVIEW OF SYSTEMS\"], \\\n",
    "\"medications\": [\"medications\"], \\\n",
    "\"assessment\": [\"assessment\"], \\\n",
    "\"exam\": [\"exam\"], \\\n",
    "\"diagnosis\": [\"diagnosis\"], \\\n",
    "\"disposition\": [\"disposition\"], \\\n",
    "\"plan\": [\"plan\"], \\\n",
    "\"edcourse\": [\"EMERGENCY DEPARTMENT COURSE\"], \\\n",
    "\"immunizations\": [\"immunizations\"], \\\n",
    "\"imaging\": [\"imaging\"], \\\n",
    "\"gynhx\": [\"GYNECOLOGIC HISTORY\"], \\\n",
    "\"procedures\": [\"procedures\"], \\\n",
    "\"other_history\": [\"other_history\"], \\\n",
    "\"labs\": [\"labs\"]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd4220a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_path = Path.cwd().joinpath(\"mediqa-chat-data\", \"TaskA\", \"TaskA-TrainingSet.csv\")\n",
    "validation_path = Path.cwd().joinpath(\n",
    "    \"mediqa-chat-data\", \"TaskA\", \"TaskA-ValidationSet.csv\"\n",
    ")\n",
    "\n",
    "train_df = pd.read_csv(train_path, index_col=\"ID\")\n",
    "valid_df = pd.read_csv(validation_path, index_col=\"ID\")\n",
    "merge_df = pd.concat([train_df, valid_df], axis=0, ignore_index=True)\n",
    "merge_df[\"dialogue_wo_whitespaces\"] = merge_df[\"dialogue\"].apply(\n",
    "    lambda x: re.sub(r\"[\\r\\n\\s]+\", \" \", x)\n",
    ")\n",
    "merge_df.reset_index(inplace=True)\n",
    "merge_df.rename(mapper={\"index\": \"ID\"}, axis=1, inplace=True)\n",
    "\n",
    "with open(\"TaskA-label2idx.json\",\"r\") as f:\n",
    "    label2idx = json.load(f)\n",
    "    \n",
    "with open(\"TaskA-idx2label.json\",\"r\") as f:\n",
    "    idx2label = json.load(f)\n",
    "\n",
    "merge_df[\"label\"] = merge_df[\"section_header\"].apply(lambda x: label2idx[x])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a71dc1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def seed_everything(seed):\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e058fccd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def postprocess_text(preds,labels):\n",
    "    seed_everything(code_config.TASKA_SUMMARY_SEED)\n",
    "    preds = [pred.strip() for pred in preds]\n",
    "    labels = [label.strip() for label in labels]\n",
    "    \n",
    "    preds = [\"\\n\".join(nltk.sent_tokenize(pred)) for pred in preds]\n",
    "    labels = [\"\\n\".join(nltk.sent_tokenize(label)) for label in labels]\n",
    "    \n",
    "    return preds, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c5f4372",
   "metadata": {},
   "outputs": [],
   "source": [
    "scorers = {\n",
    "        'rouge': (\n",
    "            evaluate.load('rouge'),\n",
    "            {'use_aggregator': False},\n",
    "            ['rouge1', 'rouge2', 'rougeL', 'rougeLsum'],\n",
    "            ['rouge1', 'rouge2', 'rougeL', 'rougeLsum']\n",
    "        ),\n",
    "        'bert_scorer': (\n",
    "            evaluate.load('bertscore'),\n",
    "            {'model_type': 'microsoft/deberta-xlarge-mnli'},\n",
    "            ['precision', 'recall', 'f1'],\n",
    "            ['bertscore_precision', 'bertscore_recall', 'bertscore_f1']\n",
    "        ),\n",
    "        'bluert': (\n",
    "            evaluate.load('bleurt', config_name='BLEURT-20'),\n",
    "            {},\n",
    "            ['scores'],\n",
    "            ['bleurt']\n",
    "        ),\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c22c600b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_metrics(references,predictions,scorer,key,save_key,**kwargs):\n",
    "        scores = scorer.compute(references=references, predictions=predictions, **kwargs)\n",
    "        if isinstance(scores[key],list):\n",
    "            if len(scores[key]) > 1:\n",
    "                raise Exception(\"scores[key] have more than one elements\")\n",
    "            return scores[key][0]\n",
    "        return scores[key]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "725671da",
   "metadata": {},
   "outputs": [],
   "source": [
    "TASKA_SUMMARY_TOKENIZER_MODEL_MAPPING = \\\n",
    "{\n",
    "    \"GanjinZero/biobart-v2-base\": \\\n",
    "    \"suryakiran786/5-fold-stratified-cv-biobart-v2-base-with-section-description-complete-data\",\n",
    "    \"google/flan-t5-large\": \\\n",
    "    \"suryakiran786/5-fold-stratified-cv-flan-t5-large-with-section-description-complete-data\", \n",
    "    \"MingZhong/DialogLED-large-5120\": \\\n",
    "    \"suryakiran786/5-fold-stratified-cv-dialogled-large-with-section-description-complete-data\",\n",
    "    \"MingZhong/DialogLED-base-16384\": \\\n",
    "    \"suryakiran786/5-fold-stratified-cv-dialogled-base-with-section-description-complete-data\"\n",
    "}\n",
    "TASKA_MULTI_CLASS_MODEL_MAPPING = \\\n",
    "{\n",
    "    \"emilyalsentzer/Bio_ClinicalBERT\": \\\n",
    "    \"suryakiran786/5-stratified-cv-bio-clinicalbert-multiclass-focal-loss-seed-42-complete-data\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58403dd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb_kwargs = {\"project\": \"metric-optimization-hpo\",\"group\":\"contrastive-search\"}\n",
    "wandbc = WeightsAndBiasesCallback(metric_name=\"rouge_bertscore_bleurt_score\", \\\n",
    "                                  wandb_kwargs=wandb_kwargs, \\\n",
    "                                  as_multirun=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "296a8d6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def metric_calculation(summary_tokenizer, summary_model, \\\n",
    "                       classifier_tokenizer, classifier_model, \\\n",
    "                       df, split):\n",
    "    \n",
    "    @wandbc.track_in_wandb()\n",
    "    def objective(trial):\n",
    "        \n",
    "        max_length = trial.suggest_int(\"max_length\", \\\n",
    "                                       code_config.TASKA_SUMMARY_MIN_TARGET_LENGTH, \\\n",
    "                                       code_config.TASKA_SUMMARY_MAX_TARGET_LENGTH)\n",
    "        penalty_alpha = trial.suggest_float(\"penalty_alpha\",0.1,10.)\n",
    "        top_k = trial.suggest_int(\"top_k\",1,15)\n",
    "        \n",
    "        generate_kwargs = {\n",
    "            \"max_length\": max_length, \\\n",
    "            \"penalty_alpha\": penalty_alpha, \\\n",
    "            \"top_k\": top_k\n",
    "        }\n",
    "    \n",
    "        TASKA_SUMMARY_TOKENIZER = \\\n",
    "        f\"{summary_tokenizer}\"\n",
    "        TASKA_SUMMARY_CHECKPOINT = \\\n",
    "        f\"{summary_model}-{split}\"\n",
    "\n",
    "        MULTI_CLASS_TOKENIZER = f\"{classifier_tokenizer}\"\n",
    "        MULTI_CLASS_CHECKPOINT = f\"{classifier_model}-{split}\"\n",
    "\n",
    "        multi_class_config = AutoConfig.from_pretrained(MULTI_CLASS_CHECKPOINT)\n",
    "        multi_class_config.num_labels = 20\n",
    "\n",
    "        taska_summary_config = AutoConfig.from_pretrained(TASKA_SUMMARY_CHECKPOINT)\n",
    "\n",
    "        multi_class_tokenizer = AutoTokenizer.from_pretrained(MULTI_CLASS_TOKENIZER, \\\n",
    "                                                                do_lower_case=True, \\\n",
    "                                                                force_download=True)\n",
    "\n",
    "        taska_summary_tokenizer = AutoTokenizer.from_pretrained(TASKA_SUMMARY_TOKENIZER, \\\n",
    "                                                                do_lower_case=True, \\\n",
    "                                                                force_download=True)\n",
    "\n",
    "        multi_class_model = \\\n",
    "        AutoModelForSequenceClassification.from_pretrained(MULTI_CLASS_CHECKPOINT, \\\n",
    "                                                           config=multi_class_config, \\\n",
    "                                                           force_download=True)\n",
    "\n",
    "        taska_summary_model = \\\n",
    "        AutoModelForSeq2SeqLM.from_pretrained(TASKA_SUMMARY_CHECKPOINT, \\\n",
    "                                              config=taska_summary_config, \\\n",
    "                                              force_download=True)\n",
    "\n",
    "        multi_class_model = multi_class_model.to(device)\n",
    "        multi_class_model.eval()\n",
    "\n",
    "        taska_summary_model = taska_summary_model.to(device)\n",
    "        taska_summary_model.eval()\n",
    "\n",
    "        test_df = df\n",
    "        test_df[\"predicted_section_header\"] = None\n",
    "        test_df[\"predicted_section_text_postprocessed\"] = None\n",
    "        test_df[\"reference_section_text_postprocessed\"] = None\n",
    "\n",
    "        for idx in test_df.index:\n",
    "            sentence = test_df.loc[idx,\"dialogue_wo_whitespaces\"]\n",
    "\n",
    "            tokenized_sentence = \\\n",
    "            multi_class_tokenizer.encode_plus(sentence,\n",
    "                                            add_special_tokens=True,\n",
    "                                            padding=\"max_length\",\n",
    "                                            truncation=True,\n",
    "                                            max_length=code_config.MULTI_CLASS_MAX_LENGTH,\n",
    "                                            verbose=False,\n",
    "                                            return_tensors=\"pt\",\n",
    "                                            return_attention_mask=True)\n",
    "\n",
    "            input_ids = tokenized_sentence[\"input_ids\"].to(device)\n",
    "            token_type_ids = tokenized_sentence[\"token_type_ids\"].to(device)\n",
    "            attention_mask = tokenized_sentence[\"attention_mask\"].to(device)\n",
    "\n",
    "            with torch.no_grad():\n",
    "                preds = multi_class_model(input_ids=input_ids, \\\n",
    "                                          token_type_ids=token_type_ids, \\\n",
    "                                          attention_mask=attention_mask)\n",
    "\n",
    "                preds = preds.logits.detach().cpu().numpy().squeeze(0)\n",
    "\n",
    "                best_idx = np.argmax(preds)\n",
    "                section_header = idx2label[str(best_idx)]\n",
    "                test_df.loc[idx,\"predicted_section_header\"] = section_header\n",
    "\n",
    "        test_df[\"predicted_section_header_desription\"] = \\\n",
    "        test_df[\"predicted_section_header\"].apply(lambda x: \" and \".join(section_header_mapping[x.lower()]))\n",
    "        test_df[\"predicted_section_header_desription\"] = \\\n",
    "        test_df[\"predicted_section_header_desription\"].str.lower()\n",
    "\n",
    "        summary_column = \"section_text\"\n",
    "        text_column = \"dialogue_w_section_header_desc\"\n",
    "        test_df[text_column] = \\\n",
    "        test_df[\"predicted_section_header_desription\"] + \\\n",
    "        f\" {str(taska_summary_tokenizer.sep_token)} \" + \\\n",
    "        test_df[\"dialogue_wo_whitespaces\"]\n",
    "\n",
    "        for idx in test_df.index:\n",
    "            sentence = test_df.loc[idx,text_column]\n",
    "            summary = test_df.loc[idx,summary_column]        \n",
    "\n",
    "            model_inputs = \\\n",
    "            taska_summary_tokenizer(sentence, \\\n",
    "                                    padding=code_config.TASKA_SUMMARY_PADDING, \\\n",
    "                                    truncation=True, \\\n",
    "                                    max_length=code_config.TASKA_SUMMARY_MAX_SOURCE_LENGTH, \\\n",
    "                                    return_tensors=\"pt\")\n",
    "\n",
    "            labels = \\\n",
    "            taska_summary_tokenizer(text_target=summary, \\\n",
    "                                    padding=code_config.TASKA_SUMMARY_PADDING, \\\n",
    "                                    truncation=True, \\\n",
    "                                    max_length=code_config.TASKA_SUMMARY_MAX_TARGET_LENGTH, \\\n",
    "                                    return_tensors=\"pt\")\n",
    "\n",
    "            model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
    "\n",
    "            with torch.no_grad():\n",
    "\n",
    "                input_ids = model_inputs[\"input_ids\"].to(device)\n",
    "                attention_mask = model_inputs[\"attention_mask\"].to(device)\n",
    "                labels = model_inputs[\"labels\"].to(device)\n",
    "\n",
    "                generated_tokens = \\\n",
    "                taska_summary_model.generate(inputs=input_ids, \\\n",
    "                                             attention_mask=attention_mask, \\\n",
    "                                             **generate_kwargs)\n",
    "\n",
    "                if isinstance(generated_tokens,tuple):\n",
    "                    generated_tokens = generated_tokens[0]\n",
    "\n",
    "                generated_tokens_decoded = \\\n",
    "                taska_summary_tokenizer.batch_decode(generated_tokens,skip_special_tokens=True)\n",
    "                labels_w_padding_tokens = \\\n",
    "                [[l.item() if l != -100 else taska_summary_tokenizer.pad_token_id for l in label] \\\n",
    "                 for label in labels.cpu()]\n",
    "                labels_decoded = \\\n",
    "                taska_summary_tokenizer.batch_decode(labels_w_padding_tokens,skip_special_tokens=True)\n",
    "\n",
    "                generated_tokens_decoded,labels_decoded = \\\n",
    "                postprocess_text(generated_tokens_decoded,labels_decoded)\n",
    "\n",
    "                test_df.loc[idx,\"predicted_section_text_postprocessed\"] = \\\n",
    "                generated_tokens_decoded\n",
    "\n",
    "                test_df.loc[idx,\"reference_section_text_postprocessed\"] = \\\n",
    "                labels_decoded\n",
    "\n",
    "        for name, (scorer, kwargs, keys, save_keys) in scorers.items():\n",
    "\n",
    "            for key, save_key in zip(keys, save_keys):\n",
    "                test_df[f\"metrics_{save_key}\"] = \\\n",
    "                test_df.progress_apply(lambda x: calculate_metrics(references=x[\"reference_section_text_postprocessed\"], \\\n",
    "                                                          predictions=x[\"predicted_section_text_postprocessed\"], \\\n",
    "                                                          scorer=scorer, \\\n",
    "                                                          key=key, \\\n",
    "                                                          save_key=save_key, \\\n",
    "                                                          **kwargs),axis=1)\n",
    "\n",
    "        return test_df[[col for col in test_df.columns if \"metrics_\" in col]].mean().mean()\n",
    "    \n",
    "    return objective"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d515d1b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "skf = StratifiedKFold(n_splits=code_config.TASKA_SUMMARY_N_SPLITS,shuffle=True,random_state=code_config.SEED)\n",
    "for split,(train_idx,valid_idx) in enumerate(skf.split(merge_df, y=merge_df[\"label\"])):\n",
    "    for summary_tokenizer, summary_model in TASKA_SUMMARY_TOKENIZER_MODEL_MAPPING.items():\n",
    "        for mc_tokenizer,mc_model in TASKA_MULTI_CLASS_MODEL_MAPPING.items():\n",
    "            \n",
    "            test_df = merge_df.iloc[valid_idx]\n",
    "            objective_fn = \\\n",
    "            metric_calculation(summary_tokenizer, \\\n",
    "                   summary_model, \\\n",
    "                   mc_tokenizer, \n",
    "                   mc_model, \\\n",
    "                   test_df, split)\n",
    "            \n",
    "            study_name = \\\n",
    "            summary_tokenizer.split(\"/\")[-1] + \\\n",
    "            \"-\" + \\\n",
    "            mc_tokenizer.split(\"/\")[-1] + \\\n",
    "            \"-\" + \\\n",
    "            str(split)\n",
    "            \n",
    "            study = optuna.create_study(study_name=study_name, \\\n",
    "                                        direction=\"maximize\")\n",
    "            study.optimize(objective_fn, n_trials=20, callbacks=[wandbc])"
   ]
  }
 ],
 "metadata": {
  "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": 5
}