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 MergeParagraphs(DocumentSplitter):
    """
    Merge paragraphs up to a maximum size. Paragraphs of a larger size will be truncated.
    """

    def __init__(self, max_tokens: int, top_n: int=None):
        self.max_tokens = max_tokens
        self.top_n = top_n

    @property
    def reads_first_n(self):
        return self.top_n

    def max_tokens(self):
        return self.max_tokens

    def split(self, doc: List[List[List[str]]]):
        all_paragraphs = []

        on_doc_token = 0  # the word in the document the current paragraph starts at
        on_paragraph = []  # text we have collect for the current paragraph
        cur_tokens = 0   # number of tokens in the current paragraph

        word_ix = 0
        for para in doc:
            para = flatten_iterable(para)
            n_words = len(para)
            if self.top_n is not None and (word_ix+self.top_n)>self.top_n:
                if word_ix == self.top_n:
                    break
                para = para[:self.top_n - word_ix]
                n_words = self.top_n - word_ix

            start_token = word_ix
            end_token = start_token + n_words
            word_ix = end_token

            if cur_tokens + n_words > self.max_tokens:
                if cur_tokens != 0:  # end the current paragraph
                    all_paragraphs.append(ExtractedParagraph(on_paragraph, on_doc_token, start_token))
                    on_paragraph = []
                    cur_tokens = 0

                if n_words >= self.max_tokens:  # either truncate the given paragraph, or begin a new paragraph
                    text = para[:self.max_tokens]
                    all_paragraphs.append(ExtractedParagraph([text], start_token,
                                                             start_token + self.max_tokens))
                    on_doc_token = end_token
                else:
                    on_doc_token = start_token
                    on_paragraph.append(para)
                    cur_tokens = n_words
            else:
                on_paragraph.append(para)
                cur_tokens += n_words

        if len(on_paragraph) > 0:
            all_paragraphs.append(ExtractedParagraph(on_paragraph, on_doc_token, word_ix))

        return all_paragraphs class TopTfIdf(ParagraphFilter):
    def __init__(self, stop, n_to_select: int, filter_dist_one: bool=False, rank=True):
        self.stop = stop
        self.rank = rank
        self.n_to_select = n_to_select
        self.filter_dist_one = filter_dist_one

    def prune(self, question, paragraphs: List[ExtractedParagraph]):
        if not self.filter_dist_one and len(paragraphs) == 1:
            return paragraphs

        tfidf = TfidfVectorizer(strip_accents="unicode", stop_words=self.stop.words)
        text = []
        for para in paragraphs:
            text.append(" ".join(" ".join(s) for s in para.text))
        try:
            para_features = tfidf.fit_transform(text)
            q_features = tfidf.transform([" ".join(question)])
        except ValueError:
            return []

        dists = pairwise_distances(q_features, para_features, "cosine").ravel()
        sorted_ix = np.lexsort(([x.start for x in paragraphs], dists))  # in case of ties, use the earlier paragraph

        if self.filter_dist_one:
            return [paragraphs[i] for i in sorted_ix[:self.n_to_select] if dists[i] < 1.0]
        else:
            return [paragraphs[i] for i in sorted_ix[:self.n_to_select]]

    def dists(self, question, paragraphs: List[ExtractedParagraph]):
        tfidf = TfidfVectorizer(strip_accents="unicode", stop_words=self.stop.words)
        text = []
        for para in paragraphs:
            text.append(" ".join(" ".join(s) for s in para.text))
        try:
            para_features = tfidf.fit_transform(text)
            q_features = tfidf.transform([" ".join(question)])
        except ValueError:
            return []

        dists = pairwise_distances(q_features, para_features, "cosine").ravel()
        sorted_ix = np.lexsort(([x.start for x in paragraphs], dists))  # in case of ties, use the earlier paragraph

        if self.filter_dist_one:
            return [(paragraphs[i], dists[i]) for i in sorted_ix[:self.n_to_select] if dists[i] < 1.0]
        else:
            return [(paragraphs[i], dists[i]) for i in sorted_ix[:self.n_to_select]] class StratifyParagraphsBuilder(DatasetBuilder):
    def __init__(self, batcher: ListBatcher,  oversample_answers: Union[int, List[int]],
                 only_answers: bool=False):
        self.batcher = batcher
        self.oversample_answers = oversample_answers
        self.only_answers = only_answers

    def build_dataset(self, data, evidence) -> Dataset:
        if isinstance(data, FilteredData):
            l = data.true_len
            data = data.data
        else:
            l = len(data)
        if self.only_answers:
            for q in data:
                q.paragraphs = [x for x in q.paragraphs if len(x.answer_spans) > 0]
            data = [x for x in data if len(x.paragraphs) > 0]

        if isinstance(self.oversample_answers, int):
            ov = [self.oversample_answers]
        else:
            ov = self.oversample_answers
        return StratifyParagraphsDataset(data, l, ov, self.batcher)

    @property
    def version(self):
        # Changed how sampling works
        return 2

    def build_stats(self, data) -> object:
        if isinstance(data, FilteredData):
            return multi_paragraph_word_counts(data.data)
        else:
            return multi_paragraph_word_counts(data)

    def __setstate__(self, state):
        if "only_answers" not in state:
            state["only_answers"] = False
        print(state)
        if state.get("oversample", 0) != 0:
            raise NotImplementedError()
        if "oversample_first" in state:
            state["oversample_answers"] = [state["oversample_first"]]
            del state["oversample_first"]

        super().__setstate__(state) class StratifyParagraphSetsBuilder(DatasetBuilder):
    def __init__(self, batch_size: int, merge: bool, force_answer: bool,
                 oversample_first_answer: Union[int, List[int]]):
        self.batch_size = batch_size
        self.merge = merge
        self.force_answer = force_answer
        self.oversample_first_answer = oversample_first_answer

    def build_stats(self, data: Union[List, FilteredData]):
        if isinstance(data, FilteredData):
            return multi_paragraph_word_counts(data.data)
        else:
            return multi_paragraph_word_counts(data)

    def build_dataset(self, data: Union[FilteredData, List], corpus) -> Dataset:
        if isinstance(data, FilteredData):
            l = data.true_len
            data = data.data
        else:
            l = len(data)
        if isinstance(self.oversample_first_answer, int):
            ov = [self.oversample_first_answer]
        else:
            ov = self.oversample_first_answer
        return StratifiedParagraphSetDataset(data, l, self.batch_size, self.force_answer,
                                             ov, self.merge) class RandomParagraphSetDatasetBuilder(DatasetBuilder):
    def __init__(self, batch_size: int, mode: str, force_answer: bool,
                 oversample_first_answer: Union[int, List[int]]):
        self.mode = mode
        self.oversample_first_answer = oversample_first_answer
        self.batch_size = batch_size
        self.force_answer = force_answer

    def build_stats(self, data: Union[FilteredData, List]):
        if isinstance(data, FilteredData):
            return multi_paragraph_word_counts(data.data)
        else:
            return multi_paragraph_word_counts(data)

    def build_dataset(self, data: Union[FilteredData, List], corpus) -> Dataset:
        if isinstance(data, FilteredData):
            l = data.true_len
            data = data.data
        else:
            l = len(data)
        if isinstance(self.oversample_first_answer, int):
            ov = [self.oversample_first_answer]
        else:
            ov = self.oversample_first_answer
        return RandomParagraphSetDataset(data, l, 2, self.batch_size, self.mode, self.force_answer, ov) class PreprocessedData(TrainingData):
    """
    Data the goes through a preprocessing pipeline, for TriviaQA this usually mean leading/choosing what
    paragraphs we want to train on, the organizing them into a dataset with the desired sampling strategy
    """

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

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

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

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

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

        print("done")

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

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

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

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

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

        print("Done")

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

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

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

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

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

    def __getstate__(self):
        state = dict(self.__dict__)
        state["_train"] = None
        state["_dev"] = None
        state["_verified_dev"] = None
        return state class ParagraphAndQuestionsBuilder(DatasetBuilder):
    """ For use with the preprocesed_corpus framework """
    def __init__(self, batching: ListBatcher, sample=None, sample_seed=None):
        if sample_seed is not None and sample is None:
            raise ValueError("Seed set, but sampling not requested")
        self.batching = batching
        self.sample_seed = sample_seed
        self.sample = sample

    def build_stats(self, data):
        if isinstance(data, FilteredData):
            return QaCorpusLazyStats(data.data)
        else:
            return QaCorpusLazyStats(data)

    def build_dataset(self, data, evidence) -> Dataset:
        if isinstance(data, FilteredData):
            data, l = data.data, data.true_len
        else:
            data, l = data, None
        if self.sample is not None:
            cur_len = len(data)
            data = np.random.RandomState(self.sample_seed).choice(data, self.sample, replace=False)
            if l is not None:
                l *= len(data) / cur_len

        if l is None:
            l = len(data)
        print("Building dataset")
        print(len(data), l)
        return ParagraphAndQuestionDataset(data, self.batching, l) class ContextLenKey(Configurable):
    def __call__(self, q: ContextAndQuestion):
        return q.n_context_words 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 NltkPlusStopWords(Configurable):
    """ Configurablable access to stop word """

    def __init__(self, punctuation=False):
        self._words = None
        self.punctuation = punctuation

    @property
    def words(self):
        if self._words is None:
            self._words = set(stopwords.words('english'))
            # Common question words we probably want to ignore, "de" was suprisingly common
            # due to its appearance in person names
            self._words.update(["many", "how", "de"])
            if self.punctuation:
                self._words.update(string.punctuation)
                self._words.update(["£", "€", "¥", "¢", "₹", "\u2212",
                                    "\u2014", "\u2013", "\ud01C", "\u2019", "\u201D", "\u2018", "\u00B0"])
        return self._words

    def __getstate__(self):
        return dict(punctuation=self.punctuation)

    def __setstate__(self, state):
        self.__init__(**state) 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 DenseMultiSpanAnswerEncoder(AnswerEncoder):
    """ Encode the answer spans into bool (span_start) and (span_end) arrays """

    def __init__(self):
        self.answer_starts = None
        self.answer_ends = None

    def get_placeholders(self) -> List:
        return [self.answer_starts, self.answer_ends]

    def init(self, batch_size, context_word_dim):
        self.answer_starts = tf.placeholder('bool', [batch_size, context_word_dim], name='answer_starts')
        self.answer_ends = tf.placeholder('bool', [batch_size, context_word_dim], name='answer_ends')

    def encode(self, batch_size, context_len, context_word_dim, batch) -> Dict:
        answer_starts = np.zeros((batch_size, context_word_dim), dtype=np.bool)
        answer_ends = np.zeros((batch_size, context_word_dim), dtype=np.bool)
        for doc_ix, doc in enumerate(batch):
            if doc.answer is None:
                continue
            answer_spans = doc.answer.answer_spans
            answer_spans = answer_spans[answer_spans[:, 1] < context_word_dim]
            answer_starts[doc_ix, answer_spans[:, 0]] = True
            answer_ends[doc_ix, answer_spans[:, 1]] = True
        return {self.answer_starts: answer_starts, self.answer_ends: answer_ends}

    def __getstate__(self):
        return {}

    def __setstate__(self, state):
        return DenseMultiSpanAnswerEncoder() class GroupedSpanAnswerEncoder(AnswerEncoder):
    """ Encode the answer spans into bool (span_start) and (span_end) arrays, and also record
    the group_id if one is present in the answer. Used for the "shared-norm" approach """

    def __init__(self):
        self.answer_starts = None
        self.answer_ends = None
        self.group_ids = None

    def get_placeholders(self) -> List:
        return [self.answer_starts, self.answer_ends, self.group_ids]

    def init(self, batch_size, context_word_dim):
        self.answer_starts = tf.placeholder('bool', [batch_size, context_word_dim], name='answer_starts')
        self.answer_ends = tf.placeholder('bool', [batch_size, context_word_dim], name='answer_ends')
        self.group_ids = tf.placeholder('int32', [batch_size], name='group_ids')

    def encode(self, batch_size, context_len, context_word_dim, batch) -> Dict:
        has_group = hasattr(batch[0].answer, "group_id")
        answer_starts = np.zeros((batch_size, context_word_dim), dtype=np.bool)
        answer_ends = np.zeros((batch_size, context_word_dim), dtype=np.bool)
        group_id = np.zeros(batch_size, dtype=np.int32)
        for doc_ix, doc in enumerate(batch):
            if has_group:
                group_id[doc_ix] = doc.answer.group_id
            if doc.answer is None:
                continue
            answer_spans = doc.answer.answer_spans
            answer_spans = answer_spans[answer_spans[:, 1] < context_word_dim]

            answer_starts[doc_ix, answer_spans[:, 0]] = True
            answer_ends[doc_ix, answer_spans[:, 1]] = True
        if not has_group:
            group_id = np.arange(0, len(batch))
        return {self.answer_starts: answer_starts, self.answer_ends: answer_ends,
                self.group_ids:group_id}

    def __getstate__(self):
        return {}

    def __setstate__(self, state):
        return GroupedSpanAnswerEncoder() 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 MultiParagraphSpanEvaluator(Evaluator):
    """
    Measure error with multiple paragraphs per a question.

    Evaluation is a bit tricky in this case, since we are generally sampling paragraphs
    each epoch we can't report exact numbers as your would see when running the
    evaluation scripts. Instead we report some numbers aimed to get an approximate idea of what is going on:

    1: question-text-{em|f1}, accuracy on questions-document pairs (or just questions if `per_doc=False`)
       using all sampled paragraphs when taking the model's highest confidence answer.
       This tends to be an overly-confident estimate since the sampled paragraphs are usually biased
       towards using paragraphs that contain the correct answer
    2: paragraph-text-{em|f1}, accuracy on answer-containing paragraphs (if `paragraph_level=True`)
    3: The Kendel Tau relation between the model's confidence and the paragraph's f1/em score,
       (if `k_tau=True`) intended to measure how valid the model's confidence score is
       when it comes to ranking.
    """

    def __init__(self, bound: int, eval, paragraph_level=True, k_tau=True,
                 per_doc=True):
        if eval not in ["squad", "triviaqa"]:
            raise ValueError()
        self.bound = bound
        self.eval = eval
        self.paragraph_level = paragraph_level
        self.k_tau = k_tau
        self.per_doc = per_doc

    def tensors_needed(self, prediction):
        span, score = prediction.get_best_span(self.bound)
        return dict(span=span, score=score)

    def evaluate(self, data: List[ContextAndQuestion], true_len, **kwargs):
        best_spans = kwargs["span"]
        span_logits = kwargs["score"]
        if self.eval == "triviaqa":
            scores = trivia_span_scores(data, best_spans)
        elif self.eval == "squad":
            scores = squad_span_scores(data, best_spans)
        else:
            raise RuntimeError()

        has_answer = np.array([len(x.answer.answer_spans) > 0 for x in data])

        selected_paragraphs = {}
        for i, point in enumerate(data):
            if self.per_doc:
                key = (point.question_id, point.doc_id)
            else:
                key = point.question_id
            if key not in selected_paragraphs:
                selected_paragraphs[key] = i
            elif span_logits[i] > span_logits[selected_paragraphs[key]]:
                selected_paragraphs[key] = i
        selected_paragraphs = list(selected_paragraphs.values())

        out = {
            "question-text-em": scores[selected_paragraphs, 2].mean(),
            "question-text-f1": scores[selected_paragraphs, 3].mean(),
        }

        if self.k_tau:
            out["text-em-k-tau"] = kendalltau(span_logits, scores[:, 2])[0]
            out["text-f1-k-tau"] = kendalltau(span_logits, scores[:, 3])[0]

        if self.paragraph_level:
            out["paragraph-text-em"] = scores[has_answer, 2].mean()
            out["paragraph-text-f1"] = scores[has_answer, 3].mean()

        prefix = "b%d/" % self.bound
        return Evaluation({prefix+k: v for k,v in out.items()})

    def __setstate__(self, state):
        if "per_doc" not in state:
            state["per_doc"] = True
        super().__setstate__(state) 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 AttentionEncoder(SequenceEncoder):
    def __init__(self, key_mapper: SequenceMapper=None,
                 post_process: Mapper=None,
                 init="glorot_uniform"):
        self.init = init
        self.key_mapper = key_mapper
        self.post_process = post_process

    def apply(self, is_train, x, mask=None):
        if self.key_mapper is not None:
            with tf.variable_scope("map_keys"):
                keys = self.key_mapper.apply(is_train, x, mask)
        else:
            keys = x

        weights = tf.get_variable("weights", keys.shape.as_list()[-1], dtype=tf.float32,
                                  initializer=get_keras_initialization(self.init))
        dist = tf.tensordot(keys, weights, axes=[[2], [0]])  # (batch, x_words)
        dist = exp_mask(dist, mask)
        dist = tf.nn.softmax(dist)

        out = tf.einsum("ajk,aj->ak", x, dist)  # (batch, x_dim)

        if self.post_process is not None:
            with tf.variable_scope("post_process"):
                out = self.post_process.apply(is_train, out)
        return out class StaticAttentionSelf(SequenceMapper):
    """ Basic non-recurrent attention a sequence and itself using the given SimilarityFunction """

    def __init__(self, attention: SimilarityFunction,
                 merge: Optional[MergeLayer]=None,
                 alignment_bias=True):
        self.alignment_bias = alignment_bias
        self.attention = attention
        self.merge = merge

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

        # (batch, x_word, key_word)
        dist_matrix = self.attention.get_scores(x, x)
        dist_matrix += tf.expand_dims(tf.eye(x_word_dim) * VERY_NEGATIVE_NUMBER, 0)  # Mask out self

        joint_mask = compute_attention_mask(x_mask, x_mask, x_word_dim, x_word_dim)
        if joint_mask is not None:
            dist_matrix += VERY_NEGATIVE_NUMBER * (1 - tf.cast(joint_mask, dist_matrix.dtype))

        if not self.alignment_bias:
            select_probs = tf.nn.softmax(dist_matrix)
        else:
            # Allow zero-attention by adding a learned bias to the normalizer
            bias = tf.exp(tf.get_variable("no-alignment-bias", initializer=tf.constant(-1.0, dtype=tf.float32)))
            dist_matrix = tf.exp(dist_matrix)
            select_probs = dist_matrix / (tf.reduce_sum(dist_matrix, axis=2, keep_dims=True) + bias)

        response = tf.matmul(select_probs, x)  # (batch, x_words, q_dim)

        if self.merge is not None:
            with tf.variable_scope("merge"):
                response = self.merge.apply(is_train, response, x)
            return response
        else:
            return response

    def __setstate__(self, state):
        if "state" in state:
            state["state"]["alignment_bias"] = 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 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 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 FullyConnected(Mapper):
    def __init__(self, n_out,
                 w_init="glorot_uniform",
                 activation: Union[str, Updater, None]="relu",
                 bias=True):
        self.w_init = w_init
        self.activation = activation
        self.n_out = n_out
        self.bias = bias

    def apply(self, is_train, x, mask=None):
        bias = (self.bias is None) or self.bias  # for backwards compat
        if isinstance(self.activation, Updater):
            out = fully_connected(x, self.n_out,
                                   use_bias=bias,
                                   activation=None,
                                   kernel_initializer=get_keras_initialization(self.w_init))
            with tf.variable_scope("activation"):
                return self.activation.apply(is_train, out, mask)
        else:
            return fully_connected(x, self.n_out,
                                   use_bias=bias,
                                   activation=get_keras_activation(self.activation),
                                   kernel_initializer=get_keras_initialization(self.w_init)) class ChainBiMapper(SequenceBiMapper):
    def __init__(self, first_layer: SequenceMapper, second_layer: SequenceMapper):
        self.first_layer = first_layer
        self.second_layer = second_layer

    def apply(self, is_train, x, mask=None):
        with tf.variable_scope("out"):
            m1 = self.first_layer.apply(is_train, x, mask)
        with tf.variable_scope("chained-out"):
            m2 = self.second_layer.apply(is_train, tf.concat([x, m1], axis=2), mask)
        return m1, m2 class ConcatWithProduct(MergeLayer):
    def apply(self, is_train, tensor1, tensor2) -> tf.Tensor:
        return tf.concat([tensor1, tensor2, tensor1 * tensor2], axis=len(tensor1.shape) - 1) class ResidualLayer(Mapper):

    def __init__(self, other: Union[Mapper, SequenceMapper]):
        self.other = other

    def apply(self, is_train, x, mask=None):
        return x + self.other.apply(is_train, x, mask) class VariationalDropoutLayer(SequenceMapper):
    """
    `VariationalDropout` is an overload term, but this is in particular referring to
    https://arxiv.org/pdf/1506.02557.pdf were the dropout mask is consistent across the time dimension
    """

    def __init__(self, keep_prob: float):
        self.keep_prob = keep_prob

    def apply(self, is_train, x, mask=None):
        shape = tf.shape(x)
        return dropout(x, self.keep_prob, is_train, [shape[0], 1, shape[2]]) class MaxPool(Encoder):
    def __init__(self, map_layer: Optional[Mapper]=None, min_val: float=0):
        self.map_layer = map_layer
        self.min_val = min_val

    def apply(self, is_train, x, mask=None):
        if self.map_layer is not None:
            x = self.map_layer.apply(is_train, x, mask)

        rank = len(x.shape) - 2
        if mask is not None:
            shape = tf.shape(x)
            mask = tf.sequence_mask(tf.reshape(mask, (-1,)), shape[-2])
            mask = tf.cast(tf.reshape(mask, (shape[0], shape[1], shape[2], 1)), tf.float32)
            # this min_val thing is kind of a hack, really we should do something like compute the
            # min val over the entire batch, or maybe just pick a very negative values, or maybe
            # do something a bit more finicky with tf.bool_mask
            # In practice it doesn't seem to be problem, and some of the earlier models used these
            # scheme so I have been sticking with it.
            if self.min_val == 0:
                x *= mask
            else:
                x = x * mask + self.min_val * (1 - mask)
            return tf.maximum(tf.reduce_max(x, axis=rank), tf.fill([1] * (len(x.shape)-1),
                                                                   float(self.min_val)))
        else:
            return tf.reduce_max(x, axis=rank)

    def __setstate__(self, state):
        if "min_val" not in state:
            state["min_val"] = 0
        if "state" in state:
            if "min_val" not in state["state"]:
                state["state"]["min_val"] = 0
        super().__setstate__(state) class CudnnGru(CudnnRnnMapper, SequenceMapper):
    def __init__(self,
                 n_units,
                 n_layers=1,
                 keep_recurrent=1,
                 w_init=TruncatedNormal(stddev=0.05),
                 recurrent_init=None,
                 bidirectional=True,
                 learn_initial_states=False):
        super().__init__("GRU", n_units, n_layers, w_init, recurrent_init, bidirectional,
                         learn_initial_states, 1, keep_recurrent)

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

    def __setstate__(self, state):
        if "state" in state:
            if "_kind" not in state["state"]:
                state["state"]["_kind"] = "GRU"
            if "learn_initial_states" not in state["state"]:
                state["state"]["learn_initial_states"] = False
            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 ConfidencePredictor(SequencePredictionLayer):
    """
    Bound prediction where we compute a non-answer logit/option using soft attention over
    the start/end logit and a `SequenceEncoder`.
    """

    def __init__(self,
                 predictor: SequenceBiMapper,
                 encoder: Union[SequenceEncoder, SequenceMultiEncoder],
                 confidence_predictor: Mapper,
                 init: str="glorot_uniform",
                 aggregate=None):
        self.predictor = predictor
        self.init = init
        self.aggregate = aggregate
        self.confidence_predictor = confidence_predictor
        self.encoder = encoder

    @property
    def version(self):
        return 1  # Fix masking

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

        if m1.shape.as_list()[-1] != 1:
            with tf.variable_scope("start_pred"):
                start_logits = fully_connected(m1, 1, activation_fn=None,
                                          weights_initializer=init_fn)
        else:
            start_logits = m1
        start_logits = tf.squeeze(start_logits, squeeze_dims=[2])

        if m1.shape.as_list()[-1] != 1:
            with tf.variable_scope("end_pred"):
                end_logits = fully_connected(m2, 1, activation_fn=None, weights_initializer=init_fn)
        else:
            end_logits = m2
        end_logits = tf.squeeze(end_logits, squeeze_dims=[2])

        masked_start_logits = exp_mask(start_logits, context_mask)
        masked_end_logits = exp_mask(end_logits, context_mask)

        start_atten = tf.einsum("ajk,aj->ak", m1, tf.nn.softmax(masked_start_logits))
        end_atten = tf.einsum("ajk,aj->ak", m2, tf.nn.softmax(masked_end_logits))
        with tf.variable_scope("encode_context"):
            enc = self.encoder.apply(is_train, context_embed, context_mask)
        if len(enc.shape) == 3:
            _, encodings, fe = enc.shape.as_list()
            enc = tf.reshape(enc, (-1, encodings*fe))

        with tf.variable_scope("confidence"):
            conf = [start_atten, end_atten, enc]
            none_logit = self.confidence_predictor.apply(is_train, tf.concat(conf, axis=1))
        with tf.variable_scope("confidence_logits"):
            none_logit = fully_connected(none_logit, 1, activation_fn=None,
                                   weights_initializer=init_fn)
            none_logit = tf.squeeze(none_logit, axis=1)

        batch_dim = tf.shape(start_logits)[0]

        # (batch, (l * l)) logits for each (start, end) pair
        all_logits = tf.reshape(tf.expand_dims(masked_start_logits, 1) +
                                tf.expand_dims(masked_end_logits, 2),
                                (batch_dim, -1))

        # (batch, (l * l) + 1) logits including the none option
        all_logits = tf.concat([all_logits, tf.expand_dims(none_logit, 1)], axis=1)
        log_norms = tf.reduce_logsumexp(all_logits, axis=1)

        # Now build a "correctness" mask in the same format
        correct_mask = tf.logical_and(tf.expand_dims(answer[0], 1), tf.expand_dims(answer[1], 2))
        correct_mask = tf.reshape(correct_mask, (batch_dim, -1))
        correct_mask = tf.concat([correct_mask, tf.logical_not(tf.reduce_any(answer[0], axis=1, keep_dims=True))], axis=1)

        # Note we are happily allowing the model to place weights on "backwards" spans, and also giving
        # it points for predicting spans that start and end at different answer spans. It would be easy to
        # fix by masking out some of the `all_logit` matrix and specify a more accuracy correct_mask, but I
        # in general left it this way to be consistent with the independent bound models that do the same.
        # Some early tests found properly masking things to not make much difference (or even to hurt), but it
        # still could be an avenue for improvement

        log_correct = tf.reduce_logsumexp(all_logits + VERY_NEGATIVE_NUMBER * (1 - tf.cast(correct_mask, tf.float32)), axis=1)
        loss = tf.reduce_mean(-(log_correct - log_norms))
        probs = tf.nn.softmax(all_logits)
        tf.add_to_collection(tf.GraphKeys.LOSSES, loss)
        return ConfidencePrediction(probs[:, :-1], masked_start_logits, masked_end_logits,
                                    probs[:, -1], none_logit, context_mask) 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) class IndependentBoundsGrouped(SpanFromBoundsPredictor):
    """ The shared norm loss, where the normalizer is shared between paragraph with the same group id """

    def __init__(self, aggregate="sum"):
        self.aggregate = aggregate

    def predict(self, answer, start_logits, end_logits, mask) -> Prediction:
        masked_start_logits = exp_mask(start_logits, mask)
        masked_end_logits = exp_mask(end_logits, mask)

        if len(answer) == 3:
            group_ids = answer[2]
            # Turn the ids into segment ids using tf.unique
            _, group_segments = tf.unique(group_ids, out_idx=tf.int32)

            losses = []
            for answer_mask, logits in zip(answer, [masked_start_logits, masked_end_logits]):
                group_norms = segment_logsumexp(logits, group_segments)
                if self.aggregate == "sum":
                    log_score = segment_logsumexp(logits + VERY_NEGATIVE_NUMBER * (1 - tf.cast(answer_mask, tf.float32)),
                                                  group_segments)
                else:
                    raise ValueError()
                losses.append(tf.reduce_mean(-(log_score - group_norms)))
            loss = tf.add_n(losses)
        else:
            raise NotImplemented()
        tf.add_to_collection(tf.GraphKeys.LOSSES, loss)
        return BoundaryPrediction(tf.nn.softmax(masked_start_logits),
                                  tf.nn.softmax(masked_end_logits),
                                  masked_start_logits, masked_end_logits, mask) class IndependentBoundsSigmoidLoss(SpanFromBoundsPredictor):
    """ Independent sigmoid loss for each start/end span """

    def __init__(self, aggregate="sum"):
        self.aggregate = aggregate

    def predict(self, answer, start_logits, end_logits, mask) -> Prediction:
        masked_start_logits = exp_mask(start_logits, mask)
        masked_end_logits = exp_mask(end_logits, mask)

        if len(answer) == 1:
            raise NotImplementedError()
        elif len(answer) == 2 and all(x.dtype == tf.bool for x in answer):
            losses = []
            for answer_mask, logits in zip(answer, [masked_start_logits, masked_end_logits]):
                answer_mask = tf.cast(answer_mask, tf.float32)
                loss = tf.nn.sigmoid_cross_entropy_with_logits(
                    labels=tf.cast(answer_mask, tf.float32),
                    logits=logits
                )
                losses.append(loss)
            loss = tf.add_n(losses)
        else:
            raise NotImplemented()
        tf.add_to_collection(tf.GraphKeys.LOSSES, tf.reduce_mean(loss, name="sigmoid-loss"))
        return BoundaryPrediction(tf.nn.sigmoid(masked_start_logits),
                                  tf.nn.sigmoid(masked_end_logits),
                                  masked_start_logits, masked_end_logits, mask) class WithIndicators(TextPreprocessor):
    """
    Adds a document or group start token before the text, and a paragraph token between each
    between in each paragraph.
    """

    PARAGRAPH_TOKEN = "%%PARAGRAPH%%"
    DOCUMENT_START_TOKEN = "%%DOCUMENT%%"
    PARAGRAPH_GROUP = "%%PARAGRAPH_GROUP%%"

    def __init__(self, remove_cross_answer: bool=True, para_tokens: bool=True, doc_start_token: bool=True):
        self.remove_cross_answer = remove_cross_answer
        self.doc_start_token = doc_start_token
        self.para_tokens = para_tokens

    def special_tokens(self) -> List[str]:
        tokens = [self.PARAGRAPH_GROUP]
        if self.doc_start_token:
            tokens.append(self.DOCUMENT_START_TOKEN)
        if self.para_tokens:
            tokens.append(self.PARAGRAPH_TOKEN)
        return tokens

    def encode_paragraph(self, question: List[str], paragraphs: List[List[str]], is_first, answer_spans: np.ndarray, inver=None):
        out = []

        offset = 0
        if self.doc_start_token and is_first:
            out.append(self.DOCUMENT_START_TOKEN)
        else:
            out.append(self.PARAGRAPH_GROUP)

        if inver is not None:
            inv_out = [np.zeros((1, 2), dtype=np.int32)]
        else:
            inv_out = None

        offset += 1
        spans = answer_spans + offset

        out += paragraphs[0]
        offset += len(paragraphs[0])
        on_ix = len(paragraphs[0])
        if inv_out is not None:
            inv_out.append(inver[:len(paragraphs[0])])

        for sent in paragraphs[1:]:
            if self.remove_cross_answer:
                remove = np.logical_and(spans[:, 0] < offset, spans[:, 1] >= offset)
                spans = spans[np.logical_not(remove)]

            if self.para_tokens:
                spans[spans[:, 0] >= offset, 0] += 1
                spans[spans[:, 1] >= offset, 1] += 1

                out.append(self.PARAGRAPH_TOKEN)
                if inv_out is not None:
                    if len(inv_out) == 0 or len(inv_out[-1]) == 0:
                        inv_out.append(np.zeros((1, 2), dtype=np.int32))
                    else:
                        inv_out.append(np.full((1, 2), inv_out[-1][-1][1], dtype=np.int32))
                offset += 1

            out += sent
            offset += len(sent)
            if inv_out is not None:
                inv_out.append(inver[on_ix:on_ix+len(sent)])
            on_ix += len(sent)

        return out, spans, None if inv_out is None else np.concatenate(inv_out)

    def __setstate__(self, state):
        if "state" in state:
            state["state"]["doc_start_token"] = True
            state["state"]["para_tokens"] = True
        else:
            if "doc_start_token" not in state:
                state["doc_start_token"] = True
            if "para_tokens" not in state:
                state["para_tokens"] = True
        super().__setstate__(state) 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 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 class TriviaQaWebDataset(TriviaQaSpanCorpus):
    def __init__(self):
        super().__init__("web") class ExtractSingleParagraph(Preprocessor):
    """ Grab a single paragraph for each (document, question) pair, builds a list of
     (filtered) `DocumentParagraphQuestion` objects """

    def __init__(self, splitter: DocumentSplitter,
                 para_filter: Optional[ParagraphFilter],
                 text_preprocess: Optional[TextPreprocessor],
                 intern,
                 require_answer=True):
        self.splitter = splitter
        self.para_filter = para_filter
        self.text_preprocess = text_preprocess
        self.intern = intern
        self.require_answer = require_answer

    def preprocess(self, questions: List[TriviaQaQuestion], evidence) -> FilteredData:
        splitter = self.splitter
        paragraph_filter = self.para_filter
        output = []
        read_only = splitter.reads_first_n
        for q in questions:
            for doc in q.all_docs:
                text = evidence.get_document(doc.doc_id, n_tokens=read_only)
                if text is None:
                    raise ValueError(doc.doc_id, doc.doc_id)

                paragraphs = splitter.split_annotated(text, doc.answer_spans)
                if paragraph_filter is not None:
                    paragraphs = paragraph_filter.prune(q.question, paragraphs)
                if self.require_answer:
                    paragraphs = [x for x in paragraphs if len(x.answer_spans) > 0]
                if len(paragraphs) == 0:
                    continue
                paragraph = paragraphs[0]
                if self.text_preprocess is not None:
                    ex = self.text_preprocess.encode_extracted_paragraph(q.question, paragraph)
                    if not self.require_answer or len(ex.answer_spans) > 0:
                        output.append(DocumentParagraphQuestion(q.question_id, doc.doc_id, (paragraph.start, paragraph.end),
                                                                q.question, ex.text,
                                                                TokenSpans(q.answer.all_answers, ex.answer_spans), 1))
                else:
                    output.append(DocumentParagraphQuestion(q.question_id, doc.doc_id, (paragraph.start, paragraph.end),
                                                            q.question, flatten_iterable(paragraph.text),
                                                            TokenSpans(q.answer.all_answers, paragraph.answer_spans), 1))
        return FilteredData(output, sum(len(x.all_docs) for x in questions))

    def finalize_chunk(self, x: FilteredData):
        if self.intern:
            question_map = {}
            for q in x.data:
                q.question_id = sys.intern(q.question_id)
                if q.question_id in question_map:
                    q.question = question_map[q.question_id]
                else:
                    q.question = tuple(sys.intern(w) for w in q.question)
                    question_map[q.question_id] = q.question
                q.doc_id = sys.intern(q.doc_id)
                q.context = [sys.intern(w) for w in q.context]

    def __setstate__(self, state):
        if "state" in state:
            if "require_answer" not in state["state"]:
                state["state"]["require_answer"] = True
        super().__setstate__(state) class ExtractMultiParagraphs(Preprocessor):
    """
    Grab multiple paragraphs per (document, question) pair, return a list of (filtered) `MultiParagraphQuestion`
    """

    def __init__(self, splitter: DocumentSplitter, ranker: Optional[ParagraphFilter],
                 text_process: Optional[TextPreprocessor], intern: bool=False, require_an_answer=True):
        self.intern = intern
        self.splitter = splitter
        self.ranker = ranker
        self.text_process = text_process
        self.require_an_answer = require_an_answer

    def preprocess(self, questions: List[TriviaQaQuestion], evidence):
        true_len = 0
        splitter = self.splitter
        para_filter = self.ranker

        with_paragraphs = []
        for q in questions:
            true_len += len(q.all_docs)
            for doc in q.all_docs:
                if self.require_an_answer and len(doc.answer_spans) == 0:
                    continue
                text = evidence.get_document(doc.doc_id, splitter.reads_first_n)
                if text is None:
                    raise ValueError("No evidence text found document: " + doc.doc_id)
                if doc.answer_spans is not None:
                    paras = splitter.split_annotated(text, doc.answer_spans)
                else:
                    # this is kind of a hack to make the rest of the pipeline work, only
                    # needed for test cases
                    paras = splitter.split_annotated(text, np.zeros((0, 2), dtype=np.int32))

                if para_filter is not None:
                    paras = para_filter.prune(q.question, paras)

                if len(paras) == 0:
                    continue
                if self.require_an_answer:
                    if all(len(x.answer_spans) == 0 for x in paras):
                        continue
                if self.text_process is not None:
                    prepped = [self.text_process.encode_extracted_paragraph(q.question, p) for p in paras]
                    if self.require_an_answer:
                        if all(len(x.answer_spans) == 0 for x in prepped):
                            continue
                    doc_paras = []
                    for i, (preprocessed, para) in enumerate(zip(prepped, paras)):
                        doc_paras.append(DocumentParagraph(doc.doc_id, para.start, para.end,
                                                           i, preprocessed.answer_spans, preprocessed.text))

                else:
                    doc_paras = [DocumentParagraph(doc.doc_id, x.start, x.end,
                                                   i, x.answer_spans, flatten_iterable(x.text))
                                 for i, x in enumerate(paras)]
                with_paragraphs.append(MultiParagraphQuestion(q.question_id, q.question,
                                                              None if q.answer is None else q.answer.all_answers,
                                                              doc_paras))

        return FilteredData(with_paragraphs, true_len)

    def finalize_chunk(self, q: FilteredData):
        if self.intern:
            intern_mutli_question(q.data)
import argparse
from datetime import datetime
from typing import Optional
from tensorflow.contrib.keras.python.keras.initializers import TruncatedNormal
from docqa import model_dir
from docqa import trainer
from docqa.data_processing.document_splitter import MergeParagraphs, TopTfIdf
from docqa.data_processing.multi_paragraph_qa import StratifyParagraphsBuilder, \
    StratifyParagraphSetsBuilder, RandomParagraphSetDatasetBuilder
from docqa.data_processing.preprocessed_corpus import PreprocessedData
from docqa.data_processing.qa_training_data import ParagraphAndQuestionsBuilder, ContextLenKey, ContextLenBucketedKey
from docqa.data_processing.text_utils import NltkPlusStopWords
from docqa.dataset import ClusteredBatcher
from docqa.doc_qa_models import Attention
from docqa.encoder import DocumentAndQuestionEncoder, DenseMultiSpanAnswerEncoder, GroupedSpanAnswerEncoder
from docqa.evaluator import LossEvaluator, MultiParagraphSpanEvaluator, SpanEvaluator
from docqa.nn.attention import BiAttention, AttentionEncoder, StaticAttentionSelf
from docqa.nn.embedder import FixedWordEmbedder, CharWordEmbedder, LearnedCharEmbedder
from docqa.nn.layers import NullBiMapper, SequenceMapperSeq, Conv1d, FullyConnected, \
    ChainBiMapper, ConcatWithProduct, ResidualLayer, VariationalDropoutLayer, MaxPool
from docqa.nn.recurrent_layers import CudnnGru
from docqa.nn.similarity_layers import TriLinear
from docqa.nn.span_prediction import ConfidencePredictor, BoundsPredictor, IndependentBoundsGrouped, \
    IndependentBoundsSigmoidLoss
from docqa.text_preprocessor import WithIndicators, TextPreprocessor
from docqa.trainer import SerializableOptimizer, TrainParams
from docqa.triviaqa.build_span_corpus import TriviaQaWebDataset
from docqa.triviaqa.training_data import ExtractSingleParagraph, ExtractMultiParagraphs




def get_triviaqa_train_params(n_epochs, n_dev, n_train):
    return TrainParams(
        SerializableOptimizer("Adadelta", dict(learning_rate=1)),
        num_epochs=n_epochs, ema=0.9999, max_checkpoints_to_keep=2,
        async_encoding=10, log_period=30, eval_period=1800, save_period=1800,
        eval_samples=dict(dev=n_dev, train=n_train))


def get_model(char_th: int, dim: int, mode: str, preprocess: Optional[TextPreprocessor]):
    recurrent_layer = CudnnGru(dim, w_init=TruncatedNormal(stddev=0.05))

    if mode.startswith("shared-norm"):
        answer_encoder = GroupedSpanAnswerEncoder()
        predictor = BoundsPredictor(
            ChainBiMapper(
                first_layer=recurrent_layer,
                second_layer=recurrent_layer
            ),
            span_predictor=IndependentBoundsGrouped(aggregate="sum")
        )
    elif mode == "confidence":
        answer_encoder = DenseMultiSpanAnswerEncoder()
        predictor = ConfidencePredictor(
            ChainBiMapper(
                first_layer=recurrent_layer,
                second_layer=recurrent_layer,
            ),
            AttentionEncoder(),
            FullyConnected(80, activation="tanh"),
            aggregate="sum"
        )
    elif mode == "sigmoid":
        answer_encoder = DenseMultiSpanAnswerEncoder()
        predictor = BoundsPredictor(
            ChainBiMapper(
                first_layer=recurrent_layer,
                second_layer=recurrent_layer
            ),
            span_predictor=IndependentBoundsSigmoidLoss()
        )
    elif mode == "paragraph" or mode == "merge":
        answer_encoder = DenseMultiSpanAnswerEncoder()
        predictor = BoundsPredictor(ChainBiMapper(
            first_layer=recurrent_layer,
            second_layer=recurrent_layer
        ))
    else:
        raise NotImplementedError(mode)

    return Attention(
        encoder=DocumentAndQuestionEncoder(answer_encoder),
        word_embed=FixedWordEmbedder(vec_name="glove.840B.300d", word_vec_init_scale=0, learn_unk=False, cpu=True),
        char_embed=CharWordEmbedder(
            LearnedCharEmbedder(word_size_th=14, char_th=char_th, char_dim=20, init_scale=0.05, force_cpu=True),

            MaxPool(Conv1d(100, 5, 0.8)),