Use `_process_doc` like the other tasks do

7ab64701 · Dirk Groeneveld · a080b5ad · 7ab64701
Commit 7ab64701 authored Apr 21, 2022 by Dirk Groeneveld
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 4 deletions

lm_eval/tasks/wikitext.py lm_eval/tasks/wikitext.py +4 -4

No files found.
--- a/lm_eval/tasks/wikitext.py
+++ b/lm_eval/tasks/wikitext.py
@@ -76,15 +76,15 @@ class WikiText(PerplexityTask):
        return True

    def training_docs(self):
-        return map(self._load_doc, self.dataset["train"])
+        return map(self._process_doc, self.dataset["train"])

    def validation_docs(self):
-        return map(self._load_doc, self.dataset["validation"])
+        return map(self._process_doc, self.dataset["validation"])

    def test_docs(self):
-        return map(self._load_doc, self.dataset["test"])
+        return map(self._process_doc, self.dataset["test"])

-    def _load_doc(self, doc):
+    def _process_doc(self, doc):
        return doc["page"]

    def doc_to_target(self, doc):