|
| 1 | +# coding=utf-8 |
| 2 | +# Copyright 2020 The TensorFlow Datasets Authors. |
| 3 | +# |
| 4 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +# you may not use this file except in compliance with the License. |
| 6 | +# You may obtain a copy of the License at |
| 7 | +# |
| 8 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +# |
| 10 | +# Unless required by applicable law or agreed to in writing, software |
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +# See the License for the specific language governing permissions and |
| 14 | +# limitations under the License. |
| 15 | + |
| 16 | +"""Natural Questions: A Benchmark for Question Answering Research.""" |
| 17 | + |
| 18 | +from __future__ import absolute_import |
| 19 | +from __future__ import division |
| 20 | +from __future__ import print_function |
| 21 | + |
| 22 | +import json |
| 23 | +import re |
| 24 | + |
| 25 | +import six |
| 26 | +import tensorflow.compat.v2 as tf |
| 27 | +import tensorflow_datasets.public_api as tfds |
| 28 | + |
| 29 | +if six.PY2: |
| 30 | + import HTMLParser as html_parser # pylint:disable=g-import-not-at-top |
| 31 | + html_unescape = html_parser.HTMLParser().unescape |
| 32 | +else: |
| 33 | + import html # pylint:disable=g-import-not-at-top |
| 34 | + html_unescape = html.unescape |
| 35 | + |
| 36 | +_CITATION = """ |
| 37 | +@article{47761, |
| 38 | +title = {Natural Questions: a Benchmark for Question Answering Research}, |
| 39 | +author = {Tom Kwiatkowski and Jennimaria Palomaki and Olivia Redfield and Michael Collins and Ankur Parikh and Chris Alberti and Danielle Epstein and Illia Polosukhin and Matthew Kelcey and Jacob Devlin and Kenton Lee and Kristina N. Toutanova and Llion Jones and Ming-Wei Chang and Andrew Dai and Jakob Uszkoreit and Quoc Le and Slav Petrov}, |
| 40 | +year = {2019}, |
| 41 | +journal = {Transactions of the Association of Computational Linguistics} |
| 42 | +} |
| 43 | +""" |
| 44 | + |
| 45 | +_DESCRIPTION = """ |
| 46 | +The NQ corpus contains questions from real users, and it requires QA systems to |
| 47 | +read and comprehend an entire Wikipedia article that may or may not contain the |
| 48 | +answer to the question. The inclusion of real user questions, and the |
| 49 | +requirement that solutions should read an entire page to find the answer, cause |
| 50 | +NQ to be a more realistic and challenging task than prior QA datasets. |
| 51 | +""" |
| 52 | + |
| 53 | +_URL = 'https://ai.google.com/research/NaturalQuestions/dataset' |
| 54 | + |
| 55 | +_BASE_DOWNLOAD_URL = 'https://storage.googleapis.com/natural_questions/v1.0' |
| 56 | +_DOWNLOAD_URLS = { |
| 57 | + 'train': [ |
| 58 | + '%s/train/nq-train-%02d.jsonl.gz' % (_BASE_DOWNLOAD_URL, i) |
| 59 | + for i in range(50) |
| 60 | + ], |
| 61 | + 'validation': [ |
| 62 | + '%s/dev/nq-dev-%02d.jsonl.gz' % (_BASE_DOWNLOAD_URL, i) |
| 63 | + for i in range(5) |
| 64 | + ] |
| 65 | +} |
| 66 | + |
| 67 | + |
| 68 | +class NaturalQuestions(tfds.core.BeamBasedBuilder): |
| 69 | + """Natural Questions: A Benchmark for Question Answering Research.""" |
| 70 | + |
| 71 | + VERSION = tfds.core.Version('0.0.1') |
| 72 | + |
| 73 | + def _info(self): |
| 74 | + return tfds.core.DatasetInfo( |
| 75 | + builder=self, |
| 76 | + description=_DESCRIPTION, |
| 77 | + # TODO(adarob): Pull in NQ token features if needed. |
| 78 | + features=tfds.features.FeaturesDict({ |
| 79 | + 'id': tf.string, |
| 80 | + 'document': { |
| 81 | + 'title': tfds.features.Text(), |
| 82 | + 'url': tfds.features.Text(), |
| 83 | + 'html': tfds.features.Text(), |
| 84 | + }, |
| 85 | + 'question': { |
| 86 | + 'text': tfds.features.Text(), |
| 87 | + }, |
| 88 | + 'annotations': tfds.features.Sequence({ |
| 89 | + 'id': tf.string, |
| 90 | + 'long_answer': { |
| 91 | + 'start_byte': tf.int64, |
| 92 | + 'end_byte': tf.int64, |
| 93 | + }, |
| 94 | + 'short_answers': tfds.features.Sequence({ |
| 95 | + 'start_byte': tf.int64, |
| 96 | + 'end_byte': tf.int64, |
| 97 | + 'text': tfds.features.Text(), |
| 98 | + }), |
| 99 | + 'yes_no_answer': tfds.features.ClassLabel( |
| 100 | + names=['NO', 'YES']) # Can also be -1 for NONE. |
| 101 | + }), |
| 102 | + }), |
| 103 | + supervised_keys=None, |
| 104 | + homepage=_URL, |
| 105 | + citation=_CITATION, |
| 106 | + ) |
| 107 | + |
| 108 | + def _split_generators(self, dl_manager): |
| 109 | + """Returns SplitGenerators.""" |
| 110 | + |
| 111 | + files = dl_manager.download(_DOWNLOAD_URLS) |
| 112 | + |
| 113 | + return [ |
| 114 | + tfds.core.SplitGenerator( |
| 115 | + name=tfds.Split.TRAIN, |
| 116 | + gen_kwargs={'filepaths': files['train']}, |
| 117 | + ), |
| 118 | + tfds.core.SplitGenerator( |
| 119 | + name=tfds.Split.VALIDATION, |
| 120 | + num_shards=8, |
| 121 | + gen_kwargs={'filepaths': files['validation']}, |
| 122 | + ), |
| 123 | + ] |
| 124 | + |
| 125 | + def _build_pcollection(self, pipeline, filepaths): |
| 126 | + """Build PCollection of examples.""" |
| 127 | + beam = tfds.core.lazy_imports.apache_beam |
| 128 | + |
| 129 | + def _parse_example(line): |
| 130 | + """Parse a single json line and emit an example dict.""" |
| 131 | + ex_json = json.loads(line) |
| 132 | + html_bytes = ex_json['document_html'].encode('utf-8') |
| 133 | + |
| 134 | + def _parse_short_answer(short_ans): |
| 135 | + """"Extract text of short answer.""" |
| 136 | + ans_bytes = html_bytes[ |
| 137 | + short_ans['start_byte']:short_ans['end_byte']] |
| 138 | + # Remove non-breaking spaces. |
| 139 | + ans_bytes = ans_bytes.replace(b'\xc2\xa0', b' ') |
| 140 | + text = ans_bytes.decode('utf-8') |
| 141 | + # Remove HTML markup. |
| 142 | + text = re.sub('<([^>]*)>', '', html_unescape(text)) |
| 143 | + # Replace \xa0 characters with spaces. |
| 144 | + return { |
| 145 | + 'start_byte': short_ans['start_byte'], |
| 146 | + 'end_byte': short_ans['end_byte'], |
| 147 | + 'text': text |
| 148 | + } |
| 149 | + |
| 150 | + def _parse_annotation(an_json): |
| 151 | + return { |
| 152 | + # Convert to str since some IDs cannot be represented by tf.int64. |
| 153 | + 'id': str(an_json['annotation_id']), |
| 154 | + 'long_answer': { |
| 155 | + 'start_byte': an_json['long_answer']['start_byte'], |
| 156 | + 'end_byte': an_json['long_answer']['end_byte'], |
| 157 | + }, |
| 158 | + 'short_answers': [ |
| 159 | + _parse_short_answer(ans) for ans in an_json['short_answers']], |
| 160 | + 'yes_no_answer': ( |
| 161 | + -1 if an_json['yes_no_answer'] == 'NONE' |
| 162 | + else an_json['yes_no_answer']) |
| 163 | + } |
| 164 | + |
| 165 | + beam.metrics.Metrics.counter('nq', 'examples').inc() |
| 166 | + # Convert to str since some IDs cannot be represented by tf.int64. |
| 167 | + id_ = str(ex_json['example_id']) |
| 168 | + return id_, { |
| 169 | + 'id': id_, |
| 170 | + 'document': { |
| 171 | + 'title': ex_json['document_title'], |
| 172 | + 'url': ex_json['document_url'], |
| 173 | + 'html': html_bytes, |
| 174 | + }, |
| 175 | + 'question': { |
| 176 | + 'text': ex_json['question_text'] |
| 177 | + }, |
| 178 | + 'annotations': [ |
| 179 | + _parse_annotation(an_json) for an_json in ex_json['annotations'] |
| 180 | + ] |
| 181 | + } |
| 182 | + |
| 183 | + return ( |
| 184 | + pipeline |
| 185 | + | beam.Create(filepaths) |
| 186 | + | beam.io.ReadAllFromText() |
| 187 | + | beam.Map(_parse_example)) |
0 commit comments