refactor MathQA as MultipleChoiceTask, uses full words for Q. and A.

lm_eval/tasks/headqa.py

@@ -19,7 +19,7 @@ class HeadQA(HFTask):
         return ""
 
     def doc_to_text(self, doc):
-        return "Q: " + doc['qtext'] + '\nA:'
+        return "Question: " + doc['qtext'] + '\nAnswer:'
 
     def doc_to_target(self, doc):
         # this picks one answer to be the "correct" one, despite sometimes 

lm_eval/tasks/mathqa.py

 from . common import HFTask
-from lm_eval.base import mean, rf
+from lm_eval.base import mean, rf, MultipleChoiceTask
 
-class MathQA(HFTask):
+
+class MathQA(HFTask, MultipleChoiceTask):
     DATASET_PATH = "math_qa"
     DATASET_NAME = None
 
@@ -14,60 +15,34 @@ class MathQA(HFTask):
     def has_test_docs(self):
         return True
 
-    def fewshot_description(self):
-        # TODO: figure out description
-        return ""
-
-    def doc_to_text(self, doc):
-        return "Q: " + doc['Problem'] + '\nA:'
-
-    def doc_to_target(self, doc):
-        # this picks one answer to be the "correct" one, despite sometimes 
-        # multiple correct answers being possible.
-        # TODO: make sure we're actually handling multi-answer correctly
-        return " " + doc['correct']
-        
-    def _remove_prefixes(self, aliases):
-        # Optimization: Remove any alias that has a strict prefix elsewhere in the list
-        # we can do this because if the prefix is acceptable by isgreedy, we can stop looking
-        aliases.sort()
-        ret = [aliases[0]]
-        for alias in aliases[1:]:
-            if not alias.startswith(ret[-1]):
-                ret.append(alias)
-
-        return ret
-        
+    def _convert_standard(self, doc):
 
-    def construct_requests(self, doc, ctx):
-
-        self.answer_options = ['a', 'b', 'c', 'd', 'e']
-        
-        ret = []
-        for i in range(len(self.answer_options)):
-            ll, _ =rf.loglikelihood(ctx, ' ' + self.answer_options[i])
-            ret.append(ll)
+        out_doc = {
+            "query": "Question: " + doc['Problem'] +" "+ doc["options"] + "\nAnswer:",
+            "choices": ['a', 'b', 'c', 'd', 'e'],
+            "gold": ['a', 'b', 'c', 'd', 'e'].index(doc['correct']),
+        }
+        return out_doc
 
-        return ret
+    def _load_docs(self, docs):
+        for record in docs:
+            yield self._convert_standard(record)
 
-    def process_results(self, doc, results):
-        max_result_idx = max(enumerate(results), key=lambda x: x[1])[0]
+    def training_docs(self):
+        docs = super().training_docs()
+        return self._load_docs(docs)
 
-        if doc['correct'] == self.answer_options[max_result_idx]:
-            result = 1.0
-        else:
-            result = 0.0
+    def validation_docs(self):
+        docs = super().validation_docs()
+        return self._load_docs(docs)
 
-        return {
-            "acc": result
-        }
+    def test_docs(self):
+        docs = super().test_docs()
+        return self._load_docs(docs)
 
-    def aggregation(self):
-        return {
-            "acc": mean,
-        }
+    def fewshot_description(self):
+        # TODO: figure out description
+        return ""
 
-    def higher_is_better(self):
-        return {
-            "acc": True
-        }
+    def doc_to_text(self, doc):
+        return doc["query"]
\ No newline at end of file