Added natural questions. Not yet validated or tested with a call to write.

lm_eval/tasks/__init__.py

@@ -9,6 +9,7 @@ from . import quac
 from . import hellaswag
 from . import openbookqa
 from . import squad
+from . import naturalqs
 
 TASK_REGISTRY = {
     # GLUE
@@ -36,6 +37,7 @@ TASK_REGISTRY = {
     "openbookqa": openbookqa.OpenBookQA,
     "squad": squad.SQuAD,
     "race": race.RACE,
+    "naturalqs": naturalqs.NaturalQs,
     "webqs": webqs.WebQs,
     "winogrande": winogrande.Winogrande,
     "anli_r1": anli.ANLIRound1,

lm_eval/tasks/naturalqs.py

+from . common import HFTask
+import apache_beam
+
+class NaturalQs(HFTask):
+    DATASET_PATH = "natural_questions"
+    DATASET_NAME = None
+
+    def has_training_docs(self):
+        return True
+
+    def has_validation_docs(self):
+        return True
+
+    def has_test_docs(self):
+        return False
+
+    def fewshot_description(self):
+        # TODO: figure out description
+        return ""
+
+    def doc_to_text(self, doc, include_target=True):
+        question = doc['question']['text']
+        short_answer = doc['annotations']['short_answers'][0]['text']
+        long_answer_start = doc['annotations']['long_answer'][0]['start_token']
+        long_answer_end = doc['annotations']['long_answer'][0]['end_token']
+        passage = " ".join(doc['document']['tokens']['token'][long_answer_start:long_answer_end])
+        
+        text = 'Q: ' + question + '\n\n' + 'A: '
+
+        if include_target:
+            # What if there is no short answer? This will be an empty string. Currently, default to the long answer otherwise.
+            if short_answer:
+                text += short_answer[0]
+            else:
+                text += long_answer
+
+        return text
+
+    def evaluate(self, docs, lm, provide_description, num_fewshot):
+        # TODO: implement
+        raise NotImplementedError()
\ No newline at end of file