BertSum_Test.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "BertSum Test",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "evij7LOqc2ry",
        "outputId": "4939755f-3302-4f7c-fb95-d491f02964e8"
      },
      "source": [
        "!pip3 install torch==1.8.1+cpu torchvision==0.9.1+cpu torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html"
      ],
      "execution_count": 48,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Looking in links: https://download.pytorch.org/whl/torch_stable.html\n",
            "Requirement already satisfied: torch==1.8.1+cpu in /usr/local/lib/python3.7/dist-packages (1.8.1+cpu)\n",
            "Requirement already satisfied: torchvision==0.9.1+cpu in /usr/local/lib/python3.7/dist-packages (0.9.1+cpu)\n",
            "Requirement already satisfied: torchaudio==0.8.1 in /usr/local/lib/python3.7/dist-packages (0.8.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch==1.8.1+cpu) (3.7.4.3)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torch==1.8.1+cpu) (1.19.5)\n",
            "Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.7/dist-packages (from torchvision==0.9.1+cpu) (7.1.2)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hVuplEqWg2Om"
      },
      "source": [
        "# Load pyrouge"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8JOBFcBn_OLz",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "47201698-7d33-4379-b772-73e04ebeaa9a"
      },
      "source": [
        "!pip install pyrouge --upgrade\n",
        "!pip install https://github.com/bheinzerling/pyrouge/archive/master.zip\n",
        "!pip install pyrouge\n",
        "!pip show pyrouge\n",
        "!git clone https://github.com/andersjo/pyrouge.git\n",
        "from pyrouge import Rouge155\n",
        "!pyrouge_set_rouge_path 'pyrouge/tools/ROUGE-1.5.5'"
      ],
      "execution_count": 49,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already up-to-date: pyrouge in /usr/local/lib/python3.7/dist-packages (0.1.3)\n",
            "Collecting https://github.com/bheinzerling/pyrouge/archive/master.zip\n",
            "  Using cached https://github.com/bheinzerling/pyrouge/archive/master.zip\n",
            "Requirement already satisfied (use --upgrade to upgrade): pyrouge==0.1.3 from https://github.com/bheinzerling/pyrouge/archive/master.zip in /usr/local/lib/python3.7/dist-packages\n",
            "Building wheels for collected packages: pyrouge\n",
            "  Building wheel for pyrouge (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pyrouge: filename=pyrouge-0.1.3-cp37-none-any.whl size=191921 sha256=7fcdcf8cf8926902e84b4606b2e960097c42b14ae231d20fd4e2e10fc076440e\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-viois7v2/wheels/70/02/b4/a23b5feb5980a5eb940441cb04ec1e17d5f18344138efbecf8\n",
            "Successfully built pyrouge\n",
            "Requirement already satisfied: pyrouge in /usr/local/lib/python3.7/dist-packages (0.1.3)\n",
            "Name: pyrouge\n",
            "Version: 0.1.3\n",
            "Summary: A Python wrapper for the ROUGE summarization evaluation package.\n",
            "Home-page: https://github.com/noutenki/pyrouge\n",
            "Author: Benjamin Heinzerling, Anders Johannsen\n",
            "Author-email: benjamin.heinzerling@h-its.org\n",
            "License: LICENSE.txt\n",
            "Location: /usr/local/lib/python3.7/dist-packages\n",
            "Requires: \n",
            "Required-by: \n",
            "fatal: destination path 'pyrouge' already exists and is not an empty directory.\n",
            "2021-06-15 06:32:45,897 [MainThread  ] [INFO ]  Set ROUGE home directory to pyrouge/tools/ROUGE-1.5.5.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3BxQRUc5KeGG",
        "outputId": "a74c4e2b-0c99-4fe7-8220-87bc17b3c3cb"
      },
      "source": [
        "!pip install transformers\n",
        "!pip install tensorboardX\n",
        "!pip install easydict"
      ],
      "execution_count": 50,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: transformers in /usr/local/lib/python3.7/dist-packages (4.6.1)\n",
            "Requirement already satisfied: sacremoses in /usr/local/lib/python3.7/dist-packages (from transformers) (0.0.45)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.7/dist-packages (from transformers) (2019.12.20)\n",
            "Requirement already satisfied: huggingface-hub==0.0.8 in /usr/local/lib/python3.7/dist-packages (from transformers) (0.0.8)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from transformers) (2.23.0)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from transformers) (20.9)\n",
            "Requirement already satisfied: tokenizers<0.11,>=0.10.1 in /usr/local/lib/python3.7/dist-packages (from transformers) (0.10.3)\n",
            "Requirement already satisfied: importlib-metadata; python_version < \"3.8\" in /usr/local/lib/python3.7/dist-packages (from transformers) (4.5.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.7/dist-packages (from transformers) (3.0.12)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.7/dist-packages (from transformers) (4.41.1)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.7/dist-packages (from transformers) (1.19.5)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.7/dist-packages (from sacremoses->transformers) (1.0.1)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from sacremoses->transformers) (1.15.0)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.7/dist-packages (from sacremoses->transformers) (7.1.2)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->transformers) (3.0.4)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->transformers) (1.24.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->transformers) (2021.5.30)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->transformers) (2.10)\n",
            "Requirement already satisfied: pyparsing>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->transformers) (2.4.7)\n",
            "Requirement already satisfied: typing-extensions>=3.6.4; python_version < \"3.8\" in /usr/local/lib/python3.7/dist-packages (from importlib-metadata; python_version < \"3.8\"->transformers) (3.7.4.3)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata; python_version < \"3.8\"->transformers) (3.4.1)\n",
            "Requirement already satisfied: tensorboardX in /usr/local/lib/python3.7/dist-packages (2.2)\n",
            "Requirement already satisfied: protobuf>=3.8.0 in /usr/local/lib/python3.7/dist-packages (from tensorboardX) (3.12.4)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from tensorboardX) (1.19.5)\n",
            "Requirement already satisfied: six>=1.9 in /usr/local/lib/python3.7/dist-packages (from protobuf>=3.8.0->tensorboardX) (1.15.0)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from protobuf>=3.8.0->tensorboardX) (57.0.0)\n",
            "Requirement already satisfied: easydict in /usr/local/lib/python3.7/dist-packages (1.9)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "amSkThJD-Bxx",
        "outputId": "163038e4-0519-4e3d-aea4-9521dc68818d"
      },
      "source": [
        "%cd /content\n",
        "!git clone https://github.com/HaloKim/KorBertSum.git\n",
        "%cd /content/KorBertSum"
      ],
      "execution_count": 51,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/content\n",
            "fatal: destination path 'KorBertSum' already exists and is not an empty directory.\n",
            "/content/KorBertSum\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m-yEOMhAAPWs"
      },
      "source": [
        "# Code"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WkUneUqUTtZ9"
      },
      "source": [
        "import os\n",
        "\n",
        "os.chdir('/content/KorBertSum/src')"
      ],
      "execution_count": 52,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2ggKByDgUfp8"
      },
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "from models import data_loader, model_builder\n",
        "from models.model_builder import Summarizer\n",
        "from others.logging import logger, init_logger\n",
        "from models.data_loader import load_dataset\n",
        "from transformers import BertConfig, BertTokenizer\n",
        "from tensorboardX import SummaryWriter\n",
        "from models.reporter import ReportMgr\n",
        "from models.stats import Statistics\n",
        "import easydict"
      ],
      "execution_count": 53,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u9qjyO4fND_h"
      },
      "source": [
        "# Preprocessing"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PeRgtmVDLgRy"
      },
      "source": [
        "def _tally_parameters(model):\n",
        "    n_params = sum([p.nelement() for p in model.parameters()])\n",
        "    return n_params\n",
        "\n",
        "def build_trainer(args, device_id, model,\n",
        "                  optim):\n",
        "    \"\"\"\n",
        "    Simplify `Trainer` creation based on user `opt`s*\n",
        "    Args:\n",
        "        opt (:obj:`Namespace`): user options (usually from argument parsing)\n",
        "        model (:obj:`onmt.models.NMTModel`): the model to train\n",
        "        fields (dict): dict of fields\n",
        "        optim (:obj:`onmt.utils.Optimizer`): optimizer used during training\n",
        "        data_type (str): string describing the type of data\n",
        "            e.g. \"text\", \"img\", \"audio\"\n",
        "        model_saver(:obj:`onmt.models.ModelSaverBase`): the utility object\n",
        "            used to save the model\n",
        "    \"\"\"\n",
        "    device = \"cpu\" if args.visible_gpus == '-1' else \"cuda\"\n",
        "\n",
        "\n",
        "    grad_accum_count = args.accum_count\n",
        "    n_gpu = args.world_size\n",
        "\n",
        "    if device_id >= 0:\n",
        "        gpu_rank = int(args.gpu_ranks[device_id])\n",
        "    else:\n",
        "        gpu_rank = 0\n",
        "        n_gpu = 0\n",
        "\n",
        "    print('gpu_rank %d' % gpu_rank)\n",
        "\n",
        "    tensorboard_log_dir = args.model_path\n",
        "\n",
        "    writer = SummaryWriter(tensorboard_log_dir, comment=\"Unmt\")\n",
        "\n",
        "    report_manager = ReportMgr(args.report_every, start_time=-1, tensorboard_writer=writer)\n",
        "\n",
        "    trainer = Trainer(args, model, optim, grad_accum_count, n_gpu, gpu_rank, report_manager)\n",
        "\n",
        "    # print(tr)\n",
        "    if (model):\n",
        "        n_params = _tally_parameters(model)\n",
        "        logger.info('* number of parameters: %d' % n_params)\n",
        "\n",
        "    return trainer\n",
        "\n",
        "\n",
        "class Trainer(object):\n",
        "    \"\"\"\n",
        "    Class that controls the training process.\n",
        "\n",
        "    Args:\n",
        "            model(:py:class:`onmt.models.model.NMTModel`): translation model\n",
        "                to train\n",
        "            train_loss(:obj:`onmt.utils.loss.LossComputeBase`):\n",
        "               training loss computation\n",
        "            valid_loss(:obj:`onmt.utils.loss.LossComputeBase`):\n",
        "               training loss computation\n",
        "            optim(:obj:`onmt.utils.optimizers.Optimizer`):\n",
        "               the optimizer responsible for update\n",
        "            trunc_size(int): length of truncated back propagation through time\n",
        "            shard_size(int): compute loss in shards of this size for efficiency\n",
        "            data_type(string): type of the source input: [text|img|audio]\n",
        "            norm_method(string): normalization methods: [sents|tokens]\n",
        "            grad_accum_count(int): accumulate gradients this many times.\n",
        "            report_manager(:obj:`onmt.utils.ReportMgrBase`):\n",
        "                the object that creates reports, or None\n",
        "            model_saver(:obj:`onmt.models.ModelSaverBase`): the saver is\n",
        "                used to save a checkpoint.\n",
        "                Thus nothing will be saved if this parameter is None\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self,  args, model,  optim,\n",
        "                  grad_accum_count=1, n_gpu=1, gpu_rank=1,\n",
        "                  report_manager=None):\n",
        "        # Basic attributes.\n",
        "        self.args = args\n",
        "        self.save_checkpoint_steps = args.save_checkpoint_steps\n",
        "        self.model = model\n",
        "        self.optim = optim\n",
        "        self.grad_accum_count = grad_accum_count\n",
        "        self.n_gpu = n_gpu\n",
        "        self.gpu_rank = gpu_rank\n",
        "        self.report_manager = report_manager\n",
        "\n",
        "        self.loss = torch.nn.BCELoss(reduction='none')\n",
        "        assert grad_accum_count > 0\n",
        "        # Set model in training mode.\n",
        "        if (model):\n",
        "            self.model.train()\n",
        "\n",
        "    def summ(self, test_iter, step, cal_lead=False, cal_oracle=False):\n",
        "      \"\"\" Validate model.\n",
        "          valid_iter: validate data iterator\n",
        "      Returns:\n",
        "          :obj:`nmt.Statistics`: validation loss statistics\n",
        "      \"\"\"\n",
        "      # Set model in validating mode.\n",
        "      def _get_ngrams(n, text):\n",
        "          ngram_set = set()\n",
        "          text_length = len(text)\n",
        "          max_index_ngram_start = text_length - n\n",
        "          for i in range(max_index_ngram_start + 1):\n",
        "              ngram_set.add(tuple(text[i:i + n]))\n",
        "          return ngram_set\n",
        "\n",
        "      def _block_tri(c, p):\n",
        "          tri_c = _get_ngrams(3, c.split())\n",
        "          for s in p:\n",
        "              tri_s = _get_ngrams(3, s.split())\n",
        "              if len(tri_c.intersection(tri_s))>0:\n",
        "                  return True\n",
        "          return False\n",
        "\n",
        "      if (not cal_lead and not cal_oracle):\n",
        "          self.model.eval()\n",
        "      stats = Statistics()\n",
        "\n",
        "      with torch.no_grad():\n",
        "          for batch in test_iter:\n",
        "              src = batch.src\n",
        "              labels = batch.labels\n",
        "              segs = batch.segs\n",
        "              clss = batch.clss\n",
        "              mask = batch.mask\n",
        "              mask_cls = batch.mask_cls\n",
        "\n",
        "              if (cal_lead):\n",
        "                  selected_ids = [list(range(batch.clss.size(1)))] * batch.batch_size\n",
        "              elif (cal_oracle):\n",
        "                  selected_ids = [[j for j in range(batch.clss.size(1)) if labels[i][j] == 1] for i in\n",
        "                                  range(batch.batch_size)]\n",
        "              else:\n",
        "                  sent_scores, mask = self.model(src, segs, clss, mask, mask_cls)\n",
        "                  sent_scores = sent_scores + mask.float()\n",
        "                  sent_scores = sent_scores.cpu().data.numpy()\n",
        "                  selected_ids = np.argsort(-sent_scores, 1)\n",
        "      return selected_ids\n",
        "\n",
        "\n",
        "\n",
        "    def _gradient_accumulation(self, true_batchs, normalization, total_stats,\n",
        "                               report_stats):\n",
        "        if self.grad_accum_count > 1:\n",
        "            self.model.zero_grad()\n",
        "\n",
        "        for batch in true_batchs:\n",
        "            if self.grad_accum_count == 1:\n",
        "                self.model.zero_grad()\n",
        "\n",
        "            src = batch.src\n",
        "            labels = batch.labels\n",
        "            segs = batch.segs\n",
        "            clss = batch.clss\n",
        "            mask = batch.mask\n",
        "            mask_cls = batch.mask_cls\n",
        "\n",
        "            sent_scores, mask = self.model(src, segs, clss, mask, mask_cls)\n",
        "\n",
        "            loss = self.loss(sent_scores, labels.float())\n",
        "            loss = (loss*mask.float()).sum()\n",
        "            (loss/loss.numel()).backward()\n",
        "            # loss.div(float(normalization)).backward()\n",
        "\n",
        "            batch_stats = Statistics(float(loss.cpu().data.numpy()), normalization)\n",
        "\n",
        "\n",
        "            total_stats.update(batch_stats)\n",
        "            report_stats.update(batch_stats)\n",
        "\n",
        "            # 4. Update the parameters and statistics.\n",
        "            if self.grad_accum_count == 1:\n",
        "                # Multi GPU gradient gather\n",
        "                if self.n_gpu > 1:\n",
        "                    grads = [p.grad.data for p in self.model.parameters()\n",
        "                             if p.requires_grad\n",
        "                             and p.grad is not None]\n",
        "                    distributed.all_reduce_and_rescale_tensors(\n",
        "                        grads, float(1))\n",
        "                self.optim.step()\n",
        "\n",
        "        # in case of multi step gradient accumulation,\n",
        "        # update only after accum batches\n",
        "        if self.grad_accum_count > 1:\n",
        "            if self.n_gpu > 1:\n",
        "                grads = [p.grad.data for p in self.model.parameters()\n",
        "                         if p.requires_grad\n",
        "                         and p.grad is not None]\n",
        "                distributed.all_reduce_and_rescale_tensors(\n",
        "                    grads, float(1))\n",
        "            self.optim.step()\n",
        "\n",
        "    def _save(self, step):\n",
        "        real_model = self.model\n",
        "        # real_generator = (self.generator.module\n",
        "        #                   if isinstance(self.generator, torch.nn.DataParallel)\n",
        "        #                   else self.generator)\n",
        "\n",
        "        model_state_dict = real_model.state_dict()\n",
        "        # generator_state_dict = real_generator.state_dict()\n",
        "        checkpoint = {\n",
        "            'model': model_state_dict,\n",
        "            # 'generator': generator_state_dict,\n",
        "            'opt': self.args,\n",
        "            'optim': self.optim,\n",
        "        }\n",
        "        checkpoint_path = os.path.join(self.args.model_path, 'model_step_%d.pt' % step)\n",
        "        logger.info(\"Saving checkpoint %s\" % checkpoint_path)\n",
        "        # checkpoint_path = '%s_step_%d.pt' % (FLAGS.model_path, step)\n",
        "        if (not os.path.exists(checkpoint_path)):\n",
        "            torch.save(checkpoint, checkpoint_path)\n",
        "            return checkpoint, checkpoint_path\n",
        "\n",
        "    def _start_report_manager(self, start_time=None):\n",
        "        \"\"\"\n",
        "        Simple function to start report manager (if any)\n",
        "        \"\"\"\n",
        "        if self.report_manager is not None:\n",
        "            if start_time is None:\n",
        "                self.report_manager.start()\n",
        "            else:\n",
        "                self.report_manager.start_time = start_time\n",
        "\n",
        "    def _maybe_gather_stats(self, stat):\n",
        "        \"\"\"\n",
        "        Gather statistics in multi-processes cases\n",
        "\n",
        "        Args:\n",
        "            stat(:obj:onmt.utils.Statistics): a Statistics object to gather\n",
        "                or None (it returns None in this case)\n",
        "\n",
        "        Returns:\n",
        "            stat: the updated (or unchanged) stat object\n",
        "        \"\"\"\n",
        "        if stat is not None and self.n_gpu > 1:\n",
        "            return Statistics.all_gather_stats(stat)\n",
        "        return stat\n",
        "\n",
        "    def _maybe_report_training(self, step, num_steps, learning_rate,\n",
        "                               report_stats):\n",
        "        \"\"\"\n",
        "        Simple function to report training stats (if report_manager is set)\n",
        "        see `onmt.utils.ReportManagerBase.report_training` for doc\n",
        "        \"\"\"\n",
        "        if self.report_manager is not None:\n",
        "            return self.report_manager.report_training(\n",
        "                step, num_steps, learning_rate, report_stats,\n",
        "                multigpu=self.n_gpu > 1)\n",
        "\n",
        "    def _report_step(self, learning_rate, step, train_stats=None,\n",
        "                     valid_stats=None):\n",
        "        \"\"\"\n",
        "        Simple function to report stats (if report_manager is set)\n",
        "        see `onmt.utils.ReportManagerBase.report_step` for doc\n",
        "        \"\"\"\n",
        "        if self.report_manager is not None:\n",
        "            return self.report_manager.report_step(\n",
        "                learning_rate, step, train_stats=train_stats,\n",
        "                valid_stats=valid_stats)\n",
        "\n",
        "    def _maybe_save(self, step):\n",
        "        \"\"\"\n",
        "        Save the model if a model saver is set\n",
        "        \"\"\"\n",
        "        if self.model_saver is not None:\n",
        "            self.model_saver.maybe_save(step)\n",
        "\n",
        "class BertData():\n",
        "    def __init__(self):\n",
        "        self.tokenizer = BertTokenizer.from_pretrained('bert-base-multilingual-cased')\n",
        "        self.sep_vid = self.tokenizer.vocab['[SEP]']\n",
        "        self.cls_vid = self.tokenizer.vocab['[CLS]']\n",
        "        self.pad_vid = self.tokenizer.vocab['[PAD]']\n",
        "\n",
        "    def preprocess(self, src):\n",
        "\n",
        "        if (len(src) == 0):\n",
        "            return None\n",
        "\n",
        "        original_src_txt = [' '.join(s) for s in src]\n",
        "        idxs = [i for i, s in enumerate(src) if (len(s) > 1)]\n",
        "\n",
        "        src = [src[i][:2000] for i in idxs]\n",
        "        src = src[:1000]\n",
        "\n",
        "        if (len(src) < 3):\n",
        "            return None\n",
        "\n",
        "        src_txt = [' '.join(sent) for sent in src]\n",
        "        text = ' [SEP] [CLS] '.join(src_txt)\n",
        "        src_subtokens = self.tokenizer.tokenize(text)\n",
        "        src_subtokens = src_subtokens[:510]\n",
        "        src_subtokens = ['[CLS]'] + src_subtokens + ['[SEP]']\n",
        "\n",
        "        src_subtoken_idxs = self.tokenizer.convert_tokens_to_ids(src_subtokens)\n",
        "        _segs = [-1] + [i for i, t in enumerate(src_subtoken_idxs) if t == self.sep_vid]\n",
        "        segs = [_segs[i] - _segs[i - 1] for i in range(1, len(_segs))]\n",
        "        segments_ids = []\n",
        "        for i, s in enumerate(segs):\n",
        "            if (i % 2 == 0):\n",
        "                segments_ids += s * [0]\n",
        "            else:\n",
        "                segments_ids += s * [1]\n",
        "        cls_ids = [i for i, t in enumerate(src_subtoken_idxs) if t == self.cls_vid]\n",
        "        labels = None\n",
        "        src_txt = [original_src_txt[i] for i in idxs]\n",
        "        tgt_txt = None\n",
        "        return src_subtoken_idxs, labels, segments_ids, cls_ids, src_txt, tgt_txt\n",
        "\n",
        "def _lazy_dataset_loader(pt_file):\n",
        "  yield  pt_file"
      ],
      "execution_count": 54,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pC_QmVS6NGvy"
      },
      "source": [
        "# Params"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6tzZYPbLbRl7"
      },
      "source": [
        "args = easydict.EasyDict({\n",
        "    \"encoder\":'classifier',\n",
        "    \"mode\":'test',\n",
        "    \"bert_data_path\":'/content/drive/MyDrive/Colab_Notebooks/Study/BertSum-master/bert_data/korean',\n",
        "    \"model_path\":'./models/bert_classifier',\n",
        "    \"result_path\":'./results',\n",
        "    \"temp_dir\":'./temp',\n",
        "    \"batch_size\":1000,\n",
        "    \"use_interval\":True,\n",
        "    \"hidden_size\":128,\n",
        "    \"ff_size\":512,\n",
        "    \"heads\":4,\n",
        "    \"inter_layers\":2,\n",
        "    \"rnn_size\":512,\n",
        "    \"param_init\":0,\n",
        "    \"param_init_glorot\":True,\n",
        "    \"dropout\":0.1,\n",
        "    \"optim\":'adam',\n",
        "    \"lr\":2e-3,\n",
        "    \"report_every\":1,\n",
        "    \"save_checkpoint_steps\":5,\n",
        "    \"block_trigram\":True,\n",
        "    \"recall_eval\":False,\n",
        "    \n",
        "    \"accum_count\":1,\n",
        "    \"world_size\":1,\n",
        "    \"visible_gpus\":'-1',\n",
        "    \"gpu_ranks\":'0',\n",
        "    \"log_file\":'/content/drive/MyDrive/Colab_Notebooks/Study/BertSum-master/logs/log.log',\n",
        "    \"test_from\":'/content/drive/MyDrive/Colab_Notebooks/Study/BertSum-master/models/bert_classifier/model_step_10000.pt'\n",
        "})\n",
        "model_flags = ['hidden_size', 'ff_size', 'heads', 'inter_layers','encoder','ff_actv', 'use_interval','rnn_size']\n"
      ],
      "execution_count": 61,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aJ1pvUQBNITZ"
      },
      "source": [
        "# Test code"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WYV7oXExF2je"
      },
      "source": [
        "def test(args, input_list, device_id, pt, step):\n",
        "  init_logger(args.log_file)\n",
        "  device = \"cpu\" if args.visible_gpus == '-1' else \"cuda\"\n",
        "  device_id = 0 if device == \"cuda\" else -1\n",
        "\n",
        "  cp = args.test_from\n",
        "  try:\n",
        "    step = int(cp.split('.')[-2].split('_')[-1])\n",
        "  except:\n",
        "    step = 0\n",
        "\n",
        "  device = \"cpu\" if args.visible_gpus == '-1' else \"cuda\"\n",
        "  if (pt != ''):\n",
        "      test_from = pt\n",
        "  else:\n",
        "      test_from = args.test_from\n",
        "  logger.info('Loading checkpoint from %s' % test_from)\n",
        "  checkpoint = torch.load(test_from, map_location=lambda storage, loc: storage)\n",
        "  opt = vars(checkpoint['opt'])\n",
        "  for k in opt.keys():\n",
        "      if (k in model_flags):\n",
        "        setattr(args, k, opt[k])\n",
        "\n",
        "  config = BertConfig.from_pretrained('bert-base-multilingual-cased')\n",
        "  model = Summarizer(args, device, load_pretrained_bert=False, bert_config = config)\n",
        "  model.load_cp(checkpoint)\n",
        "  model.eval()\n",
        "\n",
        "  test_iter = data_loader.Dataloader(args, _lazy_dataset_loader(input_list),\n",
        "                                args.batch_size, device,\n",
        "                                shuffle=False, is_test=True)\n",
        "  trainer = build_trainer(args, device_id, model, None)\n",
        "  result = trainer.summ(test_iter,step)\n",
        "  return result, input_list\n",
        "\n",
        "args.gpu_ranks = [int(i) for i in args.gpu_ranks.split(',')]\n",
        "os.environ[\"CUDA_VISIBLE_DEVICES\"] = args.visible_gpus"
      ],
      "execution_count": 62,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AywahUrUkPzB"
      },
      "source": [
        "def txt2input(text):\n",
        "  data = list(filter(None, text.split('\\n')))\n",
        "  bertdata = BertData()\n",
        "  txt_data = bertdata.preprocess(data)\n",
        "  data_dict = {\"src\":txt_data[0],\n",
        "               \"labels\":[0,1,2],\n",
        "               \"segs\":txt_data[2],\n",
        "               \"clss\":txt_data[3],\n",
        "               \"src_txt\":txt_data[4],\n",
        "               \"tgt_txt\":None}\n",
        "  input_data = []\n",
        "  input_data.append(data_dict)\n",
        "  return input_data"
      ],
      "execution_count": 63,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uXZk2tZvePTK"
      },
      "source": [
        "text = '''\n",
        "일본 정부가 문재인 대통령이 도쿄올림픽을 계기로 일본을 방문하는 방향으로 한일 양국이 조율하고 있다는 15일 요미우리신문의 보도를 부인했다.\n",
        "\n",
        "정부 대변인인 가토 가쓰노부(加藤勝信) 관방장관은 이날 오전 정례 기자회견에서 관련 질문에 “말씀하신 보도와 같은 사실이 없는 것으로 안다”고 밝혔다.\n",
        "\n",
        "앞서 요미우리는 한국 측이 도쿄올림픽을 계기로 한 문 대통령의 방일을 타진했고, 일본 측은 수용하는 방향이라고 이날 보도했다.\n",
        "\n",
        "한국 측은 문 대통령의 방일 때 스가 요시히데(菅義偉) 총리와 처음으로 정상회담을 하겠다는 생각이라고 요미우리는 전했다.\n",
        "\n",
        "가토 장관은 한일 정상회담에 대한 일본 정부의 자세에 대해 “그런 사실이 없기 때문에 가정의 질문에 대해 답하는 것을 삼가겠다”고 말했다.\n",
        "\n",
        "그는 한국 측의 독도방어훈련에 ‘어떤 대항 조치를 생각하고 있느냐’는 질문에는 “한국 해군의 훈련에 대해 정부로서는 강한 관심을 가지고 주시하는 상황이어서 지금 시점에선 논평을 삼가겠다”고 말을 아꼈다.\n",
        "\n",
        "가토 장관은 “다케시마(竹島·일본이 주장하는 독도의 명칭)는 역사적 사실에 비춰봐도, 국제법상으로도 명백한 일본 고유의 영토”라며 독도 영유권 주장을 되풀이했다.\n",
        "\n",
        "그러면서 “다케시마 문제에 대해서는 계속 우리나라의 영토, 영해, 영공을 단호히 지키겠다는 결의로 냉정하고 의연하게 대응해갈 생각”이라고 밝혔다.\n",
        "'''"
      ],
      "execution_count": 98,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7IlXBocEMzS_",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "9749408a-322e-4744-8ebb-b0f031c83184"
      },
      "source": [
        "input_data = txt2input(text)\n",
        "sum_list = test(args, input_data, -1, '', None)\n",
        "sum_list[0]"
      ],
      "execution_count": 99,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[2021-06-15 06:40:50,285 INFO] Loading checkpoint from /content/drive/MyDrive/Colab_Notebooks/Study/BertSum-master/models/bert_classifier/model_step_10000.pt\n",
            "[2021-06-15 06:40:56,568 INFO] * number of parameters: 177854209\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "gpu_rank 0\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[0, 1, 3, 4, 6, 7, 5, 2]])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 99
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3uOvhjs8M1-s"
      },
      "source": [
        "# Result"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nIhOEG0HlMpE",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "47f54014-9d72-4a10-a244-4441e321bfcb"
      },
      "source": [
        "[list(filter(None, text.split('\\n')))[i] for i in sum_list[0][0][:3]]"
      ],
      "execution_count": 100,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['일본 정부가 문재인 대통령이 도쿄올림픽을 계기로 일본을 방문하는 방향으로 한일 양국이 조율하고 있다는 15일 요미우리신문의 보도를 부인했다.',\n",
              " '정부 대변인인 가토 가쓰노부(加藤勝信) 관방장관은 이날 오전 정례 기자회견에서 관련 질문에 “말씀하신 보도와 같은 사실이 없는 것으로 안다”고 밝혔다.',\n",
              " '한국 측은 문 대통령의 방일 때 스가 요시히데(菅義偉) 총리와 처음으로 정상회담을 하겠다는 생각이라고 요미우리는 전했다.']"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 100
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CSVtNDS2Jiaa"
      },
      "source": [
        " "
      ],
      "execution_count": 100,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QLnyCkviYuIv"
      },
      "source": [
        ""
      ],
      "execution_count": 100,
      "outputs": []
    }
  ]
}