Fix MBPP+ grader scoring 0.0 (unconditional check(entry_point) call)

src/lighteval/tasks/tasks/evalplus/execution.py

@@ -74,11 +74,29 @@ def _swallow_io():
         yield
 
 
+_CHECK_DEF = re.compile(r"^\s*def\s+check\s*\(", re.MULTILINE)
+
+
 def _build_program(sample: dict, generation: str) -> str:
-    """Assemble the program: model code + test harness + a check(entry_point) call."""
+    """Assemble the executable program: model code followed by the test harness.
+
+    The two EvalPlus benchmarks structure their ``test`` field differently:
+
+      - HumanEval+ defines ``def check(candidate): ...`` and must be *invoked*
+        as ``check(entry_point)``.
+      - MBPP+ has no ``check`` wrapper; its assertions call the target function
+        directly at module scope, so they execute as soon as the test runs.
+
+    Appending ``check(entry_point)`` unconditionally raises ``NameError`` for
+    every MBPP+ problem (there is no ``check`` to call), scoring the whole task
+    0. Only emit the invocation when the harness actually defines ``check``.
+    """
     harness = sample.get("test", "")
     entry = sample.get("entry_point", "")
-    return f"{generation}\n\n{harness}\n\ncheck({entry})\n"
+    program = f"{generation}\n\n{harness}\n"
+    if entry and _CHECK_DEF.search(harness):
+        program += f"\ncheck({entry})\n"
+    return program
 
 
 def _unsafe_execute(program: str, timeout: float) -> bool:

tests/unit/tasks/test_evalplus.py

@@ -7,15 +7,21 @@
 from lighteval.tasks.tasks.evalplus import execution
 
 
-# A tiny HumanEval-style problem: implement add(a, b).
+# A tiny HumanEval-style problem: implement add(a, b). HumanEval+ tests define a
+# ``check(candidate)`` function that must be invoked with the entry point.
 TEST_HARNESS = "def check(candidate):\n    assert candidate(2, 3) == 5\n    assert candidate(0, 0) == 0\n"
 ENTRY_POINT = "add"
 CORRECT = "def add(a, b):\n    return a + b"
 WRONG = "def add(a, b):\n    return a - b"
 
+# The same problem in MBPP+ style: assertions call the target function directly
+# at module scope, with no ``check`` wrapper. Appending ``check(add)`` here would
+# raise NameError and fail every problem (the bug this guards against).
+MBPP_HARNESS = "assert add(2, 3) == 5\nassert add(0, 0) == 0\n"
 
-def _grade(code: str) -> bool:
-    sample = execution.make_sample(TEST_HARNESS, ENTRY_POINT)
+
+def _grade(code: str, harness: str = TEST_HARNESS, entry: str = ENTRY_POINT) -> bool:
+    sample = execution.make_sample(harness, entry)
     result = execution.check_correctness(sample, code, timeout=10)
     return bool(result) and all(x == 1 for x in result)
 
@@ -40,6 +46,29 @@ def test_missing_function_fails():
     assert _grade("x = 1") is False
 
 
+def test_mbpp_style_direct_assert_correct_passes():
+    # MBPP+ harness: module-level asserts, no check() wrapper. Regression test
+    # for the bug where check(entry_point) was appended unconditionally.
+    assert _grade(CORRECT, harness=MBPP_HARNESS) is True
+
+
+def test_mbpp_style_direct_assert_wrong_fails():
+    assert _grade(WRONG, harness=MBPP_HARNESS) is False
+
+
+def test_mbpp_harness_does_not_append_check_call():
+    sample = execution.make_sample(MBPP_HARNESS, ENTRY_POINT)
+    program = execution._build_program(sample, CORRECT)
+    assert "check(add)" not in program
+    assert "assert add(2, 3) == 5" in program
+
+
+def test_humaneval_harness_appends_check_call():
+    sample = execution.make_sample(TEST_HARNESS, ENTRY_POINT)
+    program = execution._build_program(sample, CORRECT)
+    assert "check(add)" in program
+
+
 def test_extract_code_prefers_python_fence():
     text = "Here is the solution:\n```python\ndef add(a, b):\n    return a + b\n```\nDone."
     assert execution.extract_code(text) == "def add(a, b):\n    return a + b"