Spaces:
No application file
No application file
| # coding=utf-8 | |
| # Copyright 2018 The Google AI Team Authors. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # Lint as: python2, python3 | |
| """Utility functions for SQuAD v1.1/v2.0 datasets.""" | |
| from __future__ import absolute_import | |
| from __future__ import division | |
| # from __future__ import google_type_annotations | |
| from __future__ import print_function | |
| import collections | |
| import json | |
| import math | |
| import re | |
| import string | |
| import sys | |
| from albert import fine_tuning_utils | |
| from albert import modeling | |
| from albert import optimization | |
| from albert import tokenization | |
| import numpy as np | |
| import six | |
| from six.moves import map | |
| from six.moves import range | |
| import tensorflow.compat.v1 as tf | |
| from tensorflow.contrib import data as contrib_data | |
| from tensorflow.contrib import layers as contrib_layers | |
| from tensorflow.contrib import tpu as contrib_tpu | |
| _PrelimPrediction = collections.namedtuple( # pylint: disable=invalid-name | |
| "PrelimPrediction", | |
| ["feature_index", "start_index", "end_index", | |
| "start_log_prob", "end_log_prob"]) | |
| _NbestPrediction = collections.namedtuple( # pylint: disable=invalid-name | |
| "NbestPrediction", ["text", "start_log_prob", "end_log_prob"]) | |
| RawResult = collections.namedtuple("RawResult", | |
| ["unique_id", | |
| "start_log_prob", | |
| "end_log_prob"]) | |
| RawResultV2 = collections.namedtuple( | |
| "RawResultV2", | |
| ["unique_id", "start_top_log_probs", "start_top_index", | |
| "end_top_log_probs", "end_top_index", "cls_logits"]) | |
| class SquadExample(object): | |
| """A single training/test example for simple sequence classification. | |
| For examples without an answer, the start and end position are -1. | |
| """ | |
| def __init__(self, | |
| qas_id, | |
| question_text, | |
| paragraph_text, | |
| orig_answer_text=None, | |
| start_position=None, | |
| end_position=None, | |
| is_impossible=False): | |
| self.qas_id = qas_id | |
| self.question_text = question_text | |
| self.paragraph_text = paragraph_text | |
| self.orig_answer_text = orig_answer_text | |
| self.start_position = start_position | |
| self.end_position = end_position | |
| self.is_impossible = is_impossible | |
| def __str__(self): | |
| return self.__repr__() | |
| def __repr__(self): | |
| s = "" | |
| s += "qas_id: %s" % (tokenization.printable_text(self.qas_id)) | |
| s += ", question_text: %s" % ( | |
| tokenization.printable_text(self.question_text)) | |
| s += ", paragraph_text: [%s]" % (" ".join(self.paragraph_text)) | |
| if self.start_position: | |
| s += ", start_position: %d" % (self.start_position) | |
| if self.start_position: | |
| s += ", end_position: %d" % (self.end_position) | |
| if self.start_position: | |
| s += ", is_impossible: %r" % (self.is_impossible) | |
| return s | |
| class InputFeatures(object): | |
| """A single set of features of data.""" | |
| def __init__(self, | |
| unique_id, | |
| example_index, | |
| doc_span_index, | |
| tok_start_to_orig_index, | |
| tok_end_to_orig_index, | |
| token_is_max_context, | |
| tokens, | |
| input_ids, | |
| input_mask, | |
| segment_ids, | |
| paragraph_len, | |
| p_mask=None, | |
| start_position=None, | |
| end_position=None, | |
| is_impossible=None): | |
| self.unique_id = unique_id | |
| self.example_index = example_index | |
| self.doc_span_index = doc_span_index | |
| self.tok_start_to_orig_index = tok_start_to_orig_index | |
| self.tok_end_to_orig_index = tok_end_to_orig_index | |
| self.token_is_max_context = token_is_max_context | |
| self.tokens = tokens | |
| self.input_ids = input_ids | |
| self.input_mask = input_mask | |
| self.segment_ids = segment_ids | |
| self.paragraph_len = paragraph_len | |
| self.start_position = start_position | |
| self.end_position = end_position | |
| self.is_impossible = is_impossible | |
| self.p_mask = p_mask | |
| def read_squad_examples(input_file, is_training): | |
| """Read a SQuAD json file into a list of SquadExample.""" | |
| with tf.gfile.Open(input_file, "r") as reader: | |
| input_data = json.load(reader)["data"] | |
| examples = [] | |
| for entry in input_data: | |
| for paragraph in entry["paragraphs"]: | |
| paragraph_text = paragraph["context"] | |
| for qa in paragraph["qas"]: | |
| qas_id = qa["id"] | |
| question_text = qa["question"] | |
| start_position = None | |
| orig_answer_text = None | |
| is_impossible = False | |
| if is_training: | |
| is_impossible = qa.get("is_impossible", False) | |
| if (len(qa["answers"]) != 1) and (not is_impossible): | |
| raise ValueError( | |
| "For training, each question should have exactly 1 answer.") | |
| if not is_impossible: | |
| answer = qa["answers"][0] | |
| orig_answer_text = answer["text"] | |
| start_position = answer["answer_start"] | |
| else: | |
| start_position = -1 | |
| orig_answer_text = "" | |
| example = SquadExample( | |
| qas_id=qas_id, | |
| question_text=question_text, | |
| paragraph_text=paragraph_text, | |
| orig_answer_text=orig_answer_text, | |
| start_position=start_position, | |
| is_impossible=is_impossible) | |
| examples.append(example) | |
| return examples | |
| def _convert_index(index, pos, m=None, is_start=True): | |
| """Converts index.""" | |
| if index[pos] is not None: | |
| return index[pos] | |
| n = len(index) | |
| rear = pos | |
| while rear < n - 1 and index[rear] is None: | |
| rear += 1 | |
| front = pos | |
| while front > 0 and index[front] is None: | |
| front -= 1 | |
| assert index[front] is not None or index[rear] is not None | |
| if index[front] is None: | |
| if index[rear] >= 1: | |
| if is_start: | |
| return 0 | |
| else: | |
| return index[rear] - 1 | |
| return index[rear] | |
| if index[rear] is None: | |
| if m is not None and index[front] < m - 1: | |
| if is_start: | |
| return index[front] + 1 | |
| else: | |
| return m - 1 | |
| return index[front] | |
| if is_start: | |
| if index[rear] > index[front] + 1: | |
| return index[front] + 1 | |
| else: | |
| return index[rear] | |
| else: | |
| if index[rear] > index[front] + 1: | |
| return index[rear] - 1 | |
| else: | |
| return index[front] | |
| def convert_examples_to_features(examples, tokenizer, max_seq_length, | |
| doc_stride, max_query_length, is_training, | |
| output_fn, do_lower_case): | |
| """Loads a data file into a list of `InputBatch`s.""" | |
| cnt_pos, cnt_neg = 0, 0 | |
| unique_id = 1000000000 | |
| max_n, max_m = 1024, 1024 | |
| f = np.zeros((max_n, max_m), dtype=np.float32) | |
| for (example_index, example) in enumerate(examples): | |
| if example_index % 100 == 0: | |
| tf.logging.info("Converting {}/{} pos {} neg {}".format( | |
| example_index, len(examples), cnt_pos, cnt_neg)) | |
| query_tokens = tokenization.encode_ids( | |
| tokenizer.sp_model, | |
| tokenization.preprocess_text( | |
| example.question_text, lower=do_lower_case)) | |
| if len(query_tokens) > max_query_length: | |
| query_tokens = query_tokens[0:max_query_length] | |
| paragraph_text = example.paragraph_text | |
| para_tokens = tokenization.encode_pieces( | |
| tokenizer.sp_model, | |
| tokenization.preprocess_text( | |
| example.paragraph_text, lower=do_lower_case), | |
| return_unicode=False) | |
| chartok_to_tok_index = [] | |
| tok_start_to_chartok_index = [] | |
| tok_end_to_chartok_index = [] | |
| char_cnt = 0 | |
| para_tokens = [six.ensure_text(token, "utf-8") for token in para_tokens] | |
| for i, token in enumerate(para_tokens): | |
| new_token = six.ensure_text(token).replace( | |
| tokenization.SPIECE_UNDERLINE.decode("utf-8"), " ") | |
| chartok_to_tok_index.extend([i] * len(new_token)) | |
| tok_start_to_chartok_index.append(char_cnt) | |
| char_cnt += len(new_token) | |
| tok_end_to_chartok_index.append(char_cnt - 1) | |
| tok_cat_text = "".join(para_tokens).replace( | |
| tokenization.SPIECE_UNDERLINE.decode("utf-8"), " ") | |
| n, m = len(paragraph_text), len(tok_cat_text) | |
| if n > max_n or m > max_m: | |
| max_n = max(n, max_n) | |
| max_m = max(m, max_m) | |
| f = np.zeros((max_n, max_m), dtype=np.float32) | |
| g = {} | |
| def _lcs_match(max_dist, n=n, m=m): | |
| """Longest-common-substring algorithm.""" | |
| f.fill(0) | |
| g.clear() | |
| ### longest common sub sequence | |
| # f[i, j] = max(f[i - 1, j], f[i, j - 1], f[i - 1, j - 1] + match(i, j)) | |
| for i in range(n): | |
| # note(zhiliny): | |
| # unlike standard LCS, this is specifically optimized for the setting | |
| # because the mismatch between sentence pieces and original text will | |
| # be small | |
| for j in range(i - max_dist, i + max_dist): | |
| if j >= m or j < 0: continue | |
| if i > 0: | |
| g[(i, j)] = 0 | |
| f[i, j] = f[i - 1, j] | |
| if j > 0 and f[i, j - 1] > f[i, j]: | |
| g[(i, j)] = 1 | |
| f[i, j] = f[i, j - 1] | |
| f_prev = f[i - 1, j - 1] if i > 0 and j > 0 else 0 | |
| if (tokenization.preprocess_text( | |
| paragraph_text[i], lower=do_lower_case, | |
| remove_space=False) == tok_cat_text[j] | |
| and f_prev + 1 > f[i, j]): | |
| g[(i, j)] = 2 | |
| f[i, j] = f_prev + 1 | |
| max_dist = abs(n - m) + 5 | |
| for _ in range(2): | |
| _lcs_match(max_dist) | |
| if f[n - 1, m - 1] > 0.8 * n: break | |
| max_dist *= 2 | |
| orig_to_chartok_index = [None] * n | |
| chartok_to_orig_index = [None] * m | |
| i, j = n - 1, m - 1 | |
| while i >= 0 and j >= 0: | |
| if (i, j) not in g: break | |
| if g[(i, j)] == 2: | |
| orig_to_chartok_index[i] = j | |
| chartok_to_orig_index[j] = i | |
| i, j = i - 1, j - 1 | |
| elif g[(i, j)] == 1: | |
| j = j - 1 | |
| else: | |
| i = i - 1 | |
| if (all(v is None for v in orig_to_chartok_index) or | |
| f[n - 1, m - 1] < 0.8 * n): | |
| tf.logging.info("MISMATCH DETECTED!") | |
| continue | |
| tok_start_to_orig_index = [] | |
| tok_end_to_orig_index = [] | |
| for i in range(len(para_tokens)): | |
| start_chartok_pos = tok_start_to_chartok_index[i] | |
| end_chartok_pos = tok_end_to_chartok_index[i] | |
| start_orig_pos = _convert_index(chartok_to_orig_index, start_chartok_pos, | |
| n, is_start=True) | |
| end_orig_pos = _convert_index(chartok_to_orig_index, end_chartok_pos, | |
| n, is_start=False) | |
| tok_start_to_orig_index.append(start_orig_pos) | |
| tok_end_to_orig_index.append(end_orig_pos) | |
| if not is_training: | |
| tok_start_position = tok_end_position = None | |
| if is_training and example.is_impossible: | |
| tok_start_position = 0 | |
| tok_end_position = 0 | |
| if is_training and not example.is_impossible: | |
| start_position = example.start_position | |
| end_position = start_position + len(example.orig_answer_text) - 1 | |
| start_chartok_pos = _convert_index(orig_to_chartok_index, start_position, | |
| is_start=True) | |
| tok_start_position = chartok_to_tok_index[start_chartok_pos] | |
| end_chartok_pos = _convert_index(orig_to_chartok_index, end_position, | |
| is_start=False) | |
| tok_end_position = chartok_to_tok_index[end_chartok_pos] | |
| assert tok_start_position <= tok_end_position | |
| def _piece_to_id(x): | |
| if six.PY2 and isinstance(x, six.text_type): | |
| x = six.ensure_binary(x, "utf-8") | |
| return tokenizer.sp_model.PieceToId(x) | |
| all_doc_tokens = list(map(_piece_to_id, para_tokens)) | |
| # 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_is_max_context = {} | |
| segment_ids = [] | |
| p_mask = [] | |
| cur_tok_start_to_orig_index = [] | |
| cur_tok_end_to_orig_index = [] | |
| tokens.append(tokenizer.sp_model.PieceToId("[CLS]")) | |
| segment_ids.append(0) | |
| p_mask.append(0) | |
| for token in query_tokens: | |
| tokens.append(token) | |
| segment_ids.append(0) | |
| p_mask.append(1) | |
| tokens.append(tokenizer.sp_model.PieceToId("[SEP]")) | |
| segment_ids.append(0) | |
| p_mask.append(1) | |
| for i in range(doc_span.length): | |
| split_token_index = doc_span.start + i | |
| cur_tok_start_to_orig_index.append( | |
| tok_start_to_orig_index[split_token_index]) | |
| cur_tok_end_to_orig_index.append( | |
| tok_end_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(1) | |
| p_mask.append(0) | |
| tokens.append(tokenizer.sp_model.PieceToId("[SEP]")) | |
| segment_ids.append(1) | |
| p_mask.append(1) | |
| paragraph_len = len(tokens) | |
| input_ids = tokens | |
| # The mask has 1 for real tokens and 0 for padding tokens. Only real | |
| # tokens are attended to. | |
| input_mask = [1] * len(input_ids) | |
| # Zero-pad up to the sequence length. | |
| while len(input_ids) < max_seq_length: | |
| input_ids.append(0) | |
| input_mask.append(0) | |
| segment_ids.append(0) | |
| 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 | |
| if is_training and not span_is_impossible: | |
| # For training, if our document chunk does not contain an annotation | |
| # we throw it out, since there is nothing to predict. | |
| 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: | |
| # continue | |
| 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 = 0 | |
| end_position = 0 | |
| if example_index < 20: | |
| tf.logging.info("*** Example ***") | |
| tf.logging.info("unique_id: %s" % (unique_id)) | |
| tf.logging.info("example_index: %s" % (example_index)) | |
| tf.logging.info("doc_span_index: %s" % (doc_span_index)) | |
| tf.logging.info("tok_start_to_orig_index: %s" % " ".join( | |
| [str(x) for x in cur_tok_start_to_orig_index])) | |
| tf.logging.info("tok_end_to_orig_index: %s" % " ".join( | |
| [str(x) for x in cur_tok_end_to_orig_index])) | |
| tf.logging.info("token_is_max_context: %s" % " ".join([ | |
| "%d:%s" % (x, y) for (x, y) in six.iteritems(token_is_max_context) | |
| ])) | |
| tf.logging.info("input_pieces: %s" % " ".join( | |
| [tokenizer.sp_model.IdToPiece(x) for x in tokens])) | |
| tf.logging.info("input_ids: %s" % " ".join([str(x) for x in input_ids])) | |
| tf.logging.info( | |
| "input_mask: %s" % " ".join([str(x) for x in input_mask])) | |
| tf.logging.info( | |
| "segment_ids: %s" % " ".join([str(x) for x in segment_ids])) | |
| if is_training and span_is_impossible: | |
| tf.logging.info("impossible example span") | |
| if is_training and not span_is_impossible: | |
| pieces = [tokenizer.sp_model.IdToPiece(token) for token in | |
| tokens[start_position: (end_position + 1)]] | |
| answer_text = tokenizer.sp_model.DecodePieces(pieces) | |
| tf.logging.info("start_position: %d" % (start_position)) | |
| tf.logging.info("end_position: %d" % (end_position)) | |
| tf.logging.info( | |
| "answer: %s" % (tokenization.printable_text(answer_text))) | |
| # note(zhiliny): With multi processing, | |
| # the example_index is actually the index within the current process | |
| # therefore we use example_index=None to avoid being used in the future. | |
| # The current code does not use example_index of training data. | |
| if is_training: | |
| feat_example_index = None | |
| else: | |
| feat_example_index = example_index | |
| feature = InputFeatures( | |
| unique_id=unique_id, | |
| example_index=feat_example_index, | |
| doc_span_index=doc_span_index, | |
| tok_start_to_orig_index=cur_tok_start_to_orig_index, | |
| tok_end_to_orig_index=cur_tok_end_to_orig_index, | |
| token_is_max_context=token_is_max_context, | |
| tokens=[tokenizer.sp_model.IdToPiece(x) for x in tokens], | |
| input_ids=input_ids, | |
| input_mask=input_mask, | |
| segment_ids=segment_ids, | |
| paragraph_len=paragraph_len, | |
| start_position=start_position, | |
| end_position=end_position, | |
| is_impossible=span_is_impossible, | |
| p_mask=p_mask) | |
| # Run callback | |
| output_fn(feature) | |
| unique_id += 1 | |
| if span_is_impossible: | |
| cnt_neg += 1 | |
| else: | |
| cnt_pos += 1 | |
| tf.logging.info("Total number of instances: {} = pos {} neg {}".format( | |
| cnt_pos + cnt_neg, cnt_pos, cnt_neg)) | |
| 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 | |
| def _get_best_indexes(logits, n_best_size): | |
| """Get the n-best logits from a list.""" | |
| index_and_score = sorted(enumerate(logits), key=lambda x: x[1], reverse=True) | |
| best_indexes = [] | |
| for i in range(len(index_and_score)): | |
| if i >= n_best_size: | |
| break | |
| best_indexes.append(index_and_score[i][0]) | |
| return best_indexes | |
| def _compute_softmax(scores): | |
| """Compute softmax probability over raw logits.""" | |
| if not scores: | |
| return [] | |
| max_score = None | |
| for score in scores: | |
| if max_score is None or score > max_score: | |
| max_score = score | |
| exp_scores = [] | |
| total_sum = 0.0 | |
| for score in scores: | |
| x = math.exp(score - max_score) | |
| exp_scores.append(x) | |
| total_sum += x | |
| probs = [] | |
| for score in exp_scores: | |
| probs.append(score / total_sum) | |
| return probs | |
| class FeatureWriter(object): | |
| """Writes InputFeature to TF example file.""" | |
| def __init__(self, filename, is_training): | |
| self.filename = filename | |
| self.is_training = is_training | |
| self.num_features = 0 | |
| self._writer = tf.python_io.TFRecordWriter(filename) | |
| def process_feature(self, feature): | |
| """Write a InputFeature to the TFRecordWriter as a tf.train.Example.""" | |
| self.num_features += 1 | |
| def create_int_feature(values): | |
| feature = tf.train.Feature( | |
| int64_list=tf.train.Int64List(value=list(values))) | |
| return feature | |
| features = collections.OrderedDict() | |
| features["unique_ids"] = create_int_feature([feature.unique_id]) | |
| features["input_ids"] = create_int_feature(feature.input_ids) | |
| features["input_mask"] = create_int_feature(feature.input_mask) | |
| features["segment_ids"] = create_int_feature(feature.segment_ids) | |
| features["p_mask"] = create_int_feature(feature.p_mask) | |
| if self.is_training: | |
| features["start_positions"] = create_int_feature([feature.start_position]) | |
| features["end_positions"] = create_int_feature([feature.end_position]) | |
| impossible = 0 | |
| if feature.is_impossible: | |
| impossible = 1 | |
| features["is_impossible"] = create_int_feature([impossible]) | |
| tf_example = tf.train.Example(features=tf.train.Features(feature=features)) | |
| self._writer.write(tf_example.SerializeToString()) | |
| def close(self): | |
| self._writer.close() | |
| def input_fn_builder(input_file, seq_length, is_training, | |
| drop_remainder, use_tpu, bsz, is_v2): | |
| """Creates an `input_fn` closure to be passed to TPUEstimator.""" | |
| name_to_features = { | |
| "unique_ids": tf.FixedLenFeature([], tf.int64), | |
| "input_ids": tf.FixedLenFeature([seq_length], tf.int64), | |
| "input_mask": tf.FixedLenFeature([seq_length], tf.int64), | |
| "segment_ids": tf.FixedLenFeature([seq_length], tf.int64), | |
| } | |
| # p_mask is not required for SQuAD v1.1 | |
| if is_v2: | |
| name_to_features["p_mask"] = tf.FixedLenFeature([seq_length], tf.int64) | |
| if is_training: | |
| name_to_features["start_positions"] = tf.FixedLenFeature([], tf.int64) | |
| name_to_features["end_positions"] = tf.FixedLenFeature([], tf.int64) | |
| name_to_features["is_impossible"] = tf.FixedLenFeature([], tf.int64) | |
| def _decode_record(record, name_to_features): | |
| """Decodes a record to a TensorFlow example.""" | |
| example = tf.parse_single_example(record, name_to_features) | |
| # tf.Example only supports tf.int64, but the TPU only supports tf.int32. | |
| # So cast all int64 to int32. | |
| for name in list(example.keys()): | |
| t = example[name] | |
| if t.dtype == tf.int64: | |
| t = tf.to_int32(t) | |
| example[name] = t | |
| return example | |
| def input_fn(params): | |
| """The actual input function.""" | |
| if use_tpu: | |
| batch_size = params["batch_size"] | |
| else: | |
| batch_size = bsz | |
| # For training, we want a lot of parallel reading and shuffling. | |
| # For eval, we want no shuffling and parallel reading doesn't matter. | |
| d = tf.data.TFRecordDataset(input_file) | |
| if is_training: | |
| d = d.repeat() | |
| d = d.shuffle(buffer_size=100) | |
| d = d.apply( | |
| contrib_data.map_and_batch( | |
| lambda record: _decode_record(record, name_to_features), | |
| batch_size=batch_size, | |
| drop_remainder=drop_remainder)) | |
| return d | |
| return input_fn | |
| def create_v1_model(albert_config, is_training, input_ids, input_mask, | |
| segment_ids, use_one_hot_embeddings, use_einsum, | |
| hub_module): | |
| """Creates a classification model.""" | |
| (_, final_hidden) = fine_tuning_utils.create_albert( | |
| albert_config=albert_config, | |
| is_training=is_training, | |
| input_ids=input_ids, | |
| input_mask=input_mask, | |
| segment_ids=segment_ids, | |
| use_one_hot_embeddings=use_one_hot_embeddings, | |
| use_einsum=use_einsum, | |
| hub_module=hub_module) | |
| final_hidden_shape = modeling.get_shape_list(final_hidden, expected_rank=3) | |
| batch_size = final_hidden_shape[0] | |
| seq_length = final_hidden_shape[1] | |
| hidden_size = final_hidden_shape[2] | |
| output_weights = tf.get_variable( | |
| "cls/squad/output_weights", [2, hidden_size], | |
| initializer=tf.truncated_normal_initializer(stddev=0.02)) | |
| output_bias = tf.get_variable( | |
| "cls/squad/output_bias", [2], initializer=tf.zeros_initializer()) | |
| final_hidden_matrix = tf.reshape(final_hidden, | |
| [batch_size * seq_length, hidden_size]) | |
| logits = tf.matmul(final_hidden_matrix, output_weights, transpose_b=True) | |
| logits = tf.nn.bias_add(logits, output_bias) | |
| logits = tf.reshape(logits, [batch_size, seq_length, 2]) | |
| logits = tf.transpose(logits, [2, 0, 1]) | |
| unstacked_logits = tf.unstack(logits, axis=0) | |
| (start_logits, end_logits) = (unstacked_logits[0], unstacked_logits[1]) | |
| return (start_logits, end_logits) | |
| def v1_model_fn_builder(albert_config, init_checkpoint, learning_rate, | |
| num_train_steps, num_warmup_steps, use_tpu, | |
| use_one_hot_embeddings, use_einsum, hub_module): | |
| """Returns `model_fn` closure for TPUEstimator.""" | |
| def model_fn(features, labels, mode, params): # pylint: disable=unused-argument | |
| """The `model_fn` for TPUEstimator.""" | |
| tf.logging.info("*** Features ***") | |
| for name in sorted(features.keys()): | |
| tf.logging.info(" name = %s, shape = %s" % (name, features[name].shape)) | |
| if "unique_ids" in features: | |
| unique_ids = features["unique_ids"] | |
| else: | |
| unique_ids = None | |
| input_ids = features["input_ids"] | |
| input_mask = features["input_mask"] | |
| segment_ids = features["segment_ids"] | |
| is_training = (mode == tf.estimator.ModeKeys.TRAIN) | |
| (start_logits, end_logits) = create_v1_model( | |
| albert_config=albert_config, | |
| is_training=is_training, | |
| input_ids=input_ids, | |
| input_mask=input_mask, | |
| segment_ids=segment_ids, | |
| use_one_hot_embeddings=use_one_hot_embeddings, | |
| use_einsum=use_einsum, | |
| hub_module=hub_module) | |
| # Assign names to the logits so that we can refer to them as output tensors. | |
| start_logits = tf.identity(start_logits, name="start_logits") | |
| end_logits = tf.identity(end_logits, name="end_logits") | |
| tvars = tf.trainable_variables() | |
| initialized_variable_names = {} | |
| scaffold_fn = None | |
| if init_checkpoint: | |
| (assignment_map, initialized_variable_names | |
| ) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint) | |
| if use_tpu: | |
| def tpu_scaffold(): | |
| tf.train.init_from_checkpoint(init_checkpoint, assignment_map) | |
| return tf.train.Scaffold() | |
| scaffold_fn = tpu_scaffold | |
| else: | |
| tf.train.init_from_checkpoint(init_checkpoint, assignment_map) | |
| tf.logging.info("**** Trainable Variables ****") | |
| for var in tvars: | |
| init_string = "" | |
| if var.name in initialized_variable_names: | |
| init_string = ", *INIT_FROM_CKPT*" | |
| tf.logging.info(" name = %s, shape = %s%s", var.name, var.shape, | |
| init_string) | |
| output_spec = None | |
| if mode == tf.estimator.ModeKeys.TRAIN: | |
| seq_length = modeling.get_shape_list(input_ids)[1] | |
| def compute_loss(logits, positions): | |
| one_hot_positions = tf.one_hot( | |
| positions, depth=seq_length, dtype=tf.float32) | |
| log_probs = tf.nn.log_softmax(logits, axis=-1) | |
| loss = -tf.reduce_mean( | |
| tf.reduce_sum(one_hot_positions * log_probs, axis=-1)) | |
| return loss | |
| start_positions = features["start_positions"] | |
| end_positions = features["end_positions"] | |
| start_loss = compute_loss(start_logits, start_positions) | |
| end_loss = compute_loss(end_logits, end_positions) | |
| total_loss = (start_loss + end_loss) / 2.0 | |
| train_op = optimization.create_optimizer( | |
| total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu) | |
| output_spec = contrib_tpu.TPUEstimatorSpec( | |
| mode=mode, | |
| loss=total_loss, | |
| train_op=train_op, | |
| scaffold_fn=scaffold_fn) | |
| elif mode == tf.estimator.ModeKeys.PREDICT: | |
| predictions = { | |
| "start_log_prob": start_logits, | |
| "end_log_prob": end_logits, | |
| } | |
| if unique_ids is not None: | |
| predictions["unique_ids"] = unique_ids | |
| output_spec = contrib_tpu.TPUEstimatorSpec( | |
| mode=mode, predictions=predictions, scaffold_fn=scaffold_fn) | |
| else: | |
| raise ValueError( | |
| "Only TRAIN and PREDICT modes are supported: %s" % (mode)) | |
| return output_spec | |
| return model_fn | |
| def accumulate_predictions_v1(result_dict, all_examples, all_features, | |
| all_results, n_best_size, max_answer_length): | |
| """accumulate predictions for each positions in a dictionary.""" | |
| 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 | |
| all_predictions = collections.OrderedDict() | |
| all_nbest_json = collections.OrderedDict() | |
| scores_diff_json = collections.OrderedDict() | |
| for (example_index, example) in enumerate(all_examples): | |
| if example_index not in result_dict: | |
| result_dict[example_index] = {} | |
| features = example_index_to_features[example_index] | |
| prelim_predictions = [] | |
| min_null_feature_index = 0 # the paragraph slice with min mull 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): | |
| if feature.unique_id not in result_dict[example_index]: | |
| result_dict[example_index][feature.unique_id] = {} | |
| result = unique_id_to_result[feature.unique_id] | |
| start_indexes = _get_best_indexes(result.start_log_prob, n_best_size) | |
| end_indexes = _get_best_indexes(result.end_log_prob, n_best_size) | |
| for start_index in start_indexes: | |
| for end_index in end_indexes: | |
| doc_offset = feature.tokens.index("[SEP]") + 1 | |
| # 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 - doc_offset >= len(feature.tok_start_to_orig_index): | |
| continue | |
| if end_index - doc_offset >= len(feature.tok_end_to_orig_index): | |
| 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 | |
| start_log_prob = result.start_log_prob[start_index] | |
| end_log_prob = result.end_log_prob[end_index] | |
| start_idx = start_index - doc_offset | |
| end_idx = end_index - doc_offset | |
| if (start_idx, end_idx) not in result_dict[example_index][feature.unique_id]: | |
| result_dict[example_index][feature.unique_id][(start_idx, end_idx)] = [] | |
| result_dict[example_index][feature.unique_id][(start_idx, end_idx)].append((start_log_prob, end_log_prob)) | |
| def write_predictions_v1(result_dict, all_examples, all_features, | |
| all_results, n_best_size, max_answer_length, | |
| output_prediction_file, output_nbest_file): | |
| """Write final predictions to the json file and log-odds of null if needed.""" | |
| tf.logging.info("Writing predictions to: %s" % (output_prediction_file)) | |
| tf.logging.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 | |
| all_predictions = collections.OrderedDict() | |
| all_nbest_json = collections.OrderedDict() | |
| scores_diff_json = collections.OrderedDict() | |
| for (example_index, example) in enumerate(all_examples): | |
| features = example_index_to_features[example_index] | |
| prelim_predictions = [] | |
| # keep track of the minimum score of null start+end of position 0 | |
| score_null = 1000000 # large and positive | |
| min_null_feature_index = 0 # the paragraph slice with min mull 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): | |
| for ((start_idx, end_idx), logprobs) in \ | |
| result_dict[example_index][feature.unique_id].items(): | |
| start_log_prob = 0 | |
| end_log_prob = 0 | |
| for logprob in logprobs: | |
| start_log_prob += logprob[0] | |
| end_log_prob += logprob[1] | |
| prelim_predictions.append( | |
| _PrelimPrediction( | |
| feature_index=feature_index, | |
| start_index=start_idx, | |
| end_index=end_idx, | |
| start_log_prob=start_log_prob / len(logprobs), | |
| end_log_prob=end_log_prob / len(logprobs))) | |
| prelim_predictions = sorted( | |
| prelim_predictions, | |
| key=lambda x: (x.start_log_prob + x.end_log_prob), | |
| reverse=True) | |
| 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_start_to_orig_index = feature.tok_start_to_orig_index | |
| tok_end_to_orig_index = feature.tok_end_to_orig_index | |
| start_orig_pos = tok_start_to_orig_index[pred.start_index] | |
| end_orig_pos = tok_end_to_orig_index[pred.end_index] | |
| paragraph_text = example.paragraph_text | |
| final_text = paragraph_text[start_orig_pos: end_orig_pos + 1].strip() | |
| 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_log_prob=pred.start_log_prob, | |
| end_log_prob=pred.end_log_prob)) | |
| # 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="empty", start_log_prob=0.0, end_log_prob=0.0)) | |
| assert len(nbest) >= 1 | |
| total_scores = [] | |
| best_non_null_entry = None | |
| for entry in nbest: | |
| total_scores.append(entry.start_log_prob + entry.end_log_prob) | |
| if not best_non_null_entry: | |
| if entry.text: | |
| best_non_null_entry = entry | |
| 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_log_prob"] = entry.start_log_prob | |
| output["end_log_prob"] = entry.end_log_prob | |
| nbest_json.append(output) | |
| assert len(nbest_json) >= 1 | |
| all_predictions[example.qas_id] = nbest_json[0]["text"] | |
| all_nbest_json[example.qas_id] = nbest_json | |
| with tf.gfile.GFile(output_prediction_file, "w") as writer: | |
| writer.write(json.dumps(all_predictions, indent=4) + "\n") | |
| with tf.gfile.GFile(output_nbest_file, "w") as writer: | |
| writer.write(json.dumps(all_nbest_json, indent=4) + "\n") | |
| return all_predictions | |
| ####### following are from official SQuAD v1.1 evaluation scripts | |
| def normalize_answer_v1(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 f1_score(prediction, ground_truth): | |
| prediction_tokens = normalize_answer_v1(prediction).split() | |
| ground_truth_tokens = normalize_answer_v1(ground_truth).split() | |
| common = ( | |
| collections.Counter(prediction_tokens) | |
| & collections.Counter(ground_truth_tokens)) | |
| num_same = sum(common.values()) | |
| if num_same == 0: | |
| return 0 | |
| precision = 1.0 * num_same / len(prediction_tokens) | |
| recall = 1.0 * num_same / len(ground_truth_tokens) | |
| f1 = (2 * precision * recall) / (precision + recall) | |
| return f1 | |
| def exact_match_score(prediction, ground_truth): | |
| return (normalize_answer_v1(prediction) == normalize_answer_v1(ground_truth)) | |
| def metric_max_over_ground_truths(metric_fn, prediction, ground_truths): | |
| scores_for_ground_truths = [] | |
| for ground_truth in ground_truths: | |
| score = metric_fn(prediction, ground_truth) | |
| scores_for_ground_truths.append(score) | |
| return max(scores_for_ground_truths) | |
| def evaluate_v1(dataset, predictions): | |
| f1 = exact_match = total = 0 | |
| for article in dataset: | |
| for paragraph in article["paragraphs"]: | |
| for qa in paragraph["qas"]: | |
| total += 1 | |
| if qa["id"] not in predictions: | |
| message = ("Unanswered question " + six.ensure_str(qa["id"]) + | |
| " will receive score 0.") | |
| print(message, file=sys.stderr) | |
| continue | |
| ground_truths = [x["text"] for x in qa["answers"]] | |
| # ground_truths = list(map(lambda x: x["text"], qa["answers"])) | |
| prediction = predictions[qa["id"]] | |
| exact_match += metric_max_over_ground_truths(exact_match_score, | |
| prediction, ground_truths) | |
| f1 += metric_max_over_ground_truths(f1_score, prediction, ground_truths) | |
| exact_match = 100.0 * exact_match / total | |
| f1 = 100.0 * f1 / total | |
| return {"exact_match": exact_match, "f1": f1} | |
| ####### above are from official SQuAD v1.1 evaluation scripts | |
| ####### following are from official SQuAD v2.0 evaluation scripts | |
| def make_qid_to_has_ans(dataset): | |
| qid_to_has_ans = {} | |
| for article in dataset: | |
| for p in article['paragraphs']: | |
| for qa in p['qas']: | |
| qid_to_has_ans[qa['id']] = bool(qa['answers']) | |
| return qid_to_has_ans | |
| def normalize_answer_v2(s): | |
| """Lower text and remove punctuation, articles and extra whitespace.""" | |
| def remove_articles(text): | |
| regex = re.compile(r'\b(a|an|the)\b', re.UNICODE) | |
| return re.sub(regex, ' ', 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 get_tokens(s): | |
| if not s: return [] | |
| return normalize_answer_v2(s).split() | |
| def compute_exact(a_gold, a_pred): | |
| return int(normalize_answer_v2(a_gold) == normalize_answer_v2(a_pred)) | |
| def compute_f1(a_gold, a_pred): | |
| gold_toks = get_tokens(a_gold) | |
| pred_toks = get_tokens(a_pred) | |
| common = collections.Counter(gold_toks) & collections.Counter(pred_toks) | |
| num_same = sum(common.values()) | |
| if len(gold_toks) == 0 or len(pred_toks) == 0: | |
| # If either is no-answer, then F1 is 1 if they agree, 0 otherwise | |
| return int(gold_toks == pred_toks) | |
| if num_same == 0: | |
| return 0 | |
| precision = 1.0 * num_same / len(pred_toks) | |
| recall = 1.0 * num_same / len(gold_toks) | |
| f1 = (2 * precision * recall) / (precision + recall) | |
| return f1 | |
| def get_raw_scores(dataset, preds): | |
| exact_scores = {} | |
| f1_scores = {} | |
| for article in dataset: | |
| for p in article['paragraphs']: | |
| for qa in p['qas']: | |
| qid = qa['id'] | |
| gold_answers = [a['text'] for a in qa['answers'] | |
| if normalize_answer_v2(a['text'])] | |
| if not gold_answers: | |
| # For unanswerable questions, only correct answer is empty string | |
| gold_answers = [''] | |
| if qid not in preds: | |
| print('Missing prediction for %s' % qid) | |
| continue | |
| a_pred = preds[qid] | |
| # Take max over all gold answers | |
| exact_scores[qid] = max(compute_exact(a, a_pred) for a in gold_answers) | |
| f1_scores[qid] = max(compute_f1(a, a_pred) for a in gold_answers) | |
| return exact_scores, f1_scores | |
| def apply_no_ans_threshold(scores, na_probs, qid_to_has_ans, na_prob_thresh): | |
| new_scores = {} | |
| for qid, s in scores.items(): | |
| pred_na = na_probs[qid] > na_prob_thresh | |
| if pred_na: | |
| new_scores[qid] = float(not qid_to_has_ans[qid]) | |
| else: | |
| new_scores[qid] = s | |
| return new_scores | |
| def make_eval_dict(exact_scores, f1_scores, qid_list=None): | |
| if not qid_list: | |
| total = len(exact_scores) | |
| return collections.OrderedDict([ | |
| ('exact', 100.0 * sum(exact_scores.values()) / total), | |
| ('f1', 100.0 * sum(f1_scores.values()) / total), | |
| ('total', total), | |
| ]) | |
| else: | |
| total = len(qid_list) | |
| return collections.OrderedDict([ | |
| ('exact', 100.0 * sum(exact_scores[k] for k in qid_list) / total), | |
| ('f1', 100.0 * sum(f1_scores[k] for k in qid_list) / total), | |
| ('total', total), | |
| ]) | |
| def find_best_thresh(preds, scores, na_probs, qid_to_has_ans): | |
| num_no_ans = sum(1 for k in qid_to_has_ans if not qid_to_has_ans[k]) | |
| cur_score = num_no_ans | |
| best_score = cur_score | |
| best_thresh = 0.0 | |
| qid_list = sorted(na_probs, key=lambda k: na_probs[k]) | |
| for i, qid in enumerate(qid_list): | |
| if qid not in scores: continue | |
| if qid_to_has_ans[qid]: | |
| diff = scores[qid] | |
| else: | |
| if preds[qid]: | |
| diff = -1 | |
| else: | |
| diff = 0 | |
| cur_score += diff | |
| if cur_score > best_score: | |
| best_score = cur_score | |
| best_thresh = na_probs[qid] | |
| return 100.0 * best_score / len(scores), best_thresh | |
| def find_all_best_thresh(main_eval, preds, exact_raw, f1_raw, na_probs, qid_to_has_ans): | |
| best_exact, exact_thresh = find_best_thresh(preds, exact_raw, na_probs, qid_to_has_ans) | |
| best_f1, f1_thresh = find_best_thresh(preds, f1_raw, na_probs, qid_to_has_ans) | |
| main_eval['best_exact'] = best_exact | |
| main_eval['best_exact_thresh'] = exact_thresh | |
| main_eval['best_f1'] = best_f1 | |
| main_eval['best_f1_thresh'] = f1_thresh | |
| def merge_eval(main_eval, new_eval, prefix): | |
| for k in new_eval: | |
| main_eval['%s_%s' % (prefix, k)] = new_eval[k] | |
| ####### above are from official SQuAD v2.0 evaluation scripts | |
| def accumulate_predictions_v2(result_dict, cls_dict, all_examples, | |
| all_features, all_results, n_best_size, | |
| max_answer_length, start_n_top, end_n_top): | |
| """accumulate predictions for each positions in a dictionary.""" | |
| 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 | |
| all_predictions = collections.OrderedDict() | |
| all_nbest_json = collections.OrderedDict() | |
| scores_diff_json = collections.OrderedDict() | |
| for (example_index, example) in enumerate(all_examples): | |
| if example_index not in result_dict: | |
| result_dict[example_index] = {} | |
| features = example_index_to_features[example_index] | |
| prelim_predictions = [] | |
| # keep track of the minimum score of null start+end of position 0 | |
| score_null = 1000000 # large and positive | |
| for (feature_index, feature) in enumerate(features): | |
| if feature.unique_id not in result_dict[example_index]: | |
| result_dict[example_index][feature.unique_id] = {} | |
| result = unique_id_to_result[feature.unique_id] | |
| cur_null_score = result.cls_logits | |
| # if we could have irrelevant answers, get the min score of irrelevant | |
| score_null = min(score_null, cur_null_score) | |
| doc_offset = feature.tokens.index("[SEP]") + 1 | |
| for i in range(start_n_top): | |
| for j in range(end_n_top): | |
| start_log_prob = result.start_top_log_probs[i] | |
| start_index = result.start_top_index[i] | |
| j_index = i * end_n_top + j | |
| end_log_prob = result.end_top_log_probs[j_index] | |
| end_index = result.end_top_index[j_index] | |
| # 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 - doc_offset >= len(feature.tok_start_to_orig_index): | |
| continue | |
| if start_index - doc_offset < 0: | |
| continue | |
| if end_index - doc_offset >= len(feature.tok_end_to_orig_index): | |
| 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 | |
| start_idx = start_index - doc_offset | |
| end_idx = end_index - doc_offset | |
| if (start_idx, end_idx) not in result_dict[example_index][feature.unique_id]: | |
| result_dict[example_index][feature.unique_id][(start_idx, end_idx)] = [] | |
| result_dict[example_index][feature.unique_id][(start_idx, end_idx)].append((start_log_prob, end_log_prob)) | |
| if example_index not in cls_dict: | |
| cls_dict[example_index] = [] | |
| cls_dict[example_index].append(score_null) | |
| def write_predictions_v2(result_dict, cls_dict, all_examples, all_features, | |
| all_results, n_best_size, max_answer_length, | |
| output_prediction_file, | |
| output_nbest_file, output_null_log_odds_file, | |
| null_score_diff_threshold): | |
| """Write final predictions to the json file and log-odds of null if needed.""" | |
| tf.logging.info("Writing predictions to: %s" % (output_prediction_file)) | |
| tf.logging.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 | |
| all_predictions = collections.OrderedDict() | |
| all_nbest_json = collections.OrderedDict() | |
| scores_diff_json = collections.OrderedDict() | |
| for (example_index, example) in enumerate(all_examples): | |
| features = example_index_to_features[example_index] | |
| prelim_predictions = [] | |
| # keep track of the minimum score of null start+end of position 0 | |
| # score_null = 1000000 # large and positive | |
| for (feature_index, feature) in enumerate(features): | |
| for ((start_idx, end_idx), logprobs) in \ | |
| result_dict[example_index][feature.unique_id].items(): | |
| start_log_prob = 0 | |
| end_log_prob = 0 | |
| for logprob in logprobs: | |
| start_log_prob += logprob[0] | |
| end_log_prob += logprob[1] | |
| prelim_predictions.append( | |
| _PrelimPrediction( | |
| feature_index=feature_index, | |
| start_index=start_idx, | |
| end_index=end_idx, | |
| start_log_prob=start_log_prob / len(logprobs), | |
| end_log_prob=end_log_prob / len(logprobs))) | |
| prelim_predictions = sorted( | |
| prelim_predictions, | |
| key=lambda x: (x.start_log_prob + x.end_log_prob), | |
| reverse=True) | |
| seen_predictions = {} | |
| nbest = [] | |
| for pred in prelim_predictions: | |
| if len(nbest) >= n_best_size: | |
| break | |
| feature = features[pred.feature_index] | |
| tok_start_to_orig_index = feature.tok_start_to_orig_index | |
| tok_end_to_orig_index = feature.tok_end_to_orig_index | |
| start_orig_pos = tok_start_to_orig_index[pred.start_index] | |
| end_orig_pos = tok_end_to_orig_index[pred.end_index] | |
| paragraph_text = example.paragraph_text | |
| final_text = paragraph_text[start_orig_pos: end_orig_pos + 1].strip() | |
| if final_text in seen_predictions: | |
| continue | |
| seen_predictions[final_text] = True | |
| nbest.append( | |
| _NbestPrediction( | |
| text=final_text, | |
| start_log_prob=pred.start_log_prob, | |
| end_log_prob=pred.end_log_prob)) | |
| # 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_log_prob=-1e6, | |
| end_log_prob=-1e6)) | |
| total_scores = [] | |
| best_non_null_entry = None | |
| for entry in nbest: | |
| total_scores.append(entry.start_log_prob + entry.end_log_prob) | |
| if not best_non_null_entry: | |
| best_non_null_entry = entry | |
| 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_log_prob"] = entry.start_log_prob | |
| output["end_log_prob"] = entry.end_log_prob | |
| nbest_json.append(output) | |
| assert len(nbest_json) >= 1 | |
| assert best_non_null_entry is not None | |
| score_diff = sum(cls_dict[example_index]) / len(cls_dict[example_index]) | |
| scores_diff_json[example.qas_id] = score_diff | |
| # predict null answers when null threshold is provided | |
| if null_score_diff_threshold is None or score_diff < null_score_diff_threshold: | |
| all_predictions[example.qas_id] = best_non_null_entry.text | |
| else: | |
| all_predictions[example.qas_id] = "" | |
| all_nbest_json[example.qas_id] = nbest_json | |
| assert len(nbest_json) >= 1 | |
| with tf.gfile.GFile(output_prediction_file, "w") as writer: | |
| writer.write(json.dumps(all_predictions, indent=4) + "\n") | |
| with tf.gfile.GFile(output_nbest_file, "w") as writer: | |
| writer.write(json.dumps(all_nbest_json, indent=4) + "\n") | |
| with tf.gfile.GFile(output_null_log_odds_file, "w") as writer: | |
| writer.write(json.dumps(scores_diff_json, indent=4) + "\n") | |
| return all_predictions, scores_diff_json | |
| def create_v2_model(albert_config, is_training, input_ids, input_mask, | |
| segment_ids, use_one_hot_embeddings, features, | |
| max_seq_length, start_n_top, end_n_top, dropout_prob, | |
| hub_module): | |
| """Creates a classification model.""" | |
| (_, output) = fine_tuning_utils.create_albert( | |
| albert_config=albert_config, | |
| is_training=is_training, | |
| input_ids=input_ids, | |
| input_mask=input_mask, | |
| segment_ids=segment_ids, | |
| use_one_hot_embeddings=use_one_hot_embeddings, | |
| use_einsum=True, | |
| hub_module=hub_module) | |
| bsz = tf.shape(output)[0] | |
| return_dict = {} | |
| output = tf.transpose(output, [1, 0, 2]) | |
| # invalid position mask such as query and special symbols (PAD, SEP, CLS) | |
| p_mask = tf.cast(features["p_mask"], dtype=tf.float32) | |
| # logit of the start position | |
| with tf.variable_scope("start_logits"): | |
| start_logits = tf.layers.dense( | |
| output, | |
| 1, | |
| kernel_initializer=modeling.create_initializer( | |
| albert_config.initializer_range)) | |
| start_logits = tf.transpose(tf.squeeze(start_logits, -1), [1, 0]) | |
| start_logits_masked = start_logits * (1 - p_mask) - 1e30 * p_mask | |
| start_log_probs = tf.nn.log_softmax(start_logits_masked, -1) | |
| # logit of the end position | |
| with tf.variable_scope("end_logits"): | |
| if is_training: | |
| # during training, compute the end logits based on the | |
| # ground truth of the start position | |
| start_positions = tf.reshape(features["start_positions"], [-1]) | |
| start_index = tf.one_hot(start_positions, depth=max_seq_length, axis=-1, | |
| dtype=tf.float32) | |
| start_features = tf.einsum("lbh,bl->bh", output, start_index) | |
| start_features = tf.tile(start_features[None], [max_seq_length, 1, 1]) | |
| end_logits = tf.layers.dense( | |
| tf.concat([output, start_features], axis=-1), | |
| albert_config.hidden_size, | |
| kernel_initializer=modeling.create_initializer( | |
| albert_config.initializer_range), | |
| activation=tf.tanh, | |
| name="dense_0") | |
| end_logits = contrib_layers.layer_norm(end_logits, begin_norm_axis=-1) | |
| end_logits = tf.layers.dense( | |
| end_logits, | |
| 1, | |
| kernel_initializer=modeling.create_initializer( | |
| albert_config.initializer_range), | |
| name="dense_1") | |
| end_logits = tf.transpose(tf.squeeze(end_logits, -1), [1, 0]) | |
| end_logits_masked = end_logits * (1 - p_mask) - 1e30 * p_mask | |
| end_log_probs = tf.nn.log_softmax(end_logits_masked, -1) | |
| else: | |
| # during inference, compute the end logits based on beam search | |
| start_top_log_probs, start_top_index = tf.nn.top_k( | |
| start_log_probs, k=start_n_top) | |
| start_index = tf.one_hot(start_top_index, | |
| depth=max_seq_length, axis=-1, dtype=tf.float32) | |
| start_features = tf.einsum("lbh,bkl->bkh", output, start_index) | |
| end_input = tf.tile(output[:, :, None], | |
| [1, 1, start_n_top, 1]) | |
| start_features = tf.tile(start_features[None], | |
| [max_seq_length, 1, 1, 1]) | |
| end_input = tf.concat([end_input, start_features], axis=-1) | |
| end_logits = tf.layers.dense( | |
| end_input, | |
| albert_config.hidden_size, | |
| kernel_initializer=modeling.create_initializer( | |
| albert_config.initializer_range), | |
| activation=tf.tanh, | |
| name="dense_0") | |
| end_logits = contrib_layers.layer_norm(end_logits, begin_norm_axis=-1) | |
| end_logits = tf.layers.dense( | |
| end_logits, | |
| 1, | |
| kernel_initializer=modeling.create_initializer( | |
| albert_config.initializer_range), | |
| name="dense_1") | |
| end_logits = tf.reshape(end_logits, [max_seq_length, -1, start_n_top]) | |
| end_logits = tf.transpose(end_logits, [1, 2, 0]) | |
| end_logits_masked = end_logits * ( | |
| 1 - p_mask[:, None]) - 1e30 * p_mask[:, None] | |
| end_log_probs = tf.nn.log_softmax(end_logits_masked, -1) | |
| end_top_log_probs, end_top_index = tf.nn.top_k( | |
| end_log_probs, k=end_n_top) | |
| end_top_log_probs = tf.reshape( | |
| end_top_log_probs, | |
| [-1, start_n_top * end_n_top]) | |
| end_top_index = tf.reshape( | |
| end_top_index, | |
| [-1, start_n_top * end_n_top]) | |
| if is_training: | |
| return_dict["start_log_probs"] = start_log_probs | |
| return_dict["end_log_probs"] = end_log_probs | |
| else: | |
| return_dict["start_top_log_probs"] = start_top_log_probs | |
| return_dict["start_top_index"] = start_top_index | |
| return_dict["end_top_log_probs"] = end_top_log_probs | |
| return_dict["end_top_index"] = end_top_index | |
| # an additional layer to predict answerability | |
| with tf.variable_scope("answer_class"): | |
| # get the representation of CLS | |
| cls_index = tf.one_hot(tf.zeros([bsz], dtype=tf.int32), | |
| max_seq_length, | |
| axis=-1, dtype=tf.float32) | |
| cls_feature = tf.einsum("lbh,bl->bh", output, cls_index) | |
| # get the representation of START | |
| start_p = tf.nn.softmax(start_logits_masked, axis=-1, | |
| name="softmax_start") | |
| start_feature = tf.einsum("lbh,bl->bh", output, start_p) | |
| # note(zhiliny): no dependency on end_feature so that we can obtain | |
| # one single `cls_logits` for each sample | |
| ans_feature = tf.concat([start_feature, cls_feature], -1) | |
| ans_feature = tf.layers.dense( | |
| ans_feature, | |
| albert_config.hidden_size, | |
| activation=tf.tanh, | |
| kernel_initializer=modeling.create_initializer( | |
| albert_config.initializer_range), | |
| name="dense_0") | |
| ans_feature = tf.layers.dropout(ans_feature, dropout_prob, | |
| training=is_training) | |
| cls_logits = tf.layers.dense( | |
| ans_feature, | |
| 1, | |
| kernel_initializer=modeling.create_initializer( | |
| albert_config.initializer_range), | |
| name="dense_1", | |
| use_bias=False) | |
| cls_logits = tf.squeeze(cls_logits, -1) | |
| return_dict["cls_logits"] = cls_logits | |
| return return_dict | |
| def v2_model_fn_builder(albert_config, init_checkpoint, learning_rate, | |
| num_train_steps, num_warmup_steps, use_tpu, | |
| use_one_hot_embeddings, max_seq_length, start_n_top, | |
| end_n_top, dropout_prob, hub_module): | |
| """Returns `model_fn` closure for TPUEstimator.""" | |
| def model_fn(features, labels, mode, params): # pylint: disable=unused-argument | |
| """The `model_fn` for TPUEstimator.""" | |
| tf.logging.info("*** Features ***") | |
| for name in sorted(features.keys()): | |
| tf.logging.info(" name = %s, shape = %s" % (name, features[name].shape)) | |
| # unique_ids = features["unique_ids"] | |
| input_ids = features["input_ids"] | |
| input_mask = features["input_mask"] | |
| segment_ids = features["segment_ids"] | |
| is_training = (mode == tf.estimator.ModeKeys.TRAIN) | |
| outputs = create_v2_model( | |
| albert_config=albert_config, | |
| is_training=is_training, | |
| input_ids=input_ids, | |
| input_mask=input_mask, | |
| segment_ids=segment_ids, | |
| use_one_hot_embeddings=use_one_hot_embeddings, | |
| features=features, | |
| max_seq_length=max_seq_length, | |
| start_n_top=start_n_top, | |
| end_n_top=end_n_top, | |
| dropout_prob=dropout_prob, | |
| hub_module=hub_module) | |
| tvars = tf.trainable_variables() | |
| initialized_variable_names = {} | |
| scaffold_fn = None | |
| if init_checkpoint: | |
| (assignment_map, initialized_variable_names | |
| ) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint) | |
| if use_tpu: | |
| def tpu_scaffold(): | |
| tf.train.init_from_checkpoint(init_checkpoint, assignment_map) | |
| return tf.train.Scaffold() | |
| scaffold_fn = tpu_scaffold | |
| else: | |
| tf.train.init_from_checkpoint(init_checkpoint, assignment_map) | |
| tf.logging.info("**** Trainable Variables ****") | |
| for var in tvars: | |
| init_string = "" | |
| if var.name in initialized_variable_names: | |
| init_string = ", *INIT_FROM_CKPT*" | |
| tf.logging.info(" name = %s, shape = %s%s", var.name, var.shape, | |
| init_string) | |
| output_spec = None | |
| if mode == tf.estimator.ModeKeys.TRAIN: | |
| seq_length = modeling.get_shape_list(input_ids)[1] | |
| def compute_loss(log_probs, positions): | |
| one_hot_positions = tf.one_hot( | |
| positions, depth=seq_length, dtype=tf.float32) | |
| loss = - tf.reduce_sum(one_hot_positions * log_probs, axis=-1) | |
| loss = tf.reduce_mean(loss) | |
| return loss | |
| start_loss = compute_loss( | |
| outputs["start_log_probs"], features["start_positions"]) | |
| end_loss = compute_loss( | |
| outputs["end_log_probs"], features["end_positions"]) | |
| total_loss = (start_loss + end_loss) * 0.5 | |
| cls_logits = outputs["cls_logits"] | |
| is_impossible = tf.reshape(features["is_impossible"], [-1]) | |
| regression_loss = tf.nn.sigmoid_cross_entropy_with_logits( | |
| labels=tf.cast(is_impossible, dtype=tf.float32), logits=cls_logits) | |
| regression_loss = tf.reduce_mean(regression_loss) | |
| # note(zhiliny): by default multiply the loss by 0.5 so that the scale is | |
| # comparable to start_loss and end_loss | |
| total_loss += regression_loss * 0.5 | |
| train_op = optimization.create_optimizer( | |
| total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu) | |
| output_spec = contrib_tpu.TPUEstimatorSpec( | |
| mode=mode, | |
| loss=total_loss, | |
| train_op=train_op, | |
| scaffold_fn=scaffold_fn) | |
| elif mode == tf.estimator.ModeKeys.PREDICT: | |
| predictions = { | |
| "unique_ids": features["unique_ids"], | |
| "start_top_index": outputs["start_top_index"], | |
| "start_top_log_probs": outputs["start_top_log_probs"], | |
| "end_top_index": outputs["end_top_index"], | |
| "end_top_log_probs": outputs["end_top_log_probs"], | |
| "cls_logits": outputs["cls_logits"] | |
| } | |
| output_spec = contrib_tpu.TPUEstimatorSpec( | |
| mode=mode, predictions=predictions, scaffold_fn=scaffold_fn) | |
| else: | |
| raise ValueError( | |
| "Only TRAIN and PREDICT modes are supported: %s" % (mode)) | |
| return output_spec | |
| return model_fn | |
| def evaluate_v2(result_dict, cls_dict, prediction_json, eval_examples, | |
| eval_features, all_results, n_best_size, max_answer_length, | |
| output_prediction_file, output_nbest_file, | |
| output_null_log_odds_file): | |
| null_score_diff_threshold = None | |
| predictions, na_probs = write_predictions_v2( | |
| result_dict, cls_dict, eval_examples, eval_features, | |
| all_results, n_best_size, max_answer_length, | |
| output_prediction_file, output_nbest_file, | |
| output_null_log_odds_file, null_score_diff_threshold) | |
| na_prob_thresh = 1.0 # default value taken from the eval script | |
| qid_to_has_ans = make_qid_to_has_ans(prediction_json) # maps qid to True/False | |
| has_ans_qids = [k for k, v in qid_to_has_ans.items() if v] | |
| no_ans_qids = [k for k, v in qid_to_has_ans.items() if not v] | |
| exact_raw, f1_raw = get_raw_scores(prediction_json, predictions) | |
| exact_thresh = apply_no_ans_threshold(exact_raw, na_probs, qid_to_has_ans, | |
| na_prob_thresh) | |
| f1_thresh = apply_no_ans_threshold(f1_raw, na_probs, qid_to_has_ans, | |
| na_prob_thresh) | |
| out_eval = make_eval_dict(exact_thresh, f1_thresh) | |
| find_all_best_thresh(out_eval, predictions, exact_raw, f1_raw, na_probs, qid_to_has_ans) | |
| null_score_diff_threshold = out_eval["best_f1_thresh"] | |
| predictions, na_probs = write_predictions_v2( | |
| result_dict, cls_dict,eval_examples, eval_features, | |
| all_results, n_best_size, max_answer_length, | |
| output_prediction_file, output_nbest_file, | |
| output_null_log_odds_file, null_score_diff_threshold) | |
| qid_to_has_ans = make_qid_to_has_ans(prediction_json) # maps qid to True/False | |
| has_ans_qids = [k for k, v in qid_to_has_ans.items() if v] | |
| no_ans_qids = [k for k, v in qid_to_has_ans.items() if not v] | |
| exact_raw, f1_raw = get_raw_scores(prediction_json, predictions) | |
| exact_thresh = apply_no_ans_threshold(exact_raw, na_probs, qid_to_has_ans, | |
| na_prob_thresh) | |
| f1_thresh = apply_no_ans_threshold(f1_raw, na_probs, qid_to_has_ans, | |
| na_prob_thresh) | |
| out_eval = make_eval_dict(exact_thresh, f1_thresh) | |
| out_eval["null_score_diff_threshold"] = null_score_diff_threshold | |
| return out_eval | |