<commit_before>import numpy as np
import tensorflow as tf
from nltk.tokenize import word_tokenize
from termcolor import colored
import regex as re

from utils import log


def load_tokenized(text_file, preprocess=None):
    """
    Loads a tokenized text file a list of list of tokens.

    Args:

        text_file: An opened file.

        preprocess: A function/callable that (linguistically) preprocesses the
            sentences

    """

    if not preprocess:
        preprocess = lambda x: x

    return [preprocess(re.split(ur"[ ]", l.rstrip())) for l in text_file]


def tokenize_char_seq(chars):
    return word_tokenize("".join(chars))


def feed_dicts(dataset, batch_size, coders, train=False):
    """

    This function ensures all encoder and decoder objects feed their the data
    they need from the dataset.
    """
    dicts = [{} for _ in range(len(dataset) / batch_size + int(len(dataset) % batch_size > 0))]

    for coder in coders:
        coder.feed_dict(dataset, batch_size, dicts=dicts, train=train)

    return dicts

# TODO postprocess, copynet, beamsearch will be hidden in runner

def training_loop(sess, epochs, trainer, all_coders, decoder, batch_size,
                  train_dataset, val_dataset,
                  postprocess, log_directory,
                  evaluation_functions,
                  use_copynet=False,
                  test_dataset=None,
                  use_beamsearch=False,
                  initial_variables=None,
                  test_run=False):

    """

    Performs the training loop for given graph and data.

    Args:

        sess: TF Session.

        epochs: Number of epochs for which the algoritm will learn.
<commit_msg>refactor beamsearch<commit_after>import numpy as np
import tensorflow as tf
from nltk.tokenize import word_tokenize
from termcolor import colored
import regex as re

from utils import log


def load_tokenized(text_file, preprocess=None):
    """
    Loads a tokenized text file a list of list of tokens.

    Args:

        text_file: An opened file.

        preprocess: A function/callable that (linguistically) preprocesses the
            sentences

    """

    if not preprocess:
        preprocess = lambda x: x

    return [preprocess(re.split(ur"[ ]", l.rstrip())) for l in text_file]


def tokenize_char_seq(chars):
    return word_tokenize("".join(chars))


def feed_dicts(dataset, batch_size, coders, train=False):
    """

    This function ensures all encoder and decoder objects feed their the data
    they need from the dataset.
    """
    dicts = [{} for _ in range(len(dataset) / batch_size + int(len(dataset) % batch_size > 0))]

    for coder in coders:
        coder.feed_dict(dataset, batch_size, dicts=dicts, train=train)

    return dicts

# TODO postprocess, copynet, beamsearch will be hidden in runner
  
def expand(session, decoder, feed_dict, state, hypotheses):
    feed_dict[decoder.encoded] = state
    hyp_length = len(hypotheses[0][1])
    hyp_count = len(hypotheses)
    if hyp_length == 2:
        for k in feed_dict:
            sh = k.get_shape()
            if not sh == tf.TensorShape(None):
                if len(sh) == 1:
                    feed_dict[k] = np.repeat(feed_dict[k], hyp_count)
                elif len(sh) == 2:
                    feed_dict[k] = np.repeat(np.array(feed_dict[k]), hyp_count, axis=0)
                else:
                    log("ERROR in expanding beamsearch hypothesis")
    elif hyp_length > 2:
        feed_dict[decoder.encoded] = np.repeat(state, hyp_count, axis=0)
    
    for i, n in zip(decoder.gt_inputs, range(hyp_length)):
        for k in range(hyp_count):
            feed_dict[i][k] = hypotheses[k][1][n]
    probs, prob_i = session.run([decoder.top10_probs[hyp_length - 1][0],
                     decoder.top10_probs[hyp_length - 1][1]],
                     feed_dict=feed_dict)
    beam = []
    for i in range(hyp_count):
        for p, x in zip(probs[i], prob_i[i]):
            beam.append((hypotheses[i][0] + p, hypotheses[i][1] + [x]))
    return beam

def beamsearch(session, decoder, feed_dict):
    beam = [(1.0, [1])]
    state = session.run(decoder.encoded, feed_dict)
    for _ in range(len(decoder.decoded_probs)):
         new_beam = expand(session, decoder, feed_dict, state, beam)
         new_beam.sort(reverse=True)
         beam = new_beam[:10]
    return beam[0][1]

def training_loop(sess, epochs, trainer, all_coders, decoder, batch_size,
                  train_dataset, val_dataset,
                  postprocess, log_directory,
                  evaluation_functions,
                  use_copynet=False,
                  test_dataset=None,
                  use_beamsearch=False,
                  initial_variables=None,
                  test_run=False):

    """

    Performs the training loop for given graph and data.

    Args:

        sess: TF Session.

        epochs: Number of epochs for which the algoritm will learn.

{'fraction_non_alphanumeric': 0.09272727272727273, 'fraction_numerical': 0.0056565656565656566, 'mean_word_length': 2.4334257975034674, 'pattern_counts': {'":': 0, '<': 3, '<?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': '3170944', 'n_tokens_mistral': 1569, 'n_tokens_neox': 1467, 'n_words': 428}