Skip to content

Commit

Permalink
Add passing tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@plaguss
plaguss committed Nov 6, 2024
1 parent 5f5d823 commit 3fbb18b
Show file tree
Hide file tree
Showing 2 changed files with 296 additions and 33 deletions.
257 changes: 257 additions & 0 deletions tests/unit/steps/tasks/math_shepherd/test_completer.py
View file
Original file line number Diff line number Diff line change
@@ -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 +",
],
]
72 changes: 39 additions & 33 deletions tests/unit/steps/tasks/math_shepherd/test_generator.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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

Expand All @@ -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
Expand All @@ -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)]]


Expand All @@ -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()

Expand All @@ -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

0 comments on commit 3fbb18b

Please sign in to comment.