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 ShallowOpenWebRanker(ParagraphFilter):
    # Hard coded weight learned from a logistic regression classifier
    TFIDF_W = 5.13365065
    LOG_WORD_START_W = 0.46022765
    FIRST_W = -0.08611607
    LOWER_WORD_W = 0.0499123
    WORD_W = -0.15537181

    def __init__(self, n_to_select):
        self.n_to_select = n_to_select
        self._stop = NltkPlusStopWords(True).words
        self._tfidf = TfidfVectorizer(strip_accents="unicode", stop_words=self._stop)

    def get_features(self, question: List[str], paragraphs: List[List[ExtractedParagraphWithAnswers]]):
        scores = self.score_paragraphs(question, flatten_iterable(paragraphs))
        # return scores
        return np.expand_dims(scores, 1)

    def get_feature_names(self):
        return ["Score"]

    def score_paragraphs(self, question, paragraphs: List[ExtractedParagraphWithAnswers]):
        tfidf = self._tfidf
        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 []

        q_words = {x for x in question if x.lower() not in self._stop}
        q_words_lower = {x.lower() for x in q_words}
        word_matches_features = np.zeros((len(paragraphs), 2))
        for para_ix, para in enumerate(paragraphs):
            found = set()
            found_lower = set()
            for sent in para.text:
                for word in sent:
                    if word in q_words:
                        found.add(word)
                    elif word.lower() in q_words_lower:
                        found_lower.add(word.lower())
            word_matches_features[para_ix, 0] = len(found)
            word_matches_features[para_ix, 1] = len(found_lower)

        tfidf = pairwise_distances(q_features, para_features, "cosine").ravel()
        starts = np.array([p.start for p in paragraphs])
        log_word_start = np.log(starts/400.0 + 1)
        first = starts == 0
        scores = tfidf * self.TFIDF_W + self.LOG_WORD_START_W * log_word_start + self.FIRST_W * first +\
                 self.LOWER_WORD_W * word_matches_features[:, 1] + self.WORD_W * word_matches_features[:, 0]
        return scores

    def prune(self, question, paragraphs: List[ExtractedParagraphWithAnswers]):
        scores = self.score_paragraphs(question, paragraphs)
        sorted_ix = np.argsort(scores)

        return [paragraphs[i] for i in sorted_ix[:self.n_to_select]]

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

    def __setstate__(self, state):
        return self.__init__(state['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 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 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 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) 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)),
            shared_parameters=True
        ),
        preprocess=preprocess,
        word_embed_layer=None,
        embed_mapper=SequenceMapperSeq(
            VariationalDropoutLayer(0.8),
            recurrent_layer,
            VariationalDropoutLayer(0.8),
        ),
        question_mapper=None,
        context_mapper=None,
        memory_builder=NullBiMapper(),
        attention=BiAttention(TriLinear(bias=True), True),
        match_encoder=SequenceMapperSeq(FullyConnected(dim * 2, activation="relu"),
                                        ResidualLayer(SequenceMapperSeq(
                                            VariationalDropoutLayer(0.8),
                                            recurrent_layer,
                                            VariationalDropoutLayer(0.8),
                                            StaticAttentionSelf(TriLinear(bias=True), ConcatWithProduct()),
                                            FullyConnected(dim * 2, activation="relu"),
                                        )),
                                        VariationalDropoutLayer(0.8)),
        predictor=predictor
    ) 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 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 TriviaQaOpenDataset(TriviaQaSpanCorpus):
    def __init__(self):
        super().__init__("web-open") class TriviaQaWikiDataset(TriviaQaSpanCorpus):
    def __init__(self):
        super().__init__("wiki") class ExtractMultiParagraphsPerQuestion(Preprocessor):
    """
    Get multiple paragraph per question, using all document. Returns a filtered list of
    `MultiParagraphQuestion`
    """

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

    def preprocess(self, questions: List[TriviaQaQuestion], evidence) -> object:
        splitter = self.splitter
        para_filter = self.ranker

        with_paragraphs = []
        for q in questions:
            paras = []
            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 doc.answer_spans is not None:
                    split = 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
                    split = splitter.split_annotated(text, np.zeros((0, 2), dtype=np.int32))
                paras.extend([DocParagraphWithAnswers(x.text, x.start, x.end, x.answer_spans, doc.doc_id)
                              for x in split])

            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_preprocess is not None:
                prepped = [self.text_preprocess.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(para.doc_id, para.start, para.end,
                                                       i, preprocessed.answer_spans, preprocessed.text))
                with_paragraphs.append(
                    MultiParagraphQuestion(q.question_id, q.question,
                                           None if q.answer is None else q.answer.all_answers,
                                           doc_paras))
            else:
                doc_paras = [DocumentParagraph(x.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, q.answer.all_answers, doc_paras))

        return FilteredData(with_paragraphs, len(questions))

    def finalize_chunk(self, q: FilteredData):
        if self.intern:
            intern_mutli_question(q.data)
import argparse
from datetime import datetime
from docqa import model_dir
from docqa import trainer
from docqa.data_processing.document_splitter import MergeParagraphs, ShallowOpenWebRanker
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 ContextLenBucketedKey
from docqa.dataset import ClusteredBatcher
from docqa.evaluator import LossEvaluator, MultiParagraphSpanEvaluator
from docqa.scripts.ablate_triviaqa import get_model
from docqa.text_preprocessor import WithIndicators
from docqa.trainer import SerializableOptimizer, TrainParams
from docqa.triviaqa.build_span_corpus import TriviaQaOpenDataset, TriviaQaWikiDataset
from docqa.triviaqa.training_data import ExtractMultiParagraphsPerQuestion



def main():
    parser = argparse.ArgumentParser(description='Train a model on TriviaQA wiki')
    parser.add_argument('mode', choices=["confidence", "merge", "shared-norm",
                                         "sigmoid", "paragraph"])
    # Note I haven't tested modes other than `shared-norm` on this corpus, so
    # some things might need adjusting
    parser.add_argument("name", help="Where to store the model")
    parser.add_argument("-t", "--n_tokens", default=400, type=int,
                        help="Paragraph size")
    parser.add_argument('-n', '--n_processes', type=int, default=2,
                        help="Number of processes (i.e., select which paragraphs to train on) "
                             "the data with"
                        )
    args = parser.parse_args()
    mode = args.mode

    out = args.name + "-" + datetime.now().strftime("%m%d-%H%M%S")

    model = get_model(100, 140, mode, WithIndicators())

    extract = ExtractMultiParagraphsPerQuestion(MergeParagraphs(args.n_tokens),
                                                ShallowOpenWebRanker(16),
                                                model.preprocessor, intern=True)

    eval = [LossEvaluator(), MultiParagraphSpanEvaluator(8, "triviaqa", mode != "merge", per_doc=False)]
    oversample = [1] * 2  # Sample the top two answer-containing paragraphs twice

    if mode == "paragraph":
        n_epochs = 120
        test = RandomParagraphSetDatasetBuilder(120, "flatten", True, oversample)
        train = StratifyParagraphsBuilder(ClusteredBatcher(60, ContextLenBucketedKey(3), True),
                                          oversample,  only_answers=True)
    elif mode == "confidence" or mode == "sigmoid":
        if mode == "sigmoid":
            n_epochs = 640
        else:
            n_epochs = 160
        test = RandomParagraphSetDatasetBuilder(120, "flatten", True, oversample)

        train = StratifyParagraphsBuilder(ClusteredBatcher(60, ContextLenBucketedKey(3), True), oversample)