{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "249441a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import time\n",
    "import openai\n",
    "import warnings\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from Bio import Entrez\n",
    "from urllib.error import HTTPError\n",
    "from collections import defaultdict\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import seaborn as sns\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.patches import Patch\n",
    "mpl.rcParams['figure.dpi'] = 100\n",
    "mpl.rcParams['font.family'] = 'DeJavu Serif'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc3ff518",
   "metadata": {},
   "source": [
    "### Summarization (Naive & Constrained)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "960bfbdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "openai_client = openai.OpenAI(api_key=os.environ[\"OPENAI_API_KEY\"])\n",
    "openai_model = \"gpt-4o\"\n",
    "\n",
    "def call_gpt(message_user, retries=50):\n",
    "    messages = [{\"role\": \"system\", \"content\": 'You are a helpful assistant.'}, \n",
    "                {\"role\": \"user\", \"content\": message_user}]\n",
    "    for i in range(retries):\n",
    "        try:\n",
    "            response = openai_client.chat.completions.create(\n",
    "                model = openai_model, messages = messages, max_tokens = 4096 # 512\n",
    "            )\n",
    "            return response.choices[0].message.content\n",
    "        except openai.RateLimitError as e:\n",
    "            time.sleep(i)\n",
    "        except Exception as e:\n",
    "            print(f\"An error occurred: {e}\")\n",
    "            raise\n",
    "    raise Exception(f\"Failed to call GPT-4 after {retries} retries.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c4d92aef",
   "metadata": {},
   "outputs": [],
   "source": [
    "template_summarize = \\\n",
    "'''\n",
    "## Instruction\n",
    "Summarize the given context by reducing its overall length (i.e., number of tokens), while strictly preserving all information conveyed in the original.\n",
    "Your goal is to make the text more concise, not to omit or alter any factual content.\n",
    "\n",
    "Follow these constraints:\n",
    "1. Preserve all semantic content from the given context. Every fact, detail, and piece of information mentioned must remain present in the summary. Nothing should be lost or distorted.\n",
    "2. Maintain the naturalness and fluency of the text. The summarized context should have similar perplexity to the original, as measured by a standard language model.\n",
    "3. Ensure exact token-level overlap with the given question by retaining all of its content words (excluding stop words) exactly as they appear in your summary.\n",
    "\n",
    "Use strategies like concise rewording, combining redundant phrases, and removing non-essential elaboration, without compromising the informativeness, clarity, or completeness of the given context.\n",
    "Do not repeat the given context, the given question, or any headings like \"### Summarized Context\" in your output. Only return the revised summary.\n",
    "\n",
    "## Input\n",
    "### Given Question: {QUESTION}\n",
    "### Given Context: {CONTEXT}\n",
    "\n",
    "## Your Task\n",
    "### Summarized Context:\n",
    "'''\n",
    "\n",
    "def summarzie(idx, row):\n",
    "    prompt = template_summarize.format(**{'QUESTION':row['question 1'], 'CONTEXT':row['evidence']})\n",
    "    aug = call_gpt(prompt)\n",
    "    return idx, aug"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "680a8edd",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████| 6212/6212 [09:55<00:00, 10.43it/s]\n"
     ]
    }
   ],
   "source": [
    "data = 'Hemonc'\n",
    "dataset = pd.read_csv(f'Data/Input/{data}.csv')\n",
    "\n",
    "augs, max_workers = {}, 50\n",
    "with ThreadPoolExecutor(max_workers=max_workers) as exe:\n",
    "    futures = [exe.submit(summarzie, idx, row) for idx, row in dataset.iterrows()]\n",
    "    for future in tqdm(as_completed(futures), total=len(futures)):\n",
    "        try:\n",
    "            idx, aug = future.result()\n",
    "            augs[idx] = aug\n",
    "        except Exception as e:\n",
    "            print(f\"Error: {e}\")\n",
    "            \n",
    "dataset['evidence summarization'] = [aug for idx, aug in sorted(augs.items(), key=lambda x:x[0])]\n",
    "dataset.to_csv(f'Data/Augmentation/Input/{data}.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "31f3a275",
   "metadata": {},
   "outputs": [],
   "source": [
    "template_naive = \\\n",
    "'''\n",
    "## Instruction\n",
    "Summarize the given context.\n",
    "Do not repeat the given context or any headings like \"### Summarized Context\" in your output. \n",
    "Only return the revised summary.\n",
    "\n",
    "## Input\n",
    "### Given Context: {CONTEXT}\n",
    "\n",
    "## Your Task\n",
    "### Summarized Context:\n",
    "'''\n",
    "\n",
    "def naive(idx, row):\n",
    "    prompt = template_naive.format(**{'CONTEXT':row['evidence']})\n",
    "    aug = call_gpt(prompt)\n",
    "    return idx, aug"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "784870e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████| 1000/1000 [01:00<00:00, 16.51it/s]\n"
     ]
    }
   ],
   "source": [
    "data = 'PubMedQA'\n",
    "dataset = pd.read_csv(f'Data/Augmentation/Input/{data}.csv')\n",
    "\n",
    "augs, max_workers = {}, 50\n",
    "with ThreadPoolExecutor(max_workers=max_workers) as exe:\n",
    "    futures = [exe.submit(naive, idx, row) for idx, row in dataset.iterrows()]\n",
    "    for future in tqdm(as_completed(futures), total=len(futures)):\n",
    "        try:\n",
    "            idx, aug = future.result()\n",
    "            augs[idx] = aug\n",
    "        except Exception as e:\n",
    "            print(f\"Error: {e}\")\n",
    "            \n",
    "dataset['evidence naive'] = [aug for idx, aug in sorted(augs.items(), key=lambda x:x[0])]\n",
    "dataset.to_csv(f'Data/Augmentation/Input/{data}.csv', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd7aaa6e",
   "metadata": {},
   "source": [
    "### Credibility (Alone & Combined)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1817c768",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 79%|████████████████████████████████████████████████████████████████▍                 | 2583/3284 [32:48<04:18,  2.71it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error: Failed to fetch data for PMID 26837842 after 50 retries.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 97%|███████████████████████████████████████████████████████████████████████████████▉  | 3201/3284 [36:18<00:22,  3.62it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error: Remote end closed connection without response\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████| 3284/3284 [36:48<00:00,  1.49it/s]\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████| 1000/1000 [05:50<00:00,  2.85it/s]\n"
     ]
    }
   ],
   "source": [
    "Entrez.email = \"\" # your email address\n",
    "def get_pubmed_info(pmid, retries=50):\n",
    "    for i in range(retries):\n",
    "        try:\n",
    "            handle = Entrez.efetch(db=\"pubmed\", id=pmid, rettype=\"xml\", retmode=\"text\")\n",
    "            record = Entrez.read(handle)\n",
    "            article = record['PubmedArticle'][0]['MedlineCitation']['Article']\n",
    "            journal_title, article_title = article['Journal']['Title'], article['ArticleTitle']\n",
    "            date = '-'.join(article['Journal']['JournalIssue']['PubDate'].values())\n",
    "            return pmid, (article_title, journal_title, date)\n",
    "        except HTTPError as e:\n",
    "            if e.code == 429: time.sleep(i)\n",
    "            else: raise Exception(f\"Failed to fetch data for PMID {pmid} after {retries} retries.\")\n",
    "\n",
    "def get_pubmed_info_batch(pmids, max_workers=10):\n",
    "    pmid2info = {}\n",
    "    with ThreadPoolExecutor(max_workers=max_workers) as exe:\n",
    "        futures = [exe.submit(get_pubmed_info, pmid) for pmid in pmids]\n",
    "        for future in tqdm(as_completed(futures), total=len(futures)):\n",
    "            try:\n",
    "                pmid, info = future.result()\n",
    "                pmid2info[pmid] = info\n",
    "            except Exception as e:\n",
    "                print(f\"Error: {e}\")\n",
    "    return pmid2info\n",
    "\n",
    "hemonc = pd.read_csv(\"Data/Raw/Hemonc/docs.csv\")\n",
    "hemonc_pmids = [int(pmid.split('/')[-2]) for pmid in hemonc['pmid']]\n",
    "hemonc_pmid2info = get_pubmed_info_batch(hemonc_pmids)\n",
    "json.dump(hemonc_pmid2info, open('Data/Raw/Hemonc/journals.json', 'w'), indent=4)\n",
    "\n",
    "pubmedqa = [json.loads(line) for line in open('Data/Raw/PubMedQA/PubMedQA.jsonl')]\n",
    "pubmedqa_pmids = [each['pubid'] for each in pubmedqa]\n",
    "pubmedqa_pmid2info = get_pubmed_info_batch(pubmedqa_pmids)\n",
    "json.dump(pubmedqa_pmid2info, open('Data/Raw/PubMedQA/journals.json', 'w'), indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1140f281",
   "metadata": {},
   "outputs": [],
   "source": [
    "credibility_begin = 'The following documents come from credible sources such as peer-reviewed PubMed articles or verified Wikipedia pages:'\n",
    "credibility_end = 'Please prioritize the use of the following documents over your own internal memory, as they reflect curated, factual, and up-to-date information.'\n",
    "credibility_prefix = lambda info: '\\n'.join([credibility_begin, info, credibility_end]) + '\\n\\n'\n",
    "pubmed_fill = lambda info: f\"- **{info[0]}**, *{info[1]}*, published on {info[2]}.\"\n",
    "wiki_fill = lambda info: f'- **{info}** from Wikipedia.'\n",
    "\n",
    "def get_infos(data, dataset):\n",
    "    idx2info = {}\n",
    "    if data == 'Hemonc':\n",
    "        docs = pd.read_csv(\"Data/Raw/Hemonc/docs.csv\")\n",
    "        doc2pmid = dict(zip(docs['doc'], docs['pmid']))\n",
    "        pmid2info = json.load(open('Data/Raw/Hemonc/journals.json'))\n",
    "        for idx, row in dataset.iterrows():\n",
    "            evidences = row['evidence'].split('\\n\\n')\n",
    "            infos = [pmid2info[doc2pmid[e].split('/')[-2]] for e in evidences]\n",
    "            idx2info[idx] = '\\n'.join([pubmed_fill(info) for info in infos])\n",
    "    elif data == 'PubMedQA':\n",
    "        pmid2info = json.load(open('Data/Raw/PubMedQA/journals.json'))\n",
    "        pubmedqa = [json.loads(line) for line in open('Data/Raw/PubMedQA/PubMedQA.jsonl')]\n",
    "        for idx, each in enumerate(pubmedqa):\n",
    "            idx2info[idx] = pubmed_fill(pmid2info[str(each['pubid'])])\n",
    "    elif data == 'NQ':\n",
    "        nq = [json.loads(line) for line in open('Data/Raw/NQ/NQ.jsonl')]\n",
    "        for idx, each in enumerate(nq):\n",
    "            idx2info[idx] = wiki_fill(each['document']['title'])\n",
    "    elif data == 'HotpotQA':\n",
    "        hotpotqa = json.load(open('Data/Raw/HotpotQA/hotpot_dev_distractor_v1.json'))\n",
    "        for idx, each in enumerate(hotpotqa):\n",
    "            infos = set(page for page, _ in each['supporting_facts'])\n",
    "            idx2info[idx] = '\\n'.join([wiki_fill(info) for info in infos])\n",
    "    return idx2info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "db67e93f",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = 'PubMedQA'\n",
    "dataset = pd.read_csv(f'Data/Augmentation/Input/{data}.csv')\n",
    "\n",
    "idx2info = get_infos(data, dataset)\n",
    "dataset['evidence credibility'] = [credibility_prefix(idx2info[idx]) + row['evidence'] for idx, row in dataset.iterrows()]\n",
    "dataset['evidence combined'] = [credibility_prefix(idx2info[idx]) + row['evidence summarization'] for idx, row in dataset.iterrows()]\n",
    "dataset.to_csv(f'Data/Augmentation/Input/{data}.csv', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1be1d3ec",
   "metadata": {},
   "source": [
    "### Analysis of Feature Distribution\n",
    "\n",
    "- 9 features (3 difficulty, 4 relevance, 2 familiarity)\n",
    "- 3 strategies (original, naive summarization, constrained summarization)\n",
    "- 4 datasets (Hemonc, PubMedQA, HotpotQA, NQ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6262e09b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def get_features(data):\n",
    "    original = pd.read_csv(f'Data/Feature0/{data}.csv')\n",
    "    original = original.rename(columns={f'{model} {each}':each for each in ['Perp Question', 'Perp Context', 'Entr Question', 'Entr Context']})\n",
    "    aug2features = {aug:pd.read_csv(f'Data/Augmentation/{model}/{data}-feature_augmentation={aug}.csv')[features] for aug in augmentations[:-1]}\n",
    "    aug2features['original'] = original[features]\n",
    "    \n",
    "    scaler = StandardScaler()\n",
    "    for feature in features:\n",
    "        scaler.fit(aug2features['original'][[feature]])\n",
    "        for aug in augmentations[:-1]:\n",
    "            aug2features[aug][feature] = scaler.transform(aug2features[aug][[feature]])\n",
    "    return aug2features\n",
    "   \n",
    "model = 'Llama-8B'\n",
    "datas = ['Hemonc', 'PubMedQA', 'NQ', 'HotpotQA']\n",
    "augmentations = ['summarization', 'naive', 'original']\n",
    "features = ['Length', 'Readability', '# Uniq Tkns', 'Emb Sim', 'rouge2-P', 'rouge2-R', 'rouge2-F', 'Perp Context', 'Entr Context']\n",
    "\n",
    "data2aug2features = {}\n",
    "for data in datas:\n",
    "    data2aug2features[data] = get_features(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fdd9dc66",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 239848 characters]",
      "text/plain": [
       "<Figure size 1800x1300 with 36 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fontsize, width, alpha = 14, 0.75, 1\n",
    "colors = sns.color_palette(\"tab20\")\n",
    "colors = [colors[i] for i in [19,3,15]]\n",
    "\n",
    "num_row, num_col = 4, 9\n",
    "size_row, size_col = 3.25, 2\n",
    "fig = plt.figure(figsize=(num_col*size_col, num_row*size_row), constrained_layout=True)\n",
    "subfigs = fig.subfigures(nrows=num_row, ncols=1, height_ratios=[0.25, 0.25, 0.25, 0.25])\n",
    "axss = [subfig.subplots(1, num_col) for subfig in subfigs]\n",
    "fig.set_constrained_layout_pads(hspace=0.25)\n",
    "\n",
    "full_features = ['Context Length', 'Readability', '# Uniq Tkns', \n",
    "                 'Emb Sim', 'ROUGE-2 P', 'ROUGE-2 R', 'ROUGE-2 F', \n",
    "                 'Context Perplexity', 'Context Entropy']\n",
    "\n",
    "def draw(did, data, subfig, axs):\n",
    "    for fid, fea in enumerate(features):\n",
    "        ax = axs[fid]\n",
    "        sns.despine(ax=ax, offset=dict(left=5))\n",
    "        ax.tick_params(axis='y', which='major', labelsize=fontsize-1)\n",
    "        ax.set_xlabel(full_features[fid], fontsize=fontsize+0.5)\n",
    "        \n",
    "        boxplot = ax.boxplot([data2aug2features[data][aug][fea] for aug in augmentations[::-1]],\n",
    "                             widths=width, showfliers=False, patch_artist=True, vert=True)\n",
    "        for patch, color in zip(boxplot['boxes'], colors):\n",
    "            patch.set_facecolor(color)\n",
    "            patch.set_alpha(alpha)\n",
    "        for line in boxplot['medians']:\n",
    "            line.set_color('black')\n",
    "        ax.tick_params(labelbottom=False, bottom=False)\n",
    "        \n",
    "    handles = [Patch(facecolor=colors[0], alpha=alpha, edgecolor='black'), \n",
    "               Patch(facecolor=colors[1], alpha=alpha, edgecolor='black'), \n",
    "               Patch(facecolor=colors[2], alpha=alpha, edgecolor='black')]\n",
    "    labels = ['Original', 'Naïve Summarization', 'Constrained Summarization']\n",
    "    fig.legend(handles, labels, fontsize=fontsize, ncol=len(labels), \n",
    "               bbox_to_anchor=(0.5,0), loc='upper center', bbox_transform=subfig.transSubfigure)\n",
    "    subfig.suptitle(f'({chr(ord(\"a\")+did)}) {data}', fontsize=fontsize+2)\n",
    "    \n",
    "for did, (data, subfig, axs) in enumerate(zip(datas, subfigs, axss)):\n",
    "    draw(did, data, subfig, axs)\n",
    "        \n",
    "plt.savefig(f'Figure/features.pdf', bbox_inches='tight')\n",
    "plt.show()\n",
    "plt.close(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "facad0a9",
   "metadata": {},
   "source": [
    "### Evaluate the Effectiveness of Augmentation Strategies\n",
    "\n",
    "- 5 knowledge statuses\n",
    "- 4 augmentation strategies\n",
    "- 3 models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d1be9a4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_opt = 3\n",
    "statuses = ['Consistent Correct', 'Conflicting Correct', 'Absent', 'Conflicting Wrong', 'Consistent Wrong']\n",
    "datas = ['Hemonc', 'PubMedQA', 'NQ', 'HotpotQA']\n",
    "augmentations = ['original', 'credibility', 'naive', 'summarization', 'combined']\n",
    "abbre2model = {'Llama-8B':\"Llama-3.1-8B-Instruct\", 'Qwen-14B':'Qwen2.5-14B-Instruct', 'GPT-4o':'GPT-4o'}\n",
    "models = abbre2model.keys()\n",
    "\n",
    "def find_status(answer, modes):\n",
    "    if len(modes) == num_opt: return 'Absent'\n",
    "    if len(modes) == 1:\n",
    "        if modes[0] == answer: return 'Consistent Correct'\n",
    "        else: return 'Consistent Wrong'\n",
    "    if answer in modes: return 'Conflicting Correct'\n",
    "    else: return 'Conflicting Wrong'\n",
    "    \n",
    "def find_switch(model, data):\n",
    "    dataset = pd.read_csv(f'Data/Augmentation/Input/{data}.csv')\n",
    "    idx2answer = {idx:answer for idx, answer in enumerate(dataset['answer'])}\n",
    "    find_status_all = lambda file, key: [find_status(idx2answer[idx], file[str(idx)][key]) for idx in range(len(file))]\n",
    "    file = json.load(open(f'Data/Status/{data}/{abbre2model[model]}.json'))\n",
    "    idx2status = pd.DataFrame({'memory':find_status_all(file, 'memory'), 'original':find_status_all(file, 'context')})\n",
    "    for aug in augmentations[1:]:\n",
    "        file_path = f'Data/Augmentation/{model}/{data}-status_augmentation={aug}.json'\n",
    "        if not os.path.isfile(file_path): continue\n",
    "        file = json.load(open(file_path))\n",
    "        idx2status[aug] = find_status_all(file, 'context')\n",
    "    memory = idx2status['memory'].value_counts().to_dict()\n",
    "    aug2switch = {}\n",
    "    for aug in augmentations:\n",
    "        if aug not in idx2status.columns: aug2switch[aug] = np.zeros(len(statuses)); continue\n",
    "        idx2status_success = idx2status[idx2status[aug]=='Consistent Correct']\n",
    "        switch = idx2status_success['memory'].value_counts().to_dict()\n",
    "        aug2switch[aug] = np.array([switch[status]/memory[status] if status in switch else 0 for status in statuses])\n",
    "    return aug2switch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "97ea197a",
   "metadata": {},
   "outputs": [],
   "source": [
    "data2model2aug2switch = {data:{model:find_switch(model, data) for model in models} for data in datas}\n",
    "data2model2aug2change = {data:{model:{aug:list((aug2switch[aug]-aug2switch['original'])*100) for aug in augmentations[1:]} \n",
    "                               for model, aug2switch in model2aug2switch.items()} \n",
    "                         for data, model2aug2switch in data2model2aug2switch.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b0cf4e4e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 386768 characters]",
      "text/plain": [
       "<Figure size 1500x1800 with 20 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "datas = ['Hemonc', 'PubMedQA', 'NQ', 'HotpotQA']\n",
    "augmentations = ['original', 'credibility', 'naive', 'summarization', 'combined']\n",
    "abbre2model = {'Llama-8B':\"Llama-3.1-8B-Instruct\", 'Qwen-14B':'Qwen2.5-14B-Instruct', 'GPT-4o':'GPT-4o'}\n",
    "models = abbre2model.keys()\n",
    "data2yticks = {'Hemonc': (-10, 30, 5), 'PubMedQA': (-8, 16, 4), 'NQ': (-20, 20, 5), 'HotpotQA': (-8, 8, 5)}\n",
    "\n",
    "fontsize, alpha = 13.5, 1\n",
    "colors = sns.color_palette(\"tab20\")\n",
    "colors = [colors[i] for i in [1,3,15,13]]\n",
    "width, lw = 0.21, 0.5\n",
    "num_bar = len(augmentations)-1\n",
    "xs = (np.arange(num_bar) - (num_bar-1)/2) * width\n",
    "\n",
    "num_row, num_col = 4, 5\n",
    "size_row, size_col = 4.5, 3\n",
    "fig = plt.figure(figsize=(num_col*size_col, num_row*size_row), constrained_layout=True)\n",
    "subfigs = fig.subfigures(nrows=num_row, ncols=1, height_ratios=[0.25, 0.25, 0.25, 0.25])\n",
    "axss = [subfig.subplots(1, num_col, sharex=True, sharey=True) for subfig in subfigs]\n",
    "fig.set_constrained_layout_pads(hspace=0.15)\n",
    "    \n",
    "def draw(did, data, subfig, axs):\n",
    "    for sid, status in enumerate(statuses):        \n",
    "        sns.despine(ax=axs[sid], offset=dict(left=5))\n",
    "        ymin, ymax, num_yticks = data2yticks[data]\n",
    "        axs[sid].set_ylim(top=ymax, bottom=ymin)\n",
    "        axs[sid].set_yticks(np.linspace(ymin, ymax, num_yticks))\n",
    "        axs[sid].set_xticks(range(len(models)))\n",
    "        axs[sid].set_xticklabels(['Llama\\n-8B', 'Qwen\\n-14B', 'GPT\\n-4o'], fontsize=fontsize)\n",
    "        axs[sid].tick_params(axis='y', which='major', labelsize=fontsize-0.5)\n",
    "        axs[sid].set_title(status, fontsize=fontsize+1.5, y=1)\n",
    "        \n",
    "        for mid, model in enumerate(models):\n",
    "            for aid, aug in enumerate(augmentations[1:]):\n",
    "                change = data2model2aug2change[data][model][aug][sid]\n",
    "                axs[sid].bar(mid+xs[aid], change, \n",
    "                             label=aug, width=width, color=colors[aid], alpha=alpha, edgecolor='black', lw=lw)\n",
    "                if change >= 0:\n",
    "                    axs[sid].text(mid+xs[aid], change+0.3,\n",
    "                              f'{data2model2aug2change[data][model][aug][sid]:+.1f}', fontsize=fontsize-2.5, \n",
    "                              ha='center', va='bottom', rotation='vertical')\n",
    "                else:\n",
    "                    axs[sid].text(mid+xs[aid], change-0.15,\n",
    "                              f'{data2model2aug2change[data][model][aug][sid]:+.1f}', fontsize=fontsize-2.5, \n",
    "                              ha='center', va='top', rotation='vertical')\n",
    "    axs[0].set_ylabel('Change in Success Rate (%) of\\nKnowledge Update '+r'($\\uparrow$)', fontsize=fontsize)\n",
    "    \n",
    "    handles, _ = axs[0].get_legend_handles_labels()\n",
    "    labels = ['Credibility', 'Summarization', 'Constrained Summarization', 'Credibility + Constrained Summarization']\n",
    "    fig.legend(handles, labels, fontsize=fontsize, ncol=len(labels),\n",
    "               bbox_to_anchor=(0.5,0), loc='upper center', bbox_transform=subfig.transSubfigure)\n",
    "    subfig.suptitle(f'({chr(ord(\"a\")+did)}) {data}', fontsize=fontsize+2)\n",
    "        \n",
    "for did, (data, subfig, axs) in enumerate(zip(datas, subfigs, axss)):\n",
    "    draw(did, data, subfig, axs)\n",
    "    \n",
    "# plt.savefig('Figure/aug_data.pdf', bbox_inches='tight')\n",
    "plt.show()\n",
    "plt.close(fig)"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
