step_1_2_boundary_prediction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c3dbf646",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "from datasets import load_dataset, Sequence, ClassLabel\n",
    "from enum import Enum\n",
    "from datetime import datetime, timedelta\n",
    "from pytz import timezone\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.optim import AdamW\n",
    "from tqdm.auto import tqdm\n",
    "from transformers import AutoTokenizer, AutoModelForTokenClassification\n",
    "from transformers import DataCollatorForTokenClassification\n",
    "\n",
    "import evaluate\n",
    "import json\n",
    "import numpy as np\n",
    "import os\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f387311a",
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_MODEL = 'microsoft/BiomedNLP-PubMedBERT-large-uncased-abstract'\n",
    "# boundaries-PICO_NER-ebm_nlp_bioc-2023_05_25_06_51_58_EDT\n",
    "# checkpoint-1508\n",
    "BOUNDARY_MODEL_PATH = 'pico_span/boundary_models/boundaries-PICO_NER-ebm_nlp_bioc-2023_05_31_20_05_09_EDT/checkpoint-4599'\n",
    "\n",
    "input_folder = 'data/bioc/json'\n",
    "OUTPUT_PATH = 'data/bioc/json/step_1_boundary_pred/threshold'\n",
    "OUTPUT_PATH_AD = 'data/bioc/json/step_1_boundary_pred/brat/AD'\n",
    "OUTPUT_PATH_COVID = 'data/bioc/json/step_1_boundary_pred/brat/COVID'\n",
    "OUTPUT_PATH_EBM_NLP = 'data/bioc/json/step_1_boundary_pred/brat/EBM-NLP'\n",
    "\n",
    "class DatasetSplit(Enum):\n",
    "    train = 0\n",
    "    validation = 1\n",
    "    test = 2\n",
    "    \n",
    "class PicoType(Enum):\n",
    "    PARTICIPANTS = 4\n",
    "    INTERVENTIONS = 2\n",
    "    OUTCOMES = 1\n",
    "    \n",
    "class SpanBoundary(Enum):\n",
    "    outside = 0\n",
    "    start = 1\n",
    "    end = 2\n",
    "    both = 3 # both start and end boundary\n",
    "    inside = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bdf8fcce",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)\n",
    "model = AutoModelForTokenClassification.from_pretrained(BOUNDARY_MODEL_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8029b615",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset json (/home/gzhang/.cache/huggingface/datasets/json/default-c27b087f41e9392f/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "364bf065f6dd49b0ab4bdfcbc95495d5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ebm_nlp = load_dataset(\n",
    "    'json',\n",
    "    data_files = {\n",
    "        'train': os.path.join(input_folder, 'test.json'),\n",
    "        'validation': os.path.join(input_folder, 'test.json'),\n",
    "        'test': os.path.join(input_folder, 'test.json')\n",
    "    }\n",
    ")\n",
    "\n",
    "remove_features = [f for f in ebm_nlp['train'].features if f not in['pmid', 'tokens', 'labels']]\n",
    "ebm_nlp['train'] = ebm_nlp['train'].remove_columns(remove_features)\n",
    "ebm_nlp['validation'] = ebm_nlp['validation'].remove_columns(remove_features)\n",
    "ebm_nlp['test'] = ebm_nlp['test'].remove_columns(remove_features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9cee4d82",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['pmid', 'tokens', 'labels'],\n",
       "        num_rows: 2042\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['pmid', 'tokens', 'labels'],\n",
       "        num_rows: 2042\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['pmid', 'tokens', 'labels'],\n",
       "        num_rows: 2042\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ebm_nlp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b6aac51b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def merge_boundary_labels_by_word_id(boundary_pred, word_ids, num_words):\n",
    "    labels = [0 for _ in range(num_words)]\n",
    "    for pred, word_id in zip(boundary_pred, word_ids):\n",
    "        if word_id is None:\n",
    "            continue\n",
    "        labels[word_id] = labels[word_id] | int(pred)\n",
    "    return labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f7439694",
   "metadata": {},
   "outputs": [],
   "source": [
    "def merge_boundary_probability_by_word_id(y_prob, word_ids, num_words):\n",
    "    start_prob = [0 for _ in range(num_words)]\n",
    "    end_prob = [0 for _ in range(num_words)]\n",
    "    for i, t in enumerate(zip(y_prob, word_ids)):\n",
    "        prob, word_id = t\n",
    "        if not word_id:\n",
    "            continue\n",
    "        start_prob[word_id] = max(\n",
    "            start_prob[word_id],\n",
    "            float(prob[SpanBoundary.start.value]),\n",
    "            float(prob[SpanBoundary.both.value]),\n",
    "        )\n",
    "        end_prob[word_id] = max(\n",
    "            end_prob[word_id],\n",
    "            float(prob[SpanBoundary.end.value]),\n",
    "            float(prob[SpanBoundary.both.value]),\n",
    "        )\n",
    "    return start_prob, end_prob\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "12a22ba5",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BoundaryLabel:\n",
    "    def __init__(self):\n",
    "        self.value = 0\n",
    "        \n",
    "    def set_start(self):\n",
    "        self.value = self.value | SpanBoundary.start.value\n",
    "        \n",
    "    def set_end(self):\n",
    "        self.value = self.value | SpanBoundary.end.value\n",
    "\n",
    "def extract_boundary_from_prob_dist(prob_dist, threshold):\n",
    "    label = BoundaryLabel()\n",
    "    if prob_dist[SpanBoundary.start.value] > threshold or prob_dist[SpanBoundary.both.value] > threshold:\n",
    "        label.set_start()\n",
    "    if prob_dist[SpanBoundary.end.value] > threshold or prob_dist[SpanBoundary.both.value] > threshold:\n",
    "        label.set_end()\n",
    "    return label.value\n",
    "        \n",
    "\n",
    "def generate_boundary_labels(dataset_dict, dataset_split, output_path, model, tokenizer, boundary_threshold=0.5):\n",
    "    # '{}_boundary_pred_high_precision.json'\n",
    "    output_file = os.path.join(output_path, f'{boundary_threshold:.2f}_boundary_pred.json')\n",
    "    if not os.path.exists(output_path):\n",
    "        os.makedirs(output_path)\n",
    "        \n",
    "    dataset = dataset_dict[dataset_split.name]\n",
    "    progress_bar = tqdm(range(len(dataset)))\n",
    "    with open(output_file, 'w+') as fout:\n",
    "        for i in range(len(dataset)):\n",
    "            row = {}\n",
    "            row['pmid'] = dataset['pmid'][i]\n",
    "            row['tokens'] = dataset['tokens'][i]\n",
    "            row['original_labels'] = dataset['labels'][i]\n",
    "            x = tokenizer(row['tokens'], padding=True, return_tensors='pt', is_split_into_words=True)\n",
    "            y = model(**x)\n",
    "            y_prob = np.squeeze(\n",
    "                torch.nn.functional.softmax(y.logits, dim=-1).detach().numpy())\n",
    "            y_pred = [extract_boundary_from_prob_dist(p, boundary_threshold) for p in y_prob]\n",
    "            # y_pred = np.argmax(y.logits.detach().numpy(), axis=-1)\n",
    "            row['boundary_pred'] = merge_boundary_labels_by_word_id(\n",
    "                y_pred,\n",
    "                x.word_ids(),\n",
    "                len(row['tokens']),\n",
    "            )\n",
    "            start_prob, end_prob = merge_boundary_probability_by_word_id(\n",
    "                y_prob,\n",
    "                x.word_ids(),\n",
    "                len(row['tokens']),\n",
    "            )\n",
    "            row['start_confidence'] = start_prob\n",
    "            row['end_confidence'] = end_prob\n",
    "            fout.write('{}\\n'.format(json.dumps(row)))\n",
    "            progress_bar.update(1)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "294154e1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6fcdb4de02aa4c5ebd8515f3b2ab3df9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1895 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# generate_boundary_labels(ebm_nlp, DatasetSplit.test, OUTPUT_PATH, model, tokenizer, boundary_threshold=0.25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f65836e1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4d51aa97dea7438a89dcfad4509ccc0e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2042 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "generate_boundary_labels(ebm_nlp, DatasetSplit.test, OUTPUT_PATH, model, tokenizer, boundary_threshold=0.20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4f126f0f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2c61130373d54c19bc8a242262df2392",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2042 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "generate_boundary_labels(ebm_nlp, DatasetSplit.test, OUTPUT_PATH, model, tokenizer, boundary_threshold=0.25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ed898cac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2bed462d2a4f4a4c8d569d8cb4e3d3cd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2042 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "generate_boundary_labels(ebm_nlp, DatasetSplit.test, OUTPUT_PATH, model, tokenizer, boundary_threshold=0.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e5f4b6d6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8cd27bc8cea743059fd9b2006e0b15ef",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2042 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "generate_boundary_labels(ebm_nlp, DatasetSplit.test, OUTPUT_PATH, model, tokenizer, boundary_threshold=0.35)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "00247056",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8d59816fc9aa491fb94494cfed69add1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2042 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "generate_boundary_labels(ebm_nlp, DatasetSplit.test, OUTPUT_PATH, model, tokenizer, boundary_threshold=0.4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "71e0120b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a1b88eaf5a294aee97d261a231c81987",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2042 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "generate_boundary_labels(ebm_nlp, DatasetSplit.test, OUTPUT_PATH, model, tokenizer, boundary_threshold=0.45)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c4713949",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "775a498940914d21a6f5da9b71a096c4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2042 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "generate_boundary_labels(ebm_nlp, DatasetSplit.test, OUTPUT_PATH, model, tokenizer, boundary_threshold=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ed5d211",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}