class ParagraphAndQuestionDataset(ListDataset):
    """ Dataset with q/a pairs and that exposes some meta-data about its elements """
    def get_spec(self) -> ParagraphAndQuestionSpec:
        return build_spec(self.batching.get_fixed_batch_size(),
                          self.batching.get_max_batch_size(),
                          self.data)

    def get_vocab(self) -> Set[str]:
        return compute_voc(self.data) class ParagraphAndQuestionSpec(object):
    """ Bound on the size of `ParagraphAndQuestion` objects """
    def __init__(self, batch_size, max_question_words=None,
                 max_num_context_words=None, max_word_size=None, max_batch_size=None):
        if batch_size is not None:
            if max_batch_size is None:
                max_batch_size = batch_size
            elif max_batch_size != batch_size:
                raise ValueError()
        self.batch_size = batch_size
        self.max_num_quesiton_words = max_question_words
        self.max_num_context_words = max_num_context_words
        self.max_word_size = max_word_size
        self.max_batch_size = max_batch_size

    def __add__(self, o):
        return ParagraphAndQuestionSpec(
            max_or_none(self.batch_size, o.batch_size),
            max_or_none(self.max_num_quesiton_words, o.max_num_quesiton_words),
            max_or_none(self.max_num_context_words, o.max_num_context_words),
            max_or_none(self.max_word_size, o.max_word_size),
            max_or_none(self.max_batch_size, o.max_batch_size)
        ) class DocumentAndQuestionEncoder(Configurable):
    """
    Uses a WordEmbedder/CharEmbedder (passed in by the client in `init`) to encode text into padded batches of arrays.
    It should really be called "ParagraphAndQuestionEncoder", but we are stuck with this for now
    """

    def __init__(self,
                 answer_encoder: AnswerEncoder,
                 doc_size_th: Optional[int]=None,
                 word_featurizer: Optional[QaTextFeautrizer]=None):
        # Parameters
        self.answer_encoder = answer_encoder
        self.doc_size_th = doc_size_th

        self.word_featurizer = word_featurizer

        self._word_embedder = None
        self._char_emb = None

        # Internal stuff we need to set on `init`
        self.len_opt = None
        self.batch_size = None
        self.max_context_word_dim = None
        self.max_ques_word_dim = None
        self.max_char_dim = None

        self.context_features = None
        self.context_words = None
        self.context_chars = None
        self.context_len = None
        self.context_word_len = None
        self.question_features = None
        self.question_words = None
        self.question_chars = None
        self.question_len = None
        self.question_word_len = None

    @property
    def version(self):
        return 3

    def init(self, input_spec: ParagraphAndQuestionSpec, len_op: bool,
             word_emb: WordEmbedder, char_emb: Optional[CharEmbedder]):

        self._word_embedder = word_emb
        self._char_emb = char_emb

        self.batch_size = input_spec.batch_size
        self.len_opt = len_op

        if self._char_emb is not None:
            if input_spec.max_word_size is not None:
                self.max_char_dim = min(self._char_emb.get_word_size_th(), input_spec.max_word_size)
            else:
                self.max_char_dim = self._char_emb.get_word_size_th()
        else:
            self.max_char_dim = 1

        if not self.len_opt:
            self.max_ques_word_dim = input_spec.max_num_quesiton_words
            self.max_context_word_dim = input_spec.max_num_context_words
            if self.max_ques_word_dim is None or self.max_context_word_dim is None:
                raise ValueError()
            if self.doc_size_th is not None:
                self.max_context_word_dim = min(self.max_context_word_dim, self.doc_size_th)
        else:
            self.max_ques_word_dim = None
            self.max_context_word_dim = None

        n_question_words = self.max_ques_word_dim
        n_context_words = self.max_context_word_dim
        batch_size = self.batch_size

        self.context_words = tf.placeholder('int32', [batch_size, n_context_words], name='context_words')
        self.context_len = tf.placeholder('int32', [batch_size], name='context_len')

        self.question_words = tf.placeholder('int32', [batch_size, n_question_words], name='question_words')
        self.question_len = tf.placeholder('int32', [batch_size], name='question_len')

        if self._char_emb:
            self.context_chars = tf.placeholder('int32', [batch_size, n_context_words, self.max_char_dim], name='context_chars')
            self.question_chars = tf.placeholder('int32', [batch_size, n_question_words, self.max_char_dim], name='question_chars')
            self.question_word_len = tf.placeholder('int32', [batch_size, n_question_words], name='question_len')
            self.context_word_len = tf.placeholder('int32', [batch_size, n_context_words], name='context_len')
        else:
            self.context_chars = None
            self.question_chars = None
            self.context_word_len = None
            self.question_word_len = None

        if self.word_featurizer is not None:
            self.question_features = tf.placeholder('float32',
                                                    [batch_size, n_question_words,
                                                     self.word_featurizer.n_question_features()],
                                                    name='question_features')
            self.context_features = tf.placeholder('float32', [batch_size, n_context_words,
                                                               self.word_featurizer.n_context_features()],
                                                   name='context_features')
        else:
            self.question_features = None
            self.context_features = None

        self.answer_encoder.init(batch_size, n_context_words)

    def get_placeholders(self):
        return [x for x in
                [self.question_len, self.question_words, self.question_chars, self.question_features,
                 self.context_len, self.context_words, self.context_chars, self.context_features,
                 self.question_word_len, self.context_word_len]
                if x is not None] + self.answer_encoder.get_placeholders()

    def encode(self, batch: List[ContextAndQuestion], is_train: bool):
        batch_size = len(batch)
        if self.batch_size is not None:
            if self.batch_size < batch_size:
                raise ValueError("Batch sized we pre-specified as %d, "
                                 "but got a batch of %d" % (self.batch_size, batch_size))
            # We have a fixed batch size, so we will pad our inputs with zeros along the batch dimension
            batch_size = self.batch_size

        context_word_dim, ques_word_dim, max_char_dim = \
            self.max_context_word_dim, self.max_ques_word_dim, self.max_char_dim

        feed_dict = {}

        # compute the question/word lengths
        if context_word_dim is not None:
            context_len = np.array([min(doc.n_context_words, context_word_dim)  # Might truncate context
                                    for doc in batch], dtype='int32')
        else:
            context_len = np.array([doc.n_context_words for doc in batch], dtype='int32')
            context_word_dim = context_len.max()

        question_len = np.array([len(x.question) for x in batch], dtype='int32')
        if ques_word_dim is not None:
            if question_len.max() > ques_word_dim:
                raise ValueError("Have a question of len %d but max ques dim is %d" %
                                 (question_len.max(), ques_word_dim))
        else:
            ques_word_dim = question_len.max()

        feed_dict[self.context_len] = context_len
        feed_dict[self.question_len] = question_len

        # Setup word placeholders
        if self._word_embedder is not None:
            context_words = np.zeros([batch_size, context_word_dim], dtype='int32')
            question_words = np.zeros([batch_size, ques_word_dim], dtype='int32')
            feed_dict[self.context_words] = context_words
            feed_dict[self.question_words] = question_words
        else:
            question_words, context_words = None, None

        # Setup char placeholders
        if self._char_emb is not None:
            context_chars = np.zeros([batch_size, context_word_dim, max_char_dim], dtype='int32')
            question_chars = np.zeros([batch_size, ques_word_dim, max_char_dim], dtype='int32')
            context_word_len = np.zeros([batch_size, context_word_dim], dtype='int32')
            question_word_len = np.zeros([batch_size, ques_word_dim], dtype='int32')
            feed_dict[self.question_chars] = question_chars
            feed_dict[self.context_chars] = context_chars
            feed_dict[self.question_word_len] = question_word_len
            feed_dict[self.context_word_len] = context_word_len
        else:
            context_chars, question_chars, question_word_len, context_word_len = None, None, None, None

        query_once = self._word_embedder.query_once()

        # Now fill in the place holders by iterating through the data
        for doc_ix, doc in enumerate(batch):
            doc_mapping = {}  # word->ix mapping if `query_once` is True

            for word_ix, word in enumerate(doc.question):
                if self._word_embedder is not None:
                    if query_once:
                        ix = doc_mapping.get(word)
                        if ix is None:
                            ix = self._word_embedder.context_word_to_ix(word, is_train)
                            doc_mapping[word] = ix
                    else:
                        ix = self._word_embedder.context_word_to_ix(word, is_train)
                    question_words[doc_ix, word_ix] = ix
                if self._char_emb is not None:
                    question_word_len[doc_ix, word_ix] = min(self.max_char_dim, len(word))
                    for char_ix, char in enumerate(word):
                        if char_ix == self.max_char_dim:
                            break
                        question_chars[doc_ix, word_ix, char_ix] = self._char_emb.char_to_ix(char)

            for word_ix, word in enumerate(doc.get_context()):
                if self._word_embedder is not None:
                    if query_once:
                        ix = doc_mapping.get(word)
                        if ix is None:
                            ix = self._word_embedder.context_word_to_ix(word, is_train)
                            doc_mapping[word] = ix
                    else:
                        ix = self._word_embedder.context_word_to_ix(word, is_train)
                    context_words[doc_ix, word_ix] = ix

                if self._char_emb is not None:
                    context_word_len[doc_ix, word_ix] = min(self.max_char_dim, len(word))
                    for char_ix, char in enumerate(word):
                        if char_ix == self.max_char_dim:
                            break
                        context_chars[doc_ix, word_ix, char_ix] = self._char_emb.char_to_ix(char)

        # Answer placeholders
        feed_dict.update(self.answer_encoder.encode(batch_size, context_len, context_word_dim, batch))

        # Features placeholders
        if self.word_featurizer is not None:
            question_word_features = np.zeros((batch_size, ques_word_dim, self.word_featurizer.n_question_features()))
            context_word_features = np.zeros((batch_size, context_word_dim, self.word_featurizer.n_context_features()))
            for doc_ix, doc in enumerate(batch):
                q_f, c_f = self.word_featurizer.get_features(doc.question, doc.get_context())
                question_word_features[doc_ix, :q_f.shape[0]] = q_f
                context_word_features[doc_ix, :c_f.shape[0]] = c_f
            feed_dict[self.context_features] = context_word_features
            feed_dict[self.question_features] = question_word_features

        return feed_dict

    def __getstate__(self):
        # The placeholders are considered transient, the model
        # will be re-initailized when re-loaded
        state = dict(
            answer_encoder=self.answer_encoder,
            doc_size_th=self.doc_size_th,
            word_featurizer=self.word_featurizer,
            version=self.version
        )
        return state

    def __setstate__(self, state):
        if state["version"] == 0:
            if "word_featurizer" in state["state"]:
                raise ValueError()
            state["state"]["word_featurizer"] = None
        if state["version"] <= 1:
            if "answer_encoder" in state["state"]:
                raise ValueError()
            state["state"]["answer_encoder"] = SingleSpanAnswerEncoder()
        elif state["version"] <= 2:
            super().__setstate__(state)
        else:
            del state["version"]
            return self.__init__(**state) 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 Prediction(object):
    """ Prediction from a model, subclasses should provide access to a tensor
     representations of the model's output """
    pass class WordEmbedder(Configurable):
    """
    Responsible for mapping words -> ids, ids -> words, and ids -> embeddings
    matrices. Needs to be initialized from a corpus by `set_vocab` after construction
    """

    def set_vocab(self, corpus, word_vec_loader: ResourceLoader, special_tokens: List[str]):
        raise NotImplementedError()

    def is_vocab_set(self):
        raise NotImplementedError()

    def question_word_to_ix(self, word, is_train) -> int:
        raise NotImplementedError()

    def context_word_to_ix(self, word, is_train) -> int:
        raise NotImplementedError()

    def query_once(self) -> bool:
        """
        Should the embedder be queried once for each unique word in the input, or once for each word.
        Intended to support placeholders, although I ended up not experimenting much w/that route
        """
        return False

    def init(self, word_vec_loader, voc: Iterable[str]):
        raise NotImplementedError()

    def embed(self, is_train, *word_ix_and_mask):
        """ [(word_ix, mask)...] -> [word_embed, ...]"""
        raise NotImplemented() class CharWordEmbedder(Configurable):
    """
    Derives word embeddings from character embeddings by combining a character embedder and a reduce layer
    """
    def __init__(self, embedder: CharEmbedder, layer: Encoder,
                 shared_parameters: bool):
        self.embeder = embedder
        self.layer = layer
        self.shared_parameters = shared_parameters

    def embed(self, is_train, *char_ix):
        embeds = self.embeder.embed(is_train, *char_ix)
        if self.shared_parameters:
            with tf.variable_scope("embedding"):
                output = [self.layer.apply(is_train, embeds[0], char_ix[0][1])]
            with tf.variable_scope("embedding", reuse=True):
                for i in range(1, len(embeds)):
                    output.append(self.layer.apply(is_train, embeds[i], char_ix[i][1]))
        else:
            output = []
            for i, emb in enumerate(embeds):
                with tf.variable_scope("embedding%d_%s" % (i, emb.name)):
                    output.append(self.layer.apply(is_train, emb, char_ix[i][1]))
        return output

    def __setstate__(self, state):
        if "state" in state:
            state["state"]["version"] = state["version"]
            state = state["state"]
        if "share" in state:
            state["shared_parameters"] = state["share"]
            del state["share"]
        super().__setstate__(state) class SequenceMapper(Configurable):
    """ (batch, time, in_dim) -> (batch, time, out_dim) """
    def apply(self, is_train, x, mask=None):
        raise NotImplementedError() class SequenceBiMapper(Configurable):
    """ (batch, time, in_dim) -> (batch, time, out_dim1), (batch, time, out_dim2) """
    def apply(self, is_train, x, mask=None):
        raise NotImplementedError() class AttentionMapper(Configurable):
    """ (batch, time1, dim1), (batch, time1, dim2) (batch, time2, dim3) -> (batch, time1, out_dim) """
    def apply(self, is_train, x, keys, memories, mask=None, memory_mask=None):
        raise NotImplementedError() class SequenceEncoder(Configurable):
    """ (batch, time, in_dim) -> (batch, out_dim) """
    def apply(self, is_train, x, mask=None):
        raise NotImplementedError() class SequenceMapperWithContext(Configurable):
    """ (batch, time, in_dim) (batch, in_dim) -> (batch, time, out_dim) """
    def apply(self, is_train, x, c, mask=None):
        raise NotImplementedError() class MapMulti(Configurable):
    """ Applies a layer to multiple inputs, possibly sharing parameters """

    def __init__(self, layer: SequenceMapper, share: bool=True):
        self.layer = layer
        self.share = share

    def apply(self, is_train, *inputs):
        if self.share:
            with tf.variable_scope("map"):
                output = [self.layer.apply(is_train, inputs[0][0], inputs[0][1])]
            with tf.variable_scope("map", reuse=True):
                for i in range(1, len(inputs)):
                    output.append(self.layer.apply(is_train, inputs[i][0], inputs[1][1]))
        else:
            output = []
            for i,(x,mask) in enumerate(inputs):
                with tf.variable_scope("map%d_%s" % (i, x.name)):
                    output.append(self.layer.apply(is_train, x, mask))
        return output class SequencePredictionLayer(Configurable):
    """ (batch, time, vec), answer -> Prediction """
    def apply(self, is_train, x, answer: List, mask=None) -> Prediction:
        return NotImplemented() class AttentionPredictionLayer(Configurable):
    """ (batch, time1, vec1), (batch, time2, vec2), answers  -> Prediction """
    def apply(self, is_train, keys, context, answer: List, mask=None, memory_mask=None) -> Prediction:
        raise NotImplementedError() class TextPreprocessor(Configurable):
    """ Preprocess text input, must be deterministic. Only used thus far adding special indicator tokens """

    def encode_extracted_paragraph(self, question: List[str], paragraph: ExtractedParagraphWithAnswers):
        text, answers, _ = self.encode_paragraph(question, paragraph.text,
                                                 paragraph.start == 0, paragraph.answer_spans)
        return ParagraphWithAnswers(text, answers)

    def encode_text(self, question: List[str], paragraph: ExtractedParagraph):
        text, _, _ = self.encode_paragraph(question, paragraph.text, paragraph.start == 0,
                                           np.zeros((0, 2), dtype=np.int32))
        return text

    def encode_paragraph(self, question: List[str], paragraphs: List[List[str]],
                         is_first, answer_spans: np.ndarray,
                         token_spans=None) -> Tuple[List[str], np.ndarray, Optional[np.ndarray]]:
        """
        Returns updated (and flattened) text, answer_spans, and token_spans
        """
        raise NotImplementedError()

    def special_tokens(self) -> List[str]:
        return [] class ResourceLoader(object):
    """
    Abstraction for models the need access to external resources to setup, currently just
    for word-vectors.
    """

    def __init__(self, load_vec_fn=load_word_vectors):
        self.load_vec_fn = load_vec_fn

    def load_word_vec(self, vec_name, voc=None):
        return self.load_vec_fn(vec_name, voc)
from typing import List, Optional, Dict, Union, Set
from tensorflow import Tensor
from docqa.data_processing.qa_training_data import ParagraphAndQuestionDataset, ParagraphAndQuestionSpec
from docqa.encoder import DocumentAndQuestionEncoder
from docqa.model import Model, Prediction
from docqa.nn.embedder import WordEmbedder, CharWordEmbedder
from docqa.nn.layers import SequenceMapper, SequenceBiMapper, AttentionMapper, SequenceEncoder, \
    SequenceMapperWithContext, MapMulti, SequencePredictionLayer, AttentionPredictionLayer
from docqa.text_preprocessor import TextPreprocessor
from docqa.utils import ResourceLoader
import tensorflow as tf
                             question_embed, question_mask,
                             context_embed, context_mask,
                             answer) -> Prediction:
        raise NotImplemented()

    def encode(self, batch: List, is_train: bool):
        data = self.encoder.encode(batch, is_train)
        data[self._is_train_placeholder] = is_train
        return data

    def __getstate__(self):
        state = super().__getstate__()
        state["_is_train_placeholder"] = None
        return state

    def __setstate__(self, state):
        if "state" in state:
            if "preprocessor" not in state["state"]:
                state["state"]["preprocessor"] = None
        super().__setstate__(state)


class ContextOnly(ParagraphQuestionModel):

    def __init__(self, encoder: DocumentAndQuestionEncoder,
                 word_embed: Optional[WordEmbedder],
                 char_embed: Optional[CharWordEmbedder],
                 context_encoder: SequenceMapper,
                 prediction: SequencePredictionLayer):
        super().__init__(encoder, word_embed, char_embed)
        self.context_encoder = context_encoder
        self.prediction = prediction

    def _get_predictions_for(self, is_train,
                             question_embed, question_mask,
                             context_embed, context_mask,
                             answer) -> Prediction:
        with tf.variable_scope("encode"):
            self.context_encoder.apply(is_train, context_embed, context_mask)

        with tf.variable_scope("predict"):
            return self.prediction.apply(is_train, context_embed, answer, context_mask)


class Attention(ParagraphQuestionModel):
    """
    Model that encodes the question and context, then applies an attention mechanism
    between the two to produce a query-aware context representation, which is used to make a prediction.
    """
    def __init__(self, encoder: DocumentAndQuestionEncoder,
                 preprocess: Optional[TextPreprocessor],
                 word_embed: Optional[WordEmbedder],
                 word_embed_layer: Optional[MapMulti],
                 char_embed: Optional[CharWordEmbedder],
                 embed_mapper: Optional[SequenceMapper],
                 question_mapper: Optional[SequenceMapper],
                 context_mapper: Optional[SequenceMapper],
                 memory_builder: SequenceBiMapper,
                 attention: AttentionMapper,
                 match_encoder: SequenceMapper,
                 predictor: Union[SequencePredictionLayer, AttentionPredictionLayer]):
        super().__init__(encoder, word_embed, char_embed, word_embed_layer, preprocess)
        self.embed_mapper = embed_mapper
        self.question_mapper = question_mapper
        self.context_mapper = context_mapper
        self.memory_builder = memory_builder
        self.attention = attention
        self.match_encoder = match_encoder
        self.predictor = predictor

    def _get_predictions_for(self, is_train,
                             question_rep, question_mask,
                             context_rep, context_mask,
                             answer) -> Prediction:
        if self.embed_mapper is not None:
            with tf.variable_scope("map_embed"):
                context_rep = self.embed_mapper.apply(is_train, context_rep, context_mask)
            with tf.variable_scope("map_embed", reuse=True):
                question_rep = self.embed_mapper.apply(is_train, question_rep, question_mask)

        if self.question_mapper is not None:
            with tf.variable_scope("map_question"):
                question_rep = self.question_mapper.apply(is_train, question_rep, question_mask)

        if self.context_mapper is not None:
            with tf.variable_scope("map_context"):
                context_rep = self.context_mapper.apply(is_train, context_rep, context_mask)

        with tf.variable_scope("buid_memories"):
            keys, memories = self.memory_builder.apply(is_train, question_rep, question_mask)

        with tf.variable_scope("apply_attention"):
            context_rep = self.attention.apply(is_train, context_rep, keys, memories, context_mask, question_mask)

        if self.match_encoder is not None:
            with tf.variable_scope("process_attention"):
                context_rep = self.match_encoder.apply(is_train, context_rep, context_mask)

        with tf.variable_scope("predict"):
            if isinstance(self.predictor, AttentionPredictionLayer):
                return self.predictor.apply(is_train, context_rep, question_rep, answer, context_mask, question_mask)
            else:
                return self.predictor.apply(is_train, context_rep, answer, context_mask)


class AttentionAndEncode(ParagraphQuestionModel):

    def __init__(self, encoder: DocumentAndQuestionEncoder,
                 word_embed: Optional[WordEmbedder],
                 word_embed_layer: Optional[MapMulti],
                 char_embed: Optional[CharWordEmbedder],
                 embed_mapper: Optional[SequenceMapper],
                 question_mapper: Optional[SequenceMapper],
                 question_encoder: SequenceEncoder,
                 context_mapper: Optional[SequenceMapper],
                 memory_builder: SequenceBiMapper,
                 attention: AttentionMapper,
                 post_attention_mapper: Optional[SequenceMapper],
                 contextual_mapper: SequenceMapperWithContext,
                 post_context_mapper: Optional[SequenceMapper],

                 predictor: SequencePredictionLayer):