examples_token_level_ensemble.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3a47caaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM\n",
    "from ensemble_utils import ensemble_generate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4a659134",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ee5ea11a",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "aa403a06",
   "metadata": {},
   "outputs": [],
   "source": [
    "T5_CLIN_LARGE_PATH = '../../../BioNLP2023/al826/supervised/physionet.org/files/clinical-t5/1.0.0/Clinical-T5-Large'\n",
    "TEST_PATH = '../../../BioNLP2023/released_14.04.2023/bionlp-workshop-2023-shared-task-1a-problem-list-summarization-1.1.0/BioNLP2023-1A-Test.csv'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3924de8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_model(path, device):\n",
    "    model = AutoModelForSeq2SeqLM.from_pretrained(T5_CLIN_LARGE_PATH, return_dict=True)\n",
    "    state = torch.load(path, map_location=device)\n",
    "    model.load_state_dict(state)\n",
    "    model.eval()\n",
    "    model.to(device)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1fdfa1e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an ensemble of three models\n",
    "base_dir = \"../../../BioNLP2023/al826/supervised/experiments/submissions-yf\"\n",
    "ensemble = [\n",
    "    load_model(f\"{base_dir}/models/clint5-large-A-1e5-10e11/seed-3/models/model.pt\", device),\n",
    "    load_model(f\"{base_dir}/models/clint5-large-AS-1e5-10e11/seed-3/models/model.pt\", device),\n",
    "    load_model(f\"{base_dir}/models/clint5-large-AS-1e5-10e11/seed-5/models/model.pt\", device),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d39e124f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load data and create a tokenizer\n",
    "data = pd.read_csv(TEST_PATH)\n",
    "tokenizer = AutoTokenizer.from_pretrained(T5_CLIN_LARGE_PATH, return_dict=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2fd1a3ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run inference on test sample (index=0)\n",
    "\n",
    "idx = 0 \n",
    "\n",
    "# Input = {A} format \n",
    "input_text_A = data.iloc[idx]['Assessment']\n",
    "inputs_A = tokenizer(input_text_A, return_tensors=\"pt\").to(device)\n",
    "\n",
    "# Input = {A}+{S} format \n",
    "input_text1 = data.iloc[idx]['Assessment']\n",
    "input_text2 = data.iloc[idx]['Subjective Sections']\n",
    "input_text_AS = f\"{input_text1} \\\\nSSSS\\\\n{input_text2}\"\n",
    "inputs_AS = tokenizer(input_text_AS, return_tensors=\"pt\").to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1450c50b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Token-level Ensemble Generation\n",
    "summary_ids = ensemble_generate(\n",
    "        ensemble,\n",
    "        [inputs_A.input_ids, inputs_AS.input_ids, inputs_AS.input_ids],\n",
    "        num_beams=4,\n",
    "        length_penalty=0.6,\n",
    "        max_length=256,\n",
    "        min_length=5,\n",
    "        no_repeat_ngram_size=4,\n",
    ")\n",
    "summary_txt = tokenizer.decode(summary_ids.squeeze(), skip_special_tokens=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8f84fe65",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1. Coronary Artery Disease; 2. Epistaxis; 3. Acute Renal Failure; 4. Chronic Obstructive Pulmonary Disease\n"
     ]
    }
   ],
   "source": [
    "print(summary_txt)"
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}