<commit_before>"""Running of a recurrent decoder.

This module aggragates what is necessary to run efficiently a recurrent
decoder. Unlike the default runner which assumes all outputs are independent on
each other, this one does not make any of these assumptions. It implements
model ensembling and beam search.

The TensorFlow session is invoked for every single output of the decoder
separately which allows ensembling from all sessions and do the beam pruning
before the a next output is emmited.
"""


from typing import List, NamedTuple, Tuple
import numpy as np
import tensorflow as tf

from neuralmonkey.tf_manager import RunResult
from neuralmonkey.runners.base_runner import BaseRunner, \
    Executable, ExecutionResult, NextExecute
from neuralmonkey.vocabulary import END_TOKEN_INDEX

# tests: mypy,pylint


# pylint: disable=invalid-name
BeamBatch = NamedTuple('BeamBatch',
                       [('decoded', np.ndarray),
                        ('logprobs', np.ndarray)])
ExpandedBeamBatch = NamedTuple('ExpandedBeamBatch',
                               [('beam_batch', BeamBatch),
                                ('next_logprobs', np.ndarray)])

# pylint: disable=too-many-locals


def _score_expanded(n: int,
                    batch_size: int,
                    expanded: List[ExpandedBeamBatch],
                    scoring_function) -> \
        Tuple[List[np.ndarray], List[np.ndarray]]:
    next_beam_hypotheses = []
    next_beam_logprobs = []
    # agregate the expanded hypotheses hypothsis-wise
    for rank in range(batch_size):
        candidate_scores = None
        candidate_hypotheses = None
        candidate_logprobs = None

        for expanded_batch in expanded:
            next_distribution = expanded_batch.next_logprobs[rank]
            if expanded_batch.beam_batch is None:
                expanded_hypotheses = np.expand_dims(np.arange(
                    len(next_distribution)), axis=1)
                expanded_logprobs = np.expand_dims(next_distribution, 1)
            else:
                hypothesis = expanded_batch.beam_batch.decoded[rank]
                prev_logprobs = expanded_batch.beam_batch.logprobs[rank]

                expanded_hypotheses = np.array(
                    [np.append(hypothesis, index)
                     for index, _ in enumerate(next_distribution)])
                expanded_logprobs = np.array(
                    [np.append(prev_logprobs, logprob)
                     for logprob in next_distribution])

            scores = scoring_function(expanded_hypotheses, expanded_logprobs)
            n_best_indices = np.argsort(scores)[-n:]
            candidate_scores = _try_append(
                candidate_scores, scores[n_best_indices])
            candidate_hypotheses = _try_append(
                candidate_hypotheses, expanded_hypotheses[n_best_indices])
            candidate_logprobs = _try_append(
                candidate_logprobs, expanded_logprobs[n_best_indices])

        n_best_indices = np.argsort(candidate_scores)[-n:]
        n_best_hypotheses = candidate_hypotheses[n_best_indices]
        n_best_logprobs = candidate_logprobs[n_best_indices]

        next_beam_hypotheses.append(n_best_hypotheses)
        next_beam_logprobs.append(n_best_logprobs)
    return next_beam_hypotheses, next_beam_logprobs


def _try_append(first, second):
    if first is None:
        return second
    else:
        return np.append(first, second)


def likelihood_beam_score(decoded, logprobs):
    """Score the beam by normalized probaility."""

    mask = []
    for hypothesis in decoded:
        before_end = True
        hyp_mask = []
        for index in hypothesis:
            hyp_mask.append(float(before_end))
            before_end &= (index != END_TOKEN_INDEX)
        mask.append(hyp_mask)
<commit_msg>attempt to make beamsearch faster<commit_after>"""Running of a recurrent decoder.

This module aggragates what is necessary to run efficiently a recurrent
decoder. Unlike the default runner which assumes all outputs are independent on
each other, this one does not make any of these assumptions. It implements
model ensembling and beam search.

The TensorFlow session is invoked for every single output of the decoder
separately which allows ensembling from all sessions and do the beam pruning
before the a next output is emmited.
"""


from typing import List, NamedTuple, Tuple
import numpy as np
import tensorflow as tf

from neuralmonkey.tf_manager import RunResult
from neuralmonkey.runners.base_runner import BaseRunner, \
    Executable, ExecutionResult, NextExecute
from neuralmonkey.vocabulary import END_TOKEN_INDEX

# tests: mypy,pylint


# pylint: disable=invalid-name
BeamBatch = NamedTuple('BeamBatch',
                       [('decoded', np.ndarray),
                        ('logprobs', np.ndarray)])
ExpandedBeamBatch = NamedTuple('ExpandedBeamBatch',
                               [('beam_batch', BeamBatch),
                                ('next_logprobs', np.ndarray)])

# pylint: disable=too-many-locals


def _n_best_indices(scores: np.ndarray, n: int) -> np.ndarray:
    if scores.shape[0] <= n:
        unsorted_n_best_indices = np.arange(scores.shape[0])
    else:
        unsorted_n_best_indices = np.argpartition(-scores, n)[:n]
    n_best_indices = unsorted_n_best_indices[
        np.argsort(scores[unsorted_n_best_indices])]
    return n_best_indices


def _score_expanded(n: int,
                    batch_size: int,
                    expanded: List[ExpandedBeamBatch],
                    scoring_function) -> \
        Tuple[List[np.ndarray], List[np.ndarray]]:
    next_beam_hypotheses = []
    next_beam_logprobs = []
    # agregate the expanded hypotheses hypothsis-wise
    for rank in range(batch_size):
        candidate_scores = None
        candidate_hypotheses = None
        candidate_logprobs = None

        for expanded_batch in expanded:
            next_distribution = expanded_batch.next_logprobs[rank]
            if expanded_batch.beam_batch is None:
                expanded_hypotheses = np.expand_dims(np.arange(
                    len(next_distribution)), axis=1)
                expanded_logprobs = np.expand_dims(next_distribution, 1)
            else:
                hypothesis = expanded_batch.beam_batch.decoded[rank]
                prev_logprobs = expanded_batch.beam_batch.logprobs[rank]

                expanded_hypotheses = np.array(
                    [np.append(hypothesis, index)
                     for index, _ in enumerate(next_distribution)])
                expanded_logprobs = np.array(
                    [np.append(prev_logprobs, logprob)
                     for logprob in next_distribution])

            scores = scoring_function(expanded_hypotheses, expanded_logprobs)
            n_best_indices = _n_best_indices(scores, n)
            candidate_scores = _try_append(
                candidate_scores, scores[n_best_indices])
            candidate_hypotheses = _try_append(
                candidate_hypotheses, expanded_hypotheses[n_best_indices])
            candidate_logprobs = _try_append(
                candidate_logprobs, expanded_logprobs[n_best_indices])

        n_best_indices = _n_best_indices(candidate_scores, n)
        n_best_hypotheses = candidate_hypotheses[n_best_indices]
        n_best_logprobs = candidate_logprobs[n_best_indices]

        next_beam_hypotheses.append(n_best_hypotheses)
        next_beam_logprobs.append(n_best_logprobs)
    return next_beam_hypotheses, next_beam_logprobs


def _try_append(first, second):
    if first is None:
        return second
    else:
        return np.append(first, second)


def likelihood_beam_score(decoded, logprobs):
    """Score the beam by normalized probaility."""

    mask = []
    for hypothesis in decoded:
        before_end = True
        hyp_mask = []
        for index in hypothesis:
            hyp_mask.append(float(before_end))
            before_end &= (index != END_TOKEN_INDEX)
        mask.append(hyp_mask)

{'fraction_non_alphanumeric': 0.0708927231807952, 'fraction_numerical': 0.0007501875468867217, 'mean_word_length': 2.6080288678394226, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '108622', 'n_tokens_mistral': 2408, 'n_tokens_neox': 2277, 'n_words': 538}