IteraTeR_demo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "IteraTeR_demo.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "48f6f4434a7640fea1fefa91da47e2c2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_ec389ecec5f24d2f885837f49482f870",
              "IPY_MODEL_e47d659999dc46629575b82284ec631d",
              "IPY_MODEL_7934c296ad83424da0fddd09e25f3f18"
            ],
            "layout": "IPY_MODEL_5ad42743e1c44fd29c5d71a044d03251"
          }
        },
        "ec389ecec5f24d2f885837f49482f870": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_f5587e92e13c4b86b8328ac0330b9b19",
            "placeholder": "​",
            "style": "IPY_MODEL_5b8d59d326bd4a69918465274623db3a",
            "value": "100%"
          }
        },
        "e47d659999dc46629575b82284ec631d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_c99f8f9fbe5444f9b3136e9b2c072732",
            "max": 3,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_4048e4b739a643b19ab3aee8477c5694",
            "value": 3
          }
        },
        "7934c296ad83424da0fddd09e25f3f18": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_e8eb9ae38af74ea6a6c314e6842ee1c6",
            "placeholder": "​",
            "style": "IPY_MODEL_5c836fc42ea743a7b59ae504a0876f3f",
            "value": " 3/3 [00:00&lt;00:00, 47.42it/s]"
          }
        },
        "5ad42743e1c44fd29c5d71a044d03251": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f5587e92e13c4b86b8328ac0330b9b19": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "5b8d59d326bd4a69918465274623db3a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "c99f8f9fbe5444f9b3136e9b2c072732": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "4048e4b739a643b19ab3aee8477c5694": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "e8eb9ae38af74ea6a6c314e6842ee1c6": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "5c836fc42ea743a7b59ae504a0876f3f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Demo code for iterative text revision on IteraTeR dataset"
      ],
      "metadata": {
        "id": "ZmWGwxJJu4bz"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aKAk1KRmu3Et"
      },
      "outputs": [],
      "source": [
        "!pip install datasets\n",
        "!pip install transformers\n",
        "!pip install torch==1.5.1\n",
        "!pip install torchvision==0.6.1\n",
        "!pip install levenshtein"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from datasets import load_dataset\n",
        "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM\n",
        "from rapidfuzz.distance import Levenshtein\n",
        "\n",
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
      ],
      "metadata": {
        "id": "r78oJAhQyE-P"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# load dataset\n",
        "dataset = load_dataset(\"wanyu/IteraTeR_human_sent\")\n",
        "dataset"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 355,
          "referenced_widgets": [
            "48f6f4434a7640fea1fefa91da47e2c2",
            "ec389ecec5f24d2f885837f49482f870",
            "e47d659999dc46629575b82284ec631d",
            "7934c296ad83424da0fddd09e25f3f18",
            "5ad42743e1c44fd29c5d71a044d03251",
            "f5587e92e13c4b86b8328ac0330b9b19",
            "5b8d59d326bd4a69918465274623db3a",
            "c99f8f9fbe5444f9b3136e9b2c072732",
            "4048e4b739a643b19ab3aee8477c5694",
            "e8eb9ae38af74ea6a6c314e6842ee1c6",
            "5c836fc42ea743a7b59ae504a0876f3f"
          ]
        },
        "id": "65OP3mKevdhk",
        "outputId": "b9c23b36-0971-4c94-e6ff-fa2d900ccc59"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Using custom data configuration wanyu--IteraTeR_human_sent-7ef9360fff9b86ec\n",
            "Reusing dataset json (/root/.cache/huggingface/datasets/json/wanyu--IteraTeR_human_sent-7ef9360fff9b86ec/0.0.0/ac0ca5f5289a6cf108e706efcf040422dbbfa8e658dee6a819f20d76bb84d26b)\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/3 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "48f6f4434a7640fea1fefa91da47e2c2"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DatasetDict({\n",
              "    train: Dataset({\n",
              "        features: ['before_sent', 'before_sent_with_intent', 'after_sent', 'labels', 'doc_id', 'revision_depth'],\n",
              "        num_rows: 3254\n",
              "    })\n",
              "    test: Dataset({\n",
              "        features: ['before_sent', 'before_sent_with_intent', 'after_sent', 'labels', 'doc_id', 'revision_depth'],\n",
              "        num_rows: 364\n",
              "    })\n",
              "    validation: Dataset({\n",
              "        features: ['before_sent', 'before_sent_with_intent', 'after_sent', 'labels', 'doc_id', 'revision_depth'],\n",
              "        num_rows: 400\n",
              "    })\n",
              "})"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# load model\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"wanyu/IteraTeR-PEGASUS-Revision-Generator\")\n",
        "model = AutoModelForSeq2SeqLM.from_pretrained(\"wanyu/IteraTeR-PEGASUS-Revision-Generator\")\n",
        "model.to(device)\n",
        "model.eval()"
      ],
      "metadata": {
        "id": "sYPxin3yvxrO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# prepare input to the model: <intention> before_sent\n",
        "before_text = dataset['test'][0]['before_sent']\n",
        "before_input = dataset['test'][0]['before_sent_with_intent']\n",
        "edit_intent = dataset['test'][0]['labels']\n",
        "print('Model input:', before_input)\n",
        "model_input = tokenizer(before_input, return_tensors='pt')\n",
        "model_input.to(device)\n",
        "print('Tokenized model input:', model_input)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PqBdcvFowFIF",
        "outputId": "c9569415-5ed6-4649-98f2-e22984d373f4"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model input: <clarity> In this paper , we present a new sequence-to-sequence pre-training model called ProphetNet, which introduces a novel self-supervised objective named future n-gram prediction and the proposed n-stream self-attention mechanism.\n",
            "Tokenized model input: {'input_ids': tensor([[96103,   222,   136,   800,   110,   108,   145,   799,   114,   177,\n",
            "          5936,   121,   497,   121, 69987,  1133,   121, 18006,   861,   568,\n",
            "         17857,  7614,   108,   162,  9905,   114,  2794,   813,   121, 83465,\n",
            "          4129,  1729,   533,  3178,   121, 14506, 11955,   111,   109,  2962,\n",
            "          3178,   121, 12342,   813,   121, 65167,  5661,   107,     1]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
            "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
            "         1]])}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# start iterative text revision\n",
        "MAX_ITER = 10\n",
        "for t in range(MAX_ITER):\n",
        "  # get current model prediction: after_sent\n",
        "  model_outputs = model.generate(**model_input, num_beams=8, max_length=1024)\n",
        "  after_text = tokenizer.batch_decode(model_outputs, skip_special_tokens=True)[0]\n",
        "  print(f'Model output at revision depth {t+1}: {after_text}')\n",
        "  edit_dist = Levenshtein.distance(before_text, after_text)\n",
        "  print(f'Edit distance at revision depth {t+1}: {edit_dist}')\n",
        "  print('============================================')\n",
        "  # check stopping criteria\n",
        "  if edit_dist == 0:\n",
        "    print(f'Model final output: {after_text}')\n",
        "    break\n",
        "  else:\n",
        "    before_text = after_text\n",
        "    before_input = f'<{edit_intent}> {before_text}'\n",
        "    model_input = tokenizer(before_input, return_tensors='pt').to(device)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gWur7fYnxXj-",
        "outputId": "f4405111-26d5-423e-8eb8-51a8f9682c6a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model output at revision depth 1: In this paper, we present a novel sequence-to-sequence pre-training model called ProphetNet, which introduces a novel self-supervised objective named future n-gram prediction and the proposed n-stream self-attention mechanism.\n",
            "Edit distance at revision depth 1: 4\n",
            "============================================\n",
            "Model output at revision depth 2: In this paper, we present a sequence-to-sequence pre-training model called ProphetNet, which introduces a novel self-supervised objective named future n-gram prediction and the proposed n-stream self-attention mechanism.\n",
            "Edit distance at revision depth 2: 6\n",
            "============================================\n",
            "Model output at revision depth 3: In this paper, we present a sequence-to-sequence pre-training model called ProphetNet, which introduces a novel self-supervised objective named future n-gram prediction and the proposed n-stream self-attention mechanism.\n",
            "Edit distance at revision depth 3: 0\n",
            "============================================\n",
            "Model final output: In this paper, we present a sequence-to-sequence pre-training model called ProphetNet, which introduces a novel self-supervised objective named future n-gram prediction and the proposed n-stream self-attention mechanism.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "iumurC3Cx-wy"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}