# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import json
import numpy as np
import tensorflow as tf
import collections
from sklearn.cross_validation import train_test_split

# +
with open('ctexts.json','r') as fopen:
    ctexts = json.load(fopen)
    
with open('headlines.json','r') as fopen:
    headlines = json.load(fopen)

# +
import re
from unidecode import unidecode
import random

max_len = 500

def textcleaning(string):
    string = unidecode(string).replace('.', '. ').replace(',', ' , ')
    string = re.sub('[^\'"A-Za-z\- ]+', ' ', string)
    string = re.sub(r'[ ]+', ' ', string.lower()).strip().split()
    if len(string) > max_len:
        string = random.sample(string, max_len)
    return ' '.join(string)


# -

# %%time
h, c = [], []
for i in range(len(ctexts)):
    try:
        c.append(textcleaning(ctexts[i]))
        h.append(textcleaning(headlines[i]))
    except:
        pass

len(h)


def build_dataset(words, n_words):
    count = [['PAD', 0], ['GO', 1], ['EOS', 2], ['UNK', 3]]
    count.extend(collections.Counter(words).most_common(n_words))
    dictionary = dict()
    for word, _ in count:
        dictionary[word] = len(dictionary)
    data = list()
    unk_count = 0
    for word in words:
        index = dictionary.get(word, 0)
        if index == 0:
            unk_count += 1
        data.append(index)
    count[0][1] = unk_count
    reversed_dictionary = dict(zip(dictionary.values(), dictionary.keys()))
    return data, count, dictionary, reversed_dictionary


concat = ' '.join(c).split()
vocabulary_size = len(list(set(concat)))
data, count, dictionary, rev_dictionary = build_dataset(concat, vocabulary_size)
print('vocab from size: %d'%(vocabulary_size))
print('Most common words', count[4:10])
print('Sample data', data[:10], [rev_dictionary[i] for i in data[:10]])
print('filtered vocab size:',len(dictionary))
print("% of vocab used: {}%".format(round(len(dictionary)/vocabulary_size,4)*100))

for i in range(len(h)):
    h[i] = h[i] + ' EOS'
h[0]

GO = dictionary['GO']
PAD = dictionary['PAD']
EOS = dictionary['EOS']
UNK = dictionary['UNK']

train_X, test_X, train_Y, test_Y = train_test_split(c, h, test_size = 0.2)


# +
def sent2idx(sent, vocab, UNK=UNK):
    tokens = sent.split()
    oovs = []
    extend_tokens = []
    tokenized = []
    for token in tokens:
        if token not in vocab:
            tokenized.append(UNK)
            if token not in oovs:
                oovs.append(token)
            extend_tokens.append(len(vocab) + oovs.index(token))
        else:
            extend_tokens.append(vocab[token])
            tokenized.append(vocab[token])
    return tokenized, extend_tokens, oovs

def target2idx(sent, oovs, vocab,UNK=UNK):
    tokens = sent.split()
    tokenized = []
    for token in tokens:
        if token not in vocab:
            if token not in oovs:
                tokenized.append(UNK)
            else:
                tokenized.append(len(vocab) + oovs.index(token))
        else:
            tokenized.append(vocab[token])
    return tokenized


# +
import tensorflow as tf
from tensorflow.python.framework import ops, tensor_shape
from tensorflow.contrib.rnn.python.ops.core_rnn_cell import _linear
from tensorflow.python.ops.rnn_cell_impl import _zero_state_tensors
from tensorflow.python.util import nest
from tensorflow.python.ops import init_ops, array_ops, variable_scope
from tensorflow.contrib.seq2seq.python.ops.attention_wrapper import (
    _compute_attention,
)

UNK_ID = 3


class PointerGeneratorGreedyEmbeddingHelper(
    tf.contrib.seq2seq.GreedyEmbeddingHelper
):
    def __init__(self, embedding, start_tokens, end_token):
        self.vocab_size = tf.shape(embedding)[-1]
        super(PointerGeneratorGreedyEmbeddingHelper, self).__init__(
            embedding, start_tokens, end_token
        )

    def sample(self, time, outputs, state, name = None):
        """sample for PointerGeneratorGreedyEmbeddingHelper."""
        del time, state
        if not isinstance(outputs, ops.Tensor):
            raise TypeError(
                'Expected outputs to be a single Tensor, got: %s'
                % type(outputs)
            )
        sample_ids = tf.argmax(outputs, axis = -1, output_type = tf.int32)
        return sample_ids

    def next_inputs(self, time, outputs, state, sample_ids, name = None):
        """next_inputs_fn for GreedyEmbeddingHelper."""
        del time, outputs  # unused by next_inputs_fn
        finished = tf.equal(sample_ids, self._end_token)
        all_finished = tf.reduce_all(finished)

        condition = tf.less(sample_ids, self.vocab_size)
        sample_ids = tf.where(
            condition, sample_ids, tf.ones_like(sample_ids) * UNK_ID
        )

        next_inputs = tf.cond(
            all_finished,
            lambda: self._start_inputs,
            lambda: self._embedding_fn(sample_ids),
        )
        return (finished, next_inputs, state)


class PointerGeneratorDecoder(tf.contrib.seq2seq.BasicDecoder):
    """Pointer Generator sampling decoder."""

    def __init__(
        self,
        source_extend_tokens,
        source_oov_words,
        coverage,
        cell,
        helper,
        initial_state,
        output_layer = None,
    ):
        self.source_oov_words = source_oov_words
        self.source_extend_tokens = source_extend_tokens
        self.coverage = coverage
        super(PointerGeneratorDecoder, self).__init__(
            cell, helper, initial_state, output_layer
        )

    @property
    def output_size(self):
        # Return the cell output and the id
        return tf.contrib.seq2seq.BasicDecoderOutput(
            rnn_output = self._rnn_output_size() + self.source_oov_words,
            sample_id = self._helper.sample_ids_shape,
        )

    @property
    def output_dtype(self):
        # Assume the dtype of the cell is the output_size structure
        # containing the input_state's first component's dtype.
        # Return that structure and the sample_ids_dtype from the helper.
        dtype = nest.flatten(self._initial_state)[0].dtype
        return tf.contrib.seq2seq.BasicDecoderOutput(
            nest.map_structure(
                lambda _: dtype, self._rnn_output_size() + self.source_oov_words
            ),
            self._helper.sample_ids_dtype,
        )

    def step(self, time, inputs, state, name = None):
        with ops.name_scope(name, 'PGDecoderStep', (time, inputs, state)):
            cell_outputs, cell_state = self._cell(inputs, state)
            attention = cell_state[0].attention
            att_cell_state = cell_state[0].cell_state
            alignments = cell_state[0].alignments

            with tf.variable_scope('calculate_pgen'):
                p_gen = _linear([attention, inputs, att_cell_state], 1, True)
                p_gen = tf.sigmoid(p_gen)

            if self._output_layer is not None:
                cell_outputs = self._output_layer(cell_outputs)

            vocab_dist = tf.nn.softmax(cell_outputs) * p_gen
            alignments = alignments * (1 - p_gen)
            vocab_size = tf.shape(vocab_dist)[-1]
            extended_vsize = vocab_size + self.source_oov_words
            batch_size = tf.shape(vocab_dist)[0]
            extra_zeros = tf.zeros((batch_size, self.source_oov_words))
            vocab_dists_extended = tf.concat(
                axis = -1, values = [vocab_dist, extra_zeros]
            )

            batch_nums = tf.range(0, limit = batch_size)
            batch_nums = tf.expand_dims(batch_nums, 1)
            attn_len = tf.shape(self.source_extend_tokens)[1]
            batch_nums = tf.tile(
                batch_nums, [1, attn_len]
            )
            indices = tf.stack(
                (batch_nums, self.source_extend_tokens), axis = 2
            )
            shape = [batch_size, extended_vsize]
            attn_dists_projected = tf.scatter_nd(indices, alignments, shape)

            final_dists = attn_dists_projected + vocab_dists_extended
            sample_ids = self._helper.sample(
                time = time, outputs = final_dists, state = cell_state
            )

            (finished, next_inputs, next_state) = self._helper.next_inputs(
                time = time,
                outputs = cell_outputs,
                state = cell_state,
                sample_ids = sample_ids,
            )

            outputs = tf.contrib.seq2seq.BasicDecoderOutput(
                final_dists, sample_ids
            )
            return (outputs, next_state, next_inputs, finished)


class PointerGeneratorAttentionWrapper(tf.contrib.seq2seq.AttentionWrapper):
    def __init__(
        self,
        cell,
        attention_mechanism,
        attention_layer_size = None,
        alignment_history = False,
        cell_input_fn = None,
        output_attention = True,
        initial_cell_state = None,
        name = None,
        coverage = False,
    ):
        super(PointerGeneratorAttentionWrapper, self).__init__(
            cell,
            attention_mechanism,
            attention_layer_size,
            alignment_history,
            cell_input_fn,
            output_attention,
            initial_cell_state,
            name,
        )
        self.coverage = coverage

    def zero_state(self, batch_size, dtype):
        with ops.name_scope(
            type(self).__name__ + 'ZeroState', values = [batch_size]
        ):
            if self._initial_cell_state is not None:
                cell_state = self._initial_cell_state
            else:
                cell_state = self._cell.zero_state(batch_size, dtype)
            error_message = (
                'When calling zero_state of AttentionWrapper %s: '
                % self._base_name
                + 'Non-matching batch sizes between the memory '
                '(encoder output) and the requested batch size.  Are you using '
                'the BeamSearchDecoder?  If so, make sure your encoder output has '
                'been tiled to beam_width via tf.contrib.seq2seq.tile_batch, and '
                'the batch_size= argument passed to zero_state is '
                'batch_size * beam_width.'
            )
            with tf.control_dependencies(
                self._batch_size_checks(batch_size, error_message)
            ):
                cell_state = nest.map_structure(
                    lambda s: tf.identity(s, name = 'checked_cell_state'),
                    cell_state,
                )
            return tf.contrib.seq2seq.AttentionWrapperState(
                cell_state = cell_state,
                time = tf.zeros([], dtype = tf.int32),
                attention = _zero_state_tensors(
                    self._attention_layer_size, batch_size, dtype
                ),
                alignments = self._item_or_tuple(
                    attention_mechanism.initial_alignments(batch_size, dtype)
                    for attention_mechanism in self._attention_mechanisms
                ),
                attention_state = self._item_or_tuple(
                    attention_mechanism.initial_state(batch_size, dtype)
                    for attention_mechanism in self._attention_mechanisms
                ),
                # since we need to read the alignment history several times, so we need set clear_after_read to False
                alignment_history = self._item_or_tuple(
                    tf.TensorArray(
                        dtype = dtype,
                        size = 0,
                        clear_after_read = False,
                        dynamic_size = True,
                    )
                    if self._alignment_history
                    else ()
                    for _ in self._attention_mechanisms
                ),
            )

    def call(self, inputs, state):
        if not isinstance(state, tf.contrib.seq2seq.AttentionWrapperState):
            raise TypeError(
                'Expected state to be instance of AttentionWrapperState. '
                'Received type %s instead.' % type(state)
            )

        # Step 1: Calculate the true inputs to the cell based on the
        # previous attention value.
        cell_inputs = self._cell_input_fn(inputs, state.attention)
        cell_state = state.cell_state
        cell_output, next_cell_state = self._cell(cell_inputs, cell_state)

        cell_batch_size = cell_output.shape[0].value or tf.shape(cell_output)[0]
        error_message = (
            'When applying AttentionWrapper %s: ' % self.name
            + 'Non-matching batch sizes between the memory '
            '(encoder output) and the query (decoder output).  Are you using '
            'the BeamSearchDecoder?  You may need to tile your memory input via '
            'the tf.contrib.seq2seq.tile_batch function with argument '
            'multiple=beam_width.'
        )
        with tf.control_dependencies(
            self._batch_size_checks(cell_batch_size, error_message)
        ):
            cell_output = tf.identity(cell_output, name = 'checked_cell_output')

        if self._is_multi:
            previous_alignments = state.alignments
            previous_alignment_history = state.alignment_history
        else:
            previous_alignments = [state.alignments]
            previous_alignment_history = [state.alignment_history]

        all_alignments = []
        all_attentions = []
        all_histories = []

        for i, attention_mechanism in enumerate(self._attention_mechanisms):
            print(attention_mechanism)
            if self.coverage:
                # if we use coverage mode, previous alignments is coverage vector
                # alignment history stack has shape:  decoder time * batch * atten_len
                # convert it to coverage vector
                previous_alignments[i] = tf.cond(
                    previous_alignment_history[i].size() > 0,
                    lambda: tf.reduce_sum(
                        tf.transpose(
                            previous_alignment_history[i].stack(), [1, 2, 0]
                        ),
                        axis = 2,
                    ),
                    lambda: tf.zeros_like(previous_alignments[i]),
                )
            # debug
            # previous_alignments[i] = tf.Print(previous_alignments[i],[previous_alignment_history[i].size(), tf.shape(previous_alignments[i]),previous_alignments[i]],message="atten wrapper:")
            attention, alignments, next_attention_state = _compute_attention(
                attention_mechanism,
                cell_output,
                previous_alignments[i],
                self._attention_layers[i] if self._attention_layers else None,
            )
            alignment_history = (
                previous_alignment_history[i].write(state.time, alignments)
                if self._alignment_history
                else ()
            )

            all_alignments.append(alignments)
            all_histories.append(alignment_history)
            all_attentions.append(attention)

        attention = tf.concat(all_attentions, 1)
        next_state = tf.contrib.seq2seq.AttentionWrapperState(
            time = state.time + 1,
            cell_state = next_cell_state,
            attention = attention,
            alignments = self._item_or_tuple(all_alignments),
            attention_state = self._item_or_tuple(all_alignments),
            alignment_history = self._item_or_tuple(all_histories),
        )

        if self._output_attention:
            return attention, next_state
        else:
            return cell_output, next_state


def _pg_luong_score(query, keys, scale, coverage, coverage_vector):
    depth = query.get_shape()[-1]
    key_units = keys.get_shape()[-1]
    num_units = keys.shape[2].value or tf.shape(keys)[2]
    if depth != key_units:
        raise ValueError(
            'Incompatible or unknown inner dimensions between query and keys.  '
            'Query (%s) has units: %s.  Keys (%s) have units: %s.  '
            "Perhaps you need to set num_units to the keys' dimension (%s)?"
            % (query, depth, keys, key_units, key_units)
        )
    dtype = query.dtype
    query = array_ops.expand_dims(query, 1)
    score = tf.matmul(query, keys, transpose_b = True)
    score = array_ops.squeeze(score, [1])
    if scale:
        g = variable_scope.get_variable(
            'attention_g',
            dtype = dtype,
            initializer = init_ops.ones_initializer,
            shape = (),
        )
        score = g * score
    if coverage:
        w_c = tf.get_variable('coverage_w', [num_units], dtype = dtype)
        coverage_vector = tf.expand_dims(coverage_vector, -1)
        score = tf.expand_dims(score, -1)
        return tf.reduce_sum(score + coverage_vector, [2])


class PointerGeneratorLuongAttention(tf.contrib.seq2seq.LuongAttention):
    def __init__(
        self,
        num_units,
        memory,
        memory_sequence_length = None,
        probability_fn = None,
        score_mask_value = float('-inf'),
        name = 'PointerGeneratorLuongAttention',
        coverage = False,
        scale = True,
    ):
        super(PointerGeneratorLuongAttention, self).__init__(
            num_units = num_units,
            memory = memory,
            scale = scale,
            memory_sequence_length = memory_sequence_length,
            probability_fn = probability_fn,
            score_mask_value = score_mask_value,
            name = name,
        )
        self.coverage = coverage

    def __call__(self, query, state):
        with tf.variable_scope(
            None, 'pointer_generator_luong_attention', [query]
        ):

            score = _pg_luong_score(
                query, self._keys, self._scale, self.coverage, state
            )
        alignments = self._probability_fn(score, state)
        next_state = alignments
        return alignments, next_state



# -

class Summarization:
    def __init__(self, size_layer, num_layers, embedded_size, dict_size):
        
        def lstm_cell(reuse=False):
            return tf.nn.rnn_cell.GRUCell(size_layer, reuse=reuse)
        
        self.X = tf.placeholder(tf.int32, [None, None])
        self.Y = tf.placeholder(tf.int32, [None, None])
        self.X_seq_len = tf.count_nonzero(self.X, 1, dtype=tf.int32)
        self.Y_seq_len = tf.count_nonzero(self.Y, 1, dtype=tf.int32)
        batch_size = tf.shape(self.X)[0]
        self.source_oov_words = tf.placeholder(tf.int32, shape=[])
        self.source_extend_tokens = tf.placeholder(tf.int32, shape=[None, None])
        main = tf.strided_slice(self.Y, [0, 0], [batch_size, -1], [1, 1])
        decoder_input = tf.concat([tf.fill([batch_size, 1], GO), main], 1)
        
        condition = tf.less(decoder_input, dict_size)
        self.decoder_input = self.Y
        self.decoder_input_length = self.Y_seq_len
        self.predict_count = tf.reduce_sum(self.decoder_input_length)
        
        embeddings = tf.Variable(tf.random_uniform([dict_size, embedded_size], -1, 1))
        
        encoder_embedded = tf.nn.embedding_lookup(embeddings, self.X)
        encoder_cells = tf.nn.rnn_cell.MultiRNNCell([lstm_cell() for _ in range(num_layers)])
        self.encoder_out, self.encoder_state = tf.nn.dynamic_rnn(cell = encoder_cells, 
                                                                 inputs = encoder_embedded, 
                                                                 sequence_length = self.X_seq_len,
                                                                 dtype = tf.float32)
        self.decode_initial_state = self.encoder_state[-1]
        print(self.decode_initial_state)
        
        atten_mech = PointerGeneratorLuongAttention(
                size_layer, self.encoder_out, memory_sequence_length=self.X_seq_len,
        coverage = True)
        decoder_cells = [lstm_cell() for _ in range(num_layers)]
        decoder_cells[0] = PointerGeneratorAttentionWrapper(
                cell=decoder_cells[0],
                attention_mechanism=atten_mech,
                attention_layer_size=size_layer,
                alignment_history = True,
                coverage = True
            )
        initial_state = [self.decode_initial_state for i in range(num_layers)]
        attention_cell_state = decoder_cells[0].zero_state(
                dtype=tf.float32, batch_size=batch_size)
        initial_state[0] = attention_cell_state.clone(
                cell_state=initial_state[0])
        self.initial_state = tuple(initial_state)
        decoder_cells = tf.contrib.rnn.MultiRNNCell(decoder_cells)
        
        decoded = tf.nn.embedding_lookup(embeddings, self.decoder_input)
        
        training_helper = tf.contrib.seq2seq.TrainingHelper(
            decoded,
            self.decoder_input_length
        )
        dense_layer = tf.layers.Dense(dict_size)
        
        training_decoder = PointerGeneratorDecoder(
            source_extend_tokens = self.source_extend_tokens,
            source_oov_words = self.source_oov_words,
            coverage = True,
            cell=decoder_cells,
            helper=training_helper,
            initial_state=self.initial_state,
            output_layer=dense_layer
        )
        
        maxlen = tf.reduce_max(self.decoder_input_length)
        train_dec_outputs, train_dec_last_state, _ = tf.contrib.seq2seq.dynamic_decode(
            training_decoder,
            output_time_major=False,
            impute_finished=True,
            maximum_iterations=maxlen,
            swap_memory=True)
        logits = train_dec_outputs.rnn_output
        
        masks = tf.sequence_mask(
            self.decoder_input_length, maxlen, 
            dtype=tf.float32)
        
        targets = tf.slice(self.Y, [0, 0], [-1, maxlen])
        i1, i2 = tf.meshgrid(tf.range(batch_size),
                     tf.range(maxlen), indexing="ij")
        indices = tf.stack((i1,i2,targets),axis=2)
        probs = tf.gather_nd(logits, indices)
        probs = tf.where(tf.less_equal(probs,0),tf.ones_like(probs)*1e-10,probs)
        crossent = -tf.log(probs)
        self.cost = tf.reduce_sum(crossent * masks) / tf.to_float(batch_size)
        alignment_history = train_dec_last_state[0].alignment_history.stack()
        alignment_history = tf.transpose(alignment_history,[1,2,0])
        coverage_loss = tf.minimum(alignment_history,tf.cumsum(alignment_history, axis=2, exclusive=True))
        self.coverage_loss = tf.reduce_sum(coverage_loss / tf.to_float(batch_size))
        self.cost = self.cost + self.coverage_loss
        
        self.optimizer = tf.train.AdamOptimizer().minimize(self.cost)
        
        helper = PointerGeneratorGreedyEmbeddingHelper(
            embedding=embeddings,
            start_tokens=tf.tile(tf.constant([GO], dtype=tf.int32), [batch_size]),
            end_token=EOS
        )
        
        inference_decoder = PointerGeneratorDecoder(
            source_extend_tokens = self.source_extend_tokens,
            source_oov_words = self.source_oov_words,
            coverage = True,
            cell=decoder_cells,
            helper=helper,
            initial_state=self.initial_state,
            output_layer=dense_layer
        )
        
        dec_outputs, dec_last_state, _ = tf.contrib.seq2seq.dynamic_decode(
            inference_decoder,
            output_time_major=False,
            maximum_iterations=tf.reduce_max(self.X_seq_len),
            swap_memory=True)
        
        self.beam_predictions = dec_outputs.sample_id
        print(self.beam_predictions)
        
        masks = tf.sequence_mask(self.Y_seq_len, tf.reduce_max(self.Y_seq_len), dtype=tf.float32)
        y_t = tf.argmax(logits,axis=2)
        y_t = tf.cast(y_t, tf.int32)
        self.prediction = tf.boolean_mask(y_t, masks)
        mask_label = tf.boolean_mask(self.Y, masks)
        correct_pred = tf.equal(self.prediction, mask_label)
        correct_index = tf.cast(correct_pred, tf.float32)
        self.accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))


size_layer = 128
num_layers = 2
embedded_size = 128
batch_size = 8

tf.reset_default_graph()
sess = tf.InteractiveSession()
model = Summarization(size_layer, num_layers, embedded_size, len(dictionary))
sess.run(tf.global_variables_initializer())

from keras.preprocessing.sequence import pad_sequences


def batching(X, Y):
    s_, es_, oovs_, target_ = [], [], [], []
    for x, y in zip(X, Y):
        s,es,oovs = sent2idx(x, dictionary)
        target = target2idx(y, oovs,dictionary)
        s_.append(s)
        es_.append(es)
        oovs_.append(oovs)
        target_.append(target)
    s_ = pad_sequences(s_,padding='post')
    es_ = pad_sequences(es_,padding='post')
    target_ = pad_sequences(target_,padding='post')
    maxlen = max([len(o) for o in oovs_])
    return s_, es_, target_, maxlen


# +
from tqdm import tqdm
from sklearn.utils import shuffle
import time

for EPOCH in range(10):
    lasttime = time.time()
    total_loss, total_accuracy, total_loss_test, total_accuracy_test = 0, 0, 0, 0
    train_X, train_Y = shuffle(train_X, train_Y)
    test_X, test_Y = shuffle(test_X, test_Y)
    pbar = tqdm(range(0, len(train_X), batch_size), desc='train minibatch loop')
    for k in pbar:
        index = min(k+batch_size,len(train_X))
        batch_x, batch_es, batch_y, maxlen = batching(train_X[k: index],
                                                     train_Y[k: index])
        acc, loss, _ = sess.run([model.accuracy, model.cost, model.optimizer], 
                                      feed_dict={model.X:batch_x,
                                                 model.source_extend_tokens:batch_es,
                                                model.Y:batch_y,
                                                model.source_oov_words:maxlen})
        total_loss += loss
        total_accuracy += acc
        pbar.set_postfix(cost=loss, accuracy = acc)
        
    pbar = tqdm(range(0, len(test_X), batch_size), desc='test minibatch loop')
    for k in pbar:
        index = min(k+batch_size,len(test_X))
        batch_x, batch_es, batch_y, maxlen = batching(test_X[k: index],
                                                     test_Y[k: index])
        acc, loss = sess.run([model.accuracy, model.cost], 
                                      feed_dict={model.X:batch_x,
                                                 model.source_extend_tokens:batch_es,
                                                model.Y:batch_y,
                                                model.source_oov_words:maxlen})
        total_loss_test += loss
        total_accuracy_test += acc
        pbar.set_postfix(cost=loss, accuracy = acc)
        
    total_loss /= (len(train_X) / batch_size)
    total_accuracy /= (len(train_X) / batch_size)
    total_loss_test /= (len(test_X) / batch_size)
    total_accuracy_test /= (len(test_X) / batch_size)
        
    print('epoch: %d, avg loss: %f, avg accuracy: %f'%(EPOCH, total_loss, total_accuracy))
    print('epoch: %d, avg loss test: %f, avg accuracy test: %f'%(EPOCH, total_loss_test, total_accuracy_test))
# -

batch_x, batch_es, batch_y, maxlen = batching(test_X[:1], test_Y[:1])

out = [rev_dictionary[i] for i in sess.run(model.beam_predictions, feed_dict = {model.X: batch_x,
                                             model.Y: batch_y,
                                             model.source_extend_tokens:batch_es,
                                             model.source_oov_words:maxlen})[0]]


def f7(seq):
    seen = set()
    seen_add = seen.add
    return ' '.join([x for x in seq if not (x in seen or seen_add(x))])


f7(out)

test_Y[:1]



StarcoderdataJupyterScriptsDedupFiltered