Merge pull request #37 from EleutherAI/download_on_demand

download_on_demand

lm_eval/base.py

@@ -54,6 +54,12 @@ class LM(abc.ABC):
 
 
 class Dataset(abc.ABC):
+
+    @abc.abstractmethod
+    def download(self):
+        """Downloads the task dataset if necessary"""
+        pass
+
     @abc.abstractmethod
     def has_training_docs(self):
         """Whether the task has a training set"""
@@ -121,4 +127,4 @@ class Dataset(abc.ABC):
             map(self.doc_to_text, self.fewshot_examples(k=num_fewshot))
         ) + "\n\n"
         example = self.doc_to_text(doc, include_target=False).strip()
-        return description + labeled_examples + example
+        return description + labeled_examples + example
\ No newline at end of file

lm_eval/tasks/common.py

-import nlp
+import datasets
 import numpy as np
 import random
 from ..base import Dataset
@@ -11,25 +11,35 @@ class HFNLPTask(Dataset):
     def __init__(self):
         super().__init__()
         self._training_docs = None
+        self.data = datasets.load_dataset(path=self.NLP_PATH, name=self.NLP_NAME)
 
-    def _load_nlp_dataset(self):
-        return nlp.load_dataset(path=self.NLP_PATH, name=self.NLP_NAME)
+    def has_training_docs(self):
+        """Whether the task has a training set"""
+        return True if "train" in self.data.keys() else False
+
+    def has_validation_docs(self):
+        """Whether the task has a validation set"""
+        return True if "validation" in self.data.keys() else False
+
+    def has_test_docs(self):
+        """Whether the task has a test set"""
+        return True if "test" in self.data.keys() else False
 
     def training_docs(self):
         # Cache training for faster few-shot.
         # If data is too large to fit in memory, override this method.
         if self.has_training_docs():
             if self._training_docs is None:
-                self._training_docs = list(self._load_nlp_dataset()["train"])
+                self._training_docs = list(self.data["train"])
             return self._training_docs
 
     def validation_docs(self):
         if self.has_validation_docs():
-            return self._load_nlp_dataset()["validation"]
+            return self.data["validation"]
 
     def test_docs(self):
         if self.has_test_docs():
-            return self._load_nlp_dataset()["test"]
+            return self.data["test"]
 
     def fewshot_examples(self, k):
         training_docs = self.training_docs()

lm_eval/tasks/coqa.py

 import json
 import random
 from lm_eval.base import Dataset
+from ..utils import sh
 
 
 class CoQA(Dataset):
+
+    def download(self):
+        sh("""
+            mkdir -p data/coqa 
+            wget http://downloads.cs.stanford.edu/nlp/data/coqa/coqa-train-v1.0.json -O data/coqa/coqa-train-v1.0.json
+            wget http://downloads.cs.stanford.edu/nlp/data/coqa/coqa-dev-v1.0.json -O data/coqa/coqa-dev-v1.0.json
+            """)
+
     def has_training_docs(self):
         return True
 

lm_eval/utils.py

+import os
+
+
+class ExitCodeError(Exception):
+    pass
+
+
+def sh(x):
+    if os.system(x):
+        raise ExitCodeError()
+
+
 def simple_parse_args_string(args_string):
     """
     Parses something like