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']) def top_disjoint_spans(span_scores, bound: int, n_spans: int, spans):
    """
    Given a n_token x n_tokens matrix of spans scores, return the top-n non-overlapping spans
    and their scores
    """

    sorted_ux = np.argsort(span_scores.ravel())[::-1]
    sorted_ux = np.stack([sorted_ux // len(span_scores), sorted_ux % len(span_scores)], axis=1)

    lens = sorted_ux[:, 1] - sorted_ux[:, 0]
    sorted_ux = sorted_ux[np.logical_and(lens >= 0, lens < bound)]

    cur_scores = np.zeros(n_spans, dtype=np.float32)
    cur_spans = np.zeros((n_spans, 2), dtype=np.int32)
    spans_found = 0

    for s, e in sorted_ux:
        # the span must either start after or end before each existing spans
        if np.all(np.logical_or(cur_spans[:spans_found, 0] > e, cur_spans[:spans_found, 1] < s)):
            if spans[s][0] < spans[e][1]:  # Don't select zero length spans
                cur_spans[spans_found, 0] = s
                cur_spans[spans_found, 1] = e
                cur_scores[spans_found] = span_scores[s, e]
                spans_found += 1
                if spans_found == n_spans:
                    break

    return cur_spans[:spans_found], cur_scores[:spans_found] class NltkAndPunctTokenizer(Configurable):

    @staticmethod
    def convert_to_spans(raw_text: str, sentences: List[List[str]]) -> List[List[Tuple[int, int]]]:
        """ Convert a tokenized version of `raw_text` into a series character spans referencing the `raw_text` """
        cur_idx = 0
        all_spans = []
        for sent in sentences:
            spans = []
            for token in sent:
                # (our) Tokenizer might transform double quotes, for this case search over several
                # possible encodings
                if double_quote_re.match(token):
                    span = double_quote_re.search(raw_text[cur_idx:])
                    tmp = cur_idx + span.start()
                    l = span.end() - span.start()
                else:
                    tmp = raw_text.find(token, cur_idx)
                    l = len(token)
                if tmp < cur_idx:
                    raise ValueError(token)
                cur_idx = tmp
                spans.append((cur_idx, cur_idx + l))
                cur_idx += l
            all_spans.append(spans)
        return all_spans

    def __init__(self):
        self.sent_tokenzier = nltk.load('tokenizers/punkt/english.pickle')
        self.word_tokenizer = nltk.TreebankWordTokenizer()

    def clean_text(self, word):
        # be consistent with quotes, and replace \u2014 and \u2212 which I have seen being mapped to UNK
        # by glove word vecs
        return word.replace("''", "\"").replace("``", "\"").replace("\u2212", "-").replace("\u2014", "\u2013")

    def tokenize_sentence(self, sent) -> List[str]:
        tokens = self.word_tokenizer.tokenize(sent)
        return [self.clean_text(token) for token in post_split_tokens(tokens)]

    def tokenize_paragraph(self, paragraph: str) -> List[List[str]]:
        return [self.tokenize_sentence(s) for s in self.sent_tokenzier.tokenize(paragraph)]

    def tokenize_paragraph_flat(self, paragraph: str) -> List[str]:
        return flatten_iterable(self.tokenize_paragraph(paragraph))

    def tokenize_with_inverse(self, paragraph: str, is_sentence: bool=False) -> ParagraphWithInverse:
        if is_sentence:
            para = [paragraph]
        else:
            para = self.sent_tokenzier.tokenize(paragraph)

        text = [self.word_tokenizer.tokenize(s) for s in para]
        for i, sent in enumerate(text):
            text[i] = post_split_tokens(sent)

        # recover (start, end) for each token relative to the raw `context`
        text_spans = self.convert_to_spans(paragraph, text)

        # Clean text, we do this at the end so `convert_to_spans` works as expected
        for i, sent in enumerate(text):
            text[i] = [self.clean_text(x) for x in sent]

        text_spans = flatten_iterable(text_spans)
        if len(text_spans) == 0:
            text_spans = np.zeros((0, 2), dtype=np.int32)
        else:
            text_spans = np.array(text_spans, dtype=np.int32)
        return ParagraphWithInverse(text, paragraph, text_spans) class Model(Configurable):
    """
    Our most general specification of a model/neural network, for our purposes a model
    is basically a pair of functions
    1) a way to map a (unspecified) kind of python object to numpy tensors and
    2) a tensorflow function that maps those kinds of tensors to a set of (also unspecified) output tensors

    For convenience, models maintain a of set of input placeholders that clients can make use of to
    feed the tensorflow function (or reference to construct their own tensor inputs).

    Models have two stages of initialization. First it needs
    to be initialized with the training data using `init` (typically this does things like deciding what
    words/chars to train embeddings for). This should only be done once for this object's lifetime.

    Afterwards use `set_inputs` to specify the input format, this does things like determine the batch size
    or the vocabulary that will be used

    After initialiation, `encode` will produce map of placeholder -> numpy array
    which can be used directly as a feed dict for the output of `get_predictions`

    For more advanced usage, `get_predictions_for` can be used with any tensors of the
    same shape/dtype as the input place holders. Clients should pass in a dict mapping
    the placeholders to the input tensors they want to use instead.

    `get_predictions_for` methods behave like any other tensorflow function, in that it will
    load/initialize/reuse variables depending on the current tensorflow scope and can add
    to tf.collections. Our trainer method makes use of some of these collections, including:
        tf.GraphKeys.LOSSES
        tf.GraphKeys.REGULARIZATION_LOSSES
        tf.GraphKeys.SUMMARIES
        tf.GraphKeys.SAVEABLE_OBJECTS
        tf.GraphKeys.TRAINABLE_VARIABLES
        "monitor/*" collections, which will be summed, and the EMA result logged to tensorboard
    """

    @property
    def name(self):
        return self.__class__.__name__

    def init(self, train_data, resource_loader: ResourceLoader):
        raise NotImplementedError()

    def set_inputs(self, datasets: List[Dataset], resource_loader: ResourceLoader) -> List[Tensor]:
        raise NotImplementedError()

    def get_prediction(self) -> Prediction:
        return self.get_predictions_for({x: x for x in self.get_placeholders()})

    def get_placeholders(self) -> List[Tensor]:
        raise NotImplementedError()

    def get_predictions_for(self, input_tensors: Dict[Tensor, Tensor]) -> Prediction:
        raise NotImplementedError()

    def encode(self, examples, is_train: bool) -> Dict[Tensor, object]:
        raise NotImplementedError() class Prediction(object):
    """ Prediction from a model, subclasses should provide access to a tensor
     representations of the model's output """
    pass class ModelDir(object):
    """ Wrapper for accessing a folder we are storing a model in"""

    def __init__(self, name: str):
        if isabs(name):
            print("WARNING!!!, using an absolute paths for models name can break restoring "
                  "the model in different directories after being checkpointed")
            # why is this even a thing?
        self.dir = name

    def get_model(self) -> Model:
        with open(join(self.dir, "model.pkl"), "rb") as f:
            return pickle.load(f)

    def get_eval_dir(self):
        answer_dir = join(self.dir, "answers")
        if not exists(answer_dir):
            os.mkdir(answer_dir)
        return answer_dir

    def get_last_train_params(self):
        last_train_file = None
        last_train_step = -1
        for file in os.listdir(self.dir):
            if file.startswith("train_from_") and file.endswith("pkl"):
                step = int(file[11:file.rfind(".pkl")])
                if step > last_train_step:
                    last_train_step = step
                    last_train_file = join(self.dir, file)

        print("Resuming using the parameters stored in: " + last_train_file)
        with open(last_train_file, "rb") as f:
            return pickle.load(f)

    def get_latest_checkpoint(self):
        return tf.train.latest_checkpoint(self.save_dir)

    def get_checkpoint(self, step):
        # I cant find much formal documentation on how to do this, but this seems to work
        return join(self.save_dir, "checkpoint-%d-%d" % (step, step))

    def get_best_weights(self):
        if exists(self.best_weight_dir):
            return tf.train.latest_checkpoint(self.best_weight_dir)
        return None

    def restore_checkpoint(self, sess, var_list=None, load_ema=True):
        """
        Restores either the best weights or the most recent checkpoint, assuming the correct
        variables have already been added to the tf default graph e.g., .get_prediction()
        has been called the model stored in `self`.
        Automatically detects if EMA weights exists, and if they do loads them instead
        """
        checkpoint = self.get_best_weights()
        if checkpoint is None:
            print("Loading most recent checkpoint")
            checkpoint = self.get_latest_checkpoint()
        else:
            print("Loading best weights")

        if load_ema:
            if var_list is None:
                # Same default used by `Saver`
                var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) + \
                           tf.get_collection(tf.GraphKeys.SAVEABLE_OBJECTS)

            # Automatically check if there are EMA variables, if so use those
            reader = tf.train.NewCheckpointReader(checkpoint)
            ema = tf.train.ExponentialMovingAverage(0)
            ema_names = {ema.average_name(x): x for x in var_list
                         if reader.has_tensor(ema.average_name(x))}
            if len(ema_names) > 0:
                print("Found EMA weights, loading them")
                ema_vars = set(x for x in ema_names.values())
                var_list = {v.op.name: v for v in var_list if v not in ema_vars}
                var_list.update(ema_names)

        saver = tf.train.Saver(var_list)
        saver.restore(sess, checkpoint)


    @property
    def save_dir(self):
        # Stores training checkpoint
        return join(self.dir, "save")

    @property
    def best_weight_dir(self):
        # Stores training checkpoint
        return join(self.dir, "best-weights")

    @property
    def log_dir(self):
        return join(self.dir, "log") class BoundaryPrediction(Prediction):
    """ Individual logits for the span start/end """

    def __init__(self, start_prob, end_prob,
                 start_logits, end_logits, mask):
        self.start_probs = start_prob
        self.end_probs = end_prob
        self.start_logits = start_logits
        self.end_logits = end_logits
        self.mask = mask
        self._bound_predictions = {}

    def get_best_span(self, bound: int):
        if bound in self._bound_predictions:
            return self._bound_predictions[bound]
        else:
            pred = best_span_from_bounds(self.start_logits, self.end_logits, bound)
            self._bound_predictions[bound] = pred
            return pred

    def get_span_scores(self):
        return tf.exp(tf.expand_dims(self.start_logits, 2) + tf.expand_dims(self.end_logits, 1))

    def get_mean_logit(self):
        logits = (self.start_logits + self.end_logits) / 2.0
        bol_mask = tf.sequence_mask(self.mask, tf.shape(self.start_logits)[1])
        bol_mask = tf.cast(bol_mask, tf.float32)
        return tf.reduce_sum(logits*bol_mask, axis=[1]) / tf.reduce_sum(bol_mask, axis=[1]) class WebParagraph(ParagraphWithInverse):
    """ Paragraph the includes a url and source """

    def __init__(self, text: List[List[str]], original_text: str, token_spans: np.ndarray,
                 paragraph_num: int, start: int, end: int, source_name, source_url):
        super().__init__(text, original_text, token_spans)
        self.source_name = source_name
        self.paragraph_num = paragraph_num
        self.source_url = source_url
        self.start = start
        self.end = end class QaSystem(object):
    """
    End-to-end QA system, uses web-requests to get relevant documents and a model
    to score candidate answer spans.
    """
    # TODO fix logging level

    _split_regex = re.compile("\s*\n\s*")  # split includes whitespace to avoid empty paragraphs

    def __init__(self,
                 wiki_cache: str,
                 paragraph_splitter: DocumentSplitter,
                 paragraph_selector: ParagraphFilter,
                 vocab: Union[str, None, Set[str]],
                 model: Union[ParagraphQuestionModel, ModelDir],
                 loader: ResourceLoader=ResourceLoader(),
                 bing_api_key=None,
                 bing_version="v5.0",
                 tagme_api_key=None,
                 blacklist_trivia_sites: bool=False,
                 n_dl_threads: int=5,
                 span_bound: int=8,
                 tagme_threshold: Optional[float]=0.2,
                 download_timeout: int=None,
                 n_web_docs=10,
                 loop=None):
        self.log = logging.getLogger('qa_system')
        self.tagme_threshold = tagme_threshold
        self.n_web_docs = n_web_docs
        self.blacklist_trivia_sites = blacklist_trivia_sites
        self.tagme_api_key = tagme_api_key

        self.client_sess = ClientSession(loop=loop)

        if bing_api_key is not None:
            if bing_version is None:
                raise ValueError("Must specify a Bing version if using a bing_api key")
            self.searcher = AsyncWebSearcher(bing_api_key, bing_version, loop=loop)
            self.text_extractor = AsyncBoilerpipeCliExtractor(n_dl_threads, download_timeout)
        else:
            self.text_extractor = None
            self.searcher = None

        if self.tagme_threshold is not None:
            self.wiki_corpus = WikiCorpus(wiki_cache, keep_inverse_mapping=True, loop=loop)
        else:
            self.wiki_corpus = None

        self.paragraph_splitter = paragraph_splitter
        self.paragraph_selector = paragraph_selector
        self.model_dir = model

        voc = None
        if vocab is not None:
            if isinstance(vocab, str):
                voc = set()
                with open(vocab, "r") as f:
                    for line in f:
                        voc.add(line.strip())
            else:
                voc = vocab
            self.log.info("Using preset vocab of size %d", len(voc))

        self.log.info("Setting up model...")
        if isinstance(model, ModelDir):
            self.model = model.get_model()
        else:
            self.model = model

        self.model.set_input_spec(ParagraphAndQuestionSpec(None), voc, loader)

        self.sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
        with self.sess.as_default():
            pred = self.model.get_prediction()

        if isinstance(model, ModelDir):
            model.restore_checkpoint(self.sess)

        self.span_scores = pred.get_span_scores()
        self.span, self.score = pred.get_best_span(span_bound)
        self.tokenizer = NltkAndPunctTokenizer()
        self.sess.graph.finalize()

    def _preprocess(self, paragraphs: List[WebParagraph]) -> List[WebParagraph]:
        if self.model.preprocessor is not None:
            prepped = []
            for para in paragraphs:
                if hasattr(para, "spans"):
                    spans = para.spans
                else:
                    spans = None

                text, _, inv = self.model.preprocessor.encode_paragraph([], para.text, para.start == 0,
                                                                          np.zeros((0, 2), dtype=np.int32), spans)
                prepped.append(WebParagraph(
                    [text], para.original_text, inv, para.paragraph_num,
                    para.start, para.end,
                    para.source_name, para.source_url
                ))
            return prepped
        else:
            return paragraphs

    async def answer_question_spans(self, question: str) -> Tuple[np.ndarray, np.ndarray, List[WebParagraph]]:
        """
        Answer a question using web search, return the top spans and confidence scores for each paragraph
        """

        paragraphs = await self.get_question_context(question)
        question = self.tokenizer.tokenize_paragraph_flat(question)
        paragraphs = self._preprocess(paragraphs)
        t0 = time.perf_counter()
        qa_pairs = [ParagraphAndQuestion(c.get_context(), question, None, "") for c in paragraphs]
        encoded = self.model.encode(qa_pairs, False)
        spans, scores = self.sess.run([self.span, self.score], encoded)
        self.log.info("Computing answer spans took %.5f seconds" % (time.perf_counter() - t0))
        return spans, scores, paragraphs

    async def answer_question(self, question: str) -> Tuple[np.ndarray, List[WebParagraph]]:
        """
        Answer a question using web search, return the paragraphs and per-span confidence scores
        in the form of a (batch, max_num_context_tokens, max_num_context_tokens) array
        """

        self.log.info("Answering question \"%s\" with web search" % question)
        context = await self.get_question_context(question)
        question = self.tokenizer.tokenize_paragraph_flat(question)
        t0 = time.perf_counter()
        out = self._get_span_scores(question, context)
        self.log.info("Computing answer spans took %.5f seconds" % (time.perf_counter() - t0))
        return out

    def answer_with_doc(self, question: str, doc: str) -> Tuple[np.ndarray, List[WebParagraph]]:
        """ Answer a question using the given text as a document """

        self.log.info("Answering question \"%s\" with a given document" % question)
        # Tokenize
        question = self.tokenizer.tokenize_paragraph_flat(question)
        context = [self.tokenizer.tokenize_with_inverse(x, False) for x in self._split_regex.split(doc)]

        # Split into super-paragraphs
        context = self._split_document(context, "User", None)

        # Select top paragraphs
        context = self.paragraph_selector.prune(question, context)
        if len(context) == 0:
            raise ValueError("Unable to process documents")

        # Select the top answer span
        t0 = time.perf_counter()
        span_scores = self._get_span_scores(question, context)
        self.log.info("Computing answer spans took %.5f seconds" % (time.perf_counter() - t0))
        return span_scores

    def _get_span_scores(self, question: List[str], paragraphs: List[WebParagraph]):
        paragraphs = self._preprocess(paragraphs)
        qa_pairs = [ParagraphAndQuestion(c.get_context(), question, None, "") for c in paragraphs]
        encoded = self.model.encode(qa_pairs, False)
        return self.sess.run(self.span_scores, encoded), paragraphs

    def _split_document(self, para: List[ParagraphWithInverse], source_name: str,
                        source_url: Optional[str]):
        tokenized_paragraphs = []
        on_token = 0
        for i, para in enumerate(self.paragraph_splitter.split_inverse(para)):
            n_tokens = para.n_tokens
            tokenized_paragraphs.append(WebParagraph(
                para.text, para.original_text, para.spans, i + 1,
                on_token, on_token + n_tokens,
                source_name, source_url
            ))
            on_token += n_tokens
        return tokenized_paragraphs

    async def _tagme(self, question: str):
        payload = {"text": question,
                   "long_text": 3,
                   "lang": "en",
                   "gcube-token": self.tagme_api_key}
        async with self.client_sess.get(url=TAGME_API, params=payload) as resp:
            data = await resp.json()
        return [ann_json for ann_json in data["annotations"] if "title" in ann_json]

    async def get_question_context(self, question: str) -> List[WebParagraph]:
        """
        Find a set of paragraphs from the web that are relevant to the given question
        """

        tokenized_paragraphs = []
        if self.tagme_threshold is not None:
            self.log.info("Query tagme for %s", question)
            tags = await self._tagme(question)
            t0 = time.perf_counter()
            found = set()
            for tag in tags:
                if tag["rho"] >= self.tagme_threshold:
                    title = tag["title"]
                    if title in found:
                        continue
                    found.add(title)
                    doc = await self.wiki_corpus.get_wiki_article(title)
                    tokenized_paragraphs += self._split_document(doc.paragraphs,
                                                                 "Wikipedia: " + doc.title, doc.url)
            if len(tokenized_paragraphs) > 0:
                self.log.info("Getting wiki docs took %.5f seconds" % (time.perf_counter() - t0))

        if self.n_web_docs > 0:
            t0 = time.perf_counter()
            self.log.info("Running bing search for %s", question)
            search_results = await self.searcher.run_search(question, self.n_web_docs)
            t1 = time.perf_counter()
            self.log.info("Completed bing search, took %.5f seconds" % (t1 - t0))
            t0 = t1
            url_to_result = {x["url"]: x for x in search_results}
            self.log.info("Extracting text for %d results", len(search_results))
            text_docs = await self.text_extractor.get_text([x["url"] for x in search_results])

            for doc in text_docs:
                if len(doc.text) == 0:
                    continue
                search_r = url_to_result[doc.url]
                if self.blacklist_trivia_sites:
                    lower = search_r["displayUrl"].lower()
                    if 'quiz' in lower or 'trivia' in lower or 'answer' in lower:
                        # heuristic to ignore trivia sites, recommend by Mandar
                        self.log.debug("Skipping trivia site: " + lower)
                        continue

                paras_text = self._split_regex.split(doc.text.strip())

                paras_tokenized = [self.tokenizer.tokenize_with_inverse(x) for x in paras_text]

                tokenized_paragraphs += self._split_document(paras_tokenized, search_r["displayUrl"], doc.url)

            self.log.info("Completed extracting text, took %.5f seconds." % (time.perf_counter() - t0))

        self.log.info("Have %d paragraphs", len(tokenized_paragraphs))
        if len(tokenized_paragraphs) == 0:
            return []
        question = self.tokenizer.tokenize_sentence(question)
        return self.paragraph_selector.prune(question, tokenized_paragraphs)

    def close(self):
        if self.wiki_corpus is not None:
            self.wiki_corpus.close()
        if self.searcher is not None:
            self.searcher.close()
        self.sess.close()
        self.client_sess.close() 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 ResourceLoader(object):
    """
    Abstraction for models the need access to external resources to setup, currently just
    for word-vectors.
    """

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

    def load_word_vec(self, vec_name, voc=None):
        return self.load_vec_fn(vec_name, voc) class LoadFromPath(object):
    def __init__(self, path):
        self.path = path

    def load_word_vec(self, vec_name, voc=None):
        return load_word_vectors(join(self.path, vec_name), voc, True)
import argparse
import logging
import sys
import time
import ujson
import numpy as np
import tensorflow as tf
from os import environ
from typing import List
from sanic import Sanic, response
from sanic.config import LOGGING
from sanic.exceptions import ServerError
from sanic.response import json
from docqa.data_processing.document_splitter import MergeParagraphs, ShallowOpenWebRanker
from docqa.data_processing.span_data import top_disjoint_spans
from docqa.data_processing.text_utils import NltkAndPunctTokenizer
from docqa.model import Model, Prediction
from docqa.model_dir import ModelDir
from docqa.nn.span_prediction import BoundaryPrediction
from docqa.server.qa_system import WebParagraph, QaSystem
from docqa.text_preprocessor import WithIndicators
from docqa.utils import ResourceLoader, LoadFromPath



"""
Server for the demo. The server uses the async/await framework so the API calls and 
web downloads can be done asynchronously, allowing the server to answer other queries in the
meantime. 
"""


class AnswerSpan(object):
    def __init__(self, conf: float, start: int, end: int):
        self.conf = conf
        self.start = start
        self.end = end

    def to_json(self):
        # make sure we are using python types
        return dict(conf=float(self.conf), start=int(self.start), end=int(self.end))


class WebParagraphWithSelectedAnswers(object):
    def __init__(self, url: str, source_name: str, original_text: str, answers: List[AnswerSpan]):
        self.source_url = url
        self.source_name = source_name
        self.original_text = original_text
        self.answers = answers

    def to_json(self):
        return dict(source_url=self.source_url, source_name=self.source_name,
                    text=self.original_text, answers=[x.to_json() for x in self.answers])


def select_answers(paras: List[WebParagraph], span_scores, bound, n_spans) -> List[WebParagraphWithSelectedAnswers]:
    """
    Selects the top `n_spans` non-overlapping spans of at most length `bound` from each paragraph,
    returns the resulting paragraphs sorted by most confidence answer
    """
    out = []
    for para, score in zip(paras, span_scores):
        n_tokens = len(para.spans)
        # Score can contain other stuff due to padding
        score = score[:n_tokens, :n_tokens]

        top_n, top_n_scores = top_disjoint_spans(score, bound, n_spans, para.spans)
        answers = []
        for score, (s, e) in zip(top_n_scores, top_n):
            s = para.spans[s][0]
            e = para.spans[e][1]
            answers.append(AnswerSpan(score, s, e))
        out.append(WebParagraphWithSelectedAnswers(para.source_url, para.source_name, para.original_text,
                                                   answers))
    out.sort(key=lambda x:-x.answers[0].conf)
    return out

ipso = """
Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore 
magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea 
commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat 
nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
"""


class DummyQa():
    async def answer_question(self, question: str):
        return self.get_random_answer()

    def get_random_answer(self):
        time.sleep(1)

        para = NltkAndPunctTokenizer().tokenize_with_inverse(ipso)