add prompts and task class

src/tot/prompts/math_reasoning.py

@@ -0,0 +1,140 @@
+standard_prompt = '''Solve the following math problem.
+
+<Example>
+Question: Let \[f(x) = \left\{\n\begin{array}{cl} ax+3, &\text{ if }x>2, \\\nx-5 &\text{ if } -2 \le x \le 2, \\\n2x-b &\text{ if } x <-2.\n\end{array}\n\right.\]Find $a+b$ if the piecewise function is continuous (which means that its graph can be drawn without lifting your pencil from the paper).
+Answer: 0.
+</Example>
+
+<Task>
+Question: {question}
+Answer: 
+</Task>
+'''
+
+cot_prompt = r'''Solve the following math problem. Perform the solution by going step by step, and write down each step.
+
+<Example>
+Question: Let \[f(x) = \left\{\n\begin{array}{cl} ax+3, &\text{ if }x>2, \\\nx-5 &\text{ if } -2 \le x \le 2, \\\n2x-b &\text{ if } x <-2.\n\end{array}\n\right.\]Find $a+b$ if the piecewise function is continuous (which means that its graph can be drawn without lifting your pencil from the paper). Let's think step by step.
+Solution: For the piecewise function to be continuous, the cases must "meet" at $2$ and $-2$. For example, $ax+3$ and $x-5$ must be equal when $x=2$. This implies $a(2)+3=2-5$, which we solve to get $2a=-6 \Rightarrow a=-3$. Similarly, $x-5$ and $2x-b$ must be equal when $x=-2$. Substituting, we get $-2-5=2(-2)-b$, which implies $b=3$. So $a+b=-3+3=\boxed{0}$.
+</Example>
+
+<Example>
+Question: Sixteen is 64$\%$ of what number? Let's think step by step.
+Solution: If the number is $x$, we can set up the equation $\frac{16}{x}=\frac{64}{100}$. We divide both sides by $4$ to get $\frac{1}{x}=\frac{4}{100}=\frac{1}{25}$, so $x=\boxed{25}$.
+</Example>
+
+<Example>
+Question: There are 3 complex numbers $a+bi$, $c+di$, and $e+fi$. If $b=1$, $e=-a-c$, and the sum of the numbers is $-i$, find $d+f$. Let's think step by step.
+Solution: We know that $a+bi+c+di+e+fi=-i$. Thus, the real parts add up to 0 and the imaginary parts add up to -1. We then have  \begin{align}\na+c+e&=0\\\nb+d+f&=-1\\\n\end{align}We know that $b=1$, therefore $d+f=\boxed{-2}$
+</Example>
+
+<Task>
+Question: {question} Let's think step by step.
+Solution: 
+</Task>
+'''
+
+propose_first_step_prompt = '''Write the first step to solve the following math problem. Only write down the first step. Do not write down the answer. Do not write down the question. Do not continue the solution. Do not write down the second step.
+<Example>
+Question: Let \[f(x) = \left\{\n\begin{array}{cl} ax+3, &\text{ if }x>2, \\\nx-5 &\text{ if } -2 \le x \le 2, \\\n2x-b &\text{ if } x <-2.\n\end{array}\n\right.\]Find $a+b$ if the piecewise function is continuous (which means that its graph can be drawn without lifting your pencil from the paper).
+First step: For the piecewise function to be continuous, the cases must "meet" at $2$ and $-2$.
+</Example>
+
+<Example>
+Question: Sixteen is 64$\%$ of what number? 
+First step: If the number is $x$, we can set up the equation $\frac{16}{x}=\frac{64}{100}$.
+</Example>
+
+<Example>
+Question: There are 3 complex numbers $a+bi$, $c+di$, and $e+fi$. If $b=1$, $e=-a-c$, and the sum of the numbers is $-i$, find $d+f$.
+First step:We know that $a+bi+c+di+e+fi=-i$. Thus, the real parts add up to 0 and the imaginary parts add up to -1.
+</Example>
+
+<Task>
+Question: {question}
+First step: 
+</Task>
+'''
+
+propose_next_step_prompt = '''Write the next step to solve the following math problem. If the solution is reached, write "Problem solved" at the end. If final answer is not yet reached, write down the next step. 
+<Example>
+Question: Let \[f(x) = \left\{\n\begin{array}{cl} ax+3, &\text{ if }x>2, \\\nx-5 &\text{ if } -2 \le x \le 2, \\\n2x-b &\text{ if } x <-2.\n\end{array}\n\right.\]Find $a+b$ if the piecewise function is continuous (which means that its graph can be drawn without lifting your pencil from the paper).
+Steps until now: For the piecewise function to be continuous, the cases must "meet" at $2$ and $-2$.
+Next step: For example, $ax+3$ and $x-5$ must be equal when $x=2$. This implies $a(2)+3=2-5$, which we solve to get $2a=-6 \Rightarrow a=-3$.
+</Example>
+
+<Example>
+Question: Sixteen is 64$\%$ of what number? 
+Steps until now: If the number is $x$, we can set up the equation $\frac{16}{x}=\frac{64}{100}$.
+Next step: We divide both sides by $4$ to get $\frac{1}{x}=\frac{4}{100}=\frac{1}{25}$, so $x=\boxed{25}$. Problem solved.
+</Example>
+
+<Example>
+Question: There are 3 complex numbers $a+bi$, $c+di$, and $e+fi$. If $b=1$, $e=-a-c$, and the sum of the numbers is $-i$, find $d+f$.
+Steps until now: We know that $a+bi+c+di+e+fi=-i$. Thus, the real parts add up to 0 and the imaginary parts add up to -1.
+Next step: We then have  \begin{align}\na+c+e&=0\\\nb+d+f&=-1\\\n\end{align}We know that $b=1$, therefore $d+f=\boxed{-2}$. Problem solved.
+</Example>
+
+<Task>
+Question: {question}
+Steps until now: {steps}
+Next step: 
+</Task>
+'''
+
+propose_final_step_prompt = ''' Extract the final answer given the question and solution steps.
+<Example>
+Question: Let \[f(x) = \left\{\n\begin{array}{cl} ax+3, &\text{ if }x>2, \\\nx-5 &\text{ if } -2 \le x \le 2, \\\n2x-b &\text{ if } x <-2.\n\end{array}\n\right.\]Find $a+b$ if the piecewise function is continuous (which means that its graph can be drawn without lifting your pencil from the paper). Let's think step by step.
+Solution: For the piecewise function to be continuous, the cases must "meet" at $2$ and $-2$. For example, $ax+3$ and $x-5$ must be equal when $x=2$. This implies $a(2)+3=2-5$, which we solve to get $2a=-6 \Rightarrow a=-3$. Similarly, $x-5$ and $2x-b$ must be equal when $x=-2$. Substituting, we get $-2-5=2(-2)-b$, which implies $b=3$. So $a+b=-3+3=\boxed{0}$.
+Final answer: 0
+</Example>
+
+<Example>
+Question: Sixteen is 64$\%$ of what number? Let's think step by step.
+Solution: If the number is $x$, we can set up the equation $\frac{16}{x}=\frac{64}{100}$. We divide both sides by $4$ to get $\frac{1}{x}=\frac{4}{100}=\frac{1}{25}$, so $x=\boxed{25}$.
+Final answer: 25
+</Example>
+
+<Example>
+Question: There are 3 complex numbers $a+bi$, $c+di$, and $e+fi$. If $b=1$, $e=-a-c$, and the sum of the numbers is $-i$, find $d+f$. Let's think step by step.
+Solution: We know that $a+bi+c+di+e+fi=-i$. Thus, the real parts add up to 0 and the imaginary parts add up to -1. We then have  \begin{align}\na+c+e&=0\\\nb+d+f&=-1\\\n\end{align}We know that $b=1$, therefore $d+f=\boxed{-2}$
+Final answer: -2
+</Example>
+
+<Task>
+Question: {question}
+Solution: {solution}
+Final answer: 
+</Task>
+'''
+
+value_prompt = ''' Evaluate if given steps are likely to solve the question. (likely, unlikely, or unsure)
+<Example>
+Question: Let \[f(x) = \left\{\n\begin{array}{cl} ax+3, &\text{ if }x>2, \\\nx-5 &\text{ if } -2 \le x \le 2, \\\n2x-b &\text{ if } x <-2.\n\end{array}\n\right.\]Find $a+b$ if the piecewise function is continuous (which means that its graph can be drawn without lifting your pencil from the paper).
+Steps: For the piecewise function to be continuous, the cases must "meet" at $2$ and $-2$.
+Reasoning: criteria to make piecewise function is correct.
+Evaluation: likely
+</Example>
+
+<Example>
+Question: Sixteen is 64$\%$ of what number? 
+Steps: If the number is $x$, we can set up the equation $\frac{16}{x}=\frac{64}{100}$.
+Reasoning: equation is correct
+Evaluation: likely
+</Example>
+
+<Example>
+Question: There are 3 complex numbers $a+bi$, $c+di$, and $e+fi$. If $b=1$, $e=-a-c$, and the sum of the numbers is $-i$, find $d+f$.
+Steps:We know that $a+bi+c+di+e+fi=-i$. Thus, the real parts add up to 1 and the imaginary parts add up to 0.
+Reasoning: equation is incorrect
+Evaluation: unlikely
+</Example>
+
+<Task>
+Question: {question}
+Steps: {steps}
+Reasoning: 
+Evaluation: 
+</Task>
+'''
+

src/tot/tasks/math_reasoning.py

@@ -0,0 +1,81 @@
+import re
+import os
+import pandas as pd
+from tot.tasks.base import Task, DATA_PATH
+from tot.prompts.math_reasoning import *
+
+
+class MathTask(Task):
+    """
+    Input (x)   : a math problem
+    Output (y)  : answer of the problem
+    Reward (r)  : 0 or 1, depending on whether the answer is correct
+    Input Example:
+        Sixteen is 64$% of what number?
+    Output Example:
+        25
+    """
+
+    def __init__(self, file='mathDAO.csv'):
+        super().__init__()
+        path = os.path.join(DATA_PATH, 'math', file)
+        self.data = list(pd.read_csv(path))
+        self.value_cache = {}
+        self.steps = 4
+        self.stops = ['\n'] * 4
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def get_input(self, idx: int) -> str:
+        return self.data[idx]['question']
+
+    @staticmethod
+    def extract_answer(answer_raw):
+        answer = answer_raw.replace(",", "")
+        answer = [s for s in re.findall(r'-?\d+\.?\d*', answer)]
+        if answer:
+            return answer[0]
+        return "Error:" + answer_raw
+
+    def test_output(self, idx: int, output: str):
+        expression = output.strip().split('\n')[-1].lower().replace('Final answer: ', '')
+        model_answer = self.extract_answer(expression)
+        correct_answer = self.data[idx]['answer']
+        if float(model_answer) == float(correct_answer):
+            return {'r': 1}
+        else:
+            return {'r': 0}
+
+    @staticmethod
+    def standard_prompt_wrap(question: str, y: str = '') -> str:
+        return standard_prompt.format(question=question) + y
+
+    @staticmethod
+    def cot_prompt_wrap(question: str, y: str = '') -> str:
+        return cot_prompt.format(question=question) + y
+
+    @staticmethod
+    def propose_prompt_wrap(question: str, steps: str = '') -> str:
+        if not steps:
+            return propose_first_step_prompt.format(question=question)
+        elif "problem solved" in steps[-1].lower():
+            return propose_final_step_prompt.format(question=question, steps=steps)
+        return propose_next_step_prompt.format(question=question, steps=steps)
+
+    @staticmethod
+    def value_prompt_wrap(question: str, steps: str = '') -> str:
+
+        if "final answer" not in steps[-1].lower():
+            return value_prompt.format(question=question, steps=steps)
+        return ''
+
+    @staticmethod
+    def value_outputs_unwrap(question: str, steps: str, value_outputs: list) -> float:
+        if len(steps) == 4 and "final answer" not in steps[-1].lower():
+            return 0
+
+        value_names = [_.split('\n')[-1].replace('Evaluation: ', '') for _ in value_outputs]
+        value_map = {'unlikely': 0.001, 'unsure': 1, 'likely': 20}  # TODO: ad hoc
+        value = sum(value * value_names.count(name) for name, value in value_map.items())
+        return value