File size: 2,132 Bytes
d91766b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | # Copyright (c) Meta Platforms, Inc. and affiliates.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import os
import json
from parsers import Parser, evaluate_equation, validate_equation
from gsm8k import GSM8KDataset
import warnings
CTD_SYSTEM_PROMPT = (
"Using only the provided numbers, create an arithmetic expression that evaluates to exactly the provided target number. You may use the operations +, -, *, and / as needed, but each number must be used exactly once. Think step-by-step. After reasoning, provide only your final expression inside \\boxed"
+ "{}"
+ " tags without including an equals sign or the target number. For example: \\boxed{a + b * c}"
+ """Respond in the following format:
<reasoning>
Your reasoning here
</reasoning>
<answer>
\\boxed{...}
</answer>"""
)
class CTDDataset(GSM8KDataset):
def __init__(
self,
tokenizer,
num_examples=0,
add_reasoning=True,
system_prompt=CTD_SYSTEM_PROMPT,
subsample=256,
):
if num_examples > 0:
warnings.warn("num_examples must be 0 for Countdown dataset. Overriding num_examples to 0.")
super().__init__(
tokenizer,
0,
add_reasoning,
system_prompt,
subsample,
) # num_examples = always 0
def load_test_dataset(self):
self.dataset = []
cur_path = os.path.dirname(os.path.abspath(__file__))
with open(f"{cur_path}/../dataset/countdown_cd3_test.jsonl", "r") as f:
for line in f:
self.dataset.append(json.loads(line))
print(len(self.dataset), "examples loaded")
def __getitem__(self, idx):
target = int(self.dataset[self.subsample[idx].item()]["output"])
numbers_str = self.dataset[self.subsample[idx].item()]["input"]
numbers = [int(num) for num in numbers_str.split(",")]
question = f"Numbers: {numbers}\nTarget: {target}"
prompt = self.create_prompt(question)
return prompt, question, (numbers, target)
|