{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fbc1d023",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\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": "code",
   "execution_count": null,
   "id": "44015bfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_opt = 3\n",
    "versions = ['memory', 'context']\n",
    "statuses = ['Consistent Correct', 'Conflicting Correct', 'Absent', 'Conflicting Wrong', 'Consistent Wrong']\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 summarize_status(idx2answer, idx2modes):\n",
    "    idx2statuses = {}\n",
    "    for idx, modes in idx2modes.items():\n",
    "        idx2statuses[int(idx)] = {key:find_status(idx2answer[int(idx)], modes[key]) for key in versions}\n",
    "    memory, context, switch = [], [], []\n",
    "    for status in statuses:\n",
    "        indicies_mem = [idx for idx in range(len(idx2answer)) if idx2statuses[idx]['memory']==status]\n",
    "        indicies_con = [idx for idx in range(len(idx2answer)) if idx2statuses[idx]['context']==status]\n",
    "        memory.append(len(indicies_mem)); context.append(len(indicies_con)); switch.append([])\n",
    "        for status_swi in statuses:\n",
    "            indicies_swi = [idx for idx in indicies_mem if idx2statuses[idx]['context']==status_swi]\n",
    "            switch[-1].append(len(indicies_swi))\n",
    "    memory, context, switch = np.array(memory), np.array(context), np.array(switch)\n",
    "    memory, context, switch = memory / memory.sum(), context / context.sum(), switch / (switch.sum(axis=1, keepdims=True) + 1e-10)\n",
    "    return memory, context, switch\n",
    "\n",
    "def get_results(data, abbre2model):\n",
    "    dataset = pd.read_csv(f'Data/Input/{data}.csv')\n",
    "    idx2answer = {idx:answer for idx, answer in enumerate(dataset['answer'])}\n",
    "    model2memory, model2context, model2switch = {}, {}, {}\n",
    "    for abbre, model in abbre2model.items():\n",
    "        file = f'Data/Status/{data}/{model}.json'\n",
    "        if not os.path.exists(file): continue\n",
    "        idx2modes = json.load(open(file))\n",
    "        model2memory[abbre], model2context[abbre], model2switch[abbre] = summarize_status(idx2answer, idx2modes)\n",
    "    return model2memory, model2context, model2switch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "41118bc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "datas = ['Hemonc', 'PubMedQA', 'NQ', 'HotpotQA']\n",
    "families = ['Gemma', 'Llama', 'Qwen']\n",
    "abbre2model = {'Gemma-2B':'gemma-2-2b-it', 'Gemma-9B':'gemma-2-9b-it', 'Gemma-27B':'gemma-2-27b-it', \n",
    "               'Llama-3B':'Llama-3.2-3B-Instruct', 'Llama-8B':'Llama-3.1-8B-Instruct', 'Llama-70B':'Llama-3.3-70B-Instruct', \n",
    "               'Qwen-3B':'Qwen2.5-3B-Instruct', 'Qwen-7B':'Qwen2.5-7B-Instruct', 'Qwen-14B':'Qwen2.5-14B-Instruct'}\n",
    "data2model2memory, data2model2context, data2model2switch = {}, {}, {}\n",
    "for data in datas:\n",
    "    data2model2memory[data], data2model2context[data], data2model2switch[data] = get_results(data, abbre2model)\n",
    "    \n",
    "results = []\n",
    "for version, data2model2ratio in zip(versions, [data2model2memory, data2model2context]):\n",
    "    for data, model2ratio in data2model2ratio.items():\n",
    "        for model, ratio in model2ratio.items():\n",
    "            family, size = model.split('-')\n",
    "            size = int(size[:-1])\n",
    "            results.append((version, data, family, size, *ratio))\n",
    "results = pd.DataFrame(results, columns=['version', 'data', 'family', 'size']+statuses)\n",
    "average_data = results.groupby(['version', 'family', 'size'])[statuses].mean().reset_index()\n",
    "average_model = results.groupby(['version', 'data'])[statuses].mean().reset_index()\n",
    "average_switch = np.mean([[data2model2switch[data][model][:,0], data2model2switch[data][model][:,1:].sum(-1)] \n",
    "                          for data in datas for model in abbre2model], axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3c940b5",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 313292 characters]",
      "text/plain": [
       "<Figure size 1500x750 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fontsize, alpha, lw = 13, 0.75, 0.5\n",
    "colors_top = sns.color_palette(\"Set2\", len(datas))\n",
    "colors_bot = sns.color_palette(\"tab20c\")\n",
    "colors_bot = [colors_bot[i] for i in [1,2,3,5,6,7,9,10,11]]\n",
    "\n",
    "num_row, num_col = 2, 5\n",
    "size_row, size_col = 3.75, 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=[1, 1])\n",
    "axs_top, axs_bot = subfigs[0].subplots(1, num_col), subfigs[1].subplots(1, num_col)\n",
    "\n",
    "fig.set_constrained_layout_pads(hspace=0.15)\n",
    "subfigs[0].suptitle(\"(a) Context-Induced Shifts from a Dataset Perspective (Multi-Choice Setting, Gold Context)\", fontsize=fontsize+3.5)\n",
    "subfigs[1].suptitle(\"(b) Context-Induced Shifts from a Model Perspective (Multi-Choice Setting, Gold Context)\", fontsize=fontsize+3.5)\n",
    "plt.rcParams['hatch.linewidth'] = lw\n",
    "\n",
    "width = 0.75\n",
    "sid2top = {0:1, 1:0.2, 2:0.2, 3:0.075, 4:0.5}\n",
    "for sid, status in enumerate(statuses):\n",
    "    sns.despine(ax=axs_top[sid], offset=dict(left=5))\n",
    "    axs_top[sid].set_ylim(top=sid2top[sid]*100)\n",
    "    axs_top[sid].set_yticks(np.linspace(0, sid2top[sid]*100, 6))\n",
    "    axs_top[sid].tick_params(axis='both', which='major', labelsize=fontsize-1, labelbottom=False, bottom=False)\n",
    "    axs_top[sid].set_title(status.replace(' ', '\\n') + (r' ($\\downarrow$)' if sid else r' ($\\uparrow$)'), fontsize=fontsize+1.5)\n",
    "    \n",
    "    for did, data in enumerate(datas):        \n",
    "        average_memory = average_model[(average_model['version']=='memory') & (average_model['data']==data)][status].to_numpy()\n",
    "        average_context = average_model[(average_model['version']=='context') & (average_model['data']==data)][status].to_numpy()\n",
    "        average_max = np.max((average_memory, average_context), axis=0)\n",
    "        diffs = average_context - average_memory\n",
    "        hatches = ['//' if diff >= 0 else 'xx' for diff in diffs]\n",
    "        \n",
    "        axs_top[sid].bar(did, average_max*100, label=data,\n",
    "                         width=width, color=colors_top[did], alpha=alpha, edgecolor='black', lw=lw)\n",
    "        axs_top[sid].bar(did, diffs*100, bottom=average_memory*100, \n",
    "                         width=width, facecolor='none', edgecolor='black', hatch=hatches, lw=lw)\n",
    "        for top, diff in zip(average_context, diffs):\n",
    "            axs_top[sid].text(did, (sid2top[sid]*0.01+max(top,top-diff))*100, f'{diff*100:+.1f}', \n",
    "                              fontsize=fontsize-1, ha='center', va='bottom')\n",
    "\n",
    "axs_top[0].set_ylabel(f'Percentage', fontsize=fontsize+0.5)\n",
    "handles, labels = axs_top[0].get_legend_handles_labels()\n",
    "handles += [\n",
    "    Patch(facecolor='white', edgecolor='black', hatch='//'),\n",
    "    Patch(facecolor='white', edgecolor='black', hatch='xx')\n",
    "]\n",
    "labels += [\n",
    "    'Increase w/ Context', 'Decrease w/ Context'\n",
    "]\n",
    "subfigs[0].legend(handles, labels, fontsize=fontsize, ncol=len(labels),\n",
    "                  bbox_to_anchor=(0.5,0), loc='upper center', bbox_transform=subfigs[0].transSubfigure)\n",
    "\n",
    "num_bar = 3\n",
    "width = 0.275\n",
    "xs = (np.arange(num_bar) - (num_bar-1)/2) * width\n",
    "\n",
    "sid2top = {0:1, 1:0.25, 2:0.25, 3:0.075, 4:0.45}\n",
    "for sid, status in enumerate(statuses):\n",
    "    sns.despine(ax=axs_bot[sid], offset=dict(left=5))\n",
    "    axs_bot[sid].set_ylim(top=sid2top[sid]*100)\n",
    "    axs_bot[sid].set_yticks(np.linspace(0, sid2top[sid]*100, 6))\n",
    "    axs_bot[sid].tick_params(axis='both', which='major', labelsize=fontsize-1, labelbottom=False, bottom=False)\n",
    "    axs_bot[sid].set_title(status.replace(\" \", \"\\n\") + (r' ($\\downarrow$)' if sid else r' ($\\uparrow$)'), fontsize=fontsize+1.5)\n",
    "        \n",
    "    for fid, family in enumerate(families):        \n",
    "        average_memory = average_data[(average_data['version']=='memory') & (average_data['family']==family)][status].to_numpy()\n",
    "        average_context = average_data[(average_data['version']=='context') & (average_data['family']==family)][status].to_numpy()\n",
    "        average_max = np.max((average_memory, average_context), axis=0)\n",
    "        diffs = average_context - average_memory\n",
    "        hatches = ['//' if diff >= 0 else 'xx' for diff in diffs]\n",
    "        \n",
    "        models = list(abbre2model.keys())[fid*num_bar:(fid+1)*num_bar]\n",
    "        axs_bot[sid].bar(fid+xs, average_max*100, label=models, \n",
    "                         width=width, color=colors_bot[fid*num_bar:(fid+1)*num_bar], edgecolor='black', lw=lw)\n",
    "        axs_bot[sid].bar(fid+xs, diffs*100, bottom=average_memory*100, \n",
    "                         width=width, facecolor='none', edgecolor='black', hatch=hatches, lw=lw)\n",
    "        for x, top, diff in zip(xs, average_context, diffs):\n",
    "            axs_bot[sid].text(fid+x, (sid2top[sid]*0.02+max(top,top-diff))*100, f'{diff*100:+.1f}', \n",
    "                              fontsize=fontsize-1, ha='center', va='bottom', rotation='vertical')\n",
    "\n",
    "axs_bot[0].set_ylabel(f'Percentage', fontsize=fontsize+0.5)\n",
    "handles, labels = axs_bot[0].get_legend_handles_labels()\n",
    "empty = Patch(facecolor='none', edgecolor='none')\n",
    "handles = handles[:3] + [empty] + handles[3:6] + [empty] + handles[6:] + [empty]\n",
    "labels = labels[:3] + [''] + labels[3:6] + [''] + labels[6:] + ['']\n",
    "handles += [\n",
    "    Patch(facecolor='white', edgecolor='black', hatch='//'),\n",
    "    Patch(facecolor='white', edgecolor='black', hatch='xx')\n",
    "]\n",
    "labels += [\n",
    "    'Increase w/ Context', 'Decrease w/ Context'\n",
    "]\n",
    "subfigs[1].legend(handles, labels, fontsize=fontsize, ncol=7,\n",
    "                  bbox_to_anchor=(0.5,0), loc='upper center', bbox_transform=subfigs[1].transSubfigure)\n",
    "\n",
    "plt.savefig('Figure/ave_shift.pdf', bbox_inches='tight')\n",
    "plt.show()\n",
    "plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5d72304",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 432724 characters]",
      "text/plain": [
       "<Figure size 1400x900 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def draw_subfig(sfid, version, data2model2status):\n",
    "    \n",
    "    axs = subfigs[sfid].subplots(1, num_col, gridspec_kw={'wspace':0.04})\n",
    "    subfigs[sfid].suptitle(f\"({'b' if sfid else 'a'}) Distribution of {version} Knowledge Statuses (Multi-Choice Setting, {'No' if version == 'Parametric' else 'Gold'} Context)\",\n",
    "                           fontsize=fontsize+3.5)\n",
    "    \n",
    "    for idx, data in enumerate(datas):\n",
    "        status = np.array([data2model2status[data].get(model, np.zeros(len(statuses))) for model in abbre2model])\n",
    "        sns.heatmap(status*100, ax=axs[idx], linewidth=0.5, vmin=vmin, vmax=vmax, \n",
    "                    yticklabels=idx==0, annot=True, fmt=\".1f\", annot_kws={'fontsize':fontsize-0.5}, \n",
    "                    cmap=\"Reds\", cbar=idx==len(datas)-1, cbar_kws={'label':'Percentage'})\n",
    "        axs[idx].set_title(data, fontsize=fontsize+1.5)\n",
    "        axs[idx].set_xticklabels(xticklabels, fontsize=fontsize, ha='right', va='center', \n",
    "                                 rotation=90, rotation_mode=\"anchor\")\n",
    "\n",
    "    axs[0].set_yticklabels(abbre2model.keys(), rotation=0, ha='right', va='center', fontsize=fontsize)\n",
    "    cbar = axs[-1].collections[0].colorbar\n",
    "    cbar.ax.tick_params(labelsize=fontsize)\n",
    "    cbar.ax.yaxis.label.set_size(fontsize+1)\n",
    "    cbar.ax.yaxis.labelpad = -1\n",
    "\n",
    "fontsize = 12\n",
    "vmin, vmax = 0, 100\n",
    "num_col, num_row = len(datas), 2\n",
    "sub_col, sub_row = len(statuses), len(abbre2model)\n",
    "size_col, size_row = 0.7, 0.5\n",
    "\n",
    "fig = plt.figure(figsize=(num_col*sub_col*size_col, num_row*sub_row*size_row), constrained_layout=True)\n",
    "subfigs = fig.subfigures(nrows=num_row, ncols=1, height_ratios=[1, 1])\n",
    "xticklabels = ['Consistent\\nCorrect', 'Conflicting\\nCorrect', 'Absent', 'Conflicting\\nWrong', 'Consistent\\nWrong']\n",
    "\n",
    "draw_subfig(0, 'Parametric', data2model2memory)\n",
    "draw_subfig(1, 'Contextual', data2model2context)\n",
    "\n",
    "plt.savefig('Figure/status_dist.pdf', bbox_inches='tight')\n",
    "plt.show()\n",
    "plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45fe23c3",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 216128 characters]",
      "text/plain": [
       "<Figure size 750x780 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 198956 characters]",
      "text/plain": [
       "<Figure size 750x780 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 209176 characters]",
      "text/plain": [
       "<Figure size 750x780 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 191376 characters]",
      "text/plain": [
       "<Figure size 750x780 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def draw_switch(data):\n",
    "    fontsize = 10 # 11.5\n",
    "    vmin, vmax = 0, 100\n",
    "    num_col, num_row = 3, 3\n",
    "    sub_col, sub_row = len(statuses), len(statuses)\n",
    "    size_col, size_row = 0.5, 0.52 # 0.65, 0.65\n",
    "    fig, axs = plt.subplots(num_row, num_col, constrained_layout=True, gridspec_kw={'wspace':0.04, 'hspace':0.01},\n",
    "                            figsize=(num_col*sub_col*size_col, num_row*sub_row*size_row))\n",
    "\n",
    "    for idx, model in enumerate(abbre2model):\n",
    "        idx_r, idx_c = idx//num_col, idx%num_col\n",
    "        sns.heatmap(100*data2model2switch[data].get(model, np.zeros((len(statuses), len(statuses)))), \n",
    "                    ax=axs[idx_r][idx_c], linewidth=0.5, vmin=vmin, vmax=vmax, \n",
    "                    yticklabels=idx_c==0, xticklabels=idx_r==num_row-1, \n",
    "                    annot=True, fmt=\".1f\", annot_kws={'fontsize':fontsize-1.75}, \n",
    "                    cmap=\"Reds\", cbar=idx_c==num_col-1, cbar_kws={'label':'Percentage'})\n",
    "        axs[idx_r][idx_c].set_title(model, fontsize=fontsize-1)\n",
    "\n",
    "        if idx_r == num_row-1:\n",
    "            axs[idx_r][idx_c].set_xlabel('Contextual Knowledge Status', fontsize=fontsize-1.5)\n",
    "            axs[idx_r][idx_c].set_xticklabels(['\\n'.join(c.split(' ')) for c in statuses], rotation=90, ha='right', \n",
    "                                              va='center', rotation_mode='anchor', fontsize=fontsize-2)\n",
    "\n",
    "        if idx_c == 0:\n",
    "            axs[idx_r][idx_c].set_ylabel('Parametric Knowledge Status', fontsize=fontsize-1.5)\n",
    "            axs[idx_r][idx_c].set_yticklabels(['\\n'.join(c.split(' ')) for c in statuses], rotation=0, ha='right', \n",
    "                                              va='center', fontsize=fontsize-2)\n",
    "\n",
    "        if idx_c == num_col-1:\n",
    "            cbar = axs[idx_r][idx_c].collections[0].colorbar\n",
    "            cbar.ax.tick_params(labelsize=fontsize-2)\n",
    "            cbar.ax.yaxis.label.set_size(fontsize-0.5)\n",
    "\n",
    "    fig.suptitle(f'Context-Induced Shifts (%) in LLM Knowledge Statuses on {data}\\n(Multi-Choice Setting, Gold Context)', fontsize=fontsize)\n",
    "    plt.savefig(f'Figure/switch_{data}.pdf', bbox_inches='tight')\n",
    "    plt.show()\n",
    "    plt.close(fig)\n",
    "    \n",
    "for data in datas:\n",
    "    draw_switch(data)"
   ]
  }
 ],
 "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
