diff --git a/reasoning_gym/cognition/rectangle_count.py b/reasoning_gym/cognition/rectangle_count.py index 959dc25f..5eecd8bf 100644 --- a/reasoning_gym/cognition/rectangle_count.py +++ b/reasoning_gym/cognition/rectangle_count.py @@ -4,6 +4,42 @@ from typing import Dict, Optional from ..factory import ProceduralDataset, register_dataset +QUESTION_TEMPLATE = """Your task is to count how many rectangles are present in an ASCII grid. + +Single rectangles are outlined with a '#', overlapping rectangles (max 2) are shown with '█'. + +Example: +- Input: How many rectangles are in the grid below? + + #### + # # + #### + + + + + + + + + + + ######### + # █## + # █ # + ########█ # + # # + ### +- Output: 3 +- Explanation: + - The first rectangle is the 3x4 rectangle in the top right. + - The other two rectangles are overlapping in the bottom left corner. + - Therefore, the final answer is 3. + +Now, it's your turn. How many rectangles do you see in the grid below? +{puzzle} +""" + def draw_rectangles_with_overlap(n, width, height, rng): # Create a grid that holds a count of how many times a cell is drawn. @@ -103,12 +139,10 @@ class RectangleCountDataset(ProceduralDataset): target = rng.randint(1, self.config.max_rectangles) puzzle, answer = draw_rectangles_with_overlap(target, self.config.width, self.config.height, rng) - puzz = f"How many rectangles do you see? Single rectangles are outlined with a '#', overlapping rectangles (max 2) are shown with '█'. \n\n {puzzle}" - return { - "question": puzz, + "question": QUESTION_TEMPLATE.format(puzzle=puzzle), "answer": str(answer), - "metadata": {}, + "metadata": {"puzzle": puzzle, "solution": answer}, } def score_answer(self, answer: Optional[str], entry: Dict[str, any]) -> float: