step_1_1_boundary_detector_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "74f84948",
   "metadata": {},
   "source": [
    "## Explore boundary detection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1bdf88da",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "from datasets import load_dataset, Sequence, ClassLabel\n",
    "from enum import Enum\n",
    "from huggingface_hub import Repository\n",
    "from huggingface_hub import get_full_repo_name, notebook_login\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",
    "from transformers import Trainer, TrainingArguments\n",
    "from transformers import get_scheduler\n",
    "\n",
    "import evaluate\n",
    "import numpy as np\n",
    "import os\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6d9fd918",
   "metadata": {},
   "outputs": [],
   "source": [
    "PICO_NER_LABELS = [\n",
    "    0, # O\n",
    "    1, # Outcomes\n",
    "    2, # Interventions\n",
    "    3, # Interventions + Outcomes\n",
    "    4, # Participants\n",
    "    5, # Participants + Outcomes\n",
    "    6, # Pariticpants + Interventions\n",
    "    7, # Pariticpants + Interventions + Outcomes\n",
    "]\n",
    "\n",
    "BOUNDARY_LABELS = [\n",
    "    0, # 'OUT',\n",
    "    1, # 'START',\n",
    "    2, # 'END',\n",
    "    3, # 'BOTH',\n",
    "    4, # 'IN',\n",
    "]\n",
    "\n",
    "class PicoType(Enum):\n",
    "    PARTICIPANTS = 4\n",
    "    INTERVENTIONS = 2\n",
    "    OUTCOMES = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1efef763",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_folder = 'data/bioc/json'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ed5ca4f8",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset json (/home/gzhang/.cache/huggingface/datasets/json/default-50d59fa552f11522/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "67761d5c63814e60a8c0cb5553ea097d",
       "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, 'train.json'),\n",
    "        'validation': os.path.join(input_folder, 'validation.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": "46689bc9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DatasetDict({\n",
      "    train: Dataset({\n",
      "        features: ['pmid', 'tokens', 'labels'],\n",
      "        num_rows: 49031\n",
      "    })\n",
      "    validation: Dataset({\n",
      "        features: ['pmid', 'tokens', 'labels'],\n",
      "        num_rows: 2471\n",
      "    })\n",
      "    test: Dataset({\n",
      "        features: ['pmid', 'tokens', 'labels'],\n",
      "        num_rows: 2042\n",
      "    })\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "print(ebm_nlp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f240148",
   "metadata": {},
   "source": [
    "### Align Labels with Word Piece Tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "75213b18",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at microsoft/BiomedNLP-PubMedBERT-large-uncased-abstract were not used when initializing BertForTokenClassification: ['cls.seq_relationship.weight', 'cls.predictions.decoder.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.seq_relationship.bias', 'cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.decoder.bias']\n",
      "- This IS expected if you are initializing BertForTokenClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForTokenClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForTokenClassification were not initialized from the model checkpoint at microsoft/BiomedNLP-PubMedBERT-large-uncased-abstract and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "model_checkpoint = 'microsoft/BiomedNLP-PubMedBERT-large-uncased-abstract'\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)\n",
    "\n",
    "id2label = {i: label for i, label in enumerate(BOUNDARY_LABELS)}\n",
    "label2id = {v: k for k, v in id2label.items()}\n",
    "\n",
    "model = AutoModelForTokenClassification.from_pretrained(\n",
    "    model_checkpoint,\n",
    "    id2label=id2label,\n",
    "    label2id=label2id,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8f9ec844",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original input has 21 words.\n",
      "Tokenized input has 29 tokens.\n",
      "\n",
      "['[CLS]', 'similarly', ',', 'post', '-', 'operative', 'increments', 'in', 'urinary', 'excretion', 'of', 'ammonia', ',', 'creatinine', 'and', '3', '-', 'methyl', '##his', '##tidine', 'were', 'not', 'altered', 'by', 'addition', 'of', 'insulin', '.', '[SEP]']\n",
      "[None, 0, 1, 2, 2, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 12, 12, 12, 12, 13, 14, 15, 16, 17, 18, 19, 20, None]\n",
      "\n",
      "Term \"post-operative\" and \"3-methylhistidine\" are split into word pieces.\n"
     ]
    }
   ],
   "source": [
    "'''\n",
    "Check the pretrained tokenizer's output.\n",
    "\n",
    "Example sentence from PMID 6420374:\n",
    "\n",
    "\"Similarly, post-operative increments in urinary excretion of ammonia, \"\n",
    "\"creatinine and 3-methylhistidine were not altered by addition of insulin.\"\n",
    "'''\n",
    "\n",
    "input_words = [\n",
    "    'Similarly', ',', 'post-operative', 'increments', 'in', 'urinary', 'excretion', 'of', 'ammonia', ',',\n",
    "    'creatinine', 'and', '3-methylhistidine', 'were', 'not', 'altered', 'by', 'addition', 'of', 'insulin', '.',\n",
    "]\n",
    "outcome_labels = [\n",
    "    0, 0, 0, 0, 0, 1, 1, 1, 1, 1,\n",
    "    1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0,\n",
    "]\n",
    "print('Original input has {} words.'.format(len(input_words)))\n",
    "\n",
    "tokenized_input = tokenizer(input_words, is_split_into_words=True)\n",
    "print('Tokenized input has {} tokens.\\n'.format(len(tokenized_input['input_ids'])))\n",
    "print(tokenized_input.tokens())\n",
    "print(tokenized_input.word_ids())\n",
    "\n",
    "print('\\nTerm \"{}\" and \"{}\" are split into word pieces.'.format(input_words[2], input_words[12]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f8955e18",
   "metadata": {},
   "outputs": [],
   "source": [
    "def start_appears(a, b):\n",
    "    for p in list(PicoType):\n",
    "        if (a & p.value) < (b & p.value):\n",
    "            return True\n",
    "    return False\n",
    "\n",
    "def end_appears(a, b):\n",
    "    for p in list(PicoType):\n",
    "        if (a & p.value) > (b & p.value):\n",
    "            return True\n",
    "    return False\n",
    "\n",
    "def extract_boundary_labels(labels):\n",
    "    boundary_labels = []\n",
    "    for i, l in enumerate(labels):\n",
    "        b = 0\n",
    "        if l > 0:\n",
    "            start = (i == 0) or start_appears(labels[i-1], l)\n",
    "            end = (i == len(labels) -1) or end_appears(l, labels[i+1])\n",
    "            if start:\n",
    "                b = b | 1\n",
    "            if end:\n",
    "                b = b | 2\n",
    "            if not start and not end:\n",
    "                b = 4\n",
    "            boundary_labels.append(b)\n",
    "        else:\n",
    "            boundary_labels.append(l)\n",
    "    return boundary_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a0c24c53",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "Aligns labels with tokenized inputs.\n",
    "\n",
    "Pre-trained tokenizers in transformers may break a single word into word pieces.\n",
    "The labels should be aligned with the tokenized outputs. Inserted special tokens\n",
    "are given -100.\n",
    "\n",
    "Code modified from https://huggingface.co/course/chapter7/2.\n",
    "\n",
    "Parameters:\n",
    "    labels (List(int)): labels for each words before transformer tokenization.\n",
    "    word_ids (List(int)): The id of the word where the word piece comes from.\n",
    "\n",
    "Returns:\n",
    "    The list of labels after alignment.\n",
    "'''\n",
    "\n",
    "def align_labels_with_tokens(labels, word_ids):\n",
    "    new_labels = []\n",
    "    current_word = None\n",
    "    for word_id in word_ids:\n",
    "        if word_id != current_word:\n",
    "            # Start of a new word!\n",
    "            current_word = word_id\n",
    "            label = -100 if word_id is None else labels[word_id]\n",
    "            new_labels.append(label)\n",
    "        elif word_id is None:\n",
    "            # Special token\n",
    "            new_labels.append(-100)\n",
    "        else:\n",
    "            # Same word as previous token\n",
    "            label = labels[word_id]\n",
    "            new_labels.append(label)\n",
    "\n",
    "    return new_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "10111b1e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "[0, 0, 0, 0, 0, 1, 4, 4, 4, 4, 4, 4, 2, 0, 0, 0, 0, 0, 0, 0, 0]\n",
      "[-100, 0, 0, 0, 0, 0, 0, 0, 1, 4, 4, 4, 4, 4, 4, 2, 2, 2, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, -100]\n",
      "[None, 0, 1, 2, 2, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 12, 12, 12, 12, 13, 14, 15, 16, 17, 18, 19, 20, None]\n"
     ]
    }
   ],
   "source": [
    "# Test label alignment on the example.\n",
    "print(outcome_labels)\n",
    "\n",
    "boundary_labels = extract_boundary_labels(outcome_labels)\n",
    "print(boundary_labels)\n",
    "\n",
    "aligned_labels = align_labels_with_tokens(boundary_labels, tokenized_input.word_ids())\n",
    "print(aligned_labels)\n",
    "\n",
    "print(tokenized_input.word_ids())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "757f6145",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/gzhang/.local/lib/python3.10/site-packages/dill/_dill.py:1705: PicklingWarning: Cannot locate reference to <enum 'PicoType'>.\n",
      "  warnings.warn('Cannot locate reference to %r.' % (obj,), PicklingWarning)\n",
      "/home/gzhang/.local/lib/python3.10/site-packages/dill/_dill.py:1707: PicklingWarning: Cannot pickle <enum 'PicoType'>: __main__.PicoType has recursive self-references that trigger a RecursionError.\n",
      "  warnings.warn('Cannot pickle %r: %s.%s has recursive self-references that trigger a RecursionError.' % (obj, obj.__module__, obj_name), PicklingWarning)\n",
      "Loading cached processed dataset at /home/gzhang/.cache/huggingface/datasets/json/default-50d59fa552f11522/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4/cache-9acfed8fad3625a9.arrow\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2471 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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",
      "Loading cached processed dataset at /home/gzhang/.cache/huggingface/datasets/json/default-50d59fa552f11522/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4/cache-750b4ab9381db7e0.arrow\n",
      "Loading cached processed dataset at /home/gzhang/.cache/huggingface/datasets/json/default-50d59fa552f11522/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4/cache-5f33e300656b3fa1.arrow\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Casting the dataset:   0%|          | 0/2471 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /home/gzhang/.cache/huggingface/datasets/json/default-50d59fa552f11522/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4/cache-5633caf09f03cc31.arrow\n"
     ]
    }
   ],
   "source": [
    "'''Tokenizes input and align tokens with labels in a batch.'''\n",
    "def tokenize_and_align_labels(dataset):\n",
    "    tokenized_inputs = tokenizer(\n",
    "        dataset['tokens'],\n",
    "        truncation=True,\n",
    "        is_split_into_words=True,\n",
    "    )\n",
    "    all_labels = dataset['labels']\n",
    "    new_labels = []\n",
    "    word_ids_list = []\n",
    "    for i, labels in enumerate(all_labels):\n",
    "        word_ids = tokenized_inputs.word_ids(i)\n",
    "        new_labels.append(\n",
    "            align_labels_with_tokens(\n",
    "                extract_boundary_labels(labels), \n",
    "                word_ids\n",
    "            ))\n",
    "        word_ids_list.append(word_ids)\n",
    "\n",
    "    tokenized_inputs['labels'] = new_labels\n",
    "    tokenized_inputs['word_ids'] = word_ids_list\n",
    "    return tokenized_inputs\n",
    "\n",
    "tokenized_dataset = ebm_nlp.map(\n",
    "    tokenize_and_align_labels,\n",
    "    batched=True,\n",
    "    remove_columns=['pmid', 'tokens'],\n",
    ")\n",
    "\n",
    "tokenized_dataset = tokenized_dataset.cast_column(\n",
    "    'labels',\n",
    "    Sequence(ClassLabel(names = BOUNDARY_LABELS))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "661d8616",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['labels', 'input_ids', 'token_type_ids', 'attention_mask', 'word_ids'],\n",
       "        num_rows: 49031\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['labels', 'input_ids', 'token_type_ids', 'attention_mask', 'word_ids'],\n",
       "        num_rows: 2471\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['labels', 'input_ids', 'token_type_ids', 'attention_mask', 'word_ids'],\n",
       "        num_rows: 2042\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenized_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37b06836",
   "metadata": {},
   "source": [
    "### Fine tune models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "20cd01a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b2a2c844",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_start = datetime.now(tz = timezone('US/Eastern'))\n",
    "task = 'PICO_NER'\n",
    "dataset_name = 'ebm_nlp_bioc'\n",
    "model_name = 'boundaries-{}-{}-{}'.format(\n",
    "    task,\n",
    "    dataset_name,\n",
    "    datetime.now(timezone('US/Eastern')).strftime('%Y_%m_%d_%H_%M_%S_%Z')\n",
    ")\n",
    "\n",
    "data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)\n",
    "output_dir = os.path.join('pico_span/boundary_models', model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "00fc7e7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "args = TrainingArguments(\n",
    "    output_dir,\n",
    "    evaluation_strategy='epoch',\n",
    "    save_strategy='epoch',\n",
    "    learning_rate=5e-5,\n",
    "    num_train_epochs=3,\n",
    "    weight_decay=0.01,\n",
    "    push_to_hub=False,\n",
    "    do_predict=True,\n",
    ")\n",
    "\n",
    "\n",
    "token_precision_metric = evaluate.load('precision')\n",
    "token_recall_metric = evaluate.load('recall')\n",
    "token_f1_metric = evaluate.load('f1')\n",
    "\n",
    "\n",
    "'''Calculates precision, recall and F1 scores.'''\n",
    "def compute_metrics(eval_preds):\n",
    "    logits, labels = eval_preds\n",
    "    predictions = np.argmax(logits, axis=-1)\n",
    "\n",
    "    # Remove ignored index (special tokens) and convert to labels\n",
    "    decoded_labels = [\n",
    "        [l for l in label if l != -100] \n",
    "        for label in labels\n",
    "    ]\n",
    "    \n",
    "    decoded_predictions = [\n",
    "        [p for (p, l) in zip(prediction, label) if l != -100]\n",
    "        for prediction, label in zip(predictions, labels)\n",
    "    ]\n",
    "#     print(decoded_labels, decoded_predictions)\n",
    "    \n",
    "    # Token level\n",
    "    flat_labels = [l for dl in decoded_labels for l in dl]\n",
    "    flat_predictions = [p for dp in decoded_predictions for p in dp]\n",
    "    \n",
    "    token_precision = token_precision_metric.compute(\n",
    "        predictions=flat_predictions,\n",
    "        references=flat_labels,\n",
    "        average='macro',\n",
    "    )\n",
    "    token_recall = token_recall_metric.compute(\n",
    "        predictions=flat_predictions,\n",
    "        references=flat_labels,\n",
    "        average='macro',\n",
    "    )\n",
    "    token_f1 = token_f1_metric.compute(\n",
    "        predictions=flat_predictions,\n",
    "        references=flat_labels,\n",
    "        average='macro',\n",
    "    )\n",
    "    \n",
    "    start_tp, start_fp, start_fn = 0, 0, 0\n",
    "    end_tp, end_fp, end_fn = 0, 0, 0\n",
    "    for label, pred in zip(flat_labels, flat_predictions):\n",
    "        if label == 0 or label == 4:\n",
    "            if pred == 1:\n",
    "                start_fp += 1\n",
    "            elif pred == 2:\n",
    "                end_fp += 1\n",
    "            elif pred == 3:\n",
    "                start_fp += 1\n",
    "                end_fp += 1\n",
    "                \n",
    "        elif label == 1:\n",
    "            if pred == 0 or pred == 4:\n",
    "                start_fn += 1\n",
    "            elif pred == 1:\n",
    "                start_tp += 1\n",
    "            elif pred == 2:\n",
    "                start_fn += 1\n",
    "                end_fp += 1\n",
    "            elif pred == 3:\n",
    "                start_tp += 1\n",
    "                end_fp += 1\n",
    "    \n",
    "        elif label == 2:\n",
    "            if pred == 0 or pred == 4:\n",
    "                end_fn += 1\n",
    "            elif pred == 1:\n",
    "                start_fp += 1\n",
    "                end_fn += 1\n",
    "            elif pred == 2:\n",
    "                end_tp += 1\n",
    "            elif pred == 3:\n",
    "                start_fp += 1\n",
    "                end_tp += 1\n",
    "                \n",
    "        elif label == 3:\n",
    "            if pred == 0 or pred == 4:\n",
    "                start_fn += 1\n",
    "                end_fn += 1\n",
    "            elif pred == 1:\n",
    "                start_tp += 1\n",
    "                end_fp += 1\n",
    "            elif pred == 2:\n",
    "                start_fp += 1\n",
    "                end_tp += 1\n",
    "            elif pred == 3:\n",
    "                start_tp += 1\n",
    "                end_tp += 1\n",
    "            \n",
    "    start_precision = start_tp / (start_tp + start_fp)\n",
    "    start_recall = start_tp / (start_tp + start_fn)\n",
    "    start_f1 = 2 * start_precision * start_recall / (start_precision + start_recall) if start_tp else 0\n",
    "    \n",
    "    end_precision = end_tp / (end_tp + end_fp)\n",
    "    end_recall = end_tp / (end_tp + end_fn)\n",
    "    end_f1 = 2 * end_precision * end_recall / (end_precision + end_recall) if end_tp else 0\n",
    "    \n",
    "    return {\n",
    "        'overall_precision': token_precision['precision'],\n",
    "        'overall_recall': token_recall['recall'],\n",
    "        'overall_f1': token_f1['f1'],\n",
    "\n",
    "        'start_precision': start_precision,\n",
    "        'start_recall': start_recall,\n",
    "        'start_f1': start_f1,\n",
    "        \n",
    "        'end_precision': end_precision,\n",
    "        'end_recall': end_recall,\n",
    "        'end_f1': end_f1,\n",
    "    }\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=tokenized_dataset['train'],\n",
    "#     eval_dataset=tokenized_dataset['train'],\n",
    "    eval_dataset=tokenized_dataset['validation'],\n",
    "    data_collator=data_collator,\n",
    "    compute_metrics=compute_metrics,\n",
    "    tokenizer=tokenizer,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ca832b74",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/gzhang/.local/lib/python3.10/site-packages/transformers/optimization.py:391: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n",
      "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",
      "/home/gzhang/.local/lib/python3.10/site-packages/torch/nn/parallel/_functions.py:68: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='4599' max='4599' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [4599/4599 28:07, Epoch 3/3]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Overall Precision</th>\n",
       "      <th>Overall Recall</th>\n",
       "      <th>Overall F1</th>\n",
       "      <th>Start Precision</th>\n",
       "      <th>Start Recall</th>\n",
       "      <th>Start F1</th>\n",
       "      <th>End Precision</th>\n",
       "      <th>End Recall</th>\n",
       "      <th>End F1</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.482400</td>\n",
       "      <td>0.489319</td>\n",
       "      <td>0.664416</td>\n",
       "      <td>0.648209</td>\n",
       "      <td>0.655783</td>\n",
       "      <td>0.610699</td>\n",
       "      <td>0.596283</td>\n",
       "      <td>0.603404</td>\n",
       "      <td>0.616008</td>\n",
       "      <td>0.628420</td>\n",
       "      <td>0.622152</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.401500</td>\n",
       "      <td>0.501913</td>\n",
       "      <td>0.676164</td>\n",
       "      <td>0.644614</td>\n",
       "      <td>0.658917</td>\n",
       "      <td>0.620579</td>\n",
       "      <td>0.575932</td>\n",
       "      <td>0.597423</td>\n",
       "      <td>0.626598</td>\n",
       "      <td>0.619078</td>\n",
       "      <td>0.622815</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.306500</td>\n",
       "      <td>0.566453</td>\n",
       "      <td>0.668451</td>\n",
       "      <td>0.658631</td>\n",
       "      <td>0.663079</td>\n",
       "      <td>0.609404</td>\n",
       "      <td>0.604198</td>\n",
       "      <td>0.606790</td>\n",
       "      <td>0.616939</td>\n",
       "      <td>0.638572</td>\n",
       "      <td>0.627569</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/gzhang/.local/lib/python3.10/site-packages/torch/nn/parallel/_functions.py:68: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/home/gzhang/.local/lib/python3.10/site-packages/torch/nn/parallel/_functions.py:68: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n"
     ]
    }
   ],
   "source": [
    "trainer.train()\n",
    "\n",
    "save_path = 'pico_span/boundary_models'\n",
    "model.save_pretrained(os.path.join(save_path, model_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc8d752c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# trainer"
   ]
  }
 ],
 "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
}