From 3fbb18bb308636e8f91971f5376bc8c7a4b72922 Mon Sep 17 00:00:00 2001 From: plaguss Date: Wed, 6 Nov 2024 16:49:54 +0100 Subject: [PATCH] Add passing tests --- .../tasks/math_shepherd/test_completer.py | 257 ++++++++++++++++++ .../tasks/math_shepherd/test_generator.py | 72 ++--- 2 files changed, 296 insertions(+), 33 deletions(-) create mode 100644 tests/unit/steps/tasks/math_shepherd/test_completer.py diff --git a/tests/unit/steps/tasks/math_shepherd/test_completer.py b/tests/unit/steps/tasks/math_shepherd/test_completer.py new file mode 100644 index 000000000..88e5119e1 --- /dev/null +++ b/tests/unit/steps/tasks/math_shepherd/test_completer.py @@ -0,0 +1,257 @@ +# Copyright 2023-present, Argilla, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import textwrap +from typing import TYPE_CHECKING, Any, Dict, List + +import pytest + +from distilabel.steps.tasks.math_shepherd.completer import MathShepherdCompleter +from tests.unit.conftest import DummyLLM + +if TYPE_CHECKING: + from distilabel.models.llms.typing import GenerateOutput + + +class MathShepherdCompleterLLM(DummyLLM): + N: int = 3 + + def load(self) -> None: + pass + + @property + def model_name(self) -> str: + return "math-shepherd-completer" + + def generate( + self, inputs: Dict[str, Any], num_generations: int = 1 + ) -> List["GenerateOutput"]: + if self.N == 1: + response = textwrap.dedent(""" + Step 1: Determine the total number of eggs Janet collects per day: Janet's ducks lay 16 eggs per day. + Step 2: Calculate the number of eggs Janet uses for herself per day: She eats three for breakfast and bakes muffins with four eggs, for a total of 3 + 4 = <<3+4=7>>7 eggs. + Step 3: Calculate the number of eggs Janet has left to sell per day: 16 - 7 = <<16-7=9>>9 eggs. + Step 4: Calculate the total amount Janet makes at the farmers' market per day: 9 * $2 = <<9*2=18>>18 dollars. + + The answer is: $18""") + else: + response = textwrap.dedent(""" + Step 2: Janet's ducks lay 16 eggs per day, and she uses 7 for eating and baking. So the number of eggs she has left is 16 - 7 = <<16-7=9>>9. + Step 3: Janet sells the remaining 9 eggs for $2 each, so she makes 9 * 2 = <<9*2=18>>18 dollars every day at the farmers' market. + The answer is: 18 + + --- + + Step 2: Janet's ducks lay 16 eggs per day, and she uses 3 for eating and bakes 4 for her friends, so she has 16 - 7 = <<16-7=9>>9 eggs left. + Step 3: Selling the 9 eggs at $2 each, she makes 9 * 2 = <<9*2=18>>18 dollars every day. + The answer is: 18 + + --- + + Step 2: Janets ducks lay 16 eggs per day. She eats 3 and bakes 4, so she has 16 - (3 + 4) = 16 - 7 = 9 eggs left. + Step 3: She sells the 9 eggs for $2 each, which means she makes 9 * $2 = $<<9*2=18>>18. + The answer is: 18""") + return [[response for _ in range(num_generations)] for _ in range(len(inputs))] + + +DUMMY_STEPS = [ + "Step 1: Determine the total number of eggs Janet collects per day: Janet's ducks lay 16 eggs per day.", + "Step 2: Calculate the number of eggs Janet uses for herself per day: She eats three for breakfast and bakes muffins with four eggs, for a total of 3 + 4 = <<3+4=7>>7 eggs.", + "Step 3: Calculate the number of eggs Janet has left to sell per day: 16 - 7 = <<16-7=9>>9 eggs.", + "Step 4: Calculate the total amount Janet makes at the farmers' market per day: 9 * $2 = <<9*2=18>>18 dollars.", + "The answer is: $18", +] + + +class TestMathShepherdCompleter: + @pytest.mark.parametrize( + "steps, num_completions", + [ + (DUMMY_STEPS, 3), + # This would be the same case as having the problem already solved in a single step, + # there's nothing else we have to do + (DUMMY_STEPS[-2:], 0), + ], + ) + def test_prepare_completions(self, steps: List[str], num_completions: int) -> None: + task = MathShepherdCompleter(llm=MathShepherdCompleterLLM(N=1), N=1) + task.load() + instruction = "Krystian works in the library. He borrows an average of 40 books every day. Every Friday, his number of borrowed books is about 40% higher than the daily average. How many books does he borrow in a week if the library is open from Monday to Friday?" + prepared_inputs = task._prepare_completions(instruction, steps) + assert len(prepared_inputs) == num_completions + + def test_process(self) -> None: + task = MathShepherdCompleter( + llm=MathShepherdCompleterLLM(N=3), + N=3, + ) + task.load() + result = next( + task.process( + [ + { + "instruction": "Janet’s ducks lay 16 eggs per day. She eats three for breakfast every morning and bakes muffins for her friends every day with four. She sells the remainder at the farmers' market daily for $2 per fresh duck egg. How much in dollars does she make every day at the farmers' market?", + "golden_solution": json.dumps( + [ + "Step 1: Janet sells 16 - 3 - 4 = <<16-3-4=9>>9 duck eggs a day.", + "Step 2: She makes 9 * 2 = $<<9*2=18>>18 every day at the farmer’s market.", + "The answer is: 18", + ] + ), + "solutions": json.dumps( + [ + [ + "Step 1: Janet sells 16 - 3 - 4 = <<16-3-4=9>>9 duck eggs a day.", + "Step 2: She makes 9 * 2 = $<<9*2=18>>18 every day at the farmer’s market.", + "The answer is: 18", + ], + [ + "Step 1: Janets ducks lay 16 eggs per day, and she uses 3 + 4 = <<3+4=7>>7 for eating and baking.", + "Step 2: So she sells 16 - 7 = <<16-7=9>>9 duck eggs every day.", + "Step 3: Those 9 eggs are worth 9 * $2 = $<<9*2=18>>18.", + "The answer is: 18", + ], + ] + ), + }, + ] + ) + ) + solutions = result[0]["solutions"] + solutions = json.loads(solutions) + assert solutions == [ + [ + "Step 1: Janet sells 16 - 3 - 4 = <<16-3-4=9>>9 duck eggs a day. +", + "Step 2: She makes 9 * 2 = $<<9*2=18>>18 every day at the farmer\u2019s market. The answer is: 18 +", + ], + [ + "Step 1: Janets ducks lay 16 eggs per day, and she uses 3 + 4 = <<3+4=7>>7 for eating and baking. +", + "Step 2: So she sells 16 - 7 = <<16-7=9>>9 duck eggs every day. +", + "Step 3: Those 9 eggs are worth 9 * $2 = $<<9*2=18>>18. The answer is: 18 +", + ], + ] + + def test_auto_label(self): + inputs = [ + { + "instruction": "Janet’s ducks lay 16 eggs per day. She eats three for breakfast every morning and bakes muffins for her friends every day with four. She sells the remainder at the farmers' market daily for $2 per fresh duck egg. How much in dollars does she make every day at the farmers' market?", + "golden_solution": json.dumps( + [ + "Step 1: Janet sells 16 - 3 - 4 = <<16-3-4=9>>9 duck eggs a day.", + "Step 2: She makes 9 * 2 = $<<9*2=18>>18 every day at the farmer’s market.", + "The answer is: 18", + ] + ), + "solutions": json.dumps( + [ + [ + "Step 1: Janet sells 16 - 3 - 4 = <<16-3-4=9>>9 duck eggs a day.", + "Step 2: She makes 9 * 2 = $<<9*2=18>>18 every day at the farmer’s market.", + "The answer is: 18", + ], + [ + "Step 1: Janets ducks lay 16 eggs per day, and she uses 3 + 4 = <<3+4=7>>7 for eating and baking.", + "Step 2: So she sells 16 - 7 = <<16-7=9>>9 duck eggs every day.", + "Step 3: Those 9 eggs are worth 9 * $2 = $<<9*2=18>>18.", + "The answer is: 18", + ], + ] + ), + }, + ] + N = 3 + task = MathShepherdCompleter( + llm=MathShepherdCompleterLLM(N=N), + N=N, + ) + task.load() + final_outputs = [ + [ + [ + "Step 2: Janet sells 9 duck eggs at the farmers' market, so she makes 9 * $1 = $<<9*1=9>>9 from selling the eggs.", + "The answer is: $9", + ], + [ + "Step 1: Janet lays 16 eggs per day, eats 3 for breakfast, uses 4 for baking, so she has 16 - 3 - 4 = 9 eggs left.", + "Step 2: Since Janet sells 9 eggs a day, and each egg is sold for $1, she makes 9 * $1 = $<<9*1=9>>9.", + "The answer is: $9", + ], + [ + "Step 1: Janet lays 16 eggs per day, eats 3, uses 4 for baking which leaves her with 16 - 3 - 4 = 9 eggs.", + "Step 2: Since she sells the eggs for $1 each, she makes 9 * $1 = $<<9*1=9>>9.", + "The answer is: $9", + ], + ], + [ + [ + "Step 3: To determine how many eggs Jan's sells at the market, we need to subtract the eggs she uses (7) from the total number of eggs laid (16), which is 16 - 7 = <<16-7=9>>9.", + "Step 4: Since she sells 9 eggs for $2 each, we multiply 9 * 2 = <<9*2=18>>18 to find out her daily earnings.", + "The answer is: 18", + ], + [ + "Step 2: Jan's ducks lay 16 eggs per day, and she uses 3 + 4 = <<3+4=7>>7 for eating and baking.", + "Step 3: To find the number of eggs Jan's sells at the market, we subtract the eggs she uses (7) from the total number of eggs laid (16), which is 16 - 7 = <<16-7=9>>9.", + "Step 4: Since she sells 9 eggs for $2 each, we multiply 9 * 2 = <<9*2=18>>18 to find out her daily earnings.", + "The answer is: 18", + ], + [ + "Step 2: Jan's ducks lay 16 eggs per day, and she uses 7 for eating and baking.", + "Step 3: To find the number of eggs Jan's sells at the market, we calculate 16 - 7 = <<16-7=9>>9.", + "Step 4: Since she sells 9 eggs for $2 each, we multiply 9 * 2 = <<9*2=18>>18 to find out her daily earnings.", + "The answer is: 18", + ], + ], + [ + [ + "Step 1: Janet's ducks lay 16 eggs per day. She eats 3 eggs and bakes 4 eggs.", + "Step 2: So, she uses 3 + 4 = <<3+4=7>>7 eggs for eating and baking.", + "Step 3: She sells the remaining eggs, which is 16 - 7 = <<16-7=9>>9 duck eggs every day.", + "Step 4: She sells each egg for $2, so the total amount she makes is 9 * 2 = <<9*2=18>>18 dollars every day.", + "The answer is: 18", + ], + [ + "Step 1: Janet's ducks lay 16 eggs per day.", + "Step 2: She eats 3 eggs and bakes 4 eggs, which is a total of 3 + 4 = <<3+4=7>>7 eggs.", + "Step 3: She sells the remaining eggs, which is 16 - 7 = <<16-7=9>>9 duck eggs every day.", + "Step 4: Since she sells each egg for $2, she makes 9 * 2 = <<9*2=18>>18 dollars every day.", + "The answer is: 18", + ], + [ + "Step 1: Janet's ducks lay 16 eggs per day.", + "Step 2: She consumes 7 eggs for eating and baking, which means she has 16 - 7 = <<16-7=9>>9 eggs left.", + "Step 3: She sells each egg for $2, so she makes 9 * 2 = <<9*2=18>>18 dollars every day.", + "The answer is: 18", + ], + ], + ] + + golden_answers = ["The answer is: 18", "The answer is: 18"] + input_positions = [(0, 0, 0), (0, 1, 0), (0, 1, 1)] + results = task._auto_label( + inputs, final_outputs, input_positions, golden_answers + ) + solutions = json.loads(results[0]["solutions"]) + + assert solutions == [ + [ + "Step 1: Janet sells 16 - 3 - 4 = <<16-3-4=9>>9 duck eggs a day. -", + "Step 2: She makes 9 * 2 = $<<9*2=18>>18 every day at the farmer’s market. The answer is: 18 +", + ], + [ + "Step 1: Janets ducks lay 16 eggs per day, and she uses 3 + 4 = <<3+4=7>>7 for eating and baking. +", + "Step 2: So she sells 16 - 7 = <<16-7=9>>9 duck eggs every day. +", + "Step 3: Those 9 eggs are worth 9 * $2 = $<<9*2=18>>18. The answer is: 18 +", + ], + ] diff --git a/tests/unit/steps/tasks/math_shepherd/test_generator.py b/tests/unit/steps/tasks/math_shepherd/test_generator.py index 01b7b24c6..4f8fda982 100644 --- a/tests/unit/steps/tasks/math_shepherd/test_generator.py +++ b/tests/unit/steps/tasks/math_shepherd/test_generator.py @@ -13,7 +13,7 @@ # limitations under the License. import json -from typing import Any, Dict, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional import pytest @@ -25,8 +25,13 @@ ) from tests.unit.conftest import DummyLLM +if TYPE_CHECKING: + from distilabel.models.llms.typing import GenerateOutput + class MathShepherdGeneratorLLM(DummyLLM): + M: Optional[int] = None + def load(self) -> None: pass @@ -36,7 +41,7 @@ def model_name(self) -> str: def generate( self, inputs: Dict[str, Any], num_generations: int = 1 - ) -> Union[str, None]: + ) -> List["GenerateOutput"]: response = """ Step 1: Calculate the number of books borrowed on a regular day (Monday to Thursday): 40 books per day @@ -52,6 +57,8 @@ def generate( 160 + 56 = <<160+56=216>>216 books The answer is: 216 books.""" + if self.M: + response = "---".join([response for _ in range(self.M)]) return [[response for _ in range(num_generations)]] @@ -70,23 +77,23 @@ class TestMathShepherdGenerator: (RULES_GSM8K, FEW_SHOTS_GSM8K), ], ) - @pytest.mark.parametrize("include_errors", [True, False]) - @pytest.mark.parametrize("N", [1, 3]) + @pytest.mark.parametrize( + "M", + [None, 5], + ) def test_format_input( self, system_prompt: Optional[str], extra_rules: Optional[str], few_shots: Optional[str], - include_errors: bool, - N: int, + M: Optional[int], ) -> None: task = MathShepherdGenerator( llm=MathShepherdGeneratorLLM(), system_prompt=system_prompt, extra_rules=extra_rules, few_shots=few_shots, - include_errors=include_errors, - N=N, + M=M, ) task.load() @@ -96,40 +103,39 @@ def test_format_input( } ) rendered_system_prompt = "" - if not system_prompt: - user_prompt = result[0]["content"] - else: + if system_prompt: rendered_system_prompt = result[0]["content"] - user_prompt = result[1]["content"] if extra_rules: assert RULES_GSM8K in rendered_system_prompt if few_shots: assert FEW_SHOTS_GSM8K in rendered_system_prompt - - if include_errors: - if N == 1: - assert f"Generate {N} example solution" in user_prompt - else: - assert ( - f"Generate {N} example solutions to the same problem, separated by a single `---`" - in user_prompt - ) - if system_prompt: + if M: assert ( - "Include errors to help students learn from their mistakes in any of the steps, including the final answer." - in rendered_system_prompt + "Generate 5 example solutions to the same problem" + in result[1]["content"] ) + else: + if M: assert ( - ", but trick them by including unannounced errors in the reasoning without explaining them" - in rendered_system_prompt + "Generate 5 example solutions to the same problem" + in result[0]["content"] ) - def test_process(self) -> None: - task = MathShepherdGenerator( - llm=MathShepherdGeneratorLLM(), - ) + @pytest.mark.parametrize( + "M, output_name", + [ + (None, "golden_solution"), + (3, "solutions"), + ], + ) + def test_process(self, M: Optional[int], output_name: str) -> None: + task = MathShepherdGenerator(llm=MathShepherdGeneratorLLM(M=M), M=M) task.load() - result = next(task.process([{"instruction": ""}]))[0]["steps"] - result = json.loads(result)[0] - assert len(result) == 5 + result = next(task.process([{"instruction": ""}]))[0][output_name] + result = json.loads(result) + if M: + assert len(result) == 3 + assert all(len(r) == 5 for r in result) + else: + assert len(result) == 5