<reponame>aeceou/OpenNAPE-py
#!/usr/bin/env python
""" PosetEditor Class and builder """
from __future__ import print_function
import codecs
import os
import time
import numpy as np
from itertools import count, zip_longest

import torch

import onmt.ape.model_builder as model_builder
import onmt.inputters as inputters
import onmt.ape.decoders.ensemble as ensemble
from .beam_search import BeamSearch
from onmt.translate.greedy_search import GreedySearch
from onmt.utils.misc import tile, set_random_seed, report_matrix
from onmt.utils.alignment import extract_alignment, build_align_pharaoh
from .translation import PostEditBuilder
from onmt.translate import GNMTGlobalScorer


def build_post_editor(opt, report_score=True, logger=None, out_file=None):
    if out_file is None:
        out_file = codecs.open(opt.output, 'w+', 'utf-8')

    load_test_model = ensemble.load_test_model \
        if len(opt.models) > 1 else model_builder.load_test_model
    fields, model, model_opt = load_test_model(opt)

    scorer = GNMTGlobalScorer.from_opt(opt)

    post_editor = PostEditor.from_opt(
        model,
        fields,
        opt,
        model_opt,
        global_scorer=scorer,
        out_file=out_file,
        report_align=opt.report_align,
        report_score=report_score,
        logger=logger
    )
    return post_editor


def max_tok_len(new, count, sofar):
    """
    In token batching scheme, the number of sequences is limited
    such that the total number of src, mt, and pe tokens (including padding)
    in a batch <= batch_size
    """
    # Maintains the longest src and tgt length in the current batch
    global max_src_in_batch, max_mt_in_batch  # this is a hack
    # Reset current longest length at a new batch (count=1)
    if count == 1:
        max_src_in_batch = 0
        max_mt_in_batch = 0
        # max_pe_in_batch = 0
    # Src: [<bos> w1 ... wN <eos>]
    max_src_in_batch = max(max_src_in_batch, len(new.src[0]) + 2)
    # mt: [<bos> w1 ... wN <eos>]
    max_mt_in_batch = max(max_mt_in_batch, len(new.mt[0]) + 2)
    # pe: [w1 ... wM <eos>]
    src_elements = count * max_src_in_batch
    mt_elements = count * max_mt_in_batch
    return max(src_elements, mt_elements)


class PostEditor(object):
    """Translate a batch of sentences with a saved model.

    Args:
        model (onmt.APEModel): APE model to use for post-editing
        fields (dict[str, torchtext.data.Field]): A dict
            mapping each side to its list of name-Field pairs.
        src_reader (onmt.inputters.DataReaderBase): Source reader.
        mt_reader (onmt.inputters.DataReaderBase): MT output reader.
        pe_reader (onmt.inputters.TextDataReader): Post-edit reader.
        gpu (int): GPU device. Set to negative for no GPU.
        n_best (int): How many beams to wait for.
        min_length (int): See
            :class:`onmt.translate.decode_strategy.DecodeStrategy`.
        max_length (int): See
            :class:`onmt.translate.decode_strategy.DecodeStrategy`.
        beam_size (int): Number of beams.
        random_sampling_topk (int): See
            :class:`onmt.translate.greedy_search.GreedySearch`.
        random_sampling_temp (int): See
            :class:`onmt.translate.greedy_search.GreedySearch`.
        stepwise_penalty (bool): Whether coverage penalty is applied every step
            or not.
        dump_beam (bool): Debugging option.
        block_ngram_repeat (int): See
            :class:`onmt.translate.decode_strategy.DecodeStrategy`.
        ignore_when_blocking (set or frozenset): See
            :class:`onmt.translate.decode_strategy.DecodeStrategy`.
        replace_unk (bool): Replace unknown token.
        data_type (str): Source data type.
        verbose (bool): Print/log every translation.
        report_time (bool): Print/log total time/frequency.
        copy_attn (bool): Use copy attention.
        global_scorer (onmt.translate.GNMTGlobalScorer): Translation
            scoring/reranking object.
        out_file (TextIO or codecs.StreamReaderWriter): Output file.
        report_score (bool) : Whether to report scores
        logger (logging.Logger or NoneType): Logger.
    """

    def __init__(
            self,
            model,
            fields,
            src_reader,
            mt_reader,
            pe_reader,
            gpu=-1,
            n_best=1,
            min_length=0,
            max_length=100,
            ratio=0.,
            beam_size=30,
            random_sampling_topk=1,
            random_sampling_temp=1,
            stepwise_penalty=None,
            dump_beam=False,
            block_ngram_repeat=0,
            ignore_when_blocking=frozenset(),
            replace_unk=False,
            phrase_table="",
            data_type="text",
            verbose=False,
            report_time=False,
            copy_attn=False,
            global_scorer=None,
            out_file=None,
            report_align=False,
            report_score=True,
            logger=None,
            seed=-1):
        self.model = model
        self.fields = fields
        pe_field = dict(self.fields)["pe"].base_field
        self._pe_vocab = pe_field.vocab
        self._pe_eos_idx = self._pe_vocab.stoi[pe_field.eos_token]
        self._pe_pad_idx = self._pe_vocab.stoi[pe_field.pad_token]
        self._pe_bos_idx = self._pe_vocab.stoi[pe_field.init_token]
        self._pe_unk_idx = self._pe_vocab.stoi[pe_field.unk_token]
        self._pe_vocab_len = len(self._pe_vocab)

        self._gpu = gpu
        self._use_cuda = gpu > -1
        self._dev = torch.device("cuda", self._gpu) \
            if self._use_cuda else torch.device("cpu")

        self.n_best = n_best
        self.max_length = max_length

        self.beam_size = beam_size
        self.random_sampling_temp = random_sampling_temp
        self.sample_from_topk = random_sampling_topk

        self.min_length = min_length
        self.ratio = ratio
        self.stepwise_penalty = stepwise_penalty
        self.dump_beam = dump_beam
        self.block_ngram_repeat = block_ngram_repeat
        self.ignore_when_blocking = ignore_when_blocking
        self._exclusion_idxs = {
            self._pe_vocab.stoi[t] for t in self.ignore_when_blocking}
        self.src_reader = src_reader
        self.mt_reader = mt_reader
        self.pe_reader = pe_reader
        self.replace_unk = replace_unk
        if self.replace_unk and not self.model.decoder.attentional:
            raise ValueError(
                "replace_unk requires an attentional decoder.")
        self.phrase_table = phrase_table
        self.data_type = data_type
        self.verbose = verbose
        self.report_time = report_time

        self.copy_attn = copy_attn

        self.global_scorer = global_scorer
        if self.global_scorer.has_cov_pen and \
                not self.model.decoder.attentional:
            raise ValueError(
                "Coverage penalty requires an attentional decoder.")
        self.out_file = out_file
        self.report_align = report_align
        self.report_score = report_score
        self.logger = logger

        self.use_filter_pred = False
        self._filter_pred = None

        # for debugging
        self.beam_trace = self.dump_beam != ""
        self.beam_accum = None
        if self.beam_trace:
            self.beam_accum = {
                "predicted_ids": [],
                "beam_parent_ids": [],
                "scores": [],
                "log_probs": []}

        set_random_seed(seed, self._use_cuda)

    @classmethod
    def from_opt(
            cls,
            model,
            fields,
            opt,
            model_opt,
            global_scorer=None,
            out_file=None,
            report_align=False,
            report_score=True,
            logger=None):
        """Alternate constructor.

        Args:
            model (onmt.modules.NMTModel): See :func:`__init__()`.
            fields (dict[str, torchtext.data.Field]): See
                :func:`__init__()`.
            opt (argparse.Namespace): Command line options
            model_opt (argparse.Namespace): Command line options saved with
                the model checkpoint.
            global_scorer (onmt.translate.GNMTGlobalScorer): See
                :func:`__init__()`..
            out_file (TextIO or codecs.StreamReaderWriter): See
                :func:`__init__()`.
            report_align (bool) : See :func:`__init__()`.
            report_score (bool) : See :func:`__init__()`.
            logger (logging.Logger or NoneType): See :func:`__init__()`.
        """

        src_reader = inputters.str2reader[opt.data_type].from_opt(opt)
        mt_reader = inputters.str2reader[opt.data_type].from_opt(opt)
        pe_reader = inputters.str2reader["text"].from_opt(opt)
        return cls(
            model,
            fields,
            src_reader,
            mt_reader,
            pe_reader,
            gpu=opt.gpu,
            n_best=opt.n_best,
            min_length=opt.min_length,
            max_length=opt.max_length,
            ratio=opt.ratio,
            beam_size=opt.beam_size,
            random_sampling_topk=opt.random_sampling_topk,
            random_sampling_temp=opt.random_sampling_temp,
            stepwise_penalty=opt.stepwise_penalty,
            dump_beam=opt.dump_beam,
            block_ngram_repeat=opt.block_ngram_repeat,
            ignore_when_blocking=set(opt.ignore_when_blocking),
            replace_unk=opt.replace_unk,
            phrase_table=opt.phrase_table,
            data_type=opt.data_type,
            verbose=opt.verbose,
            report_time=opt.report_time,
            copy_attn=model_opt.copy_attn,
            global_scorer=global_scorer,
            out_file=out_file,
            report_align=report_align,
            report_score=report_score,
            logger=logger,
            seed=opt.seed)

    def _log(self, msg):
        if self.logger:
            self.logger.info(msg)
        else:
            print(msg)

    def _gold_score(self, batch, memory_bank, inp_lengths, inp_vocabs,
                    use_inp_map, enc_states, batch_size, inp):
        if "pe" in batch.__dict__:
            gs = self._score_post_edit(
                batch, memory_bank, inp_lengths, inp_vocabs,
                batch.src_map if use_inp_map else None)
            self.model.decoder.init_state(inp, memory_bank, enc_states)
        else:
            gs = [0] * batch_size
        return gs

    def post_edit(
            self,
            src,
            mt,
            pe=None,
            batch_size=None,
            batch_type="sents",
            attn_debug=False,
            align_debug=False,
            phrase_table=""):
        """Post-edit contents of ``mt`` and get gold scores from ``pe``.

        Args:
            src: See :func:`self.src_reader.read()`.
            mt: See :func:`self.mt_reader.read()`.
            pe: See :func:`self.pe_reader.read()`.
            batch_size (int): size of examples per mini-batch
            attn_debug (bool): enables the attention logging
            align_debug (bool): enables the word alignment logging

        Returns:
            (`list`, `list`)

            * all_scores is a list of `batch_size` lists of `n_best` scores
            * all_predictions is a list of `batch_size` lists
                of `n_best` predictions
        """

        if batch_size is None:
            raise ValueError("batch_size must be set")

        src_data = {"reader": self.src_reader, "data": src, "dir": None}
        mt_data = {"reader": self.mt_reader, "data": mt, "dir": None}
        pe_data = {"reader": self.pe_reader, "data": pe, "dir": None}
        _readers, _data, _dir = inputters.Dataset.config(
            [('src', src_data), ('mt', mt_data), ('pe', pe_data)])

        # corpus_id field is useless here
        if self.fields.get("corpus_id", None) is not None:
            self.fields.pop('corpus_id')
        data = inputters.Dataset(
            self.fields, readers=_readers, data=_data, dirs=_dir,
            sort_key=inputters.str2sortkey[self.data_type],
            filter_pred=self._filter_pred
        )

        data_iter = inputters.OrderedIterator(
            dataset=data,
            device=self._dev,
            batch_size=batch_size,
            batch_size_fn=max_tok_len if batch_type == "tokens" else None,
            train=False,
            sort=False,
            sort_within_batch=True,
            shuffle=False
        )

        xlation_builder = PostEditBuilder(
            data, self.fields, self.n_best, self.replace_unk, pe,
            self.phrase_table
        )

        # Statistics
        counter = count(1)
        pred_score_total, pred_words_total = 0, 0
        gold_score_total, gold_words_total = 0, 0

        all_scores = []
        all_predictions = []

        start_time = time.time()

        for batch in data_iter:
            batch_data = self.translate_batch(
                batch, data.src_vocabs, attn_debug
            )
            translations = xlation_builder.from_batch(batch_data)

            for trans in translations:
                all_scores += [trans.pred_scores[:self.n_best]]
                pred_score_total += trans.pred_scores[0]
                pred_words_total += len(trans.pred_sents[0])
                if pe is not None:
                    gold_score_total += trans.gold_score
                    gold_words_total += len(trans.gold_sent) + 1

                n_best_preds = [" ".join(pred)
                                for pred in trans.pred_sents[:self.n_best]]
                if self.report_align:
                    align_pharaohs = [build_align_pharaoh(align) for align
                                      in trans.word_aligns[:self.n_best]]
                    n_best_preds_align = [" ".join(align) for align
                                          in align_pharaohs]
                    n_best_preds = [pred + " ||| " + align
                                    for pred, align in zip(
                                        n_best_preds, n_best_preds_align)]
                all_predictions += [n_best_preds]
                self.out_file.write('\n'.join(n_best_preds) + '\n')
                self.out_file.flush()

                if self.verbose:
                    sent_number = next(counter)
                    output = trans.log(sent_number)
                    if self.logger:
                        self.logger.info(output)
                    else:
                        os.write(1, output.encode('utf-8'))

        end_time = time.time()

        if self.report_score:
            msg = self._report_score('PRED', pred_score_total,
                                     pred_words_total)
            self._log(msg)
            if pe is not None:
                msg = self._report_score('GOLD', gold_score_total,
                                         gold_words_total)
                self._log(msg)

        if self.report_time:
            total_time = end_time - start_time
            self._log("Total translation time (s): %f" % total_time)
            self._log("Average translation time (s): %f" % (
                total_time / len(all_predictions)))
            self._log("Tokens per second: %f" % (
                pred_words_total / total_time))

        if self.dump_beam:
            import json
            json.dump(self.translator.beam_accum,
                      codecs.open(self.dump_beam, 'w', 'utf-8'))
        return all_scores, all_predictions

    def _align_pad_prediction(self, predictions, bos, pad):
        """
        Padding predictions in batch and add BOS.

        Args:
            predictions (List[List[Tensor]]): `(batch, n_best,)`, for each inp
                sequence contain n_best pe predictions all of which ended with
                eos id.
            bos (int): bos index to be used.
            pad (int): pad index to be used.

        Return:
            batched_nbest_predict (torch.LongTensor): `(batch, n_best, pe_l)`
        """
        dtype, device = predictions[0][0].dtype, predictions[0][0].device
        flatten_pe = [best.tolist() for bests in predictions
                       for best in bests]
        paded_pe = torch.tensor(
            list(zip_longest(*flatten_pe, fillvalue=pad)),
            dtype=dtype, device=device).T
        bos_tensor = torch.full([paded_pe.size(0), 1], bos,
                                dtype=dtype, device=device)
        full_pe = torch.cat((bos_tensor, paded_pe), dim=-1)
        batched_nbest_predict = full_pe.view(
            len(predictions), -1, full_pe.size(-1))  # (batch, n_best, pe_l)
        return batched_nbest_predict

    def _align_forward(self, batch, predictions):
        """
        For a batch of input and its prediction, return a list of batch predict
        alignment enc_inp indice Tensors in size ``(batch, n_best,)``.
        """
        # (0) add BOS and padding to pe prediction
        if hasattr(batch, 'pe'):
            batch_pe_idxs = batch.pe.transpose(1, 2).transpose(0, 2)
        else:
            batch_pe_idxs = self._align_pad_prediction(
                predictions, bos=self._pe_bos_idx, pad=self._pe_pad_idx)
        pe_mask = (batch_pe_idxs.eq(self._pe_pad_idx) |
                   batch_pe_idxs.eq(self._pe_eos_idx) |
                   batch_pe_idxs.eq(self._pe_bos_idx))

        n_best = batch_pe_idxs.size(1)
        # (1) Encoder forward.
        enc_out, enc_states, memory_bank, inp_lengths = self._run_encoder(batch)

        # (2) Repeat enc_out objects `n_best` times.
        # We use batch_size x n_best, get ``(src_len/mt_len, batch * n_best, nfeat)``
        for inp in enc_out.keys():
            enc_out[inp] = tile(enc_out[inp], n_best, dim=1)
            enc_states[inp] = tile(enc_states[inp], n_best, dim=1)
            if isinstance(memory_bank[inp], tuple):
                memory_bank[inp] = tuple(tile(x, n_best, dim=1) \
                                   for x in memory_bank[inp])
            else:
                memory_bank[inp] = tile(memory_bank[inp], n_best, dim=1)
            inp_lengths[inp] = tile(inp_lengths[inp], n_best)  # ``(batch * n_best,)``

        # (3) Init decoder with n_best enc_out,
        self.model.decoder.init_state(enc_out, memory_bank, enc_states)
        # reshape pe to ``(len, batch * n_best, nfeat)``
        pe = batch_pe_idxs.view(-1, batch_pe_idxs.size(-1)).T.unsqueeze(-1)
        dec_in = pe[:-1]  # exclude last target from inputs
        _, attns = self.model.decoder(
            dec_in, memory_bank, memory_lengths=inp_lengths, with_align=True)

        alignment_attn = attns['align']  # ``(B, pe_len-1, enc_len)`` each
        # masked_select
        align_pe_mask = pe_mask.view(-1, pe_mask.size(-1))
        prediction_mask = align_pe_mask[:, 1:]  # exclude bos to match pred
        # get aligned id for each prediction's valid pe tokens
        alignement = extract_alignment(
            alignment_attn, prediction_mask, inp_lengths, n_best)
        return alignement

    def translate_batch(self, batch, enc_vocabs, attn_debug):
        """Translate a batch of sentences."""
        with torch.no_grad():
            if self.beam_size == 1:
                decode_strategy = GreedySearch(
                    pad=self._pe_pad_idx,
                    bos=self._pe_bos_idx,
                    eos=self._pe_eos_idx,
                    batch_size=batch.batch_size,
                    min_length=self.min_length, max_length=self.max_length,
                    block_ngram_repeat=self.block_ngram_repeat,
                    exclusion_tokens=self._exclusion_idxs,
                    return_attention=attn_debug or self.replace_unk,
                    sampling_temp=self.random_sampling_temp,
                    keep_topk=self.sample_from_topk)
            else:
                # TODO: support these blacklisted features
                assert not self.dump_beam
                decode_strategy = BeamSearch(
                    self.beam_size,
                    batch_size=batch.batch_size,
                    pad=self._pe_pad_idx,
                    bos=self._pe_bos_idx,
                    eos=self._pe_eos_idx,
                    n_best=self.n_best,
                    global_scorer=self.global_scorer,
                    min_length=self.min_length, max_length=self.max_length,
                    return_attention=attn_debug or self.replace_unk,
                    block_ngram_repeat=self.block_ngram_repeat,
                    exclusion_tokens=self._exclusion_idxs,
                    stepwise_penalty=self.stepwise_penalty,
                    ratio=self.ratio)
            return self._translate_batch_with_strategy(batch, enc_vocabs,
                                                       decode_strategy)

    def _run_encoder(self, batch):
        src, src_lengths = batch.src if isinstance(batch.src, tuple) \
                           else (batch.src, None)
        mt, mt_lengths = batch.mt if isinstance(batch.mt, tuple) \
                           else (batch.mt, None)
        inp = {'src': src, 'mt': mt}
        inp_lengths = {'src': src_lengths, 'mt': mt_lengths}

        enc_states, memory_bank, inp_lengths = self.model.encoder(
            inp, inp_lengths)
        if inp_lengths['src'] == None:
            assert not isinstance(memory_bank['src'], tuple), \
                'Ensemble decoding only supported for text data'
            inp_lengths['src'] = torch.Tensor(batch.batch_size) \
                                      .type_as(memory_bank['src']) \
                                      .long() \
                                      .fill_(memory_bank['src'].size(0))
        if inp_lengths['mt'] == None:
            inp_lengths['mt'] = torch.Tensor(batch.batch_size) \
                                      .type_as(memory_bank['mt']) \
                                      .long() \
                                      .fill_(memory_bank['mt'].size(0))
        return inp, enc_states, memory_bank, inp_lengths

    def _decode_and_generate(
            self,
            decoder_in,
            memory_bank,
            batch,
            inp_vocabs,
            memory_lengths,
            inp_map=None,
            step=None,
            batch_offset=None):

        # Decoder forward, takes [pe_len, batch, nfeats] as input
        # and a dictionary of [src_len or mt_len, batch, hidden] as memory_bank
        # in case of inference pe_len = 1, batch = beam times batch_size
        # in case of Gold Scoring pe_len = actual length, batch = 1 batch
        dec_out, dec_attn = self.model.decoder(
            decoder_in, memory_bank, memory_lengths=memory_lengths, step=step
        )

        # Generator forward.
        if "std" in dec_attn:
            attn = dec_attn["std"]
        else:
            attn = None
        log_probs = self.model.generator(dec_out.squeeze(0))
        # returns [(batch_size x beam_size) , vocab ] when 1 step
        # or [ pe_len, batch_size, vocab ] when full sentence
        return log_probs, attn

    def _translate_batch_with_strategy(
            self,
            batch,
            inp_vocabs,
            decode_strategy):
        """Translate a batch of sentences step by step using cache.

        Args:
            batch: a batch of sentences, yield by data iterator.
            inp_vocabs (Dict[str, list]): A dictionary of lists of
                torchtext.data.Vocab if can_copy.
            decode_strategy (DecodeStrategy): A decode strategy to use for
                generate translation step by step.

        Returns:
            results (dict): The translation results.
        """
        # (0) Prep the components of the search.
        parallel_paths = decode_strategy.parallel_paths  # beam_size
        batch_size = batch.batch_size

        # (1) Run the encoder.
        inp, enc_states, memory_bank, inp_lengths = self._run_encoder(batch)
        self.model.decoder.init_state(inp, memory_bank, enc_states)

        results = {
            "predictions": None,
            "scores": None,
            "attention": None,
            "batch": batch,
            "gold_score": self._gold_score(
                batch, memory_bank, inp_lengths, inp_vocabs, False,
                enc_states, batch_size, inp)}

        # (2) prep decode_strategy. Possibly repeat inp objects.
        inp_map = None
        fn_map_state, memory_bank, memory_lengths, inp_map = \
            decode_strategy.initialize(memory_bank, inp_lengths, inp_map)
        if fn_map_state is not None:
            self.model.decoder.map_state(fn_map_state)

        # (3) Begin decoding step by step:
        for step in range(decode_strategy.max_length):
            decoder_input = decode_strategy.current_predictions.view(1, -1, 1)

            log_probs, attn = self._decode_and_generate(
                decoder_input,
                memory_bank,
                batch,
                inp_vocabs,
                memory_lengths=memory_lengths,
                inp_map=inp_map,
                step=step,
                batch_offset=decode_strategy.batch_offset)

            decode_strategy.advance(log_probs, attn)
            any_finished = decode_strategy.is_finished.any()
            if any_finished:
                decode_strategy.update_finished()
                if decode_strategy.done:
                    break

            select_indices = decode_strategy.select_indices

            if any_finished:
                # Reorder states.
                for inp in memory_bank.keys():
                    if isinstance(memory_bank, tuple):
                        memory_bank[inp] = tuple(x.index_select(1, select_indices)
                                            for x in memory_bank[inp])
                    else:
                        memory_bank[inp] = memory_bank[inp].index_select(1, select_indices)
                    memory_lengths[inp] = memory_lengths[inp].index_select(0, select_indices)

                    if inp_map is not None:
                        inp_map[inp] = inp_map[inp].index_select(1, select_indices)

            if parallel_paths > 1 or any_finished:
                self.model.decoder.map_state(
                    lambda state, dim: state.index_select(dim, select_indices))

        results["scores"] = decode_strategy.scores
        results["predictions"] = decode_strategy.predictions
        results["attention"] = decode_strategy.attention
        if self.report_align:
            results["alignment"] = self._align_forward(
                batch, decode_strategy.predictions)
        else:
            results["alignment"] = [[] for _ in range(batch_size)]
        return results

    def _score_post_edit(self, batch, memory_bank, inp_lengths,
                         inp_vocabs, inp_map):
        pe = batch.pe
        pe_in = pe[:-1]

        log_probs, attn = self._decode_and_generate(
            pe_in, memory_bank, batch, inp_vocabs,
            memory_lengths=inp_lengths, inp_map=inp_map)

        log_probs[:, :, self._pe_pad_idx] = 0
        gold = pe[1:]
        gold_scores = log_probs.gather(2, gold)
        gold_scores = gold_scores.sum(dim=0).view(-1)

        return gold_scores

    def _report_score(self, name, score_total, words_total):
        if words_total == 0:
            msg = "%s No words predicted" % (name,)
        else:
            avg_score = score_total / words_total
            ppl = np.exp(-score_total.item() / words_total)
            msg = ("%s AVG SCORE: %.4f, %s PPL: %.4f" % (
                name, avg_score,
                name, ppl))
        return msg

#!/usr/bin/env python
""" PosetEditor Class and builder """
from __future__ import print_function
import codecs
import os
import time
import numpy as np
from itertools import count, zip_longest

import torch

import onmt.ape.model_builder as model_builder
import onmt.inputters as inputters
import onmt.ape.decoders.ensemble as ensemble
from .beam_search import BeamSearch
from onmt.translate.greedy_search import GreedySearch
from onmt.utils.misc import tile, set_random_seed, report_matrix
from onmt.utils.alignment import extract_alignment, build_align_pharaoh
from .translation import PostEditBuilder
from onmt.translate import GNMTGlobalScorer


def build_post_editor(opt, report_score=True, logger=None, out_file=None):
    if out_file is None:
        out_file = codecs.open(opt.output, 'w+', 'utf-8')

    load_test_model = ensemble.load_test_model \
        if len(opt.models) > 1 else model_builder.load_test_model
    fields, model, model_opt = load_test_model(opt)

    scorer = GNMTGlobalScorer.from_opt(opt)

    post_editor = PostEditor.from_opt(
        model,
        fields,
        opt,
        model_opt,
        global_scorer=scorer,
        out_file=out_file,
        report_align=opt.report_align,
        report_score=report_score,
        logger=logger
    )
    return post_editor


def max_tok_len(new, count, sofar):
    """
    In token batching scheme, the number of sequences is limited
    such that the total number of src, mt, and pe tokens (including padding)
    in a batch <= batch_size
    """
    # Maintains the longest src and tgt length in the current batch
    global max_src_in_batch, max_mt_in_batch  # this is a hack
    # Reset current longest length at a new batch (count=1)
    if count == 1:
        max_src_in_batch = 0
        max_mt_in_batch = 0
        # max_pe_in_batch = 0
    # Src: [<bos> w1 ... wN <eos>]
    max_src_in_batch = max(max_src_in_batch, len(new.src[0]) + 2)
    # mt: [<bos> w1 ... wN <eos>]
    max_mt_in_batch = max(max_mt_in_batch, len(new.mt[0]) + 2)
    # pe: [w1 ... wM <eos>]
    src_elements = count * max_src_in_batch
    mt_elements = count * max_mt_in_batch
    return max(src_elements, mt_elements)


class PostEditor(object):
    """Translate a batch of sentences with a saved model.

    Args:
        model (onmt.APEModel): APE model to use for post-editing
        fields (dict[str, torchtext.data.Field]): A dict
            mapping each side to its list of name-Field pairs.
        src_reader (onmt.inputters.DataReaderBase): Source reader.
        mt_reader (onmt.inputters.DataReaderBase): MT output reader.
        pe_reader (onmt.inputters.TextDataReader): Post-edit reader.
        gpu (int): GPU device. Set to negative for no GPU.
        n_best (int): How many beams to wait for.
        min_length (int): See
            :class:`onmt.translate.decode_strategy.DecodeStrategy`.
        max_length (int): See
            :class:`onmt.translate.decode_strategy.DecodeStrategy`.
        beam_size (int): Number of beams.
        random_sampling_topk (int): See
            :class:`onmt.translate.greedy_search.GreedySearch`.
        random_sampling_temp (int): See
            :class:`onmt.translate.greedy_search.GreedySearch`.
        stepwise_penalty (bool): Whether coverage penalty is applied every step
            or not.
        dump_beam (bool): Debugging option.
        block_ngram_repeat (int): See
            :class:`onmt.translate.decode_strategy.DecodeStrategy`.
        ignore_when_blocking (set or frozenset): See
            :class:`onmt.translate.decode_strategy.DecodeStrategy`.
        replace_unk (bool): Replace unknown token.
        data_type (str): Source data type.
        verbose (bool): Print/log every translation.
        report_time (bool): Print/log total time/frequency.
        copy_attn (bool): Use copy attention.
        global_scorer (onmt.translate.GNMTGlobalScorer): Translation
            scoring/reranking object.
        out_file (TextIO or codecs.StreamReaderWriter): Output file.
        report_score (bool) : Whether to report scores
        logger (logging.Logger or NoneType): Logger.
    """

    def __init__(
            self,
            model,
            fields,
            src_reader,
            mt_reader,
            pe_reader,
            gpu=-1,
            n_best=1,
            min_length=0,
            max_length=100,
            ratio=0.,
            beam_size=30,
            random_sampling_topk=1,
            random_sampling_temp=1,
            stepwise_penalty=None,
            dump_beam=False,
            block_ngram_repeat=0,
            ignore_when_blocking=frozenset(),
            replace_unk=False,
            phrase_table="",
            data_type="text",
            verbose=False,
            report_time=False,
            copy_attn=False,
            global_scorer=None,
            out_file=None,
            report_align=False,
            report_score=True,
            logger=None,
            seed=-1):
        self.model = model
        self.fields = fields
        pe_field = dict(self.fields)["pe"].base_field
        self._pe_vocab = pe_field.vocab
        self._pe_eos_idx = self._pe_vocab.stoi[pe_field.eos_token]
        self._pe_pad_idx = self._pe_vocab.stoi[pe_field.pad_token]
        self._pe_bos_idx = self._pe_vocab.stoi[pe_field.init_token]
        self._pe_unk_idx = self._pe_vocab.stoi[pe_field.unk_token]
        self._pe_vocab_len = len(self._pe_vocab)

        self._gpu = gpu
        self._use_cuda = gpu > -1
        self._dev = torch.device("cuda", self._gpu) \
            if self._use_cuda else torch.device("cpu")

        self.n_best = n_best
        self.max_length = max_length

        self.beam_size = beam_size
        self.random_sampling_temp = random_sampling_temp
        self.sample_from_topk = random_sampling_topk

        self.min_length = min_length
        self.ratio = ratio
        self.stepwise_penalty = stepwise_penalty
        self.dump_beam = dump_beam
        self.block_ngram_repeat = block_ngram_repeat
        self.ignore_when_blocking = ignore_when_blocking
        self._exclusion_idxs = {
            self._pe_vocab.stoi[t] for t in self.ignore_when_blocking}
        self.src_reader = src_reader
        self.mt_reader = mt_reader
        self.pe_reader = pe_reader
        self.replace_unk = replace_unk
        if self.replace_unk and not self.model.decoder.attentional:
            raise ValueError(
                "replace_unk requires an attentional decoder.")
        self.phrase_table = phrase_table
        self.data_type = data_type
        self.verbose = verbose
        self.report_time = report_time

        self.copy_attn = copy_attn

        self.global_scorer = global_scorer
        if self.global_scorer.has_cov_pen and \
                not self.model.decoder.attentional:
            raise ValueError(
                "Coverage penalty requires an attentional decoder.")
        self.out_file = out_file
        self.report_align = report_align
        self.report_score = report_score
        self.logger = logger

        self.use_filter_pred = False
        self._filter_pred = None

        # for debugging
        self.beam_trace = self.dump_beam != ""
        self.beam_accum = None
        if self.beam_trace:
            self.beam_accum = {
                "predicted_ids": [],
                "beam_parent_ids": [],
                "scores": [],
                "log_probs": []}

        set_random_seed(seed, self._use_cuda)

    @classmethod
    def from_opt(
            cls,
            model,
            fields,
            opt,
            model_opt,
            global_scorer=None,
            out_file=None,
            report_align=False,
            report_score=True,
            logger=None):
        """Alternate constructor.

        Args:
            model (onmt.modules.NMTModel): See :func:`__init__()`.
            fields (dict[str, torchtext.data.Field]): See
                :func:`__init__()`.
            opt (argparse.Namespace): Command line options
            model_opt (argparse.Namespace): Command line options saved with
                the model checkpoint.
            global_scorer (onmt.translate.GNMTGlobalScorer): See
                :func:`__init__()`..
            out_file (TextIO or codecs.StreamReaderWriter): See
                :func:`__init__()`.
            report_align (bool) : See :func:`__init__()`.
            report_score (bool) : See :func:`__init__()`.
            logger (logging.Logger or NoneType): See :func:`__init__()`.
        """

        src_reader = inputters.str2reader[opt.data_type].from_opt(opt)
        mt_reader = inputters.str2reader[opt.data_type].from_opt(opt)
        pe_reader = inputters.str2reader["text"].from_opt(opt)
        return cls(
            model,
            fields,
            src_reader,
            mt_reader,
            pe_reader,
            gpu=opt.gpu,
            n_best=opt.n_best,
            min_length=opt.min_length,
            max_length=opt.max_length,
            ratio=opt.ratio,
            beam_size=opt.beam_size,
            random_sampling_topk=opt.random_sampling_topk,
            random_sampling_temp=opt.random_sampling_temp,
            stepwise_penalty=opt.stepwise_penalty,
            dump_beam=opt.dump_beam,
            block_ngram_repeat=opt.block_ngram_repeat,
            ignore_when_blocking=set(opt.ignore_when_blocking),
            replace_unk=opt.replace_unk,
            phrase_table=opt.phrase_table,
            data_type=opt.data_type,
            verbose=opt.verbose,
            report_time=opt.report_time,
            copy_attn=model_opt.copy_attn,
            global_scorer=global_scorer,
            out_file=out_file,
            report_align=report_align,
            report_score=report_score,
            logger=logger,
            seed=opt.seed)

    def _log(self, msg):
        if self.logger:
            self.logger.info(msg)
        else:
            print(msg)

    def _gold_score(self, batch, memory_bank, inp_lengths, inp_vocabs,
                    use_inp_map, enc_states, batch_size, inp):
        if "pe" in batch.__dict__:
            gs = self._score_post_edit(
                batch, memory_bank, inp_lengths, inp_vocabs,
                batch.src_map if use_inp_map else None)
            self.model.decoder.init_state(inp, memory_bank, enc_states)
        else:
            gs = [0] * batch_size
        return gs

    def post_edit(
            self,
            src,
            mt,
            pe=None,
            batch_size=None,
            batch_type="sents",
            attn_debug=False,
            align_debug=False,
            phrase_table=""):
        """Post-edit contents of ``mt`` and get gold scores from ``pe``.

        Args:
            src: See :func:`self.src_reader.read()`.
            mt: See :func:`self.mt_reader.read()`.
            pe: See :func:`self.pe_reader.read()`.
            batch_size (int): size of examples per mini-batch
            attn_debug (bool): enables the attention logging
            align_debug (bool): enables the word alignment logging

        Returns:
            (`list`, `list`)

            * all_scores is a list of `batch_size` lists of `n_best` scores
            * all_predictions is a list of `batch_size` lists
                of `n_best` predictions
        """

        if batch_size is None:
            raise ValueError("batch_size must be set")

        src_data = {"reader": self.src_reader, "data": src, "dir": None}
        mt_data = {"reader": self.mt_reader, "data": mt, "dir": None}
        pe_data = {"reader": self.pe_reader, "data": pe, "dir": None}
        _readers, _data, _dir = inputters.Dataset.config(
            [('src', src_data), ('mt', mt_data), ('pe', pe_data)])

        # corpus_id field is useless here
        if self.fields.get("corpus_id", None) is not None:
            self.fields.pop('corpus_id')
        data = inputters.Dataset(
            self.fields, readers=_readers, data=_data, dirs=_dir,
            sort_key=inputters.str2sortkey[self.data_type],
            filter_pred=self._filter_pred
        )

        data_iter = inputters.OrderedIterator(
            dataset=data,
            device=self._dev,
            batch_size=batch_size,
            batch_size_fn=max_tok_len if batch_type == "tokens" else None,
            train=False,
            sort=False,
            sort_within_batch=True,
            shuffle=False
        )

        xlation_builder = PostEditBuilder(
            data, self.fields, self.n_best, self.replace_unk, pe,
            self.phrase_table
        )

        # Statistics
        counter = count(1)
        pred_score_total, pred_words_total = 0, 0
        gold_score_total, gold_words_total = 0, 0

        all_scores = []
        all_predictions = []

        start_time = time.time()

        for batch in data_iter:
            batch_data = self.translate_batch(
                batch, data.src_vocabs, attn_debug
            )
            translations = xlation_builder.from_batch(batch_data)

            for trans in translations:
                all_scores += [trans.pred_scores[:self.n_best]]
                pred_score_total += trans.pred_scores[0]
                pred_words_total += len(trans.pred_sents[0])
                if pe is not None:
                    gold_score_total += trans.gold_score
                    gold_words_total += len(trans.gold_sent) + 1

                n_best_preds = [" ".join(pred)
                                for pred in trans.pred_sents[:self.n_best]]
                if self.report_align:
                    align_pharaohs = [build_align_pharaoh(align) for align
                                      in trans.word_aligns[:self.n_best]]
                    n_best_preds_align = [" ".join(align) for align
                                          in align_pharaohs]
                    n_best_preds = [pred + " ||| " + align
                                    for pred, align in zip(
                                        n_best_preds, n_best_preds_align)]
                all_predictions += [n_best_preds]
                self.out_file.write('\n'.join(n_best_preds) + '\n')
                self.out_file.flush()

                if self.verbose:
                    sent_number = next(counter)
                    output = trans.log(sent_number)
                    if self.logger:
                        self.logger.info(output)
                    else:
                        os.write(1, output.encode('utf-8'))

        end_time = time.time()

        if self.report_score:
            msg = self._report_score('PRED', pred_score_total,
                                     pred_words_total)
            self._log(msg)
            if pe is not None:
                msg = self._report_score('GOLD', gold_score_total,
                                         gold_words_total)
                self._log(msg)

        if self.report_time:
            total_time = end_time - start_time
            self._log("Total translation time (s): %f" % total_time)
            self._log("Average translation time (s): %f" % (
                total_time / len(all_predictions)))
            self._log("Tokens per second: %f" % (
                pred_words_total / total_time))

        if self.dump_beam:
            import json
            json.dump(self.translator.beam_accum,
                      codecs.open(self.dump_beam, 'w', 'utf-8'))
        return all_scores, all_predictions

    def _align_pad_prediction(self, predictions, bos, pad):
        """
        Padding predictions in batch and add BOS.

        Args:
            predictions (List[List[Tensor]]): `(batch, n_best,)`, for each inp
                sequence contain n_best pe predictions all of which ended with
                eos id.
            bos (int): bos index to be used.
            pad (int): pad index to be used.

        Return:
            batched_nbest_predict (torch.LongTensor): `(batch, n_best, pe_l)`
        """
        dtype, device = predictions[0][0].dtype, predictions[0][0].device
        flatten_pe = [best.tolist() for bests in predictions
                       for best in bests]
        paded_pe = torch.tensor(
            list(zip_longest(*flatten_pe, fillvalue=pad)),
            dtype=dtype, device=device).T
        bos_tensor = torch.full([paded_pe.size(0), 1], bos,
                                dtype=dtype, device=device)
        full_pe = torch.cat((bos_tensor, paded_pe), dim=-1)
        batched_nbest_predict = full_pe.view(
            len(predictions), -1, full_pe.size(-1))  # (batch, n_best, pe_l)
        return batched_nbest_predict

    def _align_forward(self, batch, predictions):
        """
        For a batch of input and its prediction, return a list of batch predict
        alignment enc_inp indice Tensors in size ``(batch, n_best,)``.
        """
        # (0) add BOS and padding to pe prediction
        if hasattr(batch, 'pe'):
            batch_pe_idxs = batch.pe.transpose(1, 2).transpose(0, 2)
        else:
            batch_pe_idxs = self._align_pad_prediction(
                predictions, bos=self._pe_bos_idx, pad=self._pe_pad_idx)
        pe_mask = (batch_pe_idxs.eq(self._pe_pad_idx) |
                   batch_pe_idxs.eq(self._pe_eos_idx) |
                   batch_pe_idxs.eq(self._pe_bos_idx))

        n_best = batch_pe_idxs.size(1)
        # (1) Encoder forward.
        enc_out, enc_states, memory_bank, inp_lengths = self._run_encoder(batch)

        # (2) Repeat enc_out objects `n_best` times.
        # We use batch_size x n_best, get ``(src_len/mt_len, batch * n_best, nfeat)``
        for inp in enc_out.keys():
            enc_out[inp] = tile(enc_out[inp], n_best, dim=1)
            enc_states[inp] = tile(enc_states[inp], n_best, dim=1)
            if isinstance(memory_bank[inp], tuple):
                memory_bank[inp] = tuple(tile(x, n_best, dim=1) \
                                   for x in memory_bank[inp])
            else:
                memory_bank[inp] = tile(memory_bank[inp], n_best, dim=1)
            inp_lengths[inp] = tile(inp_lengths[inp], n_best)  # ``(batch * n_best,)``

        # (3) Init decoder with n_best enc_out,
        self.model.decoder.init_state(enc_out, memory_bank, enc_states)
        # reshape pe to ``(len, batch * n_best, nfeat)``
        pe = batch_pe_idxs.view(-1, batch_pe_idxs.size(-1)).T.unsqueeze(-1)
        dec_in = pe[:-1]  # exclude last target from inputs
        _, attns = self.model.decoder(
            dec_in, memory_bank, memory_lengths=inp_lengths, with_align=True)

        alignment_attn = attns['align']  # ``(B, pe_len-1, enc_len)`` each
        # masked_select
        align_pe_mask = pe_mask.view(-1, pe_mask.size(-1))
        prediction_mask = align_pe_mask[:, 1:]  # exclude bos to match pred
        # get aligned id for each prediction's valid pe tokens
        alignement = extract_alignment(
            alignment_attn, prediction_mask, inp_lengths, n_best)
        return alignement

    def translate_batch(self, batch, enc_vocabs, attn_debug):
        """Translate a batch of sentences."""
        with torch.no_grad():
            if self.beam_size == 1:
                decode_strategy = GreedySearch(
                    pad=self._pe_pad_idx,
                    bos=self._pe_bos_idx,
                    eos=self._pe_eos_idx,
                    batch_size=batch.batch_size,
                    min_length=self.min_length, max_length=self.max_length,
                    block_ngram_repeat=self.block_ngram_repeat,
                    exclusion_tokens=self._exclusion_idxs,
                    return_attention=attn_debug or self.replace_unk,
                    sampling_temp=self.random_sampling_temp,
                    keep_topk=self.sample_from_topk)
            else:
                # TODO: support these blacklisted features
                assert not self.dump_beam
                decode_strategy = BeamSearch(
                    self.beam_size,
                    batch_size=batch.batch_size,
                    pad=self._pe_pad_idx,
                    bos=self._pe_bos_idx,
                    eos=self._pe_eos_idx,
                    n_best=self.n_best,
                    global_scorer=self.global_scorer,
                    min_length=self.min_length, max_length=self.max_length,
                    return_attention=attn_debug or self.replace_unk,
                    block_ngram_repeat=self.block_ngram_repeat,
                    exclusion_tokens=self._exclusion_idxs,
                    stepwise_penalty=self.stepwise_penalty,
                    ratio=self.ratio)
            return self._translate_batch_with_strategy(batch, enc_vocabs,
                                                       decode_strategy)

    def _run_encoder(self, batch):
        src, src_lengths = batch.src if isinstance(batch.src, tuple) \
                           else (batch.src, None)
        mt, mt_lengths = batch.mt if isinstance(batch.mt, tuple) \
                           else (batch.mt, None)
        inp = {'src': src, 'mt': mt}
        inp_lengths = {'src': src_lengths, 'mt': mt_lengths}

        enc_states, memory_bank, inp_lengths = self.model.encoder(
            inp, inp_lengths)
        if inp_lengths['src'] == None:
            assert not isinstance(memory_bank['src'], tuple), \
                'Ensemble decoding only supported for text data'
            inp_lengths['src'] = torch.Tensor(batch.batch_size) \
                                      .type_as(memory_bank['src']) \
                                      .long() \
                                      .fill_(memory_bank['src'].size(0))
        if inp_lengths['mt'] == None:
            inp_lengths['mt'] = torch.Tensor(batch.batch_size) \
                                      .type_as(memory_bank['mt']) \
                                      .long() \
                                      .fill_(memory_bank['mt'].size(0))
        return inp, enc_states, memory_bank, inp_lengths

    def _decode_and_generate(
            self,
            decoder_in,
            memory_bank,
            batch,
            inp_vocabs,
            memory_lengths,
            inp_map=None,
            step=None,
            batch_offset=None):

        # Decoder forward, takes [pe_len, batch, nfeats] as input
        # and a dictionary of [src_len or mt_len, batch, hidden] as memory_bank
        # in case of inference pe_len = 1, batch = beam times batch_size
        # in case of Gold Scoring pe_len = actual length, batch = 1 batch
        dec_out, dec_attn = self.model.decoder(
            decoder_in, memory_bank, memory_lengths=memory_lengths, step=step
        )

        # Generator forward.
        if "std" in dec_attn:
            attn = dec_attn["std"]
        else:
            attn = None
        log_probs = self.model.generator(dec_out.squeeze(0))
        # returns [(batch_size x beam_size) , vocab ] when 1 step
        # or [ pe_len, batch_size, vocab ] when full sentence
        return log_probs, attn

    def _translate_batch_with_strategy(
            self,
            batch,
            inp_vocabs,
            decode_strategy):
        """Translate a batch of sentences step by step using cache.

        Args:
            batch: a batch of sentences, yield by data iterator.
            inp_vocabs (Dict[str, list]): A dictionary of lists of
                torchtext.data.Vocab if can_copy.
            decode_strategy (DecodeStrategy): A decode strategy to use for
                generate translation step by step.

        Returns:
            results (dict): The translation results.
        """
        # (0) Prep the components of the search.
        parallel_paths = decode_strategy.parallel_paths  # beam_size
        batch_size = batch.batch_size

        # (1) Run the encoder.
        inp, enc_states, memory_bank, inp_lengths = self._run_encoder(batch)
        self.model.decoder.init_state(inp, memory_bank, enc_states)

        results = {
            "predictions": None,
            "scores": None,
            "attention": None,
            "batch": batch,
            "gold_score": self._gold_score(
                batch, memory_bank, inp_lengths, inp_vocabs, False,
                enc_states, batch_size, inp)}

        # (2) prep decode_strategy. Possibly repeat inp objects.
        inp_map = None
        fn_map_state, memory_bank, memory_lengths, inp_map = \
            decode_strategy.initialize(memory_bank, inp_lengths, inp_map)
        if fn_map_state is not None:
            self.model.decoder.map_state(fn_map_state)

        # (3) Begin decoding step by step:
        for step in range(decode_strategy.max_length):
            decoder_input = decode_strategy.current_predictions.view(1, -1, 1)

            log_probs, attn = self._decode_and_generate(
                decoder_input,
                memory_bank,
                batch,
                inp_vocabs,
                memory_lengths=memory_lengths,
                inp_map=inp_map,
                step=step,
                batch_offset=decode_strategy.batch_offset)

            decode_strategy.advance(log_probs, attn)
            any_finished = decode_strategy.is_finished.any()
            if any_finished:
                decode_strategy.update_finished()
                if decode_strategy.done:
                    break

            select_indices = decode_strategy.select_indices

            if any_finished:
                # Reorder states.
                for inp in memory_bank.keys():
                    if isinstance(memory_bank, tuple):
                        memory_bank[inp] = tuple(x.index_select(1, select_indices)
                                            for x in memory_bank[inp])
                    else:
                        memory_bank[inp] = memory_bank[inp].index_select(1, select_indices)
                    memory_lengths[inp] = memory_lengths[inp].index_select(0, select_indices)

                    if inp_map is not None:
                        inp_map[inp] = inp_map[inp].index_select(1, select_indices)

            if parallel_paths > 1 or any_finished:
                self.model.decoder.map_state(
                    lambda state, dim: state.index_select(dim, select_indices))

        results["scores"] = decode_strategy.scores
        results["predictions"] = decode_strategy.predictions
        results["attention"] = decode_strategy.attention
        if self.report_align:
            results["alignment"] = self._align_forward(
                batch, decode_strategy.predictions)
        else:
            results["alignment"] = [[] for _ in range(batch_size)]
        return results

    def _score_post_edit(self, batch, memory_bank, inp_lengths,
                         inp_vocabs, inp_map):
        pe = batch.pe
        pe_in = pe[:-1]

        log_probs, attn = self._decode_and_generate(
            pe_in, memory_bank, batch, inp_vocabs,
            memory_lengths=inp_lengths, inp_map=inp_map)

        log_probs[:, :, self._pe_pad_idx] = 0
        gold = pe[1:]
        gold_scores = log_probs.gather(2, gold)
        gold_scores = gold_scores.sum(dim=0).view(-1)

        return gold_scores

    def _report_score(self, name, score_total, words_total):
        if words_total == 0:
            msg = "%s No words predicted" % (name,)
        else:
            avg_score = score_total / words_total
            ppl = np.exp(-score_total.item() / words_total)
            msg = ("%s AVG SCORE: %.4f, %s PPL: %.4f" % (
                name, avg_score,
                name, ppl))
        return msg
#!/usr/bin/env python PosetEditor Class and builder In token batching scheme, the number of sequences is limited     such that the total number of src, mt, and pe tokens (including padding)     in a batch <= batch_size # Maintains the longest src and tgt length in the current batch # this is a hack # Reset current longest length at a new batch (count=1) # max_pe_in_batch = 0 # Src: [<bos> w1 ... wN <eos>] # mt: [<bos> w1 ... wN <eos>] # pe: [w1 ... wM <eos>] Translate a batch of sentences with a saved model.      Args:         model (onmt.APEModel): APE model to use for post-editing         fields (dict[str, torchtext.data.Field]): A dict             mapping each side to its list of name-Field pairs.         src_reader (onmt.inputters.DataReaderBase): Source reader.         mt_reader (onmt.inputters.DataReaderBase): MT output reader.         pe_reader (onmt.inputters.TextDataReader): Post-edit reader.         gpu (int): GPU device. Set to negative for no GPU.         n_best (int): How many beams to wait for.         min_length (int): See             :class:`onmt.translate.decode_strategy.DecodeStrategy`.         max_length (int): See             :class:`onmt.translate.decode_strategy.DecodeStrategy`.         beam_size (int): Number of beams.         random_sampling_topk (int): See             :class:`onmt.translate.greedy_search.GreedySearch`.         random_sampling_temp (int): See             :class:`onmt.translate.greedy_search.GreedySearch`.         stepwise_penalty (bool): Whether coverage penalty is applied every step             or not.         dump_beam (bool): Debugging option.         block_ngram_repeat (int): See             :class:`onmt.translate.decode_strategy.DecodeStrategy`.         ignore_when_blocking (set or frozenset): See             :class:`onmt.translate.decode_strategy.DecodeStrategy`.         replace_unk (bool): Replace unknown token.         data_type (str): Source data type.         verbose (bool): Print/log every translation.         report_time (bool): Print/log total time/frequency.         copy_attn (bool): Use copy attention.         global_scorer (onmt.translate.GNMTGlobalScorer): Translation             scoring/reranking object.         out_file (TextIO or codecs.StreamReaderWriter): Output file.         report_score (bool) : Whether to report scores         logger (logging.Logger or NoneType): Logger. # for debugging Alternate constructor.          Args:             model (onmt.modules.NMTModel): See :func:`__init__()`.             fields (dict[str, torchtext.data.Field]): See                 :func:`__init__()`.             opt (argparse.Namespace): Command line options             model_opt (argparse.Namespace): Command line options saved with                 the model checkpoint.             global_scorer (onmt.translate.GNMTGlobalScorer): See                 :func:`__init__()`..             out_file (TextIO or codecs.StreamReaderWriter): See                 :func:`__init__()`.             report_align (bool) : See :func:`__init__()`.             report_score (bool) : See :func:`__init__()`.             logger (logging.Logger or NoneType): See :func:`__init__()`. Post-edit contents of ``mt`` and get gold scores from ``pe``.          Args:             src: See :func:`self.src_reader.read()`.             mt: See :func:`self.mt_reader.read()`.             pe: See :func:`self.pe_reader.read()`.             batch_size (int): size of examples per mini-batch             attn_debug (bool): enables the attention logging             align_debug (bool): enables the word alignment logging          Returns:             (`list`, `list`)              * all_scores is a list of `batch_size` lists of `n_best` scores             * all_predictions is a list of `batch_size` lists                 of `n_best` predictions # corpus_id field is useless here # Statistics Padding predictions in batch and add BOS.          Args:             predictions (List[List[Tensor]]): `(batch, n_best,)`, for each inp                 sequence contain n_best pe predictions all of which ended with                 eos id.             bos (int): bos index to be used.             pad (int): pad index to be used.          Return:             batched_nbest_predict (torch.LongTensor): `(batch, n_best, pe_l)` # (batch, n_best, pe_l) For a batch of input and its prediction, return a list of batch predict         alignment enc_inp indice Tensors in size ``(batch, n_best,)``. # (0) add BOS and padding to pe prediction # (1) Encoder forward. # (2) Repeat enc_out objects `n_best` times. # We use batch_size x n_best, get ``(src_len/mt_len, batch * n_best, nfeat)`` # ``(batch * n_best,)`` # (3) Init decoder with n_best enc_out, # reshape pe to ``(len, batch * n_best, nfeat)`` # exclude last target from inputs # ``(B, pe_len-1, enc_len)`` each # masked_select # exclude bos to match pred # get aligned id for each prediction's valid pe tokens Translate a batch of sentences. # TODO: support these blacklisted features # Decoder forward, takes [pe_len, batch, nfeats] as input # and a dictionary of [src_len or mt_len, batch, hidden] as memory_bank # in case of inference pe_len = 1, batch = beam times batch_size # in case of Gold Scoring pe_len = actual length, batch = 1 batch # Generator forward. # returns [(batch_size x beam_size) , vocab ] when 1 step # or [ pe_len, batch_size, vocab ] when full sentence Translate a batch of sentences step by step using cache.          Args:             batch: a batch of sentences, yield by data iterator.             inp_vocabs (Dict[str, list]): A dictionary of lists of                 torchtext.data.Vocab if can_copy.             decode_strategy (DecodeStrategy): A decode strategy to use for                 generate translation step by step.          Returns:             results (dict): The translation results. # (0) Prep the components of the search. # beam_size # (1) Run the encoder. # (2) prep decode_strategy. Possibly repeat inp objects. # (3) Begin decoding step by step: # Reorder states.