class Configuration(object):
class Configurable(object):
class EncodeDescription(json.JSONEncoder):
    def __init__(self, name, version, params):
    def __str__(self):
    def __eq__(self, other):
    def _get_param_names(cls):
    def name(self):
    def version(self):
    def get_params(self):
    def get_config(self) -> Configuration:
    def __getstate__(self):
    def __setstate__(self, state):
def describe(obj):
    def default(self, obj):
def config_to_json(data, indent=None): class Configurable(object):
    """
    Configurable classes have names, versions, and a set of parameters that are either "simple" aka JSON serializable
    types or other Configurable objects. Configurable objects should also be serializable via pickle.
    Configurable classes are defined mainly to give us a human-readable way of reading of the `parameters`
    set for different objects and to attach version numbers to them.

    By default we follow the format sklearn uses for its `BaseEstimator` class, where parameters are automatically
    derived based on the constructor parameters.
    """

    @classmethod
    def _get_param_names(cls):
        # fetch the constructor or the original constructor before
        init = cls.__init__
        if init is object.__init__:
            # No explicit constructor to introspect
            return []

        init_signature = signature(init)
        parameters = [p for p in init_signature.parameters.values()
                      if p.name != 'self']
        if any(p.kind == p.VAR_POSITIONAL for p in parameters):
            raise RuntimeError()
        return sorted([p.name for p in parameters])

    @property
    def name(self):
        return self.__class__.__name__

    @property
    def version(self):
        return 0

    def get_params(self):
        out = {}
        for key in self._get_param_names():
            v = getattr(self, key, None)
            if isinstance(v, Configurable):
                out[key] = v.get_config()
            elif hasattr(v, "get_config"):  # for keras objects
                out[key] = {"name": v.__class__.__name__, "config": v.get_config()}
            else:
                out[key] = v
        return out

    def get_config(self) -> Configuration:
        params = {k: describe(v) for k,v in self.get_params().items()}
        return Configuration(self.name, self.version, params)

    def __getstate__(self):
        state = dict(self.__dict__)
        if "version" in state:
            if state["version"] != self.version:
                raise RuntimeError()
        else:
            state["version"] = self.version
        return state

    def __setstate__(self, state):
        if "version" not in state:
            raise RuntimeError("Version should be in state (%s)" % self.__class__.__name__)
        if state["version"] != self.version:
            warn(("%s loaded with version %s, but class " +
                 "version is %s") % (self.__class__.__name__, state["version"], self.version))

        if "state" in state:
            self.__dict__ = state["state"]
        else:
            del state["version"]
            self.__dict__ = state class PreprocessedData(TrainingData):
    """
    Data the goes through a preprocessing pipeline, for TriviaQA this usually mean leading/choosing what
    paragraphs we want to train on, the organizing them into a dataset with the desired sampling strategy
    """

    def __init__(self,
                 corpus,
                 preprocesser: Optional[Preprocessor],
                 builder: DatasetBuilder,
                 eval_builder: DatasetBuilder,
                 eval_on_verified: bool=True,
                 eval_on_train: bool = True,
                 hold_out_train: Optional[Tuple[int, int]]= None,
                 sample=None, sample_dev=None,
                 sample_preprocessed_train=None, sample_seed=None):
        self.hold_out_train = hold_out_train
        self.eval_on_train = eval_on_train
        self.sample = sample
        self.eval_on_verified = eval_on_verified
        self.sample_dev = sample_dev
        self.corpus = corpus
        self.preprocesser = preprocesser
        self.builder = builder
        self.eval_builder = eval_builder
        self.sample_preprocessed_train = sample_preprocessed_train
        self.sample_seed = sample_seed

        self._train = None
        self._dev = None
        self._verified_dev = None

    @property
    def name(self):
        return self.corpus.name

    def cache_preprocess(self, filename):
        if self.sample is not None or self.sample_dev is not None or self.hold_out_train is not None:
            raise ValueError()
        if filename.endswith("gz"):
            handle = lambda a,b: gzip.open(a, b, compresslevel=3)
        else:
            handle = open
        with handle(filename, "wb") as f:
            pickle.dump([self.preprocesser, self._train, self._dev, self._verified_dev], f)

    def load_preprocess(self, filename):
        print("Loading preprocessed data...")
        if filename.endswith("gz"):
            handle = gzip.open
        else:
            handle = open
        with handle(filename, "rb") as f:
            stored = pickle.load(f)
            stored_preprocesser, self._train, self._dev, self._verified_dev = stored
        if stored_preprocesser.get_config() != self.preprocesser.get_config():
            # print("WARNING")
            import code
            code.interact(local=locals())
            raise ValueError()

        print("done")

    def preprocess(self, n_processes=1, chunk_size=500):
        if self._train is not None:
            return
        print("Loading data...")
        train_questions = self.corpus.get_train()
        if self.hold_out_train is not None:
            print("Using held out train")
            train_questions.sort(key=lambda q:q.question_id)
            np.random.RandomState(self.hold_out_train[0]).shuffle(train_questions)
            dev_questions = train_questions[:self.hold_out_train[1]]
            train_questions = train_questions[self.hold_out_train[1]:]
        else:
            dev_questions = self.corpus.get_dev()
        if self.eval_on_verified and hasattr(self.corpus, "get_verified"):  # TODO this is a bit hacky
            verified_questions = self.corpus.get_verified()
            if verified_questions is not None:
                # we don't eval on verified docs w/o any valid human answer
                verified_questions = [x for x in verified_questions if
                                      any(len(ans) > 0 for ans in x.answer.human_answers)]
        else:
            verified_questions = None

        rng = np.random.RandomState(self.sample_seed)

        if self.sample is not None:
            l = len(train_questions)
            train_questions = rng.choice(train_questions, self.sample, replace=False)
            print("Sampled %d of %d (%.4f) train questions" % (len(train_questions), l, len(train_questions)/l))
        if self.sample_dev is not None:
            l = len(dev_questions)
            dev_questions = np.random.RandomState(self.sample_seed).choice(dev_questions, self.sample_dev, replace=False)
            print("Sampled %d of %d (%.4f) dev questions" % (len(dev_questions), l, len(dev_questions) / l))

        if self.preprocesser:
            print("Preprocessing with %d processes..." % n_processes)
            out = []
            for name, questions in [("verified", verified_questions),
                                    ("dev", dev_questions),
                                    ("train", train_questions)]:
                if questions is None:
                    out.append(None)
                    continue
                data = preprocess_par(questions, self.corpus.evidence, self.preprocesser, n_processes, chunk_size, name)
                out.append(data)
            self._verified_dev, self._dev, self._train = out
        else:
            self._verified_dev, self._dev, self._train = verified_questions, dev_questions, train_questions

        if self.sample_preprocessed_train:
            if isinstance(self._train, FilteredData):
                l =  len(self._train.data)
                self._train.data = rng.choice(self._train.data, self.sample_preprocessed_train, False)
                self._train.true_len *= len(self._train.data)/l
                print("Sampled %d of %d (%.4f) q-c pairs" % (len(self._train.data), l, len(self._train.data)/l))
            else:
                l = len(self._train)
                self._train = rng.choice(self._train, self.sample_preprocessed_train, False)
                print("Sampled %d of %d q-c pairs" % (len(self._train), l))

        print("Done")

    def get_train(self) -> Dataset:
        return self.builder.build_dataset(self._train, self.corpus)

    def get_train_corpus(self):
        return self.builder.build_stats(self._train)

    def get_eval(self) -> Dict[str, Dataset]:
        corpus = self.corpus
        eval_set = dict(dev=self.eval_builder.build_dataset(self._dev, corpus))
        if self.eval_on_train:
            eval_set["train"] = self.eval_builder.build_dataset(self._train, corpus)
        if self.eval_on_verified:
            eval_set["verified-dev"] = self.eval_builder.build_dataset(self._verified_dev, corpus)
        return eval_set

    def get_resource_loader(self) -> ResourceLoader:
        return self.corpus.get_resource_loader()

    def __setstate__(self, state):
        if "sample_seed" not in state:
            state["sample_seed"] = None
        if "sample_preprocessed_train" not in state:
            state["sample_preprocessed_train"] = None
        self.__dict__ = state

    def __getstate__(self):
        state = dict(self.__dict__)
        state["_train"] = None
        state["_dev"] = None
        state["_verified_dev"] = None
        return state class TrainingData(Configurable):

    def get_train(self) -> Dataset:
        raise NotImplementedError()

    def get_eval(self) -> Dict[str, Dataset]:
        raise NotImplementedError()

    def get_train_corpus(self) -> object:
        """
        Return an object derived from the training data that will be passed to the model's initialization phase,
        what exactly is returned can be arbitrary, but will need to be compatible with
        the model's requirements. Example, return word counts to be used to decide what word vecs to train.
        """
        raise NotImplementedError()

    def get_resource_loader(self) -> ResourceLoader:
        return ResourceLoader() class Dataset(object):
    """ Data iterator we can use to train or test a model on, responsible for both storing data
    and deciding how to batch it. """

    def get_epoch(self):
        """ Returns an iterator of batches/elements to train on, these elements are what will get
        passed to model.encode. Usually (but not necessarily) a list/batch of training examples """
        raise NotImplementedError(self.__class__)

    def get_batches(self, n_batches):
        if len(self) < n_batches:
            raise ValueError()
        return itertools.islice(self.get_epoch(), n_batches)

    def get_epochs(self, n_epochs: int):
        for _ in range(n_epochs):
            for batch in self.get_epoch():
                yield batch

    def get_samples(self, n_samples: int):
        """
        Sample for the data, be default we sample batches but subclasses can
        override this method to provide other kinds of sampling (like sampling individual elements).
        Must return both an iterator and the exact size of the iterator.
        """
        return self.get_batches(n_samples), n_samples

    def percent_filtered(self):
        # TODO nicer to just have "unfiltered_size"?
        """ If any filtering was done, the percent of examples that were filtered. Exposed so evaluators
         can compute percentages fairly even if some examples were removed during pre-processing """
        return None

    def __len__(self):
        """ Number of batches per an epoch """
        raise NotImplementedError(self.__class__) class Evaluator(Configurable):
    """ Class to generate statistics on a model's output for some data"""

    def tensors_needed(self, prediction: Prediction):
        """ Return all tensor variables needed by this evaluator in a dict, the results will
        be passed into `build_summary` as numpy arrays """
        raise NotImplementedError()

    def evaluate(self, input: List, true_len,  **kwargs) -> Evaluation:
        """
        Build a summary given the input data `input` and the result of the variables requested
        from `tensors_needed`. `true_len` is the total number of examples seen (or an approximation)
        excluding any pre-filtering that was done, its used for the case where some examples could not be
        processed by the model (e.g. too large) and were removed, but we still want to report
        accurate percentages on the entire dataset.
        """
        raise NotImplementedError() class Evaluation(object):
    """
    Evaluation of model, includes scalar summaries and per-example records
    """

    def __init__(self, scalars: Dict[str, Any], per_sample: Dict[str, List]=None):
        self.scalars = scalars
        self.per_sample = per_sample

    def add(self, other):
        for k in self.scalars:
            if k in other.scalars:
                raise ValueError("Two evaluations had the same scalar key: " + k)
        self.scalars.update(other.scalars)

        if self.per_sample is None:
            self.per_sample = other.per_sample
        elif other.per_sample is not None:
            for k in self.per_sample:
                if k in other.per_sample:
                    raise ValueError("Two evaluations had the same per sample key: " + k)
            self.per_sample.update(other.per_sample)

    def add_prefix(self, prefix):
        self.scalars = {prefix+k:v for k,v in self.scalars.items()}
        self.per_sample = {prefix+k:v for k,v in self.per_sample.items()}

    def to_summaries(self, prefix):
        return [tf.Summary(value=[tf.Summary.Value(tag=prefix + k, simple_value=v)]) for k,v in self.scalars.items()] class AysncEvaluatorRunner(object):
    """ Knows how to run a list of evaluators use a tf.Queue to feed in the data """

    def __init__(self, evaluators: List[Evaluator], model: Model, queue_size: int):
        placeholders = model.get_placeholders()
        self.eval_queue = tf.FIFOQueue(queue_size, [x.dtype for x in placeholders],
                                  name="eval_queue")
        self.enqueue_op = self.eval_queue.enqueue(placeholders)
        self.dequeue_op = self.eval_queue.dequeue()
        self.close_queue = self.eval_queue.close(True)

        # Queue in this form has not shape info, so we have to add it in back here
        for x, p in zip(placeholders, self.dequeue_op):
            p.set_shape(x.shape)
        self.evaluators = evaluators
        self.queue_size = self.eval_queue.size()
        self.model = model
        self.tensors_needed = None

    def set_input(self, prediction: Prediction):
        tensors_needed = []
        for ev in self.evaluators:
            tensors_needed.append(ev.tensors_needed(prediction))
        self.tensors_needed = tensors_needed

    def run_evaluators(self, sess: tf.Session, dataset, name, n_sample, feed_dict) -> Evaluation:
        all_tensors_needed = list(set(flatten_iterable(x.values() for x in self.tensors_needed)))

        tensors = {x: [] for x in all_tensors_needed}

        data_used = []
        if n_sample is None:
            batches, n_batches = dataset.get_epoch(), len(dataset)
        else:
            batches, n_batches = dataset.get_samples(n_sample)

        def enqueue_eval():
            try:
                for data in batches:
                    encoded = self.model.encode(data, False)
                    data_used.append(data)
                    sess.run(self.enqueue_op, encoded)
            except Exception as e:
                sess.run(self.close_queue)  # Crash the main thread
                raise e
            # we should run out of batches and exit gracefully

        th = Thread(target=enqueue_eval)

        th.daemon = True
        th.start()
        for _ in tqdm(range(n_batches), total=n_batches, desc=name, ncols=80):
            output = sess.run(all_tensors_needed, feed_dict=feed_dict)
            for i in range(len(all_tensors_needed)):
                tensors[all_tensors_needed[i]].append(output[i])
        th.join()

        if sess.run(self.queue_size) != 0:
            raise RuntimeError("All batches should be been consumed")

        # flatten the input
        for k in all_tensors_needed:
            v = tensors[k]
            if len(k.shape) == 0:
                v = np.array(v)  # List of scalars -> array
            elif any(x is None for x in k.shape.as_list()[1:]):
                # Variable sized tensors, so convert to flat python-list
                v = flatten_iterable(v)
            else:
                v = np.concatenate(v, axis=0)  # concat along the batch dim
            tensors[k] = v

        # flatten the data if it consists of batches
        if isinstance(data_used[0], List):
            data_used = flatten_iterable(data_used)

        if dataset.percent_filtered() is None:
            true_len = len(data_used)
        else:
            true_len = len(data_used) * 1 / (1 - dataset.percent_filtered())

        combined = None
        for ev, needed in zip(self.evaluators, self.tensors_needed):
            args = {k: tensors[v] for k, v in needed.items()}
            evaluation = ev.evaluate(data_used, true_len, **args)
            if combined is None:
                combined = evaluation
            else:
                combined.add(evaluation)

        return combined class EvaluatorRunner(object):
    """ Knows how to run a list of evaluators """

    def __init__(self, evaluators: List[Evaluator], model: Model):
        self.evaluators = evaluators
        self.tensors_needed = None
        self.model = model

    def set_input(self, prediction: Prediction):
        tensors_needed = []
        for ev in self.evaluators:
            tensors_needed.append(ev.tensors_needed(prediction))
        self.tensors_needed = tensors_needed

    def run_evaluators(self, sess: tf.Session, dataset: Dataset, name, n_sample=None, feed_dict=None) -> Evaluation:
        all_tensors_needed = list(set(flatten_iterable(x.values() for x in self.tensors_needed)))

        tensors = {x: [] for x in all_tensors_needed}

        if n_sample is None:
            batches, n_batches = dataset.get_epoch(), len(dataset)
        else:
            batches, n_batches = dataset.get_samples(n_sample)

        data_used = []

        for batch in tqdm(batches, total=n_batches, desc=name, ncols=80):
            feed_dict = self.model.encode(batch, is_train=False)
            output = sess.run(all_tensors_needed, feed_dict=feed_dict)
            data_used += batch
            for i in range(len(all_tensors_needed)):
                tensors[all_tensors_needed[i]].append(output[i])

        # flatten the input
        for k in all_tensors_needed:
            v = tensors[k]
            if len(k.shape) == 0:
                v = np.array(v)  # List of scalars
            elif any(x is None for x in k.shape.as_list()):
                # Variable sized tensors, so convert to flat python-list
                v = flatten_iterable(v)
            else:
                v = np.concatenate(v, axis=0)  # concat along the batch dim
            tensors[k] = v

        percent_filtered = dataset.percent_filtered()
        if percent_filtered is None:
            true_len = len(data_used)
        else:
            true_len = len(data_used) * 1 / (1 - percent_filtered)

        combined = None
        for ev, needed in zip(self.evaluators, self.tensors_needed):
            args = {k: tensors[v] for k, v in needed.items()}
            evaluation = ev.evaluate(data_used, true_len, **args)
            if evaluation is None:
                raise ValueError(ev)
            if combined is None:
                combined = evaluation
            else:
                combined.add(evaluation)

        return combined class Model(Configurable):
    """
    Our most general specification of a model/neural network, for our purposes a model
    is basically a pair of functions
    1) a way to map a (unspecified) kind of python object to numpy tensors and
    2) a tensorflow function that maps those kinds of tensors to a set of (also unspecified) output tensors

    For convenience, models maintain a of set of input placeholders that clients can make use of to
    feed the tensorflow function (or reference to construct their own tensor inputs).

    Models have two stages of initialization. First it needs
    to be initialized with the training data using `init` (typically this does things like deciding what
    words/chars to train embeddings for). This should only be done once for this object's lifetime.

    Afterwards use `set_inputs` to specify the input format, this does things like determine the batch size
    or the vocabulary that will be used

    After initialiation, `encode` will produce map of placeholder -> numpy array
    which can be used directly as a feed dict for the output of `get_predictions`

    For more advanced usage, `get_predictions_for` can be used with any tensors of the
    same shape/dtype as the input place holders. Clients should pass in a dict mapping
    the placeholders to the input tensors they want to use instead.

    `get_predictions_for` methods behave like any other tensorflow function, in that it will
    load/initialize/reuse variables depending on the current tensorflow scope and can add
    to tf.collections. Our trainer method makes use of some of these collections, including:
        tf.GraphKeys.LOSSES
        tf.GraphKeys.REGULARIZATION_LOSSES
        tf.GraphKeys.SUMMARIES
        tf.GraphKeys.SAVEABLE_OBJECTS
        tf.GraphKeys.TRAINABLE_VARIABLES
        "monitor/*" collections, which will be summed, and the EMA result logged to tensorboard
    """

    @property
    def name(self):
        return self.__class__.__name__

    def init(self, train_data, resource_loader: ResourceLoader):
        raise NotImplementedError()

    def set_inputs(self, datasets: List[Dataset], resource_loader: ResourceLoader) -> List[Tensor]:
        raise NotImplementedError()

    def get_prediction(self) -> Prediction:
        return self.get_predictions_for({x: x for x in self.get_placeholders()})

    def get_placeholders(self) -> List[Tensor]:
        raise NotImplementedError()

    def get_predictions_for(self, input_tensors: Dict[Tensor, Tensor]) -> Prediction:
        raise NotImplementedError()

    def encode(self, examples, is_train: bool) -> Dict[Tensor, object]:
        raise NotImplementedError() class ModelDir(object):
    """ Wrapper for accessing a folder we are storing a model in"""

    def __init__(self, name: str):
        if isabs(name):
            print("WARNING!!!, using an absolute paths for models name can break restoring "
                  "the model in different directories after being checkpointed")
            # why is this even a thing?
        self.dir = name

    def get_model(self) -> Model:
        with open(join(self.dir, "model.pkl"), "rb") as f:
            return pickle.load(f)

    def get_eval_dir(self):
        answer_dir = join(self.dir, "answers")
        if not exists(answer_dir):
            os.mkdir(answer_dir)
        return answer_dir

    def get_last_train_params(self):
        last_train_file = None
        last_train_step = -1
        for file in os.listdir(self.dir):
            if file.startswith("train_from_") and file.endswith("pkl"):
                step = int(file[11:file.rfind(".pkl")])
                if step > last_train_step:
                    last_train_step = step
                    last_train_file = join(self.dir, file)

        print("Resuming using the parameters stored in: " + last_train_file)
        with open(last_train_file, "rb") as f:
            return pickle.load(f)

    def get_latest_checkpoint(self):
        return tf.train.latest_checkpoint(self.save_dir)

    def get_checkpoint(self, step):
        # I cant find much formal documentation on how to do this, but this seems to work
        return join(self.save_dir, "checkpoint-%d-%d" % (step, step))

    def get_best_weights(self):
        if exists(self.best_weight_dir):
            return tf.train.latest_checkpoint(self.best_weight_dir)
        return None

    def restore_checkpoint(self, sess, var_list=None, load_ema=True):
        """
        Restores either the best weights or the most recent checkpoint, assuming the correct
        variables have already been added to the tf default graph e.g., .get_prediction()
        has been called the model stored in `self`.
        Automatically detects if EMA weights exists, and if they do loads them instead
        """
        checkpoint = self.get_best_weights()
        if checkpoint is None:
            print("Loading most recent checkpoint")
            checkpoint = self.get_latest_checkpoint()
        else:
            print("Loading best weights")

        if load_ema:
            if var_list is None:
                # Same default used by `Saver`
                var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) + \
                           tf.get_collection(tf.GraphKeys.SAVEABLE_OBJECTS)

            # Automatically check if there are EMA variables, if so use those
            reader = tf.train.NewCheckpointReader(checkpoint)
            ema = tf.train.ExponentialMovingAverage(0)
            ema_names = {ema.average_name(x): x for x in var_list
                         if reader.has_tensor(ema.average_name(x))}
            if len(ema_names) > 0:
                print("Found EMA weights, loading them")
                ema_vars = set(x for x in ema_names.values())
                var_list = {v.op.name: v for v in var_list if v not in ema_vars}
                var_list.update(ema_names)

        saver = tf.train.Saver(var_list)
        saver.restore(sess, checkpoint)


    @property
    def save_dir(self):
        # Stores training checkpoint
        return join(self.dir, "save")

    @property
    def best_weight_dir(self):
        # Stores training checkpoint
        return join(self.dir, "best-weights")

    @property
    def log_dir(self):
        return join(self.dir, "log")
import os
import pickle
import shutil
import time
import numpy as np
import tensorflow as tf
    import socket
from datetime import datetime
from os.path import exists, join, relpath
from threading import Thread
from typing import List, Union, Optional, Dict, Tuple
from tensorflow.python.training.adadelta import AdadeltaOptimizer
from tensorflow.python.training.adam import AdamOptimizer
from docqa import configurable
from docqa.configurable import Configurable
from docqa.data_processing.preprocessed_corpus import PreprocessedData
from docqa.dataset import TrainingData, Dataset
from docqa.evaluator import Evaluator, Evaluation, AysncEvaluatorRunner, EvaluatorRunner
from docqa.model import Model
from docqa.model_dir import ModelDir
            for x in lst:
                v = weight_ema.average(x)
                if v is not None:
                    lst.append(v)
        best_weight_saver = tf.train.Saver(var_list=lst, max_to_keep=1)
        cur_best = None
    else:
        best_weight_saver = None
        cur_best = None

    saver = tf.train.Saver(max_to_keep=train_params.max_checkpoints_to_keep)
    summary_writer = tf.summary.FileWriter(out.log_dir)

    # Load or initialize the model parameters
    if checkpoint is not None:
        print("Restoring from checkpoint...")
        saver.restore(sess, checkpoint)
        print("Loaded checkpoint: " + str(sess.run(global_step)))
    else:
        print("Initializing parameters...")
        sess.run(tf.global_variables_initializer())

    # Make sure no bugs occur that add to the graph in the train loop, that can cause (eventuall) OOMs
    tf.get_default_graph().finalize()

    if dry_run:
        return

    on_step = sess.run(global_step)

    if save_start:
        # summary_writer.add_graph(sess.graph, global_step=on_step)
        save_train_start(out.dir, data, sess.run(global_step), evaluators, train_params, notes)

    def enqueue_train():
        try:
            # feed data from the dataset iterator -> encoder -> queue
            for epoch in range(train_params.num_epochs):
                for batch in train.get_epoch():
                    feed_dict = model.encode(batch, True)
                    sess.run(train_enqeue, feed_dict)
        except tf.errors.CancelledError:
            # The queue_close operator has been called, exit gracefully
            return
        except Exception as e:
            # Crashes the main thread with a queue exception
            sess.run(train_close)
            raise e

    train_enqueue_thread = Thread(target=enqueue_train)
    train_enqueue_thread.daemon = True  # Ensure we exit the program on an excpetion

    print("Start training!")

    batch_time = 0

    train_dict = {is_train: True}
    eval_dict = {is_train: False}
    try:
        train_enqueue_thread.start()

        for epoch in range(train_params.num_epochs):
            for batch_ix in range(len(train)):
                t0 = time.perf_counter()
                on_step = sess.run(global_step) + 1
                get_summary = on_step % train_params.log_period == 0

                if get_summary:
                    summary, _, batch_loss = sess.run([summary_tensor, train_opt, loss], feed_dict=train_dict)
                else:
                    summary = None
                    _, batch_loss = sess.run([train_opt, loss], feed_dict=train_dict)

                if np.isnan(batch_loss):
                    raise RuntimeError("NaN loss!")

                batch_time += time.perf_counter() - t0
                if summary is not None:
                    print("on epoch=%d batch=%d step=%d, time=%.3f" %
                          (epoch, batch_ix + 1, on_step, batch_time))
                    summary_writer.add_summary(
                        tf.Summary(value=[tf.Summary.Value(tag="time", simple_value=batch_time)]), on_step)
                    summary_writer.add_summary(summary, on_step)
                    batch_time = 0

                # occasional saving
                if on_step % train_params.save_period == 0:
                    print("Checkpointing")
                    saver.save(sess, join(out.save_dir, "checkpoint-" + str(on_step)), global_step=global_step)

                # Occasional evaluation
                if (on_step % train_params.eval_period == 0) or start_eval:
                    print("Running evaluation...")
                    start_eval = False
                    t0 = time.perf_counter()
                    for name, data in eval_datasets.items():
                        n_samples = train_params.eval_samples.get(name)
                        evaluation = evaluator_runner.run_evaluators(sess, data, name, n_samples, eval_dict)
                        for s in evaluation.to_summaries(name + "-"):
                            summary_writer.add_summary(s, on_step)

                        # Maybe save as the best weights
                        if train_params.best_weights is not None and name == train_params.best_weights[0]:
                            val = evaluation.scalars[train_params.best_weights[1]]
                            if cur_best is None or val > cur_best:
                                print("Save weights with current best weights (%s vs %.5f)" % (
                                    "None" if cur_best is None else ("%.5f" % cur_best), val))
                                best_weight_saver.save(sess, join(out.best_weight_dir, "best"), global_step=global_step)
                                cur_best = val

                    print("Evaluation took: %.3f seconds" % (time.perf_counter() - t0))
    finally:
        sess.run(train_close)  # terminates the enqueue thread with an exception

    train_enqueue_thread.join()

    saver.save(sess, relpath(join(out.save_dir, "checkpoint-" + str(on_step))), global_step=global_step)
    sess.close()



def test(model: Model, evaluators, datasets: Dict[str, Dataset], loader, checkpoint,