{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "72c7c08f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import warnings\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from collections import Counter\n",
    "from collections import defaultdict\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "mpl.rcParams['figure.dpi'] = 100\n",
    "mpl.rcParams['font.family'] = 'DeJavu Serif'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0c9332f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "statuses = ['Consistent Correct', 'Conflicting Correct', 'Absent', 'Conflicting Wrong', 'Consistent Wrong']\n",
    "settings = [\n",
    "    'Parametric Knowledge\\n(Multi-Choice Setting,\\nNo Context)',\n",
    "    'Contextual Knowledge\\n(Multi-Choice Setting,\\nGold Context)',\n",
    "    'Contextual Knowledge\\n(Multi-Choice Setting,\\nNoisy Context)',\n",
    "    'Parametric Knowledge\\n(Open-Ended Setting,\\nNo Context)',\n",
    "    'Contextual Knowledge\\n(Open-Ended Setting,\\nGold Context)',\n",
    "]\n",
    "setting2freq = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "984c17d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "num_opt = 3\n",
    "dataset = pd.read_csv('Data/Input/HotpotQA.csv')\n",
    "size = len(dataset)\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",
    "\n",
    "cat2freq = defaultdict(list)\n",
    "for abbre, model in abbre2model.items():\n",
    "    mcq = json.load(open(f'Data/Status/HotpotQA/{model}.json'))\n",
    "    for cat in ['memory', 'context']:\n",
    "        freq = Counter(find_status_all(mcq, cat))\n",
    "        cat2freq[cat].append([freq[status] for status in statuses])\n",
    "setting2freq['Parametric Knowledge (Multi-Choice Setting, No Context)'] = np.stack(cat2freq['memory']) / size * 100\n",
    "setting2freq['Contextual Knowledge (Multi-Choice Setting, Gold Context)'] = np.stack(cat2freq['context']) / size * 100\n",
    "\n",
    "noisy2freq = []\n",
    "for abbre, model in abbre2model.items():\n",
    "    path = f'Data/Rebuttal/Noisy/{model}_status.json'\n",
    "    noisy = json.load(open(path))\n",
    "    freq = Counter(find_status_all(noisy, 'context'))\n",
    "    noisy2freq.append([freq[status] for status in statuses])\n",
    "setting2freq['Contextual Knowledge (Multi-Choice Setting, Noisy Context)'] = np.stack(noisy2freq) / size * 100\n",
    "\n",
    "openend2freq = defaultdict(list)\n",
    "for abbre, model in abbre2model.items():\n",
    "    path = f'Data/Rebuttal/Open/{model}_status.json'\n",
    "    openend = json.load(open(path))\n",
    "    for cat in ['memory', 'context']:\n",
    "        freq = Counter([openend[str(idx)][cat] for idx in range(len(openend))])\n",
    "        openend2freq[cat].append([freq[status] for status in statuses])\n",
    "setting2freq['Parametric Knowledge (Open-Ended Setting, No Context)'] = np.stack(openend2freq['memory']) / size * 100\n",
    "setting2freq['Contextual Knowledge (Open-Ended Setting, Gold Context)'] = np.stack(openend2freq['context']) / size * 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dea3b87",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 260892 characters]",
      "text/plain": [
       "<Figure size 1750x468 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fontsize = 12\n",
    "vmin, vmax = 0, 100\n",
    "num_col, num_row = 5, 1\n",
    "sub_col, sub_row = len(statuses), len(abbre2model)\n",
    "size_col, size_row = 0.7, 0.52\n",
    "\n",
    "fig, axs = plt.subplots(num_row, num_col, figsize=(num_col*sub_col*size_col, num_row*sub_row*size_row), constrained_layout=True)\n",
    "\n",
    "xticklabels = ['Consistent\\nCorrect', 'Conflicting\\nCorrect', 'Absent', 'Conflicting\\nWrong', 'Consistent\\nWrong']\n",
    "\n",
    "for sid, setting in enumerate(settings):\n",
    "    sns.heatmap(setting2freq[setting.replace('\\n', ' ')], ax=axs[sid], linewidth=0.5, vmin=vmin, vmax=vmax,\n",
    "                yticklabels=sid==0, annot=True, fmt=\".1f\", annot_kws={'fontsize':fontsize-0.5},\n",
    "                cmap=\"Reds\", cbar=sid==len(setting2freq)-1, cbar_kws={'label':'Percentage'})\n",
    "    axs[sid].set_title(f'({chr(ord(\"a\")+sid)}) ' + setting, fontsize=fontsize+1)\n",
    "    axs[sid].set_xticklabels(xticklabels, fontsize=fontsize, ha='right', va='center', rotation=90, rotation_mode=\"anchor\")\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+0.75)\n",
    "cbar.ax.yaxis.labelpad = -1\n",
    "\n",
    "plt.savefig('Figure/status_other.pdf', bbox_inches='tight')\n",
    "plt.show()\n",
    "plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8b7df72",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "num_opt = 3\n",
    "dataset = pd.read_csv('Data/Input/HotpotQA.csv')\n",
    "size = len(dataset)\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",
    "\n",
    "abbre2transition_noisy = {abbre: np.zeros((len(statuses), len(statuses))) for abbre in abbre2model.keys()}\n",
    "for abbre, model in abbre2model.items():\n",
    "    mcq = json.load(open(f'Data/Status/HotpotQA/{model}.json'))\n",
    "    cat2status = {'memory': find_status_all(mcq, 'memory')}\n",
    "\n",
    "    path = f'Data/Rebuttal/Noisy/{model}_status.json'\n",
    "    noisy = json.load(open(path))\n",
    "    cat2status['context'] = find_status_all(noisy, 'context')\n",
    "    cat2status = pd.DataFrame(cat2status)\n",
    "    status_freq = cat2status['memory'].value_counts().to_dict()\n",
    "    transition_freq = cat2status.groupby(['memory', 'context']).size().to_dict()\n",
    "\n",
    "    for i, status1 in enumerate(statuses):\n",
    "        if status1 in status_freq:\n",
    "            for j, status2 in enumerate(statuses):\n",
    "                if (status1, status2) in transition_freq:\n",
    "                    abbre2transition_noisy[abbre][i, j] = transition_freq[(status1, status2)] / status_freq[status1] * 100\n",
    "\n",
    "abbre2transition_open = {abbre: np.zeros((len(statuses), len(statuses))) for abbre in abbre2model.keys()}\n",
    "for abbre, model in abbre2model.items():\n",
    "    path = f'Data/Rebuttal/Open/{model}_status.json'\n",
    "    openend = json.load(open(path))\n",
    "    cat2status = pd.DataFrame.from_dict(openend, orient='index')\n",
    "    status_freq = cat2status['memory'].value_counts().to_dict()\n",
    "    transition_freq = cat2status.groupby(['memory', 'context']).size().to_dict()\n",
    "\n",
    "    for i, status1 in enumerate(statuses):\n",
    "        if status1 in status_freq:\n",
    "            for j, status2 in enumerate(statuses):\n",
    "                if (status1, status2) in transition_freq:\n",
    "                    abbre2transition_open[abbre][i, j] = transition_freq[(status1, status2)] / status_freq[status1] * 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b658bb0b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 216020 characters]",
      "text/plain": [
       "<Figure size 750x780 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 201828 characters]",
      "text/plain": [
       "<Figure size 750x780 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def draw_switch(abbre2transition, setting):\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(abbre2transition[model], \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(setting2title[setting], fontsize=fontsize)\n",
    "    plt.savefig(f'Figure/switch_{setting}.pdf', bbox_inches='tight')\n",
    "    plt.show()\n",
    "    plt.close(fig)\n",
    "\n",
    "\n",
    "setting2title = {\n",
    "    'noisy': 'Context-Induced Shifts (%) in LLM Knowledge Statuses on HotpotQA\\n(Multi-Choice Setting, Noisy Context)',\n",
    "    'open': 'Context-Induced Shifts (%) in LLM Knowledge Statuses on HotpotQA\\n(Open-Ended Setting, Gold Context)'\n",
    "}\n",
    "\n",
    "draw_switch(abbre2transition_noisy, 'noisy')\n",
    "draw_switch(abbre2transition_open, 'open')"
   ]
  }
 ],
 "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
}
