emnlp2021/style_transfer/condBERT/condbert_compile_vocab.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook reproduces creation of CondBERT vocabulary.\n",
    "\n",
    "The files `positive-words.txt`, `negative-words.txt` and `toxic_words.txt` are not reproduced exactly because of our internal issues. \n",
    "\n",
    "However, all other files (`token_toxicities.txt` and `word2coef.pkl` ) are reproduced accurately. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 0. Prerequisites"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "VOCAB_DIRNAME = 'vocabularies' "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from condbert import CondBertRewriter\n",
    "from choosers import EmbeddingSimilarityChooser\n",
    "from multiword.masked_token_predictor_bert import MaskedTokenPredictorBert"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Loading BERT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import BertTokenizer, BertForMaskedLM\n",
    "import numpy as np\n",
    "import pickle\n",
    "import os\n",
    "from tqdm.auto import tqdm, trange"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "device = torch.device('cuda:0')\n",
    "device = torch.device('cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'bert-base-uncased'\n",
    "tokenizer = BertTokenizer.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForMaskedLM: ['cls.seq_relationship.weight', 'cls.seq_relationship.bias']\n",
      "- This IS expected if you are initializing BertForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "model = BertForMaskedLM.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.to(device);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Preparing the vocabularires."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "- negative-words.txt\n",
    "- positive-words.txt\n",
    "- word2coef.pkl\n",
    "- token_toxicities.txt\n",
    "\n",
    "These files should be prepared once. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "tox_corpus_path = '../../data/train/train_toxic'\n",
    "norm_corpus_path = '../../data/train/train_normal'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(VOCAB_DIRNAME):\n",
    "    os.makedirs(VOCAB_DIRNAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Preparing the DRG-like vocabularies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import argparse\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "from nltk import ngrams\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer\n",
    "\n",
    "\n",
    "\n",
    "class NgramSalienceCalculator():\n",
    "    def __init__(self, tox_corpus, norm_corpus, use_ngrams=False):\n",
    "        ngrams = (1, 3) if use_ngrams else (1, 1)\n",
    "        self.vectorizer = CountVectorizer(ngram_range=ngrams)\n",
    "\n",
    "        tox_count_matrix = self.vectorizer.fit_transform(tox_corpus)\n",
    "        self.tox_vocab = self.vectorizer.vocabulary_\n",
    "        self.tox_counts = np.sum(tox_count_matrix, axis=0)\n",
    "\n",
    "        norm_count_matrix = self.vectorizer.fit_transform(norm_corpus)\n",
    "        self.norm_vocab = self.vectorizer.vocabulary_\n",
    "        self.norm_counts = np.sum(norm_count_matrix, axis=0)\n",
    "\n",
    "    def salience(self, feature, attribute='tox', lmbda=0.5):\n",
    "        assert attribute in ['tox', 'norm']\n",
    "        if feature not in self.tox_vocab:\n",
    "            tox_count = 0.0\n",
    "        else:\n",
    "            tox_count = self.tox_counts[0, self.tox_vocab[feature]]\n",
    "\n",
    "        if feature not in self.norm_vocab:\n",
    "            norm_count = 0.0\n",
    "        else:\n",
    "            norm_count = self.norm_counts[0, self.norm_vocab[feature]]\n",
    "\n",
    "        if attribute == 'tox':\n",
    "            return (tox_count + lmbda) / (norm_count + lmbda)\n",
    "        else:\n",
    "            return (norm_count + lmbda) / (tox_count + lmbda)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "88645\n"
     ]
    }
   ],
   "source": [
    "from collections import Counter\n",
    "c = Counter()\n",
    "\n",
    "for fn in [tox_corpus_path, norm_corpus_path]:\n",
    "    with open(fn, 'r') as corpus:\n",
    "        for line in corpus.readlines():\n",
    "            for tok in line.strip().split():\n",
    "                c[tok] += 1\n",
    "\n",
    "print(len(c))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "88645\n"
     ]
    }
   ],
   "source": [
    "vocab = {w for w, _ in c.most_common() if _ > 0}  # if we took words with > 1 occurences, vocabulary would be x2 smaller, but we'll survive this size\n",
    "print(len(vocab))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(tox_corpus_path, 'r') as tox_corpus, open(norm_corpus_path, 'r') as norm_corpus:\n",
    "    corpus_tox = [' '.join([w if w in vocab else '<unk>' for w in line.strip().split()]) for line in tox_corpus.readlines()]\n",
    "    corpus_norm = [' '.join([w if w in vocab else '<unk>' for w in line.strip().split()]) for line in norm_corpus.readlines()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "neg_out_name = VOCAB_DIRNAME + '/negative-words.txt'\n",
    "pos_out_name = VOCAB_DIRNAME + '/positive-words.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "threshold = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc = NgramSalienceCalculator(corpus_tox, corpus_norm, False)\n",
    "seen_grams = set()\n",
    "\n",
    "with open(neg_out_name, 'w') as neg_out, open(pos_out_name, 'w') as pos_out:\n",
    "    for gram in set(sc.tox_vocab.keys()).union(set(sc.norm_vocab.keys())):\n",
    "        if gram not in seen_grams:\n",
    "            seen_grams.add(gram)\n",
    "            toxic_salience = sc.salience(gram, attribute='tox')\n",
    "            polite_salience = sc.salience(gram, attribute='norm')\n",
    "            if toxic_salience > threshold:\n",
    "                neg_out.writelines(f'{gram}\\n')\n",
    "            elif polite_salience > threshold:\n",
    "                pos_out.writelines(f'{gram}\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.2 Evaluating word toxicities with a logistic regression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.pipeline import make_pipeline\n",
    "pipe = make_pipeline(CountVectorizer(), LogisticRegression(max_iter=1000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = corpus_tox + corpus_norm\n",
    "y_train = [1] * len(corpus_tox) + [0] * len(corpus_norm)\n",
    "pipe.fit(X_train, y_train);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(88519,)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coefs = pipe[1].coef_[0]\n",
    "coefs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "word2coef = {w: coefs[idx] for w, idx in pipe[0].vocabulary_.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open(VOCAB_DIRNAME + '/word2coef.pkl', 'wb') as f:\n",
    "    pickle.dump(word2coef, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.3 Labelling BERT tokens by toxicity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 135390/135390 [00:43<00:00, 3088.85it/s]\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 135390/135390 [00:45<00:00, 2977.70it/s]\n"
     ]
    }
   ],
   "source": [
    "from collections import defaultdict\n",
    "toxic_counter = defaultdict(lambda: 1)\n",
    "nontoxic_counter = defaultdict(lambda: 1)\n",
    "\n",
    "for text in tqdm(corpus_tox):\n",
    "    for token in tokenizer.encode(text):\n",
    "        toxic_counter[token] += 1\n",
    "for text in tqdm(corpus_norm):\n",
    "    for token in tokenizer.encode(text):\n",
    "        nontoxic_counter[token] += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "token_toxicities = [toxic_counter[i] / (nontoxic_counter[i] + toxic_counter[i]) for i in range(len(tokenizer.vocab))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(VOCAB_DIRNAME + '/token_toxicities.txt', 'w') as f:\n",
    "    for t in token_toxicities:\n",
    "        f.write(str(t))\n",
    "        f.write('\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Setting up the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1 Loading the vocabularies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(VOCAB_DIRNAME + \"/negative-words.txt\", \"r\") as f:\n",
    "    s = f.readlines()\n",
    "negative_words = list(map(lambda x: x[:-1], s))\n",
    "\n",
    "with open(VOCAB_DIRNAME + \"/positive-words.txt\", \"r\") as f:\n",
    "    s = f.readlines()\n",
    "positive_words = list(map(lambda x: x[:-1], s))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open(VOCAB_DIRNAME + '/word2coef.pkl', 'rb') as f:\n",
    "    word2coef = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "token_toxicities = []\n",
    "with open(VOCAB_DIRNAME + '/token_toxicities.txt', 'r') as f:\n",
    "    for line in f.readlines():\n",
    "        token_toxicities.append(float(line))\n",
    "token_toxicities = np.array(token_toxicities)\n",
    "token_toxicities = np.maximum(0, np.log(1/(1/token_toxicities-1)))   # log odds ratio\n",
    "\n",
    "# discourage meaningless tokens\n",
    "for tok in ['.', ',', '-']:\n",
    "    token_toxicities[tokenizer.encode(tok)][1] = 3\n",
    "\n",
    "for tok in ['you']:\n",
    "    token_toxicities[tokenizer.encode(tok)][1] = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "def adjust_logits(logits, label=0):\n",
    "    return logits - token_toxicities * 100 * (1 - 2 * label)\n",
    "\n",
    "predictor = MaskedTokenPredictorBert(model, tokenizer, max_len=250, device=device, label=0, contrast_penalty=0.0, logits_postprocessor=adjust_logits)\n",
    "\n",
    "editor = CondBertRewriter(\n",
    "    model=model,\n",
    "    tokenizer=tokenizer,\n",
    "    device=device,\n",
    "    neg_words=negative_words,\n",
    "    pos_words=positive_words,\n",
    "    word2coef=word2coef,\n",
    "    token_toxicities=token_toxicities,\n",
    "    predictor=predictor,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model below is used for reranking BERT hypotheses and helps to increase semantic similarity by choosing the hypotheses with  embeddings similar to the orignal words. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "chooser = EmbeddingSimilarityChooser(sim_coef=10, tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Finally, the inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Parallel application of the model to all tokens, fast, but dirty. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "you are mistake !\n"
     ]
    }
   ],
   "source": [
    "print(editor.translate('You are idiot!', prnt=False))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Application of the model to all the tokens sequentially, in the multiword mode. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "you are very beautiful !\n"
     ]
    }
   ],
   "source": [
    "print(editor.replacement_loop('You are stupid!', verbose=False, chooser=chooser, n_tokens=(1, 2, 3), n_top=10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Parameters that could be tuned:\n",
    "* The coeffincient in `adjust_logits` - the larger it is, the more the model avoids toxic words\n",
    "* The coefficient in `EmbeddingSimilarityChooser` - the larger it is, the more the model tries to preserve content \n",
    "* n_tokens - how many words can be generated from one\n",
    "* n_top - how many BERT hypotheses are reranked"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}