notebooks/kt_bert.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Knowledge Distillation using MedNLI as an intermediate dataset\n",
    "1) train a BERT model on AP and evaluate it \n",
    "2) use trained model to label MedNLI data with AP classes\n",
    "3) train a BERT model on the dataset from (2)\n",
    "4) evaluate model from (3) on the AP test set (as in (1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import re\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "\n",
    "from datasets import load_dataset\n",
    "\n",
    "import torch\n",
    "\n",
    "from transformers import (\n",
    "    AutoModelForSequenceClassification,\n",
    "    AutoTokenizer,\n",
    "    DataCollatorWithPadding,\n",
    "    Trainer,\n",
    "    TrainingArguments,\n",
    ")\n",
    "\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "from sklearn.metrics import classification_report\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1,2\"\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "import sys\n",
    "sys.path.append('../utils/')\n",
    "\n",
    "from utils import set_seed\n",
    "set_seed(42)\n",
    "\n",
    "lm_name = \"bert-base-cased\"\n",
    "layer_sharing_mode = 'even-6'\n",
    "\n",
    "OUT_PATH = '../output/{}/'.format(lm_name)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Models training/loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_checkpoint_ap = False\n",
    "load_checkpoint_mednli = False \n",
    "mednli_oversampled = False\n",
    "AP_BEST_DIR = './AP/best/'\n",
    "MEDNLI_BEST_DIR = './MedNLI/best/'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load AP Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: adapt path to data files\n",
    "data_files = {\n",
    "    \"train\": './ap/train.csv',\n",
    "    \"dev\": './ap/dev.csv',\n",
    "    \"test\": './ap/test.csv',\n",
    "}\n",
    "\n",
    "data = load_dataset(\"csv\", data_files=data_files, )\n",
    "\n",
    "LABEL_ENCODER = LabelEncoder()\n",
    "LABEL_ENCODER.fit(data['train']['Relation'])\n",
    "LABEL_ENCODER.classes_\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "AP Preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PHI_PATTERN = re.compile(r'\\[\\*\\*[^\\]]+\\*\\*\\]')\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    lm_name,\n",
    "    do_lower_case=False\n",
    ")\n",
    "tokenizer.add_tokens(['@@PHI@@'], special_tokens=True)\n",
    "\n",
    "def preprocess(row):   \n",
    "    d = {\n",
    "        'Assessment': PHI_PATTERN.sub('@@PHI@@', row['Assessment']),\n",
    "        'Plan Subsection': PHI_PATTERN.sub('@@PHI@@', row['Plan Subsection']),\n",
    "    }\n",
    "    \n",
    "    if row['Relation']:\n",
    "        d['label'] = LABEL_ENCODER.transform([row['Relation']])[0]\n",
    "    else:\n",
    "        d['label'] = 0 # if we have no label (during test), we just use a default label of 0\n",
    "    \n",
    "    return d\n",
    "\n",
    "def tokenize(examples):\n",
    "    return tokenizer(\n",
    "        examples['Assessment'],\n",
    "        examples['Plan Subsection'],\n",
    "        truncation=True,\n",
    "        max_length=512\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = data.map(preprocess)\n",
    "data = data.map(tokenize, batched=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.decode(data['train'][0]['input_ids'])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Train Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "model = AutoModelForSequenceClassification.from_pretrained(\n",
    "    lm_name,\n",
    "    num_labels=len(LABEL_ENCODER.classes_)\n",
    ")\n",
    "model.resize_token_embeddings(len(tokenizer))\n",
    "\n",
    "\n",
    "from sklearn.metrics import precision_recall_fscore_support\n",
    "\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",
    "    return {\n",
    "        'f1': f1,\n",
    "        'precision': precision,\n",
    "        'recall': recall\n",
    "    }\n",
    "\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=OUT_PATH,\n",
    "    num_train_epochs=3, \n",
    "    learning_rate=5e-5,\n",
    "    warmup_ratio=0,\n",
    "    weight_decay=0,\n",
    "    per_device_train_batch_size=8,\n",
    "    per_device_eval_batch_size=512,\n",
    "    logging_strategy=\"steps\",\n",
    "    logging_steps=20,\n",
    "    evaluation_strategy=\"steps\",\n",
    "    eval_steps=20,\n",
    "    save_strategy=\"steps\",\n",
    "    save_steps=20,\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model='f1',\n",
    "    greater_is_better=True,\n",
    "    save_total_limit = 2,\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=data[\"train\"],\n",
    "    eval_dataset=data[\"dev\"], \n",
    "    tokenizer=tokenizer,\n",
    "    compute_metrics=compute_metrics,\n",
    ")\n",
    "\n",
    "if load_checkpoint_ap:\n",
    "    loaded_model_ap = AutoModelForSequenceClassification.from_pretrained(AP_BEST_DIR,\n",
    "    local_files_only=True)\n",
    "    loaded_model_ap.to(device)\n",
    "    trainer.model = loaded_model_ap\n",
    "    model = loaded_model_ap\n",
    "else: \n",
    "    trainer.train()\n",
    "    trainer.save_model(AP_BEST_DIR)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluation on AP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(trainer, data, split_name: str):\n",
    "    print(f'Metrics for {split_name}')\n",
    "\n",
    "    preds = trainer.predict(data[split_name])\n",
    "    logits = torch.tensor(preds.predictions)\n",
    "    y_pred = np.argmax(preds.predictions, axis=-1)\n",
    "    y_pred = LABEL_ENCODER.inverse_transform(y_pred)\n",
    "    y_pred_proba = torch.nn.functional.softmax(logits, dim=1)\n",
    "    \n",
    "    \n",
    "    y_true = preds.label_ids\n",
    "    y_true = LABEL_ENCODER.inverse_transform(y_true)\n",
    "    print(f\"Evaluate {split_name}\\n\")\n",
    "    print(classification_report(y_true, y_pred, digits=3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predict(trainer, data, 'test')\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading MedNLI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: adapt path to data files\n",
    "mednli_data_files = {\n",
    "    \"train\": './physionet.org/files/mednli/1.0.0/mli_train_v1.jsonl',\n",
    "    \"dev\": './physionet.org/files/mednli/1.0.0/mli_dev_v1.jsonl',\n",
    "\n",
    "}\n",
    "mednli_data = load_dataset(\"json\", data_files=mednli_data_files )\n",
    "\n",
    "def mednli_preprocess(row):   \n",
    "    d = {\n",
    "        'Assessment': PHI_PATTERN.sub('@@PHI@@', row['sentence1']),\n",
    "        'Plan Subsection': PHI_PATTERN.sub('@@PHI@@', row['sentence2']),\n",
    "    }\n",
    "    \n",
    " \n",
    "    d['label'] = 0 # if we have no label (during test), we just use a default label of 0\n",
    "    \n",
    "    return d\n",
    "\n",
    "mednli_data = mednli_data.map(mednli_preprocess)\n",
    "mednli_data = mednli_data.map(tokenize, batched=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(mednli_data['train'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(mednli_data['dev'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mednli_predict(trainer, data, split_name: str, out_path):\n",
    "    print(f\"Labeling MedNLI_{split_name}...\\n\")\n",
    "    out_path = Path(out_path)\n",
    "    out_path.mkdir(exist_ok=True)\n",
    "    \n",
    "    preds = trainer.predict(data[split_name])\n",
    "    logits = torch.tensor(preds.predictions)\n",
    "    y_pred = np.argmax(preds.predictions, axis=-1)\n",
    "    y_pred = LABEL_ENCODER.inverse_transform(y_pred)\n",
    "    y_pred_proba = torch.nn.functional.softmax(logits, dim=1)\n",
    "    \n",
    "    \n",
    "\n",
    "    with open(out_path / f\"y_pred_mednli{split_name}.txt\", \"w\") as fout:\n",
    "        for i in y_pred:\n",
    "            fout.write(str(i) + \"\\n\")\n",
    "            \n",
    "    with open(out_path / f\"y_pred_proba_mednli_{split_name}.jsonl\", \"w\") as fout:\n",
    "        for ps in y_pred_proba.tolist():\n",
    "            json.dump(ps, fout)\n",
    "            fout.write('\\n')\n",
    "    \n",
    "    return y_pred, y_pred_proba\n",
    "\n",
    "# Hard/soft AP labels for training/dev sets of MedNLI\n",
    "hard_labels, soft_labels = mednli_predict(trainer, mednli_data, 'train', OUT_PATH)\n",
    "hard_labels_dev, soft_labels_dev = mednli_predict(trainer, mednli_data, 'dev', OUT_PATH)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "mednli_label_cntr = Counter(hard_labels)\n",
    "# \n",
    "for k, v in mednli_label_cntr.most_common(4):\n",
    "    print(k + '      : ', np.round(v/len(hard_labels)*100, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mednli_label_cntr_dev = Counter(hard_labels_dev)\n",
    "\n",
    "for k, v in mednli_label_cntr_dev.most_common(4):\n",
    "    print(k + '      : ', np.round(v/len(hard_labels_dev)*100, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Converting hard labels (strings) to integers\n",
    "transformed_labels_train = LABEL_ENCODER.transform(hard_labels)\n",
    "transformed_labels_dev = LABEL_ENCODER.transform(hard_labels_dev)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "updated_mednli_data = mednli_data.copy()\n",
    "# Creating a copy of the data and converting to integer labels\n",
    "updated_mednli_data['train'] = mednli_data['train'].map(lambda example, idx: {'label': transformed_labels_train[idx]}, with_indices=True)\n",
    "\n",
    "updated_mednli_data['dev'] = updated_mednli_data['dev'].map(lambda example, idx: {'label': transformed_labels_dev[idx]}, with_indices=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Oversampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from imblearn.over_sampling import RandomOverSampler\n",
    "from datasets import Dataset\n",
    "\n",
    "def oversample(data, split='train'):\n",
    "    df = data[split].to_pandas()\n",
    "    x_cols = df.columns.tolist()\n",
    "    x_cols.remove('label')\n",
    "    X_dev = df[x_cols]\n",
    "    y_dev = df['label']\n",
    "\n",
    "    oversample = RandomOverSampler(sampling_strategy='all')\n",
    "    X_over, y_over = oversample.fit_resample(X_dev, y_dev)\n",
    "    # add labels to df \n",
    "    X_over['label'] = y_over\n",
    "    \n",
    "    X_over['gold_label'] = LABEL_ENCODER.inverse_transform(X_over['label'])\n",
    "\n",
    "    data[split + '_oversampled'] = Dataset.from_pandas(X_over)\n",
    "\n",
    "    return data \n",
    "\n",
    "if mednli_oversampled:\n",
    "    updated_mednli_data = oversample(updated_mednli_data, split='dev')\n",
    "    updated_mednli_data = oversample(updated_mednli_data, split='train')\n",
    "    updated_mednli_data['dev_oversampled'].to_pandas()[['label', 'gold_label']].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if mednli_oversampled:\n",
    "    updated_mednli_data['train_oversampled'].to_pandas()[['label', 'gold_label']].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mednli_model = AutoModelForSequenceClassification.from_pretrained(\n",
    "    lm_name,\n",
    "    num_labels=len(LABEL_ENCODER.classes_)\n",
    ")\n",
    "mednli_model.resize_token_embeddings(len(tokenizer))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def share_layer(layer_num, sharing_mode):\n",
    "    '''\n",
    "    indicates if the layer with index 'layer_num' should be shared depending on the sharing_mode\n",
    "    '''\n",
    "    if sharing_mode == 'none':\n",
    "        return False\n",
    "    else:\n",
    "        num_layers = int(sharing_mode.split('-')[-1])*2\n",
    "        if layer_num % 2 == 0 and layer_num < num_layers:\n",
    "            return True\n",
    "        return False\n",
    "\n",
    "ap_model =  model.state_dict()\n",
    "\n",
    "\n",
    "for name, p in mednli_model.named_parameters():\n",
    "    if 'layer' in name:\n",
    "        n = name[19:]\n",
    "        first_dot = n.find('.') \n",
    "        layer_num = int(n[:first_dot])\n",
    "        # every_second layer ^\n",
    "        if share_layer(layer_num, layer_sharing_mode):\n",
    "            #p.data = ap_model[name]\n",
    "            p.data.copy_(ap_model[name].data)\n",
    "            # freeze layer\n",
    "            p.requires_grad = False\n",
    "\n",
    "ap_model =  model.state_dict()\n",
    "mednli_model_params = mednli_model.state_dict()\n",
    "\n",
    "# just checknig if all is good \n",
    "\n",
    "for name, p in mednli_model.named_parameters():\n",
    "    if 'layer' in name:\n",
    "        n = name[19:]\n",
    "        first_dot = n.find('.') \n",
    "        layer_num = int(n[:first_dot])\n",
    "        # every_second layer \n",
    "        if share_layer(layer_num, layer_sharing_mode):\n",
    "            #print(mednli_model_params[name])\n",
    "            #print(p.requires_grad)\n",
    "            #print(ap_model[name])\n",
    "            assert(p.requires_grad == False)\n",
    "        else:\n",
    "            assert(p.requires_grad == True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "o_dir = '../output/{}/{}/bert-mednli/'.format(lm_name, layer_sharing_mode)\n",
    "mednli_train_set = updated_mednli_data['train']\n",
    "mednli_dev_set = updated_mednli_data['dev']\n",
    "\n",
    "if mednli_oversampled:\n",
    "    o_dir = '../output/{}/{}/bert-mednli-oversampled/'.format(lm_name, layer_sharing_mode)\n",
    "    mednli_train_set = updated_mednli_data['train_oversampled']\n",
    "    mednli_dev_set = updated_mednli_data['dev_oversampled']\n",
    "\n",
    "training_args_mednli = TrainingArguments(\n",
    "    output_dir= o_dir,\n",
    "    num_train_epochs=1, \n",
    "    learning_rate=5e-5,\n",
    "    warmup_ratio=0,\n",
    "    weight_decay=0,\n",
    "    per_device_train_batch_size=8,\n",
    "    per_device_eval_batch_size=512,\n",
    "    logging_strategy=\"steps\",\n",
    "    logging_steps=20,\n",
    "    evaluation_strategy=\"steps\",\n",
    "    eval_steps=20,\n",
    "    save_strategy=\"steps\",\n",
    "    save_steps=20,\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model='f1',\n",
    "    greater_is_better=True,\n",
    "    save_total_limit = 2,\n",
    ")\n",
    "\n",
    "mednli_trainer = Trainer(\n",
    "    model = mednli_model,\n",
    "    args = training_args_mednli,\n",
    "    train_dataset = mednli_train_set,\n",
    "    eval_dataset = mednli_dev_set,\n",
    "    tokenizer = tokenizer,\n",
    "    compute_metrics = compute_metrics,\n",
    ")\n",
    "\n",
    "if load_checkpoint_mednli:\n",
    "    print('Loading checkpoint...')\n",
    "    if mednli_oversampled: \n",
    "        raise ValueError\n",
    "    else:\n",
    "        if layer_sharing_mode == 'none':\n",
    "            loaded_model_mednli = AutoModelForSequenceClassification.from_pretrained(MEDNLI_BEST_DIR + 'none/' , local_files_only=True)\n",
    "        elif layer_sharing_mode == 'even-6':\n",
    "            loaded_model_mednli = AutoModelForSequenceClassification.from_pretrained(MEDNLI_BEST_DIR + 'even-6/' , local_files_only=True)\n",
    "        elif layer_sharing_mode == 'even-3':\n",
    "            loaded_model_mednli = AutoModelForSequenceClassification.from_pretrained(MEDNLI_BEST_DIR + 'even-3/' , local_files_only=True)\n",
    "        else:\n",
    "            raise ValueError\n",
    "        \n",
    "        \n",
    "    loaded_model_mednli.to(device)\n",
    "    mednli_trainer.model = loaded_model_mednli\n",
    "else: \n",
    "    mednli_trainer.train()\n",
    "    mednli_trainer.save_model(MEDNLI_BEST_DIR + layer_sharing_mode +'/')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predict(mednli_trainer, data, 'test')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predict(trainer, data, 'test')\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.12 ('soapnotes')",
   "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.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "bf027a43a3fb0ff19c7f27db1c676521a8a29705b1546052c9305452682fdd7b"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}