Adding ROPES, RACE tasks

guides/tasks/supported_tasks.md

@@ -50,6 +50,8 @@
 | EP-POS | pos | ✅ |  | pos | Edge-Probing |
 | QNLI | qnli | ✅ | ✅ | qnli | GLUE |
 | QQP | qqp | ✅ | ✅ | qqp | GLUE |
+| ROPES | ropes | ✅ | ✅ | ropes |  |
+| RACE | race | ✅ | ✅ | race | `race`, `race_middle`, `race_high` |
 | ReCord | record | ✅ | ✅ | record | SuperGLUE |
 | RTE | rte | ✅ | ✅ | rte | GLUE, SuperGLUE |
 | SciTail | scitail | ✅ | ✅ | scitail |  |

jiant/scripts/download_data/constants.py

@@ -17,6 +17,7 @@
     "mrqa_natural_questions",
     "piqa",
     "winogrande",
+    "ropes",
 }
 
 DIRECT_DOWNLOAD_TASKS = set(
@@ -35,4 +36,7 @@
     "adversarial_nli_r3",
     "arc_easy",
     "arc_challenge",
+    "race",
+    "race_middle",
+    "race_high",
 }

jiant/scripts/download_data/dl_datasets/files_tasks.py

@@ -61,6 +61,10 @@ def download_task_data_and_write_config(task_name: str, task_data_path: str, tas
         download_winogrande_data_and_write_config(
             task_name=task_name, task_data_path=task_data_path, task_config_path=task_config_path
         )
+    elif task_name == "ropes":
+        download_ropes_data_and_write_config(
+            task_name=task_name, task_data_path=task_data_path, task_config_path=task_config_path
+        )
     else:
         raise KeyError(task_name)
 
@@ -171,7 +175,7 @@ def download_fever_nli_data_and_write_config(
 ):
     os.makedirs(task_data_path, exist_ok=True)
     download_utils.download_and_unzip(
-        ("https://www.dropbox.com/s/hylbuaovqwo2zav/nli_fever.zip?dl=1"), task_data_path,
+        "https://www.dropbox.com/s/hylbuaovqwo2zav/nli_fever.zip?dl=1", task_data_path,
     )
     # Since the FEVER NLI dataset doesn't have labels for the dev set, we also download the original
     # FEVER dev set and match example CIDs to obtain labels.
@@ -739,3 +743,34 @@ def download_winogrande_data_and_write_config(
         },
         path=task_config_path,
     )
+
+
+def download_ropes_data_and_write_config(
+    task_name: str, task_data_path: str, task_config_path: str
+):
+    os.makedirs(task_data_path, exist_ok=True)
+    download_utils.download_and_untar(
+        "https://ropes-dataset.s3-us-west-2.amazonaws.com/train_and_dev/"
+        "ropes-train-dev-v1.0.tar.gz",
+        task_data_path,
+    )
+    data_phase_list = ["train", "dev"]
+    jiant_phase_list = ["train", "val"]
+    for data_phase, jiant_phase in zip(data_phase_list, jiant_phase_list):
+        os.rename(
+            os.path.join(task_data_path, "ropes-train-dev-v1.0", f"{data_phase}-v1.0.json"),
+            os.path.join(task_data_path, f"{jiant_phase}.json"),
+        )
+    shutil.rmtree(os.path.join(task_data_path, "ropes-train-dev-v1.0"))
+    py_io.write_json(
+        data={
+            "task": task_name,
+            "paths": {
+                "train": os.path.join(task_data_path, "train.json"),
+                "val": os.path.join(task_data_path, "val.json"),
+            },
+            "name": task_name,
+            "kwargs": {"include_background": True},
+        },
+        path=task_config_path,
+    )

jiant/scripts/download_data/dl_datasets/hf_datasets_tasks.py

@@ -163,6 +163,24 @@
         "phase_list": ["train", "val", "test"],
         "jiant_task_name": "arc_challenge",
     },
+    "race": {
+        "path": "race",
+        "name": "all",
+        "phase_list": ["train", "val", "test"],
+        "jiant_task_name": "race",
+    },
+    "race_middle": {
+        "path": "race",
+        "name": "middle",
+        "phase_list": ["train", "val", "test"],
+        "jiant_task_name": "race",
+    },
+    "race_high": {
+        "path": "race",
+        "name": "high",
+        "phase_list": ["train", "val", "test"],
+        "jiant_task_name": "race",
+    },
 }
 
 # HF-Datasets uses "validation", we use "val"

jiant/tasks/evaluate/core.py

@@ -963,6 +963,7 @@ def get_evaluation_scheme_for_task(task) -> BaseEvaluationScheme:
             tasks.MnliTask,
             tasks.PawsXTask,
             tasks.QnliTask,
+            tasks.RaceTask,
             tasks.RteTask,
             tasks.SciTailTask,
             tasks.SentevalTenseTask,
@@ -1025,7 +1026,14 @@ def get_evaluation_scheme_for_task(task) -> BaseEvaluationScheme:
     elif isinstance(task, tasks.ReCoRDTask):
         return ReCordEvaluationScheme()
     elif isinstance(
-        task, (tasks.SquadTask, tasks.QuorefTask, tasks.NewsQATask, tasks.MrqaNaturalQuestionsTask,)
+        task,
+        (
+            tasks.SquadTask,
+            tasks.RopesTask,
+            tasks.QuorefTask,
+            tasks.NewsQATask,
+            tasks.MrqaNaturalQuestionsTask,
+        ),
     ):
         return SQuADEvaluationScheme()
     elif isinstance(task, (tasks.TyDiQATask, tasks.XquadTask)):

jiant/tasks/lib/race.py

@@ -0,0 +1,130 @@
+import numpy as np
+
+from dataclasses import dataclass
+
+from jiant.tasks.lib.templates.shared import (
+    labels_to_bimap,
+    create_input_set_from_tokens_and_segments,
+    add_cls_token,
+)
+from jiant.tasks.lib.templates import multiple_choice as mc_template
+from jiant.tasks.utils import truncate_sequences
+from jiant.utils.python.io import read_jsonl
+
+
+@dataclass
+class Example(mc_template.Example):
+    @property
+    def task(self):
+        return RaceTask
+
+    def tokenize(self, tokenizer):
+        return TokenizedExample(
+            guid=self.guid,
+            prompt=tokenizer.tokenize(self.prompt),
+            choice_list=[tokenizer.tokenize(choice) for choice in self.choice_list],
+            label_id=self.task.CHOICE_TO_ID[self.label],
+        )
+
+
+@dataclass
+class TokenizedExample(mc_template.TokenizedExample):
+    def featurize(self, tokenizer, feat_spec):
+        if feat_spec.sep_token_extra:
+            maybe_extra_sep = [tokenizer.sep_token]
+            maybe_extra_sep_segment_id = [feat_spec.sequence_a_segment_id]
+            special_tokens_count = 4  # CLS, SEP-SEP, SEP
+        else:
+            maybe_extra_sep = []
+            maybe_extra_sep_segment_id = []
+            special_tokens_count = 3  # CLS, SEP, SEP
+
+        input_set_ls = []
+        unpadded_inputs_ls = []
+        for choice in self.choice_list:
+            prompt, choice = truncate_sequences(
+                tokens_ls=[self.prompt, choice],
+                max_length=feat_spec.max_seq_length - special_tokens_count,
+                truncate_end=False,
+            )
+            unpadded_inputs = add_cls_token(
+                unpadded_tokens=(
+                    # prompt
+                    prompt
+                    + [tokenizer.sep_token]
+                    + maybe_extra_sep
+                    # choice
+                    + choice
+                    + [tokenizer.sep_token]
+                ),
+                unpadded_segment_ids=(
+                    # prompt
+                    [feat_spec.sequence_a_segment_id] * (len(prompt) + 1)
+                    + maybe_extra_sep_segment_id
+                    # choice + sep
+                    + [feat_spec.sequence_b_segment_id] * (len(choice) + 1)
+                ),
+                tokenizer=tokenizer,
+                feat_spec=feat_spec,
+            )
+            input_set = create_input_set_from_tokens_and_segments(
+                unpadded_tokens=unpadded_inputs.unpadded_tokens,
+                unpadded_segment_ids=unpadded_inputs.unpadded_segment_ids,
+                tokenizer=tokenizer,
+                feat_spec=feat_spec,
+            )
+            input_set_ls.append(input_set)
+            unpadded_inputs_ls.append(unpadded_inputs)
+
+        return DataRow(
+            guid=self.guid,
+            input_ids=np.stack([input_set.input_ids for input_set in input_set_ls]),
+            input_mask=np.stack([input_set.input_mask for input_set in input_set_ls]),
+            segment_ids=np.stack([input_set.segment_ids for input_set in input_set_ls]),
+            label_id=self.label_id,
+            tokens_list=[unpadded_inputs.unpadded_tokens for unpadded_inputs in unpadded_inputs_ls],
+        )
+
+
+@dataclass
+class DataRow(mc_template.DataRow):
+    pass
+
+
+@dataclass
+class Batch(mc_template.Batch):
+    pass
+
+
+class RaceTask(mc_template.AbstractMultipleChoiceTask):
+    Example = Example
+    TokenizedExample = TokenizedExample
+    DataRow = DataRow
+    Batch = Batch
+
+    CHOICE_KEYS = ["A", "B", "C", "D"]
+    CHOICE_TO_ID, ID_TO_CHOICE = labels_to_bimap(CHOICE_KEYS)
+    NUM_CHOICES = len(CHOICE_KEYS)
+
+    def get_train_examples(self):
+        return self._create_examples(lines=read_jsonl(self.train_path), set_type="train")
+
+    def get_val_examples(self):
+        return self._create_examples(lines=read_jsonl(self.val_path), set_type="val")
+
+    def get_test_examples(self):
+        return self._create_examples(lines=read_jsonl(self.test_path), set_type="test")
+
+    @classmethod
+    def _create_examples(cls, lines, set_type):
+        examples = []
+        for (i, line) in enumerate(lines):
+            examples.append(
+                Example(
+                    guid="%s-%s" % (set_type, i),
+                    prompt=line["article"] + " " + line["question"],
+                    choice_list=line["options"],
+                    label=line["answer"] if set_type != "test" else cls.CHOICE_KEYS[-1],
+                )
+            )
+        return examples