BERT_SUMM.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/santhoshkolloju/bert_summ/blob/master/BERT_SUMM.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "pbgEu1oyhPca",
    "outputId": "1ecf3d0e-75c4-4f0d-8a64-18fc6bcfb4a1"
   },
   "outputs": [
    {
     "ename": "SystemError",
     "evalue": "GPU device not found",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mSystemError\u001b[0m                               Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-2-d1680108c58e>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0mdevice_name\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtest\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgpu_device_name\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mdevice_name\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0;34m'/device:GPU:0'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m   \u001b[0;32mraise\u001b[0m \u001b[0mSystemError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'GPU device not found'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      5\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'Found GPU at: {}'\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice_name\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mSystemError\u001b[0m: GPU device not found"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "device_name = tf.test.gpu_device_name()\n",
    "if device_name != '/device:GPU:0':\n",
    "  raise SystemError('GPU device not found')\n",
    "print('Found GPU at: {}'.format(device_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "UqeTD_eGkxVd"
   },
   "outputs": [],
   "source": [
    "from google.colab import auth\n",
    "auth.authenticate_user()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 267
    },
    "colab_type": "code",
    "id": "SFCZikuuibh7",
    "outputId": "f228c486-ff0b-4740-c9c2-10a2d9b893ba"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cloning into 'texar_repo'...\n",
      "remote: Enumerating objects: 10204, done.\u001b[K\n",
      "remote: Total 10204 (delta 0), reused 0 (delta 0), pack-reused 10204\u001b[K\n",
      "Receiving objects: 100% (10204/10204), 2.30 MiB | 12.21 MiB/s, done.\n",
      "Resolving deltas: 100% (7774/7774), done.\n",
      "Collecting funcsigs\n",
      "  Downloading https://files.pythonhosted.org/packages/69/cb/f5be453359271714c01b9bd06126eaf2e368f1fddfff30818754b5ac2328/funcsigs-1.0.2-py2.py3-none-any.whl\n",
      "Installing collected packages: funcsigs\n",
      "Successfully installed funcsigs-1.0.2\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "\n",
    "!test -d texar_repo || git clone https://github.com/asyml/texar.git texar_repo\n",
    "if not 'texar_repo' in sys.path:\n",
    "  sys.path += ['texar_repo']\n",
    "!pip install funcsigs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 74
    },
    "colab_type": "code",
    "id": "_IB17muzsBd0",
    "outputId": "34a14a15-84b6-47d0-d286-2ae052c858d0"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nfrom google_drive_downloader import GoogleDriveDownloader as gdd\\n\\ngdd.download_file_from_google_drive(file_id=\\'0BwmD_VLjROrfTHk4NFg2SndKcjQ\\',\\n                                    dest_path=\\'./cnn_stories.tgz\\',\\n                                    unzip=True)\\n!tar -zxf cnn_stories.tgz\\n\\nfrom os import listdir\\nimport string\\n\\n# load doc into memory\\ndef load_doc(filename):\\n\\t# open the file as read only\\n\\tfile = open(filename, encoding=\\'utf-8\\')\\n\\t# read all text\\n\\ttext = file.read()\\n\\t# close the file\\n\\tfile.close()\\n\\treturn text\\n\\n# split a document into news story and highlights\\ndef split_story(doc):\\n\\t# find first highlight\\n\\tindex = doc.find(\\'@highlight\\')\\n\\t# split into story and highlights\\n\\tstory, highlights = doc[:index], doc[index:].split(\\'@highlight\\')\\n\\t# strip extra white space around each highlight\\n\\thighlights = [h.strip() for h in highlights if len(h) > 0]\\n\\treturn story, highlights\\n\\n# load all stories in a directory\\ndef load_stories(directory):\\n\\tstories = list()\\n\\tfor name in listdir(directory):\\n\\t\\tfilename = directory + \\'/\\' + name\\n\\t\\t# load document\\n\\t\\tdoc = load_doc(filename)\\n\\t\\t# split into story and highlights\\n\\t\\tstory, highlights = split_story(doc)\\n\\t\\t# store\\n\\t\\tstories.append({\\'story\\':story, \\'highlights\\':highlights})\\n\\treturn stories\\n\\n# clean a list of lines\\ndef clean_lines(lines):\\n\\tcleaned = list()\\n\\t# prepare a translation table to remove punctuation\\n\\ttable = str.maketrans(\\'\\', \\'\\', string.punctuation)\\n\\tfor line in lines:\\n\\t\\t# strip source cnn office if it exists\\n\\t\\tindex = line.find(\\'(CNN) -- \\')\\n\\t\\tif index > -1:\\n\\t\\t\\tline = line[index+len(\\'(CNN)\\'):]\\n\\t\\t# tokenize on white space\\n\\t\\tline = line.split()\\n\\t\\t# convert to lower case\\n\\t\\tline = [word.lower() for word in line]\\n\\t\\t# remove punctuation from each token\\n\\t\\tline = [w.translate(table) for w in line]\\n\\t\\t# remove tokens with numbers in them\\n\\t\\tline = [word for word in line if word.isalpha()]\\n\\t\\t# store as string\\n\\t\\tcleaned.append(\\' \\'.join(line))\\n\\t# remove empty strings\\n\\tcleaned = [c for c in cleaned if len(c) > 0]\\n\\treturn cleaned\\n\\n# load stories\\ndirectory = \\'cnn/stories/\\'\\nstories = load_stories(directory)\\nprint(\\'Loaded Stories %d\\' % len(stories))\\n\\n# clean stories\\nf1 = open(\"stories.txt\",\\'w\\')\\nf2 = open(\"summary.txt\",\\'w\\')\\nfor example in stories:\\n  example[\\'story\\'] = clean_lines(example[\\'story\\'].split(\\'\\n\\'))\\n  example[\\'highlights\\'] = clean_lines(example[\\'highlights\\'])\\n  f1.write(\" \".join(example[\\'story\\']))\\n  f1.write(\"\\n\")\\n  f2.write(\" \".join(example[\\'highlights\\']))\\n  f2.write(\"\\n\")\\nf1.close()\\nf2.close()\\n  \\nstory = open(\"stories.txt\").readlines()\\nsumm = open(\"summary.txt\").readlines() \\ntrain_story = story[0:90000]\\ntrain_summ = summ[0:90000]\\n\\neval_story = story[90000:91579]\\neval_summ = summ[90000:91579]\\n\\n\\ntest_story = story[91579:92579]\\ntest_summ = summ[91579:92579]\\n\\n\\nwith open(\"train_story.txt\",\\'w\\') as f:\\n  f.write(\"\\n\".join(train_story))\\n  \\nwith open(\"train_summ.txt\",\\'w\\') as f:\\n  f.write(\"\\n\".join(train_summ))\\n  \\nwith open(\"eval_story.txt\",\\'w\\') as f:\\n  f.write(\"\\n\".join(eval_story))\\n  \\n  \\nwith open(\"eval_summ.txt\",\\'w\\') as f:\\n  f.write(\"\\n\".join(eval_summ))\\n  \\n  \\nwith open(\"test_story.txt\",\\'w\\') as f:\\n  f.write(\"\\n\".join(test_story))\\n  \\n  \\nwith open(\"test_summ.txt\",\\'w\\') as f:\\n  f.write(\"\\n\".join(test_summ))  \\n  '"
      ]
     },
     "execution_count": 4,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "from google_drive_downloader import GoogleDriveDownloader as gdd\n",
    "\n",
    "gdd.download_file_from_google_drive(file_id='0BwmD_VLjROrfTHk4NFg2SndKcjQ',\n",
    "                                    dest_path='./cnn_stories.tgz',\n",
    "                                    unzip=True)\n",
    "!tar -zxf cnn_stories.tgz\n",
    "\n",
    "from os import listdir\n",
    "import string\n",
    "\n",
    "# load doc into memory\n",
    "def load_doc(filename):\n",
    "\t# open the file as read only\n",
    "\tfile = open(filename, encoding='utf-8')\n",
    "\t# read all text\n",
    "\ttext = file.read()\n",
    "\t# close the file\n",
    "\tfile.close()\n",
    "\treturn text\n",
    "\n",
    "# split a document into news story and highlights\n",
    "def split_story(doc):\n",
    "\t# find first highlight\n",
    "\tindex = doc.find('@highlight')\n",
    "\t# split into story and highlights\n",
    "\tstory, highlights = doc[:index], doc[index:].split('@highlight')\n",
    "\t# strip extra white space around each highlight\n",
    "\thighlights = [h.strip() for h in highlights if len(h) > 0]\n",
    "\treturn story, highlights\n",
    "\n",
    "# load all stories in a directory\n",
    "def load_stories(directory):\n",
    "\tstories = list()\n",
    "\tfor name in listdir(directory):\n",
    "\t\tfilename = directory + '/' + name\n",
    "\t\t# load document\n",
    "\t\tdoc = load_doc(filename)\n",
    "\t\t# split into story and highlights\n",
    "\t\tstory, highlights = split_story(doc)\n",
    "\t\t# store\n",
    "\t\tstories.append({'story':story, 'highlights':highlights})\n",
    "\treturn stories\n",
    "\n",
    "# clean a list of lines\n",
    "def clean_lines(lines):\n",
    "\tcleaned = list()\n",
    "\t# prepare a translation table to remove punctuation\n",
    "\ttable = str.maketrans('', '', string.punctuation)\n",
    "\tfor line in lines:\n",
    "\t\t# strip source cnn office if it exists\n",
    "\t\tindex = line.find('(CNN) -- ')\n",
    "\t\tif index > -1:\n",
    "\t\t\tline = line[index+len('(CNN)'):]\n",
    "\t\t# tokenize on white space\n",
    "\t\tline = line.split()\n",
    "\t\t# convert to lower case\n",
    "\t\tline = [word.lower() for word in line]\n",
    "\t\t# remove punctuation from each token\n",
    "\t\tline = [w.translate(table) for w in line]\n",
    "\t\t# remove tokens with numbers in them\n",
    "\t\tline = [word for word in line if word.isalpha()]\n",
    "\t\t# store as string\n",
    "\t\tcleaned.append(' '.join(line))\n",
    "\t# remove empty strings\n",
    "\tcleaned = [c for c in cleaned if len(c) > 0]\n",
    "\treturn cleaned\n",
    "\n",
    "# load stories\n",
    "directory = 'cnn/stories/'\n",
    "stories = load_stories(directory)\n",
    "print('Loaded Stories %d' % len(stories))\n",
    "\n",
    "# clean stories\n",
    "f1 = open(\"stories.txt\",'w')\n",
    "f2 = open(\"summary.txt\",'w')\n",
    "for example in stories:\n",
    "  example['story'] = clean_lines(example['story'].split('\\n'))\n",
    "  example['highlights'] = clean_lines(example['highlights'])\n",
    "  f1.write(\" \".join(example['story']))\n",
    "  f1.write(\"\\n\")\n",
    "  f2.write(\" \".join(example['highlights']))\n",
    "  f2.write(\"\\n\")\n",
    "f1.close()\n",
    "f2.close()\n",
    "  \n",
    "story = open(\"stories.txt\").readlines()\n",
    "summ = open(\"summary.txt\").readlines() \n",
    "train_story = story[0:90000]\n",
    "train_summ = summ[0:90000]\n",
    "\n",
    "eval_story = story[90000:91579]\n",
    "eval_summ = summ[90000:91579]\n",
    "\n",
    "\n",
    "test_story = story[91579:92579]\n",
    "test_summ = summ[91579:92579]\n",
    "\n",
    "\n",
    "with open(\"train_story.txt\",'w') as f:\n",
    "  f.write(\"\\n\".join(train_story))\n",
    "  \n",
    "with open(\"train_summ.txt\",'w') as f:\n",
    "  f.write(\"\\n\".join(train_summ))\n",
    "  \n",
    "with open(\"eval_story.txt\",'w') as f:\n",
    "  f.write(\"\\n\".join(eval_story))\n",
    "  \n",
    "  \n",
    "with open(\"eval_summ.txt\",'w') as f:\n",
    "  f.write(\"\\n\".join(eval_summ))\n",
    "  \n",
    "  \n",
    "with open(\"test_story.txt\",'w') as f:\n",
    "  f.write(\"\\n\".join(test_story))\n",
    "  \n",
    "  \n",
    "with open(\"test_summ.txt\",'w') as f:\n",
    "  f.write(\"\\n\".join(test_summ))  \n",
    "  \"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "SsWJmIfmij-_"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import csv\n",
    "import collections\n",
    "import sys\n",
    "from texar_repo.examples.bert.utils import data_utils, model_utils, tokenization\n",
    "import importlib\n",
    "import tensorflow as tf\n",
    "import texar as tx \n",
    "from texar_repo.examples.bert import config_classifier as config_downstream\n",
    "from texar_repo.texar.utils import transformer_utils\n",
    "from texar_repo.examples.transformer.utils import data_utils, utils\n",
    "from texar_repo.examples.transformer.bleu_tool import bleu_wrapper\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Ko2McfcdhbcN"
   },
   "outputs": [],
   "source": [
    "#config\n",
    "\n",
    "dcoder_config = {\n",
    "    'dim': 768,\n",
    "    'num_blocks': 6,\n",
    "    'multihead_attention': {\n",
    "        'num_heads': 8,\n",
    "        'output_dim': 768\n",
    "        # See documentation for more optional hyperparameters\n",
    "    },\n",
    "    'position_embedder_hparams': {\n",
    "        'dim': 768\n",
    "    },\n",
    "    'initializer': {\n",
    "        'type': 'variance_scaling_initializer',\n",
    "        'kwargs': {\n",
    "            'scale': 1.0,\n",
    "            'mode': 'fan_avg',\n",
    "            'distribution': 'uniform',\n",
    "        },\n",
    "    },\n",
    "    'poswise_feedforward': tx.modules.default_transformer_poswise_net_hparams(\n",
    "        output_dim=768)\n",
    "}\n",
    "\n",
    "loss_label_confidence = 0.9\n",
    "\n",
    "random_seed = 1234\n",
    "beam_width = 5\n",
    "alpha = 0.6\n",
    "hidden_dim = 768\n",
    "\n",
    "\n",
    "opt = {\n",
    "    'optimizer': {\n",
    "        'type': 'AdamOptimizer',\n",
    "        'kwargs': {\n",
    "            'beta1': 0.9,\n",
    "            'beta2': 0.997,\n",
    "            'epsilon': 1e-9\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "\n",
    "lr = {\n",
    "    'learning_rate_schedule': 'constant.linear_warmup.rsqrt_decay.rsqrt_depth',\n",
    "    'lr_constant': 2 * (hidden_dim ** -0.5),\n",
    "    'static_lr': 1e-3,\n",
    "    'warmup_steps': 2000,\n",
    "}\n",
    "\n",
    "bos_token_id =101\n",
    "eos_token_id = 102\n",
    "\n",
    "model_dir= \"./models\"\n",
    "run_mode= \"train_and_evaluate\"\n",
    "batch_size = 32\n",
    "test_batch_size = 32\n",
    "\n",
    "max_train_epoch = 20\n",
    "display_steps = 100\n",
    "eval_steps = 100000\n",
    "\n",
    "max_decoding_length = 400\n",
    "\n",
    "max_seq_length_src = 512\n",
    "max_seq_length_tgt = 400\n",
    "\n",
    "bert_pretrain_dir = 'bert_pretrained_models/uncased_L-12_H-768_A-12'\n",
    "#config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MrBw61rEiXeE"
   },
   "outputs": [],
   "source": [
    "!mkdir bert_pretrained_models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 354
    },
    "colab_type": "code",
    "id": "FH2RX773i40g",
    "outputId": "3e61c01b-e31b-485b-89e0-b51fddc341d4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2019-10-13 19:04:58--  https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-12_H-768_A-12.zip\n",
      "Resolving storage.googleapis.com (storage.googleapis.com)... 2404:6800:4005:80e::2010\n",
      "Connecting to storage.googleapis.com (storage.googleapis.com)|2404:6800:4005:80e::2010|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 407727028 (389M) [application/zip]\n",
      "Saving to: ‘bert_pretrained_models/uncased_L-12_H-768_A-12.zip’\n",
      "\n",
      "100%[======================================>] 407,727,028 12.4MB/s   in 32s    \n",
      "\n",
      "2019-10-13 19:05:31 (12.3 MB/s) - ‘bert_pretrained_models/uncased_L-12_H-768_A-12.zip’ saved [407727028/407727028]\n",
      "\n",
      "Archive:  bert_pretrained_models/uncased_L-12_H-768_A-12.zip\n",
      "   creating: bert_pretrained_models/uncased_L-12_H-768_A-12/\n",
      "  inflating: bert_pretrained_models/uncased_L-12_H-768_A-12/bert_model.ckpt.meta  \n",
      "  inflating: bert_pretrained_models/uncased_L-12_H-768_A-12/bert_model.ckpt.data-00000-of-00001  \n",
      "  inflating: bert_pretrained_models/uncased_L-12_H-768_A-12/vocab.txt  \n",
      "  inflating: bert_pretrained_models/uncased_L-12_H-768_A-12/bert_model.ckpt.index  \n",
      "  inflating: bert_pretrained_models/uncased_L-12_H-768_A-12/bert_config.json  \n"
     ]
    }
   ],
   "source": [
    "!wget https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-12_H-768_A-12.zip -P bert_pretrained_models/;\n",
    "!unzip bert_pretrained_models/uncased_L-12_H-768_A-12.zip -d bert_pretrained_models/\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "JC60vOxGjI-M"
   },
   "outputs": [],
   "source": [
    "class InputExample():\n",
    "    \"\"\"A single training/test example for simple sequence classification.\"\"\"\n",
    "\n",
    "    def __init__(self, guid, text_a, text_b=None):\n",
    "        \"\"\"Constructs a InputExample.\n",
    "        Args:\n",
    "            guid: Unique id for the example.\n",
    "            text_a: string. The untokenized text of the first sequence.\n",
    "                For single sequence tasks, only this sequence must be specified.\n",
    "            text_b: (Optional) string. The untokenized text of the second\n",
    "                sequence. Only must be specified for sequence pair tasks.\n",
    "            label: (Optional) string. The label of the example. This should be\n",
    "                specified for train and dev examples, but not for test examples.\n",
    "        \"\"\"\n",
    "        self.guid = guid\n",
    "        self.src_txt = text_a\n",
    "        self.tgt_txt = text_b\n",
    "        \n",
    "class InputFeatures():\n",
    "    \"\"\"A single set of features of data.\"\"\"\n",
    "\n",
    "    def __init__(self, src_input_ids,src_input_mask,src_segment_ids,tgt_input_ids,tgt_input_mask,tgt_labels):\n",
    "        self.src_input_ids = src_input_ids\n",
    "        self.src_input_mask = src_input_mask\n",
    "        self.src_segment_ids = src_segment_ids\n",
    "        self.tgt_input_ids = tgt_input_ids\n",
    "        self.tgt_input_mask = tgt_input_mask \n",
    "        self.tgt_labels = tgt_labels\n",
    "        \n",
    "       \n",
    "class DataProcessor(object):\n",
    "    \"\"\"Base class for data converters for sequence classification data sets.\"\"\"\n",
    "\n",
    "    def get_train_examples(self, data_dir):\n",
    "        \"\"\"Gets a collection of `InputExample`s for the train set.\"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def get_dev_examples(self, data_dir):\n",
    "        \"\"\"Gets a collection of `InputExample`s for the dev set.\"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def get_test_examples(self, data_dir):\n",
    "        \"\"\"Gets a collection of `InputExample`s for prediction.\"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def get_labels(self):\n",
    "        \"\"\"Gets the list of labels for this data set.\"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    @classmethod\n",
    "    def _read_tsv(cls, input_file, quotechar=None):\n",
    "        \"\"\"Reads a tab separated value file.\"\"\"\n",
    "        with tf.gfile.Open(input_file, \"r\") as f:\n",
    "            reader = csv.reader(f, delimiter=\"\\t\", quotechar=quotechar)\n",
    "            lines = []\n",
    "            i = 0\n",
    "            for line in reader:\n",
    "                lines.append(line)\n",
    "        return lines\n",
    "\n",
    "\n",
    "    @classmethod\n",
    "    def _read_file(cls, input_file, quotechar=None):\n",
    "        \"\"\"Reads a tab separated value file.\"\"\"\n",
    "        with tf.gfile.Open(input_file, \"r\") as f:\n",
    "            reader = csv.reader(f, delimiter=\"\\n\", quotechar=quotechar)\n",
    "            lines = []\n",
    "            i = 0\n",
    "            for line in reader:\n",
    "                lines.append(line)\n",
    "        return lines\n",
    "      \n",
    "      \n",
    "class CNNDailymail(DataProcessor):\n",
    "    \"\"\"Processor for the CoLA data set (GLUE version).\"\"\"\n",
    "\n",
    "    def get_train_examples(self, data_dir):\n",
    "        \"\"\"See base class.\"\"\"\n",
    "        return self._create_examples(\n",
    "            self._read_file(os.path.join(data_dir, \"train_story.txt\")),self._read_file(os.path.join(data_dir, \"train_summ.txt\")),\n",
    "            \"train\")\n",
    "\n",
    "    def get_dev_examples(self, data_dir):\n",
    "        \"\"\"See base class.\"\"\"\n",
    "        return self._create_examples(\n",
    "            self._read_file(os.path.join(data_dir, \"eval_story.txt\")),self._read_file(os.path.join(data_dir, \"eval_summ.txt\")),\n",
    "            \"dev\")\n",
    "\n",
    "    def get_test_examples(self, data_dir):\n",
    "        \"\"\"See base class.\"\"\"\n",
    "        return self._create_examples(\n",
    "            self._read_file(os.path.join(data_dir, \"test_story.txt\")),self._read_file(os.path.join(data_dir, \"test_summ.txt\")),\n",
    "            \"test\")\n",
    "\n",
    "    def _create_examples(self, src_lines,tgt_lines,set_type):\n",
    "        examples = [] \n",
    "        for i,data in enumerate(zip(src_lines,tgt_lines)):\n",
    "            guid = \"%s-%s\" % (set_type, i)\n",
    "            if set_type == \"test\" and i == 0:\n",
    "                continue\n",
    "            else:\n",
    "                #print(data)\n",
    "                if len(data[0])==0 or len(data[1])==0:\n",
    "                  continue\n",
    "                src_lines = tokenization.convert_to_unicode(data[0][0])\n",
    "                tgt_lines = tokenization.convert_to_unicode(data[1][0])\n",
    "                examples.append(InputExample(guid=guid, text_a=src_lines,\n",
    "                                         text_b=tgt_lines))\n",
    "        return examples\n",
    "  \n",
    "  \n",
    "def file_based_convert_examples_to_features(\n",
    "        examples, max_seq_length_src,max_seq_length_tgt,tokenizer, output_file):\n",
    "    \"\"\"Convert a set of `InputExample`s to a TFRecord file.\"\"\"\n",
    "\n",
    "    writer = tf.python_io.TFRecordWriter(output_file)\n",
    "\n",
    "    for (ex_index, example) in enumerate(examples):\n",
    "        #print(\"ex_index\",ex_index)\n",
    "\n",
    "        if (ex_index+1) %1000 == 0 :\n",
    "          print(\"------------processed..{}...examples\".format(ex_index))\n",
    "          \n",
    "        feature = convert_single_example(ex_index, example,\n",
    "                                         max_seq_length_src,max_seq_length_tgt,tokenizer)\n",
    "\n",
    "        def create_int_feature(values):\n",
    "            return tf.train.Feature(\n",
    "                int64_list=tf.train.Int64List(value=list(values)))\n",
    "\n",
    "        features = collections.OrderedDict()\n",
    "        features[\"src_input_ids\"] = create_int_feature(feature.src_input_ids)\n",
    "        features[\"src_input_mask\"] = create_int_feature(feature.src_input_mask)\n",
    "        features[\"src_segment_ids\"] = create_int_feature(feature.src_segment_ids)\n",
    "\n",
    "        features[\"tgt_input_ids\"] = create_int_feature(feature.tgt_input_ids)\n",
    "        features[\"tgt_input_mask\"] = create_int_feature(feature.tgt_input_mask)\n",
    "        features['tgt_labels'] = create_int_feature(feature.tgt_labels)\n",
    "        \n",
    "        \n",
    "        \n",
    "        #print(feature.tgt_labels)\n",
    "        \n",
    "\n",
    "        tf_example = tf.train.Example(\n",
    "            features=tf.train.Features(feature=features))\n",
    "        writer.write(tf_example.SerializeToString())\n",
    "\n",
    "\n",
    "def convert_single_example(ex_index, example, max_seq_length_src,max_seq_length_tgt,\n",
    "                           tokenizer):\n",
    "    \"\"\"Converts a single `InputExample` into a single `InputFeatures`.\"\"\"\n",
    "    \"\"\"\n",
    "    label_map = {}\n",
    "    for (i, label) in enumerate(label_list):\n",
    "        label_map[label] = i\n",
    "    \"\"\"\n",
    "    tokens_a = tokenizer.tokenize(example.src_txt)\n",
    "    tokens_b = tokenizer.tokenize(example.tgt_txt)\n",
    "\n",
    "\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",
    "    if len(tokens_a) > max_seq_length_src - 2:\n",
    "            tokens_a = tokens_a[0:(max_seq_length_src - 2)]\n",
    "    \n",
    "    if len(tokens_b) > max_seq_length_tgt - 2:\n",
    "            tokens_b = tokens_b[0:(max_seq_length_tgt - 2)]\n",
    "\n",
    "    \n",
    "    tokens_src = []\n",
    "    segment_ids_src = []\n",
    "    tokens_src.append(\"[CLS]\")\n",
    "    segment_ids_src.append(0)\n",
    "    for token in tokens_a:\n",
    "        tokens_src.append(token)\n",
    "        segment_ids_src.append(0)\n",
    "    tokens_src.append(\"[SEP]\")\n",
    "    segment_ids_src.append(0)\n",
    "  \n",
    "\n",
    "    tokens_tgt = []\n",
    "    segment_ids_tgt = []\n",
    "    tokens_tgt.append(\"[CLS]\")\n",
    "    #segment_ids_tgt.append(0)\n",
    "    for token in tokens_b:\n",
    "        tokens_tgt.append(token)\n",
    "        #segment_ids_tgt.append(0)\n",
    "    tokens_tgt.append(\"[SEP]\")\n",
    "    #segment_ids_tgt.append(0)\n",
    "\n",
    "    input_ids_src = tokenizer.convert_tokens_to_ids(tokens_src)\n",
    "   \n",
    "    \n",
    "\n",
    "    input_ids_tgt = tokenizer.convert_tokens_to_ids(tokens_tgt)\n",
    "    \n",
    "    #Adding begiining and end token\n",
    "    input_ids_tgt = input_ids_tgt[:-1] \n",
    "    \n",
    "    input_mask_src = [1] * len(input_ids_src)\n",
    "\n",
    "\n",
    "    input_mask_tgt = [1] * len(input_ids_tgt)\n",
    "    \n",
    "    labels_tgt = input_ids_tgt[1:]\n",
    "    \n",
    "    \n",
    "    labels_tgt.append(0)\n",
    "    \n",
    "    #print(len(input_ids_tgt))\n",
    "    #print(len(input_mask_tgt))\n",
    "    #print(len(labels_tgt))\n",
    "    #print(len(segment_ids_tgt))\n",
    "    \n",
    "    while len(input_ids_src) < max_seq_length_src:\n",
    "        input_ids_src.append(0)\n",
    "        input_mask_src.append(0)\n",
    "        segment_ids_src.append(0)\n",
    "\n",
    "    while len(input_ids_tgt) < max_seq_length_tgt:\n",
    "        input_ids_tgt.append(0)\n",
    "        input_mask_tgt.append(0)\n",
    "        segment_ids_tgt.append(0)\n",
    "        labels_tgt.append(0)\n",
    "\n",
    "    feature = InputFeatures( src_input_ids=input_ids_src,src_input_mask=input_mask_src,src_segment_ids=segment_ids_src,\n",
    "        tgt_input_ids=input_ids_tgt,tgt_input_mask=input_mask_tgt,tgt_labels=labels_tgt)\n",
    "\n",
    "    \n",
    "    return feature\n",
    "\n",
    "\n",
    "def file_based_input_fn_builder(input_file, max_seq_length_src,max_seq_length_tgt, is_training,\n",
    "                                drop_remainder, is_distributed=False):\n",
    "    \"\"\"Creates an `input_fn` closure to be passed to TPUEstimator.\"\"\"\n",
    "\n",
    "    name_to_features = {\n",
    "        \"src_input_ids\": tf.FixedLenFeature([max_seq_length_src], tf.int64),\n",
    "        \"src_input_mask\": tf.FixedLenFeature([max_seq_length_src], tf.int64),\n",
    "        \"src_segment_ids\": tf.FixedLenFeature([max_seq_length_src], tf.int64),\n",
    "        \"tgt_input_ids\": tf.FixedLenFeature([max_seq_length_tgt], tf.int64),\n",
    "        \"tgt_input_mask\": tf.FixedLenFeature([max_seq_length_tgt], tf.int64),\n",
    "        \"tgt_labels\" : tf.FixedLenFeature([max_seq_length_tgt], tf.int64),\n",
    "        \n",
    "        \n",
    "    }\n",
    "\n",
    "    def _decode_record(record, name_to_features):\n",
    "        \"\"\"Decodes a record to a TensorFlow example.\"\"\"\n",
    "        example = tf.parse_single_example(record, name_to_features)\n",
    "        print(example)\n",
    "        print(example.keys())\n",
    "\n",
    "        # tf.Example only supports tf.int64, but the TPU only supports tf.int32.\n",
    "        # So cast all int64 to int32.\n",
    "        for name in list(example.keys()):\n",
    "            t = example[name]\n",
    "            if t.dtype == tf.int64:\n",
    "                t = tf.to_int32(t)\n",
    "            example[name] = t\n",
    "\n",
    "        return example\n",
    "\n",
    "    def input_fn(params):\n",
    "        \"\"\"The actual input function.\"\"\"\n",
    "        batch_size = params[\"batch_size\"]\n",
    "\n",
    "        # For training, we want a lot of parallel reading and shuffling.\n",
    "        # For eval, we want no shuffling and parallel reading doesn't matter.\n",
    "        d = tf.data.TFRecordDataset(input_file)\n",
    "        if is_training:\n",
    "\n",
    "            if is_distributed:\n",
    "                import horovod.tensorflow as hvd\n",
    "                tf.logging.info('distributed mode is enabled.'\n",
    "                                'size:{} rank:{}'.format(hvd.size(), hvd.rank()))\n",
    "                # https://github.com/uber/horovod/issues/223\n",
    "                d = d.shard(hvd.size(), hvd.rank())\n",
    "\n",
    "                d = d.repeat()\n",
    "                d = d.shuffle(buffer_size=100)\n",
    "                d = d.apply(\n",
    "                    tf.contrib.data.map_and_batch(\n",
    "                        lambda record: _decode_record(record, name_to_features),\n",
    "                        batch_size=batch_size//hvd.size(),\n",
    "                        drop_remainder=drop_remainder))\n",
    "            else:\n",
    "                tf.logging.info('distributed mode is not enabled.')\n",
    "                d = d.repeat()\n",
    "                d = d.shuffle(buffer_size=100)\n",
    "                d = d.apply(\n",
    "                    tf.contrib.data.map_and_batch(\n",
    "                        lambda record: _decode_record(record, name_to_features),\n",
    "                        batch_size=batch_size,\n",
    "                        drop_remainder=drop_remainder))\n",
    "\n",
    "        else:\n",
    "            d = d.apply(\n",
    "                tf.contrib.data.map_and_batch(\n",
    "                    lambda record: _decode_record(record, name_to_features),\n",
    "                    batch_size=batch_size,\n",
    "                    drop_remainder=drop_remainder))\n",
    "\n",
    "        return d\n",
    "    return input_fn\n",
    "  \n",
    "  \n",
    "def get_dataset(processor,\n",
    "                tokenizer,\n",
    "                data_dir,\n",
    "                max_seq_length_src,\n",
    "                max_seq_length_tgt,\n",
    "                batch_size,\n",
    "                mode,\n",
    "                output_dir,\n",
    "                is_distributed=False):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        processor: Data Preprocessor, must have get_lables,\n",
    "            get_train/dev/test/examples methods defined.\n",
    "        tokenizer: The Sentence Tokenizer. Generally should be\n",
    "            SentencePiece Model.\n",
    "        data_dir: The input data directory.\n",
    "        max_seq_length: Max sequence length.\n",
    "        batch_size: mini-batch size.\n",
    "        model: `train`, `eval` or `test`.\n",
    "        output_dir: The directory to save the TFRecords in.\n",
    "    \"\"\"\n",
    "    #label_list = processor.get_labels()\n",
    "    if mode == 'train':\n",
    "        #train_examples = processor.get_train_examples(data_dir)\n",
    "        #train_file = os.path.join(output_dir, \"train.tf_record\")\n",
    "        train_file = \"gs://bert_summ/train.tf_record\"\n",
    "        #file_based_convert_examples_to_features(\n",
    "        #    train_examples, max_seq_length_src,max_seq_length_tgt,\n",
    "        #    tokenizer, train_file)\n",
    "        dataset = file_based_input_fn_builder(\n",
    "            input_file=train_file,\n",
    "            max_seq_length_src=max_seq_length_src,\n",
    "            max_seq_length_tgt =max_seq_length_tgt,\n",
    "            is_training=True,\n",
    "            drop_remainder=True,\n",
    "            is_distributed=is_distributed)({'batch_size': batch_size})\n",
    "    elif mode == 'eval':\n",
    "        #eval_examples = processor.get_dev_examples(data_dir)\n",
    "        #eval_file = os.path.join(output_dir, \"eval.tf_record\")\n",
    "        eval_file = \"gs://bert_summ/eval.tf_record\"\n",
    "        #file_based_convert_examples_to_features(\n",
    "        #    eval_examples, max_seq_length_src,max_seq_length_tgt,\n",
    "        #    tokenizer, eval_file)\n",
    "        dataset = file_based_input_fn_builder(\n",
    "            input_file=eval_file,\n",
    "            max_seq_length_src=max_seq_length_src,\n",
    "            max_seq_length_tgt =max_seq_length_tgt,\n",
    "            is_training=False,\n",
    "            drop_remainder=True,\n",
    "            is_distributed=is_distributed)({'batch_size': batch_size})\n",
    "    elif mode == 'test':\n",
    "      \n",
    "        #test_examples = processor.get_test_examples(data_dir)\n",
    "        #test_file = os.path.join(output_dir, \"predict.tf_record\")\n",
    "        test_file = \"gs://bert_summ/predict.tf_record\"\n",
    "        \n",
    "        #file_based_convert_examples_to_features(\n",
    "        #    test_examples, max_seq_length_src,max_seq_length_tgt,\n",
    "        #    tokenizer, test_file)\n",
    "        dataset = file_based_input_fn_builder(\n",
    "            input_file=test_file,\n",
    "            max_seq_length_src=max_seq_length_src,\n",
    "            max_seq_length_tgt =max_seq_length_tgt,\n",
    "            is_training=False,\n",
    "            drop_remainder=True,\n",
    "            is_distributed=is_distributed)({'batch_size': batch_size})\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 227
    },
    "colab_type": "code",
    "id": "Nn4vhTvJjT0D",
    "outputId": "f4080019-27c1-4f01-e498-a2b1d7600beb"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/Abstractive-Summarization-With-Transfer-Learning-master/texar_repo/examples/bert/utils/tokenization.py:41: The name tf.gfile.GFile is deprecated. Please use tf.io.gfile.GFile instead.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/Abstractive-Summarization-With-Transfer-Learning-master/texar_repo/examples/bert/utils/tokenization.py:41: The name tf.gfile.GFile is deprecated. Please use tf.io.gfile.GFile instead.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:distributed mode is not enabled.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:distributed mode is not enabled.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-8-3918a62c9cd3>:297: map_and_batch (from tensorflow.contrib.data.python.ops.batching) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.data.experimental.map_and_batch(...)`.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-8-3918a62c9cd3>:297: map_and_batch (from tensorflow.contrib.data.python.ops.batching) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.data.experimental.map_and_batch(...)`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/miniconda3/lib/python3.6/site-packages/tensorflow/contrib/data/python/ops/batching.py:273: map_and_batch (from tensorflow.python.data.experimental.ops.batching) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.data.Dataset.map(map_func, num_parallel_calls)` followed by `tf.data.Dataset.batch(batch_size, drop_remainder)`. Static tf.data optimizations will take care of using the fused implementation.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/miniconda3/lib/python3.6/site-packages/tensorflow/contrib/data/python/ops/batching.py:273: map_and_batch (from tensorflow.python.data.experimental.ops.batching) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.data.Dataset.map(map_func, num_parallel_calls)` followed by `tf.data.Dataset.batch(batch_size, drop_remainder)`. Static tf.data optimizations will take care of using the fused implementation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'src_input_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:0' shape=(512,) dtype=int64>, 'src_input_mask': <tf.Tensor 'ParseSingleExample/ParseSingleExample:1' shape=(512,) dtype=int64>, 'src_segment_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:2' shape=(512,) dtype=int64>, 'tgt_input_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:3' shape=(400,) dtype=int64>, 'tgt_input_mask': <tf.Tensor 'ParseSingleExample/ParseSingleExample:4' shape=(400,) dtype=int64>, 'tgt_labels': <tf.Tensor 'ParseSingleExample/ParseSingleExample:5' shape=(400,) dtype=int64>}\n",
      "dict_keys(['src_input_ids', 'src_input_mask', 'src_segment_ids', 'tgt_input_ids', 'tgt_input_mask', 'tgt_labels'])\n",
      "WARNING:tensorflow:From <ipython-input-8-3918a62c9cd3>:261: to_int32 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.cast` instead.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-8-3918a62c9cd3>:261: to_int32 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.cast` instead.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'src_input_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:0' shape=(512,) dtype=int64>, 'src_input_mask': <tf.Tensor 'ParseSingleExample/ParseSingleExample:1' shape=(512,) dtype=int64>, 'src_segment_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:2' shape=(512,) dtype=int64>, 'tgt_input_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:3' shape=(400,) dtype=int64>, 'tgt_input_mask': <tf.Tensor 'ParseSingleExample/ParseSingleExample:4' shape=(400,) dtype=int64>, 'tgt_labels': <tf.Tensor 'ParseSingleExample/ParseSingleExample:5' shape=(400,) dtype=int64>}\n",
      "dict_keys(['src_input_ids', 'src_input_mask', 'src_segment_ids', 'tgt_input_ids', 'tgt_input_mask', 'tgt_labels'])\n",
      "{'src_input_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:0' shape=(512,) dtype=int64>, 'src_input_mask': <tf.Tensor 'ParseSingleExample/ParseSingleExample:1' shape=(512,) dtype=int64>, 'src_segment_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:2' shape=(512,) dtype=int64>, 'tgt_input_ids': <tf.Tensor 'ParseSingleExample/ParseSingleExample:3' shape=(400,) dtype=int64>, 'tgt_input_mask': <tf.Tensor 'ParseSingleExample/ParseSingleExample:4' shape=(400,) dtype=int64>, 'tgt_labels': <tf.Tensor 'ParseSingleExample/ParseSingleExample:5' shape=(400,) dtype=int64>}\n",
      "dict_keys(['src_input_ids', 'src_input_mask', 'src_segment_ids', 'tgt_input_ids', 'tgt_input_mask', 'tgt_labels'])\n"
     ]
    }
   ],
   "source": [
    "bert_config = model_utils.transform_bert_to_texar_config(\n",
    "            os.path.join(bert_pretrain_dir, 'bert_config.json'))\n",
    "\n",
    "\n",
    "\n",
    "tokenizer = tokenization.FullTokenizer(\n",
    "        vocab_file=os.path.join(bert_pretrain_dir, 'vocab.txt'),\n",
    "        do_lower_case=True)\n",
    "\n",
    "vocab_size = len(tokenizer.vocab)\n",
    "\n",
    "processor = CNNDailymail()\n",
    "train_dataset = get_dataset(processor,tokenizer,\"./\",max_seq_length_src,max_seq_length_tgt,4,'train',\"./\")\n",
    "eval_dataset = get_dataset(processor,tokenizer,\"./\",max_seq_length_src,max_seq_length_tgt,4,'eval',\"./\")\n",
    "test_dataset = get_dataset(processor,tokenizer,\"./\",max_seq_length_src,max_seq_length_tgt,4,'test',\"./\")\n",
    "#del processor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kDhy0XGlIcY2"
   },
   "outputs": [],
   "source": [
    "del processor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "CBXxrBteAuVj",
    "outputId": "424fc4a8-d889-412f-ae34-766b961ce7ab"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "30522"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "lfw2JV11jsad"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-12-37df741c2270>:15: to_float (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.cast` instead.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-12-37df741c2270>:15: to_float (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.cast` instead.\n"
     ]
    }
   ],
   "source": [
    "#inputs to the model\n",
    "src_input_ids = tf.placeholder(tf.int64, shape=(None, None))\n",
    "src_segment_ids = tf.placeholder(tf.int64, shape=(None, None))\n",
    "tgt_input_ids = tf.placeholder(tf.int64, shape=(None, None))\n",
    "tgt_segment_ids = tf.placeholder(tf.int64, shape=(None, None))\n",
    "\n",
    "batch_size = tf.shape(src_input_ids)[0]\n",
    "\n",
    "src_input_length = tf.reduce_sum(1 - tf.to_int32(tf.equal(src_input_ids, 0)),\n",
    "                             axis=1)\n",
    "tgt_input_length = tf.reduce_sum(1 - tf.to_int32(tf.equal(src_input_ids, 0)),\n",
    "                             axis=1)\n",
    "\n",
    "labels = tf.placeholder(tf.int64, shape=(None, None))\n",
    "is_target = tf.to_float(tf.not_equal(labels, 0))\n",
    "\n",
    "\n",
    "global_step = tf.Variable(0, dtype=tf.int64, trainable=False)\n",
    "learning_rate = tf.placeholder(tf.float64, shape=(), name='lr')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jTFde06_kACm"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/utils/variables.py:50: The name tf.variable_scope is deprecated. Please use tf.compat.v1.variable_scope instead.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/utils/variables.py:50: The name tf.variable_scope is deprecated. Please use tf.compat.v1.variable_scope instead.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/data/data/data_iterators.py:326: The name tf.placeholder is deprecated. Please use tf.compat.v1.placeholder instead.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/data/data/data_iterators.py:326: The name tf.placeholder is deprecated. Please use tf.compat.v1.placeholder instead.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/data/data/data_iterators.py:328: The name tf.data.Iterator is deprecated. Please use tf.compat.v1.data.Iterator instead.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/data/data/data_iterators.py:328: The name tf.data.Iterator is deprecated. Please use tf.compat.v1.data.Iterator instead.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/data/data/data_iterators.py:329: DatasetV1.output_types (from tensorflow.python.data.ops.dataset_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.compat.v1.data.get_output_types(dataset)`.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/data/data/data_iterators.py:329: DatasetV1.output_types (from tensorflow.python.data.ops.dataset_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.compat.v1.data.get_output_types(dataset)`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/data/data/data_iterators.py:330: DatasetV1.output_shapes (from tensorflow.python.data.ops.dataset_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.compat.v1.data.get_output_shapes(dataset)`.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/data/data/data_iterators.py:330: DatasetV1.output_shapes (from tensorflow.python.data.ops.dataset_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.compat.v1.data.get_output_shapes(dataset)`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/data/data/data_iterators.py:334: DatasetV1.make_initializable_iterator (from tensorflow.python.data.ops.dataset_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `for ... in dataset:` to iterate over a dataset. If using `tf.estimator`, return the `Dataset` object directly from your input function. As a last resort, you can use `tf.compat.v1.data.make_initializable_iterator(dataset)`.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/data/data/data_iterators.py:334: DatasetV1.make_initializable_iterator (from tensorflow.python.data.ops.dataset_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `for ... in dataset:` to iterate over a dataset. If using `tf.estimator`, return the `Dataset` object directly from your input function. As a last resort, you can use `tf.compat.v1.data.make_initializable_iterator(dataset)`.\n"
     ]
    }
   ],
   "source": [
    "#create the iterator \n",
    "iterator = tx.data.FeedableDataIterator({\n",
    "        'train': train_dataset, 'eval': eval_dataset, 'test': test_dataset})\n",
    "\n",
    "batch = iterator.get_next()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "XSrDO5YBkPYh",
    "outputId": "a07623be-20ff-433c-a256-95f33b49b531"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Intializing the Bert Encoder Graph\n",
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/module_base.py:72: The name tf.make_template is deprecated. Please use tf.compat.v1.make_template instead.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/module_base.py:72: The name tf.make_template is deprecated. Please use tf.compat.v1.make_template instead.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/modules/embedders/embedder_utils.py:205: The name tf.get_variable is deprecated. Please use tf.compat.v1.get_variable instead.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/modules/embedders/embedder_utils.py:205: The name tf.get_variable is deprecated. Please use tf.compat.v1.get_variable instead.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/miniconda3/lib/python3.6/site-packages/tensorflow/python/ops/init_ops.py:1251: calling VarianceScaling.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Call initializer instance with the dtype argument instead of passing it to the constructor\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/miniconda3/lib/python3.6/site-packages/tensorflow/python/ops/init_ops.py:1251: calling VarianceScaling.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Call initializer instance with the dtype argument instead of passing it to the constructor\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/context.py:52: The name tf.get_collection_ref is deprecated. Please use tf.compat.v1.get_collection_ref instead.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/context.py:52: The name tf.get_collection_ref is deprecated. Please use tf.compat.v1.get_collection_ref instead.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/miniconda3/lib/python3.6/pydoc.py:1588: The name tf.layers.Dense is deprecated. Please use tf.compat.v1.layers.Dense instead.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/miniconda3/lib/python3.6/pydoc.py:1588: The name tf.layers.Dense is deprecated. Please use tf.compat.v1.layers.Dense instead.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/core/layers.py:597: The name tf.layers.Layer is deprecated. Please use tf.compat.v1.layers.Layer instead.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/core/layers.py:597: The name tf.layers.Layer is deprecated. Please use tf.compat.v1.layers.Layer instead.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/modules/encoders/transformer_encoders.py:340: dropout (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.dropout instead.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/modules/encoders/transformer_encoders.py:340: dropout (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.dropout instead.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/modules/networks/network_base.py:122: The name tf.layers.Dropout is deprecated. Please use tf.compat.v1.layers.Dropout instead.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/modules/networks/network_base.py:122: The name tf.layers.Dropout is deprecated. Please use tf.compat.v1.layers.Dropout instead.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/modules/networks/network_base.py:123: The name tf.layers.BatchNormalization is deprecated. Please use tf.compat.v1.layers.BatchNormalization instead.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/texar/texar/modules/networks/network_base.py:123: The name tf.layers.BatchNormalization is deprecated. Please use tf.compat.v1.layers.BatchNormalization instead.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-14-1bc1e0feda4d>:29: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.dense instead.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-14-1bc1e0feda4d>:29: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.dense instead.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading the bert pretrained weights\n",
      "WARNING:tensorflow:From /home/shify/paper/实验/Abstractive-Summarization-With-Transfer-Learning-master/texar_repo/examples/bert/utils/model_utils.py:179: The name tf.train.init_from_checkpoint is deprecated. Please use tf.compat.v1.train.init_from_checkpoint instead.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/paper/实验/Abstractive-Summarization-With-Transfer-Learning-master/texar_repo/examples/bert/utils/model_utils.py:179: The name tf.train.init_from_checkpoint is deprecated. Please use tf.compat.v1.train.init_from_checkpoint instead.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#encoder Bert model\n",
    "print(\"Intializing the Bert Encoder Graph\")\n",
    "with tf.variable_scope('bert'):\n",
    "        embedder = tx.modules.WordEmbedder(\n",
    "            vocab_size=bert_config.vocab_size,\n",
    "            hparams=bert_config.embed)\n",
    "        word_embeds = embedder(src_input_ids)\n",
    "\n",
    "        # Creates segment embeddings for each type of tokens.\n",
    "        segment_embedder = tx.modules.WordEmbedder(\n",
    "            vocab_size=bert_config.type_vocab_size,\n",
    "            hparams=bert_config.segment_embed)\n",
    "        segment_embeds = segment_embedder(src_segment_ids)\n",
    "\n",
    "        input_embeds = word_embeds + segment_embeds\n",
    "\n",
    "        # The BERT model (a TransformerEncoder)\n",
    "        encoder = tx.modules.TransformerEncoder(hparams=bert_config.encoder)\n",
    "        encoder_output = encoder(input_embeds, src_input_length)\n",
    "        \n",
    "        # Builds layers for downstream classification, which is also initialized\n",
    "        # with BERT pre-trained checkpoint.\n",
    "        with tf.variable_scope(\"pooler\"):\n",
    "            # Uses the projection of the 1st-step hidden vector of BERT output\n",
    "            # as the representation of the sentence\n",
    "            bert_sent_hidden = tf.squeeze(encoder_output[:, 0:1, :], axis=1)\n",
    "            bert_sent_output = tf.layers.dense(\n",
    "                bert_sent_hidden, config_downstream.hidden_dim,\n",
    "                activation=tf.tanh)\n",
    "            output = tf.layers.dropout(\n",
    "                bert_sent_output, rate=0.1, training=tx.global_mode_train())\n",
    "\n",
    "\n",
    "print(\"loading the bert pretrained weights\")\n",
    "# Loads pretrained BERT model parameters\n",
    "init_checkpoint = os.path.join(bert_pretrain_dir, 'bert_model.ckpt')\n",
    "#init_checkpoint = \"gs://cloud-tpu-checkpoints/bert/uncased_L-12_H-768_A-12/bert_model.ckpt\"\n",
    "model_utils.init_bert_checkpoint(init_checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "C5m48bu5kVXm"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/miniconda3/lib/python3.6/pydoc.py:1588: The name tf.train.AdamOptimizer is deprecated. Please use tf.compat.v1.train.AdamOptimizer instead.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/miniconda3/lib/python3.6/pydoc.py:1588: The name tf.train.AdamOptimizer is deprecated. Please use tf.compat.v1.train.AdamOptimizer instead.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#decoder part and mle losss\n",
    "tgt_embedding = tf.concat(\n",
    "    [tf.zeros(shape=[1, embedder.dim]), embedder.embedding[1:, :]], axis=0)\n",
    "\n",
    "decoder = tx.modules.TransformerDecoder(embedding=tgt_embedding,\n",
    "                             hparams=dcoder_config)\n",
    "# For training\n",
    "outputs = decoder(\n",
    "    memory=encoder_output,\n",
    "    memory_sequence_length=src_input_length,\n",
    "    inputs=embedder(tgt_input_ids),\n",
    "    sequence_length=tgt_input_length,\n",
    "    decoding_strategy='train_greedy',\n",
    "    mode=tf.estimator.ModeKeys.TRAIN\n",
    ")\n",
    "\n",
    "mle_loss = transformer_utils.smoothing_cross_entropy(\n",
    "        outputs.logits, labels, vocab_size, loss_label_confidence)\n",
    "mle_loss = tf.reduce_sum(mle_loss * is_target) / tf.reduce_sum(is_target)\n",
    "\n",
    "train_op = tx.core.get_train_op(\n",
    "        mle_loss,\n",
    "        learning_rate=learning_rate,\n",
    "        global_step=global_step,\n",
    "        hparams=opt)\n",
    "\n",
    "tf.summary.scalar('lr', learning_rate)\n",
    "tf.summary.scalar('mle_loss', mle_loss)\n",
    "summary_merged = tf.summary.merge_all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "sfDuR-SVkdhF"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/miniconda3/lib/python3.6/site-packages/tensorflow/python/ops/math_ops.py:2403: add_dispatch_support.<locals>.wrapper (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.where in 2.0, which has the same broadcast rule as np.where\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/shify/miniconda3/lib/python3.6/site-packages/tensorflow/python/ops/math_ops.py:2403: add_dispatch_support.<locals>.wrapper (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.where in 2.0, which has the same broadcast rule as np.where\n"
     ]
    }
   ],
   "source": [
    "#prediction \n",
    "start_tokens = tf.fill([tx.utils.get_batch_size(src_input_ids)],\n",
    "                       bos_token_id)\n",
    "predictions = decoder(\n",
    "    memory=encoder_output,\n",
    "    memory_sequence_length=src_input_length,\n",
    "    decoding_strategy='infer_greedy',\n",
    "    beam_width=beam_width,\n",
    "    alpha=alpha,\n",
    "    start_tokens=start_tokens,\n",
    "    end_token=eos_token_id,\n",
    "    max_decoding_length=400,\n",
    "    mode=tf.estimator.ModeKeys.PREDICT\n",
    ")\n",
    "if beam_width <= 1:\n",
    "    inferred_ids = predictions[0].sample_id\n",
    "else:\n",
    "    # Uses the best sample by beam search\n",
    "    inferred_ids = predictions['sample_id'][:, :, 0]\n",
    "\n",
    "\n",
    "saver = tf.train.Saver(max_to_keep=5)\n",
    "best_results = {'score': 0, 'epoch': -1}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TCmgMIV6kzO4"
   },
   "outputs": [],
   "source": [
    "def _train_epoch(sess, epoch, step, smry_writer):\n",
    "        \n",
    "            \n",
    "        fetches = {\n",
    "            'step': global_step,\n",
    "            'train_op': train_op,\n",
    "            'smry': summary_merged,\n",
    "            'loss': mle_loss,\n",
    "        }\n",
    "\n",
    "        while True:\n",
    "            try:\n",
    "              feed_dict = {\n",
    "                iterator.handle: iterator.get_handle(sess, 'train'),\n",
    "                tx.global_mode(): tf.estimator.ModeKeys.TRAIN,\n",
    "              }\n",
    "              op = sess.run([batch],feed_dict)\n",
    "              feed_dict = {\n",
    "                   src_input_ids:op[0]['src_input_ids'],\n",
    "                   src_segment_ids : op[0]['src_segment_ids'],\n",
    "                   tgt_input_ids:op[0]['tgt_input_ids'],\n",
    "\n",
    "                   labels:op[0]['tgt_labels'],\n",
    "                   learning_rate: utils.get_lr(step, lr),\n",
    "                   tx.global_mode(): tf.estimator.ModeKeys.TRAIN\n",
    "                }\n",
    "\n",
    "\n",
    "              fetches_ = sess.run(fetches, feed_dict=feed_dict)\n",
    "              step, loss = fetches_['step'], fetches_['loss']\n",
    "              if step and step % display_steps == 0:\n",
    "                  logger.info('step: %d, loss: %.4f', step, loss)\n",
    "                  print('step: %d, loss: %.4f' % (step, loss))\n",
    "                  smry_writer.add_summary(fetches_['smry'], global_step=step)\n",
    "\n",
    "              if step and step % 1000 == 0:\n",
    "                  model_path = \"gs://bert_summ/models/model_\"+str(step)+\".ckpt\"\n",
    "                  logger.info('saving model to %s', model_path)\n",
    "                  print('saving model to %s' % model_path)\n",
    "                  saver.save(sess, model_path)\n",
    "              if step and step % eval_steps == 0:\n",
    "                  _eval_epoch(sess, epoch, mode='eval')\n",
    "            except tf.errors.OutOfRangeError:\n",
    "                break\n",
    "\n",
    "        return step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "60_hbYdak5rd"
   },
   "outputs": [],
   "source": [
    "def _eval_epoch(sess, epoch, mode):\n",
    "\n",
    "        references, hypotheses = [], []\n",
    "        bsize = test_batch_size\n",
    "        fetches = {\n",
    "                'inferred_ids': inferred_ids,\n",
    "            }\n",
    "        bno=0\n",
    "        while True:\n",
    "            \n",
    "            #print(\"Temp\",temp)\n",
    "            try:\n",
    "              print(\"Batch\",bno)\n",
    "              feed_dict = {\n",
    "              iterator.handle: iterator.get_handle(sess, 'eval'),\n",
    "              tx.global_mode(): tf.estimator.ModeKeys.EVAL,\n",
    "              }\n",
    "              op = sess.run([batch],feed_dict)\n",
    "              feed_dict = {\n",
    "                   src_input_ids:op[0]['src_input_ids'],\n",
    "                   src_segment_ids : op[0]['src_segment_ids'],\n",
    "                   tx.global_mode(): tf.estimator.ModeKeys.EVAL\n",
    "              }\n",
    "              fetches_ = sess.run(fetches, feed_dict=feed_dict)\n",
    "              labels = op[0]['tgt_labels']\n",
    "              hypotheses.extend(h.tolist() for h in fetches_['inferred_ids'])\n",
    "              references.extend(r.tolist() for r in labels)\n",
    "              hypotheses = utils.list_strip_eos(hypotheses, eos_token_id)\n",
    "              references = utils.list_strip_eos(references, eos_token_id)\n",
    "              bno = bno+1\n",
    "              \n",
    "            except tf.errors.OutOfRangeError:\n",
    "                break\n",
    "\n",
    "\n",
    "        if mode == 'eval':\n",
    "            # Writes results to files to evaluate BLEU\n",
    "            # For 'eval' mode, the BLEU is based on token ids (rather than\n",
    "            # text tokens) and serves only as a surrogate metric to monitor\n",
    "            # the training process\n",
    "            #fname = os.path.join(model_dir, 'tmp.eval')\n",
    "            fname = \"./tmp.eval\"\n",
    "            #fname = \"gs://bert_summ/models/tmp.eval\"\n",
    "            hypotheses = tx.utils.str_join(hypotheses)\n",
    "            references = tx.utils.str_join(references)\n",
    "            hyp_fn, ref_fn = tx.utils.write_paired_text(\n",
    "                hypotheses, references, fname, mode='s')\n",
    "            eval_bleu = bleu_wrapper(ref_fn, hyp_fn, case_sensitive=True)\n",
    "            eval_bleu = 100. * eval_bleu\n",
    "            logger.info('epoch: %d, eval_bleu %.4f', epoch, eval_bleu)\n",
    "            print('epoch: %d, eval_bleu %.4f' % (epoch, eval_bleu))\n",
    "\n",
    "            if eval_bleu > best_results['score']:\n",
    "                logger.info('epoch: %d, best bleu: %.4f', epoch, eval_bleu)\n",
    "                best_results['score'] = eval_bleu\n",
    "                best_results['epoch'] = epoch\n",
    "                #model_path = os.path.join(model_dir, 'best-model.ckpt')\n",
    "                model_path = \"gs://bert_summ/models/best-model.ckpt\"\n",
    "                logger.info('saving model to %s', model_path)\n",
    "                print('saving model to %s' % model_path)\n",
    "                saver.save(sess, model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 357
    },
    "colab_type": "code",
    "id": "v9b6ghCsnT90",
    "outputId": "bc4a0e8b-e9ad-408c-a92f-c6b4455b2d03"
   },
   "outputs": [],
   "source": [
    "#tx.utils.maybe_create_dir(model_dir)\n",
    "#logging_file = os.path.join(model_dir, 'logging.txt')\n",
    "\n",
    "model_dir = \"gs://bert_summ/models/\"\n",
    "logging_file= \"logging.txt\"\n",
    "logger = utils.get_logger(logging_file)\n",
    "\n",
    "cpu_num = int(os.environ.get('CPU_NUM', 10))\n",
    "config = tf.ConfigProto(device_count={\"CPU\": cpu_num},\n",
    "            inter_op_parallelism_threads = cpu_num,\n",
    "            intra_op_parallelism_threads = cpu_num,\n",
    "            log_device_placement=True)\n",
    "\n",
    "with tf.Session(config=config) as sess:\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    sess.run(tf.local_variables_initializer())\n",
    "    sess.run(tf.tables_initializer())\n",
    "\n",
    "    smry_writer = tf.summary.FileWriter(model_dir, graph=sess.graph)\n",
    "\n",
    "    if run_mode == 'train_and_evaluate':\n",
    "        logger.info('Begin running with train_and_evaluate mode')\n",
    "\n",
    "        if tf.train.latest_checkpoint(model_dir) is not None:\n",
    "            logger.info('Restore latest checkpoint in %s' % model_dir)\n",
    "            saver.restore(sess, tf.train.latest_checkpoint(model_dir))\n",
    "        \n",
    "        iterator.initialize_dataset(sess)\n",
    "\n",
    "        step = 0\n",
    "        for epoch in range(max_train_epoch):\n",
    "          iterator.restart_dataset(sess, 'train')\n",
    "          step = _train_epoch(sess, epoch, step, smry_writer)\n",
    "\n",
    "    elif run_mode == 'test':\n",
    "        logger.info('Begin running with test mode')\n",
    "\n",
    "        logger.info('Restore latest checkpoint in %s' % model_dir)\n",
    "        saver.restore(sess, tf.train.latest_checkpoint(model_dir))\n",
    "\n",
    "        _eval_epoch(sess, 0, mode='test')\n",
    "\n",
    "    else:\n",
    "        raise ValueError('Unknown mode: {}'.format(run_mode))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "L35eRRNKSoOV"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "include_colab_link": true,
   "name": "BERT SUMM.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}