llama3.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4fc64f241826411e9e8e4c213ddfaf66",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "import torch\n",
    "import transformers\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from sklearn.metrics import f1_score\n",
    "from utils import html_parsing_ncbi, html_parsing_n2c2, get_classification_report, get_digit, get_macro_average_f1\n",
    "\n",
    "model_id = \"meta-llama/Meta-Llama-3-70B-Instruct\"\n",
    "# This will take 45GB of GPU memory loading in 4-bit precision\n",
    "pipeline = transformers.pipeline(\n",
    "    \"text-generation\",\n",
    "    model=model_id,\n",
    "    model_kwargs={\"torch_dtype\": torch.float16, \n",
    "                  \"attn_implementation\": \"flash_attention_2\",\n",
    "                  \"quantization_config\": {\"load_in_4bit\": True, \"bnb_4bit_compute_dtype\": torch.float16},\n",
    "                  \"low_cpu_mem_usage\": True},\n",
    "    token=\"hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. NER (Named Entity Recognition)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.1 NCBI-Disease Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1.1 Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "ncbi_df = pd.read_csv('data/NER/NCBI-disease/test_200.csv')\n",
    "ncbi_example_df = pd.read_csv('data/NER/NCBI-disease/examples.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_message = \"\"\"You are a helpful assistant to perform the following task.\n",
    "\"TASK: the task is to extract disease entities in a sentence.\"\n",
    "\"INPUT: the input is a sentence.\"\n",
    "\"OUTPUT: the output is an HTML that highlights all the disease entities in the sentence. The highlighting should only use HTML tags <span style=\\\"background-color: #FFFF00\\\"> and </span> and no other tags.\"\n",
    "\"\"\"\n",
    "\n",
    "def get_ner_ncbi_disease(sentence: str, shot: int = 0) -> str:\n",
    "    \"\"\"\n",
    "    Get the NER results of NCBI-disease dataset from few-shot prompting.\n",
    "    Args:\n",
    "        sentence: the input sentence\n",
    "        shot: the number of few-shot examples\n",
    "    Returns:\n",
    "        the NER results\n",
    "    \"\"\"\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": system_message},\n",
    "    ]\n",
    "    for i in range(shot):\n",
    "        messages.append({\"role\": \"user\", \"content\": ncbi_example_df. iloc[i]['text']}) \n",
    "        messages.append({\"role\": \"assistant\", \"content\": ncbi_example_df.iloc[i]['label_text']})\n",
    "    messages.append({\"role\": \"user\", \"content\": sentence})\n",
    "    prompt = pipeline.tokenizer.apply_chat_template(\n",
    "            messages, \n",
    "            tokenize=False, \n",
    "            add_generation_prompt=True\n",
    "    )\n",
    "    terminators = [\n",
    "        pipeline.tokenizer.eos_token_id,\n",
    "        pipeline.tokenizer.convert_tokens_to_ids(\"<|eot_id|>\")\n",
    "    ]\n",
    "\n",
    "    time_start = time.time()\n",
    "    outputs = pipeline(\n",
    "        prompt,\n",
    "        max_new_tokens=2048,\n",
    "        eos_token_id=terminators,\n",
    "        do_sample=False,\n",
    "        pad_token_id=pipeline.tokenizer.eos_token_id\n",
    "    )\n",
    "    time_end = time.time()\n",
    "\n",
    "    return outputs[0][\"generated_text\"][len(prompt):], time_end - time_start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|          | 2/200 [00:43<1:07:49, 20.55s/it]You seem to be using the pipelines sequentially on GPU. In order to maximize efficiency please use a dataset\n",
      "100%|██████████| 200/200 [1:29:00<00:00, 26.70s/it]\n"
     ]
    }
   ],
   "source": [
    "for i in tqdm(range(0, len(ncbi_df), 1)):\n",
    "    ncbi_df.loc[i, 'html_llama3_70b_instruct_one_shot'], ncbi_df.loc[i, 'llama3_70b_instruct_one_shot_time'] = get_ner_ncbi_disease(ncbi_df.loc[i, 'text'], 1)\n",
    "    ncbi_df.loc[i, 'html_llama3_70b_instruct_five_shot'], ncbi_df.loc[i, 'llama3_70b_instruct_five_shot_time'] = get_ner_ncbi_disease(ncbi_df.loc[i, 'text'], 5)\n",
    "    ncbi_df.loc[i, 'html_llama3_70b_instruct_ten_shot'], ncbi_df.loc[i, 'llama3_70b_instruct_ten_shot_time'] = get_ner_ncbi_disease(ncbi_df.loc[i, 'text'], 10)\n",
    "    ncbi_df.loc[i, 'html_llama3_70b_instruct_twenty_shot'], ncbi_df.loc[i, 'llama3_70b_instruct_twenty_shot_time'] = get_ner_ncbi_disease(ncbi_df.loc[i, 'text'], 20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1.2 Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional: you can just load the llm output from the csv file instead of running the above code\n",
    "# ncbi_df = pd.read_csv(\"data/NER/NCBI-disease/test_200_llama3_70b_instruct_results.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "ncbi_df['gt_labels'], ncbi_df['llama3_70b_instruct_one_shot_labels'] = html_parsing_ncbi(ncbi_df, 'html_llama3_70b_instruct_one_shot')\n",
    "_, ncbi_df['llama3_70b_instruct_five_shot_labels'] = html_parsing_ncbi(ncbi_df, 'html_llama3_70b_instruct_five_shot')\n",
    "_, ncbi_df['llama3_70b_instruct_ten_shot_labels'] = html_parsing_ncbi(ncbi_df, 'html_llama3_70b_instruct_ten_shot')\n",
    "_, ncbi_df['llama3_70b_instruct_twenty_shot_labels'] = html_parsing_ncbi(ncbi_df, 'html_llama3_70b_instruct_twenty_shot')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F1-Score One Shot (Strict): 0.6214953271028039\n",
      "F1-Score Five Shot (Strict): 0.6373333333333333\n",
      "F1-Score Ten Shot (Strict): 0.6628895184135978\n",
      "F1-Score Twenty Shot (Strict): 0.685459940652819\n"
     ]
    }
   ],
   "source": [
    "print(f\"F1-Score One Shot (Strict): {get_classification_report(ncbi_df, 'gt_labels', 'llama3_70b_instruct_one_shot_labels', 'strict')['default']['f1-score']}\")\n",
    "print(f\"F1-Score Five Shot (Strict): {get_classification_report(ncbi_df, 'gt_labels', 'llama3_70b_instruct_five_shot_labels', 'strict')['default']['f1-score']}\")\n",
    "print(f\"F1-Score Ten Shot (Strict): {get_classification_report(ncbi_df, 'gt_labels', 'llama3_70b_instruct_ten_shot_labels', 'strict')['default']['f1-score']}\")\n",
    "print(f\"F1-Score Twenty Shot (Strict): {get_classification_report(ncbi_df, 'gt_labels', 'llama3_70b_instruct_twenty_shot_labels', 'strict')['default']['f1-score']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F1-Score One Shot (Lenient): 0.7429906542056075\n",
      "F1-Score Five Shot (Lenient): 0.7573333333333333\n",
      "F1-Score Ten Shot (Lenient): 0.7790368271954673\n",
      "F1-Score Twenty Shot (Lenient): 0.7863501483679526\n"
     ]
    }
   ],
   "source": [
    "print(f\"F1-Score One Shot (Lenient): {get_classification_report(ncbi_df, 'gt_labels', 'llama3_70b_instruct_one_shot_labels', 'lenient')['default']['f1-score']}\")\n",
    "print(f\"F1-Score Five Shot (Lenient): {get_classification_report(ncbi_df, 'gt_labels', 'llama3_70b_instruct_five_shot_labels', 'lenient')['default']['f1-score']}\")\n",
    "print(f\"F1-Score Ten Shot (Lenient): {get_classification_report(ncbi_df, 'gt_labels', 'llama3_70b_instruct_ten_shot_labels', 'lenient')['default']['f1-score']}\")\n",
    "print(f\"F1-Score Twenty Shot (Lenient): {get_classification_report(ncbi_df, 'gt_labels', 'llama3_70b_instruct_twenty_shot_labels', 'lenient')['default']['f1-score']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Llama-3-70B-Instruct one-shot prediction time: 7.40 seconds\n",
      "Average Llama-3-70B-Instruct five-shot prediction time: 6.59 seconds\n",
      "Average Llama-3-70B-Instruct ten-shot prediction time: 6.41 seconds\n",
      "Average Llama-3-70B-Instruct twenty-shot prediction time: 6.29 seconds\n"
     ]
    }
   ],
   "source": [
    "print(f\"Average Llama-3-70B-Instruct one-shot prediction time: {ncbi_df['llama3_70b_instruct_one_shot_time'].mean():.2f} seconds\")\n",
    "print(f\"Average Llama-3-70B-Instruct five-shot prediction time: {ncbi_df['llama3_70b_instruct_five_shot_time'].mean():.2f} seconds\")\n",
    "print(f\"Average Llama-3-70B-Instruct ten-shot prediction time: {ncbi_df['llama3_70b_instruct_ten_shot_time'].mean():.2f} seconds\")\n",
    "print(f\"Average Llama-3-70B-Instruct twenty-shot prediction time: {ncbi_df['llama3_70b_instruct_twenty_shot_time'].mean():.2f} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the inference results\n",
    "ncbi_df.to_csv('data/NER/NCBI-disease/test_200_llama3_70b_instruct_results.csv', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1.2 2018 n2c2 Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2.1 Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "n2c2_df = pd.read_csv('data/NER/2018_n2c2/test_200.csv')\n",
    "n2c2_example_df = pd.read_csv('data/NER/2018_n2c2/examples.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_message = \"\"\"You are a helpful assistant to perform the following task.\n",
    "\"TASK: the task is to extract disease entities in a sentence. The entity type includes Form, Route, Frequency, Dosage, Strength, Duration, Reason, Ade, Drug.\"\n",
    "\"INPUT: the input is a sentence.\"\n",
    "\"OUTPUT: the output is an HTML that highlights all the disease entities in the sentence in different colors: Form(#FF0000), Route(#FFA500), Frequency(#FFFF00), Dosage(#00FF00), Strength(#0000FF), Duration(#800080), Reason(#FFC0CB), Ade(#964B00), Drug(#808080) in hex code. The highlighting should only use HTML tags <span style=\\\"background-color: #XXXXXX\\\"> and </span> and no other tags.\"\n",
    "\"\"\"\n",
    "def get_ner_2018_n2c2(sentence: str, shot: int = 0) -> str:\n",
    "    \"\"\"\n",
    "    Get the NER results of 2018 n2c2 dataset from few-shot prompting.\n",
    "    Args:\n",
    "        sentence: the input sentence\n",
    "        shot: the number of few-shot examples\n",
    "    Returns:\n",
    "        the NER results\n",
    "    \"\"\"\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": system_message},\n",
    "    ]\n",
    "    for i in range(shot):\n",
    "        messages.append({\"role\": \"user\", \"content\": n2c2_example_df. iloc[i]['text']}) \n",
    "        messages.append({\"role\": \"assistant\", \"content\": n2c2_example_df.iloc[i]['label_text']})\n",
    "    messages.append({\"role\": \"user\", \"content\": sentence})\n",
    "    prompt = pipeline.tokenizer.apply_chat_template(\n",
    "            messages, \n",
    "            tokenize=False, \n",
    "            add_generation_prompt=True\n",
    "    )\n",
    "    terminators = [\n",
    "        pipeline.tokenizer.eos_token_id,\n",
    "        pipeline.tokenizer.convert_tokens_to_ids(\"<|eot_id|>\")\n",
    "    ]\n",
    "\n",
    "    time_start = time.time()\n",
    "    outputs = pipeline(\n",
    "        prompt,\n",
    "        max_new_tokens=2048,\n",
    "        eos_token_id=terminators,\n",
    "        do_sample=False,\n",
    "        pad_token_id=pipeline.tokenizer.eos_token_id\n",
    "    )\n",
    "    time_end = time.time()\n",
    "\n",
    "    return outputs[0][\"generated_text\"][len(prompt):], time_end - time_start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [2:29:40<00:00, 44.90s/it]  \n"
     ]
    }
   ],
   "source": [
    "for i in tqdm(range(0, len(n2c2_df), 1)):\n",
    "    n2c2_df.loc[i, 'html_llama3_70b_instruct_one_shot'], n2c2_df.loc[i, 'llama3_70b_instruct_one_shot_time'] = get_ner_2018_n2c2(n2c2_df.loc[i, 'text'], 1)\n",
    "    n2c2_df.loc[i, 'html_llama3_70b_instruct_five_shot'], n2c2_df.loc[i, 'llama3_70b_instruct_five_shot_time'] = get_ner_2018_n2c2(n2c2_df.loc[i, 'text'], 5)\n",
    "    n2c2_df.loc[i, 'html_llama3_70b_instruct_ten_shot'], n2c2_df.loc[i, 'llama3_70b_instruct_ten_shot_time'] = get_ner_2018_n2c2(n2c2_df.loc[i, 'text'], 10)\n",
    "    n2c2_df.loc[i, 'html_llama3_70b_instruct_twenty_shot'], n2c2_df.loc[i, 'llama3_70b_instruct_twenty_shot_time'] = get_ner_2018_n2c2(n2c2_df.loc[i, 'text'], 20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2.2 Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional: you can just load the llm output from the csv file instead of running the above code\n",
    "n2c2_df = pd.read_csv(\"data/NER/2018_n2c2/test_200_llama3_70b_instruct_results.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "n2c2_df['gt_labels'], n2c2_df['llama3_70b_instruct_one_shot_labels'] = html_parsing_n2c2(n2c2_df, 'html_llama3_70b_instruct_one_shot')\n",
    "_, n2c2_df['llama3_70b_instruct_five_shot_labels'] = html_parsing_n2c2(n2c2_df, 'html_llama3_70b_instruct_five_shot')\n",
    "_, n2c2_df['llama3_70b_instruct_ten_shot_labels'] = html_parsing_n2c2(n2c2_df, 'html_llama3_70b_instruct_ten_shot')\n",
    "_, n2c2_df['llama3_70b_instruct_twenty_shot_labels'] = html_parsing_n2c2(n2c2_df, 'html_llama3_70b_instruct_twenty_shot')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F1 Score One Shot (Strict): 0.1814137367727111\n",
      "F1 Score Five Shot (Strict): 0.46791731342399345\n",
      "F1 Score Ten Shot (Strict): 0.4989763604796711\n",
      "F1 Score Twenty Shot (Strict): 0.5505756039667599\n"
     ]
    }
   ],
   "source": [
    "print(f\"F1 Score One Shot (Strict): {get_macro_average_f1(get_classification_report(n2c2_df, 'gt_labels', 'llama3_70b_instruct_one_shot_labels', 'strict'))}\")\n",
    "print(f\"F1 Score Five Shot (Strict): {get_macro_average_f1(get_classification_report(n2c2_df, 'gt_labels', 'llama3_70b_instruct_five_shot_labels', 'strict'))}\")\n",
    "print(f\"F1 Score Ten Shot (Strict): {get_macro_average_f1(get_classification_report(n2c2_df, 'gt_labels', 'llama3_70b_instruct_ten_shot_labels', 'strict'))}\")\n",
    "print(f\"F1 Score Twenty Shot (Strict): {get_macro_average_f1(get_classification_report(n2c2_df, 'gt_labels', 'llama3_70b_instruct_twenty_shot_labels', 'strict'))}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F1 Score One Shot (Lenient): 0.2736644463557607\n",
      "F1 Score Five Shot (Lenient): 0.5903324689808858\n",
      "F1 Score Ten Shot (Lenient): 0.6261996528827135\n",
      "F1 Score Twenty Shot (Lenient): 0.6952576079043253\n"
     ]
    }
   ],
   "source": [
    "print(f\"F1 Score One Shot (Lenient): {get_macro_average_f1(get_classification_report(n2c2_df, 'gt_labels', 'llama3_70b_instruct_one_shot_labels', 'lenient'))}\")\n",
    "print(f\"F1 Score Five Shot (Lenient): {get_macro_average_f1(get_classification_report(n2c2_df, 'gt_labels', 'llama3_70b_instruct_five_shot_labels', 'lenient'))}\")\n",
    "print(f\"F1 Score Ten Shot (Lenient): {get_macro_average_f1(get_classification_report(n2c2_df, 'gt_labels', 'llama3_70b_instruct_ten_shot_labels', 'lenient'))}\")\n",
    "print(f\"F1 Score Twenty Shot (Lenient): {get_macro_average_f1(get_classification_report(n2c2_df, 'gt_labels', 'llama3_70b_instruct_twenty_shot_labels', 'lenient'))}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Llama-3-70B-Instruct one-shot prediction time: 9.75 seconds\n",
      "Average Llama-3-70B-Instruct five-shot prediction time: 11.77 seconds\n",
      "Average Llama-3-70B-Instruct ten-shot prediction time: 11.09 seconds\n",
      "Average Llama-3-70B-Instruct twenty-shot prediction time: 12.29 seconds\n"
     ]
    }
   ],
   "source": [
    "print(f\"Average Llama-3-70B-Instruct one-shot prediction time: {n2c2_df['llama3_70b_instruct_one_shot_time'].mean():.2f} seconds\")\n",
    "print(f\"Average Llama-3-70B-Instruct five-shot prediction time: {n2c2_df['llama3_70b_instruct_five_shot_time'].mean():.2f} seconds\")\n",
    "print(f\"Average Llama-3-70B-Instruct ten-shot prediction time: {n2c2_df['llama3_70b_instruct_ten_shot_time'].mean():.2f} seconds\")\n",
    "print(f\"Average Llama-3-70B-Instruct twenty-shot prediction time: {n2c2_df['llama3_70b_instruct_twenty_shot_time'].mean():.2f} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the inference results\n",
    "n2c2_df.to_csv('data/NER/2018_n2c2/test_200_llama3_70b_instruct_results.csv', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. RE (Relation Extraction)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.1 2018 n2c2 Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1.1 Infernece"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "n2c2_df = pd.read_csv('data/RE/2018_n2c2/test_200.csv')\n",
    "n2c2_example_df = pd.read_csv('data/RE/2018_n2c2/examples.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_message = \"\"\"You are a helpful assistant to perform the following task.\n",
    "\"TASK: the task is to classify relations for a sentence.\"\n",
    "\"INPUT: the input is a sentence where the entities are labeled within [E${X}] and [E${X}/] in a sentence, where X is an integer representing an unique entity.\"\n",
    "\"OUTPUT: your task is to select one out of the nine types of relations ('STRENGTH-DRUG', 'ROUTE-DRUG', 'FREQUENCY-DRUG', 'FORM-DRUG', 'DOSAGE-DRUG', 'REASON-DRUG', 'DURATION-DRUG', 'ADE-DRUG', and 'No relation').\"\n",
    "\"\"\"\n",
    "def get_re_2018_n2c2(sentence: str, shot: int = 0) -> str:\n",
    "    \"\"\"\n",
    "    Get the RE results of 2018 n2c2 dataset from few-shot prompting.\n",
    "    Args:\n",
    "        sentence: the input sentence\n",
    "        shot: the number of few-shot examples\n",
    "    Returns:\n",
    "        the RE results\n",
    "    \"\"\"\n",
    "    \n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": system_message},\n",
    "    ]\n",
    "    for i in range(shot):\n",
    "        messages.append({\"role\": \"user\", \"content\": n2c2_example_df. iloc[i]['text']}) \n",
    "        messages.append({\"role\": \"assistant\", \"content\": n2c2_example_df.iloc[i]['labels']})\n",
    "    messages.append({\"role\": \"user\", \"content\": sentence})\n",
    "    prompt = pipeline.tokenizer.apply_chat_template(\n",
    "            messages, \n",
    "            tokenize=False, \n",
    "            add_generation_prompt=True\n",
    "    )\n",
    "    terminators = [\n",
    "        pipeline.tokenizer.eos_token_id,\n",
    "        pipeline.tokenizer.convert_tokens_to_ids(\"<|eot_id|>\")\n",
    "    ]\n",
    "\n",
    "    time_start = time.time()\n",
    "    outputs = pipeline(\n",
    "        prompt,\n",
    "        max_new_tokens=2048,\n",
    "        eos_token_id=terminators,\n",
    "        do_sample=False,\n",
    "        pad_token_id=pipeline.tokenizer.eos_token_id\n",
    "    )\n",
    "    time_end = time.time()\n",
    "\n",
    "    return outputs[0][\"generated_text\"][len(prompt):], time_end - time_start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [10:44<00:00,  3.22s/it]\n"
     ]
    }
   ],
   "source": [
    "for i in tqdm(range(0, len(n2c2_df), 1)):\n",
    "    n2c2_df.loc[i, 'llama3_70b_instruct_one_shot'], n2c2_df.loc[i, 'llama3_70b_instruct_one_shot_time'] = get_re_2018_n2c2(n2c2_df.loc[i, 'text'], 1)\n",
    "    n2c2_df.loc[i, 'llama3_70b_instruct_five_shot'], n2c2_df.loc[i, 'llama3_70b_instruct_five_shot_time'] = get_re_2018_n2c2(n2c2_df.loc[i, 'text'], 5)\n",
    "    n2c2_df.loc[i, 'llama3_70b_instruct_ten_shot'], n2c2_df.loc[i, 'llama3_70b_instruct_ten_shot_time'] = get_re_2018_n2c2(n2c2_df.loc[i, 'text'], 10)\n",
    "    n2c2_df.loc[i, 'llama3_70b_instruct_twenty_shot'], n2c2_df.loc[i, 'llama3_70b_instruct_twenty_shot_time'] = get_re_2018_n2c2(n2c2_df.loc[i, 'text'], 20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1.2 Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional: you can just load the llm output from the csv file instead of running the above code\n",
    "# n2c2_df = pd.read_csv(\"data/RE/2018_n2c2/test_200_llama3_70b_instruct_results.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get rid of ' ' if any\n",
    "n2c2_df['llama3_70b_instruct_one_shot'] = n2c2_df['llama3_70b_instruct_one_shot'].apply(lambda x: x[1:-1] if \"'\" in x else x)\n",
    "n2c2_df['llama3_70b_instruct_five_shot'] = n2c2_df['llama3_70b_instruct_five_shot'].apply(lambda x: x[1:-1] if \"'\" in x else x)\n",
    "n2c2_df['llama3_70b_instruct_ten_shot'] = n2c2_df['llama3_70b_instruct_ten_shot'].apply(lambda x: x[1:-1] if \"'\" in x else x)\n",
    "n2c2_df['llama3_70b_instruct_twenty_shot'] = n2c2_df['llama3_70b_instruct_twenty_shot'].apply(lambda x: x[1:-1] if \"'\" in x else x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get digit label while considering failed LLM outputs as 'No relation'\n",
    "n2c2_df['labels'] = n2c2_df['labels'].apply(get_digit)\n",
    "n2c2_df['llama3_70b_instruct_one_shot_labels'] = n2c2_df['llama3_70b_instruct_one_shot'].apply(get_digit)\n",
    "n2c2_df['llama3_70b_instruct_five_shot_labels'] = n2c2_df['llama3_70b_instruct_five_shot'].apply(get_digit)\n",
    "n2c2_df['llama3_70b_instruct_ten_shot_labels'] = n2c2_df['llama3_70b_instruct_ten_shot'].apply(get_digit)\n",
    "n2c2_df['llama3_70b_instruct_twenty_shot_labels'] = n2c2_df['llama3_70b_instruct_twenty_shot'].apply(get_digit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F1 Score One Shot: 0.18366807609619562\n",
      "F1 Score Five Shot: 0.24288825512056\n",
      "F1 Score Ten Shot: 0.31791798685738076\n",
      "F1 Score Twenty Shot: 0.3190275703328499\n"
     ]
    }
   ],
   "source": [
    "y_true = n2c2_df['labels'].tolist()\n",
    "y_pred = n2c2_df['llama3_70b_instruct_one_shot_labels'].tolist()\n",
    "print(f\"F1 Score One Shot: {f1_score(y_true, y_pred, average='macro')}\")\n",
    "y_pred = n2c2_df['llama3_70b_instruct_five_shot_labels'].tolist()\n",
    "print(f\"F1 Score Five Shot: {f1_score(y_true, y_pred, average='macro')}\")\n",
    "y_pred = n2c2_df['llama3_70b_instruct_ten_shot_labels'].tolist()\n",
    "print(f\"F1 Score Ten Shot: {f1_score(y_true, y_pred, average='macro')}\")\n",
    "y_pred = n2c2_df['llama3_70b_instruct_twenty_shot_labels'].tolist()\n",
    "print(f\"F1 Score Twenty Shot: {f1_score(y_true, y_pred, average='macro')}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Llama-3-70B-Instruct one-shot prediction time: 0.65 seconds\n",
      "Average Llama-3-70B-Instruct five-shot prediction time: 0.97 seconds\n",
      "Average Llama-3-70B-Instruct ten-shot prediction time: 0.69 seconds\n",
      "Average Llama-3-70B-Instruct twenty-shot prediction time: 0.91 seconds\n"
     ]
    }
   ],
   "source": [
    "print(f\"Average Llama-3-70B-Instruct one-shot prediction time: {n2c2_df['llama3_70b_instruct_one_shot_time'].mean():.2f} seconds\")\n",
    "print(f\"Average Llama-3-70B-Instruct five-shot prediction time: {n2c2_df['llama3_70b_instruct_five_shot_time'].mean():.2f} seconds\")\n",
    "print(f\"Average Llama-3-70B-Instruct ten-shot prediction time: {n2c2_df['llama3_70b_instruct_ten_shot_time'].mean():.2f} seconds\")\n",
    "print(f\"Average Llama-3-70B-Instruct twenty-shot prediction time: {n2c2_df['llama3_70b_instruct_twenty_shot_time'].mean():.2f} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the inference results\n",
    "n2c2_df.to_csv('data/RE/2018_n2c2/test_200_llama3_70b_instruct_results.csv', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.2 GAD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2.1 Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "gad_df = pd.read_csv('data/RE/GAD/test_200.csv')\n",
    "gad_example_df = pd.read_csv('data/RE/GAD/examples.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_message = \"\"\"You are a helpful assistant to perform the following task.\n",
    "\"TASK: the task is to classify relations between a disease and a gene for a sentence.\"\n",
    "\"INPUT: the input is a sentence where the disease is labeled as @DISEASE$ and the gene is labeled as @GENE$ accordingly in a sentence. \"\n",
    "\"OUTPUT: your task is to select one out of the two types of relations (0 and 1) for the gene and disease without any explanation or other characters: \n",
    "0, no relations \n",
    "1, has relations\"\n",
    "\"\"\"\n",
    "def get_re_gad(sentence: str, shot: int = 0) -> str:\n",
    "    \"\"\"\n",
    "    Get the RE results of GAD dataset from few-shot prompting.\n",
    "    Args:\n",
    "        sentence: the input sentence\n",
    "        shot: the number of few-shot examples\n",
    "    Returns:\n",
    "        the RE results\n",
    "    \"\"\"\n",
    "\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": system_message},\n",
    "    ]\n",
    "    for i in range(shot):\n",
    "        messages.append({\"role\": \"user\", \"content\": gad_example_df. iloc[i]['text']}) \n",
    "        messages.append({\"role\": \"assistant\", \"content\": gad_example_df.iloc[i]['labels']})\n",
    "    messages.append({\"role\": \"user\", \"content\": sentence})\n",
    "    prompt = pipeline.tokenizer.apply_chat_template(\n",
    "            messages, \n",
    "            tokenize=False, \n",
    "            add_generation_prompt=True\n",
    "    )\n",
    "    terminators = [\n",
    "        pipeline.tokenizer.eos_token_id,\n",
    "        pipeline.tokenizer.convert_tokens_to_ids(\"<|eot_id|>\")\n",
    "    ]\n",
    "\n",
    "    time_start = time.time()\n",
    "    outputs = pipeline(\n",
    "        prompt,\n",
    "        max_new_tokens=2048,\n",
    "        eos_token_id=terminators,\n",
    "        do_sample=False,\n",
    "        pad_token_id=pipeline.tokenizer.eos_token_id\n",
    "    )\n",
    "    time_end = time.time()\n",
    "\n",
    "    return outputs[0][\"generated_text\"][len(prompt):], time_end - time_start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [05:04<00:00,  1.52s/it]\n"
     ]
    }
   ],
   "source": [
    "for i in tqdm(range(0, len(gad_df), 1)):\n",
    "    gad_df.loc[i, 'llama3_70b_instruct_one_shot'], gad_df.loc[i, 'llama3_70b_instruct_one_shot_time'] = get_re_gad(gad_df.iloc[i]['text'], 1)\n",
    "    gad_df.loc[i, 'llama3_70b_instruct_five_shot'], gad_df.loc[i, 'llama3_70b_instruct_five_shot_time'] = get_re_gad(gad_df.iloc[i]['text'], 5)\n",
    "    gad_df.loc[i, 'llama3_70b_instruct_ten_shot'], gad_df.loc[i, 'llama3_70b_instruct_ten_shot_time'] = get_re_gad(gad_df.iloc[i]['text'], 10)\n",
    "    gad_df.loc[i, 'llama3_70b_instruct_twenty_shot'], gad_df.loc[i, 'llama3_70b_instruct_twenty_shot_time'] = get_re_gad(gad_df.iloc[i]['text'], 20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2.2 Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional: you can just load the llm output from the csv file instead of running the above code\n",
    "gad_df = pd.read_csv(\"data/RE/GAD/test_200_llama3_70b_instruct_results.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F1 Score One Shot: 0.4521130881140456\n",
      "F1 Score Five Shot: 0.4560100428915158\n",
      "F1 Score Ten Shot: 0.4575378538512179\n",
      "F1 Score Twenty Shot: 0.4522529971062422\n"
     ]
    }
   ],
   "source": [
    "y_true = gad_df['labels'].tolist()\n",
    "y_pred = gad_df['llama3_70b_instruct_one_shot'].tolist()\n",
    "print(f\"F1 Score One Shot: {f1_score(y_true, y_pred, average='macro')}\")\n",
    "y_pred = gad_df['llama3_70b_instruct_five_shot'].tolist()\n",
    "print(f\"F1 Score Five Shot: {f1_score(y_true, y_pred, average='macro')}\")\n",
    "y_pred = gad_df['llama3_70b_instruct_ten_shot'].tolist()\n",
    "print(f\"F1 Score Ten Shot: {f1_score(y_true, y_pred, average='macro')}\")\n",
    "y_pred = gad_df['llama3_70b_instruct_twenty_shot'].tolist()\n",
    "print(f\"F1 Score Twenty Shot: {f1_score(y_true, y_pred, average='macro')}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Llama-3-70B-Instruct one-shot prediction time: 0.29 seconds\n",
      "Average Llama-3-70B-Instruct five-shot prediction time: 0.34 seconds\n",
      "Average Llama-3-70B-Instruct ten-shot prediction time: 0.40 seconds\n",
      "Average Llama-3-70B-Instruct twenty-shot prediction time: 0.50 seconds\n"
     ]
    }
   ],
   "source": [
    "print(f\"Average Llama-3-70B-Instruct one-shot prediction time: {gad_df['llama3_70b_instruct_one_shot_time'].mean():.2f} seconds\")\n",
    "print(f\"Average Llama-3-70B-Instruct five-shot prediction time: {gad_df['llama3_70b_instruct_five_shot_time'].mean():.2f} seconds\")\n",
    "print(f\"Average Llama-3-70B-Instruct ten-shot prediction time: {gad_df['llama3_70b_instruct_ten_shot_time'].mean():.2f} seconds\")\n",
    "print(f\"Average Llama-3-70B-Instruct twenty-shot prediction time: {gad_df['llama3_70b_instruct_twenty_shot_time'].mean():.2f} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the inference results\n",
    "gad_df.to_csv('data/RE/GAD/test_200_llama3_70b_instruct_results.csv', index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}