{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "4ea3e1f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "from torch import tensor\n",
    "from glob import glob\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "from numpy import nan\n",
    "from numpy import array\n",
    "import numpy as np\n",
    "def load_eval_file(file):\n",
    "    try:\n",
    "        line = eval(open(file, \"r\").readlines()[-1].strip())\n",
    "        return line[\"discriminator_span_acc\"].item()\n",
    "    except:\n",
    "        return None\n",
    "\n",
    "def load_downstream_file(file, metric=\"acc\"):\n",
    "    lines = open(file, \"r\").readlines()\n",
    "    for line in lines:\n",
    "        if \"Best score\" in line:\n",
    "            re = line.strip().split(\" \")[-1]\n",
    "            return float(re)\n",
    "    return None\n",
    "\n",
    "def load_downstream_fewshot_file(file, metric=\"acc\"):\n",
    "    lines = open(file, \"r\").readlines()\n",
    "    line_num = 0\n",
    "    for i, line in enumerate(lines):\n",
    "        if \"*** Eval Results ***\" in line:\n",
    "            line_num = i\n",
    "    for line in lines[line_num:]:\n",
    "        if \"fewshot_validation:\" in line:\n",
    "            re = line.strip().split(\" \")[-1]\n",
    "            return float(re)\n",
    "        if \"fewshot_validation_matched\" in line:\n",
    "            re1 = line.strip().split(\" \")[-1]\n",
    "        if \"fewshot_validation_mismatched\" in line:\n",
    "            re2 = line.strip().split(\" \")[-1]\n",
    "            return float(re1), float(re2)\n",
    "    return None\n",
    "\n",
    "def load_best_step(file):\n",
    "    lines = open(file, \"r\").readlines()\n",
    "    b_step = None\n",
    "    for line in lines:\n",
    "        if \"Best score\" in line:\n",
    "            step = int(line.strip().split(\" \")[3][:-1])\n",
    "            if step % 10 == 0:\n",
    "                b_step = step\n",
    "    return b_step\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "4b654568",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "best eval: 38000 0.5152846574783325\n",
      "2 epochs: 0.5152846574783325\n",
      "max: 12000 0.5339556932449341\n"
     ]
    }
   ],
   "source": [
    "# COPA v2, v3\n",
    "task_to_aug_task = {\"qnli\": \"nli_v2_aug\", \"snli\": \"nli_v2_aug\", \"rte\": \"nli_v2_aug\", \"mnli\": \"nli_v2_aug\", \"sst2\": \"sst2_v2_aug\", \"mr\": \"sst2_v2_aug\", \"imdb\": \"sst2_v2_aug\", \"piqa\": \"copa-retrieval0.15\", \"hellaswag\": \"copa-retrieval0.15\"}\n",
    "task=\"qnli\"\n",
    "aug_task = task_to_aug_task[task]\n",
    "dis_res = []\n",
    "model=\"google/electra-base-discriminator\"\n",
    "# model=\"bert-base-uncased\"\n",
    "if \"electra\" in model:\n",
    "    objective = \"dis\"\n",
    "else:\n",
    "    objective = \"mlm\"\n",
    "ex_name=f\"trainfull\"\n",
    "m = 0\n",
    "m_step = 0\n",
    "base_dir=f\"/home/ubuntu/space3/out/hf/train/{objective}/{model}/{aug_task}/{ex_name}\"\n",
    "b_eval_step=load_best_step(os.path.join(base_dir, f\"log.out\"))\n",
    "max_step = 38000\n",
    "for step in range(1000, max_step+1000, 1000):\n",
    "    output=f\"{base_dir}/checkpoint-{step}\"\n",
    "    file = os.path.join(output, f\"eval_{task}.log\")\n",
    "    re = load_downstream_fewshot_file(file)\n",
    "    dis_res.append(re)\n",
    "    if step == b_eval_step:\n",
    "        print(\"best eval:\", b_eval_step, re)\n",
    "    if re is not None:\n",
    "        re = re if isinstance(re, float) else re[0]\n",
    "    if re is not None and re > m:\n",
    "        m = re\n",
    "        m_step = step\n",
    "        \n",
    "print(\"2 epochs:\", dis_res[-1])\n",
    "print(\"max:\", m_step, m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "56284b77",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.3584309730005094, 0.354759967453214)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dis_res[4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "848b0d70",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32712 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "task = \"sst5\"\n",
    "model_task=\"sst5\"\n",
    "roberta_acc = {\"sst2\": 77.75, \"imdb\": 63.22, \"rte\": 53.43, \"sst5\": 30.25, \"trec\": 45.00, \"subj\": 53.80, \"mr\": 71.85, \"mnlimatched\": 48.23}\n",
    "electra_acc = {\"sst2\": 82.8, \"imdb\": 82.44, \"rte\": 57.76, \"sst5\": 31.06, \"trec\": 35.80, \"subj\": 50.05, \"mr\": 81.45, \"mnlimatched\": 51.82}\n",
    "steps = {\"sst2\": 23500, \"nli\": 5000, \"sst5\": 5000, \"trec\": 600, \"subj\": 900}\n",
    "\n",
    "x = list(range(500, steps[model_task], 500))\n",
    "\n",
    "plt.plot(range(len(mlm_res)), mlm_res, label=\"trained roberta\")\n",
    "# plt.plot(range(len(dis_res)), dis_res, label=\"trained electra\")\n",
    "# plt.ylim([0.45, 0.55])\n",
    "if roberta_acc[task] is not None:\n",
    "    plt.axhline(roberta_acc[task], label=\"roberta base\", color=\"green\")\n",
    "# if electra_acc[task] is not None:\n",
    "#     plt.axhline(electra_acc[task], label=\"electra base\", color=\"red\")\n",
    "\n",
    "plt.legend()\n",
    "plt.title(f\"Test: {task}\")\n",
    "plt.savefig(f\"images/new_images/{task}.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3a245e4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "updated",
   "language": "python",
   "name": "updated"
  },
  "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
