import re import os import io import json import copy import regex import pickle import datetime import traceback import numpy as np from tqdm import tqdm from math import isclose from pathlib import Path from contextlib import redirect_stdout from concurrent.futures import TimeoutError from functools import partial import multiprocessing import multiprocess from multiprocess import Pool from typing import List, Tuple, Optional, Type, TypeVar, Any, Iterable, Union, Dict import dateutil.relativedelta from pebble import ProcessPool from timeout_decorator import timeout from word2number import w2n from sympy import simplify, N from sympy.parsing.sympy_parser import parse_expr from sympy.parsing.latex import parse_latex from latex2sympy2 import latex2sympy # from utils.math_examples import get_examples def is_multi_choice(answer): for c in answer: if c not in ["A", "B", "C", "D", "E"]: return False return True def load_jsonl(file: Union[str, Path]) -> Iterable[Any]: with open(file, "r", encoding="utf-8") as f: for line in f: try: yield json.loads(line) except: print("Error in loading:", line) exit() def save_jsonl(samples, save_path): # ensure path folder = os.path.dirname(save_path) os.makedirs(folder, exist_ok=True) with open(save_path, "a", encoding="utf-8") as f: for sample in samples: f.write(json.dumps(sample, ensure_ascii=False) + "\n") print("Saved to", save_path) # EXAMPLES = get_examples() def load_prompt(data_name, prompt_type, num_shots): if not num_shots: return [] if data_name in ["gsm_hard", "svamp", "tabmwp", "asdiv", "mawps"]: data_name = "gsm8k" if data_name in ["math_oai", "hungarian_exam", "math-oai", "aime24", "amc23"]: data_name = "math" if data_name in ["sat_math"]: data_name = "mmlu_stem" if data_name in [ "gaokao2024_I", "gaokao2024_II", "gaokao_math_qa", "gaokao2024_mix", "cn_middle_school", ]: data_name = "gaokao" if prompt_type in ["tool-integrated"]: prompt_type = "tora" return EXAMPLES[data_name][:num_shots] PROMPT_TEMPLATES = { "direct": ("Problem: {input}\nAnswer: ", "{output}", "\n\n"), "o1": ( "### Instruction:\n Return your final response within \\boxed{{}}. {input}\n\n### Think:\n ", "{output}", "\n\n", ), } def construct_prompt(example, data_name, args): if args.adapt_few_shot and data_name in [ "gaokao2024_I", "gaokao2024_II", "gaokao_math_qa", "gaokao2024_mix", "cn_middle_school", ]: demos = load_prompt(data_name, args.prompt_type, 5) else: demos = load_prompt(data_name, args.prompt_type, args.num_shots) prompt_type = args.prompt_type prompt_temp = PROMPT_TEMPLATES[args.prompt_type] splitter = prompt_temp[2] input_template, output_template, splitter = ( prompt_temp[0], prompt_temp[1], prompt_temp[2], ) demo_prompt = splitter.join( [ input_template.format(input=q) + output_template.format(output=a) for q, a in demos ] ) context = input_template.format(input=example["question"]) if len(demo_prompt) == 0 or ( args.adapt_few_shot and example["gt_ans"] not in ["A", "B", "C", "D", "E"] ): full_prompt = context else: full_prompt = demo_prompt + splitter + context return full_prompt.strip(" ") # important! # parser.py def _fix_fracs(string): substrs = string.split("\\frac") new_str = substrs[0] if len(substrs) > 1: substrs = substrs[1:] for substr in substrs: new_str += "\\frac" if len(substr) > 0 and substr[0] == "{": new_str += substr else: try: assert len(substr) >= 2 except: return string a = substr[0] b = substr[1] if b != "{": if len(substr) > 2: post_substr = substr[2:] new_str += "{" + a + "}{" + b + "}" + post_substr else: new_str += "{" + a + "}{" + b + "}" else: if len(substr) > 2: post_substr = substr[2:] new_str += "{" + a + "}" + b + post_substr else: new_str += "{" + a + "}" + b string = new_str return string def _fix_a_slash_b(string): if len(string.split("/")) != 2: return string a = string.split("/")[0] b = string.split("/")[1] try: if "sqrt" not in a: a = int(a) if "sqrt" not in b: b = int(b) assert string == "{}/{}".format(a, b) new_string = "\\frac{" + str(a) + "}{" + str(b) + "}" return new_string except: return string def _fix_sqrt(string): _string = re.sub(r"\\sqrt(\w+)", r"\\sqrt{\1}", string) return _string def convert_word_number(text: str) -> str: try: text = str(w2n.word_to_num(text)) except: pass return text # units mainly from MathQA unit_texts = [ "east", "degree", "mph", "kmph", "ft", "m sqaure", " m east", "sq m", "deg", "mile", "q .", "monkey", "prime", "ratio", "profit of rs", "rd", "o", "gm", "p . m", "lb", "tile", "per", "dm", "lt", "gain", "ab", "way", "west", "a .", "b .", "c .", "d .", "e .", "f .", "g .", "h .", "t", "a", "h", "no change", "men", "soldier", "pie", "bc", "excess", "st", "inches", "noon", "percent", "by", "gal", "kmh", "c", "acre", "rise", "a . m", "th", "π r 2", "sq", "mark", "l", "toy", "coin", "sq . m", "gallon", "° f", "profit", "minw", "yr", "women", "feet", "am", "pm", "hr", "cu cm", "square", "v â € ™", "are", "rupee", "rounds", "cubic", "cc", "mtr", "s", "ohm", "number", "kmph", "day", "hour", "minute", "min", "second", "man", "woman", "sec", "cube", "mt", "sq inch", "mp", "∏ cm ³", "hectare", "more", "sec", "unit", "cu . m", "cm 2", "rs .", "rs", "kg", "g", "month", "km", "m", "cm", "mm", "apple", "liter", "loss", "yard", "pure", "year", "increase", "decrease", "d", "less", "Surface", "litre", "pi sq m", "s .", "metre", "meter", "inch", ] unit_texts.extend([t + "s" for t in unit_texts]) def strip_string(string, skip_unit=False): string = str(string).strip() # linebreaks string = string.replace("\n", "") # right "." string = string.rstrip(".") # remove inverse spaces # replace \\ with \ string = string.replace("\\!", "") # string = string.replace("\\ ", "") # string = string.replace("\\\\", "\\") # matrix string = re.sub(r"\\begin\{array\}\{.*?\}", r"\\begin{pmatrix}", string) string = re.sub(r"\\end\{array\}", r"\\end{pmatrix}", string) string = string.replace("bmatrix", "pmatrix") # replace tfrac and dfrac with frac string = string.replace("tfrac", "frac") string = string.replace("dfrac", "frac") string = ( string.replace("\\neq", "\\ne") .replace("\\leq", "\\le") .replace("\\geq", "\\ge") ) # remove \left and \right string = string.replace("\\left", "") string = string.replace("\\right", "") string = string.replace("\\{", "{") string = string.replace("\\}", "}") # Remove unit: miles, dollars if after is not none _string = re.sub(r"\\text{.*?}$", "", string).strip() if _string != "" and _string != string: # print("Warning: unit not removed: '{}' -> '{}'".format(string, _string)) string = _string if not skip_unit: # Remove unit: texts for _ in range(2): for unit_text in unit_texts: # use regex, the prefix should be either the start of the string or a non-alphanumeric character # the suffix should be either the end of the string or a non-alphanumeric character _string = re.sub(r"(^|\W)" + unit_text + r"($|\W)", r"\1\2", string) if _string != "": string = _string # Remove circ (degrees) string = string.replace("^{\\circ}", "") string = string.replace("^\\circ", "") # remove dollar signs string = string.replace("\\$", "") string = string.replace("$", "") string = string.replace("\\(", "").replace("\\)", "") # convert word number to digit string = convert_word_number(string) # replace "\\text{...}" to "..." string = re.sub(r"\\text\{(.*?)\}", r"\1", string) for key in ["x=", "y=", "z=", "x\\in", "y\\in", "z\\in", "x\\to", "y\\to", "z\\to"]: string = string.replace(key, "") string = string.replace("\\emptyset", r"{}") string = string.replace("(-\\infty,\\infty)", "\\mathbb{R}") # remove percentage string = string.replace("\\%", "") string = string.replace("\%", "") string = string.replace("%", "") # " 0." equivalent to " ." and "{0." equivalent to "{." Alternatively, add "0" if "." is the start of the string string = string.replace(" .", " 0.") string = string.replace("{.", "{0.") # cdot # string = string.replace("\\cdot", "") if ( string.startswith("{") and string.endswith("}") and string.isalnum() or string.startswith("(") and string.endswith(")") and string.isalnum() or string.startswith("[") and string.endswith("]") and string.isalnum() ): string = string[1:-1] # inf string = string.replace("infinity", "\\infty") if "\\infty" not in string: string = string.replace("inf", "\\infty") string = string.replace("+\\inity", "\\infty") # and string = string.replace("and", "") string = string.replace("\\mathbf", "") # use regex to remove \mbox{...} string = re.sub(r"\\mbox{.*?}", "", string) # quote string.replace("'", "") string.replace('"', "") # i, j if "j" in string and "i" not in string: string = string.replace("j", "i") # replace a.000b where b is not number or b is end, with ab, use regex string = re.sub(r"(\d+)\.0*([^\d])", r"\1\2", string) string = re.sub(r"(\d+)\.0*$", r"\1", string) # if empty, return empty string if len(string) == 0: return string if string[0] == ".": string = "0" + string # to consider: get rid of e.g. "k = " or "q = " at beginning if len(string.split("=")) == 2: if len(string.split("=")[0]) <= 2: string = string.split("=")[1] string = _fix_sqrt(string) string = string.replace(" ", "") # \frac1b or \frac12 --> \frac{1}{b} and \frac{1}{2}, etc. Even works with \frac1{72} (but not \frac{72}1). Also does a/b --> \\frac{a}{b} string = _fix_fracs(string) # NOTE: X/Y changed to \frac{X}{Y} in dataset, but in simple cases fix in case the model output is X/Y string = _fix_a_slash_b(string) return string def extract_multi_choice_answer(pred_str): # TODO: SFT models if "Problem:" in pred_str: pred_str = pred_str.split("Problem:", 1)[0] pred_str = pred_str.replace("choice is", "answer is") patt = regex.search(r"answer is \(?(?P[abcde])\)?", pred_str.lower()) if patt is not None: return patt.group("ans").upper() return "placeholder" direct_answer_trigger_for_fewshot = ("choice is", "answer is") def choice_answer_clean(pred: str): pred = pred.strip("\n") # Determine if this is ICL, if so, use \n\n to split the first chunk. ICL = False for trigger in direct_answer_trigger_for_fewshot: if pred.count(trigger) > 1: ICL = True if ICL: pred = pred.split("\n\n")[0] # Split the trigger to find the answer. preds = re.split("|".join(direct_answer_trigger_for_fewshot), pred) if len(preds) > 1: answer_flag = True pred = preds[-1] else: answer_flag = False pred = pred.strip("\n").rstrip(".").rstrip("/").strip(" ").lstrip(":") # Clean the answer based on the dataset tmp = re.findall(r"\b(A|B|C|D|E)\b", pred.upper()) if tmp: pred = tmp else: pred = [pred.strip().strip(".")] if len(pred) == 0: pred = "" else: if answer_flag: # choose the first element in list ... pred = pred[0] else: # choose the last e pred = pred[-1] # Remove the period at the end, again! pred = pred.rstrip(".").rstrip("/") return pred def find_box(pred_str: str): ans = pred_str.split("boxed")[-1] if not ans: return "" if ans[0] == "{": stack = 1 a = "" for c in ans[1:]: if c == "{": stack += 1 a += c elif c == "}": stack -= 1 if stack == 0: break a += c else: a += c else: a = ans.split("$")[0].strip() return a def clean_units(pred_str: str): """Clean the units in the number.""" def convert_pi_to_number(code_string): code_string = code_string.replace("\\pi", "π") # Replace \pi or π not preceded by a digit or } with 3.14 code_string = re.sub(r"(? "3*3.14" code_string = re.sub(r"(\d)(\\?π)", r"\1*3.14", code_string) # Handle cases where π is within braces or followed by a multiplication symbol # This replaces "{π}" with "3.14" directly and "3*π" with "3*3.14" code_string = re.sub(r"\{(\\?π)\}", "3.14", code_string) code_string = re.sub(r"\*(\\?π)", "*3.14", code_string) return code_string pred_str = convert_pi_to_number(pred_str) pred_str = pred_str.replace("%", "/100") pred_str = pred_str.replace("$", "") pred_str = pred_str.replace("¥", "") pred_str = pred_str.replace("°C", "") pred_str = pred_str.replace(" C", "") pred_str = pred_str.replace("°", "") return pred_str def extract_theoremqa_answer(pred: str, answer_flag: bool = True): if any([option in pred.lower() for option in ["yes", "true"]]): pred = "True" elif any([option in pred.lower() for option in ["no", "false"]]): pred = "False" elif any( [ option in pred.lower() for option in ["(a)", "(b)", "(c)", "(d)", "(e)", "(f)"] ] ): pass else: # Some of the models somehow get used to boxed output from pre-training if "boxed" in pred: pred = find_box(pred) if answer_flag: # Extract the numbers out of the string pred = pred.split("=")[-1].strip() pred = clean_units(pred) try: tmp = str(latex2sympy(pred)) pred = str(eval(tmp)) except Exception: if re.match(r"-?[\d\.]+\s\D+$", pred): pred = pred.split(" ")[0] elif re.match(r"-?[\d\.]+\s[^\s]+$", pred): pred = pred.split(" ")[0] else: # desparate search over the last number preds = re.findall(r"-?\d*\.?\d+", pred) if len(preds) >= 1: pred = preds[-1] else: pred = "" return pred def extract_answer(pred_str, data_name, use_last_number=True): pred_str = pred_str.replace("\u043a\u0438", "") if data_name in ["mmlu_stem", "sat_math", "aqua", "gaokao2023"]: # TODO check multiple choice return choice_answer_clean(pred_str) if "final answer is $" in pred_str and "$. I hope" in pred_str: # minerva_math tmp = pred_str.split("final answer is $", 1)[1] pred = tmp.split("$. I hope", 1)[0].strip() elif "boxed" in pred_str: ans = pred_str.split("boxed")[-1] if len(ans) == 0: return "" elif ans[0] == "{": stack = 1 a = "" for c in ans[1:]: if c == "{": stack += 1 a += c elif c == "}": stack -= 1 if stack == 0: break a += c else: a += c else: a = ans.split("$")[0].strip() pred = a elif "he answer is" in pred_str: pred = pred_str.split("he answer is")[-1].strip() elif "final answer is" in pred_str: pred = pred_str.split("final answer is")[-1].strip() elif "答案是" in pred_str: # Handle Chinese few-shot multiple choice problem answer extraction pred = pred_str.split("答案是")[1].strip().split("\n\n")[0].strip() else: # use the last number if use_last_number: pattern = "-?\d*\.?\d+" pred = re.findall(pattern, pred_str.replace(",", "")) if len(pred) >= 1: pred = pred[-1] else: pred = "" else: pred = "" # choice answer if ( data_name in ["sat_math", "aqua"] or "mmlu" in data_name ): tmp = re.findall(r"\b(A|B|C|D|E)\b", pred.upper()) if tmp: pred = tmp[-1] else: pred = pred.strip().strip(".") # multiple line # pred = pred.split("\n")[0] pred = re.sub(r"\n\s*", "", pred) if pred != "" and pred[0] == ":": pred = pred[1:] if pred != "" and pred[-1] == ".": pred = pred[:-1] if pred != "" and pred[-1] == "/": pred = pred[:-1] pred = strip_string(pred, skip_unit=data_name in ["carp_en", "minerva_math"]) return pred STRIP_EXCEPTIONS = ["carp_en", "minerva_math"] def parse_ground_truth(example: Dict[str, Any], data_name): if "gt_cot" in example and "gt" in example: if data_name in ["math"]: gt_ans = extract_answer(example["gt_cot"], data_name) elif data_name in STRIP_EXCEPTIONS: gt_ans = example["gt"] else: gt_ans = strip_string(example["gt"]) return example["gt_cot"], gt_ans # parse ground truth if data_name in ["math", "minerva_math"]: gt_cot = example["solution"] gt_ans = extract_answer(gt_cot, data_name) elif data_name == "gsm8k": gt_cot, gt_ans = example["answer"].split("####") elif data_name == "svamp": gt_cot, gt_ans = example["Equation"], example["Answer"] elif data_name == "asdiv": gt_cot = example["formula"] gt_ans = re.sub(r"\(.*?\)", "", example["answer"]) elif data_name == "mawps": gt_cot, gt_ans = None, example["target"] elif data_name == "tabmwp": gt_cot = example["solution"] gt_ans = example["answer"] if example["ans_type"] in ["integer_number", "decimal_number"]: if "/" in gt_ans: gt_ans = int(gt_ans.split("/")[0]) / int(gt_ans.split("/")[1]) elif "," in gt_ans: gt_ans = float(gt_ans.replace(",", "")) elif "%" in gt_ans: gt_ans = float(gt_ans.split("%")[0]) / 100 else: gt_ans = float(gt_ans) elif data_name == "carp_en": gt_cot, gt_ans = example["steps"], example["answer"] elif data_name == "mmlu_stem": abcd = "ABCD" gt_cot, gt_ans = None, abcd[example["answer"]] elif data_name == "sat_math": gt_cot, gt_ans = None, example["Answer"] elif data_name == "aqua": gt_cot, gt_ans = None, example["correct"] elif data_name in ["gaokao2023en", "college_math", "gaokao_math_cloze"]: gt_cot, gt_ans = None, example["answer"].replace("$", "").strip() elif data_name == "gaokao_math_qa": gt_cot, gt_ans = None, example["label"] elif data_name in ["gaokao2024_mix", "cn_middle_school"]: if len(example["choice_answer"]) > 0: gt_cot, gt_ans = None, example["choice_answer"] else: gt_cot, gt_ans = None, example["answer"] elif data_name == "olympiadbench": gt_cot, gt_ans = None, example["final_answer"][0].strip("$") elif data_name in [ "aime24", "amc23", "cmath", "gaokao2024_I", "gaokao2024_II", "imo2024", ]: gt_cot, gt_ans = None, example["answer"] else: raise NotImplementedError(f"`{data_name}`") # post process gt_cot = str(gt_cot).strip() if data_name not in STRIP_EXCEPTIONS: gt_ans = strip_string(gt_ans, skip_unit=data_name == "carp_en") else: gt_ans = ( gt_ans.replace("\\neq", "\\ne") .replace("\\leq", "\\le") .replace("\\geq", "\\ge") ) return gt_cot, gt_ans def parse_question(example, data_name): question = "" if data_name == "asdiv": question = f"{example['body'].strip()} {example['question'].strip()}" elif data_name == "svamp": body = example["Body"].strip() if not body.endswith("."): body = body + "." question = f'{body} {example["Question"].strip()}' elif data_name == "tabmwp": title_str = ( f'regarding "{example["table_title"]}" ' if example["table_title"] else "" ) question = f"Read the following table {title_str}and answer a question:\n" question += f'{example["table"]}\n{example["question"]}' if example["choices"]: question += ( f' Please select from the following options: {example["choices"]}' ) elif data_name == "carp_en": question = example["content"] elif data_name == "mmlu_stem": options = example["choices"] assert len(options) == 4 for i, (label, option) in enumerate(zip("ABCD", options)): options[i] = f"({label}) {str(option).strip()}" options = " ".join(options) # question = f"{example['question'].strip()}\nWhat of the following is the right choice? Explain your answer.\n{options}" question = f"{example['question'].strip()}\nAnswer Choices: {options}" elif data_name == "sat_math": options = example["options"].strip() assert "A" == options[0] options = "(" + options for ch in "BCD": if f" {ch}) " in options: options = regex.sub(f" {ch}\) ", f" ({ch}) ", options) # question = f"{example['question'].strip()}\nWhat of the following is the right choice? Explain your answer.\n{options.strip()}" question = f"{example['question'].strip()}\nAnswer Choices: {options}" elif "aqua" in data_name: options = example["options"] choice = "(" + "(".join(options) choice = choice.replace("(", " (").replace(")", ") ").strip() choice = "\nAnswer Choices: " + choice question = example["question"].strip() + choice elif data_name == "gaokao_math_qa": options_dict = example["options"] options = [] for key in options_dict: options.append(f"({key}) {options_dict[key]}") options = " ".join(options) question = f"{example['question'].strip()}\n选项: {options}" else: for key in ["question", "problem", "Question", "input"]: if key in example: question = example[key] break # assert question != "" # Yes or No question _, gt_ans = parse_ground_truth(example, data_name) if isinstance(gt_ans, str): gt_lower = gt_ans.lower() if gt_lower in ["true", "false"]: question += " (True or False)" if gt_lower in ["yes", "no"]: question += " (Yes or No)" return question.strip() def run_execute(executor, result, prompt_type, data_name, execute=False): if not result or result == "error": return None, None report = None if "program_only" in prompt_type: prediction = extract_program_output(result) elif prompt_type in ["pot", "pal"] and execute: code = extract_program(result) prediction, report = executor.apply(code) else: prediction = extract_answer(result, data_name) # prediction = strip_string(prediction, skip_unit=data_name == "carp_en") prediction = strip_string(prediction, skip_unit=data_name in STRIP_EXCEPTIONS) return prediction, report # trajcectory.py """ trajcectory: [ {"role": "rationale", "content": "..."}, {"role": "program", "content": "..."}, {"role": "output", "content": "..."}, {"role": "rationale", "content": "..."}, ... ] """ def text_to_trajectory(traj_str: str) -> None: """ """ # parse the above interleaved string of raionale, program, output, raionale, program, output, ... # output a list of dict trajectory = [] cur_role = "rationale" cur_content = "" # print(traj_str) for i, line in enumerate(traj_str.split("\n")): if line == "```python": # program begin assert cur_role == "rationale" if cur_content: trajectory.append({"role": cur_role, "content": cur_content}) cur_content = "" cur_role = "program" elif cur_role == "program" and line == "```": # program end assert cur_content trajectory.append({"role": cur_role, "content": cur_content}) cur_content = "" cur_role = "output" elif cur_role == "output" and line.startswith("```output"): # output begin assert cur_content == "" elif cur_role == "output" and line == "```": # output end trajectory.append({"role": cur_role, "content": cur_content}) cur_content = "" cur_role = "rationale" else: # content cur_content += line if i < len(traj_str.split("\n")) - 1: cur_content += "\n" # the last content if cur_content: trajectory.append({"role": cur_role, "content": cur_content}) return trajectory def trajectory_to_text(trajectory: list) -> str: text = "" for item in trajectory: content = item["content"] if item["role"] == "program": content = f"```python\n{content}```\n" elif item["role"] == "output": content = f"```output\n{content}```\n" text += content return text def is_execution_success(output): error_key_words = ["error", "exception", "no algorithms", "no algorithms", "cannot", "nan", "..."] success = all([k not in output.lower() for k in error_key_words]) return success def extract_program(text:str=None, trajectory:list=None, last_only=False) -> str: assert text is not None or trajectory is not None, "Either text or trajectory should be provided." if trajectory is None: try: trajectory = text_to_trajectory(text) except: return "raise ValueError('Invalid trajectory')" program_list = [] import_lines = [] for i, item in enumerate(trajectory): if item["role"] == "program": cur_program = item["content"] if i < len(trajectory) - 1: assert trajectory[i+1]["role"] == "output" output = trajectory[i+1]["content"].strip() if is_execution_success(output): program_list.append(cur_program) else: # extract import lines only for line in cur_program.split("\n"): if line.startswith("import") or line.startswith("from"): import_lines.append(line) else: program_list.append(cur_program) # add import lines to the first program if len(program_list) == 0: program_list.append("") if len(import_lines) > 0: program_list[0] = "\n".join(import_lines) + "\n" + program_list[0] for i, program in enumerate(program_list[:-1]): program_list[i] = "\n".join([line for line in program.split("\n") if not line.strip().startswith("print(")]) if last_only: program = program_list[-1] else: program = "\n".join(program_list) return program def extract_program_output(pred_str, last_only=True): """ extract output between ```output\n...\n```, use regex, there might be multiple outputs, each output may have multiple lines """ outputs = re.findall(r"```output\n(.*?)\n```", pred_str, re.DOTALL) if last_only: return outputs[-1] if len(outputs) > 0 else "" else: return outputs # grader.py def choice_answer_clean(pred: str): pred = pred.strip("\n").rstrip(".").rstrip("/").strip(" ").lstrip(":") # Clean the answer based on the dataset tmp = re.findall(r"\b(A|B|C|D|E)\b", pred.upper()) if tmp: pred = tmp else: pred = [pred.strip().strip(".")] pred = pred[-1] # Remove the period at the end, again! pred = pred.rstrip(".").rstrip("/") return pred def parse_digits(num): num = regex.sub(",", "", str(num)) try: return float(num) except: if num.endswith("%"): num = num[:-1] if num.endswith("\\"): num = num[:-1] try: return float(num) / 100 except: pass return None def is_digit(num): # paired with parse_digits return parse_digits(num) is not None def str_to_pmatrix(input_str): input_str = input_str.strip() matrix_str = re.findall(r"\{.*,.*\}", input_str) pmatrix_list = [] for m in matrix_str: m = m.strip("{}") pmatrix = r"\begin{pmatrix}" + m.replace(",", "\\") + r"\end{pmatrix}" pmatrix_list.append(pmatrix) return ", ".join(pmatrix_list) def math_equal( prediction: Union[bool, float, str], reference: Union[float, str], include_percentage: bool = True, is_close: bool = True, timeout: bool = False, ) -> bool: """ Exact match of math if and only if: 1. numerical equal: both can convert to float and are equal 2. symbolic equal: both can convert to sympy expression and are equal """ # print("Judge:", prediction, reference) if prediction is None or reference is None: return False if str(prediction.strip().lower()) == str(reference.strip().lower()): return True if ( reference in ["A", "B", "C", "D", "E"] and choice_answer_clean(prediction) == reference ): return True try: # 1. numerical equal if is_digit(prediction) and is_digit(reference): prediction = parse_digits(prediction) reference = parse_digits(reference) # number questions if include_percentage: gt_result = [reference / 100, reference, reference * 100] else: gt_result = [reference] for item in gt_result: try: if is_close: if numeric_equal(prediction, item): return True else: if item == prediction: return True except Exception: continue return False except: pass if not prediction and prediction not in [0, False]: return False # 2. symbolic equal reference = str(reference).strip() prediction = str(prediction).strip() ## pmatrix (amps) if "pmatrix" in prediction and not "pmatrix" in reference: reference = str_to_pmatrix(reference) ## deal with [], (), {} pred_str, ref_str = prediction, reference if ( prediction.startswith("[") and prediction.endswith("]") and not reference.startswith("(") ) or ( prediction.startswith("(") and prediction.endswith(")") and not reference.startswith("[") ): pred_str = pred_str.strip("[]()") ref_str = ref_str.strip("[]()") for s in ["{", "}", "(", ")"]: ref_str = ref_str.replace(s, "") pred_str = pred_str.replace(s, "") if pred_str.lower() == ref_str.lower(): return True ## [a, b] vs. [c, d], return a==c and b==d if ( regex.match(r"(\(|\[).+(\)|\])", prediction) is not None and regex.match(r"(\(|\[).+(\)|\])", reference) is not None ): pred_parts = prediction[1:-1].split(",") ref_parts = reference[1:-1].split(",") if len(pred_parts) == len(ref_parts): if all( [ math_equal( pred_parts[i], ref_parts[i], include_percentage, is_close ) for i in range(len(pred_parts)) ] ): return True if ( ( prediction.startswith("\\begin{pmatrix}") or prediction.startswith("\\begin{bmatrix}") ) and ( prediction.endswith("\\end{pmatrix}") or prediction.endswith("\\end{bmatrix}") ) and ( reference.startswith("\\begin{pmatrix}") or reference.startswith("\\begin{bmatrix}") ) and ( reference.endswith("\\end{pmatrix}") or reference.endswith("\\end{bmatrix}") ) ): pred_lines = [ line.strip() for line in prediction[ len("\\begin{pmatrix}") : -len("\\end{pmatrix}") ].split("\\\\") if line.strip() ] ref_lines = [ line.strip() for line in reference[ len("\\begin{pmatrix}") : -len("\\end{pmatrix}") ].split("\\\\") if line.strip() ] matched = True if len(pred_lines) == len(ref_lines): for pred_line, ref_line in zip(pred_lines, ref_lines): pred_parts = pred_line.split("&") ref_parts = ref_line.split("&") if len(pred_parts) == len(ref_parts): if not all( [ math_equal( pred_parts[i], ref_parts[i], include_percentage, is_close, ) for i in range(len(pred_parts)) ] ): matched = False break else: matched = False if not matched: break else: matched = False if matched: return True if prediction.count("=") == 1 and reference.count("=") == 1: pred = prediction.split("=") pred = f"{pred[0].strip()} - ({pred[1].strip()})" ref = reference.split("=") ref = f"{ref[0].strip()} - ({ref[1].strip()})" if symbolic_equal(pred, ref) or symbolic_equal(f"-({pred})", ref): return True elif ( prediction.count("=") == 1 and len(prediction.split("=")[0].strip()) <= 2 and "=" not in reference ): if math_equal( prediction.split("=")[1], reference, include_percentage, is_close ): return True elif ( reference.count("=") == 1 and len(reference.split("=")[0].strip()) <= 2 and "=" not in prediction ): if math_equal( prediction, reference.split("=")[1], include_percentage, is_close ): return True # symbolic equal with sympy if timeout: if call_with_timeout(symbolic_equal_process, prediction, reference): return True else: if symbolic_equal(prediction, reference): return True return False def math_equal_process(param): return math_equal(param[-2], param[-1]) def numeric_equal(prediction: float, reference: float): # Note that relative tolerance has significant impact # on the result of the synthesized GSM-Hard dataset # if reference.is_integer(): # return isclose(reference, round(prediction), abs_tol=1e-4) # else: # prediction = round(prediction, len(str(reference).split(".")[-1])) return isclose(reference, prediction, rel_tol=1e-4) def symbolic_equal(a, b): def _parse(s): for f in [parse_latex, parse_expr, latex2sympy]: try: return f(s.replace("\\\\", "\\")) except: try: return f(s) except: pass return s a = _parse(a) b = _parse(b) # direct equal try: if str(a) == str(b) or a == b: return True except: pass # simplify equal try: if a.equals(b) or simplify(a - b) == 0: return True except: pass # equation equal try: if (abs(a.lhs - a.rhs)).equals(abs(b.lhs - b.rhs)): return True except: pass try: if numeric_equal(float(N(a)), float(N(b))): return True except: pass # matrix try: # if a and b are matrix if a.shape == b.shape: _a = a.applyfunc(lambda x: round(x, 3)) _b = b.applyfunc(lambda x: round(x, 3)) if _a.equals(_b): return True except: pass return False def symbolic_equal_process(a, b, output_queue): result = symbolic_equal(a, b) output_queue.put(result) def call_with_timeout(func, *args, timeout=1, **kwargs): output_queue = multiprocessing.Queue() process_args = args + (output_queue,) process = multiprocessing.Process(target=func, args=process_args, kwargs=kwargs) process.start() process.join(timeout) if process.is_alive(): process.terminate() process.join() return False return output_queue.get() # evaluate.py def evaluate(data_name, prompt_type, samples: list=None, file_path: str=None, max_num_samples=None, execute=False): assert samples or file_path, "samples or file_path must be provided" if not samples: samples = list(load_jsonl(file_path)) if 'idx' in samples[0]: samples = {sample['idx']: sample for sample in samples}.values() samples = sorted(samples, key=lambda x: x['idx']) else: samples = [dict(idx=idx, **sample) for idx, sample in enumerate(samples)] if max_num_samples: print(f"max_num_samples: {max_num_samples} / {len(samples)}") samples = samples[:max_num_samples] # parse gt for sample in samples: sample['gt_cot'], sample['gt'] = parse_ground_truth(sample, data_name) params = [(idx, pred, sample['gt']) for idx, sample in enumerate(samples) for pred in sample['pred']] scores = [] timeout_cnt = 0 with ProcessPool(max_workers=16) as pool: future = pool.map(math_equal_process, params, timeout=10) iterator = future.result() with tqdm(total=len(samples), desc="Evaluate") as progress_bar: while True: try: result = next(iterator) scores.append(result) except StopIteration: break except TimeoutError as error: print(error) scores.append(False) timeout_cnt += 1 except Exception as error: print(error.traceback) exit() progress_bar.update(1) idx = 0 score_mat = [] for sample in samples: sample['score'] = scores[idx: idx+len(sample['pred'])] assert len(sample['score']) == len(sample['pred']) score_mat.append(sample['score']) idx += len(sample['pred']) max_len = max([len(s) for s in score_mat]) for i, s in enumerate(score_mat): if len(s) < max_len: score_mat[i] = s + [s[-1]] * (max_len - len(s)) # pad # output mean of each column of scores col_means= np.array(score_mat).mean(axis=0) mean_score = list(np.round(col_means * 100, decimals=1)) result_json = { "num_samples": len(samples), "num_scores": len(scores), "timeout_samples": timeout_cnt, "empty_samples": len([s for s in samples if not s['pred'][-1]]), "acc": mean_score[0] } # each type score if "type" in samples[0]: type_scores = {} for sample in samples: if sample['type'] not in type_scores: type_scores[sample['type']] = [] type_scores[sample['type']].append(sample['score'][-1]) type_scores = {k: np.round(np.array(v).mean() * 100, decimals=1) for k, v in type_scores.items()} type_scores = {k: v for k, v in sorted(type_scores.items(), key=lambda item: item[0])} result_json['type_acc'] = type_scores return samples, result_json # python_executor.py class GenericRuntime: GLOBAL_DICT = {} LOCAL_DICT = None HEADERS = [] def __init__(self): self._global_vars = copy.copy(self.GLOBAL_DICT) self._local_vars = copy.copy(self.LOCAL_DICT) if self.LOCAL_DICT else None for c in self.HEADERS: self.exec_code(c) def exec_code(self, code_piece: str) -> None: if regex.search(r'(\s|^)?input\(', code_piece): # regex.search(r'(\s|^)?os.', code_piece): raise RuntimeError() exec(code_piece, self._global_vars) # TODO: use: https://github.com/shroominic/codebox-api # @high safe exec in sandbox # byte_code = compile_restricted( # code_piece, # filename='', # mode='exec' # ) # print("global vars:", self._global_vars) # _print_ = PrintCollector # exec(byte_code, {'__builtins__': utility_builtins}, None) def eval_code(self, expr: str) -> Any: return eval(expr, self._global_vars) def inject(self, var_dict: Dict[str, Any]) -> None: for k, v in var_dict.items(): self._global_vars[k] = v @property def answer(self): return self._global_vars['answer'] class DateRuntime(GenericRuntime): GLOBAL_DICT = { 'datetime': datetime.datetime, 'timedelta': dateutil.relativedelta.relativedelta, 'relativedelta': dateutil.relativedelta.relativedelta } class CustomDict(dict): def __iter__(self): return list(super().__iter__()).__iter__() class ColorObjectRuntime(GenericRuntime): GLOBAL_DICT = {'dict': CustomDict} class PythonExecutor: def __init__( self, runtime: Optional[Any] = None, get_answer_symbol: Optional[str] = None, get_answer_expr: Optional[str] = None, get_answer_from_stdout: bool = False, timeout_length: int = 5, ) -> None: self.runtime = runtime if runtime else GenericRuntime() self.answer_symbol = get_answer_symbol self.answer_expr = get_answer_expr self.get_answer_from_stdout = get_answer_from_stdout self.pool = Pool(multiprocess.cpu_count()) self.timeout_length = timeout_length def process_generation_to_code(self, gens: str): return [g.strip().split('\n') for g in gens] @staticmethod def execute( code, get_answer_from_stdout = None, runtime = None, answer_symbol = None, answer_expr = None, timeout_length = 10, auto_mode=False ): try: if auto_mode: if "print(" in code[-1]: program_io = io.StringIO() with redirect_stdout(program_io): timeout(timeout_length)(runtime.exec_code)('\n'.join(code)) program_io.seek(0) result = program_io.read() else: print(code) timeout(timeout_length)(runtime.exec_code)('\n'.join(code[:-1])) result = timeout(timeout_length)(runtime.eval_code)(code[-1]) else: if get_answer_from_stdout: program_io = io.StringIO() with redirect_stdout(program_io): timeout(timeout_length)(runtime.exec_code)('\n'.join(code)) program_io.seek(0) result = program_io.read() elif answer_symbol: timeout(timeout_length)(runtime.exec_code)('\n'.join(code)) result = runtime._global_vars[answer_symbol] elif answer_expr: timeout(timeout_length)(runtime.exec_code)('\n'.join(code)) result = timeout(timeout_length)(runtime.eval_code)(answer_expr) else: timeout(timeout_length)(runtime.exec_code)('\n'.join(code[:-1])) result = timeout(timeout_length)(runtime.eval_code)(code[-1]) report = "Done" str(result) pickle.dumps(result) # serialization check except: result = '' report = traceback.format_exc().split('\n')[-2] return result, report def apply(self, code): return self.batch_apply([code])[0] @staticmethod def truncate(s, max_length=400): half = max_length // 2 if len(s) > max_length: s = s[:half] + "..." + s[-half:] return s def batch_apply(self, batch_code): all_code_snippets = self.process_generation_to_code(batch_code) timeout_cnt = 0 all_exec_results = [] # with ProcessPool(max_workers=min(len(all_code_snippets), os.cpu_count())) as pool: with ProcessPool(max_workers=min(len(all_code_snippets), 1)) as pool: executor = partial( self.execute, get_answer_from_stdout=self.get_answer_from_stdout, runtime=self.runtime, answer_symbol=self.answer_symbol, answer_expr=self.answer_expr, timeout_length=self.timeout_length, # this timeout not work auto_mode=True ) future = pool.map(executor, all_code_snippets, timeout=self.timeout_length) iterator = future.result() if len(all_code_snippets) > 100: progress_bar = tqdm(total=len(all_code_snippets), desc="Execute") else: progress_bar = None while True: try: result = next(iterator) all_exec_results.append(result) except StopIteration: break except TimeoutError as error: print(error) all_exec_results.append(("", "Timeout Error")) timeout_cnt += 1 except Exception as error: print(error) exit() if progress_bar is not None: progress_bar.update(1) if progress_bar is not None: progress_bar.close() batch_results = [] for code, (res, report) in zip(all_code_snippets, all_exec_results): # post processing res, report = str(res).strip(), str(report).strip() res, report = self.truncate(res), self.truncate(report) batch_results.append((res, report)) return batch_results