Add dataset reader for DROP (#2556)

* Add dataset reader for DROP

* Add missing dependency

* Fix pylint, mypy and docs

* Add some more tests

* Address PR feedback

* mypy again

Author: matt-gardner

allennlp/data/dataset_readers/__init__.py

@@ -18,7 +18,8 @@
 from allennlp.data.dataset_readers.language_modeling import LanguageModelingReader
 from allennlp.data.dataset_readers.multiprocess_dataset_reader import MultiprocessDatasetReader
 from allennlp.data.dataset_readers.penn_tree_bank import PennTreeBankConstituencySpanDatasetReader
-from allennlp.data.dataset_readers.reading_comprehension import SquadReader, TriviaQaReader, QuACReader, QangarooReader
+from allennlp.data.dataset_readers.reading_comprehension import (
+        DropReader, SquadReader, TriviaQaReader, QuACReader, QangarooReader)
 from allennlp.data.dataset_readers.semantic_role_labeling import SrlReader
 from allennlp.data.dataset_readers.semantic_dependency_parsing import SemanticDependenciesDatasetReader
 from allennlp.data.dataset_readers.seq2seq import Seq2SeqDatasetReader

allennlp/data/dataset_readers/reading_comprehension/__init__.py

@@ -5,6 +5,7 @@
 These submodules contain readers for things that are predominantly reading comprehension datasets.
 """
 
+from allennlp.data.dataset_readers.reading_comprehension.drop import DropReader
 from allennlp.data.dataset_readers.reading_comprehension.squad import SquadReader
 from allennlp.data.dataset_readers.reading_comprehension.quac import QuACReader
 from allennlp.data.dataset_readers.reading_comprehension.triviaqa import TriviaQaReader

allennlp/data/dataset_readers/reading_comprehension/drop.py

@@ -369,3 +369,42 @@ def handle_cannot(reference_answers: List[str]):
     else:
         reference_answers = [x for x in reference_answers if x != 'CANNOTANSWER']
     return reference_answers
+
+
+def split_token_by_delimiter(token: Token, delimiter: str) -> List[Token]:
+    split_tokens = []
+    char_offset = token.idx
+    for sub_str in token.text.split(delimiter):
+        if sub_str:
+            split_tokens.append(Token(text=sub_str, idx=char_offset))
+            char_offset += len(sub_str)
+        split_tokens.append(Token(text=delimiter, idx=char_offset))
+        char_offset += len(delimiter)
+    if split_tokens:
+        split_tokens.pop(-1)
+        char_offset -= len(delimiter)
+        return split_tokens
+    else:
+        return [token]
+
+
+def split_tokens_by_hyphen(tokens: List[Token]) -> List[Token]:
+    hyphens = ["-", "–", "~"]
+    new_tokens: List[Token] = []
+
+    for token in tokens:
+        if any(hyphen in token.text for hyphen in hyphens):
+            unsplit_tokens = [token]
+            split_tokens: List[Token] = []
+            for hyphen in hyphens:
+                for unsplit_token in unsplit_tokens:
+                    if hyphen in token.text:
+                        split_tokens += split_token_by_delimiter(unsplit_token, hyphen)
+                    else:
+                        split_tokens.append(unsplit_token)
+                unsplit_tokens, split_tokens = split_tokens, []
+            new_tokens += unsplit_tokens
+        else:
+            new_tokens.append(token)
+
+    return new_tokens

allennlp/data/dataset_readers/reading_comprehension/util.py

@@ -0,0 +1,142 @@
+# pylint: disable=no-self-use,invalid-name, protected-access
+import pytest
+
+from allennlp.common import Params
+from allennlp.common.util import ensure_list
+from allennlp.data.dataset_readers import DropReader
+from allennlp.common.testing import AllenNlpTestCase
+
+
+class TestDropReader:
+    @pytest.mark.parametrize("lazy", (True, False))
+    def test_read_from_file(self, lazy):
+        reader = DropReader(lazy=lazy)
+        instances = ensure_list(reader.read(AllenNlpTestCase.FIXTURES_ROOT / 'data' / 'drop.json'))
+        assert len(instances) == 19
+
+        instance = instances[0]
+        assert set(instance.fields.keys()) == {
+                'question',
+                'passage',
+                'number_indices',
+                'numbers_in_passage',
+                'answer_as_passage_spans',
+                'answer_as_question_spans',
+                'answer_as_add_sub_expressions',
+                'answer_as_counts',
+                'metadata',
+                }
+
+        assert [t.text for t in instance["question"][:3]] == ["What", "happened", "second"]
+        assert [t.text for t in instance["passage"][:3]] == ["The", "Port", "of"]
+        assert [t.text for t in instance["passage"][-3:]] == ["cruise", "ships", "."]
+
+        # Note that the last number in here is added as padding in case we don't find any numbers
+        # in a particular passage.
+        assert [f.sequence_index for f in instance["number_indices"]] == [
+                16, 30, 36, 41, 52, 64, 80, 89, 147, 153, 166, 174, 177, 206, 245, 252, 267, 279,
+                283, 288, 296, -1
+                ]
+        assert [t.text for t in instance["numbers_in_passage"]] == [
+                "1", "25", "2014", "5", "2018", "1", "2", "1", "54", "52", "6", "60", "58", "2010",
+                "67", "2010", "1996", "3", "1", "6", "1", "0"]
+        assert len(instance["answer_as_passage_spans"]) == 1
+        assert instance["answer_as_passage_spans"][0] == (46, 47)
+        assert len(instance["answer_as_question_spans"]) == 1
+        assert instance["answer_as_question_spans"][0] == (5, 6)
+        assert len(instance["answer_as_add_sub_expressions"]) == 1
+        assert instance["answer_as_add_sub_expressions"][0].labels == [0,] * 22
+        assert len(instance["answer_as_counts"]) == 1
+        assert instance["answer_as_counts"][0].label == -1
+        assert set(instance['metadata'].metadata.keys()) == {
+                'answer_annotations',
+                'answer_info',
+                'answer_texts',
+                'number_indices',
+                'number_tokens',
+                'original_numbers',
+                'original_passage',
+                'original_question',
+                'passage_id',
+                'passage_token_offsets',
+                'passage_tokens',
+                'question_id',
+                'question_token_offsets',
+                'question_tokens',
+                }
+
+    def test_read_in_bert_format(self):
+        reader = DropReader(instance_format="bert")
+        instances = ensure_list(reader.read(AllenNlpTestCase.FIXTURES_ROOT / 'data' / 'drop.json'))
+        assert len(instances) == 19
+
+        print(instances[0])
+        instance = instances[0]
+        assert set(instance.fields.keys()) == {
+                'answer_as_passage_spans',
+                'metadata',
+                'passage',
+                'question',
+                'question_and_passage',
+                }
+
+        assert [t.text for t in instance["question"][:3]] == ["What", "happened", "second"]
+        assert [t.text for t in instance["passage"][:3]] == ["The", "Port", "of"]
+        assert [t.text for t in instance["passage"][-3:]] == ["cruise", "ships", "."]
+        question_length = len(instance['question'])
+        passage_length = len(instance['passage'])
+        assert len(instance['question_and_passage']) == question_length + passage_length + 1
+
+        assert len(instance["answer_as_passage_spans"]) == 1
+        assert instance["answer_as_passage_spans"][0] == (question_length + 1 + 46,
+                                                          question_length + 1 + 47)
+        assert set(instance['metadata'].metadata.keys()) == {
+                'answer_annotations',
+                'answer_texts',
+                'original_passage',
+                'original_question',
+                'passage_id',
+                'passage_token_offsets',
+                'passage_tokens',
+                'question_id',
+                'question_tokens',
+                }
+
+    def test_read_in_squad_format(self):
+        reader = DropReader(instance_format="squad")
+        instances = ensure_list(reader.read(AllenNlpTestCase.FIXTURES_ROOT / 'data' / 'drop.json'))
+        assert len(instances) == 19
+
+        print(instances[0])
+        instance = instances[0]
+        assert set(instance.fields.keys()) == {
+                'question',
+                'passage',
+                'span_start',
+                'span_end',
+                'metadata',
+                }
+
+        assert [t.text for t in instance["question"][:3]] == ["What", "happened", "second"]
+        assert [t.text for t in instance["passage"][:3]] == ["The", "Port", "of"]
+        assert [t.text for t in instance["passage"][-3:]] == ["cruise", "ships", "."]
+
+        assert instance["span_start"] == 46
+        assert instance["span_end"] == 47
+        assert set(instance['metadata'].metadata.keys()) == {
+                'answer_annotations',
+                'answer_texts',
+                'original_passage',
+                'original_question',
+                'passage_id',
+                'token_offsets',
+                'passage_tokens',
+                'question_id',
+                'question_tokens',
+                'valid_passage_spans',
+                }
+
+    def test_can_build_from_params(self):
+        reader = DropReader.from_params(Params({}))
+        assert reader._tokenizer.__class__.__name__ == 'WordTokenizer'
+        assert reader._token_indexers["tokens"].__class__.__name__ == 'SingleIdTokenIndexer'