Lower text and remove punctuation articles and extra whitespace .

def normalize_answer(s):
    """Lower text and remove punctuation, articles and extra whitespace."""
    def remove_articles(text):
        return re.sub(r'\b(a|an|the)\b', ' ', text)

    def white_space_fix(text):
        return ' '.join(text.split())

    def remove_punc(text):
        exclude = set(string.punctuation)
        return ''.join(ch for ch in text if ch not in exclude)

    def lower(text):
        return text.lower()

    return white_space_fix(remove_articles(remove_punc(lower(s))))

def
normalize_answer
(
s
)
:
def
remove_articles
(
text
)
:
return
re
.
sub
(
r'\b(a|an|the)\b'
,
' '
,
text
)
def
white_space_fix
(
text
)
:
return
' '
.
join
(
text
.
split
(
)
)
def
remove_punc
(
text
)
:
exclude
=
set
(
string
.
punctuation
)
return
''
.
join
(
ch
for
ch
in
text
if
ch
not
in
exclude
)
def
lower
(
text
)
:
return
text
.
lower
(
)
return
white_space_fix
(
remove_articles
(
remove_punc
(
lower
(
s
)
)
)
)


Loads a data file into a list of InputBatch s .

def convert_examples_to_features(examples, max_seq_length, max_sent_num, max_sf_num, tokenizer, train=False):
    """Loads a data file into a list of `InputBatch`s."""

    DUMMY = [0] * max_seq_length
    DUMMY_ = [0.0] * max_sent_num
    features = []
    logger.info('#### Constructing features... ####')
    for (ex_index, example) in enumerate(tqdm(examples, desc='Example')):

        tokens_q = tokenizer.tokenize(
            'Q: {} A: {}'.format(example.question, example.answer))
        tokens_q = ['[CLS]'] + tokens_q + ['[SEP]']

        input_ids = []
        input_masks = []
        segment_ids = []

        for title in example.titles:
            sents = example.context[title]
            for (i, s) in enumerate(sents):

                if len(input_ids) == max_sent_num:
                    break

                tokens_s = tokenizer.tokenize(
                    s)[:max_seq_length-len(tokens_q)-1]
                tokens_s = tokens_s + ['[SEP]']

                padding = [0] * (max_seq_length -
                                 len(tokens_s) - len(tokens_q))

                input_ids_ = tokenizer.convert_tokens_to_ids(
                    tokens_q + tokens_s)
                input_masks_ = [1] * len(input_ids_)
                segment_ids_ = [0] * len(tokens_q) + [1] * len(tokens_s)

                input_ids_ += padding
                input_ids.append(input_ids_)

                input_masks_ += padding
                input_masks.append(input_masks_)

                segment_ids_ += padding
                segment_ids.append(segment_ids_)

                assert len(input_ids_) == max_seq_length
                assert len(input_masks_) == max_seq_length
                assert len(segment_ids_) == max_seq_length

        target_ids = []
        target_offset = 0

        for title in example.titles:
            sfs = example.supporting_facts[title]
            for i in sfs:
                if i < len(example.context[title]) and i+target_offset < len(input_ids):
                    target_ids.append(i+target_offset)
                else:
                    logger.warning('')
                    logger.warning('Invalid annotation: {}'.format(sfs))
                    logger.warning('Invalid annotation: {}'.format(
                        example.context[title]))

            target_offset += len(example.context[title])

        assert len(input_ids) <= max_sent_num
        assert len(target_ids) <= max_sf_num

        num_sents = len(input_ids)
        num_sfs = len(target_ids)

        output_masks = [([1.0] * len(input_ids) + [0.0] * (max_sent_num -
                                                           len(input_ids) + 1)) for _ in range(max_sent_num + 2)]

        if train:

            for i in range(len(target_ids)):
                for j in range(len(target_ids)):
                    if i == j:
                        continue

                    output_masks[i][target_ids[j]] = 0.0

            for i in range(len(output_masks)):
                if i >= num_sfs+1:
                    for j in range(len(output_masks[i])):
                        output_masks[i][j] = 0.0

        else:
            for i in range(len(input_ids)):
                output_masks[i+1][i] = 0.0

        target_ids += [0] * (max_sf_num - len(target_ids))

        padding = [DUMMY] * (max_sent_num - len(input_ids))
        input_ids += padding
        input_masks += padding
        segment_ids += padding

        features.append(
            InputFeatures(input_ids=input_ids,
                          input_masks=input_masks,
                          segment_ids=segment_ids,
                          target_ids=target_ids,
                          output_masks=output_masks,
                          num_sents=num_sents,
                          num_sfs=num_sfs,
                          ex_index=ex_index))

    logger.info('Done!')

    return features

def
convert_examples_to_features
(
examples
,
max_seq_length
,
max_sent_num
,
max_sf_num
,
tokenizer
,
train
=
False
)
:
DUMMY
=
[
0
]
*
max_seq_length
DUMMY_
=
[
0.0
]
*
max_sent_num
features
=
[
]
logger
.
info
(
'#### Constructing features... ####'
)
for
(
ex_index
,
example
)
in
enumerate
(
tqdm
(
examples
,
desc
=
'Example'
)
)
:
tokens_q
=
tokenizer
.
tokenize
(
'Q: {} A: {}'
.
format
(
example
.
question
,
example
.
answer
)
)
tokens_q
=
[
'[CLS]'
]
+
tokens_q
+
[
'[SEP]'
]
input_ids
=
[
]
input_masks
=
[
]
segment_ids
=
[
]
for
title
in
example
.
titles
:
sents
=
example
.
context
[
title
]
for
(
i
,
s
)
in
enumerate
(
sents
)
:
if
len
(
input_ids
)
==
max_sent_num
:
break
tokens_s
=
tokenizer
.
tokenize
(
s
)
[
:
max_seq_length
len
(
tokens_q
)
1
]
tokens_s
=
tokens_s
+
[
'[SEP]'
]
padding
=
[
0
]
*
(
max_seq_length
len
(
tokens_s
)
len
(
tokens_q
)
)
input_ids_
=
tokenizer
.
convert_tokens_to_ids
(
tokens_q
+
tokens_s
)
input_masks_
=
[
1
]
*
len
(
input_ids_
)
segment_ids_
=
[
0
]
*
len
(
tokens_q
)
+
[
1
]
*
len
(
tokens_s
)
input_ids_
+=
padding
input_ids
.
append
(
input_ids_
)
input_masks_
+=
padding
input_masks
.
append
(
input_masks_
)
segment_ids_
+=
padding
segment_ids
.
append
(
segment_ids_
)
assert
len
(
input_ids_
)
==
max_seq_length
assert
len
(
input_masks_
)
==
max_seq_length
assert
len
(
segment_ids_
)
==
max_seq_length
target_ids
=
[
]
target_offset
=
0
for
title
in
example
.
titles
:
sfs
=
example
.
supporting_facts
[
title
]
for
i
in
sfs
:
if
i
<
len
(
example
.
context
[
title
]
)
and
i
+
target_offset
<
len
(
input_ids
)
:
target_ids
.
append
(
i
+
target_offset
)
else
:
logger
.
warning
(
''
)
logger
.
warning
(
'Invalid annotation: {}'
.
format
(
sfs
)
)
logger
.
warning
(
'Invalid annotation: {}'
.
format
(
example
.
context
[
title
]
)
)
target_offset
+=
len
(
example
.
context
[
title
]
)
assert
len
(
input_ids
)
<=
max_sent_num
assert
len
(
target_ids
)
<=
max_sf_num
num_sents
=
len
(
input_ids
)
num_sfs
=
len
(
target_ids
)
output_masks
=
[
(
[
1.0
]
*
len
(
input_ids
)
+
[
0.0
]
*
(
max_sent_num
len
(
input_ids
)
+
1
)
)
for
_
in
range
(
max_sent_num
+
2
)
]
if
train
:
for
i
in
range
(
len
(
target_ids
)
)
:
for
j
in
range
(
len
(
target_ids
)
)
:
if
i
==
j
:
continue
output_masks
[
i
]
[
target_ids
[
j
]
]
=
0.0
for
i
in
range
(
len
(
output_masks
)
)
:
if
i
>=
num_sfs
+
1
:
for
j
in
range
(
len
(
output_masks
[
i
]
)
)
:
output_masks
[
i
]
[
j
]
=
0.0
else
:
for
i
in
range
(
len
(
input_ids
)
)
:
output_masks
[
i
+
1
]
[
i
]
=
0.0
target_ids
+=
[
0
]
*
(
max_sf_num
len
(
target_ids
)
)
padding
=
[
DUMMY
]
*
(
max_sent_num
len
(
input_ids
)
)
input_ids
+=
padding
input_masks
+=
padding
segment_ids
+=
padding
features
.
append
(
InputFeatures
(
input_ids
=
input_ids
,
input_masks
=
input_masks
,
segment_ids
=
segment_ids
,
target_ids
=
target_ids
,
output_masks
=
output_masks
,
num_sents
=
num_sents
,
num_sfs
=
num_sfs
,
ex_index
=
ex_index
)
)
logger
.
info
(
'Done!'
)
return
features


def convert_examples_to_features(examples, max_seq_length, max_para_num, graph_retriever_config, tokenizer, train = False):
    """Loads a data file into a list of `InputBatch`s."""

    if not train and graph_retriever_config.db_save_path is not None:
        max_para_num = graph_retriever_config.max_context_size
        graph_retriever_config.max_para_num = max(graph_retriever_config.max_para_num, max_para_num)

    max_steps = graph_retriever_config.max_select_num

    DUMMY = [0] * max_seq_length
    features = []

    logger.info('#### Converting examples to features... ####')
    for (ex_index, example) in enumerate(tqdm(examples, desc='Example')):
        tokens_q = tokenize_question(example.question, tokenizer)

        ##############
        # Short gold #
        ##############
        title2index = {}
        input_ids = []
        input_masks = []
        segment_ids = []

        # Append gold and non-gold paragraphs from context
        if train and graph_retriever_config.use_redundant and len(example.redundant_gold) > 0:
            if graph_retriever_config.use_multiple_redundant:
                titles_list = example.short_gold + [redundant[0] for redundant in example.all_redundant_gold] + list(example.context.keys())
            else:
                titles_list = example.short_gold + [example.redundant_gold[0]] + list(example.context.keys())
        else:
            titles_list = example.short_gold + list(example.context.keys())
        for p in titles_list:

            if len(input_ids) == max_para_num:
                break

            # Avoid appending gold paragraphs as negative
            if p in title2index:
                continue

            # fullwiki eval
            # Gold paragraphs are not always in context
            if not train and graph_retriever_config.open and p not in example.context:
                continue

            title2index[p] = len(title2index)
            example.title_order.append(p)
            p = example.context[p]

            input_ids_, input_masks_, segment_ids_ = tokenize_paragraph(p, tokens_q, max_seq_length, tokenizer)
            input_ids.append(input_ids_)
            input_masks.append(input_masks_)
            segment_ids.append(segment_ids_)

        # Open-domain setting
        if graph_retriever_config.open:
            num_paragraphs_no_links = len(input_ids)

            for p_ in example.context:

                if not train and graph_retriever_config.db_save_path is not None:
                    break

                if len(input_ids) == max_para_num:
                    break

                if p_ not in example.all_linked_paras_dic:
                    continue

                for l in example.all_linked_paras_dic[p_]:

                    if len(input_ids) == max_para_num:
                        break

                    if l in title2index:
                        continue

                    title2index[l] = len(title2index)
                    example.title_order.append(l)
                    p = example.all_linked_paras_dic[p_][l]

                    input_ids_, input_masks_, segment_ids_ = tokenize_paragraph(p, tokens_q, max_seq_length, tokenizer)
                    input_ids.append(input_ids_)
                    input_masks.append(input_masks_)
                    segment_ids.append(segment_ids_)

        assert len(input_ids) <= max_para_num

        num_paragraphs = len(input_ids)
        num_steps = len(example.short_gold)+1 # 1 for EOE

        if train:
            assert num_steps <= max_steps

        output_masks = [([1.0] * len(input_ids) + [0.0] * (max_para_num - len(input_ids) + 1)) for _ in range(max_para_num + 2)]

        if (not train) and graph_retriever_config.open:
            assert len(example.context) == num_paragraphs_no_links
            for i in range(len(output_masks[0])):
                if i >= num_paragraphs_no_links:
                    output_masks[0][i] = 0.0

        for i in range(len(input_ids)):
            output_masks[i+1][i] = 0.0

        if train:
            size = num_steps-1

            for i in range(size):
                for j in range(size):
                    if i != j:
                        output_masks[i][j] = 0.0

            for i in range(size):
                output_masks[size][i] = 0.0

            for i in range(max_steps):
                if i > size:
                    for j in range(len(output_masks[i])):
                        output_masks[i][j] = 0.0

            # Use REDUNDANT setting
            # Avoid treating the redundant paragraph as a negative example at the first step
            if graph_retriever_config.use_redundant and len(example.redundant_gold) > 0:
                if graph_retriever_config.use_multiple_redundant:
                    for redundant in example.all_redundant_gold:
                        output_masks[0][title2index[redundant[0]]] = 0.0
                else:
                    output_masks[0][title2index[example.redundant_gold[0]]] = 0.0

        padding = [DUMMY] * (max_para_num - len(input_ids))
        input_ids += padding
        input_masks += padding
        segment_ids += padding

        features.append(
                InputFeatures(input_ids=input_ids,
                              input_masks=input_masks,
                              segment_ids=segment_ids,
                              output_masks = output_masks,
                              num_paragraphs = num_paragraphs,
                              num_steps = num_steps,
                              ex_index = ex_index))

        if not train or not graph_retriever_config.use_redundant or len(example.redundant_gold) == 0:
            continue

        ##################
        # Redundant gold #
        ##################
        for redundant_gold in example.all_redundant_gold:
            hist = set()
            input_ids_r = []
            input_masks_r = []
            segment_ids_r = []

            # Append gold and non-gold paragraphs from context
            for p in redundant_gold + list(example.context.keys()):

                if len(input_ids_r) == max_para_num:
                    break

                #assert p in title2index
                if p not in title2index:
                    assert p not in redundant_gold
                    continue

                if p in hist:
                    continue
                hist.add(p)

                index = title2index[p]
                input_ids_r.append(input_ids[index])
                input_masks_r.append(input_masks[index])
                segment_ids_r.append(segment_ids[index])

            # Open-domain setting (mainly for HotpotQA fullwiki)
            if graph_retriever_config.open:

                for p in title2index:

                    if len(input_ids_r) == max_para_num:
                        break

                    if p in hist:
                        continue
                    hist.add(p)

                    index = title2index[p]
                    input_ids_r.append(input_ids[index])
                    input_masks_r.append(input_masks[index])
                    segment_ids_r.append(segment_ids[index])

            assert len(input_ids_r) <= max_para_num

            num_paragraphs_r = len(input_ids_r)
            num_steps_r = len(redundant_gold)+1

            assert num_steps_r <= max_steps

            output_masks_r = [([1.0] * len(input_ids_r) + [0.0] * (max_para_num - len(input_ids_r) + 1)) for _ in range(max_para_num + 2)]

            size = num_steps_r-1

            for i in range(size):
                for j in range(size):
                    if i != j:
                        output_masks_r[i][j] = 0.0

                if i > 0:
                    output_masks_r[i][0] = 1.0

            for i in range(size): #size-1
                output_masks_r[size][i] = 0.0

            for i in range(max_steps):
                if i > size:
                    for j in range(len(output_masks_r[i])):
                        output_masks_r[i][j] = 0.0

            padding = [DUMMY] * (max_para_num - len(input_ids_r))
            input_ids_r += padding
            input_masks_r += padding
            segment_ids_r += padding

            features.append(
                    InputFeatures(input_ids=input_ids_r,
                                  input_masks=input_masks_r,
                                  segment_ids=segment_ids_r,
                                  output_masks = output_masks_r,
                                  num_paragraphs = num_paragraphs_r,
                                  num_steps = num_steps_r,
                                  ex_index = None))

            if not graph_retriever_config.use_multiple_redundant:
                break

    logger.info('Done!')
    return features


Limit the number of examples used . This is mainly for sanity - chacking new settings .

def _create_examples(self, file_name, graph_retriever_config, set_type):

        task = graph_retriever_config.task
        jsn = json.load(open(file_name, 'r'))

        examples = []

        '''
        Limit the number of examples used.
        This is mainly for sanity-chacking new settings.
        '''
        if graph_retriever_config.example_limit is not None:
            random.shuffle(jsn)
            jsn = sorted(jsn, key = lambda x: x['q_id'])
            jsn = jsn[:graph_retriever_config.example_limit]

        '''
        Find the mximum size of the initial context (links are not included)
        '''
        graph_retriever_config.max_context_size = 0

        logger.info('#### Loading examples... from {} ####'.format(file_name))
        for (_, data) in enumerate(tqdm(jsn, desc='Example')):

            guid = data['q_id']
            question = data['question']
            context = data['context'] # {context title: paragraph}
            all_linked_paras_dic = data['all_linked_paras_dic'] # {context title: {linked title: paragraph}}
            short_gold = data['short_gold'] # [title 1, title 2] (Both are gold)
            redundant_gold = data['redundant_gold'] # [title 1, title 2, title 3] ("title 1" is not gold)
            all_redundant_gold = data['all_redundant_gold']

            '''
            Limit the number of redundant examples
            '''
            all_redundant_gold = all_redundant_gold[:graph_retriever_config.max_redundant_num]

            '''
            Control the size of the initial TF-IDF retrieved paragraphs
            *** Training time: to take a blalance between TF-IDF-based and link-based negative examples ***
            '''
            if graph_retriever_config.tfidf_limit is not None:
                new_context = {}
                for title in context:
                    if len(new_context) == graph_retriever_config.tfidf_limit:
                        break
                    new_context[title] = context[title]
                context = new_context

            '''
            Use TagMe-based context at test time.
            '''
            if set_type == 'dev' and task == 'nq' and graph_retriever_config.tagme:
                assert 'tagged_context' in data

                '''
                Reformat "tagged_context" if needed (c.f. the "context" case above)
                '''
                if type(data['tagged_context']) == list:
                    tagged_context = {c[0]: c[1] for c in data['tagged_context']}
                    data['tagged_context'] = tagged_context

                '''
                Append valid paragraphs from "tagged_context" to "context"
                '''
                for tagged_title in data['tagged_context']:
                    tagged_text = data['tagged_context'][tagged_title]
                    if tagged_title not in context and tagged_title is not None and tagged_title.strip() != '' and tagged_text is not None and tagged_text.strip() != '':
                        context[tagged_title] = tagged_text

            '''
            Clean "context" by removing invalid paragraphs
            '''
            removed_keys = []
            for title in context:
                if title is None or title.strip() == '' or context[title] is None or context[title].strip() == '':
                    removed_keys.append(title)
            for key in removed_keys:
                context.pop(key)

            if task in ['squad', 'nq'] and set_type == 'train':
                new_context = {}

                orig_title = list(context.keys())[0].split('_')[0]

                orig_titles = []
                other_titles = []

                for title in context:
                    title_ = title.split('_')[0]

                    if title_ == orig_title:
                        orig_titles.append(title)
                    else:
                        other_titles.append(title)

                orig_index = 0
                other_index = 0

                while orig_index < len(orig_titles) or other_index < len(other_titles):
                    if orig_index < len(orig_titles):
                        new_context[orig_titles[orig_index]] = context[orig_titles[orig_index]]
                        orig_index += 1

                    if other_index < len(other_titles):
                        new_context[other_titles[other_index]] = context[other_titles[other_index]]
                        other_index += 1

                context = new_context

            '''
            Convert link format
            '''
            new_all_linked_paras_dic = {} # {context title: {linked title: paragraph}}

            all_linked_paras_dic # {linked_title: paragraph} or mixed
            all_linked_para_title_dic = data['all_linked_para_title_dic'] # {context_title: [linked_title_1, linked_title_2, ...]}

            removed_keys = []
            tmp = {}
            for key in all_linked_paras_dic:
                if type(all_linked_paras_dic[key]) == dict:
                    removed_keys.append(key)

                    for linked_title in all_linked_paras_dic[key]:
                        if linked_title not in all_linked_paras_dic:
                            tmp[linked_title] = all_linked_paras_dic[key][linked_title]

                        if key in all_linked_para_title_dic:
                            all_linked_para_title_dic[key].append(linked_title)
                        else:
                            all_linked_para_title_dic[key] = [linked_title]

            for key in removed_keys:
                all_linked_paras_dic.pop(key)

            for key in tmp:
                if key not in all_linked_paras_dic:
                    all_linked_paras_dic[key] = tmp[key]

            for context_title in context:
                if context_title not in all_linked_para_title_dic:
                    continue

                new_entry = {}

                for linked_title in all_linked_para_title_dic[context_title]:
                    if linked_title not in all_linked_paras_dic:
                        continue

                    new_entry[linked_title] = all_linked_paras_dic[linked_title]

                if len(new_entry) > 0:
                    new_all_linked_paras_dic[context_title] = new_entry

            all_linked_paras_dic = new_all_linked_paras_dic

            if set_type == 'dev':
                '''
                Clean "all_linked_paras_dic" by removing invalid paragraphs
                '''
                for c in all_linked_paras_dic:
                    removed_keys = []
                    links = all_linked_paras_dic[c]
                    for title in links:
                        if title is None or title.strip() == '' or links[title] is None or type(links[title]) != str or links[title].strip() == '':
                            removed_keys.append(title)
                    for key in removed_keys:
                        links.pop(key)

                all_paras = {}
                for title in context:
                    all_paras[title] = context[title]

                    if not graph_retriever_config.open:
                        continue

                    if title not in all_linked_paras_dic:
                        continue
                    for title_ in all_linked_paras_dic[title]:
                        if title_ not in all_paras:
                            all_paras[title_] = all_linked_paras_dic[title][title_]
            else:
                all_paras = None

            if set_type == 'dev' and graph_retriever_config.expand_links:
                expand_links(context, all_linked_paras_dic, all_paras)

            if set_type == 'dev' and graph_retriever_config.no_links:
                all_linked_paras_dic = {}

            graph_retriever_config.max_context_size = max(graph_retriever_config.max_context_size, len(context))

            '''
            Ensure that all the gold paragraphs are included in "context"
            '''
            if set_type == 'train':
                for t in short_gold + redundant_gold:
                    assert t in context

            examples.append(InputExample(guid = guid,
                                         q = question,
                                         c = context,
                                         para_dic = all_linked_paras_dic,
                                         s_g = short_gold,
                                         r_g = redundant_gold,
                                         all_r_g = all_redundant_gold,
                                         all_paras = all_paras))

        if set_type == 'dev':
            examples = sorted(examples, key = lambda x: len(x.all_paras))
        logger.info('Done!')

        return examples

'''
            Convert link format
            '''


Fetch the text of a document and compute hashed ngrams counts .

def count(ngram, hash_size, multi_para, doc_id):
    """Fetch the text of a document and compute hashed ngrams counts."""
    global DOC2IDX
    # FIXME: remove hard coding.
    row, col, data = [], [], []
    # Tokenize

    if multi_para is True:
        # 1. if multi_para is true, the doc contains multiple paragraphs separated by \n\n and with links.
        tokens = tokenize(fetch_text_multi_para(doc_id))
    else:
        # 2. if not, only intro docs are retrieved and the sentences are separated by \t.
        # remove sentence separations ("\t") (only for HotpotQA).
        tokens = tokenize(fetch_text(doc_id).replace("\t", ""))

    # Get ngrams from tokens, with stopword/punctuation filtering.
    ngrams = tokens.ngrams(
        n=ngram, uncased=True, filter_fn=filter_ngram
    )

    # Hash ngrams and count occurences
    counts = Counter([hash(gram, hash_size)
                      for gram in ngrams])

    # Return in sparse matrix data format.
    row.extend(counts.keys())
    col.extend([DOC2IDX[doc_id]] * len(counts))
    data.extend(counts.values())
    return row, col, data


Form a sparse word to document count matrix ( inverted index ) .

def get_count_matrix(args, db, db_opts):
    """Form a sparse word to document count matrix (inverted index).

    M[i, j] = # times word i appears in document j.
    """
    # Map doc_ids to indexes
    global DOC2IDX
    db_class = get_class(db)
    with db_class(**db_opts) as doc_db:
        doc_ids = doc_db.get_doc_ids()
    DOC2IDX = {doc_id: i for i, doc_id in enumerate(doc_ids)}

    # Setup worker pool
    # TODO: Add tokenizer's choice.
    tok_class = SimpleTokenizer
    workers = ProcessPool(
        args.num_workers,
        initializer=init,
        initargs=(tok_class, db_class, db_opts)
    )

    # Compute the count matrix in steps (to keep in memory)
    logger.info('Mapping...')
    row, col, data = [], [], []
    step = max(int(len(doc_ids) / 10), 1)
    batches = [doc_ids[i:i + step] for i in range(0, len(doc_ids), step)]
    _count = partial(count, args.ngram, args.hash_size, args.multi_para)
    for i, batch in enumerate(batches):
        logger.info('-' * 25 + 'Batch %d/%d' %
                    (i + 1, len(batches)) + '-' * 25)
        for b_row, b_col, b_data in workers.imap_unordered(_count, batch):
            row.extend(b_row)
            col.extend(b_col)
            data.extend(b_data)
    workers.close()
    workers.join()

    logger.info('Creating sparse matrix...')
    count_matrix = sp.csr_matrix(
        (data, (row, col)), shape=(args.hash_size, len(doc_ids))
    )
    count_matrix.sum_duplicates()
    return count_matrix, (DOC2IDX, doc_ids)


Convert the word count matrix into tfidf one .

def get_tfidf_matrix(cnts):
    """Convert the word count matrix into tfidf one.

    tfidf = log(tf + 1) * log((N - Nt + 0.5) / (Nt + 0.5))
    * tf = term frequency in document
    * N = number of documents
    * Nt = number of occurences of term in all documents
    """
    Ns = get_doc_freqs(cnts)
    idfs = np.log((cnts.shape[1] - Ns + 0.5) / (Ns + 0.5))
    idfs[idfs < 0] = 0
    idfs = sp.diags(idfs, 0)
    tfs = cnts.log1p()
    tfidfs = idfs.dot(tfs)
    return tfidfs

def get_tfidf_matrix ( cnts ) : Ns = get_doc_freqs ( cnts ) idfs = np . log ( ( cnts . shape [ 1 ] - Ns + 0.5 ) / ( Ns + 0.5 ) ) idfs [ idfs < 0 ] = 0 idfs = sp . diags ( idfs , 0 ) tfs = cnts . log1p ( ) tfidfs = idfs . dot ( tfs ) return tfidfs


Return word -- > # of docs it appears in .

def get_doc_freqs(cnts):
    """Return word --> # of docs it appears in."""
    binary = (cnts > 0).astype(int)
    freqs = np.array(binary.sum(1)).squeeze()
    return freqs

def get_doc_freqs ( cnts ) : binary = ( cnts > 0 ) . astype ( int ) freqs = np . array ( binary . sum ( 1 ) ) . squeeze ( ) return freqs


def close(self):
        """Close the connection to the database."""
        self.connection.close()

def close ( self ) : self . connection . close ( )


Fetch all ids of docs stored in the db .

def get_doc_ids(self):
        """Fetch all ids of docs stored in the db."""
        cursor = self.connection.cursor()
        cursor.execute("SELECT id FROM documents")
        results = [r[0] for r in cursor.fetchall()]
        cursor.close()
        return results

def
get_doc_ids
(
self
)
:
cursor
=
self
.
connection
.
cursor
(
)
cursor
.
execute
(
"SELECT id FROM documents"
)
results
=
[
r
[
0
]
for
r
in
cursor
.
fetchall
(
)
]
cursor
.
close
(
)
return
results


Fetch the raw text of the doc for doc_id .

def get_doc_text(self, doc_id):
        """Fetch the raw text of the doc for 'doc_id'."""
        cursor = self.connection.cursor()
        cursor.execute(
            "SELECT text FROM documents WHERE id = ?",
            (doc_id,)
        )
        result = cursor.fetchone()
        cursor.close()
        return result if result is None else result[0]


Fetch the hyper - linked titles of the doc for doc_id .

def get_hyper_linked(self, doc_id):
        """Fetch the hyper-linked titles of the doc for 'doc_id'."""
        cursor = self.connection.cursor()
        cursor.execute(
            "SELECT linked_title FROM documents WHERE id = ?",
            (doc_id,)
        )
        result = cursor.fetchone()
        cursor.close()
        return result if (result is None or len(result[0]) == 0) else [normalize(title) for title in result[0].split("\t")]


Fetch the original title name of the doc .

def get_original_title(self, doc_id):
        """Fetch the original title name  of the doc."""
        cursor = self.connection.cursor()
        cursor.execute(
            "SELECT original_title FROM documents WHERE id = ?",
            (doc_id,)
        )
        result = cursor.fetchone()
        cursor.close()
        return result if result is None else result[0]


def get_doc_text_hyper_linked_titles_for_articles(self, doc_id):
        """
        fetch all of the paragraphs with their corresponding hyperlink titles.
        e.g.,
        >>> paras, links = db.get_doc_text_hyper_linked_titles_for_articles("Tokyo Imperial Palace_0")
        >>> paras[2]
        'It is built on the site of the old Edo Castle. The total area including the gardens is . During the height of the 1980s Japanese property bubble, the palace grounds were valued by some to be more than the value of all of the real estate in the state of California.'
        >>> links[2]
        ['Edo Castle', 'Japanese asset price bubble', 'Real estate', 'California']
        """
        cursor = self.connection.cursor()
        cursor.execute(
            "SELECT text FROM documents WHERE id = ?",
            (doc_id,)
        )
        result = cursor.fetchone()
        cursor.close()
        if result is None:
            return [], []
        else:
            hyper_linked_paragraphs = result[0].split("\n\n")
            paragraphs, hyper_linked_titles = [], []

            for hyper_linked_paragraph in hyper_linked_paragraphs:
                paragraphs.append(remove_tags(hyper_linked_paragraph))
                hyper_linked_titles.append([normalize(title) for title in find_hyper_linked_titles(
                    hyper_linked_paragraph)])

            return paragraphs, hyper_linked_titles


Import a module given a full path to the file .

def import_module(filename):
    """Import a module given a full path to the file."""
    spec = importlib.util.spec_from_file_location('doc_filter', filename)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module

def import_module ( filename ) : spec = importlib . util . spec_from_file_location ( 'doc_filter' , filename ) module = importlib . util . module_from_spec ( spec ) spec . loader . exec_module ( module ) return module


Walk through all files located under a root path .

def iter_files(path):
    """Walk through all files located under a root path."""
    if os.path.isfile(path):
        yield path
    elif os.path.isdir(path):
        for dirpath, _, filenames in os.walk(path):
            for f in filenames:
                yield os.path.join(dirpath, f)
    else:
        raise RuntimeError('Path %s is invalid' % path)

def
iter_files
(
path
)
:
if
os
.
path
.
isfile
(
path
)
:
yield
path
elif
os
.
path
.
isdir
(
path
)
:
for
dirpath
,
_
,
filenames
in
os
.
walk
(
path
)
:
for
f
in
filenames
:
yield
os
.
path
.
join
(
dirpath
,
f
)
else
:
raise
RuntimeError
(
'Path %s is invalid'
%
path
)


Parse the contents of a file . Each line is a JSON encoded document .

def get_contents_hotpotqa(filename):
    """Parse the contents of a file. Each line is a JSON encoded document."""
    global PREPROCESS_FN
    documents = []
    extracted_items = process_jsonlines_hotpotqa(filename)
    for extracted_item in extracted_items:
        title = extracted_item["title"]
        text = extracted_item["plain_text"]
        original_title = extracted_item["original_title"]
        hyper_linked_titles = extracted_item["hyper_linked_titles"]

        documents.append((title, text,
                          hyper_linked_titles, original_title))
    return documents

def get_contents_hotpotqa ( filename ) : global PREPROCESS_FN documents = [ ] extracted_items = process_jsonlines_hotpotqa ( filename ) for extracted_item in extracted_items : title = extracted_item [ "title" ] text = extracted_item [ "plain_text" ] original_title = extracted_item [ "original_title" ] hyper_linked_titles = extracted_item [ "hyper_linked_titles" ] documents . append ( ( title , text , hyper_linked_titles , original_title ) ) return documents


def get_contents(filename):
    """Parse the contents of a file. Each line is a JSON encoded document."""
    global PREPROCESS_FN
    documents = []
    extracted_items = process_jsonlines(filename)
    for extracted_item in extracted_items:
        title = extracted_item["title"]
        text = extracted_item["plain_text"]
        original_title = extracted_item["original_title"]
        hyper_linked_titles = extracted_item["hyper_linked_titles"]

        documents.append((title, text,
                          hyper_linked_titles, original_title))
    return documents

def get_contents ( filename ) : global PREPROCESS_FN documents = [ ] extracted_items = process_jsonlines ( filename ) for extracted_item in extracted_items : title = extracted_item [ "title" ] text = extracted_item [ "plain_text" ] original_title = extracted_item [ "original_title" ] hyper_linked_titles = extracted_item [ "hyper_linked_titles" ] documents . append ( ( title , text , hyper_linked_titles , original_title ) ) return documents


Preprocess and store a corpus of documents in sqlite .

def store_contents(wiki_dir, save_path, preprocess, num_workers=None, hotpotqa_format=False):
    """Preprocess and store a corpus of documents in sqlite.

    Args:
        data_path: Root path to directory (or directory of directories) of files
          containing json encoded documents (must have `id` and `text` fields).
        save_path: Path to output sqlite db.
        preprocess: Path to file defining a custom `preprocess` function. Takes
          in and outputs a structured doc.
        num_workers: Number of parallel processes to use when reading docs.
    """
    filenames = [f for f in glob.glob(
        wiki_dir + "/*/wiki_*", recursive=True) if ".bz2" not in f]
    if os.path.isfile(save_path):
        raise RuntimeError('%s already exists! Not overwriting.' % save_path)

    logger.info('Reading into database...')
    conn = sqlite3.connect(save_path)
    c = conn.cursor()
    c.execute(
        "CREATE TABLE documents (id PRIMARY KEY, text, linked_title, original_title);")

    workers = ProcessPool(num_workers, initializer=init,
                          initargs=(preprocess,))
    count = 0
    # Due to the slight difference of input format between preprocessed HotpotQA wikipedia data and
    # the ones by Wikiextractor, we call different functions for data process.
    if hotpotqa_format is True:
        content_processing_method = get_contents_hotpotqa
    else:
        content_processing_method = get_contents

    with tqdm(total=len(filenames)) as pbar:
        for pairs in tqdm(workers.imap_unordered(content_processing_method, filenames)):
            count += len(pairs)
            c.executemany(
                "INSERT OR REPLACE INTO documents VALUES (?,?,?,?)", pairs)
            pbar.update()
    logger.info('Read %d docs.' % count)
    logger.info('Committing...')
    conn.commit()
    conn.close()


def __len__(self):
        """The number of tokens."""
        return len(self.data)

def __len__ ( self ) : return len ( self . data )


Return a view of the list of tokens from [ i j ) .

def slice(self, i=None, j=None):
        """Return a view of the list of tokens from [i, j)."""
        new_tokens = copy.copy(self)
        new_tokens.data = self.data[i: j]
        return new_tokens

def slice ( self , i = None , j = None ) : new_tokens = copy . copy ( self ) new_tokens . data = self . data [ i : j ] return new_tokens


Returns the original text ( with whitespace reinserted ) .

def untokenize(self):
        """Returns the original text (with whitespace reinserted)."""
        return ''.join([t[self.TEXT_WS] for t in self.data]).strip()

def untokenize ( self ) : return '' . join ( [ t [ self . TEXT_WS ] for t in self . data ] ) . strip ( )


Returns a list of the text of each token Args : uncased : lower cases text

def words(self, uncased=False):
        """Returns a list of the text of each token
        Args:
            uncased: lower cases text
        """
        if uncased:
            return [t[self.TEXT].lower() for t in self.data]
        else:
            return [t[self.TEXT] for t in self.data]

def words ( self , uncased = False ) : if uncased : return [ t [ self . TEXT ] . lower ( ) for t in self . data ] else : return [ t [ self . TEXT ] for t in self . data ]


Returns a list of [ start end ) character offsets of each token .

def offsets(self):
        """Returns a list of [start, end) character offsets of each token."""
        return [t[self.SPAN] for t in self.data]

def offsets ( self ) : return [ t [ self . SPAN ] for t in self . data ]


Returns a list of part - of - speech tags of each token . Returns None if this annotation was not included .

def pos(self):
        """Returns a list of part-of-speech tags of each token.
        Returns None if this annotation was not included.
        """
        if 'pos' not in self.annotators:
            return None
        return [t[self.POS] for t in self.data]

def pos ( self ) : if 'pos' not in self . annotators : return None return [ t [ self . POS ] for t in self . data ]


Returns a list of the lemmatized text of each token . Returns None if this annotation was not included .

def lemmas(self):
        """Returns a list of the lemmatized text of each token.
        Returns None if this annotation was not included.
        """
        if 'lemma' not in self.annotators:
            return None
        return [t[self.LEMMA] for t in self.data]

def lemmas ( self ) : if 'lemma' not in self . annotators : return None return [ t [ self . LEMMA ] for t in self . data ]


Returns a list of named - entity - recognition tags of each token . Returns None if this annotation was not included .

def entities(self):
        """Returns a list of named-entity-recognition tags of each token.
        Returns None if this annotation was not included.
        """
        if 'ner' not in self.annotators:
            return None
        return [t[self.NER] for t in self.data]

def entities ( self ) : if 'ner' not in self . annotators : return None return [ t [ self . NER ] for t in self . data ]


Returns a list of all ngrams from length 1 to n . Args : n : upper limit of ngram length uncased : lower cases text filter_fn : user function that takes in an ngram list and returns True or False to keep or not keep the ngram as_string : return the ngram as a string vs list

def ngrams(self, n=1, uncased=False, filter_fn=None, as_strings=True):
        """Returns a list of all ngrams from length 1 to n.
        Args:
            n: upper limit of ngram length
            uncased: lower cases text
            filter_fn: user function that takes in an ngram list and returns
              True or False to keep or not keep the ngram
            as_string: return the ngram as a string vs list
        """
        def _skip(gram):
            if not filter_fn:
                return False
            return filter_fn(gram)

        words = self.words(uncased)
        ngrams = [(s, e + 1)
                  for s in range(len(words))
                  for e in range(s, min(s + n, len(words)))
                  if not _skip(words[s:e + 1])]

        # Concatenate into strings
        if as_strings:
            ngrams = ['{}'.format(' '.join(words[s:e])) for (s, e) in ngrams]

        return ngrams

def
ngrams
(
self
,
n
=
1
,
uncased
=
False
,
filter_fn
=
,
as_strings
=
True
)
:
def
_skip
(
gram
)
:
if
not
filter_fn
:
return
False
return
filter_fn
(
gram
)
words
=
self
.
words
(
uncased
)
ngrams
=
[
(
s
,
e
+
1
)
for
s
in
range
(
len
(
words
)
)
for
e
in
range
(
s
,
min
(
s
+
n
,
len
(
words
)
)
)
if
not
_skip
(
words
[
s
:
e
+
1
]
)
]
# Concatenate into strings
if
as_strings
:
ngrams
=
[
'{}'
.
format
(
' '
.
join
(
words
[
s
:
e
]
)
)
for
(
s
,
e
)
in
ngrams
]
return
ngrams


Group consecutive entity tokens with the same NER tag .

def entity_groups(self):
        """Group consecutive entity tokens with the same NER tag."""
        entities = self.entities()
        if not entities:
            return None
        non_ent = self.opts.get('non_ent', 'O')
        groups = []
        idx = 0
        while idx < len(entities):
            ner_tag = entities[idx]
            # Check for entity tag
            if ner_tag != non_ent:
                # Chomp the sequence
                start = idx
                while (idx < len(entities) and entities[idx] == ner_tag):
                    idx += 1
                groups.append((self.slice(start, idx).untokenize(), ner_tag))
            else:
                idx += 1
        return groups

def
entity_groups
(
self
)
:
entities
=
self
.
entities
(
)
if
not
entities
:
return
non_ent
=
self
.
opts
.
get
(
'non_ent'
,
'O'
)
groups
=
[
]
idx
=
0
while
idx
<
len
(
entities
)
:
ner_tag
=
entities
[
idx
]
# Check for entity tag
if
ner_tag
!=
non_ent
:
# Chomp the sequence
start
=
idx
while
(
idx
<
len
(
entities
)
and
entities
[
idx
]
==
ner_tag
)
:
idx
+=
1
groups
.
append
(
(
self
.
slice
(
start
,
idx
)
.
untokenize
(
)
,
ner_tag
)
)
else
:
idx
+=
1
return
groups


Args : annotators : None or empty set ( only tokenizes ) .

def __init__(self, **kwargs):
        """
        Args:
            annotators: None or empty set (only tokenizes).
        """
        self._regexp = regex.compile(
            '(%s)|(%s)' % (self.ALPHA_NUM, self.NON_WS),
            flags=regex.IGNORECASE + regex.UNICODE + regex.MULTILINE
        )
        if len(kwargs.get('annotators', {})) > 0:
            logger.warning('%s only tokenizes! Skipping annotators: %s' %
                           (type(self).__name__, kwargs.get('annotators')))
        self.annotators = set()


Resolve different type of unicode encodings / capitarization in HotpotQA data .

def normalize(text):
    """Resolve different type of unicode encodings / capitarization in HotpotQA data."""
    text = unicodedata.normalize('NFD', text)
    return text[0].capitalize() + text[1:]

def normalize ( text ) : text = unicodedata . normalize ( 'NFD' , text ) return text [ 0 ] . capitalize ( ) + text [ 1 : ]


This is process_jsonlines method for extracted Wikipedia file . After extracting items by using Wikiextractor ( with -- json and -- links options ) you will get the files named with wiki_xx where each line contains the information of each article . e . g . { id : 316 url : https : // en . wikipedia . org / wiki?curid = 316 title : Academy Award for Best Production Design text : Academy Award for Best Production Design \ n \ nThe  Academy Award for Best Production Design recognizes achievement for  art direction \ n \ n } This function takes these input and extract items . Each article contains one or more than one paragraphs and each paragraphs are separeated by \ n \ n .

def process_jsonlines(filename):
    """
    This is process_jsonlines method for extracted Wikipedia file.
    After extracting items by using Wikiextractor (with `--json` and `--links` options),
    you will get the files named with wiki_xx, where each line contains the information of each article.
    e.g.,
    {"id": "316", "url": "https://en.wikipedia.org/wiki?curid=316", "title": "Academy Award for Best Production Design",
    "text": "Academy Award for Best Production Design\n\nThe Academy Award for
    Best Production Design recognizes achievement for art direction \n\n"}
    This function takes these input and extract items.
    Each article contains one or more than one paragraphs, and each paragraphs are separeated by \n\n.
    """
    # item should be nested list
    extracted_items = []
    with jsonlines.open(filename) as reader:
        for obj in reader:
            wiki_id = obj["id"]
            title = obj["title"]
            title_id = make_wiki_id(title, 0)
            text_with_links = obj["text"]

            hyper_linked_titles_text = ""
            # When we consider the whole article as a document unit (e.g., SQuAD Open, Natural Questions Open)
            # we'll keep the links with the original articles, and dynamically process and extract the links
            # when we process with our selector.
            extracted_items.append({"wiki_id": wiki_id, "title": title_id,
                                    "plain_text": text_with_links,
                                    "hyper_linked_titles": hyper_linked_titles_text,
                                    "original_title": title})

    return extracted_items


This is process_jsonlines method for intro - only processed_wikipedia file . The item example : { id : 45668011 url : https : // en . wikipedia . org / wiki?curid = 45668011 title : Flouch Roundabout text : [ Flouch Roundabout is a roundabout near Penistone South Yorkshire England where the A628 meets the A616 . ] charoffset : [[[ 0 6 ] ... ]] text_with_links : [ Flouch Roundabout is a roundabout near  Penistone  South Yorkshire England where the  A628 meets the  A616 . ] charoffset_with_links : [[[ 0 6 ] ... [ 213 214 ]]] }

def process_jsonlines_hotpotqa(filename):
    """
    This is process_jsonlines method for intro-only processed_wikipedia file.
    The item example:
    {"id": "45668011", "url": "https://en.wikipedia.org/wiki?curid=45668011", "title": "Flouch Roundabout",
     "text": ["Flouch Roundabout is a roundabout near Penistone, South Yorkshire, England, where the A628 meets the A616."],
     "charoffset": [[[0, 6],...]]
     "text_with_links" : ["Flouch Roundabout is a roundabout near Penistone,
     South Yorkshire, England, where the A628
     meets the A616."],
        "charoffset_with_links": [[[0, 6], ... [213, 214]]]}
    """
    # item should be nested list
    extracted_items = []
    with jsonlines.open(filename) as reader:
        for obj in reader:
            wiki_id = obj["id"]
            title = obj["title"]
            title_id = make_wiki_id(title, 0)
            plain_text = "\t".join(obj["text"])
            text_with_links = "\t".join(obj["text_with_links"])

            hyper_linked_titles = []
            hyper_linked_titles = find_hyper_linked_titles(text_with_links)
            if len(hyper_linked_titles) > 0:
                hyper_linked_titles_text = "\t".join(hyper_linked_titles)
            else:
                hyper_linked_titles_text = ""
            extracted_items.append({"wiki_id": wiki_id, "title": title_id,
                                    "plain_text": plain_text,
                                    "hyper_linked_titles": hyper_linked_titles_text,
                                    "original_title": title})

    return extracted_items

def
process_jsonlines_hotpotqa
(
filename
)
:
# item should be nested list
extracted_items
=
[
]
with
jsonlines
.
open
(
filename
)
as
reader
:
for
obj
in
reader
:
wiki_id
=
obj
[
"id"
]
title
=
obj
[
"title"
]
title_id
=
make_wiki_id
(
title
,
0
)
plain_text
=
"\t"
.
join
(
obj
[
"text"
]
)
text_with_links
=
"\t"
.
join
(
obj
[
"text_with_links"
]
)
hyper_linked_titles
=
[
]
hyper_linked_titles
=
find_hyper_linked_titles
(
text_with_links
)
if
len
(
hyper_linked_titles
)
>
0
:
hyper_linked_titles_text
=
"\t"
.
join
(
hyper_linked_titles
)
else
:
hyper_linked_titles_text
=
""
extracted_items
.
append
(
{
"wiki_id"
:
wiki_id
,
"title"
:
title_id
,
"plain_text"
:
plain_text
,
"hyper_linked_titles"
:
hyper_linked_titles_text
,
"original_title"
:
title
}
)
return
extracted_items


Unsigned 32 bit murmurhash for feature hashing .

def hash(token, num_buckets):
    """Unsigned 32 bit murmurhash for feature hashing."""
    return murmurhash3_32(token, positive=True) % num_buckets

def hash ( token , num_buckets ) : return murmurhash3_32 ( token , positive = True ) % num_buckets


Take out english stopwords punctuation and compound endings .

def filter_word(text):
    """Take out english stopwords, punctuation, and compound endings."""
    text = normalize(text)
    if regex.match(r'^\p{P}+$', text):
        return True
    if text.lower() in STOPWORDS:
        return True
    return False

def filter_word ( text ) : text = normalize ( text ) if regex . match ( r'^\p{P}+$' , text ) : return True if text . lower ( ) in STOPWORDS : return True return False


Decide whether to keep or discard an n - gram . Args : gram : list of tokens ( length N ) mode : Option to throw out ngram if any : any single token passes filter_word all : all tokens pass filter_word ends : book - ended by filterable tokens

def filter_ngram(gram, mode='any'):
    """Decide whether to keep or discard an n-gram.
    Args:
        gram: list of tokens (length N)
        mode: Option to throw out ngram if
          'any': any single token passes filter_word
          'all': all tokens pass filter_word
          'ends': book-ended by filterable tokens
    """
    filtered = [filter_word(w) for w in gram]
    if mode == 'any':
        return any(filtered)
    elif mode == 'all':
        return all(filtered)
    elif mode == 'ends':
        return filtered[0] or filtered[-1]
    else:
        raise ValueError('Invalid mode: %s' % mode)

def
filter_ngram
(
gram
,
mode
=
'any'
)
:
filtered
=
[
filter_word
(
w
)
for
w
in
gram
]
if
mode
==
'any'
:
return
any
(
filtered
)
elif
mode
==
'all'
:
return
all
(
filtered
)
elif
mode
==
'ends'
:
return
filtered
[
0
]
or
filtered
[
1
]
else
:
raise
ValueError
(
'Invalid mode: %s'
%
mode
)


get the subfield associated to a list of elastic fields E . g . [ file filename ] to d [ file ] [ filename ]

def get_field(d, field_list):
    """get the subfield associated to a list of elastic fields
        E.g. ['file', 'filename'] to d['file']['filename']
    """
    if isinstance(field_list, str):
        return d[field_list]
    else:
        idx = d.copy()
        for field in field_list:
            idx = idx[field]
        return idx

def get_field ( d , field_list ) : if isinstance ( field_list , str ) : return d [ field_list ] else : idx = d . copy ( ) for field in field_list : idx = idx [ field ] return idx


load paragraphs and hyperlinked titles from DB . This method is mainly for Natural Questions Open benchmark .

def load_para_and_linked_titles_dict_from_tfidf_id(tfidf_id, db):
    """
    load paragraphs and hyperlinked titles from DB.
    This method is mainly for Natural Questions Open benchmark.
    """
    # will be fixed in the later version; current tfidf weights use indexed titles as keys.
    if "_0" not in tfidf_id:
        tfidf_id = "{0}_0".format(tfidf_id)
    paras, linked_titles = db.get_doc_text_hyper_linked_titles_for_articles(
        tfidf_id)
    if len(paras) == 0:
        logger.warning("{0} is missing".format(tfidf_id))
        return [], []

    paras_dict = {}
    linked_titles_dict = {}
    article_name = tfidf_id.split("_0")[0]
    # store the para_dict and linked_titles_dict; skip the first para (title)
    for para_idx, (para, linked_title_list) in enumerate(zip(paras[1:], linked_titles[1:])):
        paras_dict["{0}_{1}".format(article_name, para_idx)] = para
        linked_titles_dict["{0}_{1}".format(
            article_name, para_idx)] = linked_title_list

    return paras_dict, linked_titles_dict


Args : tfidf_path : path to saved model file strict : fail on empty queries or continue ( and return empty result )

def __init__(self, tfidf_path=None, strict=True):
        """
        Args:
            tfidf_path: path to saved model file
            strict: fail on empty queries or continue (and return empty result)
        """
        # Load from disk
        tfidf_path = tfidf_path
        logger.info('Loading %s' % tfidf_path)
        matrix, metadata = load_sparse_csr(tfidf_path)
        self.doc_mat = matrix
        self.ngrams = metadata['ngram']
        self.hash_size = metadata['hash_size']
        self.tokenizer = SimpleTokenizer()
        self.doc_freqs = metadata['doc_freqs'].squeeze()
        self.doc_dict = metadata['doc_dict']
        self.num_docs = len(self.doc_dict[0])
        self.strict = strict

def
__init__
(
self
,
tfidf_path
=
,
strict
=
True
)
:
# Load from disk
tfidf_path
=
tfidf_path
logger
.
info
(
'Loading %s'
%
tfidf_path
)
matrix
,
metadata
=
load_sparse_csr
(
tfidf_path
)
self
.
doc_mat
=
matrix
self
.
ngrams
=
metadata
[
'ngram'
]
self
.
hash_size
=
metadata
[
'hash_size'
]
self
.
tokenizer
=
SimpleTokenizer
(
)
self
.
doc_freqs
=
metadata
[
'doc_freqs'
]
.
squeeze
(
)
self
.
doc_dict
=
metadata
[
'doc_dict'
]
self
.
num_docs
=
len
(
self
.
doc_dict
[
0
]
)
self
.
strict
=
strict


def get_doc_index(self, doc_id):
        """Convert doc_id --> doc_index"""
        return self.doc_dict[0][doc_id]

def get_doc_index ( self , doc_id ) : return self . doc_dict [ 0 ] [ doc_id ]


def get_doc_id(self, doc_index):
        """Convert doc_index --> doc_id"""
        return self.doc_dict[1][doc_index]

def get_doc_id ( self , doc_index ) : return self . doc_dict [ 1 ] [ doc_index ]


Closest docs by dot product between query and documents in tfidf weighted word vector space .

def closest_docs(self, query, k=1):
        """Closest docs by dot product between query and documents
        in tfidf weighted word vector space.
        """
        spvec = self.text2spvec(query)
        res = spvec * self.doc_mat

        if len(res.data) <= k:
            o_sort = np.argsort(-res.data)
        else:
            o = np.argpartition(-res.data, k)[0:k]
            o_sort = o[np.argsort(-res.data[o])]

        doc_scores = res.data[o_sort]
        doc_ids = [self.get_doc_id(i) for i in res.indices[o_sort]]
        return doc_ids, doc_scores

def closest_docs ( self , query , k = 1 ) : spvec = self . text2spvec ( query ) res = spvec * self . doc_mat if len ( res . data ) <= k : o_sort = np . argsort ( - res . data ) else : o = np . argpartition ( - res . data , k ) [ 0 : k ] o_sort = o [ np . argsort ( - res . data [ o ] ) ] doc_scores = res . data [ o_sort ] doc_ids = [ self . get_doc_id ( i ) for i in res . indices [ o_sort ] ] return doc_ids , doc_scores


Process a batch of closest_docs requests multithreaded . Note : we can use plain threads here as scipy is outside of the GIL .

def batch_closest_docs(self, queries, k=1, num_workers=None):
        """Process a batch of closest_docs requests multithreaded.
        Note: we can use plain threads here as scipy is outside of the GIL.
        """
        with ThreadPool(num_workers) as threads:
            closest_docs = partial(self.closest_docs, k=k)
            results = threads.map(closest_docs, queries)
        return results

def batch_closest_docs ( self , queries , k = 1 , num_workers = None ) : with ThreadPool ( num_workers ) as threads : closest_docs = partial ( self . closest_docs , k = k ) results = threads . map ( closest_docs , queries ) return results


Parse the query into tokens ( either ngrams or tokens ) .

def parse(self, query):
        """Parse the query into tokens (either ngrams or tokens)."""
        tokens = self.tokenizer.tokenize(query)
        return tokens.ngrams(n=self.ngrams, uncased=True,
                             filter_fn=filter_ngram)

def parse ( self , query ) : tokens = self . tokenizer . tokenize ( query ) return tokens . ngrams ( n = self . ngrams , uncased = True , filter_fn = filter_ngram )


Create a sparse tfidf - weighted word vector from query .

def text2spvec(self, query):
        """Create a sparse tfidf-weighted word vector from query.

        tfidf = log(tf + 1) * log((N - Nt + 0.5) / (Nt + 0.5))
        """
        # Get hashed ngrams
        # TODO: do we need to have normalize?
        words = self.parse(normalize(query))
        wids = [hash(w, self.hash_size) for w in words]

        if len(wids) == 0:
            if self.strict:
                raise RuntimeError('No valid word in: %s' % query)
            else:
                logger.warning('No valid word in: %s' % query)
                return sp.csr_matrix((1, self.hash_size))

        # Count TF
        wids_unique, wids_counts = np.unique(wids, return_counts=True)
        tfs = np.log1p(wids_counts)

        # Count IDF
        Ns = self.doc_freqs[wids_unique]
        idfs = np.log((self.num_docs - Ns + 0.5) / (Ns + 0.5))
        idfs[idfs < 0] = 0

        # TF-IDF
        data = np.multiply(tfs, idfs)

        # One row, sparse csr matrix
        indptr = np.array([0, len(wids_unique)])
        spvec = sp.csr_matrix(
            (data, wids_unique, indptr), shape=(1, self.hash_size)
        )

        return spvec

def
text2spvec
(
self
,
query
)
:
# Get hashed ngrams
# TODO: do we need to have normalize?
words
=
self
.
parse
(
normalize
(
query
)
)
wids
=
[
hash
(
w
,
self
.
hash_size
)
for
w
in
words
]
if
len
(
wids
)
==
0
:
if
self
.
strict
:
raise
RuntimeError
(
'No valid word in: %s'
%
query
)
else
:
logger
.
warning
(
'No valid word in: %s'
%
query
)
return
sp
.
csr_matrix
(
(
1
,
self
.
hash_size
)
)
# Count TF
wids_unique
,
wids_counts
=
np
.
unique
(
wids
,
return_counts
=
True
)
tfs
=
np
.
log1p
(
wids_counts
)
# Count IDF
Ns
=
self
.
doc_freqs
[
wids_unique
]
idfs
=
np
.
log
(
(
self
.
num_docs
Ns
+
0.5
)
/
(
Ns
+
0.5
)
)
idfs
[
idfs
<
0
]
=
0
# TF-IDF
data
=
np
.
multiply
(
tfs
,
idfs
)
# One row, sparse csr matrix
indptr
=
np
.
array
(
[
0
,
len
(
wids_unique
)
]
)
spvec
=
sp
.
csr_matrix
(
(
data
,
wids_unique
,
indptr
)
,
shape
=
(
1
,
self
.
hash_size
)
)
return
spvec


Read a SQuAD json file into a list of SquadExample .

def read_squad_examples(input_file, is_training, version_2_with_negative, max_answer_len=100000, skip_negatives=False):
    """Read a SQuAD json file into a list of SquadExample."""
    with open(input_file, "r", encoding='utf-8') as reader:
        input_data = json.load(reader)["data"]

    def is_whitespace(c):
        if c == " " or c == "\t" or c == "\r" or c == "\n" or ord(c) == 0x202F:
            return True
        return False

    examples = []
    for entry in input_data:
        for paragraph in entry["paragraphs"]:
            paragraph_text = paragraph["context"]
            doc_tokens = []
            char_to_word_offset = []
            prev_is_whitespace = True
            for c in paragraph_text:
                if is_whitespace(c):
                    prev_is_whitespace = True
                else:
                    if prev_is_whitespace:
                        doc_tokens.append(c)
                    else:
                        doc_tokens[-1] += c
                    prev_is_whitespace = False
                char_to_word_offset.append(len(doc_tokens) - 1)

            for qa in paragraph["qas"]:
                qas_id = str(qa["id"])
                # this is temporary added to see whether reducing the negatives
                # improves the performance.
                if skip_negatives is True and "_NEGATIVE_" in qas_id:
                    continue
                if "FAKE2" in qas_id:
                    continue
                question_text = qa["question"]
                start_position = None
                end_position = None
                orig_answer_text = None
                is_impossible = False
                if is_training:
                    if version_2_with_negative:
                        is_impossible = qa["is_impossible"]
                        if is_impossible is False:
                            switch = 0
                        else:
                            switch = 1
                    # if (len(qa["answers"]) != 1) and (not is_impossible):
                    #    raise ValueError(
                    #        "For training, each question should have exactly 1 answer.")
                    if not is_impossible:
                        answers = qa['answers']
                        if type(answers) == list:
                            answer = qa["answers"][0]
                        else:
                            answer = answers
                        orig_answer_text = answer["text"]
                        answer_offset = answer["answer_start"]
                        answer_length = len(orig_answer_text)
                        start_position = char_to_word_offset[answer_offset]
                        end_position = char_to_word_offset[answer_offset +
                                                           answer_length - 1]
                        # Only add answers where the text can be exactly recovered from the
                        # document. If this CAN'T happen it's likely due to weird Unicode
                        # stuff so we will just skip the example.
                        #
                        # Note that this means for training mode, every example is NOT
                        # guaranteed to be preserved.
                        actual_text = " ".join(
                            doc_tokens[start_position:(end_position + 1)])
                        cleaned_answer_text = " ".join(
                            whitespace_tokenize(orig_answer_text))
                        if actual_text.find(cleaned_answer_text) == -1:
                            logger.warning("Could not find answer: '%s' vs. '%s'",
                                           actual_text, cleaned_answer_text)
                            continue

                        if len(orig_answer_text.split()) > max_answer_len:
                            logger.info(
                                "Omitting a long answer: '%s'", orig_answer_text)
                            continue
                    else:
                        start_position = -1
                        end_position = -1
                        orig_answer_text = ""

                        if len(qa["answers"]) > 0:
                            answer = qa["answers"][0]
                            # Make sure that answer text will be preserved for
                            # yes/no.
                            if answer["text"] in ["yes", "no"]:
                                orig_answer_text = answer["text"]
                if not is_training:
                    switch = None
                example = SquadExample(
                    qas_id=qas_id,
                    question_text=question_text,
                    doc_tokens=doc_tokens,
                    orig_answer_text=orig_answer_text,
                    start_position=start_position,
                    end_position=end_position,
                    switch=switch,
                    is_impossible=is_impossible)
                examples.append(example)
    return examples


Loads a data file into a list of InputBatch s .

def convert_examples_to_features(examples, tokenizer, max_seq_length,
                                 doc_stride, max_query_length, is_training,
                                 cls_token_at_end=False,
                                 cls_token='[CLS]', sep_token='[SEP]', pad_token=0,
                                 sequence_a_segment_id=0, sequence_b_segment_id=1,
                                 cls_token_segment_id=0, pad_token_segment_id=0,
                                 mask_padding_with_zero=True,
                                 quiet=False):
    """Loads a data file into a list of `InputBatch`s."""

    unique_id = 1000000000

    features = []
    for (example_index, example) in enumerate(examples):
        query_tokens = tokenizer.tokenize(example.question_text)

        if len(query_tokens) > max_query_length:
            query_tokens = query_tokens[0:max_query_length]

        tok_to_orig_index = []
        orig_to_tok_index = []
        all_doc_tokens = []
        for (i, token) in enumerate(example.doc_tokens):
            orig_to_tok_index.append(len(all_doc_tokens))
            sub_tokens = tokenizer.tokenize(token)
            for sub_token in sub_tokens:
                tok_to_orig_index.append(i)
                all_doc_tokens.append(sub_token)

        tok_start_position = None
        tok_end_position = None
        if is_training and example.is_impossible:
            tok_start_position = -1
            tok_end_position = -1
        if is_training and not example.is_impossible:
            tok_start_position = orig_to_tok_index[example.start_position]
            if example.end_position < len(example.doc_tokens) - 1:
                tok_end_position = orig_to_tok_index[example.end_position + 1] - 1
            else:
                tok_end_position = len(all_doc_tokens) - 1
            (tok_start_position, tok_end_position) = _improve_answer_span(
                all_doc_tokens, tok_start_position, tok_end_position, tokenizer,
                example.orig_answer_text)

        # The -3 accounts for [CLS], [SEP] and [SEP]
        max_tokens_for_doc = max_seq_length - len(query_tokens) - 3

        # We can have documents that are longer than the maximum sequence length.
        # To deal with this we do a sliding window approach, where we take chunks
        # of the up to our max length with a stride of `doc_stride`.
        _DocSpan = collections.namedtuple(  # pylint: disable=invalid-name
            "DocSpan", ["start", "length"])
        doc_spans = []
        start_offset = 0
        while start_offset < len(all_doc_tokens):
            length = len(all_doc_tokens) - start_offset
            if length > max_tokens_for_doc:
                length = max_tokens_for_doc
            doc_spans.append(_DocSpan(start=start_offset, length=length))
            if start_offset + length == len(all_doc_tokens):
                break
            start_offset += min(length, doc_stride)

        for (doc_span_index, doc_span) in enumerate(doc_spans):
            tokens = []
            token_to_orig_map = {}
            token_is_max_context = {}
            segment_ids = []

            p_mask = []

            # CLS token at the beginning
            if not cls_token_at_end:
                tokens.append(cls_token)
                segment_ids.append(cls_token_segment_id)
                p_mask.append(0)
                cls_index = 0

            # Query
            for token in query_tokens:
                tokens.append(token)
                segment_ids.append(sequence_a_segment_id)
                p_mask.append(1)

            # SEP token
            tokens.append(sep_token)
            segment_ids.append(sequence_a_segment_id)
            p_mask.append(1)

            # Paragraph
            for i in range(doc_span.length):
                split_token_index = doc_span.start + i
                token_to_orig_map[len(
                    tokens)] = tok_to_orig_index[split_token_index]

                is_max_context = _check_is_max_context(doc_spans, doc_span_index,
                                                       split_token_index)
                token_is_max_context[len(tokens)] = is_max_context
                tokens.append(all_doc_tokens[split_token_index])
                segment_ids.append(sequence_b_segment_id)
                p_mask.append(0)
            paragraph_len = doc_span.length

            # SEP token
            tokens.append(sep_token)
            segment_ids.append(sequence_b_segment_id)
            p_mask.append(1)

            # CLS token at the end
            if cls_token_at_end:
                tokens.append(cls_token)
                segment_ids.append(cls_token_segment_id)
                p_mask.append(0)
                cls_index = len(tokens) - 1  # Index of classification token

            input_ids = tokenizer.convert_tokens_to_ids(tokens)

            # The mask has 1 for real tokens and 0 for padding tokens. Only real
            # tokens are attended to.
            input_mask = [1 if mask_padding_with_zero else 0] * len(input_ids)

            # Zero-pad up to the sequence length.
            while len(input_ids) < max_seq_length:
                input_ids.append(pad_token)
                input_mask.append(0 if mask_padding_with_zero else 1)
                segment_ids.append(pad_token_segment_id)
                p_mask.append(1)

            assert len(input_ids) == max_seq_length
            assert len(input_mask) == max_seq_length
            assert len(segment_ids) == max_seq_length

            span_is_impossible = example.is_impossible
            start_position = None
            end_position = None
            switch = None
            if is_training and not span_is_impossible:
                doc_start = doc_span.start
                doc_end = doc_span.start + doc_span.length - 1
                out_of_span = False
                if not (tok_start_position >= doc_start and
                        tok_end_position <= doc_end):
                    out_of_span = True
                if out_of_span:
                    start_position = 0
                    end_position = 0
                    span_is_impossible = True
                else:
                    doc_offset = len(query_tokens) + 2
                    start_position = tok_start_position - doc_start + doc_offset
                    end_position = tok_end_position - doc_start + doc_offset

            if is_training and span_is_impossible:
                start_position = cls_index
                end_position = cls_index
                switch = 1
            elif is_training and not span_is_impossible:
                switch = 0

            # The questions whose ``is_impossible'' are originally True should
            # be 1.
            if example.is_impossible is True:
                switch = 1

            if example_index < 20 and not quiet:
                logger.info("*** Example ***")
                logger.info("unique_id: %s" % (unique_id))
                logger.info("example_index: %s" % (example_index))
                logger.info("doc_span_index: %s" % (doc_span_index))
                logger.info("tokens: %s" % " ".join(tokens))
                logger.info("token_to_orig_map: %s" % " ".join([
                    "%d:%d" % (x, y) for (x, y) in token_to_orig_map.items()]))
                logger.info("token_is_max_context: %s" % " ".join([
                    "%d:%s" % (x, y) for (x, y) in token_is_max_context.items()
                ]))
                logger.info("input_ids: %s" %
                            " ".join([str(x) for x in input_ids]))
                logger.info(
                    "input_mask: %s" % " ".join([str(x) for x in input_mask]))
                logger.info(
                    "segment_ids: %s" % " ".join([str(x) for x in segment_ids]))
                if is_training and span_is_impossible:
                    logger.info("impossible example")
                if is_training and not span_is_impossible:
                    answer_text = " ".join(
                        tokens[start_position:(end_position + 1)])
                    logger.info("start_position: %d" % (start_position))
                    logger.info("end_position: %d" % (end_position))
                    logger.info(
                        "answer: %s" % (answer_text))

            features.append(
                InputFeatures(
                    unique_id=unique_id,
                    example_index=example_index,
                    doc_span_index=doc_span_index,
                    tokens=tokens,
                    token_to_orig_map=token_to_orig_map,
                    token_is_max_context=token_is_max_context,
                    input_ids=input_ids,
                    input_mask=input_mask,
                    segment_ids=segment_ids,
                    cls_index=cls_index,
                    p_mask=p_mask,
                    paragraph_len=paragraph_len,
                    start_position=start_position,
                    end_position=end_position,
                    switch=switch,
                    is_impossible=span_is_impossible))
            unique_id += 1

    return features


def convert_examples_to_features_yes_no(examples, tokenizer, max_seq_length,
                                        doc_stride, max_query_length, is_training,
                                        cls_token_at_end=False,
                                        cls_token='[CLS]', sep_token='[SEP]', pad_token=0,
                                        sequence_a_segment_id=0, sequence_b_segment_id=1,
                                        cls_token_segment_id=0, pad_token_segment_id=0,
                                        mask_padding_with_zero=True):
    """Loads a data file into a list of `InputBatch`s."""

    unique_id = 1000000000

    features = []
    for (example_index, example) in enumerate(examples):
        query_tokens = tokenizer.tokenize(example.question_text)

        if len(query_tokens) > max_query_length:
            query_tokens = query_tokens[0:max_query_length]

        tok_to_orig_index = []
        orig_to_tok_index = []
        all_doc_tokens = []
        for (i, token) in enumerate(example.doc_tokens):
            orig_to_tok_index.append(len(all_doc_tokens))
            sub_tokens = tokenizer.tokenize(token)
            for sub_token in sub_tokens:
                tok_to_orig_index.append(i)
                all_doc_tokens.append(sub_token)

        tok_start_position = None
        tok_end_position = None
        if is_training and example.is_impossible:
            tok_start_position = -1
            tok_end_position = -1
        if is_training and not example.is_impossible:
            tok_start_position = orig_to_tok_index[example.start_position]
            if example.end_position < len(example.doc_tokens) - 1:
                tok_end_position = orig_to_tok_index[example.end_position + 1] - 1
            else:
                tok_end_position = len(all_doc_tokens) - 1
            (tok_start_position, tok_end_position) = _improve_answer_span(
                all_doc_tokens, tok_start_position, tok_end_position, tokenizer,
                example.orig_answer_text)

        # The -3 accounts for [CLS], [SEP] and [SEP]
        max_tokens_for_doc = max_seq_length - len(query_tokens) - 3

        # We can have documents that are longer than the maximum sequence length.
        # To deal with this we do a sliding window approach, where we take chunks
        # of the up to our max length with a stride of `doc_stride`.
        _DocSpan = collections.namedtuple(  # pylint: disable=invalid-name
            "DocSpan", ["start", "length"])
        doc_spans = []
        start_offset = 0
        while start_offset < len(all_doc_tokens):
            length = len(all_doc_tokens) - start_offset
            if length > max_tokens_for_doc:
                length = max_tokens_for_doc
            doc_spans.append(_DocSpan(start=start_offset, length=length))
            if start_offset + length == len(all_doc_tokens):
                break
            start_offset += min(length, doc_stride)

        for (doc_span_index, doc_span) in enumerate(doc_spans):
            tokens = []
            token_to_orig_map = {}
            token_is_max_context = {}
            segment_ids = []

            # p_mask: mask with 1 for token than cannot be in the answer (0 for token which can be in an answer)
            # Original TF implem also keep the classification token (set to 0)
            # (not sure why...)
            p_mask = []

            # CLS token at the beginning
            if not cls_token_at_end:
                tokens.append(cls_token)
                segment_ids.append(cls_token_segment_id)
                p_mask.append(0)
                cls_index = 0

            # Query
            for token in query_tokens:
                tokens.append(token)
                segment_ids.append(sequence_a_segment_id)
                p_mask.append(1)

            # SEP token
            tokens.append(sep_token)
            segment_ids.append(sequence_a_segment_id)
            p_mask.append(1)

            # Paragraph
            for i in range(doc_span.length):
                split_token_index = doc_span.start + i
                token_to_orig_map[len(
                    tokens)] = tok_to_orig_index[split_token_index]

                is_max_context = _check_is_max_context(doc_spans, doc_span_index,
                                                       split_token_index)
                token_is_max_context[len(tokens)] = is_max_context
                tokens.append(all_doc_tokens[split_token_index])
                segment_ids.append(sequence_b_segment_id)
                p_mask.append(0)
            paragraph_len = doc_span.length

            # SEP token
            tokens.append(sep_token)
            segment_ids.append(sequence_b_segment_id)
            p_mask.append(1)

            # CLS token at the end
            if cls_token_at_end:
                tokens.append(cls_token)
                segment_ids.append(cls_token_segment_id)
                p_mask.append(0)
                cls_index = len(tokens) - 1  # Index of classification token

            input_ids = tokenizer.convert_tokens_to_ids(tokens)

            # The mask has 1 for real tokens and 0 for padding tokens. Only real
            # tokens are attended to.
            input_mask = [1 if mask_padding_with_zero else 0] * len(input_ids)

            # Zero-pad up to the sequence length.
            while len(input_ids) < max_seq_length:
                input_ids.append(pad_token)
                input_mask.append(0 if mask_padding_with_zero else 1)
                segment_ids.append(pad_token_segment_id)
                p_mask.append(1)

            assert len(input_ids) == max_seq_length
            assert len(input_mask) == max_seq_length
            assert len(segment_ids) == max_seq_length

            span_is_impossible = example.is_impossible
            start_position = None
            end_position = None
            switch = None

            if is_training and not span_is_impossible:
                doc_start = doc_span.start
                doc_end = doc_span.start + doc_span.length - 1
                out_of_span = False
                if not (tok_start_position >= doc_start and
                        tok_end_position <= doc_end):
                    out_of_span = True
                if out_of_span:
                    start_position = 0
                    end_position = 0
                    span_is_impossible = True
                else:
                    doc_offset = len(query_tokens) + 2
                    start_position = tok_start_position - doc_start + doc_offset
                    end_position = tok_end_position - doc_start + doc_offset

            if is_training and span_is_impossible:
                start_position = cls_index
                end_position = cls_index
                switch = 1
            elif is_training and not span_is_impossible:
                switch = 0

            # The questions whose ``is_impossible'' are originally True should
            # be 1. Change switch to 2 or 3 if the answer is yes/no.
            if example.is_impossible is True:
                if example.orig_answer_text == "yes":
                    switch = 2
                elif example.orig_answer_text == "no":
                    switch = 3
                else:
                    switch = 1

            if example_index < 20:
                logger.info("*** Example ***")
                logger.info("unique_id: %s" % (unique_id))
                logger.info("example_index: %s" % (example_index))
                logger.info("doc_span_index: %s" % (doc_span_index))
                logger.info("tokens: %s" % " ".join(tokens))
                logger.info("token_to_orig_map: %s" % " ".join([
                    "%d:%d" % (x, y) for (x, y) in token_to_orig_map.items()]))
                logger.info("token_is_max_context: %s" % " ".join([
                    "%d:%s" % (x, y) for (x, y) in token_is_max_context.items()
                ]))
                logger.info("input_ids: %s" %
                            " ".join([str(x) for x in input_ids]))
                logger.info(
                    "input_mask: %s" % " ".join([str(x) for x in input_mask]))
                logger.info(
                    "segment_ids: %s" % " ".join([str(x) for x in segment_ids]))
                if is_training and span_is_impossible:
                    logger.info("impossible example")
                if is_training and not span_is_impossible:
                    answer_text = " ".join(
                        tokens[start_position:(end_position + 1)])
                    logger.info("start_position: %d" % (start_position))
                    logger.info("end_position: %d" % (end_position))
                    logger.info(
                        "answer: %s" % (answer_text))

            features.append(
                InputFeatures(
                    unique_id=unique_id,
                    example_index=example_index,
                    doc_span_index=doc_span_index,
                    tokens=tokens,
                    token_to_orig_map=token_to_orig_map,
                    token_is_max_context=token_is_max_context,
                    input_ids=input_ids,
                    input_mask=input_mask,
                    segment_ids=segment_ids,
                    cls_index=cls_index,
                    p_mask=p_mask,
                    paragraph_len=paragraph_len,
                    start_position=start_position,
                    end_position=end_position,
                    switch=switch,
                    is_impossible=span_is_impossible))
            unique_id += 1

    return features


Returns tokenized answer spans that better match the annotated answer .

def _improve_answer_span(doc_tokens, input_start, input_end, tokenizer,
                         orig_answer_text):
    """Returns tokenized answer spans that better match the annotated answer."""

    # The SQuAD annotations are character based. We first project them to
    # whitespace-tokenized words. But then after WordPiece tokenization, we can
    # often find a "better match". For example:
    #
    #   Question: What year was John Smith born?
    #   Context: The leader was John Smith (1895-1943).
    #   Answer: 1895
    #
    # The original whitespace-tokenized answer will be "(1895-1943).". However
    # after tokenization, our tokens will be "( 1895 - 1943 ) .". So we can match
    # the exact answer, 1895.
    #
    # However, this is not always possible. Consider the following:
    #
    #   Question: What country is the top exporter of electornics?
    #   Context: The Japanese electronics industry is the lagest in the world.
    #   Answer: Japan
    #
    # In this case, the annotator chose "Japan" as a character sub-span of
    # the word "Japanese". Since our WordPiece tokenizer does not split
    # "Japanese", we just use "Japanese" as the annotation. This is fairly rare
    # in SQuAD, but does happen.
    tok_answer_text = " ".join(tokenizer.tokenize(orig_answer_text))

    for new_start in range(input_start, input_end + 1):
        for new_end in range(input_end, new_start - 1, -1):
            text_span = " ".join(doc_tokens[new_start:(new_end + 1)])
            if text_span == tok_answer_text:
                return (new_start, new_end)

    return (input_start, input_end)


Check if this is the max context doc span for the token .

def _check_is_max_context(doc_spans, cur_span_index, position):
    """Check if this is the 'max context' doc span for the token."""

    # Because of the sliding window approach taken to scoring documents, a single
    # token can appear in multiple documents. E.g.
    #  Doc: the man went to the store and bought a gallon of milk
    #  Span A: the man went to the
    #  Span B: to the store and bought
    #  Span C: and bought a gallon of
    #  ...
    #
    # Now the word 'bought' will have two scores from spans B and C. We only
    # want to consider the score with "maximum context", which we define as
    # the *minimum* of its left and right context (the *sum* of left and
    # right context will always be the same, of course).
    #
    # In the example the maximum context for 'bought' would be span C since
    # it has 1 left context and 3 right context, while span B has 4 left context
    # and 0 right context.
    best_score = None
    best_span_index = None
    for (span_index, doc_span) in enumerate(doc_spans):
        end = doc_span.start + doc_span.length - 1
        if position < doc_span.start:
            continue
        if position > end:
            continue
        num_left_context = position - doc_span.start
        num_right_context = end - position
        score = min(num_left_context, num_right_context) + \
            0.01 * doc_span.length
        if best_score is None or score > best_score:
            best_score = score
            best_span_index = span_index

    return cur_span_index == best_span_index


Write final predictions to the json file and log - odds of null if needed .

def write_predictions_yes_no_no_empty_answer(all_examples, all_features, all_results, n_best_size,
                                             max_answer_length, do_lower_case, output_prediction_file,
                                             output_nbest_file, output_null_log_odds_file, verbose_logging,
                                             version_2_with_negative, null_score_diff_threshold, no_masking=False):
    """Write final predictions to the json file and log-odds of null if needed."""
    logger.info("Writing predictions to: %s" % (output_prediction_file))
    logger.info("Writing nbest to: %s" % (output_nbest_file))

    example_index_to_features = collections.defaultdict(list)
    for feature in all_features:
        example_index_to_features[feature.example_index].append(feature)

    unique_id_to_result = {}
    for result in all_results:
        unique_id_to_result[result.unique_id] = result

    _PrelimPrediction = collections.namedtuple(  # pylint: disable=invalid-name
        "PrelimPrediction",
        ["feature_index", "start_index", "end_index", "start_logit", "end_logit"])

    all_predictions = collections.OrderedDict()
    all_nbest_json = collections.OrderedDict()

    for (example_index, example) in enumerate(all_examples):
        features = example_index_to_features[example_index]

        prelim_predictions = []
        score_null = 1000000  # large and positive
        min_null_feature_index = 0  # the paragraph slice with min null score
        null_start_logit = 0  # the start logit at the slice with min null score
        null_end_logit = 0  # the end logit at the slice with min null score

        for (feature_index, feature) in enumerate(features):
            result = unique_id_to_result[feature.unique_id]
            start_indexes = _get_best_indexes(result.start_logits, n_best_size)
            end_indexes = _get_best_indexes(result.end_logits, n_best_size)
            if no_masking is True:
                feature_null_score = result.start_logits[0] + \
                    result.end_logits[0]
                if feature_null_score < score_null:
                    score_null = feature_null_score
                    min_null_feature_index = feature_index
                    null_start_logit = result.start_logits[0]
                    null_end_logit = result.end_logits[0]
            # if we could have irrelevant answers, get the min score of
            # irrelevant
            for start_index in start_indexes:
                for end_index in end_indexes:
                    # We could hypothetically create invalid predictions, e.g., predict
                    # that the start of the span is in the question. We throw out all
                    # invalid predictions.
                    if start_index >= len(feature.tokens):
                        continue
                    if end_index >= len(feature.tokens):
                        continue
                    if start_index not in feature.token_to_orig_map:
                        continue
                    if end_index not in feature.token_to_orig_map:
                        continue
                    if not feature.token_is_max_context.get(start_index, False):
                        continue
                    if end_index < start_index:
                        continue
                    length = end_index - start_index + 1
                    if length > max_answer_length:
                        continue
                    prelim_predictions.append(
                        _PrelimPrediction(
                            feature_index=feature_index,
                            start_index=start_index,
                            end_index=end_index,
                            start_logit=result.start_logits[start_index],
                            end_logit=result.end_logits[end_index]))
        if no_masking is True:
            prelim_predictions.append(
                _PrelimPrediction(
                    feature_index=min_null_feature_index,
                    start_index=0,
                    end_index=0,
                    start_logit=null_start_logit,
                    end_logit=null_end_logit))
        prelim_predictions = sorted(
            prelim_predictions,
            key=lambda x: (x.start_logit + x.end_logit),
            reverse=True)

        _NbestPrediction = collections.namedtuple(  # pylint: disable=invalid-name
            "NbestPrediction", ["text", "start_logit", "end_logit", "no_answer_logit", "switch", "switch_logits"])
        no_answer_logit = result.switch_logits[1]

        seen_predictions = {}
        nbest = []
        for pred in prelim_predictions:
            if len(nbest) >= n_best_size:
                break
            feature = features[pred.feature_index]
            if pred.start_index > 0:  # this is a non-null prediction
                tok_tokens = feature.tokens[pred.start_index:(
                    pred.end_index + 1)]
                orig_doc_start = feature.token_to_orig_map[pred.start_index]
                orig_doc_end = feature.token_to_orig_map[pred.end_index]
                orig_tokens = example.doc_tokens[orig_doc_start:(
                    orig_doc_end + 1)]
                tok_text = " ".join(tok_tokens)

                # De-tokenize WordPieces that have been split off.
                tok_text = tok_text.replace(" ##", "")
                tok_text = tok_text.replace("##", "")

                # Clean whitespace
                tok_text = tok_text.strip()
                tok_text = " ".join(tok_text.split())
                orig_text = " ".join(orig_tokens)

                final_text = get_final_text(
                    tok_text, orig_text, do_lower_case, verbose_logging)
                if final_text in seen_predictions:
                    continue

                seen_predictions[final_text] = True
            else:
                final_text = ""
                seen_predictions[final_text] = True

            nbest.append(
                _NbestPrediction(
                    text=final_text,
                    start_logit=pred.start_logit,
                    end_logit=pred.end_logit,
                    no_answer_logit=no_answer_logit,
                    switch=np.argmax(result.switch_logits),
                    switch_logits=result.switch_logits
                ))
        # if we didn't include the empty option in the n-best, include it
        if no_masking is True:
            if "" not in seen_predictions:
                nbest.append(
                    _NbestPrediction(
                        text="",
                        start_logit=null_start_logit,
                        end_logit=null_end_logit,
                        no_answer_logit=no_answer_logit,
                        switch=np.argmax(result.switch_logits),
                        switch_logits=result.switch_logits
                    ))

            # In very rare edge cases we could only have single null prediction.
            # So we just create a nonce prediction in this case to avoid
            # failure.
            if no_masking is True:
                if len(nbest) == 1:
                    nbest.insert(0,
                                 _NbestPrediction(text="", start_logit=0.0, end_logit=0.0, no_answer_logit=1.0, switch=1, switch_logits=[0.0, 0.0, 0.0, 0.0]))
            else:
                if len(nbest) == 0:
                    nbest.insert(0,
                                 _NbestPrediction(text="", start_logit=0.0, end_logit=0.0, no_answer_logit=1.0, switch=1, switch_logits=[0.0, 0.0, 0.0, 0.0]))

        # In very rare edge cases we could have no valid predictions. So we
        # just create a nonce prediction in this case to avoid failure.
        if not nbest:
            nbest.append(
                _NbestPrediction(text="", start_logit=0.0, end_logit=0.0, no_answer_logit=1.0, switch=1, switch_logits=[0.0, 0.0, 0.0, 0.0]))

        assert len(nbest) >= 1

        total_scores = []
        for entry in nbest:
            total_scores.append(entry.start_logit + entry.end_logit)

        probs = _compute_softmax(total_scores)

        nbest_json = []
        for (i, entry) in enumerate(nbest):
            output = collections.OrderedDict()
            output["text"] = entry.text
            output["probability"] = probs[i]
            output["start_logit"] = entry.start_logit
            output["end_logit"] = entry.end_logit
            output["no_answer_prob"] = entry.no_answer_logit
            output["switch"] = entry.switch
            output["switch_scores"] = entry.switch_logits
            nbest_json.append(output)

        assert len(nbest_json) >= 1

        # if the n-best is high enough, pick up no answer.
        possible_answers = np.argsort(
            nbest_json[0]["switch_scores"])[::-1]
        if possible_answers[0] == 1:
            all_predictions[example.qas_id] = switch_answers(
                possible_answers[1], nbest_json[0]["text"])
        else:
            all_predictions[example.qas_id] = switch_answers(
                possible_answers[0], nbest_json[0]["text"])

        all_nbest_json[example.qas_id] = nbest_json

    with open(output_prediction_file, "w") as writer:
        writer.write(json.dumps(all_predictions, indent=4) + "\n")