CORPUS_DIR = join(dirname(dirname(__file__)), "data") TRIVIA_QA = join(expanduser("~"), "data", "triviaqa") TRIVIA_QA_UNFILTERED = join(expanduser("~"), "data", "triviaqa-unfiltered") class Configurable(object):
    """
    Configurable classes have names, versions, and a set of parameters that are either "simple" aka JSON serializable
    types or other Configurable objects. Configurable objects should also be serializable via pickle.
    Configurable classes are defined mainly to give us a human-readable way of reading of the `parameters`
    set for different objects and to attach version numbers to them.

    By default we follow the format sklearn uses for its `BaseEstimator` class, where parameters are automatically
    derived based on the constructor parameters.
    """

    @classmethod
    def _get_param_names(cls):
        # fetch the constructor or the original constructor before
        init = cls.__init__
        if init is object.__init__:
            # No explicit constructor to introspect
            return []

        init_signature = signature(init)
        parameters = [p for p in init_signature.parameters.values()
                      if p.name != 'self']
        if any(p.kind == p.VAR_POSITIONAL for p in parameters):
            raise RuntimeError()
        return sorted([p.name for p in parameters])

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

    @property
    def version(self):
        return 0

    def get_params(self):
        out = {}
        for key in self._get_param_names():
            v = getattr(self, key, None)
            if isinstance(v, Configurable):
                out[key] = v.get_config()
            elif hasattr(v, "get_config"):  # for keras objects
                out[key] = {"name": v.__class__.__name__, "config": v.get_config()}
            else:
                out[key] = v
        return out

    def get_config(self) -> Configuration:
        params = {k: describe(v) for k,v in self.get_params().items()}
        return Configuration(self.name, self.version, params)

    def __getstate__(self):
        state = dict(self.__dict__)
        if "version" in state:
            if state["version"] != self.version:
                raise RuntimeError()
        else:
            state["version"] = self.version
        return state

    def __setstate__(self, state):
        if "version" not in state:
            raise RuntimeError("Version should be in state (%s)" % self.__class__.__name__)
        if state["version"] != self.version:
            warn(("%s loaded with version %s, but class " +
                 "version is %s") % (self.__class__.__name__, state["version"], self.version))

        if "state" in state:
            self.__dict__ = state["state"]
        else:
            del state["version"]
            self.__dict__ = state class 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) def compute_answer_spans_par(questions: List[TriviaQaQuestion], corpus,
                             tokenizer, detector, n_processes: int):
    if n_processes == 1:
        word_tokenize = tokenizer.tokenize_paragraph_flat
        compute_answer_spans(questions, corpus, word_tokenize, detector)
        return questions
    from multiprocessing import Pool
    with Pool(n_processes) as p:
        chunks = split(questions, n_processes)
        questions = flatten_iterable(p.starmap(_compute_answer_spans_chunk,
                                               [[c, corpus, tokenizer, detector] for c in chunks]))
        return questions class FastNormalizedAnswerDetector(object):
    """ almost twice as fast and very,very close to NormalizedAnswerDetector's output """

    def __init__(self):
        # These come from the TrivaQA official evaluation script
        self.skip = {"a", "an", "the", ""}
        self.strip = string.punctuation + "".join([u"‘", u"’", u"´", u"`", "_"])

        self.answer_tokens = None

    def set_question(self, normalized_aliases):
        self.answer_tokens = normalized_aliases

    def any_found(self, para):
        # Normalize the paragraph
        words = [w.lower().strip(self.strip) for w in flatten_iterable(para)]
        occurances = []
        for answer_ix, answer in enumerate(self.answer_tokens):
            # Locations where the first word occurs
            word_starts = [i for i, w in enumerate(words) if answer[0] == w]
            n_tokens = len(answer)

            # Advance forward until we find all the words, skipping over articles
            for start in word_starts:
                end = start + 1
                ans_token = 1
                while ans_token < n_tokens and end < len(words):
                    next = words[end]
                    if answer[ans_token] == next:
                        ans_token += 1
                        end += 1
                    elif next in self.skip:
                        end += 1
                    else:
                        break
                if n_tokens == ans_token:
                    occurances.append((start, end))
        return list(set(occurances)) class TriviaQaEvidenceCorpusTxt(object):
    """
    Corpus of the tokenized text from the given TriviaQa evidence documents.
    Allows the text to be retrieved by document id
    """

    _split_all = re.compile("[\n ]")
    _split_para = re.compile("\n\n+")  # FIXME we should not have saved document w/extra spaces...

    def __init__(self, file_id_map=None):
        self.directory = join(CORPUS_DIR, "triviaqa/evidence")
        self.file_id_map = file_id_map

    def get_vocab(self):
        with open(join(self.directory, "vocab.txt"), "r") as f:
            return {x.strip() for x in f}

    def load_word_vectors(self, vec_name):
        filename = join(self.directory, vec_name + "_pruned.pkl")
        if exists(filename):
            with open(filename, "rb"):
                return pickle.load(filename)
        else:
            return None

    def list_documents(self):
        if self.file_id_map is not None:
            return list(self.file_id_map.keys())
        f = []
        for dirpath, dirnames, filenames in walk(self.directory):
            if dirpath == self.directory:
                # Exclude files in the top level dir, like the vocab file
                continue
            if self.directory != dirpath:
                rel_path = relpath(dirpath, self.directory)
                f += [join(rel_path, x[:-4]) for x in filenames]
            else:
                f += [x[:-4] for x in filenames]
        return f

    def get_document(self, doc_id, n_tokens=None, flat=False):
        if self.file_id_map is None:
            file_id = doc_id
        else:
            file_id = self.file_id_map.get(doc_id)

        if file_id is None:
            return None

        file_id = join(self.directory, file_id + ".txt")
        if not exists(file_id):
            return None

        with open(file_id, "r") as f:
            if n_tokens is None:
                text = f.read()
                if flat:
                    return [x for x in self._split_all.split(text) if len(x) > 0]
                else:
                    paragraphs = []
                    for para in self._split_para.split(text):
                        paragraphs.append([sent.split(" ") for sent in para.split("\n")])
                    return paragraphs
            else:
                paragraphs = []
                paragraph = []
                cur_tokens = 0
                for line in f:
                    if line == "\n":
                        if not flat and len(paragraph) > 0:
                            paragraphs.append(paragraph)
                            paragraph = []
                    else:
                        sent = line.split(" ")
                        sent[-1] = sent[-1].rstrip()
                        if len(sent) + cur_tokens > n_tokens:
                            if n_tokens != cur_tokens:
                                paragraph.append(sent[:n_tokens-cur_tokens])
                            break
                        else:
                            paragraph.append(sent)
                            cur_tokens += len(sent)
                if flat:
                    return flatten_iterable(paragraph)
                else:
                    if len(paragraph) > 0:
                        paragraphs.append(paragraph)
                    return paragraphs def iter_trivia_question(filename, file_map, require_filename):
    return build_questions(iter_question_json(filename), file_map, require_filename) class TriviaQaQuestion(object):
    __slots__ = ["question", "question_id", "answer", "entity_docs", "web_docs"]

    def __init__(self, question, question_id, answer, entity_docs, web_docs):
        self.question = question
        self.question_id = question_id
        self.answer = answer
        self.entity_docs = entity_docs
        self.web_docs = web_docs

    @property
    def all_docs(self):
        if self.web_docs is not None:
            return self.web_docs + self.entity_docs
        else:
            return self.entity_docs

    def to_compressed_json(self):
        return [
            self.question,
            self.question_id,
            [self.answer.__class__.__name__] + [getattr(self.answer, x) for x in self.answer.__slots__],
            [[doc.__class__.__name__] + [getattr(doc, x) for x in doc.__slots__] for doc in self.entity_docs],
            [[getattr(doc, x) for x in doc.__slots__] for doc in self.web_docs],
        ]

    @staticmethod
    def from_compressed_json(text):
        question, quid, answer, entity_docs, web_docs = json.loads(text)
        if answer[0] == "WikipediaEntity":
            answer = WikipediaEntity(*answer[1:])
        elif answer[0] == "Numerical":
            answer = Numerical(*answer[1:])
        elif answer[0] == "FreeForm":
            answer = FreeForm(*answer[1:])
        elif answer[0] == "Range":
            answer = Range(*answer[1:])
        else:
            raise ValueError()
        for i, doc in enumerate(entity_docs):
            if doc[0] == "TagMeEntityDoc":
                entity_docs[i] = TagMeEntityDoc(*doc[1:])
            elif doc[0] == "SearchEntityDoc":
                    entity_docs[i] = SearchEntityDoc(*doc[1:])
        web_docs = [SearchDoc(*x) for x in web_docs]
        return TriviaQaQuestion(question, quid, answer, entity_docs, web_docs) 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)
import argparse
import json
import pickle
import unicodedata
from itertools import islice
from os import mkdir
from os.path import join, exists
from typing import List, Optional, Dict
from docqa.config import CORPUS_DIR, TRIVIA_QA, TRIVIA_QA_UNFILTERED
from docqa.configurable import Configurable
from docqa.data_processing.text_utils import NltkAndPunctTokenizer
from docqa.triviaqa.answer_detection import compute_answer_spans_par, FastNormalizedAnswerDetector
from docqa.triviaqa.evidence_corpus import TriviaQaEvidenceCorpusTxt
from docqa.triviaqa.read_data import iter_trivia_question, TriviaQaQuestion
from docqa.utils import ResourceLoader


"""
Build span-level training data from the raw trivia-qa inputs, in particular load the questions
from the json file and annotates each question/doc with the places the question answer's occur 
within the document, and save the results in our format. Assumes the evidence corpus has 
already been preprocessed 
"""


def build_dataset(name: str, tokenizer, train_files: Dict[str, str],
                  answer_detector, n_process: int, prune_unmapped_docs=True,
                  sample=None):
    out_dir = join(CORPUS_DIR, "triviaqa", name)
    if not exists(out_dir):
        mkdir(out_dir)

    file_map = {}  # maps document_id -> filename

    for name, filename in train_files.items():
        print("Loading %s questions" % name)
        if sample is None:
            questions = list(iter_trivia_question(filename, file_map, False))
        else:
            if isinstance(sample,  int):
                questions = list(islice(iter_trivia_question(filename, file_map, False), sample))
            elif isinstance(sample, dict):
                questions = list(islice(iter_trivia_question(filename, file_map, False), sample[name]))
            else:
                raise ValueError()

        if prune_unmapped_docs:
            for q in questions:
                if q.web_docs is not None:
                    q.web_docs = [x for x in q.web_docs if x.doc_id in file_map]
                q.entity_docs = [x for x in q.entity_docs if x.doc_id in file_map]

        print("Adding answers for %s question" % name)

        corpus = TriviaQaEvidenceCorpusTxt(file_map)