{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from read_results import main as read_results\n",
    "from pathlib import Path\n",
    "import pandas as pd\n",
    "import matplotlib\n",
    "import numpy as np\n",
    "matplotlib.rcParams.update({'font.size': 12})\n",
    "save= False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[20, 20, 20, 20, 19, 18]\n",
      "[20, 20, 20, 20, 19, 18]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 41036 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "from read_logs import read_log\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "pd.options.display.max_colwidth = 200\n",
    "matplotlib.rcParams.update({'font.size': 15})\n",
    "\n",
    "\n",
    "all_n_components = [4, 8, 10, 15, 25, 30]\n",
    "all_probe_dfs = []\n",
    "all_prompt_dfs = []\n",
    "for n_components in all_n_components:\n",
    "    all_res = []\n",
    "    for seed in list(range(1,21)):\n",
    "        try:\n",
    "            all_res += read_log(f'logs/trans0.01_emit0.01_nsym10_NCOM{n_components}_nexamples5000_seed{seed}')[-2:]\n",
    "        except:\n",
    "            print(f\"Skipped: logs/trans0.01_emit0.01_nsym10_NCOM{n_components}_nexamples5000_seed{seed}\")\n",
    "            pass\n",
    "    df = pd.DataFrame(all_res)\n",
    "    prompt_df = df[['prompt_tune' in name for name in df['name']]].reset_index()\n",
    "    probe_df = df[['probe' in name for name in df['name']]].reset_index()\n",
    "    all_prompt_dfs.append(prompt_df)\n",
    "    all_probe_dfs.append(probe_df)\n",
    "\n",
    "head_means = np.asarray([df['test_acc'].mean() for df in all_probe_dfs])\n",
    "head_intervals = np.asarray([df['test_acc'].std() * 1.96 / np.sqrt(len(df)) for df in all_probe_dfs])\n",
    "\n",
    "prompt_means = np.asarray([df['test_acc'].mean() for df in all_prompt_dfs])\n",
    "prompt_intervals = np.asarray([df['test_acc'].std() * 1.96 / np.sqrt(len(df)) for df in all_prompt_dfs])\n",
    "\n",
    "print([len(df) for df in all_probe_dfs])\n",
    "print([len(df) for df in all_prompt_dfs])\n",
    "\n",
    "\n",
    "plt.plot(all_n_components, prompt_means, label='Prompt Tuning', markersize=10, marker='^', color='C0')\n",
    "plt.plot(all_n_components, head_means, label='Head Tuning', markersize=10, marker='o',  color='C1')\n",
    "plt.fill_between(all_n_components, (prompt_means-prompt_intervals), (prompt_means+prompt_intervals), color='C0', alpha=.1)\n",
    "plt.fill_between(all_n_components, (head_means-head_intervals), (head_means+head_intervals), color='C1', alpha=.1)\n",
    "\n",
    "plt.axvline(x=10, label='Vocab size', c='gray', linestyle='dashed')\n",
    "plt.legend(loc='upper right')\n",
    "# plt.legend()\n",
    "plt.xlabel('# of hidden states')\n",
    "plt.ylabel('Test Accuracy')\n",
    "plt.xticks([8, 10, 15, 20, 25, 30])\n",
    "\n",
    "plt.savefig('prompt_vs_head_tuning.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32464 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "matplotlib.rcParams.update({'font.size': 14})\n",
    "\n",
    "\n",
    "all_vanilla_dfs = []\n",
    "all_mem_slots = [2,3,5,7]\n",
    "for n_mem_slots in all_mem_slots:\n",
    "    all_res = []\n",
    "    for seed in [1,2,3,4,5]:\n",
    "        try:\n",
    "            all_res += read_log(f'logs/VANILLA_trans0.01_emit0.01_nsym10_nexamples5000_nhid4_nmem{n_mem_slots}_seed{seed}')\n",
    "        except:\n",
    "            pass\n",
    "    \n",
    "    df = pd.DataFrame(all_res)\n",
    "    probe_df = df[[('memory' not in name and 'probe' in name) for name in df['name']]].reset_index()\n",
    "    all_vanilla_dfs.append(probe_df)\n",
    "\n",
    "head_means = np.asarray([df['test_acc'].mean() for df in all_vanilla_dfs])\n",
    "head_intervals = np.asarray([df['test_acc'].std() * 1.96 / np.sqrt(len(df)) for df in all_vanilla_dfs])\n",
    "\n",
    "\n",
    "all_memory_dfs = []\n",
    "all_mem_slots = [2,3,5,7]\n",
    "for n_mem_slots in all_mem_slots:\n",
    "    all_res = []\n",
    "    for seed in [1,2,3,4,5]:\n",
    "        try:\n",
    "            all_res += read_log(f'logs/MEMORY_trans0.01_emit0.01_nsym10_nexamples5000_nhid4_nmem{n_mem_slots}_seed{seed}')\n",
    "        except:\n",
    "            pass\n",
    "    \n",
    "    df = pd.DataFrame(all_res)\n",
    "    probe_df = df[[('probe' in name) for name in df['name']]].reset_index()\n",
    "    all_memory_dfs.append(probe_df)\n",
    "\n",
    "memory_means = np.asarray([df['test_acc'].mean() for df in all_memory_dfs])\n",
    "memory_intervals = np.asarray([df['test_acc'].std() * 1.96 / np.sqrt(len(df)) for df in all_memory_dfs])\n",
    "\n",
    "\n",
    "    \n",
    "    \n",
    "plt.plot(np.asarray(all_mem_slots)*4, head_means, label='Vanilla HMM + linear', markersize=10, marker='o',  color='C3')\n",
    "plt.fill_between(np.asarray(all_mem_slots)*4, (head_means-head_intervals), (head_means+head_intervals), color='C3', alpha=.1)\n",
    "plt.plot(np.asarray(all_mem_slots)*4, memory_means, label='Memory HMM + attention', markersize=10, marker='o',  color='C2')\n",
    "plt.fill_between(np.asarray(all_mem_slots)*4, (memory_means-memory_intervals), (memory_means+memory_intervals), color='C2', alpha=.1)\n",
    "plt.axvline(x=10, label='Vocab size', c='gray', linestyle='dashed')\n",
    "plt.legend(loc='center right', bbox_to_anchor=(1.0, 0.65))\n",
    "# plt.legend()\n",
    "plt.xlabel('# of hidden states')\n",
    "plt.ylabel('Test Accuracy')\n",
    "# plt.xticks([8, 10, 15, 20, 25, 30])\n",
    "plt.savefig('memory_plot.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
