demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0205ad7a",
   "metadata": {},
   "source": [
    "# SciBERT Fine-Tuning \n",
    "\n",
    "---\n",
    "\n",
    "In this notebook, we use the 🤗 `transformers` library to fine-tune the `allenai/scibert_scivocab_uncased` model on various datasets. The goal is for the fine-tuned model to perform Named Entity Recognition of Drugs, Diseases, and Genes. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f94afa73-cca7-4cd5-8d74-8971f84eddc8",
   "metadata": {},
   "source": [
    "!pip install --target=$\"modules\" datasets transformers\n",
    "! pip install --target=$\"modules\" seqeval \n",
    "! pip install --target=$\"modules\" spacy "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "849c4d25-4f4e-4e18-b17e-665d4abf90d6",
   "metadata": {
    "tags": []
   },
   "source": [
    "! pip install --target=$\"modules\" evaluate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11034662-5c9f-45be-99d2-634d311a7d60",
   "metadata": {},
   "source": [
    "! pip install --target=$\"modules\" bioc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2f674f50",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\nicho\\miniconda3\\envs\\LGL1\\lib\\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": [
    "from datasets import Dataset, ClassLabel, Sequence, load_dataset, load_metric\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import bioc\n",
    "from spacy import displacy\n",
    "import transformers\n",
    "#import evaluate\n",
    "from transformers import (AutoModelForTokenClassification, \n",
    "                          AutoTokenizer, \n",
    "                          DataCollatorForTokenClassification,\n",
    "                          pipeline,\n",
    "                          TrainingArguments, \n",
    "                          Trainer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9c3c5157",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.31.0\n"
     ]
    }
   ],
   "source": [
    "# confirm version > 4.11.0\n",
    "print(transformers.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05ebaa92-58d0-40d4-ade9-a2ed6053827d",
   "metadata": {},
   "source": [
    "## Reading files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "101c50e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# no train-test provided, so we create our own\n",
    "cons_dataset = load_dataset(\"json\", data_files=\"./allDataShort2.json\")\n",
    "#cons_dataset = datasets\n",
    "cons_dataset = cons_dataset[\"train\"].train_test_split(test_size=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "db6c1a35-ba5f-4976-aa27-14293fe676e8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['index', 'text', 'gene', 'gene_indices_start', 'gene_indices_end', 'disease', 'disease_indices_start', 'disease_indices_end', 'drugs', 'drug_indices_start', 'drug_indices_end'],\n",
       "        num_rows: 6554\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['index', 'text', 'gene', 'gene_indices_start', 'gene_indices_end', 'disease', 'disease_indices_start', 'disease_indices_end', 'drugs', 'drug_indices_start', 'drug_indices_end'],\n",
       "        num_rows: 7\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cons_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9bfbe10",
   "metadata": {},
   "source": [
    "---\n",
    "## Token Labeling\n",
    "\n",
    "Finally, we can label each token with its entity. We use BIO tagging on two entities, `DRUG` and `EFFECT`. This results in five possible classes for each token:\n",
    "\n",
    "* `O` - outside any entity we care about\n",
    "* `B-DRUG` - the beginning of a `DRUG` entity\n",
    "* `I-DRUG` - inside a `DRUG` entity\n",
    "* `B-EFFECT` - the beginning of an `EFFECT` entity\n",
    "* `I-EFFECT` - inside an `EFFECT` entity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "89d0ec15",
   "metadata": {},
   "outputs": [],
   "source": [
    "label_list = ['O', 'B-DRUG', 'I-DRUG', 'B-DISEASE', 'I-DISEASE', 'B-GENE', 'I-GENE']\n",
    "\n",
    "custom_seq = Sequence(feature=ClassLabel(num_classes=7, \n",
    "                                         names=label_list,\n",
    "                                         names_file=None, id=None), length=-1, id=None)\n",
    "\n",
    "cons_dataset[\"train\"].features[\"ner_tags\"] = custom_seq\n",
    "cons_dataset[\"test\"].features[\"ner_tags\"] = custom_seq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "48c50247",
   "metadata": {},
   "outputs": [],
   "source": [
    "#model_checkpoint = \"allenai/scibert_scivocab_uncased\"\n",
    "model_checkpoint = './trainedSB2'\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c9481628",
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import isnan\n",
    "def generate_row_labels(row, verbose=False):\n",
    "    \"\"\" Given a row from the consolidated dataset, \n",
    "    generates BIO tags for drug and effect entities. \n",
    "    \"\"\"\n",
    "\n",
    "    text = row[\"text\"]\n",
    "\n",
    "    labels = []\n",
    "    label = \"O\"\n",
    "    prefix = \"\"\n",
    "    \n",
    "    # while iterating through tokens, increment to traverse all drug and effect spans\n",
    "    drug_index = 0\n",
    "    disease_index = 0\n",
    "    gene_index=0\n",
    "    \n",
    "    tokens = tokenizer(text, return_offsets_mapping=True)\n",
    "\n",
    "    for n in range(len(tokens[\"input_ids\"])):\n",
    "        offset_start, offset_end = tokens[\"offset_mapping\"][n]\n",
    "\n",
    "        # should only happen for [CLS] and [SEP]\n",
    "        if offset_end - offset_start == 0:\n",
    "            labels.append(-100)\n",
    "            continue\n",
    "        \n",
    "        if (type(row[\"drug_indices_start\"]) == list) and drug_index < len(row[\"drug_indices_start\"]) and offset_start == row[\"drug_indices_start\"][drug_index]:\n",
    "            label = \"DRUG\"\n",
    "            prefix = \"B-\"\n",
    "\n",
    "        elif (type(row[\"disease_indices_start\"]) == list) and disease_index < len(row[\"disease_indices_start\"]) and offset_start == row[\"disease_indices_start\"][disease_index]:\n",
    "            label = \"DISEASE\"\n",
    "            prefix = \"B-\"\n",
    "            \n",
    "        elif (type(row[\"gene_indices_start\"]) == list) and gene_index < len(row[\"gene_indices_start\"]) and offset_start == row[\"gene_indices_start\"][gene_index]:\n",
    "            label = \"GENE\"\n",
    "            prefix = \"B-\"\n",
    "        \n",
    "        labels.append(label_list.index(f\"{prefix}{label}\"))\n",
    "            \n",
    "        if (type(row[\"drug_indices_end\"]) == list) and drug_index < len(row[\"drug_indices_end\"]) and offset_end == row[\"drug_indices_end\"][drug_index]:\n",
    "            label = \"O\"\n",
    "            prefix = \"\"\n",
    "            drug_index += 1\n",
    "            \n",
    "        elif (type(row[\"disease_indices_end\"]) == list) and disease_index < len(row[\"disease_indices_end\"]) and offset_end == row[\"disease_indices_end\"][disease_index]:\n",
    "            label = \"O\"\n",
    "            prefix = \"\"\n",
    "            disease_index += 1\n",
    "            \n",
    "        elif (type(row[\"gene_indices_end\"]) == list) and gene_index < len(row[\"gene_indices_end\"]) and offset_end == row[\"gene_indices_end\"][gene_index]:\n",
    "            label = \"O\"\n",
    "            prefix = \"\"\n",
    "            gene_index += 1\n",
    "\n",
    "        # need to transition \"inside\" if we just entered an entity\n",
    "        if prefix == \"B-\":\n",
    "            prefix = \"I-\"\n",
    "    \n",
    "    if verbose:\n",
    "        print(f\"{row}\\n\")\n",
    "        orig = tokenizer.convert_ids_to_tokens(tokens[\"input_ids\"])\n",
    "        for n in range(len(labels)):\n",
    "            print(orig[n], labels[n])\n",
    "    tokens[\"labels\"] = labels\n",
    "    \n",
    "    return tokens\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4783acc7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map: 100%|█████████████████████████████████████████████████████████████████| 6554/6554 [00:14<00:00, 452.94 examples/s]\n",
      "Map: 100%|███████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 102.88 examples/s]\n"
     ]
    }
   ],
   "source": [
    "labeled_dataset = cons_dataset.map(generate_row_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb73c47b",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "---\n",
    "## SciBERT Model Fine-Tuning\n",
    "\n",
    "We are now ready to fine-tune the SciBERT model on our dataset. This section is modified from the following 🤗 notebook provided here: https://github.com/huggingface/notebooks/blob/master/examples/token_classification.ipynb\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "08cb693f",
   "metadata": {},
   "outputs": [],
   "source": [
    "task = \"ner\" # Should be one of \"ner\", \"pos\" or \"chunk\"\n",
    "#model_checkpoint = \"allenai/scibert_scivocab_uncased\"\n",
    "batch_size = 16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d7128ee4",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoModelForTokenClassification.from_pretrained(model_checkpoint, num_labels=len(label_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "ff2b60fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install transformers[torch]\n",
    "model_name = model_checkpoint.split(\"/\")[-1]\n",
    "args = TrainingArguments(\n",
    "    f\"{model_name}-finetuned-{task}\",\n",
    "    evaluation_strategy = \"epoch\",\n",
    "    learning_rate=1e-5,\n",
    "    per_device_train_batch_size=batch_size,\n",
    "    per_device_eval_batch_size=batch_size,\n",
    "    num_train_epochs=5,\n",
    "    weight_decay=0.05,\n",
    "    logging_steps=1\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "f1807139",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_collator = DataCollatorForTokenClassification(tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "560f0969",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install seqeval\n",
    "#module = evaluate.load(\"lvwerra/element_count\")\n",
    "metric = load_metric(\"seqeval\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "156e1cf4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_metrics(p):\n",
    "    predictions, labels = p\n",
    "    predictions = np.argmax(predictions, axis=2)\n",
    "\n",
    "    # Remove ignored index (special tokens)\n",
    "    true_predictions = [\n",
    "        [label_list[p] for (p, l) in zip(prediction, label) if l != -100]\n",
    "        for prediction, label in zip(predictions, labels)\n",
    "    ]\n",
    "    true_labels = [\n",
    "        [label_list[l] for (p, l) in zip(prediction, label) if l != -100]\n",
    "        for prediction, label in zip(predictions, labels)\n",
    "    ]\n",
    "\n",
    "    results = metric.compute(predictions=true_predictions, references=true_labels)\n",
    "    return {\n",
    "        \"precision\": results[\"overall_precision\"],\n",
    "        \"recall\": results[\"overall_recall\"],\n",
    "        \"f1\": results[\"overall_f1\"],\n",
    "        \"accuracy\": results[\"overall_accuracy\"],\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "bd9c7f71",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model,\n",
    "    args,\n",
    "    train_dataset=labeled_dataset[\"train\"],\n",
    "    eval_dataset=labeled_dataset[\"test\"],\n",
    "    data_collator=data_collator,\n",
    "    tokenizer=tokenizer,\n",
    "    compute_metrics=compute_metrics, \n",
    "\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "e644de18",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You're using a BertTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
     ]
    },
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "{'DRUG': {'precision': 1.0, 'recall': 1.0, 'f1': 1.0, 'number': 3},\n",
       " 'GENE': {'precision': 1.0, 'recall': 1.0, 'f1': 1.0, 'number': 13},\n",
       " 'overall_precision': 1.0,\n",
       " 'overall_recall': 1.0,\n",
       " 'overall_f1': 1.0,\n",
       " 'overall_accuracy': 1.0}"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predictions, labels, _ = trainer.predict(labeled_dataset[\"test\"])\n",
    "predictions = np.argmax(predictions, axis=2)\n",
    "\n",
    "# Remove ignored index (special tokens)\n",
    "true_predictions = [\n",
    "    [label_list[p] for (p, l) in zip(prediction, label) if l != -100]\n",
    "    for prediction, label in zip(predictions, labels)\n",
    "]\n",
    "true_labels = [\n",
    "    [label_list[l] for (p, l) in zip(prediction, label) if l != -100]\n",
    "    for prediction, label in zip(predictions, labels)\n",
    "]\n",
    "\n",
    "results = metric.compute(predictions=true_predictions, references=true_labels)\n",
    "results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c7e4324",
   "metadata": {},
   "source": [
    "---\n",
    "## See Model Outputs\n",
    "\n",
    "We load our fine-tuned model into a `pipeline` object to run arbitrary input against it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "51e55f32",
   "metadata": {},
   "outputs": [],
   "source": [
    "effect_ner_model = pipeline(task=\"ner\", model=model, tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "387dd2af",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'entity': 'LABEL_0',\n",
       "  'score': 0.9997309,\n",
       "  'index': 1,\n",
       "  'word': 'clinical',\n",
       "  'start': 0,\n",
       "  'end': 8},\n",
       " {'entity': 'LABEL_0',\n",
       "  'score': 0.9997073,\n",
       "  'index': 2,\n",
       "  'word': 'studies',\n",
       "  'start': 9,\n",
       "  'end': 16},\n",
       " {'entity': 'LABEL_0',\n",
       "  'score': 0.9996835,\n",
       "  'index': 3,\n",
       "  'word': 'have',\n",
       "  'start': 17,\n",
       "  'end': 21},\n",
       " {'entity': 'LABEL_0',\n",
       "  'score': 0.99972457,\n",
       "  'index': 4,\n",
       "  'word': 'indicated',\n",
       "  'start': 22,\n",
       "  'end': 31},\n",
       " {'entity': 'LABEL_0',\n",
       "  'score': 0.9997483,\n",
       "  'index': 5,\n",
       "  'word': 'the',\n",
       "  'start': 32,\n",
       "  'end': 35},\n",
       " {'entity': 'LABEL_0',\n",
       "  'score': 0.9997286,\n",
       "  'index': 6,\n",
       "  'word': 'following',\n",
       "  'start': 36,\n",
       "  'end': 45},\n",
       " {'entity': 'LABEL_0',\n",
       "  'score': 0.9997805,\n",
       "  'index': 7,\n",
       "  'word': 'relative',\n",
       "  'start': 46,\n",
       "  'end': 54},\n",
       " {'entity': 'LABEL_0',\n",
       "  'score': 0.9997863,\n",
       "  'index': 8,\n",
       "  'word': 'potency',\n",
       "  'start': 55,\n",
       "  'end': 62},\n",
       " {'entity': 'LABEL_0',\n",
       "  'score': 0.9997502,\n",
       "  'index': 9,\n",
       "  'word': 'differences',\n",
       "  'start': 63,\n",
       "  'end': 74},\n",
       " {'entity': 'LABEL_0',\n",
       "  'score': 0.99968636,\n",
       "  'index': 10,\n",
       "  'word': ':',\n",
       "  'start': 75,\n",
       "  'end': 76},\n",
       " {'entity': 'LABEL_1',\n",
       "  'score': 0.9793291,\n",
       "  'index': 11,\n",
       "  'word': 'flu',\n",
       "  'start': 77,\n",
       "  'end': 80},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.98558635,\n",
       "  'index': 12,\n",
       "  'word': '##ticas',\n",
       "  'start': 80,\n",
       "  'end': 85},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.99234176,\n",
       "  'index': 13,\n",
       "  'word': '##one',\n",
       "  'start': 85,\n",
       "  'end': 88},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.6478919,\n",
       "  'index': 14,\n",
       "  'word': 'prop',\n",
       "  'start': 89,\n",
       "  'end': 93},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.8109712,\n",
       "  'index': 15,\n",
       "  'word': '##ionate',\n",
       "  'start': 93,\n",
       "  'end': 99},\n",
       " {'entity': 'LABEL_0',\n",
       "  'score': 0.9969921,\n",
       "  'index': 16,\n",
       "  'word': '>',\n",
       "  'start': 100,\n",
       "  'end': 101},\n",
       " {'entity': 'LABEL_1',\n",
       "  'score': 0.9772372,\n",
       "  'index': 17,\n",
       "  'word': 'bud',\n",
       "  'start': 102,\n",
       "  'end': 105},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.9892029,\n",
       "  'index': 18,\n",
       "  'word': '##es',\n",
       "  'start': 105,\n",
       "  'end': 107},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.9902013,\n",
       "  'index': 19,\n",
       "  'word': '##oni',\n",
       "  'start': 107,\n",
       "  'end': 110},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.9839754,\n",
       "  'index': 20,\n",
       "  'word': '##de',\n",
       "  'start': 110,\n",
       "  'end': 112},\n",
       " {'entity': 'LABEL_0',\n",
       "  'score': 0.99758625,\n",
       "  'index': 21,\n",
       "  'word': '=',\n",
       "  'start': 113,\n",
       "  'end': 114},\n",
       " {'entity': 'LABEL_1',\n",
       "  'score': 0.9797254,\n",
       "  'index': 22,\n",
       "  'word': 'bec',\n",
       "  'start': 115,\n",
       "  'end': 118},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.99061537,\n",
       "  'index': 23,\n",
       "  'word': '##lo',\n",
       "  'start': 118,\n",
       "  'end': 120},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.9920591,\n",
       "  'index': 24,\n",
       "  'word': '##meth',\n",
       "  'start': 120,\n",
       "  'end': 124},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.98958415,\n",
       "  'index': 25,\n",
       "  'word': '##ason',\n",
       "  'start': 124,\n",
       "  'end': 128},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.98872375,\n",
       "  'index': 26,\n",
       "  'word': '##e',\n",
       "  'start': 128,\n",
       "  'end': 129},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.5297541,\n",
       "  'index': 27,\n",
       "  'word': 'dip',\n",
       "  'start': 130,\n",
       "  'end': 133},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.8482835,\n",
       "  'index': 28,\n",
       "  'word': '##rop',\n",
       "  'start': 133,\n",
       "  'end': 136},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.54456973,\n",
       "  'index': 29,\n",
       "  'word': '##ionate',\n",
       "  'start': 136,\n",
       "  'end': 142},\n",
       " {'entity': 'LABEL_0',\n",
       "  'score': 0.99813783,\n",
       "  'index': 30,\n",
       "  'word': '>',\n",
       "  'start': 143,\n",
       "  'end': 144},\n",
       " {'entity': 'LABEL_1',\n",
       "  'score': 0.9497072,\n",
       "  'index': 31,\n",
       "  'word': 'tri',\n",
       "  'start': 145,\n",
       "  'end': 148},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.97295266,\n",
       "  'index': 32,\n",
       "  'word': '##am',\n",
       "  'start': 148,\n",
       "  'end': 150},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.9831851,\n",
       "  'index': 33,\n",
       "  'word': '##cin',\n",
       "  'start': 150,\n",
       "  'end': 153},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.9881147,\n",
       "  'index': 34,\n",
       "  'word': '##olone',\n",
       "  'start': 153,\n",
       "  'end': 158},\n",
       " {'entity': 'LABEL_0',\n",
       "  'score': 0.4705981,\n",
       "  'index': 35,\n",
       "  'word': 'acet',\n",
       "  'start': 159,\n",
       "  'end': 163},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.8985247,\n",
       "  'index': 36,\n",
       "  'word': '##oni',\n",
       "  'start': 163,\n",
       "  'end': 166},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.81968755,\n",
       "  'index': 37,\n",
       "  'word': '##de',\n",
       "  'start': 166,\n",
       "  'end': 168},\n",
       " {'entity': 'LABEL_0',\n",
       "  'score': 0.99776864,\n",
       "  'index': 38,\n",
       "  'word': '=',\n",
       "  'start': 169,\n",
       "  'end': 170},\n",
       " {'entity': 'LABEL_1',\n",
       "  'score': 0.9653367,\n",
       "  'index': 39,\n",
       "  'word': 'flu',\n",
       "  'start': 171,\n",
       "  'end': 174},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.9850605,\n",
       "  'index': 40,\n",
       "  'word': '##nis',\n",
       "  'start': 174,\n",
       "  'end': 177},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.9890426,\n",
       "  'index': 41,\n",
       "  'word': '##oli',\n",
       "  'start': 177,\n",
       "  'end': 180},\n",
       " {'entity': 'LABEL_2',\n",
       "  'score': 0.9816441,\n",
       "  'index': 42,\n",
       "  'word': '##de',\n",
       "  'start': 180,\n",
       "  'end': 182},\n",
       " {'entity': 'LABEL_0',\n",
       "  'score': 0.9975062,\n",
       "  'index': 43,\n",
       "  'word': '.',\n",
       "  'start': 183,\n",
       "  'end': 184}]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# something from our validation set\n",
    "effect_ner_model(labeled_dataset[\"test\"][6][\"text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b67731e",
   "metadata": {},
   "source": [
    "---\n",
    "We try out the first few examples of adverse effects from the Wikipedia page on adverse effects and visualize with the displaCy library:\n",
    "\n",
    "https://en.wikipedia.org/wiki/Adverse_effect#Medications"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "872f67db",
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_entities(sentence):\n",
    "    tokens = effect_ner_model(sentence)\n",
    "    entities = []\n",
    "    # ['O', 'B-DRUG', 'I-DRUG', 'B-DISEASE', 'I-DISEASE', 'B-GENE', 'I-GENE']\n",
    "    for token in tokens:\n",
    "        label = int(token[\"entity\"][-1])\n",
    "        if label != 0:\n",
    "            token[\"label\"] = label_list[label]\n",
    "            entities.append(token)\n",
    "    \n",
    "    params = [{\"text\": sentence,\n",
    "               \"ents\": entities,\n",
    "               \"title\": None}]\n",
    "    \n",
    "    html = displacy.render(params, style=\"ent\", manual=True, options={\n",
    "        \"colors\": {\n",
    "                    \n",
    "                   \"B-DRUG\": \"#f08080\",\n",
    "                   \"I-DRUG\": \"#f08080\",\n",
    "                   \"B-DISEASE\": \"#9bddff\",\n",
    "                   \"I-DISEASE\": \"#9bddff\",\n",
    "                    \"B-GENE\": \"#008080\",\n",
    "                   \"I-GENE\": \"#008080\",\n",
    "               },\n",
    "    })\n",
    "    return html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b919e5e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on local URL:  http://127.0.0.1:7860\n",
      "\n",
      "Thanks for being a Gradio user! If you have questions or feedback, please join our Discord server and chat with us: https://discord.gg/feTf9x3ZSB\n",
      "\n",
      "To create a public link, set `share=True` in `launch()`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"http://127.0.0.1:7860/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import gradio as gr\n",
    "\n",
    "exampleList = [\n",
    "    'Famotidine is a histamine H2-receptor antagonist used in inpatient settings for prevention of stress ulcers and is showing increasing popularity because of its low cost.',\n",
    "    'A randomized Phase III trial demonstrated noninferiority of APF530 500 mg SC ( granisetron 10 mg ) to intravenous palonosetron 0.25 mg in preventing CINV in patients receiving MEC or HEC in acute ( 0 - 24 hours ) and delayed ( 24 - 120 hours ) settings , with activity over 120 hours .',\n",
    "    'What are the known interactions between Aspirin and the COX-1 enzyme?',\n",
    "    'Can you explain the mechanism of action of Metformin and its effect on the AMPK pathway?',\n",
    "    'Are there any genetic variations in the CYP2C9 gene that may influence the response to Warfarin therapy?',\n",
    "    'I am curious about the role of Herceptin in targeting the HER2/neu protein in breast cancer treatment. How does it work?',\n",
    "    'What are the common side effects associated with Lisinopril, an angiotensin-converting enzyme (ACE) inhibitor?',\n",
    "    'Can you explain the significance of the BCR-ABL fusion protein in the context of Imatinib therapy for chronic myeloid leukemia (CML)?',\n",
    "    'How does Ibuprofen affect the COX-2 enzyme compared to COX-1?',\n",
    "    'Are there any recent studies exploring the use of Pembrolizumab as an immune checkpoint inhibitor targeting PD-1?',\n",
    "    'I have heard about the SLC6A4 gene and its association with serotonin reuptake inhibitors (SSRIs) like Fluoxetine.',\n",
    "    'Could you provide insights into the BRAF mutation and its relevance in response to Vemurafenib treatment in melanoma patients?'\n",
    "]\n",
    "\n",
    "footer = \"\"\"\n",
    "LLMGeneLinker uses a domain-specific transformer like SciBERT finetuned on AllenAI drug dataset, BC5CDR disease, NCBI disease, DrugProt and GeneTAG datasets. The resulting SciBERT model performs Named Entity Recognition to tag drug, protein, gene, diseases in input text. Sentence embedding of SciBERT is then fed into BERT \n",
    "This was made during the <a target=\"_blank\" href =https://www.sginnovate.com/event/hackathon-large-language-models-bio> LLMs for Bio Hackathon</a> organised by 4Catalyzer and SGInnovate.\n",
    "<br>\n",
    "Made by Team GeneLink (<a target=\"_blank\" href=https://www.linkedin.com/in/ntkb/>Nicholas</a>, <a target=\"_blank\" href=https://www.linkedin.com/in/yewchong-sim/>Yew Chong</a>, <a target=\"_blank\" href=https://www.linkedin.com/in/lim-ting-wei-021383175/>Ting Wei</a>,  <a target=\"_blank\" href=https://www.linkedin.com/in/brendan-lim-ciwen/>Brendan</a>\n",
    "<hr>\n",
    "Note: Performance is noted to be poorer on genes, acronyms, and receptors (named entities that may be targets for drugs or genes)\n",
    "Original notebook adapted from <a target=\"_blank\" href=https://huggingface.co/jsylee/scibert_scivocab_uncased-finetuned-ner>jsylee/scibert_scivocab_uncased-finetuned-ner</a>\n",
    "\"\"\"\n",
    "\n",
    "with gr.Blocks() as demo:\n",
    "    gr.Markdown(\"## LLMGeneLinker (LGL)\")\n",
    "    gr.Markdown(footer)\n",
    "    \n",
    "    txt = gr.Textbox(label=\"Input\", lines=2)\n",
    "    txt_3 = gr.HTML(label=\"Output\")\n",
    "    btn = gr.Button(value=\"Submit\")\n",
    "    btn.click(visualize_entities, inputs=txt, outputs=txt_3)\n",
    "\n",
    "    gr.Markdown(\"## Text Examples\")\n",
    "    gr.Examples(\n",
    "        [[x] for x in exampleList],\n",
    "        txt,\n",
    "        txt_3,\n",
    "        visualize_entities,\n",
    "        cache_examples=False,\n",
    "        run_on_click=True\n",
    "    )\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    demo.launch()"
   ]
  }
 ],
 "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}