{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f79f0116",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'sns' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Input \u001b[0;32mIn [1]\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtransformers\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m AutoTokenizer\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mos\u001b[39;00m\n\u001b[0;32m----> 4\u001b[0m \u001b[43msns\u001b[49m\u001b[38;5;241m.\u001b[39mset_style(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mwhitegrid\u001b[39m\u001b[38;5;124m\"\u001b[39m, {\u001b[38;5;124m'\u001b[39m\u001b[38;5;124maxes.grid\u001b[39m\u001b[38;5;124m'\u001b[39m : \u001b[38;5;28;01mFalse\u001b[39;00m})\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtransformers\u001b[39;00m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'sns' is not defined"
     ]
    }
   ],
   "source": [
    "from datasets import load_from_disk\n",
    "from transformers import AutoTokenizer\n",
    "import os\n",
    "sns.set_style(\"whitegrid\", {'axes.grid' : False})\n",
    "import transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "bd870615",
   "metadata": {},
   "outputs": [],
   "source": [
    "# look into /n/fs/nlp-mengzhou/space3/meta/data/copa.bin/threshold0.15.list and turn it into the current format\n",
    "def transform(data):\n",
    "    tokenizer = AutoTokenizer.from_pretrained(\"roberta-base\")\n",
    "    \n",
    "    def decode(examples):\n",
    "        lens = len(examples[\"src_tokens\"])\n",
    "        prompts = []\n",
    "        char_starts = []\n",
    "        char_ends = []\n",
    "        labels = []\n",
    "        for i in range(lens):\n",
    "            src_tokens = examples[\"src_tokens\"][i]\n",
    "            prompt = tokenizer.decode(src_tokens[1:-1])\n",
    "\n",
    "            span_start_ends = examples[\"span_start_ends\"][i]\n",
    "            span_tokens = src_tokens[span_start_ends[0][0]: span_start_ends[1][0]]\n",
    "            span = tokenizer.decode(span_tokens)\n",
    "            \n",
    "            label = examples[\"targets\"][i][0]\n",
    "            \n",
    "            if span in prompt:\n",
    "                char_start = prompt.index(span)\n",
    "\n",
    "                char_end = char_start + len(span)\n",
    "                char_starts.append([char_start])\n",
    "                char_ends.append([char_end])\n",
    "                prompts.append(prompt)\n",
    "                labels.append(label)\n",
    "                \n",
    "            \n",
    "        examples[\"prompt\"] = prompts\n",
    "        examples[\"char_start\"] = char_starts\n",
    "        examples[\"char_end\"] = char_ends\n",
    "        examples[\"label\"] = labels\n",
    "        return examples\n",
    "            \n",
    "    data = data.map(decode, batched=True, remove_columns=list(data.features.keys()), num_proc=20)\n",
    "    return data\n",
    "\n",
    "def run():\n",
    "    valid_file = \"/n/fs/nlp-mengzhou/space3/meta/data/copa.bin/threshold0.15.list/valid.pt\"\n",
    "    valid_data = load_from_disk(valid_file)\n",
    "    train_file = \"/n/fs/nlp-mengzhou/space3/meta/data/copa.bin/threshold0.15.list/train.pt\"\n",
    "    train_data = load_from_disk(train_file)\n",
    "    \n",
    "    transformed_valid_data = transform(valid_data)\n",
    "    transformed_train_data = transform(train_data)\n",
    "    \n",
    "    save_dir = \"/n/fs/nlp-mengzhou/space3/data/created/copa-retrieval0.15\"\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    save_train_file = os.path.join(save_dir, \"train.pt\")\n",
    "    save_valid_file = os.path.join(save_dir, \"valid.pt\")\n",
    "    transformed_train_data.save_to_disk(save_train_file)\n",
    "    transformed_valid_data.save_to_disk(save_valid_file)\n",
    "    return transformed_train_data, transformed_valid_data\n",
    "\n",
    "train_data, valid_data = run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0e586462",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Reusing dataset glue (/home/ubuntu/.cache/huggingface/datasets/glue/sst2/1.0.0/dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2e2504d716e34297af4c0369cd2510f6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from transformers import default_data_collator\n",
    "from transformers import AutoTokenizer \n",
    "dataset_name = \"sst2\"\n",
    "dataset = load_dataset(\"glue\", dataset_name)[\"validation\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "id": "9cb376a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "main_dir=\"/home/ubuntu/space3/out/hf/analysis_v1\"\n",
    "electra_reps = torch.load(f\"{main_dir}/google/electra-base-discriminator/{dataset_name}/reps.pt\")\n",
    "bert_mask_reps = torch.load(f\"{main_dir}/bert-base-uncased/{dataset_name}/reps.pt\")\n",
    "\n",
    "# legacy_dir=\"/home/ubuntu/space3/analysis/analysis-v1/legacy\"\n",
    "# electra_legacy_reps = torch.load(f\"{legacy_dir}/sst2-electra.reps\")\n",
    "# bert_mask_legacy_reps = torch.load(f\"{legacy_dir}/sst2-bert-mask.reps\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "id": "c49cf83b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.preprocessing import normalize\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.manifold import TSNE\n",
    "from matplotlib import pyplot as plt\n",
    "def get_x(input=1):\n",
    "    return [input] * 767 + [((- electra_bias - input * sum(electra_weight[0, :-1])) / electra_weight[0, -1]).item()]\n",
    "def get_x2(input=2):\n",
    "    return [input] * 300 + [((- electra_bias - input * sum(electra_weight[0, :-468])) / electra_weight[0, 300:]).item()] * 468\n",
    "electra_reps = normalize(electra_reps)\n",
    "bert_mask_reps = normalize(bert_mask_reps)\n",
    "p1 = normalize([get_x(1)])\n",
    "p2 = normalize([get_x(4)])\n",
    "ps = np.array([p1[0], p2[0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "id": "e90c81b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "pca = PCA(n_components=2)\n",
    "pca_bert = pca.fit_transform(bert_mask_reps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "id": "ce5f70c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "pca = PCA(n_components=2).fit(electra_reps)\n",
    "pca_electra = pca.transform(electra_reps)\n",
    "ps = pca.transform(ps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "id": "32646002",
   "metadata": {},
   "outputs": [],
   "source": [
    "tse = TSNE(n_components=2, learning_rate='auto', init='random')\n",
    "tse_bert = tse.fit_transform(bert_mask_reps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "id": "9346375b",
   "metadata": {},
   "outputs": [],
   "source": [
    "tse = TSNE(n_components=2, learning_rate='auto', init='random')\n",
    "tse_electra = tse.fit_transform(electra_reps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "id": "4f1c7498",
   "metadata": {},
   "outputs": [],
   "source": [
    "electra_labels = []\n",
    "labels = dataset.features[\"label\"].names\n",
    "for i in dataset[\"label\"]:\n",
    "    if i == 0:\n",
    "        electra_labels.append(\"electra real positive\")\n",
    "        electra_labels.append(\"electra fake negative\")\n",
    "    else:\n",
    "        electra_labels.append(\"electra fake positive\")\n",
    "        electra_labels.append(\"electra real negative\")\n",
    "color_map = {\"electra real positive\": \"y\", \"electra real negative\": \"plum\", \"electra fake positive\": \"g\", \"electra fake negative\": \"b\" }\n",
    "electra_colors = [color_map[label] for label in electra_labels]\n",
    "\n",
    "bert_labels = [\"bert real positive\" if i == 0 else \"bert real negative\" for i in dataset[\"label\"]]\n",
    "bert_color_map = {\"bert real positive\": \"steelblue\", \"bert real negative\": \"orange\"}\n",
    "bert_colors = [bert_color_map[label] for label in bert_labels]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "id": "e91edd61",
   "metadata": {},
   "outputs": [],
   "source": [
    "electra = transformers.ElectraForPreTraining.from_pretrained(\"google/electra-base-discriminator\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "id": "f5417ed2",
   "metadata": {},
   "outputs": [],
   "source": [
    "electra_weight = electra.discriminator_predictions.dense_prediction.weight\n",
    "electra_bias = electra.discriminator_predictions.dense_prediction.bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "c03585ad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.19760444, 0.57395462],\n",
       "       [0.19783528, 0.57326851]])"
      ]
     },
     "execution_count": 136,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "id": "1148915d",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "ps[1, 0] = 1\n",
    "ps[1, 1] = 1.16128-2.97223"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "id": "fc455ce1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'sst2')"
      ]
     },
     "execution_count": 140,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 199124 characters]\n",
      "text/plain": [
       "<Figure size 640x640 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(8, 8), dpi=80)\n",
    "plt.scatter(pca_bert[:, 0], pca_bert[:, 1], c=bert_colors, label=bert_labels)\n",
    "plt.scatter(pca_electra[:, 0], pca_electra[:, 1], label=electra_labels, c=electra_colors)\n",
    "plt.plot(ps[:, 0], ps[:, 1])\n",
    "\n",
    "import matplotlib.lines as mlines\n",
    "patches = []\n",
    "for label in color_map:\n",
    "    patch = mlines.Line2D([], [], color=color_map[label], marker='o', linestyle='None', markersize=6, label=label)\n",
    "    patches.append(patch)\n",
    "for label in bert_color_map:\n",
    "    patch = mlines.Line2D([], [], color=bert_color_map[label], marker='o', linestyle='None', markersize=6, label=label)\n",
    "    patches.append(patch)\n",
    "plt.ylim(-0.4, 0.5)\n",
    "plt.legend(handles=patches, loc=0, ncol=2)\n",
    "plt.title(\"sst2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "435b7f95",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.1586622, 0.6805651]])"
      ]
     },
     "execution_count": 77,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "b2d1e53c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "db4843e5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "67d99655",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(sst2_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "d6025cf3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['entailment', 'neutral', 'contradiction']"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sst2_dataset[\"validation\"].features[\"label\"].names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5add583f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from datasets import load_from_disk\n",
    "\n",
    "for ob in [\"mlm\", \"dis\"]:\n",
    "    for task in [\"sst2_aug\", \"sst2_v2_aug\", \"nli_aug\", \"nli_v2_aug\", \"boolq_aug\"]:\n",
    "        file = f\"/n/fs/nlp-mengzhou/space3/meta/data/created/heuristic/keywords-v2/prompt_dataset/{ob}/{task}/train.pt\"\n",
    "        dataset = load_from_disk(file)\n",
    "        output_file = os.path.join(os.path.dirname(file), \"train.txt\")\n",
    "        with open(output_file, \"w\") as f:\n",
    "            for ex in dataset:\n",
    "                f.write(f\"{ex['prompt']}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff357bf9",
   "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
}
