diff --git a/opencompass/datasets/mbpp.py b/opencompass/datasets/mbpp.py index 3e4218246..7ccdc0ee7 100644 --- a/opencompass/datasets/mbpp.py +++ b/opencompass/datasets/mbpp.py @@ -344,6 +344,7 @@ def _filtered_get_mbpp_plus(*args, **kwargs): def _process_answer(self, text): patterns = [ + r"^\s*'(.*?)'\s*\[DONE\]", r"\[BEGIN\]\s*'(.*)'\s*\[DONE\]", r"BEGIN\s*'(.*)'\s*\[DONE\]", r"\[BEGIN\]\s*'(.*)'\s*DONE", diff --git a/tests/datasets/test_mbpp.py b/tests/datasets/test_mbpp.py new file mode 100644 index 000000000..19cf976eb --- /dev/null +++ b/tests/datasets/test_mbpp.py @@ -0,0 +1,24 @@ +import unittest + +from opencompass.datasets.mbpp import MBPPEvaluator + + +class TestMBPPEvaluator(unittest.TestCase): + + def test_process_answer_uses_first_generated_program(self): + raw = """ 'def target(x): + return x + 1' +[DONE] + +[BEGIN] + 'def wrong(x): + return x - 1' +[DONE] +""" + + self.assertEqual(MBPPEvaluator()._process_answer(raw), + 'def target(x):\n return x + 1') + + +if __name__ == '__main__': + unittest.main()