Merge pull request #833 from EleutherAI/fix-ppl

[Refactor] Fix wikitext task

lm_eval/models/huggingface.py

@@ -296,14 +296,14 @@ class HFLM(LM):
                     )
             else:
                 assert accelerator.distributed_type in [
-                    DistributedType.FSDP, 
-                    DistributedType.MULTI_GPU
+                    DistributedType.FSDP,
+                    DistributedType.MULTI_GPU,
                 ], "Unsupported distributed type provided. Only DDP and FSDP are supported."
                 if accelerator.distributed_type == DistributedType.FSDP:
                     self._model = accelerator.prepare(self.model)
                 else:
                     self._model = accelerator.prepare_model(
-                        self.model, evaluation_mode=True 
+                        self.model, evaluation_mode=True
                     )
                 self._device = torch.device(f"cuda:{accelerator.local_process_index}")
                 self.accelerator = accelerator

lm_eval/tasks/wikitext/preprocess_wikitext.py

@@ -34,3 +34,15 @@ def wikitext_detokenizer(doc):
     string = string.replace(" 's", "'s")
 
     return string
+
+
+def process_results(doc, results):
+    (loglikelihood,) = results
+    # IMPORTANT: wikitext counts number of words in *original doc before detokenization*
+    _words = len(re.split(r"\s+", doc["page"]))
+    _bytes = len(doc["page"].encode("utf-8"))
+    return {
+        "word_perplexity": (loglikelihood, _words),
+        "byte_perplexity": (loglikelihood, _bytes),
+        "bits_per_byte": (loglikelihood, _bytes),
+    }

lm_eval/tasks/wikitext/wikitext.yaml

@@ -7,6 +7,7 @@ validation_split: validation
 test_split: test
 doc_to_text: ""
 doc_to_target: !function preprocess_wikitext.wikitext_detokenizer
+process_results: !function preprocess_wikitext.process_results
 should_decontaminate: true
 doc_to_decontamination_query: "{{page}}"
 metric_list: