Added NQ-Open task

lm_eval/tasks/__init__.py

@@ -20,6 +20,7 @@ from . import swag
 from . import openbookqa
 from . import squad
 from . import naturalqs
+from . import nqopen
 from . import sat
 from . import arithmetic
 from . import lambada
@@ -150,6 +151,7 @@ TASK_REGISTRY = {
     "squad2": squad.SQuAD2,
     "race": race.RACE,
     # "naturalqs": naturalqs.NaturalQs, # not implemented yet
+    "nq_open": nqopen.NQOpen,
     "headqa": headqa.HeadQAEsDeprecated,  # for backwards compat - headqa used to default to es
     "headqa_es": headqa.HeadQAEs,
     "headqa_en": headqa.HeadQAEn,

lm_eval/tasks/nqopen.py

+"""
+Natural Questions: a Benchmark for Question Answering Research
+https://storage.googleapis.com/pub-tools-public-publication-data/pdf/1f7b46b5378d757553d3e92ead36bda2e4254244.pdf
+
+The Natural Questions (NQ) corpus is a question-answering dataset that contains
+questions from real users and requires QA systems to read and comprehend an entire
+Wikipedia article that may or may not contain the answer to the question. The
+inclusion of real user questions, and the requirement that solutions should read
+an entire page to find the answer, cause NQ to be a more realistic and challenging
+task than prior QA datasets.
+
+TODO: NaturalQS has a *really* large train set that huggingface just automatically
+downloads even if you dont use it. we should try and only download the val set and
+not even bother with the train set.
+
+Homepage: https://ai.google.com/research/NaturalQuestions
+"""
+import re
+import string
+from lm_eval.base import Task, rf
+from lm_eval.metrics import mean
+
+
+_CITATION = """
+@article{47761,
+    title={Natural Questions: a Benchmark for Question Answering Research},
+    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},
+    year={2019},
+    journal={Transactions of the Association of Computational Linguistics}
+}
+"""
+
+
+class NQOpen(Task):
+    VERSION = 0
+    DATASET_PATH = "nq_open"
+    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 training_docs(self):
+        return self.dataset["train"]
+
+    def validation_docs(self):
+        return self.dataset["validation"]
+
+    def test_docs(self):
+        raise NotImplementedError()
+    
+    def doc_to_text(self, doc):
+        return f"Question: {doc['question']}\nAnswer:"
+
+    def should_decontaminate(self):
+        return True
+
+    def doc_to_decontamination_query(self, doc):
+        return doc["question"]
+
+    def doc_to_target(self, doc):
+        return " " + doc["answer"][0]
+
+    def construct_requests(self, doc, ctx):
+        """Uses RequestFactory to construct Requests and returns an iterable of
+        Requests which will be sent to the LM.
+        :param doc:
+                The document as returned from training_docs, validation_docs, or test_docs.
+        :param ctx: str
+                The context string, generated by fewshot_context. This includes the natural
+                language description, as well as the few shot examples, and the question
+                part of the document for `doc`.
+        """
+        continuation = rf.greedy_until(ctx, {"until": ["\n", ".", ","]})
+        return continuation
+
+    def process_results(self, doc, results):
+        """Take a single document and the LM results and evaluates, returning a
+        dict where keys are the names of submetrics and values are the values of
+        the metric for that one document
+
+        :param doc:
+            The document as returned from training_docs, validation_docs, or test_docs.
+        :param results:
+            The results of the requests created in construct_requests.
+        """
+        print("raw results:", results)
+        continuation = results[0].strip().lower().translate(str.maketrans('', '', string.punctuation))
+        answers = [answer.lower().translate(str.maketrans('', '', string.punctuation)) for answer in doc["answer"]]
+        
+        # remove duplicate whitespace
+        continuation = re.sub(' +', ' ', continuation)
+        
+        # remove articles
+        continuation = re.sub('(\s+)(a|an|the)(\s+)', ' ', continuation)
+        answers = [re.sub('(\s+)(a|an|the)(\s+)', ' ', cand) for cand in answers]
+ 
+        print(float(continuation in answers), continuation, answers)
+        return {
+            "em": float(continuation in answers)
+        }
+
+    def aggregation(self):
+        """
+        :returns: {str: [float] -> float}
+            A dictionary where keys are the names of submetrics and values are
+            functions that aggregate a list of metrics
+        """
+        return {
+            "em": mean,
+        } 
+        
+    def higher_is_better(self):
+        """
+        :returns: {str: bool}
+            A dictionary where keys are the names of submetrics and values are
+            whether a higher value of the submetric is better
+        """
+        return {
+            "em": True,
+        }