generate_adversarials.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# coding=utf-8\n",
    "# Copyright 2020 FireBERT authors. All rights reserved.\n",
    "#\n",
    "# Licensed under the MIT license\n",
    "# See https://github.com/FireBERT-author/FireBERT/blob/master/LICENSE for details\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "#"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adversarial sample generation with code from the TextFooler paper\n",
    "\n",
    "#### Requirements\n",
    "\n",
    "1. TextFooler required resources: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "counter_fitting_embeddings_path = 'resources/embeddings/counter-fitted-vectors.txt'\n",
    "counter_fitting_cos_sim_path = 'resources/cos_sim_counter_fitting.npy'\n",
    "USE_cache_path = 'scratch/tf_cache'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. Path of the tuned model to fool and it's training task:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#task = 'imdb' # can be imdb or mnli\n",
    "#model_path = 'resources/models/IMDB/pytorch_model.bin'\n",
    "#model_path = 'resources/models/IMDB_on_lightning/pytorch_model.bin'\n",
    "#model_path = 'resources/models/co-tuned_IMDB_on_lightning_final_filter/pytorch_model.bin'\n",
    "\n",
    "task = 'mnli' # can be imdb or mnli\n",
    "#model_path = 'resources/models/MNLI/pytorch_model.bin'\n",
    "#model_path = 'resources/models/MNLI_on_lightning/pytorch_model.bin'\n",
    "model_path = 'resources/models/co-tuned_MNLI_on_lightning_final_filter/pytorch_model.bin'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. Path of the dataset to generate samples from and the name of the output file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# IMDB:\n",
    "\n",
    "#dataset_path = 'data/IMDB/imdb_train.txt'\n",
    "#output_path = 'data/IMDB/generated/imdb_adversarial_samples_for_train'\n",
    "#dataset_path = 'data/IMDB/imdb_dev.txt'\n",
    "#output_path = 'data/IMDB/generated/imdb_adversarial_samples_for_dev'\n",
    "#dataset_path = 'data/IMDB/imdb_test.txt'\n",
    "#output_path = 'data/IMDB/generated/imdb_adversarial_samples_for_test'\n",
    "\n",
    "# MNLI:\n",
    "\n",
    "#dataset_path = 'data/MNLI/original/multinli_1.0_train.txt'\n",
    "#output_path = 'data/MNLI/generated/mnli_adversarial_samples_for_train'\n",
    "#dataset_path = 'data/MNLI/original/multinli_1.0_dev_matched.txt'\n",
    "#output_path = 'data/MNLI/generated/mnli_adversarial_samples_for_dev'\n",
    "#dataset_path = 'data/MNLI/original/multinli_1.0_dev_mismatched.txt'\n",
    "#output_path = 'data/MNLI/generated/mnli_adversarial_samples_for_test'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. The number of samples to process from the dataset and batch size:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataset size is 40000 for imdb train and 390702 for MNLI train\n",
    "# dev and test dataset sizes are 5K for IMDB and 10K for MNLI\n",
    "#data_size = 390703 #add one for the header row that is skipped in the logic\n",
    "#batch_size = 32"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we just want to run against the TextFooler sample data (for example, to compare baselines), use these dataset paths instead"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run with just the textfooler sample data\n",
    "#dataset_path = 'data/TextFooler/imdb'\n",
    "dataset_path = 'data/TextFooler/mnli_matched'\n",
    "data_size = 1000\n",
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also run TextFooler against one of our evaluation purtubation models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import TensorDataset, Dataset, SequentialSampler, DataLoader\n",
    "import pytorch_lightning as pl\n",
    "\n",
    "from processors import MnliProcessor, ImdbProcessor\n",
    "from bert_base_model import LightningBertForSequenceClassification\n",
    "from firebert_fse import FireBERT_FSE\n",
    "from firebert_fve import FireBERT_FVE\n",
    "\n",
    "import string\n",
    "import switch\n",
    "import numpy as np\n",
    "import os\n",
    "import re\n",
    "\n",
    "import tensorflow.compat.v1 as tf\n",
    "tf.disable_v2_behavior()\n",
    "\n",
    "import tensorflow_hub as hub\n",
    "\n",
    "from nltk.tokenize.treebank import TreebankWordDetokenizer\n",
    "\n",
    "from transformers.data.processors.utils import DataProcessor, InputExample, InputFeatures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The text fooler logic that we generalized and encapsulated in a class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "The TextFooler algorithm in a class for generating adversarial texts. Adapted to work with a BERT\n",
    "classifier tuned as a PyTorch Lightning model.\n",
    "\n",
    "This class was adapted from code by TextFooler at https://github.com/jind11/TextFooler,\n",
    "a code repository in support of the paper:\n",
    "\n",
    "Jin, Di, et al. \"Is BERT Really Robust? Natural Language Attack on Text Classification and Entailment.\"\n",
    "arXiv preprint arXiv:1907.11932 (2019).\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "class PaperFooler(object):\n",
    "    def __init__(self,\n",
    "                 tokenizer,\n",
    "                 lightning_model,\n",
    "                 USE_cache_path,\n",
    "                 counter_fitting_embeddings_path,\n",
    "                 counter_fitting_cos_sim_path=None,\n",
    "                 max_seq_length=128):\n",
    "        super(PaperFooler, self).__init__()\n",
    "        \n",
    "        self.tokenizer = tokenizer\n",
    "        self.model = lightning_model.cuda()\n",
    "        self.max_seq_length = max_seq_length\n",
    "        \n",
    "        # prepare synonym extractor\n",
    "        # build dictionary via the embedding file\n",
    "        print(\"Building vocab...\")\n",
    "        self.idx2word = {}\n",
    "        self.word2idx = {}\n",
    "\n",
    "        with open(counter_fitting_embeddings_path, 'r', encoding=\"utf-8\") as ifile:\n",
    "            for line in ifile:\n",
    "                word = line.split()[0]\n",
    "                if word not in self.idx2word:\n",
    "                    self.idx2word[len(self.idx2word)] = word\n",
    "                    self.word2idx[word] = len(self.idx2word) - 1\n",
    "\n",
    "        # for cosine similarity matrix\n",
    "        print(\"Building cos sim matrix...\")\n",
    "        if counter_fitting_cos_sim_path:\n",
    "            # load pre-computed cosine similarity matrix if provided\n",
    "            print('Load pre-computed cosine similarity matrix from {}'.format(counter_fitting_cos_sim_path))\n",
    "            self.cos_sim = np.load(counter_fitting_cos_sim_path)\n",
    "        else:\n",
    "            # calculate the cosine similarity matrix\n",
    "            print('Start computing the cosine similarity matrix!')\n",
    "            embeddings = []\n",
    "            with open(counter_fitting_embeddings_path, 'r') as ifile:\n",
    "                for line in ifile:\n",
    "                    embedding = [float(num) for num in line.strip().split()[1:]]\n",
    "                    embeddings.append(embedding)\n",
    "            embeddings = np.array(embeddings)\n",
    "            product = np.dot(embeddings, embeddings.T)\n",
    "            norm = np.linalg.norm(embeddings, axis=1, keepdims=True)\n",
    "            self.cos_sim = product / np.dot(norm, norm.T)\n",
    "        print(\"Cos sim import finished!\")\n",
    "\n",
    "        # build the semantic similarity module\n",
    "        self.sim_predictor = USE(USE_cache_path)\n",
    "\n",
    "        self.stop_words_set = switch.get_stopwords()\n",
    "  \n",
    "\n",
    "    def text_pred(self, text_data, batch_size):\n",
    "        # Switch the model to eval mode.\n",
    "        self.model.eval()\n",
    "\n",
    "        # transform text data into a batch of indices\n",
    "        batch = self.transform_text(text_data, batch_size)\n",
    "\n",
    "        probs_all = []\n",
    "        for input_ids, attention_mask, token_type_ids, ex_idx in batch:\n",
    "            input_ids = input_ids.cuda()\n",
    "            attention_mask = attention_mask.cuda()\n",
    "            token_type_ids = token_type_ids.cuda()\n",
    "            ex_idx = ex_idx.cuda()\n",
    "\n",
    "            with torch.no_grad():\n",
    "                logits = self.model(input_ids=input_ids, attention_mask=attention_mask,\n",
    "                                    token_type_ids=token_type_ids, example_idx=ex_idx)\n",
    "                probs = nn.functional.softmax(logits, dim=-1)\n",
    "                probs_all.append(probs)\n",
    "\n",
    "        return torch.cat(probs_all, dim=0)\n",
    "\n",
    "\n",
    "    def _truncate_seq_pair(self, tokens_a, tokens_b, max_length):\n",
    "        \"\"\"Truncates a sequence pair in place to the maximum length.\"\"\"\n",
    "\n",
    "        # This is a simple heuristic which will always truncate the longer sequence\n",
    "        # one token at a time. This makes more sense than truncating an equal percent\n",
    "        # of tokens from each, since if one sequence is very short then each token\n",
    "        # that's truncated likely contains more information than a longer sequence.\n",
    "        while True:\n",
    "            total_length = len(tokens_a) + len(tokens_b)\n",
    "            if total_length <= max_length:\n",
    "                break\n",
    "            if len(tokens_a) > len(tokens_b):\n",
    "                tokens_a.pop()\n",
    "            else:\n",
    "                tokens_b.pop()\n",
    "\n",
    "    def convert_examples_to_features(self, examples, max_seq_length, tokenizer):\n",
    "        \"\"\"Loads a data file into a list of `InputFeature`s.\"\"\"\n",
    "\n",
    "        features = []\n",
    "        for (ex_index, (text_a, text_b)) in enumerate(examples):\n",
    "            tokens_a = tokenizer.tokenize(' '.join(text_a))\n",
    "\n",
    "            tokens_b = None\n",
    "            if text_b:\n",
    "                tokens_b = tokenizer.tokenize(' '.join(text_b))\n",
    "                # Modifies `tokens_a` and `tokens_b` in place so that the total\n",
    "                # length is less than the specified length.\n",
    "                # Account for [CLS], [SEP], [SEP] with \"- 3\"\n",
    "                self._truncate_seq_pair(tokens_a, tokens_b, max_seq_length - 3)\n",
    "            else:\n",
    "                # Account for [CLS] and [SEP] with \"- 2\"\n",
    "                if len(tokens_a) > max_seq_length - 2:\n",
    "                    tokens_a = tokens_a[:(max_seq_length - 2)]\n",
    "\n",
    "            tokens = [\"[CLS]\"] + tokens_a + [\"[SEP]\"]\n",
    "            token_type_ids = [0] * len(tokens)\n",
    "\n",
    "            if tokens_b:\n",
    "                tokens += tokens_b + [\"[SEP]\"]\n",
    "                token_type_ids += [1] * (len(tokens_b) + 1)\n",
    "\n",
    "            input_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
    "\n",
    "            # The mask has 1 for real tokens and 0 for padding tokens. Only real\n",
    "            # tokens are attended to.\n",
    "            attention_mask = [1] * len(input_ids)\n",
    "\n",
    "            # Zero-pad up to the sequence length.\n",
    "            padding = [0] * (max_seq_length - len(input_ids))\n",
    "            input_ids += padding\n",
    "            attention_mask += padding\n",
    "            token_type_ids += padding\n",
    "\n",
    "            assert len(input_ids) == max_seq_length\n",
    "            assert len(attention_mask) == max_seq_length\n",
    "            assert len(token_type_ids) == max_seq_length\n",
    "\n",
    "            features.append(\n",
    "                InputFeatures(input_ids=input_ids,\n",
    "                              attention_mask=attention_mask,\n",
    "                              token_type_ids=token_type_ids))\n",
    "        return features\n",
    "\n",
    "    def transform_text(self, data, batch_size):\n",
    "        # transform data into seq of embeddings\n",
    "        eval_features = self.convert_examples_to_features(list(zip(data['text_a'], data['text_b'])),\n",
    "                                                          self.max_seq_length, self.tokenizer)\n",
    "\n",
    "        all_input_ids = torch.tensor([f.input_ids for f in eval_features], dtype=torch.long)\n",
    "        all_attention_mask = torch.tensor([f.attention_mask for f in eval_features], dtype=torch.long)\n",
    "        all_token_type_ids = torch.tensor([f.token_type_ids for f in eval_features], dtype=torch.long)\n",
    "        all_idxs = torch.tensor([i for i in range(len(all_input_ids))], dtype=torch.long)\n",
    "        \n",
    "        eval_data = TensorDataset(all_input_ids, all_attention_mask, all_token_type_ids, all_idxs)\n",
    "\n",
    "        # Run prediction for data sequentially\n",
    "        eval_sampler = SequentialSampler(eval_data)\n",
    "        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=batch_size)\n",
    "\n",
    "        return eval_dataloader\n",
    "    \n",
    "    def pick_most_similar_words_batch(self, src_words, sim_mat, idx2word, ret_count=10, threshold=0.):\n",
    "        \"\"\"\n",
    "        embeddings is a matrix with (d, vocab_size)\n",
    "        \"\"\"\n",
    "        sim_order = np.argsort(-sim_mat[src_words, :])[:, 1:1 + ret_count]\n",
    "        sim_words, sim_values = [], []\n",
    "        for idx, src_word in enumerate(src_words):\n",
    "            sim_value = sim_mat[src_word][sim_order[idx]]\n",
    "            mask = sim_value >= threshold\n",
    "            sim_word, sim_value = sim_order[idx][mask], sim_value[mask]\n",
    "            sim_word = [idx2word[id] for id in sim_word]\n",
    "            sim_words.append(sim_word)\n",
    "            sim_values.append(sim_value)\n",
    "        return sim_words, sim_values\n",
    "\n",
    "    def pos_filter(self, ori_pos, new_pos_list):\n",
    "        same = [True if ori_pos == new_pos or (set([ori_pos, new_pos]) <= set(['NOUN', 'VERB']))\n",
    "                else False\n",
    "                for new_pos in new_pos_list]\n",
    "        return same    \n",
    "         \n",
    "    def generate_adversarial(self, task, text_a, text_b, true_label, batch_size,\n",
    "           import_score_threshold=-1., sim_score_threshold=0.7, sim_score_window=15, synonym_num=50):\n",
    "        # first check the prediction of the original text\n",
    "        orig_probs = self.text_pred({'text_a': [text_a], 'text_b': [text_b]}, batch_size).squeeze()\n",
    "        orig_label = torch.argmax(orig_probs)\n",
    "        orig_prob = orig_probs.max()\n",
    "        text_ls = text_b if task == 'mnli' else text_a\n",
    "        if true_label != orig_label:\n",
    "            return '', 0, orig_label, orig_label, 0\n",
    "        else:\n",
    "            len_text = len(text_ls)\n",
    "            if len_text < sim_score_window:\n",
    "                sim_score_threshold = 0.1  # shut down the similarity thresholding function\n",
    "            half_sim_score_window = (sim_score_window - 1) // 2\n",
    "            num_queries = 1\n",
    "\n",
    "            # get the pos and verb tense info\n",
    "            pos_ls = switch.get_pos(text_ls)\n",
    "\n",
    "            # get importance score\n",
    "            leave_1_texts = [text_ls[:ii]+['<oov>']+text_ls[min(ii+1, len_text):] for ii in range(len_text)]\n",
    "            if task == 'mnli':\n",
    "                leave_1_probs = self.text_pred({'text_a':[text_a]*len_text, 'text_b': leave_1_texts}, batch_size)\n",
    "            else:\n",
    "                leave_1_probs = self.text_pred({'text_a':leave_1_texts, 'text_b': [text_b]*len_text}, batch_size)                      \n",
    "            num_queries += len(leave_1_texts)\n",
    "            leave_1_probs_argmax = torch.argmax(leave_1_probs, dim=-1)\n",
    "            import_scores = (orig_prob - leave_1_probs[:, orig_label] + (leave_1_probs_argmax != orig_label).float() * (\n",
    "                        leave_1_probs.max(dim=-1)[0] - torch.index_select(orig_probs, 0,\n",
    "                                                                          leave_1_probs_argmax))).data.cpu().numpy()\n",
    "\n",
    "            # get words to perturb ranked by importance score for word in words_perturb\n",
    "            words_perturb = []\n",
    "            for idx, score in sorted(enumerate(import_scores), key=lambda x: x[1], reverse=True):\n",
    "                if score > import_score_threshold and text_ls[idx] not in self.stop_words_set:\n",
    "                    words_perturb.append((idx, text_ls[idx]))\n",
    "\n",
    "            # find synonyms\n",
    "            words_perturb_idx = [self.word2idx[word] for idx, word in words_perturb if word in self.word2idx]\n",
    "            #src_words, sim_mat, idx2word, ret_count=10, threshold=0.\n",
    "            synonym_words, _ = self.pick_most_similar_words_batch(words_perturb_idx, self.cos_sim, \n",
    "                                                                  self.idx2word, synonym_num, 0.5)\n",
    "            synonyms_all = []\n",
    "            for idx, word in words_perturb:\n",
    "                if word in self.word2idx:\n",
    "                    synonyms = synonym_words.pop(0)\n",
    "                    if synonyms:\n",
    "                        synonyms_all.append((idx, synonyms))\n",
    "\n",
    "            # start replacing and attacking\n",
    "            text_prime = text_ls[:]\n",
    "            text_cache = text_prime[:]\n",
    "            num_changed = 0\n",
    "            for idx, synonyms in synonyms_all:\n",
    "                new_texts = [text_prime[:idx] + [synonym] + text_prime[min(idx + 1, len_text):] for synonym in synonyms]\n",
    "                if task == 'mnli':\n",
    "                    new_probs = self.text_pred({'text_a': [text_a] * len(synonyms), 'text_b': new_texts}, batch_size)\n",
    "                else:\n",
    "                    new_probs = self.text_pred({'text_a': new_texts, 'text_b': [text_b] * len(synonyms)}, batch_size)\n",
    "                \n",
    "                # compute semantic similarity\n",
    "                if idx >= half_sim_score_window and len_text - idx - 1 >= half_sim_score_window:\n",
    "                    text_range_min = idx - half_sim_score_window\n",
    "                    text_range_max = idx + half_sim_score_window + 1\n",
    "                elif idx < half_sim_score_window and len_text - idx - 1 >= half_sim_score_window:\n",
    "                    text_range_min = 0\n",
    "                    text_range_max = sim_score_window\n",
    "                elif idx >= half_sim_score_window and len_text - idx - 1 < half_sim_score_window:\n",
    "                    text_range_min = len_text - sim_score_window\n",
    "                    text_range_max = len_text\n",
    "                else:\n",
    "                    text_range_min = 0\n",
    "                    text_range_max = len_text\n",
    "                semantic_sims = \\\n",
    "                    self.sim_predictor.semantic_sim([' '.join(text_cache[text_range_min:text_range_max])] * len(new_texts),\n",
    "                                           list(map(lambda x: ' '.join(x[text_range_min:text_range_max]), new_texts)))[0]\n",
    "                \n",
    "                num_queries += len(new_texts)\n",
    "                if len(new_probs.shape) < 2:\n",
    "                    new_probs = new_probs.unsqueeze(0)\n",
    "                new_probs_mask = (orig_label != torch.argmax(new_probs, dim=-1)).data.cpu().numpy()\n",
    "                # prevent bad synonyms\n",
    "                new_probs_mask *= (semantic_sims >= sim_score_threshold)\n",
    "                # prevent incompatible pos\n",
    "                synonyms_pos_ls = [switch.get_pos(new_text[max(idx - 4, 0):idx + 5])[min(4, idx)]\n",
    "                                   if len(new_text) > 10 else switch.get_pos(new_text)[idx] for new_text in new_texts]\n",
    "                pos_mask = np.array(self.pos_filter(pos_ls[idx], synonyms_pos_ls))\n",
    "                new_probs_mask *= pos_mask\n",
    "\n",
    "                if np.sum(new_probs_mask) > 0:\n",
    "                    text_prime[idx] = synonyms[(new_probs_mask * semantic_sims).argmax()]\n",
    "                    num_changed += 1\n",
    "                    break\n",
    "                else:\n",
    "                    new_label_probs = new_probs[:, orig_label] + torch.from_numpy(\n",
    "                        (semantic_sims < sim_score_threshold) + (1 - pos_mask).astype(float)).float().cuda()\n",
    "                    new_label_prob_min, new_label_prob_argmin = torch.min(new_label_probs, dim=-1)\n",
    "                    if new_label_prob_min < orig_prob:\n",
    "                        text_prime[idx] = synonyms[new_label_prob_argmin]\n",
    "                        num_changed += 1\n",
    "                text_cache = text_prime[:]\n",
    "            \n",
    "            if task == 'mnli':\n",
    "                new_label = torch.argmax(self.text_pred({'text_a':[text_a], 'text_b': [text_prime]}, batch_size))\n",
    "            else:\n",
    "                new_label = torch.argmax(self.text_pred({'text_a':[text_prime], 'text_b': [text_b]}, batch_size))\n",
    "\n",
    "            if true_label != new_label:\n",
    "                return TreebankWordDetokenizer().detokenize(text_prime), num_changed, orig_label, new_label, num_queries\n",
    "            else:\n",
    "                return '', num_changed, orig_label, new_label, num_queries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Universal Sentence Encoder encapsulated in a class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "USE (Universal Sentence Encoder) in a class for determining semantic similarities.\n",
    "\n",
    "This class was adapted from code by TextFooler at https://github.com/jind11/TextFooler,\n",
    "a code repository in support of the paper:\n",
    "\n",
    "Jin, Di, et al. \"Is BERT Really Robust? Natural Language Attack on Text Classification and Entailment.\"\n",
    "arXiv preprint arXiv:1907.11932 (2019).\n",
    "\"\"\"\n",
    "\n",
    "class USE(object):\n",
    "    def __init__(self, cache_path):\n",
    "        super(USE, self).__init__()\n",
    "        #config =  tf.compat.v1.ConfigProto()\n",
    "        #config.gpu_options.allow_growth = True\n",
    "        #session =  tf.compat.v1.Session(config=config)\n",
    "        config = tf.ConfigProto()\n",
    "        config.gpu_options.allow_growth = True\n",
    "        self.sess = tf.Session(config=config)\n",
    "        #tf.compat.v1.disable_eager_execution()\n",
    "        \n",
    "        os.environ['TFHUB_CACHE_DIR'] = cache_path\n",
    "        module_url = \"https://tfhub.dev/google/universal-sentence-encoder-large/3\"\n",
    "        self.embed = hub.Module(module_url)\n",
    "\n",
    "        self.build_graph()\n",
    "        self.sess.run([tf.global_variables_initializer(), tf.tables_initializer()])\n",
    "\n",
    "    def build_graph(self):\n",
    "        self.sts_input1 = tf.placeholder(tf.string, shape=(None))\n",
    "        self.sts_input2 = tf.placeholder(tf.string, shape=(None))\n",
    "\n",
    "        sts_encode1 = tf.nn.l2_normalize(self.embed(self.sts_input1), axis=1)\n",
    "        sts_encode2 = tf.nn.l2_normalize(self.embed(self.sts_input2), axis=1)\n",
    "        self.cosine_similarities = tf.reduce_sum(tf.multiply(sts_encode1, sts_encode2), axis=1)\n",
    "        clip_cosine_similarities = tf.clip_by_value(self.cosine_similarities, -1.0, 1.0)\n",
    "        self.sim_scores = 1.0 - tf.acos(clip_cosine_similarities)\n",
    "\n",
    "    def semantic_sim(self, sents1, sents2):\n",
    "        scores = self.sess.run(\n",
    "            [self.sim_scores],\n",
    "            feed_dict={\n",
    "                self.sts_input1: sents1,\n",
    "                self.sts_input2: sents2,\n",
    "            })\n",
    "        return scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Routines to read and scrub datasets (mnli and imdb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Utilities for working with local datasets for processing.\n",
    "\n",
    "These methods were adapted from code by TextFooler at https://github.com/jind11/TextFooler,\n",
    "a code repository in support of the paper:\n",
    "\n",
    "Jin, Di, et al. \"Is BERT Really Robust? Natural Language Attack on Text Classification and Entailment.\"\n",
    "arXiv preprint arXiv:1907.11932 (2019).\n",
    "\"\"\"\n",
    "\n",
    "def clean_str(string, TREC=False):\n",
    "    \"\"\"\n",
    "    Tokenization/string cleaning for all datasets except for SST.\n",
    "    Every dataset is lower cased except for TREC\n",
    "    \"\"\"\n",
    "    string = re.sub(r\"[^A-Za-z0-9(),!?\\'\\`]\", \" \", string)\n",
    "    string = re.sub(r\"\\'s\", \" \\'s\", string)\n",
    "    string = re.sub(r\"\\'ve\", \" \\'ve\", string)\n",
    "    string = re.sub(r\"n\\'t\", \" n\\'t\", string)\n",
    "    string = re.sub(r\"\\'re\", \" \\'re\", string)\n",
    "    string = re.sub(r\"\\'d\", \" \\'d\", string)\n",
    "    string = re.sub(r\"\\'ll\", \" \\'ll\", string)\n",
    "    string = re.sub(r\",\", \" , \", string)\n",
    "    string = re.sub(r\"!\", \" ! \", string)\n",
    "    string = re.sub(r\"\\(\", \" \\( \", string)\n",
    "    string = re.sub(r\"\\)\", \" \\) \", string)\n",
    "    string = re.sub(r\"\\?\", \" \\? \", string)\n",
    "    string = re.sub(r\"\\s{2,}\", \" \", string)\n",
    "    return string.strip() if TREC else string.strip().lower()\n",
    "\n",
    "def read_corpus(path, data_size, clean=True, MR=True, encoding='utf8', shuffle=False, lower=True):\n",
    "    data = []\n",
    "    labels = []\n",
    "    empty = []\n",
    "    with open(path, encoding=encoding) as fin:\n",
    "        for idx, line in enumerate(fin):\n",
    "            if idx >= data_size:\n",
    "                break\n",
    "            if MR:\n",
    "                label, sep, text = line.partition(' ')\n",
    "                label = int(label)\n",
    "            else:\n",
    "                label, sep, text = line.partition(',')\n",
    "                label = int(label) - 1\n",
    "            if clean:\n",
    "                text = clean_str(text.strip()) if clean else text.strip()\n",
    "            if lower:\n",
    "                text = text.lower()\n",
    "            labels.append(label)\n",
    "            data.append(text.split())\n",
    "            empty.append(None)\n",
    "\n",
    "    if shuffle:\n",
    "        perm = list(range(len(data)))\n",
    "        random.shuffle(perm)\n",
    "        data = [data[i] for i in perm]\n",
    "        labels = [labels[i] for i in perm]\n",
    "\n",
    "    return {\"text_a\": data,\n",
    "            \"text_b\": empty,\n",
    "            \"label\": labels}\n",
    "\n",
    "def read_data(filepath, data_size, lowercase=False, ignore_punctuation=False, stopwords=[]):\n",
    "    \"\"\"\n",
    "    Read the premises, hypotheses and labels from some NLI dataset's\n",
    "    file and return them in a dictionary. The file should be in the same\n",
    "    form as SNLI's .txt files.\n",
    "\n",
    "    Args:\n",
    "        filepath: The path to a file containing some premises, hypotheses\n",
    "            and labels that must be read. The file should be formatted in\n",
    "            the same way as the SNLI (and MultiNLI) dataset.\n",
    "\n",
    "    Returns:\n",
    "        A dictionary containing three lists, one for the premises, one for\n",
    "        the hypotheses, and one for the labels in the input data.\n",
    "    \"\"\"\n",
    "    \n",
    "    labeldict = {\"contradiction\": 0,\n",
    "                  \"entailment\": 1,\n",
    "                  \"neutral\": 2}\n",
    "\n",
    "    with open(filepath, 'r', encoding='utf8') as input_data:\n",
    "        premises, hypotheses, labels = [], [], []\n",
    "\n",
    "        # Translation tables to remove punctuation from strings.\n",
    "        punct_table = str.maketrans({key: ' '\n",
    "                                     for key in string.punctuation})\n",
    "\n",
    "        for idx, line in enumerate(input_data):\n",
    "            if idx >= data_size:\n",
    "                break\n",
    "\n",
    "            line = line.strip().split('\\t')\n",
    "\n",
    "            # Ignore sentences that have no gold label.\n",
    "            if line[0] == '-':\n",
    "                continue\n",
    "            \n",
    "            # skip the header row (if there is one)\n",
    "            if line[0] == 'gold_label':\n",
    "                continue\n",
    "\n",
    "            premise = line[1]\n",
    "            hypothesis = line[2]\n",
    "\n",
    "            if lowercase:\n",
    "                premise = premise.lower()\n",
    "                hypothesis = hypothesis.lower()\n",
    "\n",
    "            if ignore_punctuation:\n",
    "                premise = premise.translate(punct_table)\n",
    "                hypothesis = hypothesis.translate(punct_table)\n",
    "                \n",
    "            # strip ('s and )'s\n",
    "            premise = premise.translate({ord(i):None for i in '()'})\n",
    "            hypothesis = hypothesis.translate({ord(i):None for i in '()'})\n",
    "            \n",
    "            # Each premise and hypothesis is split into a list of words.\n",
    "            premises.append([w for w in premise.rstrip().split()\n",
    "                             if w not in stopwords])\n",
    "            hypotheses.append([w for w in hypothesis.rstrip().split()\n",
    "                             if w not in stopwords])\n",
    "            labels.append(labeldict[line[0]])\n",
    "\n",
    "        return {\"text_a\": premises,\n",
    "                \"text_b\": hypotheses,\n",
    "                \"label\": labels}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Routine to generate samples from the fooler and dataset path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "def elapsed_time():\n",
    "    global t_start\n",
    "\n",
    "    t_now = time.time()\n",
    "    t = t_now-t_start\n",
    "    t_start = t_now\n",
    "    return t\n",
    "\n",
    "def generate_samples(task, fooler, dataset_path, data_size, batch_size):\n",
    "    \n",
    "    # get data to attack, fetch first [data_size] data samples for adversarial attacking\n",
    "    \n",
    "    if task == 'mnli':\n",
    "        dataloader = read_data\n",
    "        labeldict = {0: \"contradiction\",\n",
    "                     1: \"entailment\",\n",
    "                     2:  \"neutral\"}\n",
    "    else:\n",
    "        #imdb\n",
    "        dataloader = read_corpus\n",
    "        labeldict = {0: 0, 1: 1}\n",
    "\n",
    "    data = dataloader(dataset_path, data_size)\n",
    "    print(\"Data import finished!\")\n",
    "        \n",
    "    test_examples = [InputExample(i, TreebankWordDetokenizer().detokenize(a), \n",
    "                                  TreebankWordDetokenizer().detokenize(b) if b is not None else None,\n",
    "                                  labeldict[label]) \\\n",
    "                     for i,(a,b,label) in \\\n",
    "                     enumerate(zip(data['text_a'], data['text_b'], data['label']))]\n",
    "    \n",
    "\n",
    "    fooler.model.set_test_dataset(None, examples=test_examples)\n",
    "    \n",
    "    orig_failures = 0.\n",
    "    adv_failures = 0.\n",
    "    changed_rates = []\n",
    "    nums_queries = []\n",
    "    \n",
    "    adv_examples=[]\n",
    "    \n",
    "    for idx, text_a in enumerate(data['text_a']):\n",
    "        if idx % 10 == 0:\n",
    "            print('elapsed time: {}s - {} samples out of {} have been finished!'.format(\n",
    "                round(elapsed_time(),2), idx, data_size))\n",
    "\n",
    "            message = 'accuracy: {:.3f}%, adv accuracy: {:.3f}%, ' \\\n",
    "              'avg changed rate: {:.3f}%, num of queries: {:.1f}\\n'.format((1-orig_failures/(idx+1))*100,\n",
    "                                                                 (1-adv_failures/(idx+1))*100,\n",
    "                                                                 np.mean(changed_rates)*100,\n",
    "                                                                 np.mean(nums_queries))\n",
    "            print(message)\n",
    "\n",
    "        text_b, true_label = data['text_b'][idx], data['label'][idx]\n",
    "                    \n",
    "        new_text, num_changed, orig_label, \\\n",
    "            new_label, num_queries  = fooler.generate_adversarial(task, text_a, text_b, true_label,\n",
    "                                                                  batch_size=batch_size,)\n",
    "        if true_label != orig_label:\n",
    "            orig_failures += 1\n",
    "        else:\n",
    "            nums_queries.append(num_queries)\n",
    "        if true_label != new_label:\n",
    "            adv_failures += 1\n",
    "            if new_text != '':\n",
    "                adv_examples.append(\n",
    "                    InputExample(guid=idx,\n",
    "                                 text_a=TreebankWordDetokenizer().detokenize(text_a) if task == 'mnli' else new_text,\n",
    "                                 text_b=new_text if task == 'mnli' else text_b,\n",
    "                                 label=labeldict[true_label] if task == 'mnli' else true_label))\n",
    "            \n",
    "        changed_rate = 1.0 * num_changed / (len(text_b) if task == 'mnli' else len(text_a))\n",
    "        if true_label == orig_label and true_label != new_label:\n",
    "            changed_rates.append(changed_rate)\n",
    "\n",
    "            \n",
    "    print('elapsed time: {}s - {} samples out of {} have been finished!'.format(\n",
    "        round(elapsed_time(),2), idx+1, data_size))            \n",
    "            \n",
    "    message = 'For target model {}: original accuracy: {:.3f}%, adv accuracy: {:.3f}%, ' \\\n",
    "              'avg changed rate: {:.3f}%, num of queries: {:.1f}\\n'.format(task,\n",
    "                                                                 (1-orig_failures/(idx+1))*100,\n",
    "                                                                 (1-adv_failures/(idx+1))*100,\n",
    "                                                                 np.mean(changed_rates)*100,\n",
    "                                                                 np.mean(nums_queries))\n",
    "    print(message)\n",
    "\n",
    "    return adv_examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Routines to create and save the samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_examples(bert_model, task, dataset_path, data_size, batch_size, max_seq_length,\n",
    "                    counter_fitting_embeddings_path, counter_fitting_cos_sim_path, USE_cache_path):\n",
    "    \n",
    "    \n",
    "    print(\"Building TextFooler...\")\n",
    "    fooler = PaperFooler(bert_model.tokenizer,\n",
    "                         bert_model,\n",
    "                         USE_cache_path,\n",
    "                         counter_fitting_embeddings_path,\n",
    "                         counter_fitting_cos_sim_path,\n",
    "                         max_seq_length = max_seq_length)\n",
    "    print(\"TextFooler built!\")\n",
    "\n",
    "    \n",
    "    return generate_samples(task, fooler, dataset_path, data_size, batch_size)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_examples(bert_model, adv_examples, output_path):\n",
    "    \n",
    "    features = bert_model.get_processor()._create_features(adv_examples)\n",
    "    \n",
    "    torch.save(features, output_path)\n",
    "    \n",
    "    with open(output_path + '.txt', 'w') as output:\n",
    "        for row in adv_examples:\n",
    "            output.write(row.label + '\\t' +\n",
    "                         row.text_a + '\\t' + \n",
    "                         row.text_b + '\\n') if task == 'mnli' else \\\n",
    "            output.write(str(row.label) + ' ' + \n",
    "                         row.text_a + '\\n')\n",
    "            \n",
    "    import pickle\n",
    "    with open(output_path + '.pkl', \"wb\") as f:\n",
    "        pickle.dump(adv_examples, f)\n",
    "    \n",
    "    print('\\nPyTorch-ready InputFeature file saved in {}'.format(output_path))\n",
    "    print('Pickled InputExample file saved in {}.pkl'.format(output_path))\n",
    "    print('Raw text saved in {}.txt'.format(output_path))\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Config options if we want to use evaluation-time perturbation models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_eval_model = False\n",
    "\n",
    "eval_model_type = 'FUSE' # can be FUSE or FIVE\n",
    "\n",
    "#FIVE best MNLI params from the paper\n",
    "# eval_model_hparams =  {'batch_size': 8, 'use_USE': False, 'stop_words': True, 'perturb_words': 1, \n",
    "#                          'verbose': False, 'vote_avg_logits': True, 'std': 8.139999999999995, 'vector_count': 8}\n",
    "\n",
    "#FuSE best MNLI params from the paper\n",
    "eval_model_hparams =  {'use_USE':True, 'USE_method':\"filter\", 'USE_multiplier':14, 'stop_words':True, 'perturb_words':2,\n",
    "            'candidates_per_word':10, 'total_alternatives':14, 'match_pos':True, 'batch_size':1,'verbose':False, \n",
    "            'vote_avg_logits':True}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize our underlying BERT model tuned on the task to fool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "processor = MnliProcessor() if task == 'mnli' else ImdbProcessor()\n",
    "\n",
    "if use_eval_model:\n",
    "    bert_model = (FireBERT_FSE(load_from=model_path, \n",
    "                               processor=processor, \n",
    "                               hparams=eval_model_hparams) if eval_model_type == 'FUSE' else\n",
    "                  FireBERT_FVE(load_from=model_path, \n",
    "                               processor=processor, \n",
    "                               hparams=eval_model_hparams))\n",
    "else:    \n",
    "    bert_model = LightningBertForSequenceClassification(\n",
    "        load_from=model_path, \n",
    "        processor=processor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate samples and benchmarks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_start = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "adv_examples = create_examples(bert_model, task, dataset_path, data_size, batch_size, 128 if task =='mnli' else 256, \n",
    "                               counter_fitting_embeddings_path, counter_fitting_cos_sim_path, USE_cache_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Uncomment if you want to see the samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "#adv_examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Comment or uncomment depending on whether you want to save them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "#save_examples(bert_model, adv_examples, output_path)"
   ]
  },
  {
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}