Fixes SQuAD v2 metric computation.

232c9ab6 · Charles Foster · 4a64031c · 232c9ab6
Commit 232c9ab6 authored Feb 13, 2021 by Charles Foster
Show whitespace changes
Inline Side-by-side

Showing with 9 additions and 4 deletions

lm_eval/tasks/squad.py lm_eval/tasks/squad.py +9 -4

No files found.
--- a/lm_eval/tasks/squad.py
+++ b/lm_eval/tasks/squad.py
@@ -64,15 +64,20 @@ class SQuAD(HFTask):

        continuation, = results

-        predictions = {
+        no_answer_probability = 0.0
+        if continuation.startswith(' unanswerable'):
+            no_answer_probability = 1.0
+
+        predictions = [{
            'id': doc['id'],
            'prediction_text': continuation,
-        }
+            'no_answer_probability': no_answer_probability,
+        }]

-        references = {
+        references = [{
            'id': doc['id'],
            'answers': doc['answers'],
-        }
+        }]

        metrics = squad_metric.compute(predictions=predictions, references=references)