{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f29da5e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "import os\n",
    "import torch\n",
    "%matplotlib inline\n",
    "from sklearn.preprocessing import normalize\n",
    "from sklearn.decomposition import PCA\n",
    "from datasets import load_dataset\n",
    "from transformers import AutoModelForPreTraining, AutoTokenizer\n",
    "from sklearn.preprocessing import normalize\n",
    "import numpy as np\n",
    "import matplotlib.ticker as ticker\n",
    "import seaborn as sns\n",
    "import sys\n",
    "sys.path.append(\"/n/fs/nlp-mengzhou/space3/hf\")\n",
    "\n",
    "from transformers import AutoModelForMaskedLM\n",
    "from templates import *\n",
    "from utils import *\n",
    "from models import *\n",
    "sns.set_style(\"whitegrid\", {'axes.grid' : False})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6e93e514",
   "metadata": {},
   "outputs": [],
   "source": [
    "task = \"sst2\"\n",
    "additional_info = {}\n",
    "names = {\"bert-base-uncased\": \"bert\", \"roberta-base\": \"roberta\", \"google/electra-base-discriminator\": \"electra\"}\n",
    "for model in [\"bert-base-uncased\", \"roberta-base\", \"google/electra-base-discriminator\"]:\n",
    "    model_dir = f\"/n/fs/nlp-mengzhou/space3/out/hf/analysis_v1/{model}/{task}\"\n",
    "    file = os.path.join(model_dir, \"additional-info.pt\")\n",
    "    if os.path.exists(file):\n",
    "        this_reps = torch.load(file)\n",
    "        additional_info[f\"zeroshot-{names[model]}\"] = this_reps\n",
    "    else:\n",
    "        print(f\"{file} does not exist\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fdd29dd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "sst2_model_dirs = [\"/n/fs/nlp-mengzhou/space3/out/hf/few_shot/sst2/bert-base-uncased/shot16/prompt_FT/1/bs4_lr2e-5\",\n",
    "    \"/n/fs/nlp-mengzhou/space3/out/hf/few_shot/sst2/roberta-base/shot16/prompt_FT/1/bs4_lr1e-5\",\n",
    "    \"/n/fs/nlp-mengzhou/space3/out/hf/few_shot/sst2/google/electra-base-discriminator/shot16/prompt_FT/1/bs8_lr1e-5\",]\n",
    "boolq_model_dirs = [\"/n/fs/nlp-mengzhou/space3/out/hf/few_shot/boolq/bert-base-uncased/shot16/prompt_FT/3/bs2_lr2e-5\",\n",
    "                   \"/n/fs/nlp-mengzhou/space3/out/hf/few_shot/boolq/roberta-base/shot16/prompt_FT/3/bs2_lr2e-5\",\n",
    "                   \"/n/fs/nlp-mengzhou/space3/out/hf/few_shot/boolq/google/electra-base-discriminator/shot16/prompt_FT/3/bs2_lr2e-5\"]\n",
    "sst5_model_dirs = [\"\", \"\", \"/n/fs/nlp-mengzhou/space3/out/hf/few_shot/sst5/google/electra-base-discriminator/shot16/prompt_FT/1/bs8_lr2e-5\"]\n",
    "agnews_model_dirs = [\"\", \"\", \"/n/fs/nlp-mengzhou/space3/out/hf/few_shot/ag_news/google/electra-base-discriminator/shot16/prompt_FT/9/bs8_lr1e-5\"]\n",
    "model_dirs = {\"sst2\": sst2_model_dirs, \"boolq\": boolq_model_dirs, \"sst5\": sst5_model_dirs, \"ag_news\": agnews_model_dirs}\n",
    "for model_dir in model_dirs[task]:\n",
    "    file = os.path.join(model_dir, \"additional-info.pt\")\n",
    "    if os.path.exists(file):\n",
    "        this_reps = torch.load(file)\n",
    "        if \"roberta\" in model_dir:\n",
    "            model = \"roberta\"\n",
    "        elif \"bert\" in model_dir:\n",
    "            model = \"bert\"\n",
    "        elif \"electra\" in model_dir:\n",
    "            model = \"electra\"\n",
    "        additional_info[f\"fewshot-{model}\"] = this_reps\n",
    "    else:\n",
    "        print(f\"{file} does not exist\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 305,
   "id": "8fd84759",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of BertForPreTraining were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['cls.predictions.decoder.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Input \u001b[0;32mIn [305]\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     28\u001b[0m static_reps \u001b[38;5;241m=\u001b[39m {}\n\u001b[1;32m     29\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m model_name \u001b[38;5;129;01min\u001b[39;00m names:\n\u001b[0;32m---> 30\u001b[0m     \u001b[43mget_static_reps\u001b[49m\u001b[43m(\u001b[49m\u001b[43mstatic_reps\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel_name\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mzeroshot\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m     31\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m model_dir \u001b[38;5;129;01min\u001b[39;00m model_dirs:\n\u001b[1;32m     32\u001b[0m     get_static_reps(static_reps, model_dir, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfewshot\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "Input \u001b[0;32mIn [305]\u001b[0m, in \u001b[0;36mget_static_reps\u001b[0;34m(static_reps, model_name, prefix)\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mget_static_reps\u001b[39m(static_reps, model_name, prefix):\n\u001b[1;32m      4\u001b[0m     model \u001b[38;5;241m=\u001b[39m AutoModelForPreTraining\u001b[38;5;241m.\u001b[39mfrom_pretrained(model_name)\n\u001b[0;32m----> 5\u001b[0m     tokenizer \u001b[38;5;241m=\u001b[39m \u001b[43mAutoTokenizer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_pretrained\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel_name\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      6\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mroberta\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m model_name:\n\u001b[1;32m      7\u001b[0m         embeddings \u001b[38;5;241m=\u001b[39m model\u001b[38;5;241m.\u001b[39mstate_dict()[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mroberta.embeddings.word_embeddings.weight\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n",
      "File \u001b[0;32m/n/fs/nlp-mengzhou/anaconda3/envs/updated/lib/python3.8/site-packages/transformers/models/auto/tokenization_auto.py:549\u001b[0m, in \u001b[0;36mAutoTokenizer.from_pretrained\u001b[0;34m(cls, pretrained_model_name_or_path, *inputs, **kwargs)\u001b[0m\n\u001b[1;32m    547\u001b[0m tokenizer_class_py, tokenizer_class_fast \u001b[38;5;241m=\u001b[39m TOKENIZER_MAPPING[\u001b[38;5;28mtype\u001b[39m(config)]\n\u001b[1;32m    548\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m tokenizer_class_fast \u001b[38;5;129;01mand\u001b[39;00m (use_fast \u001b[38;5;129;01mor\u001b[39;00m tokenizer_class_py \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[0;32m--> 549\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtokenizer_class_fast\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_pretrained\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpretrained_model_name_or_path\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    550\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    551\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m tokenizer_class_py \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
      "File \u001b[0;32m/n/fs/nlp-mengzhou/anaconda3/envs/updated/lib/python3.8/site-packages/transformers/tokenization_utils_base.py:1744\u001b[0m, in \u001b[0;36mPreTrainedTokenizerBase.from_pretrained\u001b[0;34m(cls, pretrained_model_name_or_path, *init_inputs, **kwargs)\u001b[0m\n\u001b[1;32m   1741\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1742\u001b[0m         logger\u001b[38;5;241m.\u001b[39minfo(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mloading file \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mfile_path\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m from cache at \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mresolved_vocab_files[file_id]\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m-> 1744\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mcls\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_from_pretrained\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1745\u001b[0m \u001b[43m    \u001b[49m\u001b[43mresolved_vocab_files\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1746\u001b[0m \u001b[43m    \u001b[49m\u001b[43mpretrained_model_name_or_path\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1747\u001b[0m \u001b[43m    \u001b[49m\u001b[43minit_configuration\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1748\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43minit_inputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1749\u001b[0m \u001b[43m    \u001b[49m\u001b[43muse_auth_token\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43muse_auth_token\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1750\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcache_dir\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcache_dir\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1751\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1752\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/n/fs/nlp-mengzhou/anaconda3/envs/updated/lib/python3.8/site-packages/transformers/tokenization_utils_base.py:1879\u001b[0m, in \u001b[0;36mPreTrainedTokenizerBase._from_pretrained\u001b[0;34m(cls, resolved_vocab_files, pretrained_model_name_or_path, init_configuration, use_auth_token, cache_dir, *init_inputs, **kwargs)\u001b[0m\n\u001b[1;32m   1877\u001b[0m \u001b[38;5;66;03m# Instantiate tokenizer.\u001b[39;00m\n\u001b[1;32m   1878\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1879\u001b[0m     tokenizer \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mcls\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43minit_inputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43minit_kwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1880\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mOSError\u001b[39;00m:\n\u001b[1;32m   1881\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mOSError\u001b[39;00m(\n\u001b[1;32m   1882\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mUnable to load vocabulary from file. \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1883\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mPlease check that the provided vocabulary is accessible and not corrupted.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1884\u001b[0m     )\n",
      "File \u001b[0;32m/n/fs/nlp-mengzhou/anaconda3/envs/updated/lib/python3.8/site-packages/transformers/models/bert/tokenization_bert_fast.py:177\u001b[0m, in \u001b[0;36mBertTokenizerFast.__init__\u001b[0;34m(self, vocab_file, tokenizer_file, do_lower_case, unk_token, sep_token, pad_token, cls_token, mask_token, tokenize_chinese_chars, strip_accents, **kwargs)\u001b[0m\n\u001b[1;32m    163\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\n\u001b[1;32m    164\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m    165\u001b[0m     vocab_file\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    175\u001b[0m     \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs\n\u001b[1;32m    176\u001b[0m ):\n\u001b[0;32m--> 177\u001b[0m     \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m    178\u001b[0m \u001b[43m        \u001b[49m\u001b[43mvocab_file\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    179\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtokenizer_file\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtokenizer_file\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    180\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdo_lower_case\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdo_lower_case\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    181\u001b[0m \u001b[43m        \u001b[49m\u001b[43munk_token\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43munk_token\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    182\u001b[0m \u001b[43m        \u001b[49m\u001b[43msep_token\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msep_token\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    183\u001b[0m \u001b[43m        \u001b[49m\u001b[43mpad_token\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mpad_token\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    184\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcls_token\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcls_token\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    185\u001b[0m \u001b[43m        \u001b[49m\u001b[43mmask_token\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmask_token\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    186\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtokenize_chinese_chars\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtokenize_chinese_chars\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    187\u001b[0m \u001b[43m        \u001b[49m\u001b[43mstrip_accents\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstrip_accents\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    188\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    189\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    191\u001b[0m     pre_tok_state \u001b[38;5;241m=\u001b[39m json\u001b[38;5;241m.\u001b[39mloads(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbackend_tokenizer\u001b[38;5;241m.\u001b[39mnormalizer\u001b[38;5;241m.\u001b[39m__getstate__())\n\u001b[1;32m    192\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m (\n\u001b[1;32m    193\u001b[0m         pre_tok_state\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mlowercase\u001b[39m\u001b[38;5;124m\"\u001b[39m, do_lower_case) \u001b[38;5;241m!=\u001b[39m do_lower_case\n\u001b[1;32m    194\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m pre_tok_state\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstrip_accents\u001b[39m\u001b[38;5;124m\"\u001b[39m, strip_accents) \u001b[38;5;241m!=\u001b[39m strip_accents\n\u001b[1;32m    195\u001b[0m     ):\n",
      "File \u001b[0;32m/n/fs/nlp-mengzhou/anaconda3/envs/updated/lib/python3.8/site-packages/transformers/tokenization_utils_fast.py:108\u001b[0m, in \u001b[0;36mPreTrainedTokenizerFast.__init__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    105\u001b[0m     fast_tokenizer \u001b[38;5;241m=\u001b[39m tokenizer_object\n\u001b[1;32m    106\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m fast_tokenizer_file \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m from_slow:\n\u001b[1;32m    107\u001b[0m     \u001b[38;5;66;03m# We have a serialization from tokenizers which let us directly build the backend\u001b[39;00m\n\u001b[0;32m--> 108\u001b[0m     fast_tokenizer \u001b[38;5;241m=\u001b[39m \u001b[43mTokenizerFast\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_file\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfast_tokenizer_file\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    109\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m slow_tokenizer \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    110\u001b[0m     \u001b[38;5;66;03m# We need to convert a slow tokenizer to build the backend\u001b[39;00m\n\u001b[1;32m    111\u001b[0m     fast_tokenizer \u001b[38;5;241m=\u001b[39m convert_slow_tokenizer(slow_tokenizer)\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "label_words = {\"sst2\": [\" terrible\", \" great\"]}\n",
    "static_reps = {}\n",
    "def get_static_reps(static_reps, model_name, prefix):\n",
    "    model = AutoModelForPreTraining.from_pretrained(model_name)\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "    if \"roberta\" in model_name:\n",
    "        embeddings = model.state_dict()[\"roberta.embeddings.word_embeddings.weight\"]\n",
    "        static_reps_roberta = []\n",
    "        for label_word in label_words[task]:\n",
    "            index = tokenizer.encode(label_word)[1]\n",
    "            static_reps_roberta.append(embeddings[index])\n",
    "        static_reps_roberta = torch.stack(static_reps_roberta)\n",
    "        static_reps[f\"{prefix}-roberta\"] = static_reps_roberta\n",
    "    elif \"bert\" in model_name:\n",
    "        embeddings = model.state_dict()[\"bert.embeddings.word_embeddings.weight\"]\n",
    "        static_reps_bert = []\n",
    "        for label_word in label_words[task]:\n",
    "            index = tokenizer.encode(label_word.strip())[1]\n",
    "            static_reps_bert.append(embeddings[index])\n",
    "        static_reps_bert = torch.stack(static_reps_bert)\n",
    "        static_reps[f\"{prefix}-bert\"] = static_reps_bert\n",
    "    else:\n",
    "        embeddings = model.state_dict()[\"electra.embeddings.word_embeddings.weight\"]\n",
    "        weight_params = model.state_dict()[\"discriminator_predictions.dense_prediction.weight\"]\n",
    "        bias_params = model.state_dict()[\"discriminator_predictions.dense_prediction.bias\"]\n",
    "        static_reps[f\"{prefix}-electra\"] = [weight_params, bias_params]\n",
    "\n",
    "static_reps = {}\n",
    "for model_name in names:\n",
    "    get_static_reps(static_reps, model_name, \"zeroshot\")\n",
    "for model_dir in model_dirs:\n",
    "    get_static_reps(static_reps, model_dir, \"fewshot\")\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "68b9d8bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pca(reps):\n",
    "    pca_model = PCA(n_components=2).fit(reps)\n",
    "    return pca_model\n",
    "\n",
    "def pca_infer(pca_model, reps):\n",
    "    return pca_model.transform(reps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "d109c1d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data_type(task):\n",
    "    if task in [\"sst2\", \"qnli\", \"rte\", \"mnli\"]:\n",
    "        return \"glue\"\n",
    "    elif task in [\"sst5\", \"mr\"]:\n",
    "        return \"file\"\n",
    "    else:\n",
    "        return \"other\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "2ff93218",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dataset(task):\n",
    "    data_type = get_data_type(task)\n",
    "    if data_type == \"glue\":\n",
    "        dataset = load_dataset(\"glue\", task)[\"validation\"]\n",
    "    elif data_type == \"other\":\n",
    "        dataset = load_dataset(task)\n",
    "        if task == \"imdb\":\n",
    "            dataset = dataset[\"test\"]\n",
    "        else:\n",
    "            dataset = dataset[\"validation\"]\n",
    "        if task == \"boolq\":\n",
    "            dataset = dataset.add_column(\"label\", dataset[\"answer\"])\n",
    "    else:\n",
    "        dataset = None\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "id": "cd4bd9dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extended(ax, x, y, **args):\n",
    "\n",
    "    xlim = ax.get_xlim()\n",
    "    ylim = ax.get_ylim()\n",
    "\n",
    "    x_ext = np.linspace(xlim[0], xlim[1], 100)\n",
    "    p = np.polyfit(x, y , deg=1)\n",
    "    y_ext = np.poly1d(p)(x_ext)\n",
    "    ax.plot(x_ext, y_ext, **args)\n",
    "    ax.set_xlim(xlim)\n",
    "    ax.set_ylim(ylim)\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 303,
   "id": "1a9bcdd8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration default\n",
      "Reusing dataset boolq (/n/fs/nlp-mengzhou/.cache/datasets/boolq/default/0.1.0/bf0dd57da941c50de94ae3ce3cef7fea48c08f337a4b7aac484e9dddc5aa24e5)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2365ee7a70aa48e48c6ffde85683d56c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(3270, 768)\n",
      "0.7076452599388379\n",
      "(3270, 768)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/n/fs/nlp-mengzhou/anaconda3/envs/updated/lib/python3.8/site-packages/sklearn/linear_model/_logistic.py:814: ConvergenceWarning: lbfgs failed to converge (status=1):\n",
      "STOP: TOTAL NO. of ITERATIONS REACHED LIMIT.\n",
      "\n",
      "Increase the number of iterations (max_iter) or scale the data as shown in:\n",
      "    https://scikit-learn.org/stable/modules/preprocessing.html\n",
      "Please also refer to the documentation for alternative solver options:\n",
      "    https://scikit-learn.org/stable/modules/linear_model.html#logistic-regression\n",
      "  n_iter_i = _check_optimize_result(\n",
      "/n/fs/nlp-mengzhou/anaconda3/envs/updated/lib/python3.8/site-packages/sklearn/linear_model/_logistic.py:814: ConvergenceWarning: lbfgs failed to converge (status=1):\n",
      "STOP: TOTAL NO. of ITERATIONS REACHED LIMIT.\n",
      "\n",
      "Increase the number of iterations (max_iter) or scale the data as shown in:\n",
      "    https://scikit-learn.org/stable/modules/preprocessing.html\n",
      "Please also refer to the documentation for alternative solver options:\n",
      "    https://scikit-learn.org/stable/modules/linear_model.html#logistic-regression\n",
      "  n_iter_i = _check_optimize_result(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.7006116207951071\n",
      "(6540, 768)\n",
      "0.6957186544342507\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/n/fs/nlp-mengzhou/anaconda3/envs/updated/lib/python3.8/site-packages/sklearn/linear_model/_logistic.py:814: ConvergenceWarning: lbfgs failed to converge (status=1):\n",
      "STOP: TOTAL NO. of ITERATIONS REACHED LIMIT.\n",
      "\n",
      "Increase the number of iterations (max_iter) or scale the data as shown in:\n",
      "    https://scikit-learn.org/stable/modules/preprocessing.html\n",
      "Please also refer to the documentation for alternative solver options:\n",
      "    https://scikit-learn.org/stable/modules/linear_model.html#logistic-regression\n",
      "  n_iter_i = _check_optimize_result(\n"
     ]
    }
   ],
   "source": [
    "from sklearn.linear_model import LogisticRegression\n",
    "dataset = get_dataset(task)\n",
    "def prob(reps, labels):\n",
    "    if len(reps) > len(labels):\n",
    "        reps = reps.reshape(-1, 2, reps.shape[-1])\n",
    "        reps = reps[:, 1] - reps[:, 0]\n",
    "    clf = LogisticRegression(random_state=0).fit(reps, labels)\n",
    "    preds = clf.predict(reps)\n",
    "    return (preds == labels).sum() / len(preds)\n",
    "for model in [\"bert\", \"roberta\", \"electra\"]:\n",
    "    print(reps[f\"zeroshot-{model}\"].shape)\n",
    "    print(prob(reps[f\"zeroshot-{model}\"], dataset[\"label\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 304,
   "id": "a48dff0b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_7726/2483707788.py:36: DeprecationWarning: In future, it will be an error for 'np.bool_' scalars to be interpreted as an index\n",
      "  ax.scatter(pca_reps[index, 0], pca_reps[index, 1], label=f\"{features[label]}\", c=\"plum\" if label == 0 else \"steelblue\")\n",
      "/tmp/ipykernel_7726/2483707788.py:36: DeprecationWarning: In future, it will be an error for 'np.bool_' scalars to be interpreted as an index\n",
      "  ax.scatter(pca_reps[index, 0], pca_reps[index, 1], label=f\"{features[label]}\", c=\"plum\" if label == 0 else \"steelblue\")\n",
      "/tmp/ipykernel_7726/2483707788.py:36: DeprecationWarning: In future, it will be an error for 'np.bool_' scalars to be interpreted as an index\n",
      "  ax.scatter(pca_reps[index, 0], pca_reps[index, 1], label=f\"{features[label]}\", c=\"plum\" if label == 0 else \"steelblue\")\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 51884 characters]\n",
      "text/plain": [
       "<Figure size 864x216 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def get_decision_boundary(reps_model, static_reps):\n",
    "    def get_boundary_point(index):\n",
    "        last_ele = (0 - np.dot(reps_model[index][:-1], static_reps[:-1])) / static_reps[-1]\n",
    "        boundary_point = np.concatenate([reps_model[index][:-1].tolist(), [last_ele]])\n",
    "        return boundary_point\n",
    "    p1 = get_boundary_point(0)\n",
    "    p2 = get_boundary_point(3)\n",
    "    all_points = []\n",
    "    for i in range(100):\n",
    "        all_points.append(get_boundary_point(i))\n",
    "    return p1, p2, all_points\n",
    "    \n",
    "\n",
    "def transform_point(p1, pca_model):\n",
    "    p1 = pca_model.transform(p1)\n",
    "    return p1\n",
    "\n",
    "\n",
    "prefix = \"zeroshot\"\n",
    "num_labels = 2\n",
    "labels = np.array(dataset[\"label\"])\n",
    "features = sst2_dataset.features[\"label\"].names\n",
    "fig, axes = plt.subplots(1, 3, figsize=(12, 3))\n",
    "for i, model_name in enumerate([\"bert\", \"roberta\", \"electra\"]):\n",
    "    ax = axes[i]\n",
    "    model_name = f\"{prefix}-{model_name}\"\n",
    "    reps_model = reps[model_name]\n",
    "    if \"electra\" in model_name:\n",
    "        reps_model = reps_model.reshape(-1, num_labels, reps_model.shape[-1])\n",
    "        reps_model = reps_model[:, 0] - reps_model[:, 1]\n",
    "#     reps_model = normalize(reps_model)\n",
    "    pca_model = pca(reps_model)\n",
    "    pca_reps = pca_infer(pca_model, reps_model)\n",
    "    for label in set(labels):\n",
    "        index = np.where(labels == label)[0]\n",
    "        ax.scatter(pca_reps[index, 0], pca_reps[index, 1], label=f\"{features[label]}\", c=\"plum\" if label == 0 else \"steelblue\")\n",
    "        ax.legend()\n",
    "        ax.set_title(model_name)\n",
    "fig.tight_layout()\n",
    "plt.savefig(f\"images/zeroshot-{task}.pdf\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 312,
   "id": "1eacf418",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 72676 characters]\n",
      "text/plain": [
       "<Figure size 864x216 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "prefix = \"fewshot\"\n",
    "num_labels = 2\n",
    "labels = np.array(dataset[\"label\"])\n",
    "features = [\"Yes\", \"No\"]\n",
    "fig, axes = plt.subplots(1, 3, figsize=(12, 3))\n",
    "for i, model_name in enumerate([\"bert\", \"roberta\", \"electra\"]):\n",
    "    ax = axes[i]\n",
    "    model_name = f\"{prefix}-{model_name}\"\n",
    "    reps_model = reps[model_name]\n",
    "#     reps_model = normalize(reps_model)\n",
    "    if \"electra\" in model_name:\n",
    "        reps_model = reps_model.reshape(-1, num_labels, reps_model.shape[-1])\n",
    "        reps_model = reps_model[:, 0] - reps_model[:, 1]\n",
    "    pca_model = pca(reps_model)\n",
    "    pca_reps = pca_infer(pca_model, reps_model)\n",
    "    for label in [1, 0]:\n",
    "        index = np.where(labels == label)[0]\n",
    "        ax.scatter(pca_reps[index, 0], pca_reps[index, 1], label=f\"{features[label]}\", c=\"plum\" if label == 0 else \"steelblue\")\n",
    "        ax.legend()\n",
    "        ax.set_title(model_name)\n",
    "fig.tight_layout()\n",
    "plt.savefig(f\"images/16shot-{task}.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1743e179",
   "metadata": {},
   "outputs": [],
   "source": [
    "# prediction distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "301727d9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 23848 characters]\n",
      "text/plain": [
       "<Figure size 1080x720 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "def sigmoid(z):\n",
    "    return 1/(1 + np.exp(-z))\n",
    "\n",
    "fig, axes = plt.subplots(2, 2, figsize=(15, 10))\n",
    "for i, tt in enumerate(additional_info):\n",
    "    for j, subset in enumerate([\"train\", \"fewshot_validation\"]):\n",
    "        logits = additional_info[tt][subset][\"logits\"]\n",
    "        axes[i][j].hist(sigmoid(logits))\n",
    "        axes[i][j].set_title(f\"{task} {tt} {subset if subset == 'train' else 'validation'}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d88f390b",
   "metadata": {},
   "source": [
    "## Preprocess function for text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "404e9b74",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "max_seq_length = 512\n",
    "objective = \"mlm\"\n",
    "padding=\"max_length\"\n",
    "num_labels = 2\n",
    "\n",
    "def preprocess_function(examples, template=None, label_name=\"label\"):\n",
    "    result = tokenize_and_mapping(examples, template, tokenizer, 512, \"mlm\", \"prompt\", None, padding, num_labels)\n",
    "    return result\n",
    "\n",
    "processed_dataset = encoded_dataset.map(lambda x: preprocess_function(x), batched=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73079e01",
   "metadata": {},
   "source": [
    "## Plot function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "id": "ba1b721f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(9842,)\n",
      "(9842, 3)\n"
     ]
    }
   ],
   "source": [
    "def plot_one(ax, logits):\n",
    "    sns.histplot(data=logits,  ax=ax, binwidth=0.1, )\n",
    "    lens = len(logits)\n",
    "    ax.yaxis.set_major_formatter(ticker.PercentFormatter(len(logits))) # xmax=len(logits_k_h)))\n",
    "    ax.set_xticks([0, 1])\n",
    "    ax.tick_params(axis='y', labelsize=15)\n",
    "    ax.tick_params(axis='x', labelsize=15)\n",
    "    ax.set_xlabel(\"\")\n",
    "    ax.set_ylabel(\"\")\n",
    "    \n",
    "def plot_four(objective, axes, logits, labels):\n",
    "    for k in range(num_labels):\n",
    "        for h in range(num_labels):\n",
    "            index = np.where(labels == k)[0]\n",
    "            if objective == \"dis\":\n",
    "                logits_k_h = 1 - sigmoid(logits[index, h])\n",
    "            else:\n",
    "                logits_k_h = softmax(logits)[index, h]\n",
    "            plot_one(ax=axes[k][h], logits=logits_k_h)\n",
    "            if h == 0:\n",
    "                axes[k][h].set_ylabel(y_words[k], fontsize=17)\n",
    "            if k == num_labels - 1 :\n",
    "                axes[k][h].set_xlabel(label_words[h], fontsize=17)\n",
    "           \n",
    "print(labels.shape)\n",
    "print(logits.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f97be74",
   "metadata": {},
   "source": [
    "## keyword output distribution analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "id": "7b5167af",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoModelForMaskedLM.from_pretrained(\"google/electra-base-generator\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"google/electra-base-generator\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "618ca29d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_keyword_id(keyword):\n",
    "    return tokenizer(keyword)[\"input_ids\"][1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "c434b016",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100 examples are done.\n",
      "200 examples are done.\n",
      "300 examples are done.\n",
      "400 examples are done.\n",
      "500 examples are done.\n",
      "600 examples are done.\n",
      "700 examples are done.\n",
      "800 examples are done.\n",
      "900 examples are done.\n",
      "1000 examples are done.\n"
     ]
    }
   ],
   "source": [
    "# load_data from great terrible\n",
    "file = \"/n/fs/nlp-mengzhou/space3/meta/data/sample_roberta/terrible.sents\"\n",
    "lines = open(file, \"r\").readlines()\n",
    "inputs = tokenizer(lines, truncation=True)\n",
    "great_id = get_keyword_id(\"great\")\n",
    "terrible_id = get_keyword_id(\"terrible\")\n",
    "m = nn.Softmax(dim=-1)\n",
    "\n",
    "great_logits = []\n",
    "terrible_logits = []\n",
    "great_probs = []\n",
    "terrible_probs = []\n",
    "for i in range(len(inputs[\"input_ids\"])):\n",
    "    input_ids = torch.tensor(inputs[\"input_ids\"][i])\n",
    "    keyword_index = torch.where(input_ids == terrible_id)[0]\n",
    "    if len(keyword_index) == 1:\n",
    "        keyword_index == keyword_index[0]\n",
    "    input_ids[keyword_index] = tokenizer.mask_token_id\n",
    "    input_ids = input_ids.unsqueeze(0)\n",
    "    attention_mask = torch.tensor(inputs[\"attention_mask\"][i]).unsqueeze(0)\n",
    "    logits = model(input_ids, attention_mask).logits\n",
    "    probs = m(logits)\n",
    "    great_logit = logits[0, keyword_index, great_id]\n",
    "    terrible_logit = logits[0, keyword_index, terrible_id]\n",
    "    great_prob = probs[0, keyword_index, great_id]\n",
    "    terrible_prob = probs[0, keyword_index, terrible_id]\n",
    "    \n",
    "    great_logits.append(great_logit)\n",
    "    terrible_logits.append(terrible_logit)\n",
    "    great_probs.append(great_prob)\n",
    "    terrible_probs.append(terrible_prob)\n",
    "    \n",
    "    if (i+1) % 100 == 0:\n",
    "        print(f\"{i+1} examples are done.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "4be9447a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "\n",
    "great_great_logits = torch.cat(great_logits).clone()\n",
    "great_terrible_logits = torch.cat(terrible_logits).clone()\n",
    "great_great_probs = torch.cat(great_probs).clone()\n",
    "great_terrible_probs = torch.cat(terrible_probs).clone()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "id": "51b41938",
   "metadata": {},
   "outputs": [],
   "source": [
    "terrible_great_logits = torch.cat(great_logits).clone()\n",
    "terrible_terrible_logits = torch.cat(terrible_logits).clone()\n",
    "terrible_great_probs = torch.cat(great_probs).clone()\n",
    "terrible_terrible_probs = torch.cat(terrible_probs).clone()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "id": "a00a7ca4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 23396 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "num_labels = 2\n",
    "label_words = [\"terrible\", \"great\"]\n",
    "fig, axes = plt.subplots(2, 2, figsize=(8, 6))\n",
    "great_logits = torch.cat([great_terrible_logits.unsqueeze(-1), great_great_logits.unsqueeze(-1)], dim=1)\n",
    "terrible_logits = torch.cat([terrible_terrible_logits.unsqueeze(-1), terrible_great_logits.unsqueeze(-1)], dim=1)\n",
    "logits = torch.cat([terrible_logits, great_logits, ], dim=0).detach().numpy()\n",
    "plot_four(\"mlm\", axes, logits, torch.tensor([0] * len(great_great_logits) + [1] * len(terrible_great_logits)))\n",
    "fig.tight_layout()\n",
    "plt.savefig(\"images/sample_sentences_dist.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c898fda",
   "metadata": {},
   "source": [
    "## Plot the zero-shot and few-shot predictions "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 185,
   "id": "e7a160e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/n/fs/nlp-mengzhou/space3/out/hf/analysis_v1/bert-base-uncased/snli/additional-info.pt does not exist\n",
      "/n/fs/nlp-mengzhou/space3/out/hf/analysis_v1/google/electra-base-generator/snli/additional-info.pt does not exist\n"
     ]
    }
   ],
   "source": [
    "# load the zero-shot predictions\n",
    "\n",
    "task = \"snli\"\n",
    "num_labels = 2\n",
    "additional_info = {}\n",
    "names = {\"bert-base-uncased\": \"bert\", \n",
    "         \"roberta-base\": \"roberta\", \n",
    "         \"google/electra-base-discriminator\": \"electra\",\n",
    "         \"google/electra-base-generator\": \"electra_generator\"}\n",
    "\n",
    "for model in [\"bert-base-uncased\", \"roberta-base\", \"google/electra-base-discriminator\", \"google/electra-base-generator\"]:\n",
    "    model_dir = f\"/n/fs/nlp-mengzhou/space3/out/hf/analysis_v1/{model}/{task}\"\n",
    "    file = os.path.join(model_dir, \"additional-info.pt\")\n",
    "    if os.path.exists(file):\n",
    "        this_reps = torch.load(file)\n",
    "        additional_info[f\"zeroshot-{names[model]}\"] = this_reps\n",
    "    else:\n",
    "        print(f\"{file} does not exist\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 190,
   "id": "9178b9b1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'fewshot_validation': defaultdict(list,\n",
       "             {'logits': array([17.954926, 12.519672, 15.957389, ..., 13.217457,  9.160052,\n",
       "                     11.806158], dtype=float32),\n",
       "              'labels': array([1, 0, 2, ..., 0, 2, 1])})}"
      ]
     },
     "execution_count": 190,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "additional_info[\"zeroshot-roberta\"][\"\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 187,
   "id": "a0033178",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Reusing dataset snli (/n/fs/nlp-mengzhou/.cache/datasets/snli/plain_text/1.0.0/1f60b67533b65ae0275561ff7828aad5ee4282d0e6f844fd148d05d3c6ea251b)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "adf734a1757742339e8ac1040ed99f04",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset = load_dataset(\"snli\")\n",
    "labels = np.array(dataset[\"validation\"][\"label\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "id": "6d84bd7e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ACC for LABEL 0: 99.8\n",
      "ACC for LABEL 1: 0.1\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 10012 characters]\n",
      "text/plain": [
       "<Figure size 720x216 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "key = \"zeroshot-electra_generator\"\n",
    "ori_logits = additional_info[key][\"fewshot_validation\"][\"logits\"].reshape(-1, num_labels)\n",
    "fig, axes = plt.subplots(1, 2, figsize=(10, 3))\n",
    "\n",
    "num_labels = 2\n",
    "for label in range(num_labels):\n",
    "    class_logits = ori_logits[labels == label]\n",
    "    if key == \"zeroshot-electra\":\n",
    "        pred = class_logits.argmin(-1)\n",
    "    else:\n",
    "        pred = class_logits.argmax(-1)\n",
    "    \n",
    "    plot_one(axes[label], pred.astype(float))\n",
    "    print(f\"ACC for LABEL {label}: {round((pred == label).sum() / len(pred) * 100, 1)}\")\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "id": "e9387e0e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 18092 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axes = plt.subplots(num_labels, num_labels, figsize=(8, 6), sharex=True, sharey=True)\n",
    "plt.subplots_adjust(hspace=0.1, wspace=0.1)\n",
    "for i, tt in enumerate([\"zeroshot-electra_generator\"]):\n",
    "    objective = map_to_objective(tt)\n",
    "    for j, subset in enumerate([\"fewshot_validation\"]):\n",
    "        logits = additional_info[tt][subset][\"logits\"].reshape(-1, num_labels)\n",
    "#         if objective == \"dis\":\n",
    "#             labels_labels = labels.reshape(-1, num_labels)\n",
    "#             labels = labels.argmin(-1)\n",
    "        plot_four(objective, axes[:, i*num_labels: (i+1)*num_labels], logits, labels)\n",
    "fig.suptitle('Zero-shot ELECTRA Generator', fontsize=16)\n",
    "plt.savefig(f\"images/{task}-zeroshot-electra_generator.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "id": "ae0d236e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['zeroshot-roberta', 'zeroshot-electra', 'zeroshot-electra_generator'])"
      ]
     },
     "execution_count": 141,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def sigmoid(z):\n",
    "    return 1/(1 + np.exp(-z))\n",
    "def softmax(z, temp=1):\n",
    "    return ((np.exp(z)**temp) / (np.exp(z).sum(-1).reshape(-1, 1)**temp))\n",
    "def map_to_objective(key):\n",
    "    model = key.split(\"-\")[1]\n",
    "    if model == \"electra\":\n",
    "        return \"dis\"\n",
    "    return \"mlm\"\n",
    "z = additional_info[\"zeroshot-roberta\"][\"fewshot_validation\"][\"logits\"].reshape(-1, num_labels)\n",
    "additional_info.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 164,
   "id": "1a776405",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(9842, 3)"
      ]
     },
     "execution_count": 164,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "additional_info[\"zeroshot-electra\"][\"fewshot_validation\"][\"logits\"].reshape(-1, 3).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 203,
   "id": "b1c1d2a8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0.98, 'Zero-Shot RoBERTa                               Zero-Shot ELECTRA')"
      ]
     },
     "execution_count": 203,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 30008 characters]\n",
      "text/plain": [
       "<Figure size 864x360 with 18 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "colors = sns.color_palette(as_cmap=True)\n",
    "num_labels = 3\n",
    "y_words = [\"Entailment\", \"Neutral\", \"Contradiction\"]\n",
    "label_words = [\"Yes\", \"Maybe\", \"No\"]\n",
    "fig, axes = plt.subplots(num_labels, num_labels*2, figsize=(12, 5), sharex=True, sharey=True)\n",
    "plt.subplots_adjust(hspace=0.1, wspace=0.1)\n",
    "labels = additional_info[\"zeroshot-roberta\"][subset][\"labels\"]\n",
    "for i, tt in enumerate([\"zeroshot-roberta\", \"zeroshot-electra\"]):\n",
    "    objective = map_to_objective(tt)\n",
    "    for j, subset in enumerate([\"fewshot_validation\"]):\n",
    "        logits = additional_info[tt][subset][\"logits\"].reshape(-1, num_labels)\n",
    "        logits = logits.reshape(-1, num_labels)\n",
    "        plot_four(objective, axes[:, i*num_labels: (i+1)*num_labels], logits, labels)\n",
    "        \n",
    "\n",
    "if tt == \"fewshot-electra\":\n",
    "    fig.suptitle(\"Few-Shot ELECTRA\", fontsize=17)\n",
    "elif tt == \"fewshot-roberta\":\n",
    "    fig.suptitle(\"Few-Shot RoBERTa\", fontsize=17)\n",
    "elif tt == \"zeroshot-electra\":\n",
    "    fig.suptitle(\"Zero-Shot ELECTRA\", fontsize=17)\n",
    "elif tt == \"zeroshot-roberta\":\n",
    "    fig.suptitle(\"Zero-Shot RoBERTa\", fontsize=17)\n",
    "fig.suptitle(\"Zero-Shot RoBERTa                               Zero-Shot ELECTRA\", fontsize=17)\n",
    "# plt.savefig(f\"images/{task}-zero.pdf\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 182,
   "id": "452e2b44",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'fewshot_validation': defaultdict(list,\n",
       "             {'logits': array([17.954926, 12.519672, 15.957389, ..., 13.217457,  9.160052,\n",
       "                     11.806158], dtype=float32),\n",
       "              'labels': []})}"
      ]
     },
     "execution_count": 182,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "additional_info[\"zeroshot-roberta\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "fc7962d7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['zeroshot-electra'])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load few_shot results\n",
    "\n",
    "sst2_model_dirs = [\"/n/fs/nlp-mengzhou/space3/out/hf/few_shot/sst2/bert-base-uncased/shot16/prompt_FT/1/bs4_lr2e-5\",\n",
    "    \"/n/fs/nlp-mengzhou/space3/out/hf/few_shot/sst2/roberta-base/shot16/prompt_FT/1/bs4_lr1e-5\",\n",
    "    \"/n/fs/nlp-mengzhou/space3/out/hf/few_shot/sst2/google/electra-base-discriminator/shot16/prompt_FT/1/bs8_lr1e-5\",]\n",
    "boolq_model_dirs = [\"/n/fs/nlp-mengzhou/space3/out/hf/few_shot/boolq/bert-base-uncased/shot16/prompt_FT/3/bs2_lr2e-5\",\n",
    "                   \"/n/fs/nlp-mengzhou/space3/out/hf/few_shot/boolq/roberta-base/shot16/prompt_FT/3/bs2_lr2e-5\",\n",
    "                   \"/n/fs/nlp-mengzhou/space3/out/hf/few_shot/boolq/google/electra-base-discriminator/shot16/prompt_FT/3/bs2_lr2e-5\"]\n",
    "sst5_model_dirs = [\"\", \"\", \"/n/fs/nlp-mengzhou/space3/out/hf/few_shot/sst5/google/electra-base-discriminator/shot16/prompt_FT/1/bs8_lr2e-5\"]\n",
    "agnews_model_dirs = [\"\", \"\", \"/n/fs/nlp-mengzhou/space3/out/hf/few_shot/ag_news/google/electra-base-discriminator/shot16/prompt_FT/9/bs8_lr1e-5\"]\n",
    "model_dirs = {\"sst2\": sst2_model_dirs, \"boolq\": boolq_model_dirs, \"sst5\": sst5_model_dirs, \"ag_news\": agnews_model_dirs}\n",
    "for model_dir in model_dirs[task]:\n",
    "    file = os.path.join(model_dir, \"additional-info.pt\")\n",
    "    if os.path.exists(file):\n",
    "        this_reps = torch.load(file)\n",
    "        if \"roberta\" in model_dir:\n",
    "            model = \"roberta\"\n",
    "        elif \"bert\" in model_dir:\n",
    "            model = \"bert\"\n",
    "        elif \"electra\" in model_dir:\n",
    "            model = \"electra\"\n",
    "        additional_info[f\"fewshot-{model}\"] = this_reps\n",
    "    else:\n",
    "        print(f\"{file} does not exist\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06b06ab6",
   "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
}
