<jupyter_start><jupyter_text>Setup<jupyter_code>!pip install pytorch-pretrained-bert
!pip install googledrivedownloader
!pip install pyrouge
!pip install tensorboardX
def download_from_gdrive(file_id, destination_path, unzip=False):
    from google_drive_downloader import GoogleDriveDownloader as gdd

    gdd.download_file_from_google_drive(file_id=file_id,
                                        dest_path=destination_path,
                                        unzip=unzip)
from google.colab import auth

auth.authenticate_user()
PROCESSED_DATA = '1-NJKRoNzk2ugjjB2mfnmpR15KYB_Fr0s'
CNN_STORIES = '0BwmD_VLjROrfTHk4NFg2SndKcjQ'
DAILY_MAIN_STORIES = '0BwmD_VLjROrfM1BxdkxVaTY2bWs'
CORENLP_PATH = '/home/stanford-corenlp-full-2018-10-05/stanford-corenlp-3.9.2.jar'
MAP_PATH = '/home/BertSum/urls/'
SIMPLER_JSON_PATH = '/home/json_data/cnndm'
TOKENIZED_PATH = '/home/merged_stories_tokenized/'
RAW_PATH = '/home/raw_stories/'
FINAL_JSON_PATH = '/home/json_data/'
BERT_DATA_PATH = '/home/bert_data/'
import numpy as np
import pandas as pd<jupyter_output><empty_output><jupyter_text>BERSTUM - High level implementation  Prepare sample data<jupyter_code>!wget http://nlp.stanford.edu/software/stanford-corenlp-full-2018-10-05.zip -P /home &>/dev/null
!unzip /home/stanford-corenlp-full-2018-10-05.zip -d /home &>/dev/null
ls /home/stanford-corenlp-full-2018-10-05/
#!export CLASSPATH=/root/stanford-corenlp-full-2018-10-05/stanford-corenlp-3.9.2.jar#:/root/stanford-corenlp-full-2018-10-05/stanford-corenlp-3.9.2-models.jar
#download_from_gdrive(file_id=PROCESSED_DATA, destination_path='/home/bert_data/cnn_dailymail.zip', unzip=True)
#!rm /home/bert_data/cnn_dailymail.zip
!mkdir {TOKENIZED_PATH}
!mkdir /home/raw_stories # This is where we store the individual story files
!mkdir /home/story_dump # This is where we store the story files from the gdrive download
!mkdir /home/story_dump/all_stories
download_from_gdrive(file_id=CNN_STORIES, destination_path='/home/story_dump/cnn_stories.tgz', unzip=False)
download_from_gdrive(file_id=DAILY_MAIN_STORIES, destination_path='/home/story_dump/dailymail_stories.tgz', unzip=False)

ls /home/bert_data
ls /home/story_dump
!tar xvzf /home/story_dump/cnn_stories.tgz  -C /home/story_dump/all_stories/ &>/dev/null
!tar xvzf /home/story_dump/dailymail_stories.tgz  -C /home/story_dump/all_stories/ &>/dev/null
cnn_stories = !ls /home/story_dump/all_stories/cnn/stories
dailymail_stories = !ls /home/story_dump/all_stories/dailymail/stories
len(cnn_stories), len(dailymail_stories)
all_stories = cnn_stories + dailymail_stories
all_stories[:5]
import numpy as np
# Random sample of 5 stories for both cnn and daily mail (total sample size of 10)
sample_cnn_stories = np.random.permutation(cnn_stories)[:5]
sample_dailymail_stories = np.random.permutation(dailymail_stories)[:5]
for story_ in sample_cnn_stories:
    !cp /home/story_dump/all_stories/cnn/stories/{story_} /home/story_dump/

for story_ in sample_dailymail_stories:
    !cp /home/story_dump/all_stories/dailymail/stories/{story_} /home/story_dump/
    
ls /home/story_dump/<jupyter_output>0f848ac3922e8deaf6a803b2713ed867bc1d19e3.story
5b76984e7d6c3ffb9e10b0fe9b356e90beac4c0b.story
63f019e87c5c3e7cb247a6942832c9395340e733.story
69a61a5463de5eea642fe884d019c6b4b4c2ee9f.story
6e4e310308475f1c7a47dc98074bb187bfd05e72.story
965ec9ba5b923cdf09b97fa7eadc79d7a48438ed.story
[0m[01;34mall_stories[0m/
cnn_stories.tgz
d50ea3563e1cd9c0d744d0e8b1a4ba6f8bb2c7dc.story
dailymail_stories.tgz
e25064cc1a7a29921c79bc8e1d476a4a8fbabf60.story
f5b6b1729d6088f07e7c43afbea9bc685ae054f1.story
ff0345888fafc20f762f3a74fac915a137b6a942.story<jupyter_text>Text pre-processing<jupyter_code>cd /home
!git clone https://github.com/thinkingmachines/BertSum.git
!mkdir /logs/ # Used to store log files for preprocess.py
!mkdir /home/json_data/
!mkdir {SIMPLER_JSON_PATH}
!for i in /home/story_dump/*.story; do mv "$i" /home/raw_stories/; done
!python /home/BertSum/src/preprocess.py -mode tokenize -raw_path {RAW_PATH} -save_path {TOKENIZED_PATH} -corenlp_path {CORENLP_PATH} &>/dev/null
ls -ltr {RAW_PATH}| tail -2
!tail -20 {RAW_PATH}0f848ac3922e8deaf6a803b2713ed867bc1d19e3.story
!head -20 {RAW_PATH}0f848ac3922e8deaf6a803b2713ed867bc1d19e3.story
# Notice the story starts with the article and ends with the highlighted sentences
cat {RAW_PATH}0f848ac3922e8deaf6a803b2713ed867bc1d19e3.story
ls -ltr {TOKENIZED_PATH}| tail -2
!cat {TOKENIZED_PATH}0f848ac3922e8deaf6a803b2713ed867bc1d19e3.story.json
rm /home/stanford-corenlp-full-2018-10-05.zip
ls -ltr {TOKENIZED_PATH}| tail -2
#mkdir {TOKENIZED_PATH}sample
#cp {TOKENIZED_PATH}1faf623126af4beed4ea40802bcbbdfb7922d2f1.story.json {TOKENIZED_PATH}sample
#ls {TOKENIZED_PATH}sample/*.json
ls {MAP_PATH}
!tail {MAP_PATH}mapping_valid.txt*
ls {TOKENIZED_PATH} | tail -2
!python /home/BertSum/src/preprocess.py -mode format_to_lines -raw_path {TOKENIZED_PATH} -save_path {SIMPLER_JSON_PATH} -map_path {MAP_PATH} -lower -corenlp_path {CORENLP_PATH}# &>/dev/null

ls /home/json_data/
cat /home/json_data/cnndm.train.0.json
import json

# This is a list of dictionaries where each dictionary refers to an article and has the keys 'src' and 'tgt'
# 'src' refers to the source article while 'tgt' refers to the target summaries (those shown after the @higlight tags in the original document)
# The value for both keys is a list of lists. The outer list has elements corresponding to each sentence within an article.
# The inner lists correspond to sentences where their elements correspond to indidvidual tokens found in each sentence.
test = json.load(open('/home/json_data/cnndm.train.0.json', 'r'))
len(test)
test[0]['tgt']
test[0]['src']
test[0]
mkdir {BERT_DATA_PATH}
# Colab notebook has 2 cpus (!grep -c ^processor /proc/cpuinfo)
# Convert json file to pytorch file (.pt)
# Json file is converted to a pytorch dataset object in the format that is required by the BERT module object that it will later be put into
# This includes word-piece tokenization, adding BERT tags (i.e. [CLS], [SEP]), indexing word-piece tokens,
# and deriving oracle labels (actual sentence summaries), segment Ids, [CLS] ids, source text (found in input), and target text (found in input)
!python /home/BertSum/src/preprocess.py -mode format_to_bert -raw_path /home/json_data/ -save_path {BERT_DATA_PATH} -oracle_mode greedy -n_cpus 2 -log_file /logs/preprocess.log
ls {BERT_DATA_PATH}
# Can't print insides of saved pytorch dataset
cat {BERT_DATA_PATH}cnndm.train.0.bert.pt<jupyter_output><empty_output><jupyter_text>Additional pyrouge setup<jupyter_code>cd /home
!git clone https://github.com/andersjo/pyrouge.git
cd pyrouge
ls /home/pyrouge/tools/ROUGE-1.5.5/
# Don't need to run assuming pyrogue is already installed 
#!python setup.py install
!pyrouge_set_rouge_path /home/pyrouge/tools/ROUGE-1.5.5/
# This makes sure that XML/Parser.pm is installed (required by pyrouge)
!sudo apt-get install libxml-parser-perl
cd /home/pyrouge/tools/ROUGE-1.5.5/data/
rm WordNet-2.0.exc.db
!./WordNet-2.0-Exceptions/buildExeptionDB.pl ./WordNet-2.0-Exceptions ./smart_common_words.txt ./WordNet-2.0.exc.db
!python -m pyrouge.test<jupyter_output>2019-06-02 06:05:09,469 [MainThread  ] [INFO ]  Written ROUGE configuration to /tmp/tmpque5f5v9/rouge_conf.xml
F2019-06-02 06:05:09,550 [MainThread  ] [INFO ]  Processing files in data/SL2003_models_plain_text.
2019-06-02 06:05:09,551 [MainThread  ] [INFO ]  Processing SL.P.10.R.A.SL062003-19.html.
2019-06-02 06:05:09,551 [MainThread  ] [INFO ]  Processing SL.P.10.R.A.SL062003-05.html.
2019-06-02 06:05:09,551 [MainThread  ] [INFO ]  Processing SL.P.10.R.A.SL062003-18.html.
2019-06-02 06:05:09,551 [MainThread  ] [INFO ]  Processing SL.P.10.R.A.SL062003-06.html.
2019-06-02 06:05:09,551 [MainThread  ] [INFO ]  Processing SL.P.10.R.A.SL062003-16.html.
2019-06-02 06:05:09,551 [MainThread  ] [INFO ]  Processing SL.P.10.R.A.SL062003-20.html.
2019-06-02 06:05:09,552 [MainThread  ] [INFO ]  Processing SL.P.10.R.A.SL062003-02.html.
2019-06-02 06:05:09,552 [MainThread  ] [INFO ]  Processing SL.P.10.R.A.SL062003-22.html.
2019-06-02 06:05:09,552 [MainThread  ] [INFO ]  Processing SL.P.10.R.A.SL0620[...]<jupyter_text>BERTSUM Inference  Load model and sample data from GCS<jupyter_code>def dl_gcs(object_fp, dest_fp, bucket_name='nlp-experiment-datasets', project_id='bert-experiments'):
    from googleapiclient.discovery import build
    gcs_service = build('storage', 'v1')
    from apiclient.http import MediaIoBaseDownload
    with open(dest_fp, 'wb') as f:
        # Download the file from a given Google Cloud Storage bucket.
        request = gcs_service.objects().get_media(bucket=bucket_name,
                                            object=f'{object_fp}')
        media = MediaIoBaseDownload(f, request)

        done = False
        while not done:
    # _ is a placeholder for a progress object that we ignore.
    # (Our file is small, so we skip reporting progress.)
            _, done = media.next_chunk()        
    print('Download complete')
    
MODELS_DIR = '/home/models/'
!mkdir {MODELS_DIR}
MODEL = 'model_step_43000.pt'
MODEL_FP = f'{MODELS_DIR}{MODEL}'
SAMPLE_DATA_FP = '/home/sample_data/cnndm.test.0.bert.pt'
!mkdir /home/sample_data/ # Source of data
!mkdir /home/bert_results/ # Destination of results
!mkdir /home/temp/ 
!mkdir /home/logs/
# Download trained model
dl_gcs('bertsum/cnn_dailymail/models/model_step_43000.pt', MODEL_FP)
# Download sample test data
dl_gcs('bertsum/cnn_dailymail/sample_data/cnndm.test.0.bert.pt', SAMPLE_DATA_FP)
ls {MODELS_DIR}
ls /home/sample_data/
cat /home/sample_data/cnndm.test.0.bert.pt<jupyter_output><empty_output><jupyter_text>Modules<jupyter_code># From src/models/optimizers.py
""" Optimizers class """
import torch
import torch.optim as optim
from torch.nn.utils import clip_grad_norm_


# from onmt.utils import use_gpu


def use_gpu(opt):
    """
    Creates a boolean if gpu used
    """
    return (hasattr(opt, 'gpu_ranks') and len(opt.gpu_ranks) > 0) or \
           (hasattr(opt, 'gpu') and opt.gpu > -1)

def build_optim(model, opt, checkpoint):
    """ Build optimizer """
    saved_optimizer_state_dict = None

    if opt.train_from:
        optim = checkpoint['optim']
        # We need to save a copy of optim.optimizer.state_dict() for setting
        # the, optimizer state later on in Stage 2 in this method, since
        # the method optim.set_parameters(model.parameters()) will overwrite
        # optim.optimizer, and with ith the values stored in
        # optim.optimizer.state_dict()
        saved_optimizer_state_dict = optim.optimizer.state_dict()
    else:
        optim = Optimizer(
            opt.optim, opt.learning_rate, opt.max_grad_norm,
            lr_decay=opt.learning_rate_decay,
            start_decay_steps=opt.start_decay_steps,
            decay_steps=opt.decay_steps,
            beta1=opt.adam_beta1,
            beta2=opt.adam_beta2,
            adagrad_accum=opt.adagrad_accumulator_init,
            decay_method=opt.decay_method,
            warmup_steps=opt.warmup_steps)

    # Stage 1:
    # Essentially optim.set_parameters (re-)creates and optimizer using
    # model.paramters() as parameters that will be stored in the
    # optim.optimizer.param_groups field of the torch optimizer class.
    # Importantly, this method does not yet load the optimizer state, as
    # essentially it builds a new optimizer with empty optimizer state and
    # parameters from the model.
    optim.set_parameters(model.named_parameters())

    if opt.train_from:
        # Stage 2: In this stage, which is only performed when loading an
        # optimizer from a checkpoint, we load the saved_optimizer_state_dict
        # into the re-created optimizer, to set the optim.optimizer.state
        # field, which was previously empty. For this, we use the optimizer
        # state saved in the "saved_optimizer_state_dict" variable for
        # this purpose.
        # See also: https://github.com/pytorch/pytorch/issues/2830
        optim.optimizer.load_state_dict(saved_optimizer_state_dict)
        # Convert back the state values to cuda type if applicable
        if use_gpu(opt):
            for state in optim.optimizer.state.values():
                for k, v in state.items():
                    if torch.is_tensor(v):
                        state[k] = v.cuda()

        # We want to make sure that indeed we have a non-empty optimizer state
        # when we loaded an existing model. This should be at least the case
        # for Adam, which saves "exp_avg" and "exp_avg_sq" state
        # (Exponential moving average of gradient and squared gradient values)
        if (optim.method == 'adam') and (len(optim.optimizer.state) < 1):
            raise RuntimeError(
                "Error: loaded Adam optimizer from existing model" +
                " but optimizer state is empty")

    return optim


class MultipleOptimizer(object):
    """ Implement multiple optimizers needed for sparse adam """

    def __init__(self, op):
        """ ? """
        self.optimizers = op

    def zero_grad(self):
        """ ? """
        for op in self.optimizers:
            op.zero_grad()

    def step(self):
        """ ? """
        for op in self.optimizers:
            op.step()

    @property
    def state(self):
        """ ? """
        return {k: v for op in self.optimizers for k, v in op.state.items()}

    def state_dict(self):
        """ ? """
        return [op.state_dict() for op in self.optimizers]

    def load_state_dict(self, state_dicts):
        """ ? """
        assert len(state_dicts) == len(self.optimizers)
        for i in range(len(state_dicts)):
            self.optimizers[i].load_state_dict(state_dicts[i])


class Optimizer(object):
    """
    Controller class for optimization. Mostly a thin
    wrapper for `optim`, but also useful for implementing
    rate scheduling beyond what is currently available.
    Also implements necessary methods for training RNNs such
    as grad manipulations.
    Args:
      method (:obj:`str`): one of [sgd, adagrad, adadelta, adam]
      lr (float): learning rate
      lr_decay (float, optional): learning rate decay multiplier
      start_decay_steps (int, optional): step to start learning rate decay
      beta1, beta2 (float, optional): parameters for adam
      adagrad_accum (float, optional): initialization parameter for adagrad
      decay_method (str, option): custom decay options
      warmup_steps (int, option): parameter for `noam` decay
    We use the default parameters for Adam that are suggested by
    the original paper https://arxiv.org/pdf/1412.6980.pdf
    These values are also used by other established implementations,
    e.g. https://www.tensorflow.org/api_docs/python/tf/train/AdamOptimizer
    https://keras.io/optimizers/
    Recently there are slightly different values used in the paper
    "Attention is all you need"
    https://arxiv.org/pdf/1706.03762.pdf, particularly the value beta2=0.98
    was used there however, beta2=0.999 is still arguably the more
    established value, so we use that here as well
    """

    def __init__(self, method, learning_rate, max_grad_norm,
                 lr_decay=1, start_decay_steps=None, decay_steps=None,
                 beta1=0.9, beta2=0.999,
                 adagrad_accum=0.0,
                 decay_method=None,
                 warmup_steps=4000
                 ):
        self.last_ppl = None
        self.learning_rate = learning_rate
        self.original_lr = learning_rate
        self.max_grad_norm = max_grad_norm
        self.method = method
        self.lr_decay = lr_decay
        self.start_decay_steps = start_decay_steps
        self.decay_steps = decay_steps
        self.start_decay = False
        self._step = 0
        self.betas = [beta1, beta2]
        self.adagrad_accum = adagrad_accum
        self.decay_method = decay_method
        self.warmup_steps = warmup_steps

    def set_parameters(self, params):
        """ ? """
        self.params = []
        self.sparse_params = []
        for k, p in params:
            if p.requires_grad:
                if self.method != 'sparseadam' or "embed" not in k:
                    self.params.append(p)
                else:
                    self.sparse_params.append(p)
        if self.method == 'sgd':
            self.optimizer = optim.SGD(self.params, lr=self.learning_rate)
        elif self.method == 'adagrad':
            self.optimizer = optim.Adagrad(self.params, lr=self.learning_rate)
            for group in self.optimizer.param_groups:
                for p in group['params']:
                    self.optimizer.state[p]['sum'] = self.optimizer\
                        .state[p]['sum'].fill_(self.adagrad_accum)
        elif self.method == 'adadelta':
            self.optimizer = optim.Adadelta(self.params, lr=self.learning_rate)
        elif self.method == 'adam':
            self.optimizer = optim.Adam(self.params, lr=self.learning_rate,
                                        betas=self.betas, eps=1e-9)
        elif self.method == 'sparseadam':
            self.optimizer = MultipleOptimizer(
                [optim.Adam(self.params, lr=self.learning_rate,
                            betas=self.betas, eps=1e-8),
                 optim.SparseAdam(self.sparse_params, lr=self.learning_rate,
                                  betas=self.betas, eps=1e-8)])
        else:
            raise RuntimeError("Invalid optim method: " + self.method)

    def _set_rate(self, learning_rate):
        self.learning_rate = learning_rate
        if self.method != 'sparseadam':
            self.optimizer.param_groups[0]['lr'] = self.learning_rate
        else:
            for op in self.optimizer.optimizers:
                op.param_groups[0]['lr'] = self.learning_rate

    def step(self):
        """Update the model parameters based on current gradients.
        Optionally, will employ gradient modification or update learning
        rate.
        """
        self._step += 1

        # Decay method used in tensor2tensor.
        if self.decay_method == "noam":
            self._set_rate(
                self.original_lr *

                 min(self._step ** (-0.5),
                     self._step * self.warmup_steps**(-1.5)))

            # self._set_rate(self.original_lr *self.model_size ** (-0.5) *min(1.0, self._step / self.warmup_steps)*max(self._step, self.warmup_steps)**(-0.5))
        # Decay based on start_decay_steps every decay_steps
        else:
            if ((self.start_decay_steps is not None) and (
                     self._step >= self.start_decay_steps)):
                self.start_decay = True
            if self.start_decay:
                if ((self._step - self.start_decay_steps)
                   % self.decay_steps == 0):
                    self.learning_rate = self.learning_rate * self.lr_decay

        if self.method != 'sparseadam':
            self.optimizer.param_groups[0]['lr'] = self.learning_rate

        if self.max_grad_norm:
            clip_grad_norm_(self.params, self.max_grad_norm)
        self.optimizer.step()

# From src/models/neural.py
import math

import torch
import torch.nn as nn


def gelu(x):
    return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))


class PositionwiseFeedForward(nn.Module):
    """ A two-layer Feed-Forward-Network with residual layer norm.
    Args:
        d_model (int): the size of input for the first-layer of the FFN.
        d_ff (int): the hidden layer size of the second-layer
            of the FNN.
        dropout (float): dropout probability in :math:`[0, 1)`.
    """

    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.actv = gelu
        self.dropout_1 = nn.Dropout(dropout)
        self.dropout_2 = nn.Dropout(dropout)

    def forward(self, x):
        inter = self.dropout_1(self.actv(self.w_1(self.layer_norm(x))))
        output = self.dropout_2(self.w_2(inter))
        return output + x


class MultiHeadedAttention(nn.Module):
    """
    Multi-Head Attention module from
    "Attention is All You Need"
    :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`.
    Similar to standard `dot` attention but uses
    multiple attention distributions simulataneously
    to select relevant items.
    .. mermaid::
       graph BT
          A[key]
          B[value]
          C[query]
          O[output]
          subgraph Attn
            D[Attn 1]
            E[Attn 2]
            F[Attn N]
          end
          A --> D
          C --> D
          A --> E
          C --> E
          A --> F
          C --> F
          D --> O
          E --> O
          F --> O
          B --> O
    Also includes several additional tricks.
    Args:
       head_count (int): number of parallel heads
       model_dim (int): the dimension of keys/values/queries,
           must be divisible by head_count
       dropout (float): dropout parameter
    """

    def __init__(self, head_count, model_dim, dropout=0.1, use_final_linear=True):
        assert model_dim % head_count == 0
        self.dim_per_head = model_dim // head_count
        self.model_dim = model_dim

        super(MultiHeadedAttention, self).__init__()
        self.head_count = head_count

        self.linear_keys = nn.Linear(model_dim,
                                     head_count * self.dim_per_head)
        self.linear_values = nn.Linear(model_dim,
                                       head_count * self.dim_per_head)
        self.linear_query = nn.Linear(model_dim,
                                      head_count * self.dim_per_head)
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)
        self.use_final_linear = use_final_linear
        if (self.use_final_linear):
            self.final_linear = nn.Linear(model_dim, model_dim)

    def forward(self, key, value, query, mask=None,
                layer_cache=None, type=None, predefined_graph_1=None):
        """
        Compute the context vector and the attention vectors.
        Args:
           key (`FloatTensor`): set of `key_len`
                key vectors `[batch, key_len, dim]`
           value (`FloatTensor`): set of `key_len`
                value vectors `[batch, key_len, dim]`
           query (`FloatTensor`): set of `query_len`
                 query vectors  `[batch, query_len, dim]`
           mask: binary mask indicating which keys have
                 non-zero attention `[batch, query_len, key_len]`
        Returns:
           (`FloatTensor`, `FloatTensor`) :
           * output context vectors `[batch, query_len, dim]`
           * one of the attention vectors `[batch, query_len, key_len]`
        """

        # CHECKS
        # batch, k_len, d = key.size()
        # batch_, k_len_, d_ = value.size()
        # aeq(batch, batch_)
        # aeq(k_len, k_len_)
        # aeq(d, d_)
        # batch_, q_len, d_ = query.size()
        # aeq(batch, batch_)
        # aeq(d, d_)
        # aeq(self.model_dim % 8, 0)
        # if mask is not None:
        #    batch_, q_len_, k_len_ = mask.size()
        #    aeq(batch_, batch)
        #    aeq(k_len_, k_len)
        #    aeq(q_len_ == q_len)
        # END CHECKS

        batch_size = key.size(0)
        dim_per_head = self.dim_per_head
        head_count = self.head_count
        key_len = key.size(1)
        query_len = query.size(1)

        def shape(x):
            """  projection """
            return x.view(batch_size, -1, head_count, dim_per_head) \
                .transpose(1, 2)

        def unshape(x):
            """  compute context """
            return x.transpose(1, 2).contiguous() \
                .view(batch_size, -1, head_count * dim_per_head)

        # 1) Project key, value, and query.
        if layer_cache is not None:
            if type == "self":
                query, key, value = self.linear_query(query), \
                                    self.linear_keys(query), \
                                    self.linear_values(query)

                key = shape(key)
                value = shape(value)

                if layer_cache is not None:
                    device = key.device
                    if layer_cache["self_keys"] is not None:
                        key = torch.cat(
                            (layer_cache["self_keys"].to(device), key),
                            dim=2)
                    if layer_cache["self_values"] is not None:
                        value = torch.cat(
                            (layer_cache["self_values"].to(device), value),
                            dim=2)
                    layer_cache["self_keys"] = key
                    layer_cache["self_values"] = value
            elif type == "context":
                query = self.linear_query(query)
                if layer_cache is not None:
                    if layer_cache["memory_keys"] is None:
                        key, value = self.linear_keys(key), \
                                     self.linear_values(value)
                        key = shape(key)
                        value = shape(value)
                    else:
                        key, value = layer_cache["memory_keys"], \
                                     layer_cache["memory_values"]
                    layer_cache["memory_keys"] = key
                    layer_cache["memory_values"] = value
                else:
                    key, value = self.linear_keys(key), \
                                 self.linear_values(value)
                    key = shape(key)
                    value = shape(value)
        else:
            key = self.linear_keys(key)
            value = self.linear_values(value)
            query = self.linear_query(query)
            key = shape(key)
            value = shape(value)

        query = shape(query)

        key_len = key.size(2)
        query_len = query.size(2)

        # 2) Calculate and scale scores.
        query = query / math.sqrt(dim_per_head)
        scores = torch.matmul(query, key.transpose(2, 3))

        if mask is not None:
            mask = mask.unsqueeze(1).expand_as(scores)
            scores = scores.masked_fill(mask, -1e18)

        # 3) Apply attention dropout and compute context vectors.

        attn = self.softmax(scores)

        if (not predefined_graph_1 is None):
            attn_masked = attn[:, -1] * predefined_graph_1
            attn_masked = attn_masked / (torch.sum(attn_masked, 2).unsqueeze(2) + 1e-9)

            attn = torch.cat([attn[:, :-1], attn_masked.unsqueeze(1)], 1)

        drop_attn = self.dropout(attn)
        if (self.use_final_linear):
            context = unshape(torch.matmul(drop_attn, value))
            output = self.final_linear(context)
            return output
        else:
            context = torch.matmul(drop_attn, value)
            return context

        # CHECK
        # batch_, q_len_, d_ = output.size()
        # aeq(q_len, q_len_)
        # aeq(batch, batch_)
        # aeq(d, d_)

        # Return one attn
# From /src/models/rnn.py
import torch
import torch.nn.functional as F
from torch import nn


class LayerNormLSTMCell(nn.LSTMCell):

    def __init__(self, input_size, hidden_size, bias=True):
        super().__init__(input_size, hidden_size, bias)

        self.ln_ih = nn.LayerNorm(4 * hidden_size)
        self.ln_hh = nn.LayerNorm(4 * hidden_size)
        self.ln_ho = nn.LayerNorm(hidden_size)

    def forward(self, input, hidden=None):
        self.check_forward_input(input)
        if hidden is None:
            hx = input.new_zeros(input.size(0), self.hidden_size, requires_grad=False)
            cx = input.new_zeros(input.size(0), self.hidden_size, requires_grad=False)
        else:
            hx, cx = hidden
        self.check_forward_hidden(input, hx, '[0]')
        self.check_forward_hidden(input, cx, '[1]')

        gates = self.ln_ih(F.linear(input, self.weight_ih, self.bias_ih)) \
                + self.ln_hh(F.linear(hx, self.weight_hh, self.bias_hh))
        i, f, o = gates[:, :(3 * self.hidden_size)].sigmoid().chunk(3, 1)
        g = gates[:, (3 * self.hidden_size):].tanh()

        cy = (f * cx) + (i * g)
        hy = o * self.ln_ho(cy).tanh()
        return hy, cy


class LayerNormLSTM(nn.Module):

    def __init__(self, input_size, hidden_size, num_layers=1, bias=True, bidirectional=False):
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.bidirectional = bidirectional

        num_directions = 2 if bidirectional else 1
        self.hidden0 = nn.ModuleList([
            LayerNormLSTMCell(input_size=(input_size if layer == 0 else hidden_size * num_directions),
                              hidden_size=hidden_size, bias=bias)
            for layer in range(num_layers)
        ])

        if self.bidirectional:
            self.hidden1 = nn.ModuleList([
                LayerNormLSTMCell(input_size=(input_size if layer == 0 else hidden_size * num_directions),
                                  hidden_size=hidden_size, bias=bias)
                for layer in range(num_layers)
            ])

    def forward(self, input, hidden=None):
        seq_len, batch_size, hidden_size = input.size()  # supports TxNxH only
        num_directions = 2 if self.bidirectional else 1
        if hidden is None:
            hx = input.new_zeros(self.num_layers * num_directions, batch_size, self.hidden_size, requires_grad=False)
            cx = input.new_zeros(self.num_layers * num_directions, batch_size, self.hidden_size, requires_grad=False)
        else:
            hx, cx = hidden

        ht = [[None, ] * (self.num_layers * num_directions)] * seq_len
        ct = [[None, ] * (self.num_layers * num_directions)] * seq_len

        if self.bidirectional:
            xs = input
            for l, (layer0, layer1) in enumerate(zip(self.hidden0, self.hidden1)):
                l0, l1 = 2 * l, 2 * l + 1
                h0, c0, h1, c1 = hx[l0], cx[l0], hx[l1], cx[l1]
                for t, (x0, x1) in enumerate(zip(xs, reversed(xs))):
                    ht[t][l0], ct[t][l0] = layer0(x0, (h0, c0))
                    h0, c0 = ht[t][l0], ct[t][l0]
                    t = seq_len - 1 - t
                    ht[t][l1], ct[t][l1] = layer1(x1, (h1, c1))
                    h1, c1 = ht[t][l1], ct[t][l1]
                xs = [torch.cat((h[l0], h[l1]), dim=1) for h in ht]
            y = torch.stack(xs)
            hy = torch.stack(ht[-1])
            cy = torch.stack(ct[-1])
        else:
            h, c = hx, cx
            for t, x in enumerate(input):
                for l, layer in enumerate(self.hidden0):
                    ht[t][l], ct[t][l] = layer(x, (h[l], c[l]))
                    x = ht[t][l]
                h, c = ht[t], ct[t]
            y = torch.stack([h[-1] for h in ht])
            hy = torch.stack(ht[-1])
            cy = torch.stack(ct[-1])

        return y, (hy, cy)
# From /src/models/encoder.py
import math

import torch
import torch.nn as nn

#from models.neural import MultiHeadedAttention, PositionwiseFeedForward
#from models.rnn import LayerNormLSTM


class Classifier(nn.Module):
    def __init__(self, hidden_size):
        super(Classifier, self).__init__()
        self.linear1 = nn.Linear(hidden_size, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x, mask_cls):
        h = self.linear1(x).squeeze(-1)
        sent_scores = self.sigmoid(h) * mask_cls.float()
        return sent_scores


class PositionalEncoding(nn.Module):

    def __init__(self, dropout, dim, max_len=5000):
        pe = torch.zeros(max_len, dim)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp((torch.arange(0, dim, 2, dtype=torch.float) *
                              -(math.log(10000.0) / dim)))
        pe[:, 0::2] = torch.sin(position.float() * div_term)
        pe[:, 1::2] = torch.cos(position.float() * div_term)
        pe = pe.unsqueeze(0)
        super(PositionalEncoding, self).__init__()
        self.register_buffer('pe', pe)
        self.dropout = nn.Dropout(p=dropout)
        self.dim = dim

    def forward(self, emb, step=None):
        emb = emb * math.sqrt(self.dim)
        if (step):
            emb = emb + self.pe[:, step][:, None, :]

        else:
            emb = emb + self.pe[:, :emb.size(1)]
        emb = self.dropout(emb)
        return emb

    def get_emb(self, emb):
        return self.pe[:, :emb.size(1)]


class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, heads, d_ff, dropout):
        super(TransformerEncoderLayer, self).__init__()

        self.self_attn = MultiHeadedAttention(
            heads, d_model, dropout=dropout)
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.dropout = nn.Dropout(dropout)

    def forward(self, iter, query, inputs, mask):
        if (iter != 0):
            input_norm = self.layer_norm(inputs)
        else:
            input_norm = inputs

        mask = mask.unsqueeze(1)
        context = self.self_attn(input_norm, input_norm, input_norm,
                                 mask=mask)
        out = self.dropout(context) + inputs
        return self.feed_forward(out)


class TransformerInterEncoder(nn.Module):
    def __init__(self, d_model, d_ff, heads, dropout, num_inter_layers=0):
        super(TransformerInterEncoder, self).__init__()
        self.d_model = d_model
        self.num_inter_layers = num_inter_layers
        self.pos_emb = PositionalEncoding(dropout, d_model)
        self.transformer_inter = nn.ModuleList(
            [TransformerEncoderLayer(d_model, heads, d_ff, dropout)
             for _ in range(num_inter_layers)])
        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.wo = nn.Linear(d_model, 1, bias=True)
        self.sigmoid = nn.Sigmoid()

    def forward(self, top_vecs, mask):
        """ See :obj:`EncoderBase.forward()`"""

        batch_size, n_sents = top_vecs.size(0), top_vecs.size(1)
        pos_emb = self.pos_emb.pe[:, :n_sents]
        x = top_vecs * mask[:, :, None].float()
        x = x + pos_emb

        for i in range(self.num_inter_layers):
            x = self.transformer_inter[i](i, x, x, 1 - mask)  # all_sents * max_tokens * dim

        x = self.layer_norm(x)
        sent_scores = self.sigmoid(self.wo(x))
        sent_scores = sent_scores.squeeze(-1) * mask.float()

        return sent_scores


class RNNEncoder(nn.Module):

    def __init__(self, bidirectional, num_layers, input_size,
                 hidden_size, dropout=0.0):
        super(RNNEncoder, self).__init__()
        num_directions = 2 if bidirectional else 1
        assert hidden_size % num_directions == 0
        hidden_size = hidden_size // num_directions

        self.rnn = LayerNormLSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            bidirectional=bidirectional)

        self.wo = nn.Linear(num_directions * hidden_size, 1, bias=True)
        self.dropout = nn.Dropout(dropout)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x, mask):
        """See :func:`EncoderBase.forward()`"""
        x = torch.transpose(x, 1, 0)
        memory_bank, _ = self.rnn(x)
        memory_bank = self.dropout(memory_bank) + x
        memory_bank = torch.transpose(memory_bank, 1, 0)

        sent_scores = self.sigmoid(self.wo(memory_bank))
        sent_scores = sent_scores.squeeze(-1) * mask.float()
        return sent_scores
# From src/models/model_builder.py
import torch
import torch.nn as nn
from pytorch_pretrained_bert import BertModel, BertConfig
from torch.nn.init import xavier_uniform_

#from models.encoder import TransformerInterEncoder, Classifier, RNNEncoder
#from models.optimizers import Optimizer


def build_optim(args, model, checkpoint):
    """ Build optimizer """
    saved_optimizer_state_dict = None

    if args.train_from != '':
        optim = checkpoint['optim']
        saved_optimizer_state_dict = optim.optimizer.state_dict()
    else:
        optim = Optimizer(
            args.optim, args.lr, args.max_grad_norm,
            beta1=args.beta1, beta2=args.beta2,
            decay_method=args.decay_method,
            warmup_steps=args.warmup_steps)

    optim.set_parameters(list(model.named_parameters()))

    if args.train_from != '':
        optim.optimizer.load_state_dict(saved_optimizer_state_dict)
        if args.visible_gpus != '-1':
            for state in optim.optimizer.state.values():
                for k, v in state.items():
                    if torch.is_tensor(v):
                        state[k] = v.cuda()

        if (optim.method == 'adam') and (len(optim.optimizer.state) < 1):
            raise RuntimeError(
                "Error: loaded Adam optimizer from existing model" +
                " but optimizer state is empty")

    return optim


class Bert(nn.Module):
    def __init__(self, temp_dir, load_pretrained_bert, bert_config):
        super(Bert, self).__init__()
        if(load_pretrained_bert):
            self.model = BertModel.from_pretrained('bert-base-uncased', cache_dir=temp_dir)
        else:
            self.model = BertModel(bert_config)

    def forward(self, x, segs, mask):
        encoded_layers, _ = self.model(x, segs, attention_mask =mask)
        top_vec = encoded_layers[-1]
        return top_vec



class Summarizer(nn.Module):
    def __init__(self, args, device, load_pretrained_bert = False, bert_config = None):
        super(Summarizer, self).__init__()
        self.args = args
        self.device = device
        self.bert = Bert(args.temp_dir, load_pretrained_bert, bert_config)
        if (args.encoder == 'classifier'):
            self.encoder = Classifier(self.bert.model.config.hidden_size)
        elif(args.encoder=='transformer'):
            self.encoder = TransformerInterEncoder(self.bert.model.config.hidden_size, args.ff_size, args.heads,
                                                   args.dropout, args.inter_layers)
        elif(args.encoder=='rnn'):
            self.encoder = RNNEncoder(bidirectional=True, num_layers=1,
                                      input_size=self.bert.model.config.hidden_size, hidden_size=args.rnn_size,
                                      dropout=args.dropout)
        elif (args.encoder == 'baseline'):
            bert_config = BertConfig(self.bert.model.config.vocab_size, hidden_size=args.hidden_size,
                                     num_hidden_layers=6, num_attention_heads=8, intermediate_size=args.ff_size)
            self.bert.model = BertModel(bert_config)
            self.encoder = Classifier(self.bert.model.config.hidden_size)

        if args.param_init != 0.0:
            for p in self.encoder.parameters():
                p.data.uniform_(-args.param_init, args.param_init)
        if args.param_init_glorot:
            for p in self.encoder.parameters():
                if p.dim() > 1:
                    xavier_uniform_(p)

        self.to(device)
    def load_cp(self, pt):
        self.load_state_dict(pt['model'], strict=True)

    def forward(self, x, segs, clss, mask, mask_cls, sentence_range=None):

        top_vec = self.bert(x, segs, mask)
        sents_vec = top_vec[torch.arange(top_vec.size(0)).unsqueeze(1), clss]
        sents_vec = sents_vec * mask_cls[:, :, None].float()
        sent_scores = self.encoder(sents_vec, mask_cls).squeeze(-1)
        return sent_scores, mask_cls<jupyter_output><empty_output><jupyter_text>Inference on sample dataset<jupyter_code>import torch
from pytorch_pretrained_bert import BertConfig
from argparse import Namespace
# Setting these as default arguments
# These attributes can be edited as needed
args = Namespace(
    encoder="transformer",
    mode="test",
    bert_data_path="/home/sample_data/cnndm",
    model_path=MODELS_DIR,
    result_path="/home/bert_results/cnndm",
    temp_dir="/home/temp/",
    bert_config_path='/home/BertSum/bert_config_uncased_base.json',
    batch_size=1000,
    model_fp=MODEL_FP,

    use_interval=True,
    hidden_size=128,
    ff_size=2048, # Size used during training
    heads=4,
    inter_layers=2,
    rnn_size=512,

    param_init=0.0,
    param_init_glorot=True,
    dropout=0.1,
    optim='adam',
    lr=1,
    beta1= 0.9,
    beta2=0.999,
    decay_method='',
    warmup_steps=8000,
    max_grad_norm=0,

    save_checkpoint_steps=5,
    accum_count=1,
    world_size=1,
    report_every=1,
    train_steps=1000,
    recall_eval=False,


    visible_gpus='-1',
    gpu_ranks='0',
    log_file='/home/logs/cnndm.log',
    dataset='',
    seed=666,

    test_all=False,
    test_from=MODEL,
    train_from='',
    report_rouge=True,
    block_trigram=True,
)
# Have to be in this directory for the pickle file loading to import BERTSUM in the same environment that it was trained in
cd /home/BertSum/src
# Load pretrained BERTSUM model
trained_model = torch.load(MODEL_FP, map_location=lambda storage, loc: storage)
# Read BERT config file
# This contains information about the BERT model (e.g. hidden size for the transformer layers, number of transformer layers, number of self attention heads)
config = BertConfig.from_json_file(args.bert_config_path)
# Specify device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Instantiate Summarizer and load pretrained model
model = Summarizer(args, device, load_pretrained_bert=False, bert_config = config)
model.load_cp(trained_model)
model
#test_function
def test_api(self, test_iter, step, cal_lead=False, cal_oracle=False):
    """ Validate model.
        valid_iter: validate data iterator
    Returns:
        :obj:`nmt.Statistics`: validation loss statistics
    """
    # Set model in validating mode.
    def _get_ngrams(n, text):
        ngram_set = set()
        text_length = len(text)
        max_index_ngram_start = text_length - n
        for i in range(max_index_ngram_start + 1):
            ngram_set.add(tuple(text[i:i + n]))
        return ngram_set

    def _block_tri(c, p):
        tri_c = _get_ngrams(3, c.split())
        for s in p:
            tri_s = _get_ngrams(3, s.split())
            if len(tri_c.intersection(tri_s))>0:
                return True
        return False

    if (not cal_lead and not cal_oracle):
        # Evaluate without performing backpropagation and dropout
        self.model.eval()
    _source_article = []
    with torch.no_grad():
        for batch in test_iter:
            src = batch.src
            labels = batch.labels
            segs = batch.segs
            clss = batch.clss
            mask = batch.mask
            mask_cls = batch.mask_cls

            _source_article += src
            pred = []

            if (cal_lead):
                selected_ids = [list(range(batch.clss.size(1)))] * batch.batch_size
            elif (cal_oracle):
                selected_ids = [[j for j in range(batch.clss.size(1)) if labels[i][j] == 1] for i in
                                range(batch.batch_size)]
            else:
                sent_scores, mask = self.model(src, segs, clss, mask, mask_cls)

                sent_scores = sent_scores + mask.float()
                sent_scores = sent_scores.cpu().data.numpy()
                selected_ids = np.argsort(-sent_scores, 1)
            # selected_ids = np.sort(selected_ids,1)
            for i, idx in enumerate(selected_ids):
                _pred = []
                if(len(batch.src_str[i])==0):
                    continue
                for j in selected_ids[i][:len(batch.src_str[i])]:
                    if(j>=len( batch.src_str[i])):
                        continue
                    candidate = batch.src_str[i][j].strip()
                    if(self.args.block_trigram):
                        if(not _block_tri(candidate,_pred)):
                            _pred.append(candidate)
                    else:
                        _pred.append(candidate)

                    if ((not cal_oracle) and (not self.args.recall_eval) and len(_pred) == 3):
                        break

                _pred = '<q>'.join(_pred)
                if(self.args.recall_eval):
                    _pred = ' '.join(_pred.split()[:len(batch.tgt_str[i].split())])

                pred.append(_pred)

    results = {'source_article': _source_article, 'predicted_summary': _pred}
    return results
from models import data_loader, model_builder
from models.data_loader import load_dataset
from models.trainer import build_trainer
cp = args.test_from
step = int(cp.split('.')[-2].split('_')[-1])
device_id = 0 if device == "cuda" else -1
# Read data
test_iter = data_loader.Dataloader(args, load_dataset(args, 'test', shuffle=False),
                              args.batch_size, device,
                              shuffle=False, is_test=True)
# sample is an iterator that contains a batch of data inputs to the "trainer"
# For our context, this will be for the purpose of doing a forward pass with unseen articles
sample = next(iter(test_iter))
# You'll notice the initial tokenization is performed on the article while separating it into a list of sentences
sample.src_str[0]
# Notice that the target text is abstractive, meaning it's not exactly extracted from the source document
sample.tgt_str
trainer = build_trainer(args, device_id, model, None)
trainer.test(test_iter, step)
ls /home/bert_results
# Predicted summary
!head -1 /home/bert_results/cnndm_step43000.candidate
# Ground truth summary
!head -1 /home/bert_results/cnndm_step43000.gold
# First 10 predicted summaries
!head -10 /home/bert_results/cnndm_step43000.candidate
# First 10 ground truth summaries
!head -10 /home/bert_results/cnndm_step43000.gold<jupyter_output><NAME> posted video on his instagram of the cast being removed<q>the austrian is currently out after suffering knee ligament damage<q>bayern munich face porto in the champions league on wednesday
sabrina osterkamp was last seen leaving home in naracoorte at 12pm on sunday<q>the german tourist was driving a friend 's car about 110km south to mount gambier in south australia 's southeast<q>the car was found abandoned on side of road in centre of mount gambier<q>the 25-year-old contacted authorities to say she was safe and well almost 24 hours later
<NAME> vandalised and spat on the vehicle in 2013<q>the forensic officer swabbed the spit and found philp from his dna<q>he pleaded guilty to wilful damage of police property and was fined $ 300<q><NAME> claimed it was out of character but he has similar convictions
national union of teachers expected to call for a vote on industrial action<q>ballot of its 300,000 members would take place after general election<q>fears fun[...]<jupyter_text>BERTSUM - Low level implementation  Prepare sample data<jupyter_code><jupyter_output><empty_output><jupyter_text>Text pre-processing<jupyter_code><jupyter_output><empty_output><jupyter_text>BERTSUM inference  Load model from GCS<jupyter_code><jupyter_output><empty_output><jupyter_text>Inference on sample dataset<jupyter_code>import torch
from pytorch_pretrained_bert import BertConfig
from argparse import Namespace
# Setting these as default arguments
# These attributes can be edited as needed
args = Namespace(
    encoder="transformer",
    mode="test",
    bert_data_path="/home/sample_data/cnndm",
    model_path=MODELS_DIR,
    result_path="/home/bert_results/cnndm",
    temp_dir="/home/temp/",
    bert_config_path='/home/BertSum/bert_config_uncased_base.json',
    batch_size=1000,

    use_interval=True,
    hidden_size=128,
    ff_size=2048, # Size used during training
    heads=4,
    inter_layers=2,
    rnn_size=512,

    param_init=0.0,
    param_init_glorot=True,
    dropout=0.1,
    optim='adam',
    lr=1,
    beta1= 0.9,
    beta2=0.999,
    decay_method='',
    warmup_steps=8000,
    max_grad_norm=0,

    save_checkpoint_steps=5,
    accum_count=1,
    world_size=1,
    report_every=1,
    train_steps=1000,
    recall_eval=False,


    visible_gpus='-1',
    gpu_ranks='0',
    log_file='/home/logs/cnndm.log',
    dataset='',
    seed=666,

    test_all=False,
    test_from=MODEL,
    train_from='',
    report_rouge=True,
    block_trigram=True,
)
# Have to be in this directory for the pickle file loading to import BERTSUM in the same environment that it was trained in
cd /home/BertSum/src
# Load pretrained BERTSUM model
trained_model = torch.load(MODEL_FP, map_location=lambda storage, loc: storage)
# Read BERT config file
# This contains information about the BERT model (e.g. hidden size for the transformer layers, number of transformer layers, number of self attention heads)
config = BertConfig.from_json_file(args.bert_config_path)
# Specify device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Instantiate Summarizer and load pretrained model
model = Summarizer(args, device, load_pretrained_bert=False, bert_config = config)
model.load_cp(trained_model)
model
from models import data_loader, model_builder
from models.data_loader import load_dataset
from models.trainer import build_trainer, Trainer

def test_api(self, test_iter, step, top_n_sentences=3, cal_lead=False, cal_oracle=False):
    """ Validate model.
        valid_iter: validate data iterator
    Returns:
        :obj:`nmt.Statistics`: validation loss statistics
    """
    # Set model in validating mode.
    def _get_ngrams(n, text):
        ngram_set = set()
        text_length = len(text)
        max_index_ngram_start = text_length - n
        for i in range(max_index_ngram_start + 1):
            ngram_set.add(tuple(text[i:i + n]))
        return ngram_set

    def _block_tri(c, p):
        tri_c = _get_ngrams(3, c.split())
        for s in p:
            tri_s = _get_ngrams(3, s.split())
            if len(tri_c.intersection(tri_s))>0:
                return True
        return False

    if (not cal_lead and not cal_oracle):
        # Evaluate without performing backpropagation and dropout
        self.model.eval()
    source_article = []
    pred = []
    with torch.no_grad():
        for batch in test_iter:
            src = batch.src
            labels = batch.labels
            segs = batch.segs
            clss = batch.clss
            mask = batch.mask
            mask_cls = batch.mask_cls
            src_str = batch.src_str

            source_article += [' '.join(article) for article in src_str]

            if (cal_lead):
                selected_ids = [list(range(batch.clss.size(1)))] * batch.batch_size
            elif (cal_oracle):
                selected_ids = [[j for j in range(batch.clss.size(1)) if labels[i][j] == 1] for i in
                                range(batch.batch_size)]
            else:
                sent_scores, mask = self.model(src, segs, clss, mask, mask_cls)

                sent_scores = sent_scores + mask.float()
                sent_scores = sent_scores.cpu().data.numpy()
                # Sort sentence ids in descending order based on sentence scores (representing summary importance)
                selected_ids = np.argsort(-sent_scores, 1)
            # selected_ids = np.sort(selected_ids,1)
            for i, idx in enumerate(selected_ids):
                _pred = []
                if(len(batch.src_str[i])==0):
                    continue
                # Loop through each sentence
                # len(batch.src_str[i]) refers to the number of sentences in the jth test example
                for j in selected_ids[i][:len(batch.src_str[i])]:
                    if(j>=len( batch.src_str[i])):
                        continue
                    candidate = batch.src_str[i][j].strip()
                    if(self.args.block_trigram):
                        if(not _block_tri(candidate,_pred)):
                            _pred.append(candidate)
                    else:
                        _pred.append(candidate)
                    
                    # len(_pred) == 3 means that we limit sentences to top top_n_sentences
                    if ((not cal_oracle) and (not self.args.recall_eval) and len(_pred) == top_n_sentences):
                        break

                _pred = '<q>'.join(_pred)

                pred.append(_pred)

    results = {'source_article': source_article, 'predicted_summary': pred}
    return results

def example_api(self, example, step, top_n_sentences=3, device='cpu', cal_lead=False, cal_oracle=False):
    """ 
    Runs inference on a single test example. Designed for API deployemnt.
    """
    # Set model in validating mode.
    def _get_ngrams(n, text):
        ngram_set = set()
        text_length = len(text)
        max_index_ngram_start = text_length - n
        for i in range(max_index_ngram_start + 1):
            ngram_set.add(tuple(text[i:i + n]))
        return ngram_set

    def _block_tri(c, p):
        tri_c = _get_ngrams(3, c.split())
        for s in p:
            tri_s = _get_ngrams(3, s.split())
            if len(tri_c.intersection(tri_s))>0:
                return True
        return False

    if (not cal_lead and not cal_oracle):
        # Evaluate without performing backpropagation and dropout
        self.model.eval()
    # Set model device (cuda or cpu)
    self.model.to(device=device) 
    source_article = []
    pred = []
    src = example.src
    labels = example.labels
    segs = example.segs
    clss = example.clss
    mask = example.mask
    mask_cls = example.mask_cls
    src_str = example.src_str

    source_article += [' '.join(article) for article in src_str]

    if (cal_lead):
        selected_ids = [list(range(example.clss.size(1)))] * example.batch_size
    elif (cal_oracle):
        selected_ids = [[j for j in range(example.clss.size(1)) if labels[i][j] == 1] for i in
                        range(example.batch_size)]
    else:
        sent_scores, mask = self.model(src, segs, clss, mask, mask_cls)

        sent_scores = sent_scores + mask.float()
        sent_scores = sent_scores.cpu().data.numpy()
        # Sort sentence ids in descending order based on sentence scores (representing summary importance)
        selected_ids = np.argsort(-sent_scores, 1)
    # selected_ids = np.sort(selected_ids,1)
    for i, idx in enumerate(selected_ids):
        _pred = []
        if(len(example.src_str[i])==0):
            continue
        # Loop through each sentence
        # len(example.src_str[i]) refers to the number of sentences in the jth test example
        for j in selected_ids[i][:len(example.src_str[i])]:
            if(j>=len( example.src_str[i])):
                continue
            candidate = example.src_str[i][j].strip()
            if(self.args.block_trigram):
                if(not _block_tri(candidate,_pred)):
                    _pred.append(candidate)
            else:
                _pred.append(candidate)

            # len(_pred) == 3 means that we limit sentences to top top_n_sentences
            if ((not cal_oracle) and (not self.args.recall_eval) and len(_pred) == top_n_sentences):
                break

        _pred = '<q>'.join(_pred)

        pred.append(_pred)

    results = {'source_article': source_article, 'predicted_summary': pred}
    return results
Trainer.test_api = test_api
Trainer.example_api = example_api
cp = args.test_from
step = int(cp.split('.')[-2].split('_')[-1])
device_id = 0 if device == "cuda" else -1
# Read data
test_iter = data_loader.Dataloader(args, load_dataset(args, 'test', shuffle=False),
                              args.batch_size, device,
                              shuffle=False, is_test=True)
# sample is an iterator that contains a batch of data inputs to the "trainer"
# For our context, this will be for the purpose of doing a forward pass with unseen articles
sample = next(iter(test_iter))
sample.
# You'll notice the initial tokenization is performed on the article while separating it into a list of sentences
sample.src_str[0]
# Notice that the target text is abstractive, meaning it's not exactly extracted from the source document
sample.tgt_str
sample.src[0], sample.labels[0],sample.segs[0]
sample.clss[0], sample.mask[0], sample.mask_cls[0]
# build_trainer returns a Trainer class which now contains our custom test_api method
trainer = build_trainer(args, device_id, model, None)
%%time
#%debug
results = trainer.test_api(test_iter, step, top_n_sentences=4)
results_df = pd.DataFrame(results)
results_df.head()
results_df.shape
results_df.iloc[0, 0]
results_df.iloc[0, 1]
ls /home/bert_results
# Predicted summary
!head -1 /home/bert_results/cnndm_step43000.candidate
# Ground truth summary
!head -1 /home/bert_results/cnndm_step43000.gold
# First 10 predicted summaries
!head -10 /home/bert_results/cnndm_step43000.candidate
# First 10 ground truth summaries
!head -10 /home/bert_results/cnndm_step43000.gold<jupyter_output><NAME> posted video on his instagram of the cast being removed<q>the austrian is currently out after suffering knee ligament damage<q>bayern munich face porto in the champions league on wednesday
sabrina osterkamp was last seen leaving home in naracoorte at 12pm on sunday<q>the german tourist was driving a friend 's car about 110km south to mount gambier in south australia 's southeast<q>the car was found abandoned on side of road in centre of mount gambier<q>the 25-year-old contacted authorities to say she was safe and well almost 24 hours later
<NAME> vandalised and spat on the vehicle in 2013<q>the forensic officer swabbed the spit and found philp from his dna<q>he pleaded guilty to wilful damage of police property and was fined $ 300<q>mr philp claimed it was out of character but he has similar convictions
national union of teachers expected to call for a vote on industrial action<q>ballot of its 300,000 members would take place after general election<q>fears fun[...]<jupyter_text>Example level inference<jupyter_code>from prepro import data_builder
from models import data_loader
import nltk
nltk.download('punkt')
# Preprocessing arguments
# Setting these as default arguments
# These attributes can be edited as needed
pp_args = Namespace(
    mode="",
    oracle_mode='greedy',
    shard_size=2000,
    min_nsents=3,
    max_nsents=100,
    min_src_ntokens=5,
    max_src_ntokens=200,
    lower=True,
    dataset='',
    n_cpus=2
)
src_str = "<NAME> has given bayern munich some good news after posting a video on his instagram of the cast being removed following his recent knee ligament damage . the austrian has n't had much luck in terms of injuries but is clearly delighted he can now step-up his recovery and return for the bundesliga leaders . on his instagram , he said : ` moving forward ! <NAME> is looking to get back to fitness for bayern munich now his cast has been removed the versatile star had his cast signed by all his team-mates who arrived in portugal on monday , ahead of their champions league quarter-final clash with porto . bayern thrashed shakhtar donetsk in the last round and will be looking for a repeat result when they line up on wednesday at the estadio do dragao . having bounced back from their shock defeat to borussia monchengladbach , pep guardiola 's side have won three on the trot and are in good form going into the clash . the austrian international is a key player for bayern and his versatility sees him play at full back or in midfield"
tgt_str = '<NAME> is looking to get back to fitness for bayern munich now his cast has been removed<q><NAME> has given bayern munich some good news after posting a video on his instagram of the cast being removed following his recent knee ligament damage .<q>the austrian international is a key player for bayern and his versatility sees him play at full back or in midfield<q>the versatile star had his cast signed by all his team-mates who arrived in portugal on monday , ahead of their champions league quarter-final clash with porto .'
#Separate documents into list of sentences
src = [sent.split() for sent in nltk.tokenize.sent_tokenize(src_str)]
tgt = [sent.split() for sent in nltk.tokenize.sent_tokenize(tgt_str)]

src
bert = data_builder.BertData(pp_args)
#src_str, tgt_str, oracle_ids
oracle_ids = data_builder.greedy_selection(src, tgt, 3)
oracle_ids
b_data = bert.preprocess(src, tgt, oracle_ids)
b_data
indexed_tokens, labels, segments_ids, cls_ids, src_txt, tgt_txt = b_data
b_dict = {"src": indexed_tokens, "labels": labels, "segs": segments_ids, 'clss': cls_ids,
                       'src_str': src_txt, "tgt_str": tgt_txt}
data = [[b_dict['src'], b_dict['labels'], b_dict['segs'], b_dict['clss'], b_dict['src_txt'], b_dict['tgt_txt']]]
# Dimensions of input
len(data), len(data[0])
batch = data_loader.Batch(data, is_test=True, device='cpu') #cuda
batch.src_str
batch.tgt_str
# device_id = -1 meaning it's using cpu
trainer = build_trainer(args, device_id, model, None)
trainer.example_api(batch, step, top_n_sentences=2, device='cpu')
def summarize_text(src_str, args, pp_args, top_n_sentences=3, tgt_str=''):
    '''
    Summarizes input text by returning the most important sentences based on the BERT model fine-tuned on CNN and Daily Mail articles
    '''
    #Separate documents into list of sentences
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    src = [sent.split() for sent in nltk.tokenize.sent_tokenize(src_str)]
    tgt = [sent.split() for sent in nltk.tokenize.sent_tokenize(tgt_str)]
    bert = data_builder.BertData(pp_args)
    oracle_ids = data_builder.greedy_selection(src, tgt, 3)
    b_data = bert.preprocess(src, tgt, oracle_ids)
    indexed_tokens, labels, segments_ids, cls_ids, src_txt, tgt_txt = b_data
    b_dict = {"src": indexed_tokens, "labels": labels, "segs": segments_ids,
              'clss': cls_ids, 'src_str': src_txt, "tgt_str": tgt_txt}
    data = [[b_dict['src'], b_dict['labels'], b_dict['segs'], b_dict['clss'], b_dict['src_str'], b_dict['tgt_str']]]
    batch = data_loader.Batch(data, is_test=True, device=device)
    trained_model = torch.load(args.model_fp, map_location=lambda storage, loc: storage)
    # Read BERT config file
    # This contains information about the BERT model (e.g. hidden size for the transformer layers, number of transformer layers, number of self attention heads)
    config = BertConfig.from_json_file(args.bert_config_path)
    # Instantiate Summarizer and load pretrained model
    model = Summarizer(args, device, load_pretrained_bert=False, bert_config = config)
    model.load_cp(trained_model)
    device_id = 0 if device == "cuda" else -1
    trainer = build_trainer(args, device_id, model, None)
    results = trainer.example_api(batch, step, top_n_sentences=top_n_sentences, device=device)
    return results
    

%%time
summarize_text(src_str, args, pp_args, tgt_str=tgt_str, top_n_sentences=1)
tiger_woods_wiki = '''Woods grew up in Orange County, California. He was a child prodigy who was introduced to golf before the age of two by his athletic father, <NAME>. Earl was a single-digit handicap amateur golfer who also was one of the earliest African-American college baseball players at Kansas State University. Tiger's father was a member of the military and had playing privileges at the Navy golf course beside the Joint Forces Training Base in Los Alamitos, which allowed Tiger to play there. Tiger also played at the par 3 Heartwell golf course in Long Beach, as well as some of the municipals in Long Beach. In 1978, Tiger putted against comedian <NAME> in a television appearance on The Mike Douglas Show. At age three, he shot a 48 over nine holes at the Navy course. At age five, he appeared in Golf Digest and on ABC's That's Incredible! Before turning seven, Tiger won the Under Age 10 section of the Drive, Pitch, and Putt competition, held at the Navy Golf Course in Cypress, California. In 1984 at the age of eight, he won the 9–10 boys' event, the youngest age group available, at the Junior World Golf Championships. He first broke 80 at age eight. He went on to win the Junior World Championships six times, including four consecutive wins from 1988 to 1991. Woods' father Earl wrote that Tiger first defeated him at the age of 11 years, with Earl trying his best. Earl lost to Tiger every time from then on. Woods first broke 70 on a regulation golf course at age 12. When Woods was 13 years old, he played in the 1989 Big I, which was his first major national junior tournament. In the final round, he was paired with pro <NAME>, who was then relatively unknown. The event's format placed a professional with each group of juniors who had qualified. Daly birdied three of the last four holes to beat Woods by only one stroke. As a young teenager, Woods first met <NAME> in Los Angeles at the Bel-Air Country Club, when Nicklaus was performing a clinic for the club's members. Woods was part of the show, and he impressed Nicklaus and the crowd with his skills and potential. Earl Woods had researched in detail the career accomplishments of Nicklaus and had set his young son the goals of breaking those records. Woods was 15 years old and a student at Western High School in Anaheim when he became the youngest U.S. Junior Amateur champion; this was a record that stood until it was broken by <NAME> in 2010. He was named 1991's Southern California Amateur Player of the Year (for the second consecutive year) and Golf Digest Junior Amateur Player of the Year. In 1992, he defended his title at the U.S. Junior Amateur Championship, becoming the tournament's first two-time winner. He also competed in his first PGA Tour event, the Nissan Los Angeles Open (he missed the 36-hole cut), and was named Golf Digest Amateur Player of the Year, Golf World Player of the Year, and Golfweek National Amateur of the Year. The following year, Woods won his third consecutive U.S. Junior Amateur Championship; he remains the event's only three-time winner. In 1994, at the TPC at Sawgrass in Florida, he became the youngest winner of the U.S. Amateur Championship, a record he held until 2008 when it was broken by <NAME>. He was a member of the American team at the 1994 Eisenhower Trophy World Amateur Golf Team Championships (winning), and the 1995 Walker Cup (losing). Woods graduated from Western High School at age 18 in 1994 and was voted "Most Likely to Succeed" among the graduating class. He had starred for the high school's golf team under coach <NAME>. Woods overcame difficulties with stuttering as a boy. This was not known until he wrote a letter to a boy who contemplated suicide. Woods wrote, "I know what it's like to be different and to sometimes not fit in. I also stuttered as a child and I would talk to my dog and he would sit there and listen until he fell asleep. I also took a class for two years to help me, and I finally learned to stop. '''
tiger_woods_wiki
summarize_text(tiger_woods_wiki, args, pp_args, tgt_str='', top_n_sentences=1)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09772835536501132, 'fraction_numerical': 0.020915207812313923, 'mean_word_length': 2.7028442288612062, 'pattern_counts': {'":': 11, '<': 89, '<?xml version=': 0, '>': 113, 'https://': 7, 'lorem ipsum': 0, 'www.': 1, 'xml': 2}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 37, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '345490', 'license': 'Apache-2.0', 'n_tokens_mistral': 22337, 'n_tokens_neox': 20989, 'n_words': 5718, 'path': 'experiments/BERTSUM_fine_tuning_forwardpass_experiments.ipynb', 'repo_name': 'thinkingmachines/BertSum', 'size': '67260'}