From 60d9a29d38e429714ae9a2e73972e3b15db7e671 Mon Sep 17 00:00:00 2001 From: Rich Jones Date: Tue, 11 Mar 2025 00:44:48 +0100 Subject: [PATCH] fix dice --- reasoning_gym/arithmetic/dice.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/reasoning_gym/arithmetic/dice.py b/reasoning_gym/arithmetic/dice.py index 82c430d3..b68018e0 100644 --- a/reasoning_gym/arithmetic/dice.py +++ b/reasoning_gym/arithmetic/dice.py @@ -75,14 +75,14 @@ def generate_puzzle(num_dice, max_dice_size, rng): high_target = min_total + 2 * (max_total - min_total) // 3 target = rng.randint(low_target, high_target) - # Compute probability. - (num, den) = compute_probability(dice, target) + # Compute probability. Don't change this. + (num, den), probability = compute_probability(dice, target) # Create a string representing the dice, e.g., "1d20, 1d17, 1d6" etc. dice_str = ", ".join(f"1d{s}" for s in dice) # Return the puzzle. - return {"dice_str": dice_str, "target": target, "num": num, "den": den} + return {"dice_str": dice_str, "target": target, "num": num, "den": den, "probability": probability} @dataclass @@ -127,7 +127,8 @@ class DiceDataset(ProceduralDataset): "difficulty": { "num_dice": self.config.num_dice, "max_dice_size": self.config.max_dice_size, - } + }, + "puzzle": puzzle, }, }