{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "138889b92720ce2e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-14T09:06:04.487186Z",
     "start_time": "2024-05-14T09:06:04.255111Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>runname</th>\n",
       "      <th>steps</th>\n",
       "      <th>agg_score</th>\n",
       "      <th>commonsense_qa/acc</th>\n",
       "      <th>commonsense_qa/acc_norm</th>\n",
       "      <th>hellaswag/acc</th>\n",
       "      <th>hellaswag/acc_norm</th>\n",
       "      <th>openbookqa/acc</th>\n",
       "      <th>openbookqa/acc_norm</th>\n",
       "      <th>piqa/acc</th>\n",
       "      <th>...</th>\n",
       "      <th>siqa/acc</th>\n",
       "      <th>siqa/acc_norm</th>\n",
       "      <th>winogrande/acc</th>\n",
       "      <th>winogrande/acc_norm</th>\n",
       "      <th>sciq/acc</th>\n",
       "      <th>sciq/acc_norm</th>\n",
       "      <th>arc/acc</th>\n",
       "      <th>arc/acc_norm</th>\n",
       "      <th>mmlu/acc</th>\n",
       "      <th>mmlu/acc_norm</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>C4</td>\n",
       "      <td>0</td>\n",
       "      <td>0.330893</td>\n",
       "      <td>0.186</td>\n",
       "      <td>0.233</td>\n",
       "      <td>0.272</td>\n",
       "      <td>0.258</td>\n",
       "      <td>0.166</td>\n",
       "      <td>0.286</td>\n",
       "      <td>0.542</td>\n",
       "      <td>...</td>\n",
       "      <td>0.367</td>\n",
       "      <td>0.362</td>\n",
       "      <td>0.516</td>\n",
       "      <td>0.497</td>\n",
       "      <td>0.208</td>\n",
       "      <td>0.202</td>\n",
       "      <td>0.2195</td>\n",
       "      <td>0.2510</td>\n",
       "      <td>0.230294</td>\n",
       "      <td>0.250147</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>C4</td>\n",
       "      <td>1000</td>\n",
       "      <td>0.355112</td>\n",
       "      <td>0.229</td>\n",
       "      <td>0.260</td>\n",
       "      <td>0.286</td>\n",
       "      <td>0.288</td>\n",
       "      <td>0.128</td>\n",
       "      <td>0.250</td>\n",
       "      <td>0.614</td>\n",
       "      <td>...</td>\n",
       "      <td>0.351</td>\n",
       "      <td>0.404</td>\n",
       "      <td>0.519</td>\n",
       "      <td>0.476</td>\n",
       "      <td>0.565</td>\n",
       "      <td>0.518</td>\n",
       "      <td>0.2680</td>\n",
       "      <td>0.2935</td>\n",
       "      <td>0.238951</td>\n",
       "      <td>0.250399</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>C4</td>\n",
       "      <td>2000</td>\n",
       "      <td>0.378435</td>\n",
       "      <td>0.268</td>\n",
       "      <td>0.278</td>\n",
       "      <td>0.312</td>\n",
       "      <td>0.330</td>\n",
       "      <td>0.122</td>\n",
       "      <td>0.276</td>\n",
       "      <td>0.646</td>\n",
       "      <td>...</td>\n",
       "      <td>0.375</td>\n",
       "      <td>0.400</td>\n",
       "      <td>0.509</td>\n",
       "      <td>0.500</td>\n",
       "      <td>0.676</td>\n",
       "      <td>0.577</td>\n",
       "      <td>0.3065</td>\n",
       "      <td>0.3230</td>\n",
       "      <td>0.247275</td>\n",
       "      <td>0.255482</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>C4</td>\n",
       "      <td>3000</td>\n",
       "      <td>0.387795</td>\n",
       "      <td>0.280</td>\n",
       "      <td>0.295</td>\n",
       "      <td>0.331</td>\n",
       "      <td>0.380</td>\n",
       "      <td>0.152</td>\n",
       "      <td>0.274</td>\n",
       "      <td>0.660</td>\n",
       "      <td>...</td>\n",
       "      <td>0.376</td>\n",
       "      <td>0.387</td>\n",
       "      <td>0.512</td>\n",
       "      <td>0.496</td>\n",
       "      <td>0.725</td>\n",
       "      <td>0.621</td>\n",
       "      <td>0.3175</td>\n",
       "      <td>0.3340</td>\n",
       "      <td>0.254534</td>\n",
       "      <td>0.267363</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>C4</td>\n",
       "      <td>4000</td>\n",
       "      <td>0.399320</td>\n",
       "      <td>0.296</td>\n",
       "      <td>0.298</td>\n",
       "      <td>0.351</td>\n",
       "      <td>0.406</td>\n",
       "      <td>0.168</td>\n",
       "      <td>0.282</td>\n",
       "      <td>0.676</td>\n",
       "      <td>...</td>\n",
       "      <td>0.382</td>\n",
       "      <td>0.404</td>\n",
       "      <td>0.522</td>\n",
       "      <td>0.503</td>\n",
       "      <td>0.723</td>\n",
       "      <td>0.618</td>\n",
       "      <td>0.3255</td>\n",
       "      <td>0.3470</td>\n",
       "      <td>0.254762</td>\n",
       "      <td>0.263563</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1171</th>\n",
       "      <td>The Pile</td>\n",
       "      <td>163000</td>\n",
       "      <td>0.463789</td>\n",
       "      <td>0.379</td>\n",
       "      <td>0.349</td>\n",
       "      <td>0.441</td>\n",
       "      <td>0.555</td>\n",
       "      <td>0.240</td>\n",
       "      <td>0.366</td>\n",
       "      <td>0.701</td>\n",
       "      <td>...</td>\n",
       "      <td>0.405</td>\n",
       "      <td>0.388</td>\n",
       "      <td>0.585</td>\n",
       "      <td>0.560</td>\n",
       "      <td>0.875</td>\n",
       "      <td>0.820</td>\n",
       "      <td>0.4475</td>\n",
       "      <td>0.4450</td>\n",
       "      <td>0.299378</td>\n",
       "      <td>0.326313</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1172</th>\n",
       "      <td>The Pile</td>\n",
       "      <td>164000</td>\n",
       "      <td>0.462758</td>\n",
       "      <td>0.369</td>\n",
       "      <td>0.344</td>\n",
       "      <td>0.438</td>\n",
       "      <td>0.552</td>\n",
       "      <td>0.248</td>\n",
       "      <td>0.348</td>\n",
       "      <td>0.708</td>\n",
       "      <td>...</td>\n",
       "      <td>0.395</td>\n",
       "      <td>0.401</td>\n",
       "      <td>0.577</td>\n",
       "      <td>0.567</td>\n",
       "      <td>0.874</td>\n",
       "      <td>0.806</td>\n",
       "      <td>0.4465</td>\n",
       "      <td>0.4355</td>\n",
       "      <td>0.302083</td>\n",
       "      <td>0.331563</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1173</th>\n",
       "      <td>The Pile</td>\n",
       "      <td>165000</td>\n",
       "      <td>0.465026</td>\n",
       "      <td>0.383</td>\n",
       "      <td>0.350</td>\n",
       "      <td>0.438</td>\n",
       "      <td>0.553</td>\n",
       "      <td>0.234</td>\n",
       "      <td>0.352</td>\n",
       "      <td>0.707</td>\n",
       "      <td>...</td>\n",
       "      <td>0.400</td>\n",
       "      <td>0.401</td>\n",
       "      <td>0.569</td>\n",
       "      <td>0.556</td>\n",
       "      <td>0.874</td>\n",
       "      <td>0.811</td>\n",
       "      <td>0.4460</td>\n",
       "      <td>0.4455</td>\n",
       "      <td>0.305193</td>\n",
       "      <td>0.331708</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1174</th>\n",
       "      <td>The Pile</td>\n",
       "      <td>166000</td>\n",
       "      <td>0.462349</td>\n",
       "      <td>0.377</td>\n",
       "      <td>0.346</td>\n",
       "      <td>0.440</td>\n",
       "      <td>0.557</td>\n",
       "      <td>0.228</td>\n",
       "      <td>0.346</td>\n",
       "      <td>0.711</td>\n",
       "      <td>...</td>\n",
       "      <td>0.398</td>\n",
       "      <td>0.398</td>\n",
       "      <td>0.572</td>\n",
       "      <td>0.558</td>\n",
       "      <td>0.877</td>\n",
       "      <td>0.811</td>\n",
       "      <td>0.4525</td>\n",
       "      <td>0.4385</td>\n",
       "      <td>0.301952</td>\n",
       "      <td>0.331295</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1175</th>\n",
       "      <td>The Pile</td>\n",
       "      <td>167000</td>\n",
       "      <td>0.464539</td>\n",
       "      <td>0.386</td>\n",
       "      <td>0.354</td>\n",
       "      <td>0.434</td>\n",
       "      <td>0.557</td>\n",
       "      <td>0.232</td>\n",
       "      <td>0.356</td>\n",
       "      <td>0.706</td>\n",
       "      <td>...</td>\n",
       "      <td>0.402</td>\n",
       "      <td>0.402</td>\n",
       "      <td>0.573</td>\n",
       "      <td>0.559</td>\n",
       "      <td>0.867</td>\n",
       "      <td>0.802</td>\n",
       "      <td>0.4475</td>\n",
       "      <td>0.4375</td>\n",
       "      <td>0.301934</td>\n",
       "      <td>0.330810</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>1176 rows × 21 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "       runname   steps  agg_score  commonsense_qa/acc  \\\n",
       "0           C4       0   0.330893               0.186   \n",
       "1           C4    1000   0.355112               0.229   \n",
       "2           C4    2000   0.378435               0.268   \n",
       "3           C4    3000   0.387795               0.280   \n",
       "4           C4    4000   0.399320               0.296   \n",
       "...        ...     ...        ...                 ...   \n",
       "1171  The Pile  163000   0.463789               0.379   \n",
       "1172  The Pile  164000   0.462758               0.369   \n",
       "1173  The Pile  165000   0.465026               0.383   \n",
       "1174  The Pile  166000   0.462349               0.377   \n",
       "1175  The Pile  167000   0.464539               0.386   \n",
       "\n",
       "      commonsense_qa/acc_norm  hellaswag/acc  hellaswag/acc_norm  \\\n",
       "0                       0.233          0.272               0.258   \n",
       "1                       0.260          0.286               0.288   \n",
       "2                       0.278          0.312               0.330   \n",
       "3                       0.295          0.331               0.380   \n",
       "4                       0.298          0.351               0.406   \n",
       "...                       ...            ...                 ...   \n",
       "1171                    0.349          0.441               0.555   \n",
       "1172                    0.344          0.438               0.552   \n",
       "1173                    0.350          0.438               0.553   \n",
       "1174                    0.346          0.440               0.557   \n",
       "1175                    0.354          0.434               0.557   \n",
       "\n",
       "      openbookqa/acc  openbookqa/acc_norm  piqa/acc  ...  siqa/acc  \\\n",
       "0              0.166                0.286     0.542  ...     0.367   \n",
       "1              0.128                0.250     0.614  ...     0.351   \n",
       "2              0.122                0.276     0.646  ...     0.375   \n",
       "3              0.152                0.274     0.660  ...     0.376   \n",
       "4              0.168                0.282     0.676  ...     0.382   \n",
       "...              ...                  ...       ...  ...       ...   \n",
       "1171           0.240                0.366     0.701  ...     0.405   \n",
       "1172           0.248                0.348     0.708  ...     0.395   \n",
       "1173           0.234                0.352     0.707  ...     0.400   \n",
       "1174           0.228                0.346     0.711  ...     0.398   \n",
       "1175           0.232                0.356     0.706  ...     0.402   \n",
       "\n",
       "      siqa/acc_norm  winogrande/acc  winogrande/acc_norm  sciq/acc  \\\n",
       "0             0.362           0.516                0.497     0.208   \n",
       "1             0.404           0.519                0.476     0.565   \n",
       "2             0.400           0.509                0.500     0.676   \n",
       "3             0.387           0.512                0.496     0.725   \n",
       "4             0.404           0.522                0.503     0.723   \n",
       "...             ...             ...                  ...       ...   \n",
       "1171          0.388           0.585                0.560     0.875   \n",
       "1172          0.401           0.577                0.567     0.874   \n",
       "1173          0.401           0.569                0.556     0.874   \n",
       "1174          0.398           0.572                0.558     0.877   \n",
       "1175          0.402           0.573                0.559     0.867   \n",
       "\n",
       "      sciq/acc_norm  arc/acc  arc/acc_norm  mmlu/acc  mmlu/acc_norm  \n",
       "0             0.202   0.2195        0.2510  0.230294       0.250147  \n",
       "1             0.518   0.2680        0.2935  0.238951       0.250399  \n",
       "2             0.577   0.3065        0.3230  0.247275       0.255482  \n",
       "3             0.621   0.3175        0.3340  0.254534       0.267363  \n",
       "4             0.618   0.3255        0.3470  0.254762       0.263563  \n",
       "...             ...      ...           ...       ...            ...  \n",
       "1171          0.820   0.4475        0.4450  0.299378       0.326313  \n",
       "1172          0.806   0.4465        0.4355  0.302083       0.331563  \n",
       "1173          0.811   0.4460        0.4455  0.305193       0.331708  \n",
       "1174          0.811   0.4525        0.4385  0.301952       0.331295  \n",
       "1175          0.802   0.4475        0.4375  0.301934       0.330810  \n",
       "\n",
       "[1176 rows x 21 columns]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "from matplotlib.figure import Figure\n",
    "\n",
    "df = pd.read_csv(\"../src_data/eval_results.csv\")\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b610f43caefdf01",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-14T09:06:04.563945Z",
     "start_time": "2024-05-14T09:06:04.562142Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "initial_id",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-14T09:06:37.927921Z",
     "start_time": "2024-05-14T09:06:37.588025Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "from matplotlib import pyplot as plt\n",
    "metrics = ['agg_score', 'commonsense_qa/acc_norm', 'hellaswag/acc_norm', 'openbookqa/acc_norm', 'piqa/acc_norm',\n",
    "                   'siqa/acc_norm', 'winogrande/acc_norm', 'arc/acc_norm', 'mmlu/acc_norm']\n",
    "\n",
    "def normalize_runname(runname):\n",
    "    return runname.replace(\"/\", \"_\")\n",
    "\n",
    "grouped = (\n",
    "    df.groupby([\"runname\", \"steps\"])\n",
    "    .agg(\n",
    "        {\n",
    "            key: \"mean\" for key in metrics\n",
    "        }\n",
    "    )\n",
    "    .reset_index()\n",
    ")\n",
    "\n",
    "file_id=\"../assets/data/plots/dataset_ablations\"\n",
    "files = {}\n",
    "for metric in metrics:\n",
    "    datas = {}\n",
    "    for name, group in grouped.groupby(\"runname\"):\n",
    "        group = group[[\"steps\", metric]].sort_values(by=\"steps\")\n",
    "        group = group.set_index(\"steps\")\n",
    "        rolling_avg = group\n",
    "        # rolling_avg = group.rolling(window=5).mean()\n",
    "        datas[name] = {\n",
    "            \"x\": (rolling_avg.index * 2048 * 1024 * 1e-9).tolist(),\n",
    "            \"y\": rolling_avg[metric].tolist(),\n",
    "            \"label\": name,\n",
    "        }\n",
    "    # Sort the datata based on the steps\n",
    "    datas = {k: v for k, v in sorted(datas.items(), key=lambda x: -x[1][\"y\"][-1])}\n",
    "    # Create a folder\n",
    "    os.makedirs(f\"{file_id}\", exist_ok=True)\n",
    "    with open(f\"{file_id}/{normalize_runname(metric)}.json\", \"w\") as f:\n",
    "        json.dump({\n",
    "            \"data\": datas,\n",
    "            \"layout\": {\n",
    "                \"title\": {\n",
    "                    \"text\": \"Dataset ablations\"\n",
    "                },\n",
    "            }\n",
    "        }, f)\n",
    "    files[metric] = {\"file\": f\"{normalize_runname(metric)}.json\"}\n",
    "# Create index\n",
    "with open(f\"{file_id}/index.json\", \"w\") as f:\n",
    "    json.dump({\n",
    "        \"files\": files,\n",
    "        \"settings\": {\n",
    "            \"defaultMetric\": \"agg_score\",\n",
    "            \"slider\":{\"min\":0,\"max\":30,\"default\":5}\n",
    "        }\n",
    "    }, f)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "af28ebbd054cdc33",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-04T22:25:33.206952Z",
     "start_time": "2024-05-04T22:25:33.205262Z"
    },
    "collapsed": false
   },
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}