Modified stopping criteria for gpt2

61b57292 · Tian Yun · 0b8cb8b9 · 61b57292
Commit 61b57292 authored Apr 26, 2022 by Tian Yun
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

lm_eval/models/gpt2.py lm_eval/models/gpt2.py +1 -1

No files found.
--- a/lm_eval/models/gpt2.py
+++ b/lm_eval/models/gpt2.py
@@ -139,7 +139,7 @@ class HFLM(BaseLM):
         
        return transformers.StoppingCriteriaList([
            MultitokenEOSCriteria(stopping_criteria_ids, self.tokenizer),
-            EOSCriteria(stopping_criteria_ids)
+            EOSCriteria(self.tokenizer.eos_token)
        ])

    def _model_generate(self, context, max_length, stopping_criteria_ids):