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 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 ContextAndQuestion(object):
    """
    This is our standard unit of training data, context and a single question/answer.
    The answer type is unspecified and depends on the application.
    """

    def __init__(self, question: List[str], answer: Optional[Answer],
                 question_id: object, doc_id=None):
        self.question = question
        self.answer = answer
        self.question_id = question_id
        self.doc_id = doc_id

    @property
    def n_context_words(self) -> int:
        raise NotImplementedError()

    def get_context(self) -> List[str]:
        raise NotImplementedError() class Batcher(object):
    '''
    Batch sentences of tokenized text into character id matrices.
    '''
    def __init__(self, lm_vocab_file: str, max_token_length: str):
        '''
        lm_vocab_file = the language model vocabulary file (one line per
            token)
        max_token_length = the maximum number of characters in each token
        '''
        self._lm_vocab = UnicodeCharsVocabulary(
            lm_vocab_file, max_token_length
        )
        self._max_token_length = max_token_length

    def batch_sentences(self, sentences: List[List[str]]):
        '''
        Batch the sentences as character ids
        Each sentence is a list of tokens without <s> or </s>, e.g.
        [['The', 'first', 'sentence', '.'], ['Second', '.']]
        '''
        n_sentences = len(sentences)
        max_length = max(len(sentence) for sentence in sentences) + 2

        X_char_ids = np.zeros(
            (n_sentences, max_length, self._max_token_length),
            dtype=np.int64
        )

        for k, sent in enumerate(sentences):
            length = len(sent) + 2
            char_ids_without_mask = self._lm_vocab.encode_chars(' '.join(sent))
            # add one so that 0 is the mask value
            X_char_ids[k, :length, :] = char_ids_without_mask + 1

        return X_char_ids class TokenBatcher(object):
    '''
    Batch sentences of tokenized text into token id matrices.
    '''
    def __init__(self, lm_vocab_file: str):
        '''
        lm_vocab_file = the language model vocabulary file (one line per
            token)
        '''
        self._lm_vocab = Vocabulary(lm_vocab_file)

    def batch_sentences(self, sentences: List[List[str]]):
        '''
        Batch the sentences as character ids
        Each sentence is a list of tokens without <s> or </s>, e.g.
        [['The', 'first', 'sentence', '.'], ['Second', '.']]
        '''
        n_sentences = len(sentences)
        max_length = max(len(sentence) for sentence in sentences) + 2

        X_ids = np.zeros((n_sentences, max_length), dtype=np.int64)

        for k, sent in enumerate(sentences):
            length = len(sent) + 2
            ids_without_mask = self._lm_vocab.encode(' '.join(sent))
            # add one so that 0 is the mask value
            X_ids[k, :length] = ids_without_mask + 1

        return X_ids class BidirectionalLanguageModel(object):
    def __init__(
            self,
            options_file: str,
            weight_file: str,
            ids_placeholder,
            use_character_inputs=True,
            embedding_weight_file=None,
            max_batch_size=128,
    ):
        '''
        Creates the language model computational graph and loads weights
        Two options for input type:
            (1) To use character inputs (paired with Batcher)
                pass use_character_inputs=True, and ids_placeholder
                of shape (None, None, max_characters_per_token)
            (2) To use token ids as input (paired with TokenBatcher),
                pass use_character_inputs=False and ids_placeholder
                of shape (None, None).  In this case, embedding_weight_file
                is also required input
        options_file: location of the json formatted file with
                      LM hyperparameters
        weight_file: location of the hdf5 file with LM weights
        ids_placeholder: a tf.placeholder of type int32.
            If use_character_inputs=True, it is shape
                (None, None, max_characters_per_token) and holds the input
                character ids for a batch
            If use_character_input=False, it is shape (None, None) and
                holds the input token ids for a batch
        use_character_inputs: if True, then use character ids as input,
            otherwise use token ids
        max_batch_size: the maximum allowable batch size
        '''
        with open(options_file, 'r') as fin:
            options = json.load(fin)

        if not use_character_inputs:
            if embedding_weight_file is None:
                raise ValueError(
                    "embedding_weight_file is required input with "
                    "not use_character_inputs"
                )

        self._lm_graph = BidirectionalLanguageModelGraph(
            options,
            weight_file,
            ids_placeholder,
            embedding_weight_file=embedding_weight_file,
            use_character_inputs=use_character_inputs,
            max_batch_size=max_batch_size)

        self._ops = None

    def get_ops(self):
        '''
        returns a dictionary with tensorflow ops:
            {'lm_embeddings': embedding_op,
             'lengths': sequence_lengths_op,
             'mask': op to compute mask}
        embedding_op computes the LM embeddings and is shape
            (None, 3, None, 1024)
        lengths_op computes the sequence lengths and is shape (None, )
        mask computes the sequence mask and is shape (None, None)
        '''
        if self._ops is None:
            self._build_ops()
        return self._ops

    def _build_ops(self):
        with tf.control_dependencies([self._lm_graph.update_state_op]):
            # get the LM embeddings
            token_embeddings = self._lm_graph.embedding
            layers = [
                tf.concat([token_embeddings, token_embeddings], axis=2)
            ]

            n_lm_layers = len(self._lm_graph.lstm_outputs['forward'])
            for i in range(n_lm_layers):
                layers.append(
                    tf.concat(
                        [self._lm_graph.lstm_outputs['forward'][i],
                         self._lm_graph.lstm_outputs['backward'][i]],
                        axis=-1
                    )
                )

            # The layers include the BOS/EOS tokens.  Remove them
            sequence_length_wo_bos_eos = self._lm_graph.sequence_lengths - 2
            layers_without_bos_eos = []
            for layer in layers:
                layer_wo_bos_eos = layer[:, 1:, :]
                layer_wo_bos_eos = tf.reverse_sequence(
                    layer_wo_bos_eos,
                    self._lm_graph.sequence_lengths - 1,
                    seq_axis=1,
                    batch_axis=0,
                )
                layer_wo_bos_eos = layer_wo_bos_eos[:, 1:, :]
                layer_wo_bos_eos = tf.reverse_sequence(
                    layer_wo_bos_eos,
                    sequence_length_wo_bos_eos,
                    seq_axis=1,
                    batch_axis=0,
                )
                layers_without_bos_eos.append(layer_wo_bos_eos)

            # concatenate the layers
            lm_embeddings = tf.concat(
                [tf.expand_dims(t, axis=1) for t in layers_without_bos_eos],
                axis=1
            )

            # get the mask op without bos/eos.
            # tf doesn't support reversing boolean tensors, so cast
            # to int then back
            mask_wo_bos_eos = tf.cast(self._lm_graph.mask[:, 1:], 'int32')
            mask_wo_bos_eos = tf.reverse_sequence(
                mask_wo_bos_eos,
                self._lm_graph.sequence_lengths - 1,
                seq_axis=1,
                batch_axis=0,
            )
            mask_wo_bos_eos = mask_wo_bos_eos[:, 1:]
            mask_wo_bos_eos = tf.reverse_sequence(
                mask_wo_bos_eos,
                sequence_length_wo_bos_eos,
                seq_axis=1,
                batch_axis=0,
            )
            mask_wo_bos_eos = tf.cast(mask_wo_bos_eos, 'bool')

        self._ops = {
            'lm_embeddings': lm_embeddings,
            'lengths': sequence_length_wo_bos_eos,
            'token_embeddings': self._lm_graph.embedding,
            'mask': mask_wo_bos_eos,
        } 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 SequencePredictionLayer(Configurable):
    """ (batch, time, vec), answer -> Prediction """
    def apply(self, is_train, x, answer: List, mask=None) -> Prediction:
        return NotImplemented() 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 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 Mapper(SequenceMapper):
    """ (dim1, dim2, ...., input_dim) -> (im1, dim2, ...., output_dim) """
    def apply(self, is_train, x, mask=None):
        raise NotImplementedError() 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) def flatten_iterable(listoflists: Iterable[Iterable[T]]) -> List[T]:
    return [item for sublist in listoflists for item in sublist]
from os.path import join, expanduser
from typing import Optional, List, Dict, Union
from tensorflow import Tensor
from docqa.configurable import Configurable
from docqa.data_processing.qa_training_data import ParagraphAndQuestionDataset, ContextAndQuestion
from docqa.elmo.data import Batcher, TokenBatcher
from docqa.elmo.lm_model import BidirectionalLanguageModel
from docqa.encoder import DocumentAndQuestionEncoder
from docqa.model import Model, Prediction
from docqa.nn.embedder import WordEmbedder, CharWordEmbedder
from docqa.nn.layers import SequenceMapper, SequencePredictionLayer, SequenceBiMapper, AttentionMapper, \
    AttentionPredictionLayer, Mapper
from docqa.utils import ResourceLoader, flatten_iterable
import numpy as np
import tensorflow as tf

        return self.get_placeholders()

    def get_placeholders(self):
        return self.encoder.get_placeholders() + [
            self._is_train_placeholder,
            self._question_char_ids_placeholder,
            self._context_char_ids_placeholder
        ] + ([self._context_sentence_ixs] if (self._context_sentence_ixs is not None) else [])

    def get_predictions_for(self, input_tensors: Dict[Tensor, Tensor]):
        is_train = input_tensors[self._is_train_placeholder]
        enc = self.encoder

        q_lm_model = BidirectionalLanguageModel(self.lm_model.options_file, self.lm_model.weight_file,
                                                input_tensors[self._question_char_ids_placeholder],
                                                embedding_weight_file=self.lm_model.embed_weights_file,
                                                use_character_inputs=not self.token_lookup,
                                                max_batch_size=self.max_batch_size)
        q_lm_encoding = q_lm_model.get_ops()["lm_embeddings"]

        with tf.variable_scope(tf.get_variable_scope(), reuse=True):
            c_lm_model = BidirectionalLanguageModel(self.lm_model.options_file, self.lm_model.weight_file,
                                                    input_tensors[self._context_char_ids_placeholder],
                                                    embedding_weight_file=self.lm_model.embed_weights_file,
                                                    use_character_inputs=not self.token_lookup,
                                                    max_batch_size=self._max_num_sentences)
            c_lm_encoding = c_lm_model.get_ops()["lm_embeddings"]

        if self.per_sentence:
            c_lm_encoding = tf.gather_nd(c_lm_encoding, input_tensors[self._context_sentence_ixs])

        q_mask = input_tensors[enc.question_len]
        c_mask = input_tensors[enc.context_len]

        q_embed = []
        c_embed = []

        if enc.question_chars in input_tensors:
            with tf.variable_scope("char-embed"):
                q, c = self.char_embed.embed(is_train,
                                             (input_tensors[enc.question_chars], input_tensors[enc.question_word_len]),
                                             (input_tensors[enc.context_chars], input_tensors[enc.context_word_len]))
            q_embed.append(q)
            c_embed.append(c)

        if enc.question_words in input_tensors:
            with tf.variable_scope("word-embed"):
                q, c = self.word_embed.embed(is_train,
                                            (input_tensors[enc.question_words], q_mask),
                                            (input_tensors[enc.context_words], c_mask))
            q_embed.append(q)
            c_embed.append(c)

        if enc.question_features in input_tensors:
            q_embed.append(input_tensors.get(enc.question_features))
            c_embed.append(input_tensors.get(enc.context_features))

        q_embed = tf.concat(q_embed, axis=2)
        c_embed = tf.concat(c_embed, axis=2)

        answer = [input_tensors[x] for x in enc.answer_encoder.get_placeholders()]
        return self._get_predictions_for(is_train, q_embed, q_mask, q_lm_encoding,
                                         c_embed, c_mask, c_lm_encoding, answer)

    def _get_predictions_for(self,
                             is_train,
                             question_embed, question_mask, question_lm,
                             context_embed, context_mask, context_lm,
                             answer) -> Prediction:
        raise NotImplementedError()

    def encode(self, batch: List[ContextAndQuestion], is_train: bool):
        if len(batch) > self.max_batch_size:
            raise ValueError("The model can only use a batch <= %d, but got %d" %
                             (self.max_batch_size, len(batch)))
        data = self.encoder.encode(batch, is_train)
        data[self._question_char_ids_placeholder] = self._batcher.batch_sentences([q.question for q in batch])
        data[self._is_train_placeholder] = is_train
        context_word_dim = data[self.encoder.context_words].shape[1]

        if not self.per_sentence:
            data[self._context_char_ids_placeholder] = \
                self._batcher.batch_sentences([x.get_context() for x in batch])
        else:
            data[self._context_char_ids_placeholder] = \
                self._batcher.batch_sentences(flatten_iterable([x.sentences for x in batch]))

            # Compute indices where context_sentence_ixs[sentence#, k, sentence_word#] = (batch#, k, batch_word#)
            # for each word. We use this to map the tokens built for the sentences back to
            # the format where sentences are flattened for each batch
            context_sentence_ixs = np.zeros((len(batch), 3, context_word_dim, 3), dtype=np.int32)
            total_sent_ix = 0
            for ix, point in enumerate(batch):
                word_ix = 0
                for sent_ix, sent in enumerate(point.sentences):
                    for w_ix in range(len(sent)):
                        for k in range(3):
                            context_sentence_ixs[ix, k, word_ix] = [total_sent_ix, k, w_ix]
                        word_ix += 1
                    total_sent_ix += 1
            data[self._context_sentence_ixs] = context_sentence_ixs
        return data


class AttentionWithElmo(ElmoQaModel):
    """ Elmo model that uses attention """

    def __init__(self,
                 encoder: DocumentAndQuestionEncoder,
                 lm_model: LanguageModel,
                 max_batch_size: int,
                 per_sentence: bool,
                 append_embed: bool,
                 append_before_atten: bool,
                 word_embed: Optional[WordEmbedder],
                 char_embed: Optional[CharWordEmbedder],
                 embed_mapper: Optional[SequenceMapper],
                 lm_reduce_shared: Optional[Mapper],
                 lm_reduce: Optional[Mapper],

                 memory_builder: SequenceBiMapper,