{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import transformers\n",
    "import metrics\n",
    "import spacy\n",
    "import pandas as pd\n",
    "\n",
    "nlp = spacy.load(\"en_core_web_trf\")\n",
    "\n",
    "class ParaphraseModel(object):\n",
    "    def __init__(self, model_name, device=\"cuda\"):\n",
    "        self.model_name = model_name\n",
    "        self.device = device\n",
    "        self.tokenizer = transformers.AutoTokenizer.from_pretrained(self.model_name, cache_dir=\"./cache/models/\")\n",
    "        if \"pegasus\" in self.model_name.lower():\n",
    "            self.model = transformers.PegasusForConditionalGeneration.from_pretrained(self.model_name, cache_dir=\"./cache/models/\")\n",
    "        elif \"bart\" in self.model_name.lower():\n",
    "            self.model = transformers.BartForConditionalGeneration.from_pretrained(self.model_name, cache_dir=\"./cache/models/\")\n",
    "        elif \"t5\" in self.model_name.lower():\n",
    "            self.model = transformers.AutoModelForSeq2SeqLM.from_pretrained(self.model_name, cache_dir=\"./cache/models/\")\n",
    "        self.model = self.model.to(self.device)\n",
    "        \n",
    "    def post_process(self, output_texts):\n",
    "        output_texts = [t.strip(' •-\"') for t in output_texts]\n",
    "        return output_texts\n",
    "        \n",
    "    def generate(self, input_text, output_count=10, max_length=256, beam_width=2, num_beam_groups=2):\n",
    "        if \"t5\" in self.model_name.lower():\n",
    "            input_text = \"paraphrase: \" + input_text\n",
    "        model_inputs = self.tokenizer([input_text],\n",
    "                                      truncation=\"longest_first\",\n",
    "                                      padding=\"longest\",\n",
    "                                      max_length=max_length,\n",
    "                                      return_tensors=\"pt\").to(self.device)\n",
    "        translated = self.model.generate(**model_inputs,\n",
    "                                         num_return_sequences=output_count,\n",
    "                                         # number of beams for beam search\n",
    "                                         num_beams=int(output_count*beam_width),\n",
    "                                         # number of groups to divide num_beams into in order to ensure diversity\n",
    "                                         num_beam_groups=num_beam_groups,\n",
    "                                         repetition_penalty=1.2,\n",
    "                                         # higher the penalty, the more diverse are the outputs\n",
    "                                         diversity_penalty=0.3,\n",
    "                                         early_stopping=True,)\n",
    "        tgt_text = self.tokenizer.batch_decode(translated, skip_special_tokens=True)\n",
    "        output_texts = self.post_process(tgt_text)\n",
    "        output_texts = list(set(output_texts))\n",
    "        output_texts.sort()\n",
    "        return output_texts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#model = ParaphraseModel(\"./pegasus-model\")\n",
    "#model = ParaphraseModel(\"./bart-model\")\n",
    "model = ParaphraseModel(\"./t5-large-model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate(input_text, output_count=32, keep_word=None, sort_by=\"wpd\"):\n",
    "    paraphrase_list = model.generate(input_text, output_count=output_count, max_length=64)\n",
    "    nlp_input_text = nlp(input_text)\n",
    "    out_para = []\n",
    "    out_wpd = []\n",
    "    out_ld = []\n",
    "    if sort_by==\"wpd\":\n",
    "        wpd_list = [metrics.wpd(nlp_input_text, nlp(t)) for t in paraphrase_list]\n",
    "        sorted_wpd = [[wpd,para] for wpd,para in sorted(zip(wpd_list, paraphrase_list))]\n",
    "        max_wpd = max(wpd_list)\n",
    "        for wpd, para in sorted_wpd:\n",
    "            nlp_para = nlp(para)\n",
    "            wpd = metrics.wpd(nlp_input_text, nlp_para)\n",
    "            ld = metrics.ld(nlp_input_text, nlp_para)\n",
    "            if keep_word:\n",
    "                para_tokens = [token.text.lower() for token in nlp_para]\n",
    "                if keep_word in para_tokens:\n",
    "                    if wpd > 0.0:\n",
    "                        out_para.append(para)\n",
    "                        out_wpd.append(wpd)\n",
    "                        out_ld.append(ld)\n",
    "            else:\n",
    "                if wpd > 0.0:\n",
    "                    out_para.append(para)\n",
    "                    out_wpd.append(wpd)\n",
    "                    out_ld.append(ld)\n",
    "    else:\n",
    "        ld_list = [metrics.ld(nlp_input_text, nlp(t)) for t in paraphrase_list]\n",
    "        sorted_ld = [[ld,para] for ld,para in sorted(zip(ld_list, paraphrase_list))]\n",
    "        max_ld = max(ld_list)\n",
    "        for ld, para in sorted_ld:\n",
    "            nlp_para = nlp(para)\n",
    "            wpd = metrics.wpd(nlp_input_text, nlp_para)\n",
    "            ld = metrics.ld(nlp_input_text, nlp_para)\n",
    "            if keep_word:\n",
    "                para_tokens = [token.text.lower() for token in nlp_para]\n",
    "                if keep_word in para_tokens:\n",
    "                    if wpd > 0.0:\n",
    "                        out_para.append(para)\n",
    "                        out_wpd.append(wpd)\n",
    "                        out_ld.append(ld)\n",
    "            else:\n",
    "                if wpd > 0.0:\n",
    "                    out_para.append(para)\n",
    "                    out_wpd.append(wpd)\n",
    "                    out_ld.append(ld)\n",
    "    return out_para, out_wpd, out_ld"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "bed\n",
    "* There's a lot of trash on the bed of the river.\n",
    "* I keep a glass of water next to my bed when I sleep."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_text = \"I keep a glass of water next to my bed when I sleep.\"\n",
    "out_para, out_wpd, out_ld = generate(input_text, output_count=128, sort_by=\"wpd\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format='retina'\n",
    "import seaborn as sns\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.title(\"T5 Model Output WPD\")\n",
    "plt.xlim([0, 1.0])\n",
    "sns.histplot(out_wpd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.title(\"T5 Model Output LD\")\n",
    "plt.xlim([0, 1.0])\n",
    "sns.histplot(out_ld)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#input_text = \"I keep a glass of water next to my bed when I sleep.\"\n",
    "#out_para, out_wpd, out_ld = generate(input_text, output_count=128, sort_by=\"wpd\")\n",
    "for i, wpd in zip(out_para, out_wpd):\n",
    "    wpd = round(wpd, 3)\n",
    "    print(i, wpd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_text = \"I keep a glass of water next to my bed when I sleep.\"\n",
    "out_para, out_wpd, out_ld= generate(input_text, output_count=128, sort_by=\"ld\")\n",
    "for i, ld in zip(out_para, out_ld):\n",
    "    ld = round(ld, 3)\n",
    "    print(i, ld)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s1 = \"I keep a glass of water next to my bed when I sleep.\"\n",
    "s2 = \"During the night, I keep a glass of water next to my bed.\"\n",
    "s1, s2 = nlp(s1), nlp(s2)\n",
    "metrics.ld(s1, s2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s1 = \"I keep a glass of water next to my bed when I sleep.\"\n",
    "s2 = \"When I sleep, I keep a glass of water by my bedside.\"\n",
    "s1, s2 = nlp(s1), nlp(s2)\n",
    "metrics.ld(s1, s2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_text = \"There's a lot of trash on the bed of the river.\"\n",
    "out_para, _, _ = generate(input_text, keep_word=\"bed\")\n",
    "for i in out_para[:3]:\n",
    "    print(i)\n",
    "for i in out_para[-3:]:\n",
    "    print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_text = \"I keep a glass of water next to my bed when I sleep.\"\n",
    "out_para, _, _ = generate(input_text, keep_word=\"bed\")\n",
    "for i in out_para[:3]:\n",
    "    print(i)\n",
    "for i in out_para[-3:]:\n",
    "    print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_text = \"The expanded window will give us time to catch the thieves.\"\n",
    "out_para, _, _ = generate(input_text, keep_word=\"window\")\n",
    "for i in out_para[:3]:\n",
    "    print(i)\n",
    "for i in out_para[-3:]:\n",
    "    print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_text = \"You have a two-hour window of clear weather to finish working on the lawn.\"\n",
    "out_para, _, _ = generate(input_text, keep_word=\"window\")\n",
    "for i in out_para[:3]:\n",
    "    print(i)\n",
    "for i in out_para[-3:]:\n",
    "    print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_text = \"The pilot managed to land the airplane safely.\"\n",
    "out_para, _, _ = generate(input_text, keep_word=\"land\")\n",
    "for i in out_para[:3]:\n",
    "    print(i)\n",
    "for i in out_para[-3:]:\n",
    "    print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_text = \"The enemy landed several of our aircrafts.\"\n",
    "out_para, _, _ = generate(input_text, keep_word=\"land\")\n",
    "for i in out_para[:3]:\n",
    "    print(i)\n",
    "for i in out_para[-3:]:\n",
    "    print(i)"
   ]
  },
  {
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
