class ModelDir(object):
    def __init__(self, name: str):
    def get_model(self) -> Model:
    def get_eval_dir(self):
    def get_last_train_params(self):
    def get_latest_checkpoint(self):
    def get_checkpoint(self, step):
    def get_best_weights(self):
    def restore_checkpoint(self, sess, var_list=None, load_ema=True):
    def save_dir(self):
    def best_weight_dir(self):
    def log_dir(self): class SerializableOptimizer(Configurable):
class TrainParams(Configurable):
    def __init__(self, opt_name, params=None):
    def get_params(self):
    def get(self, name=None):
def init(out: ModelDir, model: Model, override=False):
    def __init__(self,
                 opt: SerializableOptimizer,
                 num_epochs: int,
                 eval_period: int,
                 log_period: int,
                 save_period: int,
                 eval_samples: Dict[str, Optional[int]],
                 regularization_weight: Optional[float] = None,
                 async_encoding: Optional[int] = None,
                 max_checkpoints_to_keep: int = 5,
                 loss_ema: Optional[float] = .999,
                 eval_at_zero: bool = False,
                 monitor_ema: float = .999,
                 ema: Optional[float] = None,
                 best_weights: Optional[Tuple[str, str]] = None
                 ):
def save_train_start(out,
                     data: TrainingData,
                     global_step: int,
                     evaluators: List[Evaluator],
                     train_params: TrainParams,
                     notes: str):
def _build_train_ops(train_params):
def continue_training(
        data: TrainingData,
        model: Model,
        train_params: TrainParams,
        evaluators: List[Evaluator],
        out: ModelDir,
        notes: str = None,
        dry_run=False):
def start_training(
        data: TrainingData,
        model: Model,
        train_params: TrainParams,
        evaluators: List[Evaluator],
        out: ModelDir,
        notes: str = None,
        initialize_from=None,
        dry_run=False):
def resume_training(out: ModelDir, notes: str = None, dry_run=False, start_eval=False):
def resume_training_with(
        data: TrainingData,
        out: ModelDir,
        train_params: TrainParams,
        evaluators: List[Evaluator],
        notes: str = None,
        dry_run: bool = False):
def _train(model: Model,
           data: TrainingData,
           checkpoint: Union[str, None],
           parameter_checkpoint: Union[str, None],
           save_start: bool,
           train_params: TrainParams,
           evaluators: List[Evaluator],
           out: ModelDir,
           notes=None,
           dry_run=False,
           start_eval=False):
def _train_async(model: Model,
                 data: TrainingData,
                 checkpoint: Union[str, None],
                 parameter_checkpoint: Union[str, None],
                 save_start: bool,
                 train_params: TrainParams,
                 evaluators: List[Evaluator],
                 out: ModelDir,
                 notes=None,
                 dry_run=False,
                 start_eval=False):
    def enqueue_train():
def test(model: Model, evaluators, datasets: Dict[str, Dataset], loader, checkpoint,
         ema=True, aysnc_encoding=None, sample=None) -> Dict[str, Evaluation]: class ContextLenBucketedKey(Configurable):
    def __init__(self, bucket_size: int):
        self.bucket_size = bucket_size

    def __call__(self, q: SentencesAndQuestion):
        return q.n_context_words//self.bucket_size class ContextLenKey(Configurable):
    def __call__(self, q: ContextAndQuestion):
        return q.n_context_words class ClusteredBatcher(ListBatcher):
    def __init__(self,
                 batch_size: int,
                 clustering: Callable,
                 shuffle_buckets=False,
                 truncate_batches=False):
        self.batch_size = batch_size
        self.clustering = clustering
        self.shuffle_buckets = shuffle_buckets
        self.truncate_batches = truncate_batches

    def get_fixed_batch_size(self):
        return None if self.truncate_batches else self.batch_size

    def get_max_batch_size(self):
        return self.batch_size

    def get_epoch(self, data: List):
        data = sorted(data, key=self.clustering)
        if self.shuffle_buckets:
            shuffle_list_buckets(data, self.clustering, np.random)
        n_batches = len(data) // self.batch_size
        intervals = [(i*self.batch_size, (i + 1)*self.batch_size) for i in range(0, n_batches)]
        remainder = len(data) % self.batch_size
        if self.truncate_batches and remainder > 0:
            intervals.append((len(data) - remainder, len(data)))
        np.random.shuffle(intervals)
        for i, j in intervals:
            yield data[i:j]

    def epoch_size(self, n_elements):
        size = n_elements // self.batch_size
        if self.truncate_batches and (n_elements % self.batch_size) > 0:
            size += 1
        return size 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 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 SingleSpanAnswerEncoder(AnswerEncoder):
    """ Encode the answer as integer coordinates, pick a random answer span if multiple spans exists """

    def __init__(self):
        self.answer_spans = None

    def get_placeholders(self) -> List:
        return [self.answer_spans]

    def init(self, batch_size, context_word_dim):
        self.answer_spans = tf.placeholder('int32', [batch_size, 2], name='answer_spans')

    def encode(self, batch_size, context_len, context_word_dim, batch) -> Dict:
        answer_spans = np.zeros([batch_size, 2], dtype='int32')

        for doc_ix, doc in enumerate(batch):
            answer = doc.answer

            if answer is None:
                continue

            if isinstance(answer, ParagraphSpans):
                answer = doc.answer[np.random.randint(0, len(doc.answer))]
                word_start = answer.para_word_start
                word_end = answer.para_word_end
            elif isinstance(answer, TokenSpans):
                candidates = np.where(answer.answer_spans[:, 1] < context_len[doc_ix])[0]
                if len(candidates) == 0:
                    continue
                ix = candidates[np.random.randint(0, len(candidates))]
                word_start, word_end = answer.answer_spans[ix]
            else:
                raise NotImplementedError()

            if word_start > word_end:
                raise ValueError()
            if word_end >= context_len[doc_ix]:
                raise ValueError(word_end)

            answer_spans[doc_ix, 0] = word_start
            answer_spans[doc_ix, 1] = word_end
        return {self.answer_spans: answer_spans}

    def __getstate__(self):
        return {}

    def __setstate__(self, state):
        return self.__init__() class LossEvaluator(Evaluator):

    def tensors_needed(self, _):
        return dict(loss=tf.add_n(tf.get_collection(tf.GraphKeys.LOSSES)))

    def evaluate(self, data, true_len, loss):
        return Evaluation({"loss": np.mean(loss)}) class SpanEvaluator(Evaluator):
    """
    Evaluate span based models, if text_eval is a set it should produce exactly
    the scores returned by the corresponding official evaluation scripts
    """

    def __init__(self, bound: List[int], text_eval: str=None):
        if text_eval is not None and text_eval not in ["squad", "triviaqa"]:
            raise ValueError()
        self.bound = bound
        self.text_eval = text_eval

    def tensors_needed(self, prediction):
        return {str(b): prediction.get_best_span(b)[0] for b in self.bound}

    def evaluate(self, data: List[ContextAndQuestion], true_len, **kwargs):
        ev = Evaluation({})
        for b in self.bound:
            best_spans = kwargs[str(b)]
            if self.text_eval is None:
                scores = span_scores(data, best_spans)
            elif self.text_eval == "triviaqa":
                scores = trivia_span_scores(data, best_spans)
            elif self.text_eval == "squad":
                scores = squad_span_scores(data, best_spans)
            else:
                raise RuntimeError()

            scores = scores.sum(axis=0) / true_len

            prefix = "b%d/"%b
            out = {
                prefix + "accuracy": scores[0],
                prefix + "f1": scores[1],
            }
            if self.text_eval is not None:
                out[prefix + "text-em"] = scores[2]
                out[prefix + "text-f1"] = scores[3]

            ev.add(Evaluation(out))
        return ev class BiAttention(AttentionMapper):
    """ Bi-Directonal Attention from https://arxiv.org/abs/1611.01603 """

    def __init__(self, sim: SimilarityFunction, q2c: bool, query_dots: bool=True):
        self.sim = sim
        self.q2c = q2c
        self.query_dots = query_dots

    def apply(self, is_train, x, keys, memories, x_mask=None, mem_mask=None):
        x_word_dim = tf.shape(x)[1]
        key_word_dim = tf.shape(keys)[1]

        dist_matrix = self.sim.get_scores(x, keys)
        joint_mask = compute_attention_mask(x_mask, mem_mask, x_word_dim, key_word_dim)
        if joint_mask is not None:
            dist_matrix += VERY_NEGATIVE_NUMBER * (1 - tf.cast(joint_mask, dist_matrix.dtype))
        query_probs = tf.nn.softmax(dist_matrix)  # probability of each mem_word per x_word

        # Batch matrix multiplication to get the attended vectors
        select_query = tf.matmul(query_probs, memories)  # (batch, x_words, q_dim)

        if not self.q2c:
            if self.query_dots:
                return tf.concat([x, select_query, x * select_query], axis=2)
            else:
                return tf.concat([x, select_query], axis=2)

        # select query-to-context
        context_dist = tf.reduce_max(dist_matrix, axis=2)  # (batch, x_word``s)
        context_probs = tf.nn.softmax(context_dist)  # (batch, x_words)
        select_context = tf.einsum("ai,aik->ak", context_probs, x)  # (batch, x_dim)
        select_context = tf.expand_dims(select_context, 1)

        if self.query_dots:
            return tf.concat([x, select_query, x * select_query, x * select_context], axis=2)
        else:
            return tf.concat([x, select_query, x * select_context], axis=2)

    def __setstate__(self, state):
        if "state" in state:
            if "query_dots" not in state["state"]:
                state["state"]["query_dots"] = True
        super().__setstate__(state) class FixedWordEmbedder(WordEmbedder):

    def __init__(self,
                 vec_name: str,
                 word_vec_init_scale: float = 0.05,
                 learn_unk: bool = True,
                 keep_probs: float = 1,
                 keep_word: float= 1,
                 shrink_embed: bool=False,
                 cpu=False):
        self.keep_word = keep_word
        self.keep_probs = keep_probs
        self.word_vec_init_scale = word_vec_init_scale
        self.learn_unk = learn_unk
        self.vec_name = vec_name
        self.cpu = cpu
        self.shrink_embed = shrink_embed

        # Built in `init`
        self._word_to_ix = None
        self._word_emb_mat = None
        self._special_tokens = None

    def set_vocab(self, _, loader: ResourceLoader, special_tokens: List[str]):
        if special_tokens is not None:
            self._special_tokens = sorted(special_tokens)

    def is_vocab_set(self):
        return True

    def question_word_to_ix(self, word, is_train):
        ix = self._word_to_ix.get(word, 1)
        if ix == 1:
            return self._word_to_ix.get(word.lower(), 1)
        else:
            return ix

    def context_word_to_ix(self, word, is_train):
        # print(word)
        ix = self._word_to_ix.get(word, 1)
        if ix == 1:
            return self._word_to_ix.get(word.lower(), 1)
        else:
            return ix

    @property
    def version(self):
        # added `cpu`
        return 1

    def init(self, loader: ResourceLoader, voc: Iterable[str]):
        if self.cpu:
            with tf.device("/cpu:0"):
                self._init(loader, voc)
        else:
            self._init(loader, voc)

    def _init(self, loader: ResourceLoader, voc: Iterable[str]):
        # TODO we should not be building variables here
        if voc is not None:
            word_to_vec = loader.load_word_vec(self.vec_name, voc)
        else:
            word_to_vec = loader.load_word_vec(self.vec_name)
            voc = set(word_to_vec.keys())

        self._word_to_ix = {}

        dim = next(iter(word_to_vec.values())).shape[0]

        null_embed = tf.zeros((1, dim), dtype=tf.float32)
        unk_embed = tf.get_variable(shape=(1, dim), name="unk_embed",
                                    dtype=np.float32, trainable=self.learn_unk,
                                    initializer=tf.random_uniform_initializer(-self.word_vec_init_scale,
                                                                              self.word_vec_init_scale))
        ix = 2
        matrix_list = [null_embed, unk_embed]

        if self._special_tokens is not None and len(self._special_tokens) > 0:
            print("Building embeddings for %d special_tokens" % (len(self._special_tokens)))
            tok_embed = tf.get_variable(shape=(len(self._special_tokens), dim), name="token_embed",
                                        dtype=np.float32, trainable=True,
                                        initializer=tf.random_uniform_initializer(-self.word_vec_init_scale,
                                                                                  self.word_vec_init_scale))
            matrix_list.append(tok_embed)
            for token in self._special_tokens:
                self._word_to_ix[token] = ix
                ix += 1

        mat = []
        for word in voc:
            if word in self._word_to_ix:
                continue  # in case we already added due after seeing a capitalized version of `word`
            if word in word_to_vec:
                mat.append(word_to_vec[word])
                self._word_to_ix[word] = ix
                ix += 1
            else:
                lower = word.lower()  # Full back to the lower-case version
                if lower in word_to_vec and lower not in self._word_to_ix:
                    mat.append(word_to_vec[lower])
                    self._word_to_ix[lower] = ix
                    ix += 1

        print("Had pre-trained word embeddings for %d of %d words" % (len(mat), len(voc)))

        matrix_list.append(tf.constant(value=np.vstack(mat)))

        self._word_emb_mat = tf.concat(matrix_list, axis=0)

    def embed(self, is_train, *word_ix):
        if any(len(x) != 2 for x in word_ix):
            raise ValueError()
        mat = self._word_emb_mat

        if self.keep_probs < 1:
            mat = tf.cond(is_train,
                          lambda: tf.nn.dropout(mat, self.keep_probs),
                          lambda: mat)
        if self.keep_word < 1:
            mat = tf.cond(is_train,
                          lambda: tf.nn.dropout(mat, self.keep_word, (mat.shape.as_list()[0], 1)),
                          lambda: mat)
        if self.cpu:
            with tf.device("/cpu:0"):
                return [tf.nn.embedding_lookup(self._word_emb_mat, x[0]) for x in word_ix]
        else:
            return [tf.nn.embedding_lookup(self._word_emb_mat, x[0]) for x in word_ix]

    def __getstate__(self):
        state = dict(self.__dict__)
        state["_word_emb_mat"] = None  # we will rebuild these anyway
        state["_word_to_ix"] = None
        return dict(version=self.version, state=state)

    def __setstate__(self, state):
        if "state" in state:
            if "cpu" not in state["state"]:
                state["state"]["cpu"] = False
            if "keep_probs" not in state["state"]:
                state["state"]["keep_probs"] = 1.0
            if "keep_word" not in state["state"]:
                state["state"]["keep_word"] = 1.0
            if "_special_tokens" not in state["state"]:
                state["state"]["_special_tokens"] = []
        super().__setstate__(state) 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 LearnedCharEmbedder(CharEmbedder):

    def __init__(self, word_size_th: int, char_th: int, char_dim: int,
                 init_scale: float=0.1, force_cpu: bool=False):
        self.force_cpu = force_cpu
        self.word_size_th = word_size_th
        self.char_th = char_th
        self.char_dim = char_dim
        self.init_scale = init_scale

        self._char_to_ix = None

    def get_word_size_th(self):
        return self.word_size_th

    def set_vocab(self, corpus):
        w_counts = corpus.get_word_counts()
        counts = Counter()
        for w,count in w_counts.items():
            for c in w:
                counts[c] += count

        self._char_to_ix = {c:i+2 for i,c in enumerate(c for c,count in counts.items() if count >= self.char_th)}
        print("Learning an embedding for %d characters" % (len(self._char_to_ix)-1))

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

    def char_to_ix(self, char):
        return self._char_to_ix.get(char, 1)

    def embed(self, is_train, *char_ix):
        if self.force_cpu:
            with tf.device('/cpu:0'):
                return self._embed(*char_ix)
        else:
            return self._embed(*char_ix)

    def _embed(self, *char_ix):
        zero = tf.zeros((1, self.char_dim), dtype=np.float32)
        mat = tf.get_variable("char_emb_mat", (len(self._char_to_ix)+1, self.char_dim),
                              tf.float32, initializer=tf.random_uniform_initializer(-self.init_scale, self.init_scale))
        emb_mat = tf.concat([zero, mat], axis=0)

        x = char_ix[0]
        tf.nn.embedding_lookup(emb_mat, x[0])

        return [tf.nn.embedding_lookup(emb_mat, x[0]) for x in char_ix]

    def __setstate__(self, state):
        if "state" in state:
            if "force_cpu" not in state["state"]:
                state["state"]["force_cpu"] = False
        super().__setstate__(state) class NullBiMapper(SequenceBiMapper):
    def apply(self, is_train, x, mask=None):
        return x, x class NullMapper(Updater):
    def apply(self, is_train, x, mask=None):
        return x class SequenceMapperSeq(SequenceMapper):
    def __init__(self, *layers: SequenceMapper):
        self.layers = layers

    def apply(self, is_train, x, mask=None):
        for i, layer in enumerate(self.layers):
            with tf.variable_scope("layer_" + str(i)):
                x = layer.apply(is_train, x, mask)
        return x

    def get_params(self):
        return dict(layers=[x.get_config() for x in self.layers]) class ReduceLayer(Encoder):
    def __init__(self, reduce: str, map_layer: Optional[Mapper]=None, mask=True):
        self.map_layer = map_layer
        self.reduce = reduce
        self.mask = mask

    def apply(self, is_train, x, mask=None):
        if not self.mask:
            # Ignore the mask, this option exists since BiDaF does not account for masking here as well
            mask = None

        if mask is not None:
            valid_mask = tf.cast(tf.sequence_mask(mask, tf.shape(x)[1]), tf.float32)
            for i in range(len(x.shape) - 2):
                tf.expand_dims(valid_mask, len(valid_mask.shape))
        else:
            valid_mask = None

        if self.map_layer is not None:
            x = self.map_layer.apply(is_train, x, mask)
        rank = len(x.shape) - 2

        if self.reduce == "max":
            if mask is not None:
                raise NotImplementedError()
            else:
                return tf.reduce_max(x, axis=rank)
        elif self.reduce == "mean":
            if valid_mask is not None:
                x *= valid_mask
                return tf.reduce_sum(x, axis=rank) / tf.cast(tf.expand_dims(mask, 1), tf.float32)
            else:
                return tf.reduce_mean(x, axis=rank)
        elif self.reduce == "sum":
            if valid_mask is not None:
                x *= valid_mask
            return tf.reduce_sum(x, axis=rank)
        else:
            raise ValueError()

    def __setstate__(self, state):
        if "state" in state:
            if "mask" not in state["state"]:
                state["state"]["mask"] = False
        return super().__setstate__(state) class Conv1d(Mapper):
    def __init__(self, num_filters, filter_size, keep_probs, activation="relu"):
        self.keep_probs = keep_probs
        self.num_filters = num_filters
        self.filter_size = filter_size
        self.activation = activation

    def apply(self, is_train, x, mask=None):
        num_channels = x.get_shape()[3]
        filter_ = tf.get_variable("conv1d/filters", shape=[1, self.filter_size, num_channels, self.num_filters], dtype='float')
        bias = tf.get_variable("conv1d/bias", shape=[self.num_filters], dtype='float')
        strides = [1, 1, 1, 1]
        if self.keep_probs < 1.0:
            x = dropout(x, self.keep_probs, is_train)
        fn = get_keras_activation(self.activation)
        return fn(tf.nn.conv2d(x, filter_, strides, "VALID") + bias) class HighwayLayer(Updater):
    def __init__(self, init="glorot_uniform", activation="tanh"):
        self.init = init
        self.activation = activation

    def apply(self, is_train, x, mask=None):
        n_out = x.shape.as_list()[-1]
        init = get_keras_initialization(self.init)
        activation = get_keras_activation(self.activation)
        trans = fully_connected(x, n_out, activation=activation, kernel_initializer=init, name="highway")
        gate = fully_connected(x, n_out, activation=tf.nn.sigmoid, kernel_initializer=init, name="gate")
        return gate * trans + (1 - gate) * x class ChainConcat(SequenceBiMapper):
    """ How bidaf builds start/end logits """

    def __init__(self, start_layer: SequenceMapper, end_layer: SequenceMapper,
                 soft_select_start_word: bool=True, use_original: bool=True,
                 use_start_layer: bool=True, init: str="glorot_uniform"):
        self.init = init
        self.use_original = use_original
        self.start_layer = start_layer
        self.use_start_layer = use_start_layer
        self.end_layer = end_layer
        self.soft_select_start_word = soft_select_start_word

    def apply(self, is_train, context_embed, context_mask=None):
        init_fn = get_keras_initialization(self.init)
        with tf.variable_scope("start_layer"):
            m1 = self.start_layer.apply(is_train, context_embed, context_mask)

        with tf.variable_scope("start_pred"):
            logits1 = fully_connected(tf.concat([m1, context_embed], axis=2), 1,
                                      activation=None, kernel_initializer=init_fn)
            masked_logits1 = exp_mask(tf.squeeze(logits1, squeeze_dims=[2]), context_mask)
            prediction1 = tf.nn.softmax(masked_logits1)

        m2_input = []
        if self.use_original:
            m2_input.append(context_embed)
        if self.use_start_layer:
            m2_input.append(m1)
        if self.soft_select_start_word:
            soft_select = tf.einsum("ai,aik->ak", prediction1, m1)
            soft_select_tiled = tf.tile(tf.expand_dims(soft_select, axis=1), [1, tf.shape(m1)[1], 1])
            m2_input += [soft_select_tiled, soft_select_tiled * m1]

        with tf.variable_scope("end_layer"):
            m2 = self.end_layer.apply(is_train, tf.concat(m2_input, axis=2), context_mask)

        with tf.variable_scope("end_pred"):
            logits2 = fully_connected(tf.concat([m2, context_embed], axis=2), 1,
                                      activation=None, kernel_initializer=init_fn)

        return logits1, logits2

    def __setstate__(self, state):
        if "state" in state:
            if "aggregate" not in state["state"]:
                state["state"]["aggregate"] = None
        return super().__setstate__(state) class DropoutLayer(Updater):
    def __init__(self, keep_probs: float):
        self.keep_prob = keep_probs

    def apply(self, is_train, x, mask=None):
        return dropout(x, self.keep_prob, is_train) class CudnnLstm(CudnnRnnMapper, SequenceMapper):
    def __init__(self,
                 n_units,
                 n_layers=1,
                 lstm_bias=1,
                 w_init=TruncatedNormal(stddev=0.05),
                 recurrent_init=None,
                 bidirectional=True,
                 learn_initial_states=False):
        super().__init__("LSTM", n_units, n_layers, w_init, recurrent_init, bidirectional,
                         learn_initial_states, lstm_bias)

    def apply(self, is_train, x, mask=None):
        return super().map(is_train, x, mask)

    def __setstate__(self, state):
        if "state" in state:
            if "recurrent_init" not in state["state"]:
                state["state"]["recurrent_init"] = None
            if "keep_recurrent" not in state["state"]:
                state["state"]["keep_recurrent"] = 1
        super().__setstate__(state) class TriLinear(_WithBias):
    """ Function used by BiDaF, bi-linear with an extra component for the dots of the vectors """
    def __init__(self, init="glorot_uniform", bias=False):
        super().__init__(bias)
        self.init = init

    def _distance_logits(self, x, keys):
        init = get_keras_initialization(self.init)

        key_w = tf.get_variable("key_w", shape=keys.shape.as_list()[-1], initializer=init, dtype=tf.float32)
        key_logits = tf.tensordot(keys, key_w, axes=[[2], [0]])  # (batch, key_len)

        x_w = tf.get_variable("input_w", shape=x.shape.as_list()[-1], initializer=init, dtype=tf.float32)
        x_logits = tf.tensordot(x, x_w, axes=[[2], [0]])  # (batch, x_len)

        dot_w = tf.get_variable("dot_w", shape=x.shape.as_list()[-1], initializer=init, dtype=tf.float32)

        # Compute x * dot_weights first, the batch mult with x
        x_dots = x * tf.expand_dims(tf.expand_dims(dot_w, 0), 0)
        dot_logits = tf.matmul(x_dots, keys, transpose_b=True)

        return dot_logits + tf.expand_dims(key_logits, 1) + tf.expand_dims(x_logits, 2)

    @property
    def version(self):
        return 1 class BoundsPredictor(SequencePredictionLayer):
    """ Standard start/end bound prediction """

    def __init__(self, predictor: SequenceBiMapper, init: str="glorot_uniform",
                 span_predictor: SpanFromBoundsPredictor = IndependentBounds()):
        self.predictor = predictor
        self.init = init
        self.span_predictor = span_predictor

    def apply(self, is_train, context_embed, answer, context_mask=None):
        init_fn = get_keras_initialization(self.init)
        with tf.variable_scope("bounds_encoding"):
            m1, m2 = self.predictor.apply(is_train, context_embed, context_mask)

        with tf.variable_scope("start_pred"):
            logits1 = fully_connected(m1, 1, activation_fn=None,
                                      weights_initializer=init_fn)
            logits1 = tf.squeeze(logits1, squeeze_dims=[2])

        with tf.variable_scope("end_pred"):
            logits2 = fully_connected(m2, 1, activation_fn=None, weights_initializer=init_fn)
            logits2 = tf.squeeze(logits2, squeeze_dims=[2])

        with tf.variable_scope("predict_span"):
            return self.span_predictor.predict(answer, logits1, logits2, context_mask)

    def __setstate__(self, state):
        if "state" in state:
            if "aggregate" in state["state"]:
                state["state"]["bound_predictor"] = IndependentBounds(state["state"]["aggregate"])
            elif "bound_predictor" not in state:
                state["state"]["bound_predictor"] = IndependentBounds()
        super().__setstate__(state) def clean_title(title):
def parse_squad_data(source, name, tokenizer, use_tqdm=True) -> List[Document]:
def main(): class DocumentQaTrainingData(ParagraphQaTrainingData):
    def _preprocess(self, x):
        data = split_docs(x)
        return data, len(data) class SerializableOptimizer(Configurable):
    """ So we can record what tensorflow optimizer we used """

    def __init__(self, opt_name, params=None):
        self.params = params
        self.opt_name = opt_name

    def get_params(self):
        return dict(opt_name=self.opt_name, params=self.params)

    def get(self, name=None):
        params = {} if self.params is None else self.params
        if self.opt_name == "Adam":
            if name is None:
                return AdamOptimizer(**params)
            else:
                return AdamOptimizer(name=name, **params)
        elif self.opt_name == "Adadelta":
            if name is None:
                return AdadeltaOptimizer(**params)
            else:
                return AdadeltaOptimizer(name=name, **params)
        else:
            raise NotImplemented() class TrainParams(Configurable):
    """ Parameters related to training """

    def __init__(self,
                 opt: SerializableOptimizer,
                 num_epochs: int,
                 eval_period: int,
                 log_period: int,
                 save_period: int,
                 eval_samples: Dict[str, Optional[int]],
                 regularization_weight: Optional[float] = None,
                 async_encoding: Optional[int] = None,
                 max_checkpoints_to_keep: int = 5,
                 loss_ema: Optional[float] = .999,
                 eval_at_zero: bool = False,
                 monitor_ema: float = .999,
                 ema: Optional[float] = None,
                 best_weights: Optional[Tuple[str, str]] = None
                 ):
        """
        :param opt: Optimizer to use
        :param num_epochs: Number of epochs to train for
        :param eval_period: How many batches to train on between evaluations
        :param log_period: How many batches to train on between logging
        :param save_period: How many batches to train on between checkpointing
        :param eval_samples: How many samples to draw during evaluation, None of a full epoch
        :param regularization_weight: How highly to weight regulraization, defaults to 1
        :param async_encoding: Encoding batches in a seperate thread, and store in a queue of this size
        :param max_checkpoints_to_keep: Max number of checkpoints to keep during training
        :param loss_ema: EMA weights for monitoring the loss during training
        :param eval_at_zero: Run an evaluation cycle before any training
        :param monitor_ema: EMA weights for monitor functions
        :param ema: EMA to use on the trainable parameters
        :param best_weights: Store the weights with the highest scores on the given eval dataset/metric
        """
        self.async_encoding = async_encoding
        self.regularization_weight = regularization_weight
        self.max_checkpoints_to_keep = max_checkpoints_to_keep
        self.opt = opt
        self.eval_at_zero = eval_at_zero
        self.ema = ema
        self.loss_ema = loss_ema
        self.monitor_ema = monitor_ema
        self.num_epochs = num_epochs
        self.eval_period = eval_period
        self.log_period = log_period
        self.save_period = save_period
        self.eval_samples = eval_samples
        self.best_weights = best_weights def get_output_name_from_cli():
    parser = argparse.ArgumentParser(description='')
    parser.add_argument('--name', '-n', nargs=1, help='name of output to exmaine')

    args = parser.parse_args()
    if args.name:
        out = join(args.name[0] + "-" + datetime.now().strftime("%m%d-%H%M%S"))
        print("Starting run on: " + out)
    else:
        out = "out/run-" + datetime.now().strftime("%m%d-%H%M%S")
        print("Starting run on: " + out)
    return out
from docqa import model_dir
from docqa import trainer
from docqa.data_processing.qa_training_data import ContextLenBucketedKey, ContextLenKey
from docqa.dataset import ClusteredBatcher
from docqa.doc_qa_models import Attention
from docqa.encoder import DocumentAndQuestionEncoder, SingleSpanAnswerEncoder
from docqa.evaluator import LossEvaluator, SpanEvaluator
from docqa.nn.attention import BiAttention
from docqa.nn.embedder import FixedWordEmbedder, CharWordEmbedder, LearnedCharEmbedder
from docqa.nn.layers import NullBiMapper, NullMapper, SequenceMapperSeq, ReduceLayer, Conv1d, HighwayLayer, ChainConcat, \
    DropoutLayer
from docqa.nn.recurrent_layers import CudnnLstm
from docqa.nn.similarity_layers import TriLinear
from docqa.nn.span_prediction import BoundsPredictor
from docqa.squad.build_squad_dataset import SquadCorpus
from docqa.squad.squad_data import DocumentQaTrainingData
from docqa.trainer import SerializableOptimizer, TrainParams
from docqa.utils import get_output_name_from_cli


def main():
    """
    A close-as-possible impelemntation of BiDaF, its based on the `dev` tensorflow 1.1 branch of Ming's repo
    which, in particular, uses Adam not Adadelta. I was not able to replicate the results in paper using Adadelta,
    but with Adam i was able to get to 78.0 F1 on the dev set with this scripts. I believe this approach is
    an exact reproduction up the code in the repo, up to initializations.

    Notes: Exponential Moving Average is very important, as is early stopping. This is also in particualr best run
    on a GPU due to the large number of parameters and batch size involved.
    """
    out = get_output_name_from_cli()

    train_params = TrainParams(SerializableOptimizer("Adam", dict(learning_rate=0.001)),
                               num_epochs=12, ema=0.999, async_encoding=10,
                               log_period=30, eval_period=1000, save_period=1000,
                               eval_samples=dict(dev=None, train=8000))

    # recurrent_layer = BiRecurrentMapper(LstmCellSpec(100, keep_probs=0.8))
    # recurrent_layer = FusedLstm()
    recurrent_layer = SequenceMapperSeq(DropoutLayer(0.8), CudnnLstm(100))

    model = Attention(
        encoder=DocumentAndQuestionEncoder(SingleSpanAnswerEncoder()),
        word_embed=FixedWordEmbedder(vec_name="glove.6B.100d", word_vec_init_scale=0, learn_unk=False),
        char_embed=CharWordEmbedder(
            embedder=LearnedCharEmbedder(16, 49, 8),
            layer=ReduceLayer("max", Conv1d(100, 5, 0.8), mask=False),
            shared_parameters=True
        ),
        word_embed_layer=None,
        embed_mapper=SequenceMapperSeq(
            HighwayLayer(activation="relu"), HighwayLayer(activation="relu"),
            recurrent_layer),
        preprocess=None,
        question_mapper=None,
        context_mapper=None,
        memory_builder=NullBiMapper(),
        attention=BiAttention(TriLinear(bias=True), True),
        match_encoder=NullMapper(),
        predictor= BoundsPredictor(
            ChainConcat(
                start_layer=SequenceMapperSeq(
                    recurrent_layer,
                    recurrent_layer),
                end_layer=recurrent_layer
            )
        ),

    )

    with open(__file__, "r") as f:
        notes = f.read()


    eval = [LossEvaluator(), SpanEvaluator(bound=[17], text_eval="squad")]